Send tasks that require the pipelines to same worker

Author: gabegma

azimuth/modules/base_classes/dask_module.py

@@ -15,7 +15,7 @@
 from datasets import Dataset
 from distributed import Client, Event, Future, rejoin, secede
 
-from azimuth.config import CommonFieldsConfig
+from azimuth.config import CommonFieldsConfig, ModelContractConfig
 from azimuth.modules.base_classes.caching import HDF5CacheMixin
 from azimuth.types import DatasetSplitName, ModuleResponse
 from azimuth.utils.logs import TimerLogging
@@ -94,6 +94,11 @@ def _get_config_scope(self, config) -> ConfigScope:
             scoped_config = base.__args__[0]
         return cast(ConfigScope, scoped_config.parse_obj(config.dict(by_alias=True)))
 
+    @property
+    def can_load_model(self) -> bool:
+        # TODO Not all modules that inherit from ModelContractConfig load the model. Smarter way?
+        return isinstance(self.config, ModelContractConfig)
+
     def start_task_on_dataset_split(
         self, client: Client, dependencies: List["DaskModule"] = None
     ) -> "DaskModule":
@@ -118,6 +123,7 @@ def start_task_on_dataset_split(
             pure=False,
             dependencies=deps,
             key=f"{self.task_id}_{uuid.uuid4()}",  # Unique identifier
+            workers=0 if self.can_load_model else 1,
         )
         # Tell that this future is used on which indices.
         self.future.indices = self.get_caching_indices()

azimuth/utils/cluster.py

@@ -40,7 +40,7 @@ def default_cluster(large=False) -> SpecCluster:
     tmp_file = pjoin(str(tempfile.mkdtemp()), "dask-worker-space")
     with dask.config.set({"distributed.worker.daemon": False}):
         cluster = distributed.LocalCluster(
-            n_workers=2,
+            n_workers=2,  # Assignment to workers is hard-coded, so it needs to stay 2.
             local_directory=tmp_file,
             threads_per_worker=1,
             memory_limit=memory_limit,  # "auto" doesnt work well.

tests/test_routers/conftest.py

@@ -8,7 +8,7 @@
 from fastapi import FastAPI
 from starlette.testclient import TestClient
 
-import azimuth.app as me_app
+import azimuth.app as azimuth_app
 from azimuth.app import get_ready_flag
 from azimuth.config import AzimuthConfig
 from tests.utils import DATASET_CFG, SIMPLE_PERTURBATION_TESTING_CONFIG
@@ -24,7 +24,7 @@ def is_set(self):
 
 def create_test_app(config) -> FastAPI:
     json.dump(config.dict(by_alias=True), open("/tmp/config.json", "w"))
-    return me_app.start_app("/tmp/config.json", load_config_history=False, debug=False)
+    return azimuth_app.start_app("/tmp/config.json", load_config_history=False, debug=False)
 
 
 FAST_TEST_CFG = {
@@ -44,7 +44,7 @@ def wait_for_startup_after(app):
     while resp.json()["startupTasksReady"] is not True:
         time.sleep(1)
         resp = client.get("/status")
-    task_manager = me_app.get_task_manager()
+    task_manager = azimuth_app.get_task_manager()
     while task_manager.is_locked:
         time.sleep(1)
 
@@ -70,7 +70,7 @@ def app() -> FastAPI:
     while resp.json()["startupTasksReady"] is not True:
         time.sleep(1)
         resp = client.get("/status")
-    task_manager = me_app.get_task_manager()
+    task_manager = azimuth_app.get_task_manager()
     while task_manager.is_locked:
         time.sleep(1)
     yield _app
@@ -79,7 +79,7 @@ def app() -> FastAPI:
 @pytest.fixture(scope="function")
 def app_not_started(app) -> FastAPI:
 
-    startup_tasks = me_app.get_startup_tasks()
+    startup_tasks = azimuth_app.get_startup_tasks()
 
     class ModuleThatWillNeverEnd:
         def status(self):