continue

Author: giancarloromeo

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

@@ -40,6 +40,7 @@ def _init(startup_complete_event: threading.Event) -> None:
         shutdown_event = asyncio.Event()
 
         fastapi_app = app_factory()
+        assert isinstance(fastapi_app, FastAPI)  # nosec
 
         async def setup_task_worker():
             redis_client_sdk = RedisClientSDK(

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

@@ -205,4 +205,5 @@ def register_task(  # type: ignore[misc]
         bind=True,
         base=AbortableTask,
         time_limit=None if timeout is None else timeout.total_seconds(),
+        pydantic=True,
     )(wrapped_fn)

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

@@ -1,9 +1,6 @@
-from functools import partial
 from pathlib import Path
-from typing import Any
 
 from kombu.utils.json import register_type  # type: ignore[import-untyped]
-from pydantic import BaseModel
 
 
 def _path_encoder(obj):
@@ -23,14 +20,6 @@ def _class_full_name(clz: type) -> str:
     return ".".join([clz.__module__, clz.__qualname__])
 
 
-def _pydantic_model_encoder(obj: BaseModel, *args, **kwargs) -> dict[str, Any]:
-    return obj.model_dump(*args, **kwargs, mode="json")
-
-
-def _pydantic_model_decoder(clz: type[BaseModel], data: dict[str, Any]) -> BaseModel:
-    return clz(**data)
-
-
 def register_celery_types() -> None:
     register_type(
         Path,
@@ -39,13 +28,3 @@ def register_celery_types() -> None:
         _path_decoder,
     )
     register_type(set, _class_full_name(set), encoder=list, decoder=set)
-
-
-def register_pydantic_types(*models: type[BaseModel]) -> None:
-    for model in models:
-        register_type(
-            model,
-            _class_full_name(model),
-            encoder=_pydantic_model_encoder,
-            decoder=partial(_pydantic_model_decoder, model),
-        )

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

@@ -2,15 +2,10 @@
 
 from celery import Celery  # type: ignore[import-untyped]
 from celery_library.task import register_task
-from celery_library.types import register_celery_types, register_pydantic_types
+from celery_library.types import register_celery_types
 from models_library.api_schemas_storage.export_data_async_jobs import AccessRightError
-from models_library.api_schemas_storage.storage_schemas import (
-    FileUploadCompletionBody,
-    FoldersBody,
-)
 from servicelib.logging_utils import log_context
 
-from ...models import FileMetaData
 from ._files import complete_upload_file
 from ._paths import compute_path_size, delete_paths
 from ._simcore_s3 import deep_copy_files_from_project, export_data
@@ -20,9 +15,6 @@
 
 def setup_worker_tasks(app: Celery) -> None:
     register_celery_types()
-    register_pydantic_types(FileMetaData)
-    register_pydantic_types(FileUploadCompletionBody)
-    register_pydantic_types(FoldersBody)
 
     with log_context(_logger, logging.INFO, msg="worker task registration"):
         register_task(app, export_data, dont_autoretry_for=(AccessRightError,))

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

@@ -26,10 +26,18 @@
 )
 
 
-assert _settings.STORAGE_CELERY
+assert _settings.STORAGE_CELERY  # nosec
 app = create_celery_app(_settings.STORAGE_CELERY)
-app_factory = partial(create_app(_settings))
-worker_init.connect(partial(on_worker_init, app_factory, _settings.STORAGE_CELERY))
+app_factory = partial(create_app, _settings)
+
+
+def worker_init_wrapper(sender, **_kwargs):
+    return partial(on_worker_init, app_factory, _settings.STORAGE_CELERY)(
+        sender, **_kwargs
+    )
+
+
+worker_init.connect(worker_init_wrapper)
 worker_shutdown.connect(on_worker_shutdown)
 
 

services/storage/tests/conftest.py

@@ -1001,9 +1001,13 @@ async def with_storage_celery_worker_controller(
     monkeypatch.setenv("STORAGE_WORKER_MODE", "true")
     app_settings = ApplicationSettings.create_from_envs()
     app_factory = partial(create_app, app_settings)
-    worker_init.connect(
-        partial(on_worker_init, app_factory, app_settings.STORAGE_CELERY)
-    )
+
+    def _on_worker_init_wrapper(sender: Celery, **_kwargs) -> None:
+        return partial(on_worker_init, app_factory, app_settings.STORAGE_CELERY)(
+            sender, **_kwargs
+        )
+
+    worker_init.connect(_on_worker_init_wrapper)
     worker_shutdown.connect(on_worker_shutdown)
 
     setup_worker_tasks(celery_app)