move set_progress

Author: giancarloromeo

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

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

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

@@ -3,6 +3,7 @@
 from typing import Final
 
 from celery.result import AsyncResult  # type: ignore[import-untyped]
+from models_library.progress_bar import ProgressReport
 from servicelib.redis._client import RedisClientSDK
 
 from ..models import TaskContext, TaskID, TaskMetadata, TaskUUID, build_task_id_prefix
@@ -12,6 +13,7 @@
 _CELERY_TASK_ID_KEY_SEPARATOR: Final[str] = ":"
 _CELERY_TASK_SCAN_COUNT_PER_BATCH: Final[int] = 10000
 _CELERY_TASK_METADATA_KEY: Final[str] = "metadata"
+_CELERY_TASK_PROGRESS_KEY: Final[str] = "progress"
 
 _logger = logging.getLogger(__name__)
 
@@ -33,6 +35,10 @@ 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_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 = (
             _CELERY_TASK_INFO_PREFIX
@@ -68,3 +74,10 @@ async def set_metadata(
             _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),
+            key=_CELERY_TASK_PROGRESS_KEY,
+            value=report.model_dump_json(),
+        )  # type: ignore

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

@@ -49,6 +49,8 @@ class TaskMetadata(BaseModel):
 class TaskInfoStore(Protocol):
     async def exists(self, task_id: TaskID) -> bool: ...
 
+    async def get_progress(self, task_id: TaskID) -> ProgressReport | None: ...
+
     async def get_metadata(self, task_id: TaskID) -> TaskMetadata | None: ...
 
     async def get_uuids(self, task_context: TaskContext) -> set[TaskUUID]: ...
@@ -59,6 +61,8 @@ async def set_metadata(
         self, task_id: TaskID, task_data: TaskMetadata, expiry: timedelta
     ) -> None: ...
 
+    async def set_progress(self, task_id: TaskID, report: ProgressReport) -> None: ...
+
 
 class TaskStatus(BaseModel):
     task_uuid: TaskUUID

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

@@ -17,7 +17,7 @@
     set_celery_worker,
     set_fastapi_app,
 )
-from ...modules.celery.worker import CeleryTaskQueueWorker
+from ...modules.celery.worker import CeleryTaskWorker
 
 _logger = logging.getLogger(__name__)
 
