Merge branch 'master' into make-api-server-celery-concurrency-configurable

Author: bisgaard-itis

.env-devel

@@ -130,6 +130,7 @@ DYNAMIC_SIDECAR_PROMETHEUS_MONITORING_NETWORKS=[]
 DYNAMIC_SIDECAR_PROMETHEUS_SERVICE_LABELS={}
 DYNAMIC_SIDECAR_API_SAVE_RESTORE_STATE_TIMEOUT=01:00:00
 DIRECTOR_V2_TRACING={}
+DIRECTOR_V2_DYNAMIC_SCHEDULER_ENABLED=1
 
 # DYNAMIC_SCHEDULER ----
 DYNAMIC_SCHEDULER_LOGLEVEL=INFO

packages/models-library/src/models_library/functions.py

@@ -190,21 +190,12 @@ class FunctionJobBase(BaseModel):
     function_class: FunctionClass
 
 
-class RegisteredFunctionJobBase(FunctionJobBase):
-    uid: FunctionJobID
-    created_at: datetime.datetime
-
-
 class ProjectFunctionJob(FunctionJobBase):
     function_class: Literal[FunctionClass.PROJECT] = FunctionClass.PROJECT
     project_job_id: ProjectID | None
     job_creation_task_id: TaskID | None
 
 
-class RegisteredProjectFunctionJob(ProjectFunctionJob, RegisteredFunctionJobBase):
-    pass
-
-
 class RegisteredProjectFunctionJobPatch(BaseModel):
     function_class: Literal[FunctionClass.PROJECT] = FunctionClass.PROJECT
     title: str | None
@@ -221,10 +212,6 @@ class SolverFunctionJob(FunctionJobBase):
     job_creation_task_id: TaskID | None
 
 
-class RegisteredSolverFunctionJob(SolverFunctionJob, RegisteredFunctionJobBase):
-    pass
-
-
 class RegisteredSolverFunctionJobPatch(BaseModel):
     function_class: Literal[FunctionClass.SOLVER] = FunctionClass.SOLVER
     title: str | None
@@ -239,10 +226,6 @@ class PythonCodeFunctionJob(FunctionJobBase):
     function_class: Literal[FunctionClass.PYTHON_CODE] = FunctionClass.PYTHON_CODE
 
 
-class RegisteredPythonCodeFunctionJob(PythonCodeFunctionJob, RegisteredFunctionJobBase):
-    pass
-
-
 class RegisteredPythonCodeFunctionJobPatch(BaseModel):
     function_class: Literal[FunctionClass.PYTHON_CODE] = FunctionClass.PYTHON_CODE
     title: str | None
@@ -256,6 +239,24 @@ class RegisteredPythonCodeFunctionJobPatch(BaseModel):
     Field(discriminator="function_class"),
 ]
 
