refactor

Author: sanderegg

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/__init__.py

@@ -5,7 +5,7 @@
 from fastapi import FastAPI
 from servicelib.logging_utils import log_context
 
-from ._constants import MODULE_NAME
+from ._constants import MODULE_NAME_SCHEDULER
 from ._manager import run_new_pipeline, setup_manager, shutdown_manager, stop_pipeline
 from ._worker import setup_worker, shutdown_worker
 
@@ -14,7 +14,9 @@
 
 def on_app_startup(app: FastAPI) -> Callable[[], Coroutine[Any, Any, None]]:
     async def start_scheduler() -> None:
-        with log_context(_logger, level=logging.INFO, msg=f"starting {MODULE_NAME}"):
+        with log_context(
+            _logger, level=logging.INFO, msg=f"starting {MODULE_NAME_SCHEDULER}"
+        ):
             await setup_worker(app)
             await setup_manager(app)
 
@@ -23,7 +25,9 @@ async def start_scheduler() -> None:
 
 def on_app_shutdown(app: FastAPI) -> Callable[[], Coroutine[Any, Any, None]]:
     async def stop_scheduler() -> None:
-        with log_context(_logger, level=logging.INFO, msg=f"stopping {MODULE_NAME}"):
+        with log_context(
+            _logger, level=logging.INFO, msg=f"stopping {MODULE_NAME_SCHEDULER}"
+        ):
             await shutdown_manager(app)
             await shutdown_worker(app)
 

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_constants.py

@@ -1,7 +1,7 @@
 import datetime
 from typing import Final
 
-MODULE_NAME: Final[str] = "computational-distributed-scheduler"
+MODULE_NAME_SCHEDULER: Final[str] = "computational-distributed-scheduler"
 MODULE_NAME_WORKER: Final[str] = "computational-distributed-worker"
 SCHEDULER_INTERVAL: Final[datetime.timedelta] = datetime.timedelta(seconds=5)
 MAX_CONCURRENT_PIPELINE_SCHEDULING: Final[int] = 10

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py

@@ -8,25 +8,22 @@
 from models_library.users import UserID
 from servicelib.background_task import start_periodic_task, stop_periodic_task
 from servicelib.logging_utils import log_context
-from servicelib.redis import RedisClientSDK
 from servicelib.redis_utils import exclusive
 from servicelib.utils import limited_gather
-from settings_library.redis import RedisDatabase
 
 from ...models.comp_runs import RunMetadataDict
 from ...utils.rabbitmq import publish_project_log
 from ..db import get_db_engine
 from ..db.repositories.comp_pipelines import CompPipelinesRepository
 from ..db.repositories.comp_runs import CompRunsRepository
 from ..rabbitmq import get_rabbitmq_client
-from ..redis import get_redis_client_manager
 from ._constants import (
     MAX_CONCURRENT_PIPELINE_SCHEDULING,
-    MODULE_NAME,
+    MODULE_NAME_SCHEDULER,
     SCHEDULER_INTERVAL,
 )
 from ._publisher import request_pipeline_scheduling
-from ._utils import SCHEDULED_STATES
+from ._utils import SCHEDULED_STATES, get_redis_client_from_app, get_redis_lock_key
 
 _logger = logging.getLogger(__name__)
 
@@ -99,28 +96,13 @@ async def stop_pipeline(
     if updated_comp_run:
         # ensure the scheduler starts right away
         rabbitmq_client = get_rabbitmq_client(app)
-        await request_pipeline_scheduling(updated_comp_run, rabbitmq_client, db_engine)
-
-
-def _get_app_from_args(*args, **kwargs) -> FastAPI:
-    assert kwargs is not None  # nosec
-    if args:
-        app = args[0]
-    else:
-        assert "app" in kwargs  # nosec
-        app = kwargs["app"]
-    assert isinstance(app, FastAPI)  # nosec
-    return app
-
-
-def _redis_client_getter(*args, **kwargs) -> RedisClientSDK:
-    app = _get_app_from_args(*args, **kwargs)
-    return get_redis_client_manager(app).client(RedisDatabase.LOCKS)
-
-
-def _redis_lock_key_builder(*args, **kwargs) -> str:
-    app = _get_app_from_args(*args, **kwargs)
-    return f"{app.title}_{MODULE_NAME}"
+        await request_pipeline_scheduling(
+            rabbitmq_client,
+            db_engine,
+            user_id=updated_comp_run.user_id,
+            project_id=updated_comp_run.project_uuid,
+            iteration=updated_comp_run.iteration,
+        )
 
 
 async def _get_pipeline_dag(project_id: ProjectID, db_engine: Engine) -> nx.DiGraph:
@@ -129,7 +111,12 @@ async def _get_pipeline_dag(project_id: ProjectID, db_engine: Engine) -> nx.DiGr
     return pipeline_at_db.get_graph()
 
 
-@exclusive(_redis_client_getter, lock_key=_redis_lock_key_builder)
+@exclusive(
+    get_redis_client_from_app,
+    lock_key=get_redis_lock_key(
+        MODULE_NAME_SCHEDULER, unique_lock_key_builder=lambda: ""
+    ),
+)
 async def schedule_pipelines(app: FastAPI) -> None:
     with log_context(_logger, logging.DEBUG, msg="scheduling pipelines"):
         db_engine = get_db_engine(app)
@@ -159,7 +146,7 @@ async def setup_manager(app: FastAPI) -> None:
     app.state.scheduler_manager = start_periodic_task(
         schedule_pipelines,
         interval=SCHEDULER_INTERVAL,
-        task_name=MODULE_NAME,
+        task_name=MODULE_NAME_SCHEDULER,
         app=app,
     )