unify interfaces

Author: giancarloromeo

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

@@ -7,8 +7,8 @@
 from settings_library.redis import RedisDatabase
 
 from .backends._redis import RedisTaskInfoStore
-from .client import CeleryTaskClient
 from .common import create_app
+from .task_manager import CeleryTaskManager
 from .types import register_celery_types
 
 _logger = logging.getLogger(__name__)
@@ -24,7 +24,7 @@ async def on_startup() -> None:
             client_name=f"{app.title}.celery_tasks",
         )
 
-        app.state.celery_client = CeleryTaskClient(
+        app.state.celery_client = CeleryTaskManager(
             celery_app,
             celery_settings,
             RedisTaskInfoStore(redis_client_sdk),
@@ -35,10 +35,10 @@ async def on_startup() -> None:
     app.add_event_handler("startup", on_startup)
 
 
-def get_celery_client(app: FastAPI) -> CeleryTaskClient:
+def get_celery_client(app: FastAPI) -> CeleryTaskManager:
     assert hasattr(app.state, "celery_client")  # nosec
     celery_client = app.state.celery_client
-    assert isinstance(celery_client, CeleryTaskClient)
+    assert isinstance(celery_client, CeleryTaskManager)
     return celery_client
 
 

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

@@ -16,12 +16,12 @@
 
 from . import set_event_loop
 from .backends._redis import RedisTaskInfoStore
+from .task_manager import CeleryTaskManager
 from .utils import (
     get_fastapi_app,
     set_celery_worker,
     set_fastapi_app,
 )
-from .worker import CeleryTaskWorker
 
 _logger = logging.getLogger(__name__)
 
@@ -57,7 +57,9 @@ async def setup_task_worker():
             assert isinstance(sender.app, Celery)  # nosec
             set_celery_worker(
                 sender.app,
-                CeleryTaskWorker(
+                CeleryTaskManager(
+                    sender.app,
+                    celery_settings,
                     RedisTaskInfoStore(redis_client_sdk),
                 ),
             )

packages/celery-library/src/celery_library/client.py renamed to packages/celery-library/src/celery_library/task_manager.py

@@ -31,8 +31,8 @@
 _MAX_PROGRESS_VALUE = 1.0
 
 
-@dataclass
-class CeleryTaskClient:
+@dataclass(frozen=True)
+class CeleryTaskManager:
     _celery_app: Celery
     _celery_settings: CelerySettings
     _task_info_store: TaskInfoStore
@@ -155,3 +155,9 @@ async def list_tasks(self, task_context: TaskContext) -> list[Task]:
             msg=f"Listing tasks: {task_context=}",
         ):
             return await self._task_info_store.list_tasks(task_context)
+
+    async def set_task_progress(self, task_id: TaskID, report: ProgressReport) -> None:
+        await self._task_info_store.set_task_progress(
+            task_id=task_id,
+            report=report,
+        )

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

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

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

@@ -2,8 +2,8 @@
 from typing import Annotated, Final, cast
 from urllib.parse import quote
 
-from celery_library.client import CeleryTaskClient
 from celery_library.models import TaskMetadata, TaskUUID
+from celery_library.task_manager import CeleryTaskManager
 from fastapi import APIRouter, Depends, Header, Request
 from models_library.api_schemas_rpc_async_jobs.async_jobs import AsyncJobNameData
 from models_library.api_schemas_storage.storage_schemas import (
@@ -270,7 +270,7 @@ async def abort_upload_file(
     status_code=status.HTTP_202_ACCEPTED,
 )
 async def complete_upload_file(
-    celery_client: Annotated[CeleryTaskClient, Depends(get_celery_client)],
+    celery_client: Annotated[CeleryTaskManager, Depends(get_celery_client)],
     query_params: Annotated[StorageQueryParamsBase, Depends()],
     location_id: LocationID,
     file_id: StorageFileID,
@@ -326,7 +326,7 @@ async def complete_upload_file(
     response_model=Envelope[FileUploadCompleteFutureResponse],
 )
 async def is_completed_upload_file(
-    celery_client: Annotated[CeleryTaskClient, Depends(get_celery_client)],
+    celery_client: Annotated[CeleryTaskManager, Depends(get_celery_client)],
     query_params: Annotated[StorageQueryParamsBase, Depends()],
     location_id: LocationID,
     file_id: StorageFileID,

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

@@ -1,12 +1,12 @@
 from typing import Annotated
 
 from celery_library import get_celery_client as _get_celery_client_from_app
-from celery_library.client import CeleryTaskClient
+from celery_library.task_manager import CeleryTaskManager
 from fastapi import Depends, FastAPI
 from servicelib.fastapi.dependencies import get_app
 
 
 def get_celery_client(
     app: Annotated[FastAPI, Depends(get_app)],
-) -> CeleryTaskClient:
+) -> CeleryTaskManager:
     return _get_celery_client_from_app(app)

services/storage/src/simcore_service_storage/api/rest/dependencies/celery.py

@@ -27,8 +27,8 @@
 from celery.signals import worker_init, worker_shutdown
 from celery.worker.worker import WorkController
 from celery_library.signals import on_worker_init, on_worker_shutdown
+from celery_library.task_manager import CeleryTaskManager
 from celery_library.utils import get_celery_worker
-from celery_library.worker import CeleryTaskWorker
 from faker import Faker
 from fakeredis.aioredis import FakeRedis
 from fastapi import FastAPI
@@ -364,7 +364,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: CeleryTaskWorker,
+    with_storage_celery_worker: CeleryTaskManager,
 ) -> Callable[
     [ByteSize, str, SimcoreS3FileID | None], Awaitable[tuple[Path, SimcoreS3FileID]]
 ]:
@@ -479,7 +479,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: CeleryTaskWorker,
+    with_storage_celery_worker: CeleryTaskManager,
 ) -> Callable[[str, ProjectID, NodeID], Awaitable[SimcoreS3FileID]]:
     async def _directory_creator(
         dir_name: str, project_id: ProjectID, node_id: NodeID
@@ -1029,7 +1029,7 @@ def _on_worker_init_wrapper(sender: WorkController, **_kwargs) -> None:
 @pytest.fixture
 def with_storage_celery_worker(
     with_storage_celery_worker_controller: TestWorkController,
-) -> CeleryTaskWorker:
+) -> CeleryTaskManager:
     assert isinstance(with_storage_celery_worker_controller.app, Celery)
     return get_celery_worker(with_storage_celery_worker_controller.app)
 

services/storage/tests/conftest.py

@@ -11,10 +11,9 @@
 import pytest
 from celery import Celery, Task
 from celery_library import get_celery_client
-from celery_library.client import TaskMetadata
-from celery_library.models import TaskID
+from celery_library.models import TaskID, TaskMetadata
 from celery_library.task import register_task
-from celery_library.worker import CeleryTaskWorker
+from celery_library.task_manager import CeleryTaskManager
 from fastapi import FastAPI
 from models_library.api_schemas_rpc_async_jobs.async_jobs import (
     AsyncJobGet,
@@ -203,7 +202,7 @@ async def test_async_jobs_workflow(
     initialized_app: FastAPI,
     register_rpc_routes: None,
     storage_rabbitmq_rpc_client: RabbitMQRPCClient,
-    with_storage_celery_worker: CeleryTaskWorker,
+    with_storage_celery_worker: CeleryTaskManager,
     user_id: UserID,
     product_name: ProductName,
     exposed_rpc_start: str,
@@ -251,7 +250,7 @@ async def test_async_jobs_cancel(
     initialized_app: FastAPI,
     register_rpc_routes: None,
     storage_rabbitmq_rpc_client: RabbitMQRPCClient,
-    with_storage_celery_worker: CeleryTaskWorker,
+    with_storage_celery_worker: CeleryTaskManager,
     user_id: UserID,
     product_name: ProductName,
     exposed_rpc_start: str,
@@ -316,7 +315,7 @@ async def test_async_jobs_raises(
     initialized_app: FastAPI,
     register_rpc_routes: None,
     storage_rabbitmq_rpc_client: RabbitMQRPCClient,
-    with_storage_celery_worker: CeleryTaskWorker,
+    with_storage_celery_worker: CeleryTaskManager,
     user_id: UserID,
     product_name: ProductName,
     exposed_rpc_start: str,

services/storage/tests/unit/test_async_jobs.py

@@ -23,7 +23,7 @@
 from aiohttp import ClientSession
 from aws_library.s3 import S3KeyNotFoundError, S3ObjectKey, SimcoreS3API
 from aws_library.s3._constants import MULTIPART_UPLOADS_MIN_TOTAL_SIZE
-from celery_library.worker import CeleryTaskWorker
+from celery_library.task_manager import CeleryTaskManager
 from faker import Faker
 from fastapi import FastAPI
 from models_library.api_schemas_storage.storage_schemas import (
@@ -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: CeleryTaskWorker,
+    with_storage_celery_worker: CeleryTaskManager,
 ):
     file_size = TypeAdapter(ByteSize).validate_python("500Mib")
     file_name = faker.file_name()