+
+class RegisteredFunctionJobBase(FunctionJobBase):
+    uid: FunctionJobID
+    created_at: datetime.datetime
+
+
+class RegisteredProjectFunctionJob(ProjectFunctionJob, RegisteredFunctionJobBase):
+    pass
+
+
+class RegisteredSolverFunctionJob(SolverFunctionJob, RegisteredFunctionJobBase):
+    pass
+
+
+class RegisteredPythonCodeFunctionJob(PythonCodeFunctionJob, RegisteredFunctionJobBase):
+    pass
+
+
 RegisteredFunctionJob: TypeAlias = Annotated[
     RegisteredProjectFunctionJob
     | RegisteredPythonCodeFunctionJob
@@ -275,6 +276,36 @@ class FunctionJobStatus(BaseModel):
     status: str
 
 
+class RegisteredFunctionJobWithStatusBase(RegisteredFunctionJobBase, FunctionJobBase):
+    status: FunctionJobStatus
+
+
+class RegisteredProjectFunctionJobWithStatus(
+    RegisteredProjectFunctionJob, RegisteredFunctionJobWithStatusBase
+):
+    pass
+
+
+class RegisteredSolverFunctionJobWithStatus(
+    RegisteredSolverFunctionJob, RegisteredFunctionJobWithStatusBase
+):
+    pass
+
+
+class RegisteredPythonCodeFunctionJobWithStatus(
+    RegisteredPythonCodeFunctionJob, RegisteredFunctionJobWithStatusBase
+):
+    pass
+
+
+RegisteredFunctionJobWithStatus: TypeAlias = Annotated[
+    RegisteredProjectFunctionJobWithStatus
+    | RegisteredPythonCodeFunctionJobWithStatus
+    | RegisteredSolverFunctionJobWithStatus,
+    Field(discriminator="function_class"),
+]
+
+
 class FunctionJobCollection(BaseModel):
     """Model for a collection of function jobs"""
 
@@ -309,6 +340,12 @@ class RegisteredFunctionJobDB(FunctionJobDB):
     created: datetime.datetime
 
 
+class RegisteredFunctionJobWithStatusDB(FunctionJobDB):
+    uuid: FunctionJobID
+    created: datetime.datetime
+    status: str
+
+
 class FunctionDB(BaseModel):
     function_class: FunctionClass
     title: str = ""

packages/service-library/src/servicelib/deferred_tasks/__init__.py

@@ -22,14 +22,16 @@
 - `cancel`: (called by the user) [optional]:
     send a message to cancel the current task. A warning will be logged but no call to either
     `on_result` or `on_finished_with_error` will occur.
+-  `on_cancelled` (called by state `ManuallyCancelled`) [optional] {can be overwritten by the user}:
+    called after the cancellation is handled by the worker executing the `run`
 
 
 ## DeferredHandler lifecycle
 
 ```mermaid
 stateDiagram-v2
-    * --> Scheduled: via [start]
-    ** --> ManuallyCancelled: via [cancel]
+    (1) --> Scheduled: via [start]
+    (2) --> ManuallyCancelled: via [cancel]
 
     ManuallyCancelled --> Worker: attempts to cancel task in
 
@@ -41,9 +43,10 @@
     ErrorResult --> FinishedWithError: gives up when out of retries or if cancelled
     Worker --> DeferredResult: success
 
-    DeferredResult --> °: calls [on_result]
-    FinishedWithError --> °°: calls [on_finished_with_error]
-    Worker --> °°°: task cancelled
+    DeferredResult --> (3): calls [on_result]
+    FinishedWithError --> (4): calls [on_finished_with_error]
+    Worker --> Removed*: task cancelled
+    Removed* --> (5): calls [on_cancelled]
 ```
 
 ### States
@@ -57,6 +60,7 @@
 - `FinishedWIthError`: logs error, invokes `on_finished_with_error` and removes the schedule
 - `DeferredResult`: invokes `on_result` and removes the schedule
 - `ManuallyCancelled`: sends message to all instances to cancel. The instance handling the task will cancel the task and remove the schedule
+- `Removed*`: is a fake state that does not exist only used to convey the information that the cancellation event is triggered after removal
 """
 
 from ._base_deferred_handler import (

packages/service-library/src/servicelib/deferred_tasks/_base_deferred_handler.py

@@ -45,6 +45,21 @@ async def get_retries(cls, context: DeferredContext) -> NonNegativeInt:
         assert context  # nosec
         return 0
 
+    @classmethod
+    async def get_retry_delay(
+        cls,
+        context: DeferredContext,
+        remaining_attempts: NonNegativeInt,
+        total_attempts: NonNegativeInt,
+    ) -> timedelta:
+        """
+        returns: the delay between eatch retry attempt (default: 0s)
+        """
+        assert context  # nosec
+        assert remaining_attempts  # nosec
+        assert total_attempts  # nosec
+        return timedelta(seconds=0)
+
     @classmethod
     @abstractmethod
     async def get_timeout(cls, context: DeferredContext) -> timedelta:
@@ -84,6 +99,11 @@ async def on_finished_with_error(
         NOTE: by design the default action is to do nothing
         """
 
+    @classmethod
+    @abstractmethod
+    async def on_cancelled(cls, context: DeferredContext) -> None:
+        """called after handling ``cancel`` request by the copy executing ``run``"""
+
     @classmethod
     async def cancel(cls, task_uid: TaskUID) -> None:
         """cancels a deferred"""

packages/service-library/src/servicelib/deferred_tasks/_deferred_manager.py

@@ -2,7 +2,7 @@
 import inspect
 import logging
 from collections.abc import Awaitable, Callable, Iterable
-from datetime import timedelta
+from datetime import datetime, timedelta
 from enum import Enum
 from typing import Any, Final
 
@@ -118,6 +118,14 @@ def _raise_if_not_type(task_result: Any, expected_types: Iterable[type]) -> None
         raise TypeError(msg)
 
 
+async def _wait_until_future_date(possible_future_date: datetime) -> None:
+    while True:
+        now = arrow.utcnow().datetime
+        if now >= possible_future_date:
+            return
+        await asyncio.sleep(1)
+
+
 class DeferredManager:  # pylint:disable=too-many-instance-attributes
     def __init__(
         self,
@@ -297,18 +305,21 @@ async def __start(
         subclass = self.__get_subclass(class_unique_reference)
         deferred_context = self.__get_deferred_context(start_context)
 
+        retry_count = await subclass.get_retries(deferred_context)
         task_schedule = TaskScheduleModel(
             timeout=await subclass.get_timeout(deferred_context),
-            execution_attempts=await subclass.get_retries(deferred_context) + 1,
+            total_attempts=retry_count,
+            execution_attempts=retry_count + 1,
             class_unique_reference=class_unique_reference,
             start_context=start_context,
             state=TaskState.SCHEDULED,
         )
 
+        await self._task_tracker.save(task_uid, task_schedule)
+
         with log_catch(_logger, reraise=False):
             await subclass.on_created(task_uid, deferred_context)
 
-        await self._task_tracker.save(task_uid, task_schedule)
         _logger.debug("Scheduled task '%s' with entry: %s", task_uid, task_schedule)
         await self.__publish_to_queue(task_uid, _FastStreamRabbitQueue.SCHEDULED)
 
@@ -460,7 +471,29 @@ async def _fs_handle_error_result(  # pylint:disable=method-hidden
             task_schedule.result, TaskResultCancelledError
         ):
             _logger.debug("Schedule retry attempt for task_uid '%s'", task_uid)
-            # does not retry if task was cancelled
+
+            # resilenet wait before retrying
+            if task_schedule.wait_cancellation_until is None:
+                # save the new one
+                subclass = self.__get_subclass(task_schedule.class_unique_reference)
+                deferred_context = self.__get_deferred_context(
+                    task_schedule.start_context
+                )
+                sleep_interval = await subclass.get_retry_delay(
+                    context=deferred_context,
+                    remaining_attempts=task_schedule.execution_attempts,
+                    total_attempts=task_schedule.total_attempts,
+                )
+                task_schedule.wait_cancellation_until = (
+                    arrow.utcnow().datetime + sleep_interval
+                )
+                await self._task_tracker.save(task_uid, task_schedule)
+
+            await _wait_until_future_date(task_schedule.wait_cancellation_until)
+            task_schedule.wait_cancellation_until = None
+            await self._task_tracker.save(task_uid, task_schedule)
+
+            # waiting is done can proceed with retry
             task_schedule.state = TaskState.SUBMIT_TASK
             await self._task_tracker.save(task_uid, task_schedule)
             await self.__publish_to_queue(task_uid, _FastStreamRabbitQueue.SUBMIT_TASK)
@@ -570,6 +603,11 @@ async def _fs_handle_manually_cancelled(  # pylint:disable=method-hidden
         _logger.info("Found and cancelled run for '%s'", task_uid)
         await self.__remove_task(task_uid, task_schedule)
 
+        subclass = self.__get_subclass(task_schedule.class_unique_reference)
+        deferred_context = self.__get_deferred_context(task_schedule.start_context)
+        with log_catch(_logger, reraise=False):
+            await subclass.on_cancelled(deferred_context)
+
     async def __is_present(self, task_uid: TaskUID) -> bool:
         task_schedule: TaskScheduleModel | None = await self._task_tracker.get(task_uid)
         return task_schedule is not None

packages/service-library/src/servicelib/deferred_tasks/_task_schedule.py

@@ -1,7 +1,9 @@
 from datetime import datetime, timedelta
 from enum import Enum
+from typing import Annotated
 
 import arrow
+from common_library.basic_types import DEFAULT_FACTORY
 from pydantic import BaseModel, Field, NonNegativeInt
 
 from ._base_deferred_handler import StartContext
@@ -23,37 +25,63 @@ class TaskState(str, Enum):
 
 
 class TaskScheduleModel(BaseModel):
-    timeout: timedelta = Field(
-        ..., description="Amount of time after which the task execution will time out"
-    )
-    class_unique_reference: ClassUniqueReference = Field(
-        ...,
-        description="reference to the class containing the code and handlers for the execution of the task",
-    )
-    start_context: StartContext = Field(
-        ...,
-        description="data used to assemble the ``StartContext``",
-    )
-
-    state: TaskState = Field(
-        ..., description="represents the execution step of the task"
-    )
-
-    execution_attempts: NonNegativeInt = Field(
-        ...,
-        description="remaining attempts to run the code, only retries if this is > 0",
-    )
-
-    time_started: datetime = Field(
-        default_factory=lambda: arrow.utcnow().datetime,
-        description="time when task schedule was created, used for statistics",
-    )
-
-    result: TaskExecutionResult | None = Field(
-        default=None,
-        description=(
-            f"Populated by {TaskState.WORKER}. It always has a value after worker handles it."
-            "Will be used "
+    timeout: Annotated[
+        timedelta,
+        Field(
+            description="Amount of time after which the task execution will time out"
         ),
-        discriminator="result_type",
-    )
+    ]
+    class_unique_reference: Annotated[
+        ClassUniqueReference,
+        Field(
+            description="reference to the class containing the code and handlers for the execution of the task",
+        ),
+    ]
+    start_context: Annotated[
+        StartContext,
+        Field(
+            description="data used to assemble the ``StartContext``",
+        ),
+    ]
+
+    state: Annotated[
+        TaskState, Field(description="represents the execution step of the task")
+    ]
+
+    total_attempts: Annotated[
+        NonNegativeInt,
+        Field(
+            description="maximum number of attempts before giving up (0 means no retries)"
+        ),
+    ]
+
+    execution_attempts: Annotated[
+        NonNegativeInt,
+        Field(
+            description="remaining attempts to run the code, only retries if this is > 0",
+        ),
+    ]
+
+    wait_cancellation_until: Annotated[
+        datetime | None,
+        Field(description="when set has to wait till this before cancelling the task"),
+    ] = None
+
+    time_started: Annotated[
+        datetime,
+        Field(
+            default_factory=lambda: arrow.utcnow().datetime,
+            description="time when task schedule was created, used for statistics",
+        ),
+    ] = DEFAULT_FACTORY
+
+    result: Annotated[
+        TaskExecutionResult | None,
+        Field(
+            description=(
+                f"Populated by {TaskState.WORKER}. It always has a value after worker handles it."
+                "Will be used "
+            ),
+            discriminator="result_type",
+        ),
+    ] = None