🎨 Store and retrieve task_name when listing Celery tasks (#7538)

Co-authored-by: Odei Maiz <[email protected]>

Author: giancarloromeo

packages/models-library/src/models_library/api_schemas_rpc_async_jobs/async_jobs.py

@@ -1,12 +1,16 @@
-from typing import Any, TypeAlias
+from typing import Annotated, Any, TypeAlias
 from uuid import UUID
 
-from models_library.users import UserID
-from pydantic import BaseModel
+from pydantic import BaseModel, StringConstraints
 
+from ..products import ProductName
 from ..progress_bar import ProgressReport
+from ..users import UserID
 
 AsyncJobId: TypeAlias = UUID
+AsyncJobName: TypeAlias = Annotated[
+    str, StringConstraints(strip_whitespace=True, min_length=1)
+]
 
 
 class AsyncJobStatus(BaseModel):
@@ -21,6 +25,7 @@ class AsyncJobResult(BaseModel):
 
 class AsyncJobGet(BaseModel):
     job_id: AsyncJobId
+    job_name: AsyncJobName
 
 
 class AsyncJobAbort(BaseModel):
@@ -31,5 +36,5 @@ class AsyncJobAbort(BaseModel):
 class AsyncJobNameData(BaseModel):
     """Data for controlling access to an async job"""
 
+    product_name: ProductName
     user_id: UserID
-    product_name: str

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

@@ -26,7 +26,7 @@ async def _task_progress_cb(
 ) -> None:
     worker = get_celery_worker(task.app)
     assert task.name  # nosec
-    await worker.set_progress(
+    await worker.set_task_progress(
         task_id=task_id,
         report=report,
     )
@@ -87,7 +87,7 @@ async def export_data(
 
         async def _progress_cb(report: ProgressReport) -> None:
             assert task.name  # nosec
-            await get_celery_worker(task.app).set_progress(task_id, report)
+            await get_celery_worker(task.app).set_task_progress(task_id, report)
             _logger.debug("'%s' progress %s", task_id, report.percent_value)
 
         async with ProgressBarData(

services/storage/src/simcore_service_storage/api/rest/_files.py

@@ -35,7 +35,7 @@
     UploadLinks,
 )
 from ...modules.celery.client import CeleryTaskClient
-from ...modules.celery.models import TaskUUID
+from ...modules.celery.models import TaskMetadata, TaskUUID
 from ...simcore_s3_dsm import SimcoreS3DataManager
 from .._worker_tasks._files import complete_upload_file as remote_complete_upload_file
 from .dependencies.celery import get_celery_client
@@ -284,8 +284,10 @@ async def complete_upload_file(
         user_id=query_params.user_id,
         product_name=_UNDEFINED_PRODUCT_NAME_FOR_WORKER_TASKS,  # NOTE: I would need to change the API here
     )
-    task_uuid = await celery_client.send_task(
-        remote_complete_upload_file.__name__,
+    task_uuid = await celery_client.submit_task(
+        TaskMetadata(
+            name=remote_complete_upload_file.__name__,
+        ),
         task_context=async_job_name_data.model_dump(),
         user_id=async_job_name_data.user_id,
         location_id=location_id,

services/storage/src/simcore_service_storage/api/rpc/_async_jobs.py

@@ -127,10 +127,12 @@ async def list_jobs(
     _ = filter_
     assert app  # nosec
     try:
-        task_uuids = await get_celery_client(app).get_task_uuids(
+        tasks = await get_celery_client(app).list_tasks(
             task_context=job_id_data.model_dump(),
         )
     except CeleryError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc
 
-    return [AsyncJobGet(job_id=task_uuid) for task_uuid in task_uuids]
+    return [
+        AsyncJobGet(job_id=task.uuid, job_name=task.metadata.name) for task in tasks
+    ]

services/storage/src/simcore_service_storage/api/rpc/_paths.py

@@ -10,6 +10,7 @@
 from servicelib.rabbitmq import RPCRouter
 
 from ...modules.celery import get_celery_client
+from ...modules.celery.models import TaskMetadata
 from .._worker_tasks._paths import compute_path_size as remote_compute_path_size
 from .._worker_tasks._paths import delete_paths as remote_delete_paths
 
@@ -24,15 +25,18 @@ async def compute_path_size(
     location_id: LocationID,
     path: Path,
 ) -> AsyncJobGet:
-    task_uuid = await get_celery_client(app).send_task(
-        remote_compute_path_size.__name__,
+    task_name = remote_compute_path_size.__name__
+    task_uuid = await get_celery_client(app).submit_task(
+        task_metadata=TaskMetadata(
+            name=task_name,
+        ),
         task_context=job_id_data.model_dump(),
         user_id=job_id_data.user_id,
         location_id=location_id,
         path=path,
     )
 
-    return AsyncJobGet(job_id=task_uuid)
+    return AsyncJobGet(job_id=task_uuid, job_name=task_name)
 
 
 @router.expose(reraise_if_error_type=None)
@@ -42,11 +46,14 @@ async def delete_paths(
     location_id: LocationID,
     paths: set[Path],
 ) -> AsyncJobGet:
-    task_uuid = await get_celery_client(app).send_task(
-        remote_delete_paths.__name__,
+    task_name = remote_delete_paths.__name__
+    task_uuid = await get_celery_client(app).submit_task(
+        task_metadata=TaskMetadata(
+            name=task_name,
+        ),
         task_context=job_id_data.model_dump(),
         user_id=job_id_data.user_id,
         location_id=location_id,
         paths=paths,
     )
-    return AsyncJobGet(job_id=task_uuid)
+    return AsyncJobGet(job_id=task_uuid, job_name=task_name)

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

@@ -20,28 +20,32 @@ async def copy_folders_from_project(
     job_id_data: AsyncJobNameData,
     body: FoldersBody,
 ) -> AsyncJobGet:
-    task_uuid = await get_celery_client(app).send_task(
-        deep_copy_files_from_project.__name__,
+    task_name = deep_copy_files_from_project.__name__
+    task_uuid = await get_celery_client(app).submit_task(
+        task_metadata=TaskMetadata(
+            name=task_name,
+        ),
         task_context=job_id_data.model_dump(),
         user_id=job_id_data.user_id,
         body=body,
     )
 
-    return AsyncJobGet(job_id=task_uuid)
+    return AsyncJobGet(job_id=task_uuid, job_name=task_name)
 
 
 @router.expose()
 async def start_export_data(
     app: FastAPI, job_id_data: AsyncJobNameData, paths_to_export: list[PathToExport]
 ) -> AsyncJobGet:
-    task_uuid = await get_celery_client(app).send_task(
-        export_data.__name__,
-        task_context=job_id_data.model_dump(),
+    task_name = export_data.__name__
+    task_uuid = await get_celery_client(app).submit_task(
         task_metadata=TaskMetadata(
+            name=task_name,
             ephemeral=False,
             queue=TasksQueue.CPU_BOUND,
         ),
+        task_context=job_id_data.model_dump(),
         user_id=job_id_data.user_id,
         paths_to_export=paths_to_export,
     )
-    return AsyncJobGet(job_id=task_uuid)
+    return AsyncJobGet(job_id=task_uuid, job_name=task_name)

services/storage/src/simcore_service_storage/modules/celery/backends/_redis.py

@@ -1,12 +1,20 @@
+import contextlib
 import logging
 from datetime import timedelta
 from typing import Final
 
-from celery.result import AsyncResult  # type: ignore[import-untyped]
 from models_library.progress_bar import ProgressReport
+from pydantic import ValidationError
 from servicelib.redis._client import RedisClientSDK
 
-from ..models import TaskContext, TaskID, TaskMetadata, TaskUUID, build_task_id_prefix
+from ..models import (
+    Task,
+    TaskContext,
+    TaskID,
+    TaskMetadata,
+    TaskUUID,
+    build_task_id_prefix,
+)
 
 _CELERY_TASK_INFO_PREFIX: Final[str] = "celery-task-info-"
 _CELERY_TASK_ID_KEY_ENCODING = "utf-8"
@@ -26,26 +34,64 @@ class RedisTaskInfoStore:
     def __init__(self, redis_client_sdk: RedisClientSDK) -> None:
         self._redis_client_sdk = redis_client_sdk
 
-    async def exists(self, task_id: TaskID) -> bool:
+    async def create_task(
+        self,
+        task_id: TaskID,
+        task_metadata: TaskMetadata,
+        expiry: timedelta,
+    ) -> None:
+        task_key = _build_key(task_id)
+        await self._redis_client_sdk.redis.hset(
+            name=task_key,
+            key=_CELERY_TASK_METADATA_KEY,
+            value=task_metadata.model_dump_json(),
+        )  # type: ignore
+        await self._redis_client_sdk.redis.expire(
+            task_key,
+            expiry,
+        )
+
+    async def exists_task(self, task_id: TaskID) -> bool:
         n = await self._redis_client_sdk.redis.exists(_build_key(task_id))
         assert isinstance(n, int)  # nosec
         return n > 0
 
-    async def get_metadata(self, task_id: TaskID) -> TaskMetadata | None:
-        result = await self._redis_client_sdk.redis.hget(_build_key(task_id), _CELERY_TASK_METADATA_KEY)  # type: ignore
-        return TaskMetadata.model_validate_json(result) if result else None
+    async def get_task_metadata(self, task_id: TaskID) -> TaskMetadata | None:
+        raw_result = await self._redis_client_sdk.redis.hget(_build_key(task_id), _CELERY_TASK_METADATA_KEY)  # type: ignore
+        if not raw_result:
+            return None
+
+        try:
+            return TaskMetadata.model_validate_json(raw_result)
+        except ValidationError as exc:
+            _logger.debug(
+                "Failed to deserialize task metadata for task %s: %s", task_id, f"{exc}"
+            )
+            return None
 
-    async def get_progress(self, task_id: TaskID) -> ProgressReport | None:
-        result = await self._redis_client_sdk.redis.hget(_build_key(task_id), _CELERY_TASK_PROGRESS_KEY)  # type: ignore
-        return ProgressReport.model_validate_json(result) if result else None
+    async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None:
+        raw_result = await self._redis_client_sdk.redis.hget(_build_key(task_id), _CELERY_TASK_PROGRESS_KEY)  # type: ignore
+        if not raw_result:
+            return None
 
-    async def get_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
+        try:
+            return ProgressReport.model_validate_json(raw_result)
+        except ValidationError as exc:
+            _logger.debug(
+                "Failed to deserialize task progress for task %s: %s", task_id, f"{exc}"
+            )
+            return None
+
+    async def list_tasks(self, task_context: TaskContext) -> list[Task]:
         search_key = (
             _CELERY_TASK_INFO_PREFIX
             + build_task_id_prefix(task_context)
             + _CELERY_TASK_ID_KEY_SEPARATOR
         )
-        keys = set()
+        search_key_len = len(search_key)
+
+        keys: list[str] = []
+        pipeline = self._redis_client_sdk.redis.pipeline()
         async for key in self._redis_client_sdk.redis.scan_iter(
             match=search_key + "*", count=_CELERY_TASK_SCAN_COUNT_PER_BATCH
         ):
@@ -55,27 +101,31 @@ async def get_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
                 if isinstance(key, bytes)
                 else key
             )
-            keys.add(TaskUUID(_key.removeprefix(search_key)))
-        return keys
+            keys.append(_key)
+            pipeline.hget(_key, _CELERY_TASK_METADATA_KEY)
 
-    async def remove(self, task_id: TaskID) -> None:
-        await self._redis_client_sdk.redis.delete(_build_key(task_id))
-        AsyncResult(task_id).forget()
+        results = await pipeline.execute()
 
-    async def set_metadata(
-        self, task_id: TaskID, task_metadata: TaskMetadata, expiry: timedelta
-    ) -> None:
-        await self._redis_client_sdk.redis.hset(
-            name=_build_key(task_id),
-            key=_CELERY_TASK_METADATA_KEY,
-            value=task_metadata.model_dump_json(),
-        )  # type: ignore
-        await self._redis_client_sdk.redis.expire(
-            _build_key(task_id),
-            expiry,
-        )
+        tasks = []
+        for key, raw_metadata in zip(keys, results, strict=True):
+            if raw_metadata is None:
+                continue
+
+            with contextlib.suppress(ValidationError):
+                task_metadata = TaskMetadata.model_validate_json(raw_metadata)
+                tasks.append(
+                    Task(
+                        uuid=TaskUUID(key[search_key_len:]),
+                        metadata=task_metadata,
+                    )
+                )
+
+        return tasks
+
+    async def remove_task(self, task_id: TaskID) -> None:
+        await self._redis_client_sdk.redis.delete(_build_key(task_id))
 
-    async def set_progress(self, task_id: TaskID, report: ProgressReport) -> None:
+    async def set_task_progress(self, task_id: TaskID, report: ProgressReport) -> None:
         await self._redis_client_sdk.redis.hset(
             name=_build_key(task_id),
             key=_CELERY_TASK_PROGRESS_KEY,