feat: add cli annotation subcommand (#2539)

add `cellxgene annotate` subcommand for invoking MLflow model to generate new `obs` annotations, initially intended for cell type annotations.

test/unit/cli/mlflow_model_fixture.py

@@ -0,0 +1,20 @@
+import shutil
+from tempfile import TemporaryDirectory, mkstemp
+
+import mlflow
+
+
+def write_model(model) -> str:
+    with TemporaryDirectory() as mlflow_model_dir:
+        mlflow.pyfunc.save_model(mlflow_model_dir, python_model=model)
+        return shutil.make_archive(mkstemp()[1], "zip", mlflow_model_dir)
+
+
+class FakeModel(mlflow.pyfunc.PythonModel):
+    def __init__(self, input_to_output: dict = {}):
+        self.input_to_output = input_to_output
+
+    def predict(self, context, model_input) -> None:
+        # this stdout output is useful for validating the input in a test, noting that this model will be invoked in a
+        # subprocess, so stdout is one means of communicating information back to the test code
+        print(f"__MODEL_INPUT__={model_input.iloc[0][0]}")

test/unit/cli/test_annotate.py

@@ -0,0 +1,97 @@
+import unittest
+from tempfile import mkstemp
+
+from click.testing import CliRunner
+
+from server.cli.annotate import annotate
+from test.unit.cli.mlflow_model_fixture import FakeModel, write_model
+
+
+class TestCliAnnotate(unittest.TestCase):
+    def test__annotate__loads_and_runs(self):
+        """
+        Invokes the `annotate` subcommand of cellxgene CLI, using a CliRunner() programmatic invocation.
+
+        This tests the happy path case:
+        1) Command line options are parsed;
+        2) An MLflow model zip archive can be read in (from local disk), unpacked, and invoked;
+        3) The correct options are passed to the MLflow model.
+        4) The annotate subcommand exits successfully.
+
+        This does not verify model output or predictions (it's a fake MLflow model, after all); it's up to the real model
+        to output its predictions as it wants, but this is specific to the model and so not tested here.
+
+        The  CliRunner() invokes the subcommand in a subprocess, and the annotate subcommand itself invokes the MLflow
+        model in yet another subprocess. So while this test can help determine if everything is working, it is not a
+        simple matter to debug in the case of a failure. However, the stdout/stderr of the MLflow process is captured
+        by the CliRunner() subprocess, so errors can be inspected in result.stdout when debugging this test. Hope this
+        helps!
+        """
+
+        _, query_dataset_file_path = mkstemp()
+        model = FakeModel()
+        model_file_path = write_model(model)
+
+        result = CliRunner().invoke(
+            annotate,
+            [
+                "--input-h5ad-file",
+                query_dataset_file_path,
+                "--model-url",
+                model_file_path,
+                "--output-h5ad-file",
+                f"{query_dataset_file_path}.output",
+            ],
+        )
+
+        # to help debugging, show the output from the CliRunner and MLflow stdout
+        if result.exit_code:
+            print(result.stdout)
+
+        self.assertEqual(0, result.exit_code, "runs successfully")
+
+        # The FakeModel will print it inputs to stdout, as "__MODEL_INPUT__={...}", allowing us to assert that it received valid inputs.
+        self.assertIn(
+            "__MODEL_INPUT__={"
+            f'"query_dataset_h5ad_path": "{query_dataset_file_path}", '
+            f'"output_h5ad_path": "{query_dataset_file_path}.output", '
+            '"annotation_prefix": "cxg_cell_type", "classifier": "default", '
+            '"organism": "Homo sapiens", "use_gpu": true}',
+            result.stdout,
+            "inputs passed correctly",
+        )
+
+    def test__annotate__verifies_mutually_exclusive_options(self):
+        required_options = ["--input-h5ad-file", "some.h5ad", "--model-url", "some_url"]
+        result = CliRunner().invoke(
+            annotate,
+            required_options + [],
+        )
+
+        self.assertNotEqual(0, result.exit_code, "aborts with non-success code")
+        self.assertIn(
+            "--update_h5ad_file or --output_h5ad_file must be specified",
+            result.stdout,
+            "error message displayed",
+        )
+
+        result = CliRunner().invoke(
+            annotate, required_options + ["--output-h5ad-file", "some_arg", "--update-h5ad-file"]
+        )
+
+        self.assertNotEqual(0, result.exit_code, "aborts with non-success code")
+        self.assertIn(
+            "--update_h5ad_file and --output_h5ad_file are mutually exclusive",
+            result.stdout,
+            "error message displayed",
+        )
+
+
+# TODO:
+# Test annotate cli args more comprehensively
+# Test server.cli.annotate._validate_options
+# Test model caching feature works
+# Test model loading from s3 works (maybe w/just a real model)
+
+if __name__ == "__main__":
+    unittest.main()