first working version

Author: Andrei Neagu

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

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 from collections.abc import AsyncIterator, Iterable
+from contextlib import asynccontextmanager
 from datetime import datetime, timedelta
 from typing import Final
 from uuid import uuid4
@@ -14,7 +15,7 @@
 from servicelib.utils import limited_gather
 
 from ._deferred_runner import DeferredRunner
-from ._dependencies import enqueue_event
+from ._dependencies import enqueue_schedule_event
 from ._errors import UnexpectedStepHandlingError
 from ._models import (
     OperationContext,
@@ -28,6 +29,7 @@
 from ._store import (
     KeyNotFoundInHashError,
     ScheduleDataStoreProxy,
+    StepGroupProxy,
     StepStoreProxy,
     Store,
     get_store,
@@ -92,14 +94,18 @@ async def _is_operation_in_progress_status(
 
 
 async def _start_and_mark_as_started(
-    step_proxy: StepStoreProxy, *, is_creating: bool
+    step_proxy: StepStoreProxy,
+    *,
+    is_creating: bool,
+    expected_steps_count: NonNegativeInt,
 ) -> None:
     await DeferredRunner.start(
         schedule_id=step_proxy.schedule_id,
         operation_name=step_proxy.operation_name,
         step_group_name=step_proxy.step_group_name,
         step_name=step_proxy.step_name,
         is_creating=is_creating,
+        expected_steps_count=expected_steps_count,
     )
     await step_proxy.set_multiple(
         {
@@ -150,16 +156,13 @@ async def create(
                 "is_creating": True,
             }
         )
-        await enqueue_event(self.app, schedule_id)
+        await enqueue_schedule_event(self.app, schedule_id)
         return schedule_id
 
-    async def safe_on_schedule_event(self, schedule_id: ScheduleId) -> None:
-        """
-        NOTE: do not call this directly, you are doing something wrong
-        runs scheduling actions for a given schedule_id
-        """
+    @asynccontextmanager
+    async def _safe_event(self, schedule_id: ScheduleId) -> AsyncIterator[None]:
         try:
-            await self._on_schedule_event(schedule_id)
+            yield
         except KeyNotFoundInHashError as err:
             _logger.debug(
                 "Cannot process schedule_id='%s' since it's data was not found: %s",
@@ -182,6 +185,11 @@ async def safe_on_schedule_event(self, schedule_id: ScheduleId) -> None:
                 message=log_kwargs["msg"],
             )
 
+    async def safe_on_schedule_event(self, schedule_id: ScheduleId) -> None:
+        # NOTE: do not call this directly, you are doing something wrong
+        async with self._safe_event(schedule_id):
+            await self._on_schedule_event(schedule_id)
+
     async def cancel_schedule(self, schedule_id: ScheduleId) -> None:
         """
         Cancels and runs destruction of the operation
@@ -282,6 +290,7 @@ async def _on_schedule_event(self, schedule_id: ScheduleId) -> None:
             step_group=step_group,
             is_creating=is_creating,
         )
+        group_step_count = len(step_group)
 
         # get steps to start
         if to_start_step_proxies := await _get_steps_to_start(
@@ -297,7 +306,11 @@ async def _on_schedule_event(self, schedule_id: ScheduleId) -> None:
             ):
                 await limited_gather(
                     *(
-                        _start_and_mark_as_started(step_proxy, is_creating=is_creating)
+                        _start_and_mark_as_started(
+                            step_proxy,
+                            is_creating=is_creating,
+                            expected_steps_count=group_step_count,
+                        )
                         for step_proxy in to_start_step_proxies
                     ),
                     limit=_PARALLEL_STATUS_REQUESTS,
@@ -356,7 +369,7 @@ async def _continue_handling_as_creation(
                 # does a next group exist?
                 _ = operation[next_group_index]
                 await schedule_data_proxy.set("group_index", value=next_group_index)
-                await enqueue_event(self.app, schedule_id)
+                await enqueue_schedule_event(self.app, schedule_id)
             except IndexError:
                 # does the step need repeating?
                 if current_step_group.repeat_steps is True:
@@ -381,7 +394,18 @@ async def _continue_handling_as_creation(
                         )
                         await step_rproxy.remove()
 
-                    await enqueue_event(self.app, schedule_id)
+                    group_proxy = StepGroupProxy(
+                        store=self._store,
+                        schedule_id=schedule_id,
+                        operation_name=operation_name,
+                        step_group_name=current_step_group.get_step_group_name(
+                            index=group_index
+                        ),
+                        is_creating=True,
+                    )
+                    await group_proxy.remove()
+
+                    await enqueue_schedule_event(self.app, schedule_id)
                 else:
                     # TODO: the end has bean reached, do nothing from now on
                     _logger.debug(
@@ -439,7 +463,7 @@ async def _continue_handling_as_creation(
                 f"{operation_name=} was not successfull: {steps_statuses=}, moving to revert",
             ):
                 await schedule_data_proxy.set("is_creating", value=False)
-                await enqueue_event(self.app, schedule_id)
+                await enqueue_schedule_event(self.app, schedule_id)
             return
 
         raise UnexpectedStepHandlingError(
@@ -461,7 +485,7 @@ async def _continue_handling_as_reverting(
                 return
 
             await schedule_data_proxy.set("group_index", value=previous_group_index)
-            await enqueue_event(self.app, schedule_id)
+            await enqueue_schedule_event(self.app, schedule_id)
             return
 
         # if any in FAILED this is unexpected falg to be investigated

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

@@ -1,13 +1,14 @@
 from datetime import timedelta
 
 from fastapi import FastAPI
+from pydantic import NonNegativeInt
 from servicelib.deferred_tasks import BaseDeferredHandler, DeferredContext, TaskUID
 from servicelib.deferred_tasks._models import TaskResultError
 
-from ._dependencies import enqueue_event
+from ._dependencies import enqueue_schedule_event
 from ._models import OperationName, ScheduleId, StepGroupName, StepName, StepStatus
 from ._operation import BaseStep, OperationRegistry
-from ._store import StepStoreProxy, Store, get_store
+from ._store import StepGroupProxy, StepStoreProxy, Store, get_store
 
 
 def _get_step_store_proxy(context: DeferredContext) -> StepStoreProxy:
@@ -29,16 +30,39 @@ def _get_step_store_proxy(context: DeferredContext) -> StepStoreProxy:
     )
 
 
+def _get_step_group_proxy(context: DeferredContext) -> StepGroupProxy:
+    app: FastAPI = context["app"]
+    schedule_id: ScheduleId = context["schedule_id"]
+    operation_name: OperationName = context["operation_name"]
+    step_group_name: StepGroupName = context["step_group_name"]
+    is_creating = context["is_creating"]
+
+    store: Store = get_store(app)
+    return StepGroupProxy(
+        store=store,
+        schedule_id=schedule_id,
+        operation_name=operation_name,
+        step_group_name=step_group_name,
+        is_creating=is_creating,
+    )
+
+
 def _get_step(context: DeferredContext) -> type[BaseStep]:
     operation_name: OperationName = context["operation_name"]
     step_name: StepName = context["step_name"]
     return OperationRegistry.get_step(operation_name, step_name)
 
 
-async def _send_schedule_event(context: DeferredContext) -> None:
+async def _send_group_done_check_event(context: DeferredContext) -> None:
     app: FastAPI = context["app"]
     schedule_id: ScheduleId = context["schedule_id"]
-    await enqueue_event(app, schedule_id)
+    expected_steps_count: NonNegativeInt = context["expected_steps_count"]
+
+    if (
+        await _get_step_group_proxy(context).increment_and_get_done_steps_count()
+        == expected_steps_count
+    ):
+        await enqueue_schedule_event(app, schedule_id)
 
 
 class DeferredRunner(BaseDeferredHandler[None]):
@@ -51,13 +75,15 @@ async def start(  # type:ignore[override] # pylint:disable=arguments-differ
         step_group_name: StepGroupName,
         step_name: StepName,
         is_creating: bool,
+        expected_steps_count: NonNegativeInt,
     ) -> DeferredContext:
         return {
             "schedule_id": schedule_id,
             "operation_name": operation_name,
             "step_group_name": step_group_name,
             "step_name": step_name,
             "is_creating": is_creating,
+            "expected_steps_count": expected_steps_count,
         }
 
     @classmethod
@@ -104,7 +130,8 @@ async def run(cls, context: DeferredContext) -> None:
     async def on_result(cls, result: None, context: DeferredContext) -> None:
         _ = result
         await _get_step_store_proxy(context).set("status", StepStatus.SUCCESS)
-        await _send_schedule_event(context)
+
+        await _send_group_done_check_event(context)
 
     @classmethod
     async def on_finished_with_error(
@@ -113,9 +140,11 @@ async def on_finished_with_error(
         await _get_step_store_proxy(context).set_multiple(
             {"status": StepStatus.FAILED, "error_traceback": error.format_error()}
         )
-        await _send_schedule_event(context)
+
+        await _send_group_done_check_event(context)
 
     @classmethod
     async def on_cancelled(cls, context: DeferredContext) -> None:
         await _get_step_store_proxy(context).set("status", StepStatus.CANCELLED)
-        await _send_schedule_event(context)
+
+        await _send_group_done_check_event(context)

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

@@ -8,6 +8,6 @@
     from ._event_scheduler import EventScheduler
 
 
-async def enqueue_event(app: FastAPI, schedule_id: ScheduleId) -> None:
+async def enqueue_schedule_event(app: FastAPI, schedule_id: ScheduleId) -> None:
     event_scheduler: EventScheduler = app.state.generic_scheduler_event_scheduler
-    await event_scheduler.enqueue_event(schedule_id)
+    await event_scheduler.enqueue_schedule_event(schedule_id)

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

@@ -1,7 +1,7 @@
 import functools
 import logging
 from collections.abc import AsyncIterator
-from typing import Final
+from typing import Final, NotRequired, TypedDict
 
 from fastapi import FastAPI
 from fastapi_lifespan_manager import State
@@ -13,6 +13,7 @@
     RabbitQueue,
     RabbitRouter,
 )
+from faststream.rabbit.schemas.queue import ClassicQueueArgs
 
 from ...core.settings import ApplicationSettings
 from ._core import get_core
@@ -21,13 +22,20 @@
 _logger = logging.getLogger(__name__)
 
 
+class QueuConfiguration(TypedDict):
+    queue_name: str
+    arguments: NotRequired[ClassicQueueArgs | None]
+
+
 _EXCHANGE_NAME: Final[str] = __name__
-_QUEUE_NAME: Final[str] = "event_scheduler"
-_NO_PARALLE_MESSAGE_CONSUMING: Final[int] = 1
 
 
-def _get_global_queue(queue_name: str) -> RabbitQueue:
-    return RabbitQueue(f"{_EXCHANGE_NAME}_{queue_name}", durable=True)
+def _get_global_queue(
+    queue_name: str, arguments: ClassicQueueArgs | None = None
+) -> RabbitQueue:
+    return RabbitQueue(
+        f"{_EXCHANGE_NAME}_{queue_name}", durable=True, arguments=arguments
+    )
 
 
 def _stop_retry_for_unintended_errors(func):
@@ -66,47 +74,45 @@ def __init__(self, app: FastAPI) -> None:
 
         settings: ApplicationSettings = app.state.settings
 
-        self.broker: RabbitBroker = RabbitBroker(
-            settings.DYNAMIC_SCHEDULER_RABBITMQ.dsn,
-            log_level=logging.DEBUG,
-            max_consumers=_NO_PARALLE_MESSAGE_CONSUMING,
+        self._broker: RabbitBroker = RabbitBroker(
+            settings.DYNAMIC_SCHEDULER_RABBITMQ.dsn, log_level=logging.DEBUG
         )
-        self.router: RabbitRouter = RabbitRouter()
-        self.exchange = RabbitExchange(
+        self._router: RabbitRouter = RabbitRouter()
+        self._exchange = RabbitExchange(
             _EXCHANGE_NAME, durable=True, type=ExchangeType.DIRECT
         )
+        self._queue_schedule_event = _get_global_queue(queue_name="schedule_queue")
 
     @_stop_retry_for_unintended_errors
-    async def _on_secure_schedule_event(  # pylint:disable=method-hidden
+    async def _on_schedule_secure_event(  # pylint:disable=method-hidden
         self, schedule_id: ScheduleId
     ) -> None:
-        # advance operation
-        # NOTE: should no longer forward operation if nothing needs doing
-        # an unexpected error might be raised
         await get_core(self.app).safe_on_schedule_event(schedule_id)
 
-    async def enqueue_event(self, schedule_id: ScheduleId) -> None:
-        await self.broker.publish(
-            schedule_id, queue=_get_global_queue(_QUEUE_NAME), exchange=self.exchange
+    async def enqueue_schedule_event(self, schedule_id: ScheduleId) -> None:
+        await self._broker.publish(
+            schedule_id,
+            queue=self._queue_schedule_event,
+            exchange=self._exchange,
         )
 
     def _register_subscribers(self) -> None:
         # pylint:disable=unexpected-keyword-arg
         # pylint:disable=no-value-for-parameter
-        self._on_secure_schedule_event = self.router.subscriber(
-            queue=_get_global_queue(_QUEUE_NAME),
-            exchange=self.exchange,
+        self._on_schedule_secure_event = self._router.subscriber(
+            queue=self._queue_schedule_event,
+            exchange=self._exchange,
             retry=True,
-        )(self._on_secure_schedule_event)
+        )(self._on_schedule_secure_event)
 
     async def setup(self) -> None:
         self._register_subscribers()
-        self.broker.include_router(self.router)
+        self._broker.include_router(self._router)
 
-        await self.broker.start()
+        await self._broker.start()
 
     async def shutdown(self) -> None:
-        await self.broker.close()
+        await self._broker.close()
 
 
 async def lifespan(app: FastAPI) -> AsyncIterator[State]:

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

@@ -91,6 +91,10 @@ def __init__(self, *, repeat_steps: bool, wait_before_repeat: timedelta) -> None
         self.repeat_steps = repeat_steps
         self.wait_before_repeat = wait_before_repeat
 
+    @abstractmethod
+    def __len__(self) -> int:
+        """number of steps in this group"""
+
     @abstractmethod
     def get_step_group_name(self, *, index: NonNegativeInt) -> StepGroupName:
         """returns the name of this step group"""
@@ -121,6 +125,9 @@ def __init__(
             repeat_steps=repeat_steps, wait_before_repeat=wait_before_repeat
         )
 
+    def __len__(self) -> int:
+        return 1
+
     def get_step_group_name(self, *, index: NonNegativeInt) -> StepGroupName:
         return f"{index}S{'R' if self.repeat_steps else ''}"
 
@@ -148,6 +155,9 @@ def __init__(
             repeat_steps=repeat_steps, wait_before_repeat=wait_before_repeat
         )
 
+    def __len__(self) -> int:
+        return len(self._steps)
+
     @property
     def steps(self) -> list[type[BaseStep]]:
         return self._steps