Add initial executor CLI (#72)

* Add initial executor CLI

* Add simple tests

* Add run tests

* Clean-up launch CLI options

* Fix tests

* change worker default to 1, add from file

Co-authored-by: Josh Horton <[email protected]>

devtools/conda-envs/test-env.yaml

@@ -19,6 +19,7 @@ dependencies:
   - tqdm
   - rich
   - click
+  - click-option-group
   - rdkit
   - openff-utilities
   - openff-toolkit-base >=0.10.0

openff/bespokefit/__init__.py

@@ -2,10 +2,23 @@
 BespokeFit
 Creating bespoke parameters for individual molecules.
 """
+import logging
+import sys
 
 from ._version import get_versions
 
 versions = get_versions()
 __version__ = versions["version"]
 __git_revision__ = versions["full-revisionid"]
 del get_versions, versions
+
+
+# Silence verbose messages when running the CLI otherwise you can't read the output
+# without seeing tens of 'Unable to load AmberTools' or don't import simtk warnings...
+if sys.argv[0].endswith("openff-bespoke"):
+
+    from openff.bespokefit.utilities.logging import DeprecationWarningFilter
+
+    # if "openff-bespoke"
+    logging.getLogger("openff.toolkit").setLevel(logging.ERROR)
+    logging.getLogger().addFilter(DeprecationWarningFilter())

openff/bespokefit/cli/cli.py

@@ -1,5 +1,6 @@
 import click
 
+from openff.bespokefit.cli.executor import executor_cli
 from openff.bespokefit.cli.prepare import prepare_cli
 
 
@@ -8,4 +9,5 @@ def cli():
     """The root group for all CLI commands."""
 
 
+cli.add_command(executor_cli)
 cli.add_command(prepare_cli)

openff/bespokefit/cli/executor/__init__.py

@@ -0,0 +1,3 @@
+from openff.bespokefit.cli.executor.executor import executor_cli
+
+__all__ = [executor_cli]

openff/bespokefit/cli/executor/executor.py

@@ -0,0 +1,19 @@
+import click
+
+from openff.bespokefit.cli.executor.launch import launch_cli
+from openff.bespokefit.cli.executor.list import list_cli
+from openff.bespokefit.cli.executor.run import run_cli
+from openff.bespokefit.cli.executor.submit import submit_cli
+from openff.bespokefit.cli.executor.watch import watch_cli
+
+
+@click.group("executor")
+def executor_cli():
+    """Commands for interacting with a bespoke executor."""
+
+
+executor_cli.add_command(launch_cli)
+executor_cli.add_command(submit_cli)
+executor_cli.add_command(run_cli)
+executor_cli.add_command(watch_cli)
+executor_cli.add_command(list_cli)

openff/bespokefit/cli/executor/launch.py

@@ -0,0 +1,114 @@
+import time
+from typing import Optional
+
+import click
+import rich
+from click_option_group import optgroup
+from rich import pretty
+
+from openff.bespokefit.cli.utilities import create_command, print_header
+
+
+# The run command inherits these options so be sure to take that into account when
+# making changes here.
+def launch_options(
+    directory: str = "bespoke-executor",
+    n_fragmenter_workers: Optional[int] = 1,
+    n_qc_compute_workers: Optional[int] = 1,
+    n_optimizer_workers: Optional[int] = 1,
+    launch_redis_if_unavailable: Optional[bool] = True,
+):
+
+    return [
+        optgroup("Executor configuration"),
+        optgroup.option(
+            "--directory",
+            type=click.Path(exists=False, file_okay=False, dir_okay=True),
+            help="The directory to store any working and log files in",
+            required=True,
+            default=directory,
+            show_default=directory is not None,
+        ),
+        optgroup.group("Worker configuration"),
+        optgroup.option(
+            "--n-fragmenter-workers",
+            "n_fragmenter_workers",
+            type=click.INT,
+            help="The number of fragmentation workers to spawn",
+            required=n_fragmenter_workers is None,
+            default=n_fragmenter_workers,
+            show_default=n_fragmenter_workers is not None,
+        ),
+        optgroup.option(
+            "--n-qc-compute-workers",
+            "n_qc_compute_workers",
+            type=click.INT,
+            help="The number of QC compute workers to spawn",
+            required=n_qc_compute_workers is None,
+            default=n_qc_compute_workers,
+            show_default=n_qc_compute_workers is not None,
+        ),
+        optgroup.option(
+            "--n-optimizer-workers",
+            "n_optimizer_workers",
+            type=click.INT,
+            help="The number of optimizer workers to spawn",
+            required=n_optimizer_workers is None,
+            default=n_optimizer_workers,
+            show_default=n_optimizer_workers is not None,
+        ),
+        optgroup.group("Storage configuration"),
+        optgroup.option(
+            "--launch-redis/--no-launch-redis",
+            "launch_redis_if_unavailable",
+            help="Whether to launch a redis server if an already running one cannot be "
+            "found.",
+            required=launch_redis_if_unavailable is None,
+            default=launch_redis_if_unavailable,
+            show_default=launch_redis_if_unavailable is not None,
+        ),
+    ]
+
+
+def _launch_cli(
+    directory: str,
+    n_fragmenter_workers: int,
+    n_qc_compute_workers: int,
+    n_optimizer_workers: int,
+    launch_redis_if_unavailable: bool,
+):
+    """Launch a bespoke executor."""
+
+    pretty.install()
+
+    console = rich.get_console()
+    print_header(console)
+
+    from openff.bespokefit.executor import BespokeExecutor
+
+    executor_status = console.status("launching the bespoke executor")
+    executor_status.start()
+
+    with BespokeExecutor(
+        directory=directory,
+        n_fragmenter_workers=n_fragmenter_workers,
+        n_qc_compute_workers=n_qc_compute_workers,
+        n_optimizer_workers=n_optimizer_workers,
+        launch_redis_if_unavailable=launch_redis_if_unavailable,
+    ):
+
+        executor_status.stop()
+        console.print("[[green]✓[/green]] bespoke executor launched")
+
+        try:
+            while True:
+                time.sleep(5)
+        except KeyboardInterrupt:
+            pass
+
+
+launch_cli = create_command(
+    click_command=click.command("launch"),
+    click_options=launch_options(),
+    func=_launch_cli,
+)

openff/bespokefit/cli/executor/list.py

@@ -0,0 +1,51 @@
+import click
+import requests
+import rich
+from rich import pretty
+from rich.padding import Padding
+
+from openff.bespokefit.cli.utilities import print_header
+
+
+@click.command("list")
+def list_cli():
+    """List the ids of any bespoke optimizations."""
+
+    pretty.install()
+
+    console = rich.get_console()
+    print_header(console)
+
+    from openff.bespokefit.executor.services import settings
+    from openff.bespokefit.executor.services.coordinator.models import (
+        CoordinatorGETPageResponse,
+    )
+
+    href = (
+        f"http://127.0.0.1:"
+        f"{settings.BEFLOW_GATEWAY_PORT}"
+        f"{settings.BEFLOW_API_V1_STR}/"
+        f"{settings.BEFLOW_COORDINATOR_PREFIX}"
+    )
+
+    try:
+
+        request = requests.get(href)
+        request.raise_for_status()
+
+    except requests.ConnectionError:
+        console.print(
+            "A connection could not be made to the bespoke executor. Please make sure "
+            "there is a bespoke executor running."
+        )
+        return
+
+    response = CoordinatorGETPageResponse.parse_raw(request.content)
+    response_ids = [item.id for item in response.contents]
+
+    if len(response_ids) == 0:
+        console.print("No optimizations were found.")
+        return
+
+    console.print(Padding("The following optimizations were found:", (0, 0, 1, 0)))
+    console.print("\n".join(response_ids))

openff/bespokefit/cli/executor/run.py

@@ -0,0 +1,92 @@
+from typing import Optional
+
+import click
+import rich
+from rich import pretty
+from rich.padding import Padding
+
+from openff.bespokefit.cli.executor.launch import launch_options
+from openff.bespokefit.cli.executor.submit import _submit, submit_options
+from openff.bespokefit.cli.utilities import create_command, print_header
+
+
+def _run_cli(
+    input_file_path: str,
+    output_file_path: str,
+    force_field_path: str,
+    spec_name: Optional[str],
+    spec_file_name: Optional[str],
+    directory: str,
+    n_fragmenter_workers: int,
+    n_qc_compute_workers: int,
+    n_optimizer_workers: int,
+    launch_redis_if_unavailable: bool,
+):
+    """Run bespoke optimization using a temporary executor.
+
+    If you are running many bespoke optimizations it is recommended that you first launch
+    a bespoke executor using the `launch` command and then submit the optimizations to it
+    using the `submit` command.
+    """
+
+    pretty.install()
+
+    console = rich.get_console()
+    print_header(console)
+
+    from openff.bespokefit.executor import BespokeExecutor, wait_until_complete
+
+    executor_status = console.status("launching the bespoke executor")
+    executor_status.start()
+
+    with BespokeExecutor(
+        directory=directory,
+        n_fragmenter_workers=n_fragmenter_workers,
+        n_qc_compute_workers=n_qc_compute_workers,
+        n_optimizer_workers=n_optimizer_workers,
+        launch_redis_if_unavailable=launch_redis_if_unavailable,
+    ):
+
+        executor_status.stop()
+        console.print("[[green]✓[/green]] bespoke executor launched")
+        console.line()
+
+        response = _submit(
+            console,
+            input_file_path,
+            force_field_path,
+            spec_name,
+            spec_file_name,
+        )
+
+        if response is None:
+            return
+
+        console.print(Padding("3. running the fitting pipeline", (1, 0, 1, 0)))
+
+        results = wait_until_complete(response.id)
+
+        if results is None:
+            return
+
+        with open(output_file_path, "w") as file:
+            file.write(results.json())
+
+
+__run_options = [*submit_options()]
+__run_options.insert(
+    1,
+    click.option(
+        "--output",
+        "output_file_path",
+        type=click.Path(exists=False, file_okay=True, dir_okay=False),
+        help="The JSON file to save the results to",
+        default="output.json",
+        show_default=True,
+    ),
+)
+__run_options.extend(launch_options())
+
+run_cli = create_command(
+    click_command=click.command("run"), click_options=__run_options, func=_run_cli
+)