✨ dynamic-scheduler can chain operations

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

@@ -9,9 +9,14 @@
     get_step_group_proxy,
     get_step_store_proxy,
 )
+from ._event_after_registration import (
+    register_to_start_after_on_created_completed,
+    register_to_start_after_on_undo_completed,
+)
 from ._lifespan import generic_scheduler_lifespan
 from ._models import (
     OperationName,
+    OperationToStart,
     ProvidedOperationContext,
     RequiredOperationContext,
     ScheduleId,
@@ -36,8 +41,11 @@
     "OperationContextProxy",
     "OperationName",
     "OperationRegistry",
+    "OperationToStart",
     "ParallelStepGroup",
     "ProvidedOperationContext",
+    "register_to_start_after_on_created_completed",
+    "register_to_start_after_on_undo_completed",
     "RequiredOperationContext",
     "restart_operation_step_stuck_during_undo",
     "restart_operation_step_stuck_in_manual_intervention_during_create",

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

@@ -30,16 +30,27 @@
 from ._errors import (
     CannotCancelWhileWaitingForManualInterventionError,
     NoDataFoundError,
+    OperationNotCancellableError,
     StepNameNotInCurrentGroupError,
     StepNotInErrorStateError,
     StepNotWaitingForManualInterventionError,
     UnexpectedStepHandlingError,
 )
-from ._event import enqueue_schedule_event
+from ._event import (
+    enqueue_create_completed_event,
+    enqueue_schedule_event,
+    enqueue_undo_completed_event,
+)
+from ._event_after_registration import (
+    register_to_start_after_on_created_completed,
+    register_to_start_after_on_undo_completed,
+)
 from ._models import (
+    EventType,
     OperationContext,
     OperationErrorType,
     OperationName,
+    OperationToStart,
     ScheduleId,
     StepName,
     StepStatus,
@@ -52,6 +63,7 @@
 from ._store import (
     DeleteStepKeys,
     OperationContextProxy,
+    OperationEventsProxy,
     ScheduleDataStoreProxy,
     StepGroupProxy,
     StepStoreProxy,
@@ -79,7 +91,11 @@ def __init__(
         self._store: Store = Store.get_from_app_state(app)
 
     async def start_operation(
-        self, operation_name: OperationName, initial_operation_context: OperationContext
+        self,
+        operation_name: OperationName,
+        initial_operation_context: OperationContext,
+        on_create_completed: OperationToStart | None,
+        on_undo_completed: OperationToStart | None,
     ) -> ScheduleId:
         """start an operation by it's given name and providing an initial context"""
         schedule_id: ScheduleId = f"{uuid4()}"
@@ -112,6 +128,16 @@ async def start_operation(
         )
         await operation_content_proxy.create_or_update(initial_operation_context)
 
+        if on_create_completed:
+            await register_to_start_after_on_created_completed(
+                self.app, schedule_id, to_start=on_create_completed
+            )
+
+        if on_undo_completed:
+            await register_to_start_after_on_undo_completed(
+                self.app, schedule_id=schedule_id, to_start=on_undo_completed
+            )
+
         await enqueue_schedule_event(self.app, schedule_id)
         return schedule_id
 
@@ -140,6 +166,10 @@ async def cancel_operation(self, schedule_id: ScheduleId) -> None:
         group_index = await schedule_data_proxy.read("group_index")
 
         operation = OperationRegistry.get_operation(operation_name)
+
+        if operation.is_cancellable is False:
+            raise OperationNotCancellableError(operation_name=operation_name)
+
         group = operation[group_index]
 
         group_step_proxies = get_group_step_proxies(
@@ -152,15 +182,14 @@ async def cancel_operation(self, schedule_id: ScheduleId) -> None:
         )
 
         # not allowed to cancel while waiting for manual intervention
-        if any(
-            await limited_gather(
-                *(
-                    get_requires_manual_intervention(step)
-                    for step in group_step_proxies.values()
-                ),
-                limit=PARALLEL_REQUESTS,
-            )
-        ):
+        require_manual_intervention = await limited_gather(
+            *(
+                get_requires_manual_intervention(step)
+                for step in group_step_proxies.values()
+            ),
+            limit=PARALLEL_REQUESTS,
+        )
+        if any(require_manual_intervention):
             raise CannotCancelWhileWaitingForManualInterventionError(
                 schedule_id=schedule_id
             )
@@ -465,7 +494,7 @@ async def _advance_as_creating(
         # CREATION logic:
         # 1) if all steps in group in SUUCESS
         # - 1a) -> move to next group
-        # - 1b) if reached the end of the CREATE operation -> remove all created data
+        # - 1b) if reached the end of the CREATE operation -> remove all created data [EMIT create complete event]
         # 2) if manual intervention is required -> do nothing else
         # 3) if any step in CANCELLED or FAILED (and not in manual intervention) -> move to undo
 
@@ -483,10 +512,33 @@ async def _advance_as_creating(
                 await enqueue_schedule_event(self.app, schedule_id)
             except IndexError:
 
-                # 1b) if reached the end of the CREATE operation -> remove all created data
+                # 1b) if reached the end of the CREATE operation -> remove all created data [EMIT create complete event]
+                on_created_proxy = OperationEventsProxy(
+                    self._store, schedule_id, EventType.ON_CREATED_COMPLETED
+                )
+                on_create_operation_name: OperationName | None = None
+                on_create_initial_context: OperationContext | None = None
+                if await on_created_proxy.exists():
+                    on_create_operation_name = await on_created_proxy.read(
+                        "operation_name"
+                    )
+                    on_create_initial_context = await on_created_proxy.read(
+                        "initial_context"
+                    )
+
                 await cleanup_after_finishing(
                     self._store, schedule_id=schedule_id, is_creating=True
                 )
+                if (
+                    on_create_operation_name is not None
+                    and on_create_initial_context is not None
+                ):
+                    await enqueue_create_completed_event(
+                        self.app,
+                        schedule_id,
+                        on_create_operation_name,
+                        on_create_initial_context,
+                    )
 
             return
 
@@ -551,7 +603,7 @@ async def _advance_as_undoing(
     ) -> None:
         # UNDO logic:
         # 1) if all steps in group in SUCCESS
-        # - 1a) if reached the end of the UNDO operation -> remove all created data
+        # - 1a) if reached the end of the UNDO operation -> remove all created data [EMIT undo complete event]
         # - 1b) -> move to previous group
         # 2) it is unexpected to have a FAILED step -> do nothing else
         # 3) it is unexpected to have a CANCELLED step -> do nothing else
@@ -561,10 +613,31 @@ async def _advance_as_undoing(
             previous_group_index = group_index - 1
             if previous_group_index < 0:
 
-                # 1a) if reached the end of the UNDO operation -> remove all created data
+                # 1a) if reached the end of the UNDO operation -> remove all created data [EMIT undo complete event]
+                on_undo_proxy = OperationEventsProxy(
+                    self._store, schedule_id, EventType.ON_UNDO_COMPLETED
+                )
+                on_undo_operation_name: OperationName | None = None
+                on_undo_initial_context: OperationContext | None = None
+                if await on_undo_proxy.exists():
+                    on_undo_operation_name = await on_undo_proxy.read("operation_name")
+                    on_undo_initial_context = await on_undo_proxy.read(
+                        "initial_context"
+                    )
+
                 await cleanup_after_finishing(
                     self._store, schedule_id=schedule_id, is_creating=False
                 )
+                if (
+                    on_undo_operation_name is not None
+                    and on_undo_initial_context is not None
+                ):
+                    await enqueue_undo_completed_event(
+                        self.app,
+                        schedule_id,
+                        on_undo_operation_name,
+                        on_undo_initial_context,
+                    )
                 return
 
             # 1b) -> move to previous group
@@ -634,9 +707,15 @@ async def start_operation(
     app: FastAPI,
     operation_name: OperationName,
     initial_operation_context: OperationContext,
+    *,
+    on_create_completed: OperationToStart | None = None,
+    on_undo_completed: OperationToStart | None = None,
 ) -> ScheduleId:
     return await Core.get_from_app_state(app).start_operation(
-        operation_name, initial_operation_context
+        operation_name,
+        initial_operation_context,
+        on_create_completed,
+        on_undo_completed,
     )
 
 

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

@@ -49,6 +49,10 @@ class InitialOperationContextKeyNotAllowedError(BaseGenericSchedulerError):
     )
 
 
+class OperationNotCancellableError(BaseGenericSchedulerError):
+    msg_template: str = "Operation '{operation_name}' is not cancellable"
+
+
 class CannotCancelWhileWaitingForManualInterventionError(BaseGenericSchedulerError):
     msg_template: str = (
         "Cannot cancel schedule_id='{schedule_id}' while one or more steps are waiting for manual intervention."

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

@@ -2,12 +2,40 @@
 
 from fastapi import FastAPI
 
-from ._models import ScheduleId
+from ._models import OperationContext, OperationName, ScheduleId
 
 if TYPE_CHECKING:
     from ._event_scheduler import EventScheduler
 
 
-async def enqueue_schedule_event(app: FastAPI, schedule_id: ScheduleId) -> None:
+def _get_event_scheduler(app: FastAPI) -> "EventScheduler":
+    # NOTE: could not use EventScheduler.get_from_app_state(app)
+    # due to circular dependency
     event_scheduler: EventScheduler = app.state.generic_scheduler_event_scheduler
-    await event_scheduler.enqueue_schedule_event(schedule_id)
+    return event_scheduler
+
+
+async def enqueue_schedule_event(app: FastAPI, schedule_id: ScheduleId) -> None:
+    await _get_event_scheduler(app).enqueue_schedule_event(schedule_id)
+
+
+async def enqueue_create_completed_event(
+    app: FastAPI,
+    schedule_id: ScheduleId,
+    operation_name: OperationName,
+    initial_context: OperationContext,
+) -> None:
+    await _get_event_scheduler(app).enqueue_create_completed_event(
+        schedule_id, operation_name, initial_context
+    )
+
+
+async def enqueue_undo_completed_event(
+    app: FastAPI,
+    schedule_id: ScheduleId,
+    operation_name: OperationName,
+    initial_context: OperationContext,
+) -> None:
+    await _get_event_scheduler(app).enqueue_undo_completed_event(
+        schedule_id, operation_name, initial_context
+    )

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

@@ -0,0 +1,92 @@
+import logging
+
+from fastapi import FastAPI
+from servicelib.fastapi.app_state import SingletonInAppStateMixin
+from servicelib.logging_utils import log_context
+
+from ._core import start_operation
+from ._models import (
+    EventType,
+    OperationContext,
+    OperationName,
+    OperationToStart,
+    ScheduleId,
+)
+from ._operation import OperationRegistry
+from ._store import OperationEventsProxy, Store
+
+_logger = logging.getLogger(__name__)
+
+
+class AfterEventManager(SingletonInAppStateMixin):
+    """
+    Allows to register an operation to be started after
+    another operation ends the CREATED or UNDONE successfully.
+    """
+
+    app_state_name: str = "after_event_manager"
+
+    def __init__(self, app: FastAPI) -> None:
+        self.app = app
+        self._store = Store.get_from_app_state(app)
+
+    async def register_to_start_after(
+        self,
+        schedule_id: ScheduleId,
+        event_type: EventType,
+        *,
+        to_start: OperationToStart,
+    ) -> None:
+        # ensure operation exists
+        OperationRegistry.get_operation(to_start.operation_name)
+
+        events_proxy = OperationEventsProxy(self._store, schedule_id, event_type)
+        await events_proxy.create_or_update_multiple(
+            {
+                "initial_context": to_start.initial_context,
+                "operation_name": to_start.operation_name,
+            }
+        )
+        _logger.debug(
+            "Registered event_type='%s' to_start='%s' for schedule_id='%s'",
+            event_type,
+            to_start,
+            schedule_id,
+        )
+
+    async def safe_on_event_type(
+        self,
+        event_type: EventType,
+        schedule_id: ScheduleId,
+        operation_name: OperationName,
+        initial_context: OperationContext,
+    ) -> None:
+        with log_context(
+            _logger,
+            logging.DEBUG,
+            f"processing {event_type=} for {schedule_id=} {operation_name=} {initial_context=}",
+            log_duration=True,
+        ):
+            await self._on_event_type(
+                event_type, schedule_id, operation_name, initial_context
+            )
+
+    async def _on_event_type(
+        self,
+        event_type: EventType,
+        schedule_id: ScheduleId,
+        operation_name: OperationName,
+        initial_context: OperationContext,
+    ) -> None:
+        new_schedule_id = await start_operation(
+            self.app, operation_name, initial_context
+        )
+        _logger.debug(
+            "Finished execution of event_type='%s' for schedule_id='%s'. "
+            "Started new_schedule_id='%s' from operation_name='%s' with initial_context='%s'",
+            event_type,
+            schedule_id,
+            new_schedule_id,
+            operation_name,
+            initial_context,
+        )