🎨 Enhance detection of stalled streamed results (#8685)

Author: giancarloromeo

packages/celery-library/src/celery_library/backends/redis.py

@@ -179,7 +179,7 @@ async def push_task_stream_items(
         pipe = self._redis_client_sdk.redis.pipeline()
         pipe.rpush(stream_key, *(r.model_dump_json(by_alias=True) for r in result))
         pipe.hset(
-            stream_meta_key, mapping={"last_update": datetime.now(UTC).isoformat()}
+            stream_meta_key, mapping={"last_update": datetime.now(tz=UTC).isoformat()}
         )
         pipe.expire(stream_key, _CELERY_TASK_STREAM_EXPIRY)
         pipe.expire(stream_meta_key, _CELERY_TASK_STREAM_EXPIRY)
@@ -195,6 +195,16 @@ async def set_task_stream_done(self, task_key: TaskKey) -> None:
             )
         )
 
+    async def set_task_stream_last_update(self, task_key: TaskKey) -> None:
+        stream_meta_key = _build_redis_stream_meta_key(task_key)
+        await handle_redis_returns_union_types(
+            self._redis_client_sdk.redis.hset(
+                name=stream_meta_key,
+                key=_CELERY_TASK_STREAM_LAST_UPDATE_KEY,
+                value=datetime.now(tz=UTC).isoformat(),
+            )
+        )
+
     async def pull_task_stream_items(
         self, task_key: TaskKey, limit: int = 20
     ) -> tuple[list[TaskStreamItem], bool, datetime | None]:

packages/celery-library/src/celery_library/task_manager.py

@@ -186,9 +186,7 @@ async def get_task_status(
     @handle_celery_errors
     async def list_tasks(self, owner_metadata: OwnerMetadata) -> list[Task]:
         with log_context(
-            _logger,
-            logging.DEBUG,
-            msg=f"Listing tasks: {owner_metadata=}",
+            _logger, logging.DEBUG, "Listing tasks: owner_metadata=%s", owner_metadata
         ):
             return await self._task_store.list_tasks(owner_metadata)
 
@@ -204,23 +202,33 @@ async def set_task_progress(
     @handle_celery_errors
     async def set_task_stream_done(self, task_key: TaskKey) -> None:
         with log_context(
-            _logger,
-            logging.DEBUG,
-            msg=f"Set task stream done: {task_key=}",
+            _logger, logging.DEBUG, "Set task stream done: task_key= %s", task_key
         ):
             if not await self.task_exists(task_key):
                 raise TaskNotFoundError(task_key=task_key)
 
             await self._task_store.set_task_stream_done(task_key)
 
+    @handle_celery_errors
+    async def set_task_stream_last_update(self, task_key: TaskKey) -> None:
+        with log_context(
+            _logger, logging.DEBUG, "Set task stream last update: task_key=%s", task_key
+        ):
+            if not await self.task_exists(task_key):
+                raise TaskNotFoundError(task_key=task_key)
+
+            await self._task_store.set_task_stream_last_update(task_key)
+
     @handle_celery_errors
     async def push_task_stream_items(
         self, task_key: TaskKey, *items: TaskStreamItem
     ) -> None:
         with log_context(
             _logger,
             logging.DEBUG,
-            msg=f"Push task stream items: {task_key=} {items=}",
+            "Push task stream items: task_key=%s items=%s",
+            task_key,
+            items,
         ):
             if not await self.task_exists(task_key):
                 raise TaskNotFoundError(task_key=task_key)
@@ -238,7 +246,11 @@ async def pull_task_stream_items(
         with log_context(
             _logger,
             logging.DEBUG,
-            msg=f"Pull task results: {owner_metadata=} {task_uuid=} {offset=} {limit=}",
+            "Pull task results: owner_metadata=%s task_uuid=%s offset=%s limit=%s",
+            owner_metadata,
+            task_uuid,
+            offset,
+            limit,
         ):
             task_key = owner_metadata.model_dump_task_key(task_uuid=task_uuid)
             if not await self.task_exists(task_key):

packages/service-library/src/servicelib/celery/models.py

@@ -213,6 +213,8 @@ async def set_task_progress(
 
     async def set_task_stream_done(self, task_key: TaskKey) -> None: ...
 
+    async def set_task_stream_last_update(self, task_key: TaskKey) -> None: ...
+
     async def push_task_stream_items(
         self, task_key: TaskKey, *item: TaskStreamItem
     ) -> None: ...

packages/service-library/src/servicelib/celery/task_manager.py

@@ -56,4 +56,6 @@ async def pull_task_stream_items(
 
     async def set_task_stream_done(self, task_key: TaskKey) -> None: ...
 
+    async def set_task_stream_last_update(self, task_key: TaskKey) -> None: ...
+
     async def task_exists(self, task_key: TaskKey) -> bool: ...

services/storage/src/simcore_service_storage/api/_worker_tasks/_simcore_s3.py

@@ -84,7 +84,10 @@ async def export_data(
     with log_context(
         _logger,
         logging.INFO,
-        f"'{task_key}' export data (for {user_id=}) fom selection: {paths_to_export}",
+        "export data task (%s) (for user=%s) from selection: %s",
+        task_key,
+        user_id,
+        paths_to_export,
     ):
         dsm = get_dsm_provider(get_app_server(task.app).app).get(
             SimcoreS3DataManager.get_location_id()
@@ -166,6 +169,8 @@ async def search(
             limit=1,  # NOTE: yield items as they come
         ):
             if not items:
+                # NOTE: still set the last update time to signal progress in search
+                await app_server.task_manager.set_task_stream_last_update(task_key)
                 continue
 
             data = [