Improve API notebook example (#70)

* add data config path to the server cli * Improve api example and set more restrictive pinning on cupy * Drop `DataConfig` in favor of `MainConfig` * Update tests for ultrack server cli * fix error of pickling a local object on Windows CI * minor changes to notebook text --------- Co-authored-by: Jordao Bragantini <[email protected]>
royerlab · Jun 27, 2024 · 8f10963 · 8f10963
1 parent 4d200ca
commit 8f10963
Show file tree

Hide file tree

Showing 11 changed files with 270 additions and 151,170 deletions.
diff --git a/examples/api/api_example.ipynb b/examples/api/api_example.ipynb
diff --git a/examples/api/environment_gpu.yml b/examples/api/environment_gpu.yml
@@ -7,7 +7,7 @@ channels:
 dependencies:
   - coin-or-cbc
   - cucim
-  - cupy!=13.0.*
+  - cupy!=13.*
   - gurobi
   - jupyter
   - pip

diff --git a/examples/flow_field_3d/environment_gpu.yml b/examples/flow_field_3d/environment_gpu.yml
@@ -9,7 +9,7 @@ channels:
 dependencies:
   - coin-or-cbc
   - cucim
-  - cupy!=13.0.*
+  - cupy!=13.*
   - gurobi
   - jupyter
   - napari

diff --git a/examples/micro_sam/environment_gpu.yml b/examples/micro_sam/environment_gpu.yml
@@ -9,7 +9,7 @@ channels:
 dependencies:
   - coin-or-cbc
   - cucim
-  - cupy!=13.0.*
+  - cupy!=13.*
   - gurobi
   - jupyter
   - pip

diff --git a/examples/multi_color_ensemble/environment_gpu.yml b/examples/multi_color_ensemble/environment_gpu.yml
@@ -8,7 +8,7 @@ dependencies:
   - cellpose
   - coin-or-cbc
   - cucim
-  - cupy!=13.0.*
+  - cupy!=13.*
   - gurobi
   - jupyter
   - pip

diff --git a/examples/neuromast_plantseg/environment_gpu.yml b/examples/neuromast_plantseg/environment_gpu.yml
@@ -9,7 +9,7 @@ channels:
 dependencies:
   - coin-or-cbc
   - cucim
-  - cupy!=13.0.*
+  - cupy!=13.*
   - gurobi
   - jupyter
   - pyqt

diff --git a/examples/stardist_2d/environment_gpu.yml b/examples/stardist_2d/environment_gpu.yml
@@ -10,7 +10,7 @@ dependencies:
   - conda-forge::cudatoolkit=11.2
   - conda-forge::cudnn=8.4.1
   - conda-forge::tensorflow-gpu=2.11.0
-  - conda-forge::cupy!=13.0.*
+  - conda-forge::cupy!=13.*
   - coin-or-cbc
   - gurobi
   - jupyter

diff --git a/examples/zebrahub/environment_gpu.yml b/examples/zebrahub/environment_gpu.yml
@@ -7,7 +7,7 @@ channels:
 dependencies:
   - coin-or-cbc
   - cucim
-  - cupy!=13.0.*
+  - cupy!=13.*
   - gurobi
   - jupyter
   - pip

diff --git a/ultrack/api/main.py b/ultrack/api/main.py
@@ -1,13 +1,12 @@
 import os
 from multiprocessing import Process
 from pathlib import Path
-from threading import Thread
 from typing import Union
 
 import uvicorn
 
+from ultrack import MainConfig
 from ultrack.api import app
-from ultrack.config import DataConfig
 
 
 def _in_notebook():
@@ -23,7 +22,7 @@ def _in_notebook():
 
 def start_server(
     api_results_path: Union[Path, str, None] = None,
-    ultrack_data_config: Union[DataConfig, None] = None,
+    ultrack_data_config: Union[MainConfig, None] = None,
     host: str = "0.0.0.0",
     port: int = 8000,
 ):

diff --git a/ultrack/cli/_test/test_cli.py b/ultrack/cli/_test/test_cli.py
@@ -13,6 +13,9 @@
 from ultrack.utils.data import make_config_content
 
 
+def _run_server(instance_config_path: str):
+    _run_command(["server", "--port", "54123", "-cfg", instance_config_path])
+
 def _run_command(command_and_args: List[str]) -> None:
     try:
         main(command_and_args)
@@ -161,6 +164,25 @@ def test_clear_database(self, instance_config_path: str, mode: str) -> None:
             ]
         )
 
+    def test_server(self, instance_config_path: str) -> None:
+        # Start server in a background thread
+        process = Process(target=_run_server, args=(instance_config_path,))
+        process.start()
+
+        # Wait for server to start
+        import time
+
+        time.sleep(10)
+
+        response = requests.get("http://127.0.0.1:54123")
+        print(response.content)
+
+        assert process.is_alive()
+        assert response.status_code == 200
+
+        process.terminate()
+        process.join()
+
 
 def test_create_config(tmp_path: Path) -> None:
     _run_command(["create_config", str(tmp_path / "config.toml")])
@@ -178,28 +200,3 @@ def test_labels_to_contours(zarr_dataset_paths: List[str], tmp_path: Path) -> No
 
 def test_check_gurobi() -> None:
     _run_command(["check_gurobi"])
-
-
-def _run_server():
-    _run_command(["server", "--port", "54123"])
-
-
-def test_server() -> None:
-
-    # Start server in a background thread
-    process = Process(target=_run_server)
-    process.start()
-
-    # Wait for server to start
-    import time
-
-    time.sleep(10)
-
-    response = requests.get("http://127.0.0.1:54123")
-    print(response.content)
-
-    assert process.is_alive()
-    assert response.status_code == 200
-
-    process.terminate()
-    process.join()
diff --git a/ultrack/cli/server.py b/ultrack/cli/server.py
@@ -1,9 +1,12 @@
-import os
+from typing import Optional
 
 import click
-import uvicorn
+import toml
 
-from ultrack.api.app import app
+from ultrack import MainConfig
+from ultrack.api import start_server
+from ultrack.cli.utils import config_option
+from ultrack.config import DataConfig
 
 
 @click.command("server")
@@ -22,9 +25,9 @@
     default=None,
     show_default=True,
 )
-def server_cli(host: str, port: int, api_results_path: str) -> None:
+@config_option()
+def server_cli(
+    host: str, port: int, api_results_path: Optional[str], config: MainConfig
+) -> None:
     """Start the websockets ultrack API."""
-    if api_results_path is not None:
-        os.environ["API_RESULTS_PATH"] = str(api_results_path)
-
-    uvicorn.run(app, host=host, port=port)
+    start_server(api_results_path, ultrack_data_config=config, host=host, port=port)