✨ Adding generic scheduling capability to dynamic-scheduler [part 1/2]

packages/pytest-simcore/src/pytest_simcore/helpers/docker.py

@@ -1,19 +1,28 @@
+import contextlib
 import json
 import logging
 import os
 import re
 import subprocess
+from collections.abc import AsyncIterator, Callable
+from contextlib import AbstractAsyncContextManager
 from enum import Enum
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any, Protocol
 
 import docker
 import yaml
 from tenacity import retry
 from tenacity.after import after_log
-from tenacity.stop import stop_after_attempt
+from tenacity.asyncio import AsyncRetrying
+from tenacity.retry import retry_if_exception_type
+from tenacity.stop import stop_after_attempt, stop_after_delay
 from tenacity.wait import wait_fixed
 
+if TYPE_CHECKING:
+    from servicelib.rabbitmq import RabbitMQClient
+    from servicelib.redis import RedisClientSDK
+
 
 # NOTE: CANNOT use models_library.generated_models.docker_rest_api.Status2 because some of the
 # packages tests installations do not include this library!!
@@ -272,3 +281,56 @@ def save_docker_infos(destination_dir: Path):
             f"wrote docker log and json files for {len(all_containers)} containers in ",
             destination_dir,
         )
+
+
+class _ClientWithPingProtocol(Protocol):
+    async def ping(self) -> bool: ...
+
+
+class ServiceManager:
+    def __init__(
+        self,
+        redis_client: "RedisClientSDK",
+        rabbit_client: "RabbitMQClient",
+        paused_container: Callable[[str], AbstractAsyncContextManager[None]],
+    ) -> None:
+        self.redis_client = redis_client
+        self.rabbit_client = rabbit_client
+        self.paused_container = paused_container
+
+    @contextlib.asynccontextmanager
+    async def _paused_container(
+        self, container_name: str, client: _ClientWithPingProtocol
+    ) -> AsyncIterator[None]:
+        async with self.paused_container(container_name):
+            async for attempt in AsyncRetrying(
+                wait=wait_fixed(0.1),
+                stop=stop_after_delay(10),
+                reraise=True,
+                retry=retry_if_exception_type(AssertionError),
+            ):
+                with attempt:
+                    assert await client.ping() is False
+            yield
+
+        async for attempt in AsyncRetrying(
+            wait=wait_fixed(0.1),
+            stop=stop_after_delay(10),
+            reraise=True,
+            retry=retry_if_exception_type(AssertionError),
+        ):
+            with attempt:
+                assert await client.ping() is True
+
+    @contextlib.asynccontextmanager
+    async def pause_rabbit(self) -> AsyncIterator[None]:
+        async with self._paused_container("rabbit", self.rabbit_client):
+            yield
+
+    @contextlib.asynccontextmanager
+    async def pause_redis(self) -> AsyncIterator[None]:
+        # save db for clean restore point
+        await self.redis_client.redis.save()
+
+        async with self._paused_container("redis", self.redis_client):
+            yield

packages/service-library/tests/deferred_tasks/test_deferred_tasks.py

@@ -2,16 +2,15 @@
 # pylint:disable=unused-argument
 
 import asyncio
-import contextlib
 import datetime
 import itertools
 import json
 import random
 import sys
-from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
+from collections.abc import AsyncIterable, Awaitable, Callable
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, suppress
 from pathlib import Path
-from typing import Any, Protocol
+from typing import Any
 
 import psutil
 import pytest
@@ -20,6 +19,7 @@
 from common_library.serialization import model_dump_with_secrets
 from pydantic import NonNegativeFloat, NonNegativeInt
 from pytest_mock import MockerFixture
+from pytest_simcore.helpers.docker import ServiceManager
 from servicelib.rabbitmq import RabbitMQClient
 from servicelib.redis import RedisClientSDK
 from servicelib.sequences_utils import partition_gen
