Merge branch 'master' into 1973-add-celery-worker-to-api-server

Author: bisgaard-itis

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

@@ -13,11 +13,11 @@
 from models_library.users import UserID
 from models_library.utils.enums import StrAutoEnum
 from pydantic import BaseModel, ConfigDict, Field
-from servicelib.celery.models import TaskID
 
 from .projects import ProjectID
 from .utils.change_case import snake_to_camel
 
+TaskID: TypeAlias = str
 FunctionID: TypeAlias = UUID
 FunctionJobID: TypeAlias = UUID
 FileID: TypeAlias = UUID
@@ -206,9 +206,11 @@ class RegisteredProjectFunctionJob(ProjectFunctionJob, RegisteredFunctionJobBase
 
 
 class RegisteredProjectFunctionJobPatch(BaseModel):
-    function_class: FunctionClass
+    function_class: Literal[FunctionClass.PROJECT] = FunctionClass.PROJECT
     title: str | None
     description: str | None
+    inputs: FunctionInputs
+    outputs: FunctionOutputs
     project_job_id: ProjectID | None
     job_creation_task_id: TaskID | None
 
@@ -224,9 +226,11 @@ class RegisteredSolverFunctionJob(SolverFunctionJob, RegisteredFunctionJobBase):
 
 
 class RegisteredSolverFunctionJobPatch(BaseModel):
-    function_class: FunctionClass
+    function_class: Literal[FunctionClass.SOLVER] = FunctionClass.SOLVER
     title: str | None
     description: str | None
+    inputs: FunctionInputs
+    outputs: FunctionOutputs
     solver_job_id: ProjectID | None
     job_creation_task_id: TaskID | None
 
@@ -240,8 +244,10 @@ class RegisteredPythonCodeFunctionJob(PythonCodeFunctionJob, RegisteredFunctionJ
 
 
 class RegisteredPythonCodeFunctionJobPatch(BaseModel):
-    function_class: FunctionClass
+    function_class: Literal[FunctionClass.PYTHON_CODE] = FunctionClass.PYTHON_CODE
     title: str | None
+    inputs: FunctionInputs
+    outputs: FunctionOutputs
     description: str | None
 
 

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

@@ -158,3 +158,8 @@ class FunctionJobCollectionsExecuteApiAccessDeniedError(FunctionBaseError):
         "User {user_id} does not have the permission to execute function job collections"
     )
     status_code: int = 403  # Forbidden
+
+
+class FunctionJobPatchModelIncompatibleError(FunctionBaseError):
+    msg_template = "Incompatible patch model for Function '{function_id}' in product '{product_name}'."
+    status_code: int = 422

packages/service-library/src/servicelib/aiohttp/long_running_tasks/_manager.py

@@ -1,12 +1,12 @@
 from aiohttp import web
 
-from ...long_running_tasks.base_long_running_manager import BaseLongRunningManager
+from ...long_running_tasks.manager import LongRunningManager
 from ...long_running_tasks.models import TaskContext
 from ._constants import APP_LONG_RUNNING_MANAGER_KEY
 from ._request import get_task_context
 
 
-class AiohttpLongRunningManager(BaseLongRunningManager):
+class AiohttpLongRunningManager(LongRunningManager):
 
     @staticmethod
     def get_task_context(request: web.Request) -> TaskContext:

packages/service-library/src/servicelib/fastapi/long_running_tasks/_manager.py

@@ -1,10 +1,10 @@
 from fastapi import Request
 
-from ...long_running_tasks.base_long_running_manager import BaseLongRunningManager
+from ...long_running_tasks.manager import LongRunningManager
 from ...long_running_tasks.models import TaskContext
 
 
-class FastAPILongRunningManager(BaseLongRunningManager):
+class FastAPILongRunningManager(LongRunningManager):
     @staticmethod
     def get_task_context(request: Request) -> TaskContext:
         _ = request

packages/service-library/src/servicelib/long_running_tasks/_rpc_server.py

@@ -15,15 +15,15 @@
 _logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from .base_long_running_manager import BaseLongRunningManager
+    from .manager import LongRunningManager
 
 
 router = RPCRouter()
 
 
 @router.expose(reraise_if_error_type=(BaseLongRunningError,))
 async def start_task(
-    long_running_manager: "BaseLongRunningManager",
+    long_running_manager: "LongRunningManager",
     *,
     registered_task_name: RegisteredTaskName,
     unique: bool = False,
@@ -44,7 +44,7 @@ async def start_task(
 
 @router.expose(reraise_if_error_type=(BaseLongRunningError,))
 async def list_tasks(
-    long_running_manager: "BaseLongRunningManager", *, task_context: TaskContext
+    long_running_manager: "LongRunningManager", *, task_context: TaskContext
 ) -> list[TaskBase]:
     return await long_running_manager.tasks_manager.list_tasks(
         with_task_context=task_context
@@ -53,7 +53,7 @@ async def list_tasks(
 
 @router.expose(reraise_if_error_type=(BaseLongRunningError,))
 async def get_task_status(
-    long_running_manager: "BaseLongRunningManager",
+    long_running_manager: "LongRunningManager",
     *,
     task_context: TaskContext,
     task_id: TaskId,
@@ -65,7 +65,7 @@ async def get_task_status(
 
 @router.expose(reraise_if_error_type=(BaseLongRunningError, RPCTransferrableTaskError))
 async def get_task_result(
-    long_running_manager: "BaseLongRunningManager",
+    long_running_manager: "LongRunningManager",
     *,
     task_context: TaskContext,
     task_id: TaskId,
@@ -94,7 +94,7 @@ async def get_task_result(
 
 @router.expose(reraise_if_error_type=(BaseLongRunningError,))
 async def remove_task(
-    long_running_manager: "BaseLongRunningManager",
+    long_running_manager: "LongRunningManager",
     *,
     task_context: TaskContext,
     task_id: TaskId,

packages/service-library/src/servicelib/long_running_tasks/long_running_client_helper.py

@@ -0,0 +1,46 @@
+import logging
+
+import redis.asyncio as aioredis
+from settings_library.redis import RedisDatabase, RedisSettings
+
+from ..logging_utils import log_context
+from ..redis._client import RedisClientSDK
+from .models import LRTNamespace
+
+_logger = logging.getLogger(__name__)
+
+
+class LongRunningClientHelper:
+    def __init__(self, redis_settings: RedisSettings):
+        self.redis_settings = redis_settings
+
+        self._client: RedisClientSDK | None = None
+
+    async def setup(self) -> None:
+        self._client = RedisClientSDK(
+            self.redis_settings.build_redis_dsn(RedisDatabase.LONG_RUNNING_TASKS),
+            client_name="long_running_tasks_cleanup_client",
+        )
+        await self._client.setup()
+
+    async def shutdown(self) -> None:
+        if self._client:
+            await self._client.shutdown()
+
+    @property
+    def _redis(self) -> aioredis.Redis:
+        assert self._client  # nosec
+        return self._client.redis
+
+    async def cleanup(self, lrt_namespace: LRTNamespace) -> None:
+        """removes Redis keys associated to the LRTNamespace if they exist"""
+        keys_to_remove: list[str] = [
+            x async for x in self._redis.scan_iter(f"{lrt_namespace}*")
+        ]
+        with log_context(
+            _logger,
+            logging.DEBUG,
+            msg=f"Removing {keys_to_remove=} from Redis for {lrt_namespace=}",
+        ):
+            if len(keys_to_remove) > 0:
+                await self._redis.delete(*keys_to_remove)

packages/service-library/src/servicelib/long_running_tasks/base_long_running_manager.py renamed to packages/service-library/src/servicelib/long_running_tasks/manager.py

@@ -11,7 +11,7 @@
 from .task import TasksManager
 
 
-class BaseLongRunningManager(ABC):
+class LongRunningManager(ABC):
     """
     Provides a commond inteface for aiohttp and fastapi services
     """

packages/service-library/src/servicelib/long_running_tasks/task.py

@@ -11,6 +11,7 @@
 from common_library.async_tools import cancel_wait_task
 from models_library.api_schemas_long_running_tasks.base import TaskProgress
 from pydantic import NonNegativeFloat, PositiveFloat
+from servicelib.utils import limited_gather
 from settings_library.redis import RedisDatabase, RedisSettings
 from tenacity import (
     AsyncRetrying,
@@ -21,7 +22,7 @@
 
 from ..background_task import create_periodic_task
 from ..logging_errors import create_troubleshootting_log_kwargs
-from ..logging_utils import log_context
+from ..logging_utils import log_catch, log_context
 from ..redis import RedisClientSDK, exclusive
 from ._redis_store import RedisStore
 from ._serialization import dumps
@@ -50,6 +51,7 @@
 _STATUS_UPDATE_CHECK_INTERNAL: Final[datetime.timedelta] = datetime.timedelta(seconds=1)
 _MAX_EXCLUSIVE_TASK_CANCEL_TIMEOUT: Final[NonNegativeFloat] = 5
 _TASK_REMOVAL_MAX_WAIT: Final[NonNegativeFloat] = 60
+_PARALLEL_TASKS_CANCELLATION: Final[int] = 5
 
 AllowedErrrors: TypeAlias = tuple[type[BaseException], ...]
 
@@ -205,34 +207,40 @@ async def setup(self) -> None:
         await self._started_event_task_tasks_monitor.wait()
 
     async def teardown(self) -> None:
-        # ensure all created tasks are cancelled
-        for tracked_task in await self._tasks_data.list_tasks_data():
-            with suppress(TaskNotFoundError):
+        # stop cancelled_tasks_removal
+        if self._task_cancelled_tasks_removal:
+            await cancel_wait_task(self._task_cancelled_tasks_removal)
+
+        # stopping only tasks that are handled by this manager
+        # otherwise it will cancel long running tasks that were running in diffierent processes
+        async def _remove_local_task(task_data: TaskData) -> None:
+            with log_catch(_logger, reraise=False):
                 await self.remove_task(
-                    tracked_task.task_id,
-                    tracked_task.task_context,
-                    wait_for_removal=True,
+                    task_data.task_id,
+                    task_data.task_context,
+                    wait_for_removal=False,
                 )
+                await self._attempt_to_remove_local_task(task_data.task_id)
+
+        await limited_gather(
+            *[
+                _remove_local_task(tracked_task)
+                for task_id in self._created_tasks
+                if (tracked_task := await self._tasks_data.get_task_data(task_id))
+                is not None
+            ],
+            log=_logger,
+            limit=_PARALLEL_TASKS_CANCELLATION,
+        )
 
-        for task in self._created_tasks.values():
-            _logger.warning(
-                "Task %s was not completed before shutdown, cancelling it",
-                task.get_name(),
-            )
-            await cancel_wait_task(task)
-
-        # stale_tasks_monitor
+        # stop stale_tasks_monitor
         if self._task_stale_tasks_monitor:
             await cancel_wait_task(
                 self._task_stale_tasks_monitor,
                 max_delay=_MAX_EXCLUSIVE_TASK_CANCEL_TIMEOUT,
             )
 
-        # cancelled_tasks_removal
-        if self._task_cancelled_tasks_removal:
-            await cancel_wait_task(self._task_cancelled_tasks_removal)
-
-        # tasks_monitor
+        # stop tasks_monitor
         if self._task_tasks_monitor:
             await cancel_wait_task(self._task_tasks_monitor)
 

packages/service-library/tests/long_running_tasks/conftest.py

@@ -11,8 +11,8 @@
 from faker import Faker
 from pytest_mock import MockerFixture
 from servicelib.logging_utils import log_catch
-from servicelib.long_running_tasks.base_long_running_manager import (
-    BaseLongRunningManager,
+from servicelib.long_running_tasks.manager import (
+    LongRunningManager,
 )
 from servicelib.long_running_tasks.models import LRTNamespace, TaskContext
 from servicelib.long_running_tasks.task import TasksManager
@@ -24,7 +24,7 @@
 _logger = logging.getLogger(__name__)
 
 
-class _TestingLongRunningManager(BaseLongRunningManager):
+class _TestingLongRunningManager(LongRunningManager):
     @staticmethod
     def get_task_context(request) -> TaskContext:
         _ = request
@@ -37,16 +37,16 @@ async def get_long_running_manager(
 ) -> AsyncIterator[
     Callable[
         [RedisSettings, RabbitSettings, LRTNamespace | None],
-        Awaitable[BaseLongRunningManager],
+        Awaitable[LongRunningManager],
     ]
 ]:
-    managers: list[BaseLongRunningManager] = []
+    managers: list[LongRunningManager] = []
 
     async def _(
         redis_settings: RedisSettings,
         rabbit_settings: RabbitSettings,
         lrt_namespace: LRTNamespace | None,
-    ) -> BaseLongRunningManager:
+    ) -> LongRunningManager:
         manager = _TestingLongRunningManager(
             stale_task_check_interval=timedelta(seconds=TEST_CHECK_STALE_INTERVAL_S),
             stale_task_detect_timeout=timedelta(seconds=TEST_CHECK_STALE_INTERVAL_S),