@@ -54,7 +54,7 @@ async def lifespan(
         set_event_loop(fastapi_app, loop)
 
         set_fastapi_app(sender.app, fastapi_app)
-        set_celery_worker(sender.app, CeleryTaskQueueWorker(sender.app))
+        set_celery_worker(sender.app, CeleryTaskWorker(sender.app))
         loop.run_until_complete(lifespan(startup_complete_event, shutdown_event))
 
     thread = threading.Thread(

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

@@ -1,19 +1,19 @@
 from celery import Celery  # type: ignore[import-untyped]
 from fastapi import FastAPI
 
-from .worker import CeleryTaskQueueWorker
+from .worker import CeleryTaskWorker
 
 _WORKER_KEY = "celery_worker"
 _FASTAPI_APP_KEY = "fastapi_app"
 
 
-def set_celery_worker(celery_app: Celery, worker: CeleryTaskQueueWorker) -> None:
+def set_celery_worker(celery_app: Celery, worker: CeleryTaskWorker) -> None:
     celery_app.conf[_WORKER_KEY] = worker
 
 
-def get_celery_worker(celery_app: Celery) -> CeleryTaskQueueWorker:
+def get_celery_worker(celery_app: Celery) -> CeleryTaskWorker:
     worker = celery_app.conf[_WORKER_KEY]
-    assert isinstance(worker, CeleryTaskQueueWorker)
+    assert isinstance(worker, CeleryTaskWorker)
     return worker
 
 

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

@@ -1,28 +1,19 @@
 import logging
+from dataclasses import dataclass
 
-from celery import Celery  # type: ignore[import-untyped]
 from models_library.progress_bar import ProgressReport
-from servicelib.logging_utils import log_context
 
-from ..celery.models import TaskID
+from ..celery.models import TaskID, TaskInfoStore
 
 _logger = logging.getLogger(__name__)
 
 
-class CeleryTaskQueueWorker:
-    def __init__(self, celery_app: Celery) -> None:
-        self.celery_app = celery_app
+@dataclass
+class CeleryTaskWorker:
+    _task_info_store: TaskInfoStore
 
-    def set_task_progress(
-        self, task_name: str, task_id: TaskID, report: ProgressReport
-    ) -> None:
-        with log_context(
-            _logger,
-            logging.DEBUG,
-            msg=f"Setting progress for {task_name}: {report.model_dump_json()}",
-        ):
-            self.celery_app.tasks[task_name].update_state(
-                task_id=task_id,
-                state="RUNNING",
-                meta=report.model_dump(mode="json"),
-            )
+    async def set_progress(self, task_id: TaskID, report: ProgressReport) -> None:
+        await self._task_info_store.set_progress(
+            task_id=task_id,
+            report=report,
+        )

services/storage/tests/conftest.py

@@ -76,7 +76,7 @@
     on_worker_shutdown,
 )
 from simcore_service_storage.modules.celery.utils import get_celery_worker
-from simcore_service_storage.modules.celery.worker import CeleryTaskQueueWorker
+from simcore_service_storage.modules.celery.worker import CeleryTaskWorker
 from simcore_service_storage.modules.s3 import get_s3_client
 from simcore_service_storage.simcore_s3_dsm import SimcoreS3DataManager
 from sqlalchemy import literal_column
@@ -365,7 +365,7 @@ def upload_file(
     create_upload_file_link_v2: Callable[..., Awaitable[FileUploadSchema]],
     create_file_of_size: Callable[[ByteSize, str | None], Path],
     create_simcore_file_id: Callable[[ProjectID, NodeID, str], SimcoreS3FileID],
-    with_storage_celery_worker: CeleryTaskQueueWorker,
+    with_storage_celery_worker: CeleryTaskWorker,
 ) -> Callable[
     [ByteSize, str, SimcoreS3FileID | None], Awaitable[tuple[Path, SimcoreS3FileID]]
 ]:
@@ -480,7 +480,7 @@ async def create_empty_directory(
     create_simcore_file_id: Callable[[ProjectID, NodeID, str], SimcoreS3FileID],
     create_upload_file_link_v2: Callable[..., Awaitable[FileUploadSchema]],
     client: httpx.AsyncClient,
-    with_storage_celery_worker: CeleryTaskQueueWorker,
+    with_storage_celery_worker: CeleryTaskWorker,
 ) -> Callable[[str, ProjectID, NodeID], Awaitable[SimcoreS3FileID]]:
     async def _directory_creator(
         dir_name: str, project_id: ProjectID, node_id: NodeID
@@ -1024,7 +1024,7 @@ async def with_storage_celery_worker_controller(
 @pytest.fixture
 def with_storage_celery_worker(
     with_storage_celery_worker_controller: TestWorkController,
-) -> CeleryTaskQueueWorker:
+) -> CeleryTaskWorker:
     assert isinstance(with_storage_celery_worker_controller.app, Celery)
     return get_celery_worker(with_storage_celery_worker_controller.app)
 

services/storage/tests/unit/test_async_jobs.py

@@ -30,7 +30,7 @@
 from simcore_service_storage.modules.celery import get_celery_client
 from simcore_service_storage.modules.celery._task import register_task
 from simcore_service_storage.modules.celery.models import TaskID
-from simcore_service_storage.modules.celery.worker import CeleryTaskQueueWorker
+from simcore_service_storage.modules.celery.worker import CeleryTaskWorker
 from tenacity import (
     AsyncRetrying,
     retry_if_exception_type,
@@ -199,7 +199,7 @@ async def test_async_jobs_workflow(
     initialized_app: FastAPI,
     register_rpc_routes: None,
     storage_rabbitmq_rpc_client: RabbitMQRPCClient,
-    with_storage_celery_worker: CeleryTaskQueueWorker,
+    with_storage_celery_worker: CeleryTaskWorker,
     user_id: UserID,
     product_name: ProductName,
     exposed_rpc_start: str,
@@ -247,7 +247,7 @@ async def test_async_jobs_cancel(
     initialized_app: FastAPI,
     register_rpc_routes: None,
     storage_rabbitmq_rpc_client: RabbitMQRPCClient,
-    with_storage_celery_worker: CeleryTaskQueueWorker,
+    with_storage_celery_worker: CeleryTaskWorker,
     user_id: UserID,
     product_name: ProductName,
     exposed_rpc_start: str,
@@ -328,7 +328,7 @@ async def test_async_jobs_raises(
     initialized_app: FastAPI,
     register_rpc_routes: None,
     storage_rabbitmq_rpc_client: RabbitMQRPCClient,
-    with_storage_celery_worker: CeleryTaskQueueWorker,
+    with_storage_celery_worker: CeleryTaskWorker,
     user_id: UserID,
     product_name: ProductName,
     exposed_rpc_start: str,

services/storage/tests/unit/test_handlers_files.py

@@ -53,7 +53,7 @@
 from servicelib.aiohttp import status
 from simcore_service_storage.constants import S3_UNDEFINED_OR_EXTERNAL_MULTIPART_ID
 from simcore_service_storage.models import FileDownloadResponse, S3BucketName, UploadID
-from simcore_service_storage.modules.celery.worker import CeleryTaskQueueWorker
+from simcore_service_storage.modules.celery.worker import CeleryTaskWorker
 from simcore_service_storage.simcore_s3_dsm import SimcoreS3DataManager
 from sqlalchemy.ext.asyncio import AsyncEngine
 from tenacity.asyncio import AsyncRetrying
@@ -683,7 +683,7 @@ async def test_upload_real_file_with_s3_client(
     node_id: NodeID,
     faker: Faker,
     s3_client: S3Client,
-    with_storage_celery_worker: CeleryTaskQueueWorker,
+    with_storage_celery_worker: CeleryTaskWorker,
 ):
     file_size = TypeAdapter(ByteSize).validate_python("500Mib")
     file_name = faker.file_name()

services/storage/tests/unit/test_modules_celery.py

@@ -29,7 +29,7 @@
     get_celery_worker,
     get_fastapi_app,
 )
-from simcore_service_storage.modules.celery.worker import CeleryTaskQueueWorker
+from simcore_service_storage.modules.celery.worker import CeleryTaskWorker
 from tenacity import Retrying, retry_if_exception_type, stop_after_delay, wait_fixed
 
 _logger = logging.getLogger(__name__)
@@ -41,7 +41,7 @@
 @pytest.fixture
 def celery_client(
     initialized_app: FastAPI,
-    with_storage_celery_worker: CeleryTaskQueueWorker,
+    with_storage_celery_worker: CeleryTaskWorker,
 ) -> CeleryTaskQueueClient:
     return get_celery_client(initialized_app)
 
@@ -56,8 +56,7 @@ def sleep_for(seconds: float) -> None:
 
     for n, file in enumerate(files, start=1):
         with log_context(_logger, logging.INFO, msg=f"Processing file {file}"):
-            worker.set_task_progress(
-                task_name=task_name,
+            await worker.set_progress(
                 task_id=task_id,
                 report=ProgressReport(actual_value=n / len(files)),
             )