initial commit

Author: giancarloromeo

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

@@ -6,7 +6,14 @@
 from models_library.progress_bar import ProgressReport
 from servicelib.redis._client import RedisClientSDK
 
-from ..models import TaskContext, TaskID, TaskMetadata, TaskUUID, build_task_id_prefix
+from ..models import (
+    TaskContext,
+    TaskID,
+    TaskMetadata,
+    TaskUUID,
+    build_task_id,
+    build_task_id_prefix,
+)
 
 _CELERY_TASK_INFO_PREFIX: Final[str] = "celery-task-info-"
 _CELERY_TASK_ID_KEY_ENCODING = "utf-8"
@@ -26,17 +33,39 @@ 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:
-        n = await self._redis_client_sdk.redis.exists(_build_key(task_id))
+    async def create(
+        self,
+        task_context: TaskContext,
+        task_uuid: TaskUUID,
+        task_metadata: TaskMetadata,
+        expiry: timedelta,
+    ) -> None:
+        task_key = _build_key(build_task_id(task_context, task_uuid))
+        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(self, task_context: TaskContext, task_uuid: TaskUUID) -> bool:
+        n = await self._redis_client_sdk.redis.exists(_build_key(build_task_id(task_context, task_uuid)))  # type: ignore
         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
+    async def get_metadata(
+        self, task_context: TaskContext, task_uuid: TaskUUID
+    ) -> TaskMetadata | None:
+        result = await self._redis_client_sdk.redis.hget(_build_key(build_task_id(task_context, task_uuid)), _CELERY_TASK_METADATA_KEY)  # type: ignore
         return TaskMetadata.model_validate_json(result) if result else 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
+    async def get_progress(
+        self, task_context: TaskContext, task_uuid: TaskUUID
+    ) -> ProgressReport | None:
+        result = await self._redis_client_sdk.redis.hget(_build_key(build_task_id(task_context, task_uuid)), _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]:
@@ -62,19 +91,6 @@ async def remove(self, task_id: TaskID) -> None:
         await self._redis_client_sdk.redis.delete(_build_key(task_id))
         AsyncResult(task_id).forget()
 
-    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,
-        )
-
     async def set_progress(self, task_id: TaskID, report: ProgressReport) -> None:
         await self._redis_client_sdk.redis.hset(
             name=_build_key(task_id),

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

@@ -38,10 +38,9 @@ class CeleryTaskClient:
 
     async def send_task(
         self,
-        task_name: str,
+        task_metadata: TaskMetadata,
         *,
         task_context: TaskContext,
-        task_metadata: TaskMetadata | None = None,
         **task_params,
     ) -> TaskUUID:
         with log_context(
@@ -50,11 +49,10 @@ async def send_task(
             msg=f"Submit {task_name=}: {task_context=} {task_params=}",
         ):
             task_uuid = uuid4()
-            task_id = build_task_id(task_context, task_uuid)
             task_metadata = task_metadata or TaskMetadata()
             self._celery_app.send_task(
                 task_name,
-                task_id=task_id,
+                task_id=build_task_id(task_context, task_uuid),
                 kwargs=task_params,
                 queue=task_metadata.queue.value,
             )
@@ -64,7 +62,9 @@ async def send_task(
                 if task_metadata.ephemeral
                 else self._celery_settings.CELERY_RESULT_EXPIRES
             )
-            await self._task_store.set_metadata(task_id, task_metadata, expiry=expiry)
+            await self._task_store.create(
+                task_context, task_uuid, task_metadata, expiry=expiry
+            )
             return task_uuid
 
     @make_async()

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

@@ -8,6 +8,7 @@
 
 TaskContext: TypeAlias = dict[str, Any]
 TaskID: TypeAlias = str
+TaskName: TypeAlias = str
 TaskUUID: TypeAlias = UUID
 
 _CELERY_TASK_ID_KEY_SEPARATOR: Final[str] = ":"
@@ -40,6 +41,7 @@ class TasksQueue(StrEnum):
 
 
 class TaskMetadata(BaseModel):
+    name: TaskName
     ephemeral: bool = True
     queue: TasksQueue = TasksQueue.DEFAULT
 
@@ -50,6 +52,14 @@ class TaskMetadata(BaseModel):
 class TaskInfoStore(Protocol):
     async def exists(self, task_id: TaskID) -> bool: ...
 
+    async def create(
+        self,
+        task_context: TaskContext,
+        task_uuid: TaskUUID,
+        task_metadata: TaskMetadata,
+        expiry: timedelta,
+    ) -> None: ...
+
     async def get_progress(self, task_id: TaskID) -> ProgressReport | None: ...
 
     async def get_metadata(self, task_id: TaskID) -> TaskMetadata | None: ...
@@ -58,10 +68,6 @@ async def get_uuids(self, task_context: TaskContext) -> set[TaskUUID]: ...
 
     async def remove(self, task_id: TaskID) -> None: ...
 
-    async def set_metadata(
-        self, task_id: TaskID, task_metadata: TaskMetadata, expiry: timedelta
-    ) -> None: ...
-
     async def set_progress(self, task_id: TaskID, report: ProgressReport) -> None: ...