@@ -332,59 +332,6 @@ async def rabbit_client(
     return create_rabbitmq_client("pinger")
 
 
-class ClientWithPingProtocol(Protocol):
-    async def ping(self) -> bool: ...
-
-
-class ServiceManager:
-    def __init__(
-        self,
-        redis_client: RedisClientSDK,
-        rabbit_client: RabbitMQClient,
-        paused_container: Callable[[str], AbstractAsyncContextManager[None]],
-    ) -> None:
-        self.redis_client = redis_client
-        self.rabbit_client = rabbit_client
-        self.paused_container = paused_container
-
-    @contextlib.asynccontextmanager
-    async def _paused_container(
-        self, container_name: str, client: ClientWithPingProtocol
-    ) -> AsyncIterator[None]:
-        async with self.paused_container(container_name):
-            async for attempt in AsyncRetrying(
-                wait=wait_fixed(0.1),
-                stop=stop_after_delay(10),
-                reraise=True,
-                retry=retry_if_exception_type(AssertionError),
-            ):
-                with attempt:
-                    assert await client.ping() is False
-            yield
-
-        async for attempt in AsyncRetrying(
-            wait=wait_fixed(0.1),
-            stop=stop_after_delay(10),
-            reraise=True,
-            retry=retry_if_exception_type(AssertionError),
-        ):
-            with attempt:
-                assert await client.ping() is True
-
-    @contextlib.asynccontextmanager
-    async def pause_rabbit(self) -> AsyncIterator[None]:
-        async with self._paused_container("rabbit", self.rabbit_client):
-            yield
-
-    @contextlib.asynccontextmanager
-    async def pause_redis(self) -> AsyncIterator[None]:
-        # save db for clean restore point
-        await self.redis_client.redis.save()
-
-        async with self._paused_container("redis", self.redis_client):
-            yield
-
-
 @pytest.fixture
 def mock_default_socket_timeout(mocker: MockerFixture) -> None:
     mocker.patch(

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/api/rest/_ops.py

@@ -5,7 +5,7 @@
     DynamicServiceGet,
 )
 
-from ...services import scheduler_interface
+from ...services import common_interface
 from ._dependencies import (
     get_app,
 )
@@ -19,6 +19,6 @@ async def running_services(
 ) -> list[DynamicServiceGet]:
     """returns all running dynamic services. Used by ops internall to determine
     when it is safe to shutdown the platform"""
-    return await scheduler_interface.list_tracked_dynamic_services(
+    return await common_interface.list_tracked_dynamic_services(
         app, user_id=None, project_id=None
     )

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/api/rpc/_services.py

@@ -20,7 +20,7 @@
     ServiceWasNotFoundError,
 )
 
-from ...services import scheduler_interface
+from ...services import common_interface
 
 router = RPCRouter()
 
@@ -29,7 +29,7 @@
 async def list_tracked_dynamic_services(
     app: FastAPI, *, user_id: UserID | None = None, project_id: ProjectID | None = None
 ) -> list[DynamicServiceGet]:
-    return await scheduler_interface.list_tracked_dynamic_services(
+    return await common_interface.list_tracked_dynamic_services(
         app, user_id=user_id, project_id=project_id
     )
 
@@ -38,14 +38,14 @@ async def list_tracked_dynamic_services(
 async def get_service_status(
     app: FastAPI, *, node_id: NodeID
 ) -> NodeGet | DynamicServiceGet | NodeGetIdle:
-    return await scheduler_interface.get_service_status(app, node_id=node_id)
+    return await common_interface.get_service_status(app, node_id=node_id)
 
 
 @router.expose()
 async def run_dynamic_service(
     app: FastAPI, *, dynamic_service_start: DynamicServiceStart
 ) -> NodeGet | DynamicServiceGet:
-    return await scheduler_interface.run_dynamic_service(
+    return await common_interface.run_dynamic_service(
         app, dynamic_service_start=dynamic_service_start
     )
 
@@ -59,7 +59,7 @@ async def run_dynamic_service(
 async def stop_dynamic_service(
     app: FastAPI, *, dynamic_service_stop: DynamicServiceStop
 ) -> None:
-    return await scheduler_interface.stop_dynamic_service(
+    return await common_interface.stop_dynamic_service(
         app, dynamic_service_stop=dynamic_service_stop
     )
 
@@ -68,25 +68,25 @@ async def stop_dynamic_service(
 async def get_project_inactivity(
     app: FastAPI, *, project_id: ProjectID, max_inactivity_seconds: NonNegativeInt
 ) -> GetProjectInactivityResponse:
-    return await scheduler_interface.get_project_inactivity(
+    return await common_interface.get_project_inactivity(
         app, project_id=project_id, max_inactivity_seconds=max_inactivity_seconds
     )
 
 
 @router.expose()
 async def restart_user_services(app: FastAPI, *, node_id: NodeID) -> None:
-    await scheduler_interface.restart_user_services(app, node_id=node_id)
+    await common_interface.restart_user_services(app, node_id=node_id)
 
 
 @router.expose()
 async def retrieve_inputs(
     app: FastAPI, *, node_id: NodeID, port_keys: list[ServicePortKey]
 ) -> RetrieveDataOutEnveloped:
-    return await scheduler_interface.retrieve_inputs(
+    return await common_interface.retrieve_inputs(
         app, node_id=node_id, port_keys=port_keys
     )
 
 
 @router.expose()
 async def update_projects_networks(app: FastAPI, *, project_id: ProjectID) -> None:
-    await scheduler_interface.update_projects_networks(app, project_id=project_id)
+    await common_interface.update_projects_networks(app, project_id=project_id)

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/core/events.py

@@ -23,6 +23,7 @@
 from ..services.deferred_manager import deferred_manager_lifespan
 from ..services.director_v0 import director_v0_lifespan
 from ..services.director_v2 import director_v2_lifespan
+from ..services.generic_scheduler import get_generic_scheduler_lifespans
 from ..services.notifier import get_notifier_lifespans
 from ..services.rabbitmq import rabbitmq_lifespan
 from ..services.redis import redis_lifespan
@@ -79,6 +80,9 @@ def create_app_lifespan(
     for lifespan in get_notifier_lifespans():
         app_lifespan.add(lifespan)
 
+    for lifespan in get_generic_scheduler_lifespans():
+        app_lifespan.add(lifespan)
+
     app_lifespan.add(service_tracker_lifespan)
     app_lifespan.add(deferred_manager_lifespan)
     app_lifespan.add(status_monitor_lifespan)

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/scheduler_interface.py → services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/common_interface.py

@@ -50,6 +50,8 @@ async def get_service_status(
 async def run_dynamic_service(
     app: FastAPI, *, dynamic_service_start: DynamicServiceStart
 ) -> NodeGet | DynamicServiceGet:
+    await set_request_as_running(app, dynamic_service_start)
+
     settings: ApplicationSettings = app.state.settings
     if settings.DYNAMIC_SCHEDULER_USE_INTERNAL_SCHEDULER:
         raise NotImplementedError
@@ -59,13 +61,14 @@ async def run_dynamic_service(
         await director_v2_client.run_dynamic_service(dynamic_service_start)
     )
 
-    await set_request_as_running(app, dynamic_service_start)
     return response
 
 
 async def stop_dynamic_service(
     app: FastAPI, *, dynamic_service_stop: DynamicServiceStop
 ) -> None:
+    await set_request_as_stopped(app, dynamic_service_stop)
+
     settings: ApplicationSettings = app.state.settings
     if settings.DYNAMIC_SCHEDULER_USE_INTERNAL_SCHEDULER:
         raise NotImplementedError
@@ -78,8 +81,6 @@ async def stop_dynamic_service(
         timeout=settings.DYNAMIC_SCHEDULER_STOP_SERVICE_TIMEOUT,
     )
 
-    await set_request_as_stopped(app, dynamic_service_stop)
-
 
 async def get_project_inactivity(
     app: FastAPI, *, project_id: ProjectID, max_inactivity_seconds: NonNegativeInt

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/generic_scheduler/__init__.py

@@ -0,0 +1,49 @@
+from ._api import (
+    cancel_operation,
+    restart_operation_stuck_in_error_during_revert,
+    restart_operation_stuck_in_manual_intervention_during_create,
+    start_operation,
+)
+from ._deferred_runner import (
+    get_operation_context_proxy,
+    get_step_group_proxy,
+    get_step_store_proxy,
+)
+from ._lifespan import get_generic_scheduler_lifespans
+from ._models import (
+    OperationName,
+    ProvidedOperationContext,
+    RequiredOperationContext,
+    ScheduleId,
+)
+from ._operation import (
+    BaseStep,
+    Operation,
+    OperationRegistry,
+    ParallelStepGroup,
+    SingleStepGroup,
+)
+from ._store import OperationContextProxy, StepGroupProxy, StepStoreProxy
+
+__all__: tuple[str, ...] = (
+    "BaseStep",
+    "cancel_operation",
+    "get_generic_scheduler_lifespans",
+    "get_operation_context_proxy",
+    "get_step_group_proxy",
+    "get_step_store_proxy",
+    "Operation",
+    "OperationContextProxy",
+    "OperationName",
+    "OperationRegistry",
+    "ParallelStepGroup",
+    "ProvidedOperationContext",
+    "RequiredOperationContext",
+    "restart_operation_stuck_in_error_during_revert",
+    "restart_operation_stuck_in_manual_intervention_during_create",
+    "ScheduleId",
+    "SingleStepGroup",
+    "start_operation",
+    "StepGroupProxy",
+    "StepStoreProxy",
+)