progress updates are now async and saved to storage

Author: Andrei Neagu

packages/models-library/src/models_library/api_schemas_long_running_tasks/base.py

@@ -1,4 +1,5 @@
 import logging
+from collections.abc import Awaitable, Callable
 from typing import Annotated, TypeAlias
 
 from pydantic import BaseModel, Field, field_validator, validate_call
@@ -22,8 +23,16 @@ class TaskProgress(BaseModel):
     message: ProgressMessage = ""
     percent: ProgressPercent = 0.0
 
+    # used to propagate progress updates internally
+    _update_callback: Callable[["TaskProgress"], Awaitable[None]] | None = None
+
+    def set_update_callback(
+        self, callback: Callable[["TaskProgress"], Awaitable[None]]
+    ) -> None:
+        self._update_callback = callback
+
     @validate_call
-    def update(
+    async def update(
         self,
         *,
         message: ProgressMessage | None = None,
@@ -40,6 +49,14 @@ def update(
 
         _logger.debug("Progress update: %s", f"{self}")
 
+        if self._update_callback is not None:
+            await self._update_callback(self)
+        else:
+            _logger.warning(
+                "No update callback set for TaskProgress %s, progress will not be propagated",
+                self.task_id,
+            )
+
     @classmethod
     def create(cls, task_id: TaskId | None = None) -> "TaskProgress":
         return cls(task_id=task_id)

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

@@ -1,5 +1,6 @@
 import asyncio
 import datetime
+import functools
 import inspect
 import logging
 import traceback
@@ -378,6 +379,16 @@ def _get_task_id(self, task_name: str, *, is_unique: bool) -> TaskId:
         unique_part = "unique" if is_unique else f"{uuid4()}"
         return f"{self.namespace}.{task_name}.{unique_part}"
 
+    async def _update_progress(
+        self,
+        task_id: TaskId,
+        task_context: TaskContext | None,
+        task_progress: TaskProgress,
+    ) -> None:
+        tracked_data = await self._get_tracked_task(task_id, task_context)
+        tracked_data.task_progress = task_progress
+        await self._tasks_data.set_task_data(task_id=task_id, value=tracked_data)
+
     async def start_task(
         self,
         registered_task_name: RegisteredTaskName,
@@ -410,14 +421,18 @@ async def start_task(
             )
 
         task_progress = TaskProgress.create(task_id=task_id)
+        # set update callback
+        task_progress.set_update_callback(
+            functools.partial(self._update_progress, task_id, task_context)
+        )
 
         # bind the task with progress 0 and 1
         async def _progress_task(progress: TaskProgress, handler: TaskProtocol):
-            progress.update(message="starting", percent=0)
+            await progress.update(message="starting", percent=0)
             try:
                 return await handler(progress, **task_kwargs)
             finally:
-                progress.update(message="finished", percent=1)
+                await progress.update(message="finished", percent=1)
 
         async_task = asyncio.create_task(
             _progress_task(task_progress, task), name=task_name

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

@@ -36,7 +36,7 @@ async def _string_list_task(
     for index in range(num_strings):
         generated_strings.append(f"{index}")
         await asyncio.sleep(sleep_time)
-        progress.update(message="generated item", percent=index / num_strings)
+        await progress.update(message="generated item", percent=index / num_strings)
         if fail:
             msg = "We were asked to fail!!"
             raise RuntimeError(msg)

packages/service-library/tests/aiohttp/long_running_tasks/test_long_running_tasks.py

@@ -33,6 +33,10 @@
     "redis",
 ]
 
+pytest_simcore_ops_services_selection = [
+    "redis-commander",
+]
+
 
 @pytest.fixture
 def app(

packages/service-library/tests/fastapi/long_running_tasks/test_long_running_tasks.py

@@ -52,7 +52,7 @@ async def _string_list_task(
     for index in range(num_strings):
         generated_strings.append(f"{index}")
         await asyncio.sleep(sleep_time)
-        progress.update(message="generated item", percent=index / num_strings)
+        await progress.update(message="generated item", percent=index / num_strings)
         if fail:
             msg = "We were asked to fail!!"
             raise RuntimeError(msg)

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

@@ -51,7 +51,7 @@ async def a_background_task(
     """sleeps and raises an error or returns 42"""
     for i in range(total_sleep):
         await asyncio.sleep(1)
-        progress.update(percent=(i + 1) / total_sleep)
+        await progress.update(percent=(i + 1) / total_sleep)
     if raise_when_finished:
         msg = "raised this error as instructed"
         raise RuntimeError(msg)

services/director-v2/src/simcore_service_director_v2/api/routes/dynamic_scheduler.py

@@ -106,7 +106,7 @@ async def _task_remove_service_containers(
         async def _progress_callback(
             message: ProgressMessage, percent: ProgressPercent | None, _: TaskId
         ) -> None:
-            progress.update(message=message, percent=percent)
+            await progress.update(message=message, percent=percent)
 
         await dynamic_sidecars_scheduler.remove_service_containers(
             node_uuid=node_uuid, progress_callback=_progress_callback
@@ -171,7 +171,7 @@ async def _task_save_service_state(
         async def _progress_callback(
             message: ProgressMessage, percent: ProgressPercent | None, _: TaskId
         ) -> None:
-            progress.update(message=message, percent=percent)
+            await progress.update(message=message, percent=percent)
 
         await dynamic_sidecars_scheduler.save_service_state(
             node_uuid=node_uuid, progress_callback=_progress_callback
@@ -218,7 +218,7 @@ async def _task_push_service_outputs(
         async def _progress_callback(
             message: ProgressMessage, percent: ProgressPercent | None, _: TaskId
         ) -> None:
-            progress.update(message=message, percent=percent)
+            await progress.update(message=message, percent=percent)
 
         await dynamic_sidecars_scheduler.push_service_outputs(
             node_uuid=node_uuid, progress_callback=_progress_callback

services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_events_utils.py

@@ -217,7 +217,7 @@ async def service_remove_sidecar_proxy_docker_networks_and_volumes(
     if set_were_state_and_outputs_saved is not None:
         scheduler_data.dynamic_sidecar.were_state_and_outputs_saved = True
 
-    task_progress.update(
+    await task_progress.update(
         message="removing dynamic sidecar stack", percent=ProgressPercent(0.1)
     )
     await remove_dynamic_sidecar_stack(
@@ -232,7 +232,7 @@ async def service_remove_sidecar_proxy_docker_networks_and_volumes(
             node_id=scheduler_data.node_uuid,
         )
 
-    task_progress.update(message="removing network", percent=ProgressPercent(0.2))
+    await task_progress.update(message="removing network", percent=ProgressPercent(0.2))
     await remove_dynamic_sidecar_network(scheduler_data.dynamic_sidecar_network_name)
 
     if scheduler_data.dynamic_sidecar.were_state_and_outputs_saved:
@@ -243,7 +243,7 @@ async def service_remove_sidecar_proxy_docker_networks_and_volumes(
             )
         else:
             # Remove all dy-sidecar associated volumes from node
-            task_progress.update(
+            await task_progress.update(
                 message="removing volumes", percent=ProgressPercent(0.3)
             )
             with log_context(_logger, logging.DEBUG, f"removing volumes '{node_uuid}'"):
@@ -265,7 +265,7 @@ async def service_remove_sidecar_proxy_docker_networks_and_volumes(
         scheduler_data.service_name,
     )
 
-    task_progress.update(
+    await task_progress.update(
         message="removing project networks", percent=ProgressPercent(0.8)
     )
     used_projects_networks = await get_projects_networks_containers(
@@ -284,7 +284,7 @@ async def service_remove_sidecar_proxy_docker_networks_and_volumes(
     await app.state.dynamic_sidecar_scheduler.scheduler.remove_service_from_observation(
         scheduler_data.node_uuid
     )
-    task_progress.update(
+    await task_progress.update(
         message="finished removing resources", percent=ProgressPercent(1)
     )
 

services/director-v2/tests/unit/test_modules_dynamic_sidecar_docker_service_specs_sidecar.py

@@ -55,6 +55,11 @@
     "RABBIT_PORT",
     "RABBIT_SECURE",
     "RABBIT_USER",
+    "REDIS_HOST",
+    "REDIS_PASSWORD",
+    "REDIS_PORT",
+    "REDIS_SECURE",
+    "REDIS_USER",
     "S3_ACCESS_KEY",
     "S3_BUCKET_NAME",
     "S3_ENDPOINT",