continue

Author: giancarloromeo

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

@@ -1,12 +1,14 @@
 from typing import Any, TypeAlias
 from uuid import UUID
 
-from models_library.users import UserID
 from pydantic import BaseModel
 
+from ..products import ProductName
 from ..progress_bar import ProgressReport
+from ..users import UserID
 
 AsyncJobId: TypeAlias = UUID
+AsyncJobName: TypeAlias = str
 
 
 class AsyncJobStatus(BaseModel):
@@ -21,6 +23,7 @@ class AsyncJobResult(BaseModel):
 
 class AsyncJobGet(BaseModel):
     job_id: AsyncJobId
+    job_name: AsyncJobName
 
 
 class AsyncJobAbort(BaseModel):
@@ -31,5 +34,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,8 +25,10 @@ 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_uuid = await get_celery_client(app).submit_task(
+        task_metadata=TaskMetadata(
+            name=remote_compute_path_size.__name__,
+        ),
         task_context=job_id_data.model_dump(),
         user_id=job_id_data.user_id,
         location_id=location_id,
@@ -42,8 +45,10 @@ 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_uuid = await get_celery_client(app).submit_task(
+        task_metadata=TaskMetadata(
+            name=remote_delete_paths.__name__,
+        ),
         task_context=job_id_data.model_dump(),
         user_id=job_id_data.user_id,
         location_id=location_id,

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

@@ -20,8 +20,10 @@ 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_uuid = await get_celery_client(app).submit_task(
+        task_metadata=TaskMetadata(
+            name=deep_copy_files_from_project.__name__,
+        ),
         task_context=job_id_data.model_dump(),
         user_id=job_id_data.user_id,
         body=body,
@@ -34,13 +36,13 @@ async def copy_folders_from_project(
 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_uuid = await get_celery_client(app).submit_task(
         task_metadata=TaskMetadata(
+            name=export_data.__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,
     )

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

@@ -7,10 +7,11 @@
 from servicelib.redis._client import RedisClientSDK
 
 from ..models import (
+    Task,
     TaskContext,
+    TaskID,
     TaskMetadata,
     TaskUUID,
-    build_task_id,
     build_task_id_prefix,
 )
 
@@ -24,24 +25,21 @@
 _logger = logging.getLogger(__name__)
 
 
-def _build_key(task_context: TaskContext, task_uuid: TaskUUID | None = None) -> str:
-    if task_uuid is None:
-        return _CELERY_TASK_INFO_PREFIX + build_task_id_prefix(task_context)
-    return _CELERY_TASK_INFO_PREFIX + build_task_id(task_context, task_uuid)
+def _build_key(task_id: TaskID) -> str:
+    return _CELERY_TASK_INFO_PREFIX + task_id
 
 
 class RedisTaskInfoStore:
     def __init__(self, redis_client_sdk: RedisClientSDK) -> None:
         self._redis_client_sdk = redis_client_sdk
 
-    async def create(
+    async def create_task(
         self,
-        task_context: TaskContext,
-        task_uuid: TaskUUID,
+        task_id: TaskID,
         task_metadata: TaskMetadata,
         expiry: timedelta,
     ) -> None:
-        task_key = _build_key(task_context, task_uuid)
+        task_key = _build_key(task_id)
         await self._redis_client_sdk.redis.hset(
             name=task_key,
             key=_CELERY_TASK_METADATA_KEY,
@@ -52,26 +50,28 @@ async def create(
             expiry,
         )
 
-    async def exists(self, task_context: TaskContext, task_uuid: TaskUUID) -> bool:
-        n = await self._redis_client_sdk.redis.exists(_build_key(task_context, task_uuid))  # type: ignore
+    async def exists_task(self, task_id: TaskID) -> bool:
+        n = await self._redis_client_sdk.redis.exists(_build_key(task_id))  # type: ignore
         assert isinstance(n, int)  # nosec
         return n > 0
 
-    async def get_metadata(
-        self, task_context: TaskContext, task_uuid: TaskUUID
-    ) -> TaskMetadata | None:
-        result = await self._redis_client_sdk.redis.hget(_build_key(task_context, task_uuid), _CELERY_TASK_METADATA_KEY)  # type: ignore
+    async def get_task_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_progress(
-        self, task_context: TaskContext, task_uuid: TaskUUID
-    ) -> ProgressReport | None:
-        result = await self._redis_client_sdk.redis.hget(_build_key(task_context, task_uuid), _CELERY_TASK_PROGRESS_KEY)  # type: ignore
+    async def get_task_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_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
-        search_key = _build_key(task_context) + _CELERY_TASK_ID_KEY_SEPARATOR
-        keys = set()
+    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: list[str] = []
+        tasks = []
+        pipe = 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
         ):
@@ -81,18 +81,28 @@ 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)
 
-    async def remove(self, task_context: TaskContext, task_uuid: TaskUUID) -> None:
-        await self._redis_client_sdk.redis.delete(_build_key(task_context, task_uuid))  # type: ignore
-        AsyncResult(build_task_id(task_context, task_uuid)).forget()
+        for key in keys:
+            pipe.hget(key, _CELERY_TASK_METADATA_KEY)
 
-    async def set_progress(
-        self, task_context: TaskContext, task_uuid: TaskUUID, report: ProgressReport
-    ) -> None:
+        results = await pipe.execute()
+        for key, task_metadata in zip(keys, results, strict=False):
+            tasks.append(
+                Task(
+                    uuid=TaskUUID(key.removeprefix(search_key)),
+                    metadata=TaskMetadata.model_validate_json(task_metadata),
+                )
+            )
+        return tasks
+
+    async def remove_task(self, task_id: TaskID) -> None:
+        await self._redis_client_sdk.redis.delete(_build_key(task_id))  # type: ignore
+        AsyncResult(task_id).forget()
+
+    async def set_task_progress(self, task_id: TaskID, report: ProgressReport) -> None:
         await self._redis_client_sdk.redis.hset(
-            name=_build_key(task_context, task_uuid),
+            name=_build_key(task_id),
             key=_CELERY_TASK_PROGRESS_KEY,
             value=report.model_dump_json(),
         )  # type: ignore

services/storage/src/simcore_service_storage/modules/celery/client.py

@@ -13,6 +13,7 @@
 from settings_library.celery import CelerySettings
 
 from .models import (
+    Task,
     TaskContext,
     TaskInfoStore,
     TaskMetadata,
@@ -35,7 +36,7 @@ class CeleryTaskClient:
     _celery_settings: CelerySettings
     _task_store: TaskInfoStore
 
-    async def send_task(
+    async def submit_task(
         self,
         task_metadata: TaskMetadata,
         *,
@@ -48,9 +49,10 @@ async def send_task(
             msg=f"Submit {task_metadata.name=}: {task_context=} {task_params=}",
         ):
             task_uuid = uuid4()
+            task_id = build_task_id(task_context, task_uuid)
             self._celery_app.send_task(
                 task_metadata.name,
-                task_id=build_task_id(task_context, task_uuid),
+                task_id=task_id,
                 kwargs=task_params,
                 queue=task_metadata.queue.value,
             )
@@ -60,9 +62,7 @@ async def send_task(
                 if task_metadata.ephemeral
                 else self._celery_settings.CELERY_RESULT_EXPIRES
             )
-            await self._task_store.create(
-                task_context, task_uuid, task_metadata, expiry=expiry
-            )
+            await self._task_store.create_task(task_id, task_metadata, expiry=expiry)
             return task_uuid
 
     @make_async()
@@ -91,18 +91,17 @@ async def get_task_result(
             async_result = self._celery_app.AsyncResult(task_id)
             result = async_result.result
             if async_result.ready():
-                task_metadata = await self._task_store.get_metadata(
-                    task_context, task_uuid
-                )
+                task_metadata = await self._task_store.get_task_metadata(task_id)
                 if task_metadata is not None and task_metadata.ephemeral:
-                    await self._task_store.remove(task_context, task_uuid)
+                    await self._task_store.remove_task(task_id)
             return result
 
-    async def _get_progress_report(
+    async def _get_task_progress_report(
         self, task_context: TaskContext, task_uuid: TaskUUID, state: TaskState
     ) -> ProgressReport:
         if state in (TaskState.STARTED, TaskState.RETRY, TaskState.ABORTED):
-            progress = await self._task_store.get_progress(task_context, task_uuid)
+            task_id = build_task_id(task_context, task_uuid)
+            progress = await self._task_store.get_task_progress(task_id)
             if progress is not None:
                 return progress
         if state in (
@@ -119,7 +118,9 @@ async def _get_progress_report(
         )
 
     @make_async()
-    def _get_state(self, task_context: TaskContext, task_uuid: TaskUUID) -> TaskState:
+    def _get_task_celery_state(
+        self, task_context: TaskContext, task_uuid: TaskUUID
+    ) -> TaskState:
         task_id = build_task_id(task_context, task_uuid)
         return TaskState(self._celery_app.AsyncResult(task_id).state)
 
@@ -131,18 +132,19 @@ async def get_task_status(
             logging.DEBUG,
             msg=f"Getting task status: {task_context=} {task_uuid=}",
         ):
-            task_state = await self._get_state(task_context, task_uuid)
-            task_id = build_task_id(task_context, task_uuid)
+            task_state = await self._get_task_celery_state(task_context, task_uuid)
             return TaskStatus(
                 task_uuid=task_uuid,
                 task_state=task_state,
-                progress_report=await self._get_progress_report(task_id, task_state),
+                progress_report=await self._get_task_progress_report(
+                    task_context, task_uuid, task_state
+                ),
             )
 
-    async def get_task_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
+    async def list_tasks(self, task_context: TaskContext) -> list[Task]:
         with log_context(
             _logger,
             logging.DEBUG,
-            msg=f"Getting task uuids: {task_context=}",
+            msg=f"Listing tasks: {task_context=}",
         ):
-            return await self._task_store.get_uuids(task_context)
+            return await self._task_store.list_tasks(task_context)