Merge branch 'master' into pr/giancarloromeo/8141

Author: pcrespov

packages/aws-library/src/aws_library/s3/__init__.py

@@ -22,10 +22,10 @@
 )
 
 __all__: tuple[str, ...] = (
-    "CopiedBytesTransferredCallback",
-    "MultiPartUploadLinks",
     "PRESIGNED_LINK_MAX_SIZE",
     "S3_MAX_FILE_SIZE",
+    "CopiedBytesTransferredCallback",
+    "MultiPartUploadLinks",
     "S3AccessError",
     "S3BucketInvalidError",
     "S3DestinationNotEmptyError",
@@ -37,8 +37,8 @@
     "S3RuntimeError",
     "S3UploadNotFoundError",
     "SimcoreS3API",
-    "UploadedBytesTransferredCallback",
     "UploadID",
+    "UploadedBytesTransferredCallback",
 )
 
 # nopycln: file

packages/celery-library/src/celery_library/backends/_redis.py renamed to packages/celery-library/src/celery_library/backends/redis.py

@@ -1,25 +1,23 @@
 import contextlib
 import logging
 from datetime import timedelta
-from typing import Final
+from typing import TYPE_CHECKING, Final
 
 from models_library.progress_bar import ProgressReport
 from pydantic import ValidationError
 from servicelib.celery.models import (
     Task,
     TaskFilter,
     TaskID,
+    TaskInfoStore,
     TaskMetadata,
-    TaskUUID,
+    Wildcard,
 )
-from servicelib.redis import RedisClientSDK
-
-from ..utils import build_task_id_prefix
+from servicelib.redis import RedisClientSDK, handle_redis_returns_union_types
 
 _CELERY_TASK_INFO_PREFIX: Final[str] = "celery-task-info-"
 _CELERY_TASK_ID_KEY_ENCODING = "utf-8"
-_CELERY_TASK_ID_KEY_SEPARATOR: Final[str] = ":"
-_CELERY_TASK_SCAN_COUNT_PER_BATCH: Final[int] = 10000
+_CELERY_TASK_SCAN_COUNT_PER_BATCH: Final[int] = 1000
 _CELERY_TASK_METADATA_KEY: Final[str] = "metadata"
 _CELERY_TASK_PROGRESS_KEY: Final[str] = "progress"
 
@@ -41,23 +39,24 @@ async def create_task(
         expiry: timedelta,
     ) -> None:
         task_key = _build_key(task_id)
-        await self._redis_client_sdk.redis.hset(
-            name=task_key,
-            key=_CELERY_TASK_METADATA_KEY,
-            value=task_metadata.model_dump_json(),
-        )  # type: ignore
+        await handle_redis_returns_union_types(
+            self._redis_client_sdk.redis.hset(
+                name=task_key,
+                key=_CELERY_TASK_METADATA_KEY,
+                value=task_metadata.model_dump_json(),
+            )
+        )
         await self._redis_client_sdk.redis.expire(
             task_key,
             expiry,
         )
 
-    async def exists_task(self, task_id: TaskID) -> bool:
-        n = await self._redis_client_sdk.redis.exists(_build_key(task_id))
-        assert isinstance(n, int)  # nosec
-        return n > 0
-
     async def get_task_metadata(self, task_id: TaskID) -> TaskMetadata | None:
-        raw_result = await self._redis_client_sdk.redis.hget(_build_key(task_id), _CELERY_TASK_METADATA_KEY)  # type: ignore
+        raw_result = await handle_redis_returns_union_types(
+            self._redis_client_sdk.redis.hget(
+                _build_key(task_id), _CELERY_TASK_METADATA_KEY
+            )
+        )
         if not raw_result:
             return None
 
@@ -70,7 +69,11 @@ async def get_task_metadata(self, task_id: TaskID) -> TaskMetadata | None:
             return None
 
     async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None:
-        raw_result = await self._redis_client_sdk.redis.hget(_build_key(task_id), _CELERY_TASK_PROGRESS_KEY)  # type: ignore
+        raw_result = await handle_redis_returns_union_types(
+            self._redis_client_sdk.redis.hget(
+                _build_key(task_id), _CELERY_TASK_PROGRESS_KEY
+            )
+        )
         if not raw_result:
             return None
 
@@ -83,17 +86,14 @@ async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None:
             return None
 
     async def list_tasks(self, task_filter: TaskFilter) -> list[Task]:
-        search_key = (
-            _CELERY_TASK_INFO_PREFIX
-            + build_task_id_prefix(task_filter)
-            + _CELERY_TASK_ID_KEY_SEPARATOR
+        search_key = _CELERY_TASK_INFO_PREFIX + task_filter.create_task_id(
+            task_uuid=Wildcard()
         )
-        search_key_len = len(search_key)
 
         keys: list[str] = []
         pipeline = 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
+            match=search_key, count=_CELERY_TASK_SCAN_COUNT_PER_BATCH
         ):
             # fake redis (tests) returns bytes, real redis returns str
             _key = (
@@ -115,7 +115,7 @@ async def list_tasks(self, task_filter: TaskFilter) -> list[Task]:
                 task_metadata = TaskMetadata.model_validate_json(raw_metadata)
                 tasks.append(
                     Task(
-                        uuid=TaskUUID(key[search_key_len:]),
+                        uuid=TaskFilter.get_task_uuid(key),
                         metadata=task_metadata,
                     )
                 )
@@ -126,8 +126,19 @@ async def remove_task(self, task_id: TaskID) -> None:
         await self._redis_client_sdk.redis.delete(_build_key(task_id))
 
     async def set_task_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
+        await handle_redis_returns_union_types(
+            self._redis_client_sdk.redis.hset(
+                name=_build_key(task_id),
+                key=_CELERY_TASK_PROGRESS_KEY,
+                value=report.model_dump_json(),
+            )
+        )
+
+    async def task_exists(self, task_id: TaskID) -> bool:
+        n = await self._redis_client_sdk.redis.exists(_build_key(task_id))
+        assert isinstance(n, int)  # nosec
+        return n > 0
+
+
+if TYPE_CHECKING:
+    _: type[TaskInfoStore] = RedisTaskInfoStore

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

@@ -2,13 +2,9 @@
 from typing import Any
 
 from celery import Celery  # type: ignore[import-untyped]
-from servicelib.redis import RedisClientSDK
 from settings_library.celery import CelerySettings
 from settings_library.redis import RedisDatabase
 
-from .backends._redis import RedisTaskInfoStore
-from .task_manager import CeleryTaskManager
-
 
 def _celery_configure(celery_settings: CelerySettings) -> dict[str, Any]:
     base_config = {
@@ -36,22 +32,3 @@ def create_app(settings: CelerySettings) -> Celery:
         ),
         **_celery_configure(settings),
     )
-
-
-async def create_task_manager(
-    app: Celery, settings: CelerySettings
-) -> CeleryTaskManager:
-    redis_client_sdk = RedisClientSDK(
-        settings.CELERY_REDIS_RESULT_BACKEND.build_redis_dsn(
-            RedisDatabase.CELERY_TASKS
-        ),
-        client_name="celery_tasks",
-    )
-    await redis_client_sdk.setup()
-    # GCR please address https://github.com/ITISFoundation/osparc-simcore/issues/8159
-
-    return CeleryTaskManager(
-        app,
-        settings,
-        RedisTaskInfoStore(redis_client_sdk),
-    )

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

@@ -1,6 +1,8 @@
 import base64
 import pickle
 
+from common_library.errors_classes import OsparcErrorMixin
+
 
 class TransferrableCeleryError(Exception):
     def __repr__(self) -> str:
@@ -22,3 +24,7 @@ def decode_celery_transferrable_error(error: TransferrableCeleryError) -> Except
     assert isinstance(error, TransferrableCeleryError)  # nosec
     result: Exception = pickle.loads(base64.b64decode(error.args[0]))  # noqa: S301
     return result
+
+
+class TaskNotFoundError(OsparcErrorMixin, Exception):
+    msg_template = "Task with id '{task_id}' was not found"

packages/celery-library/src/celery_library/rpc/_async_jobs.py

@@ -13,6 +13,7 @@
 from models_library.api_schemas_rpc_async_jobs.exceptions import (
     JobAbortedError,
     JobError,
+    JobMissingError,
     JobNotDoneError,
     JobSchedulerError,
 )
@@ -22,6 +23,7 @@
 from servicelib.rabbitmq import RPCRouter
 
 from ..errors import (
+    TaskNotFoundError,
     TransferrableCeleryError,
     decode_celery_transferrable_error,
 )
@@ -30,7 +32,7 @@
 router = RPCRouter()
 
 
-@router.expose(reraise_if_error_type=(JobSchedulerError,))
+@router.expose(reraise_if_error_type=(JobSchedulerError, JobMissingError))
 async def cancel(
     task_manager: TaskManager, job_id: AsyncJobId, job_filter: AsyncJobFilter
 ):
@@ -42,11 +44,13 @@ async def cancel(
             task_filter=task_filter,
             task_uuid=job_id,
         )
+    except TaskNotFoundError as exc:
+        raise JobMissingError(job_id=job_id) from exc
     except CeleryError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc
 
 
-@router.expose(reraise_if_error_type=(JobSchedulerError,))
+@router.expose(reraise_if_error_type=(JobSchedulerError, JobMissingError))
 async def status(
     task_manager: TaskManager, job_id: AsyncJobId, job_filter: AsyncJobFilter
 ) -> AsyncJobStatus:
@@ -59,6 +63,8 @@ async def status(
             task_filter=task_filter,
             task_uuid=job_id,
         )
+    except TaskNotFoundError as exc:
+        raise JobMissingError(job_id=job_id) from exc
     except CeleryError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc
 
@@ -71,9 +77,10 @@ async def status(
 
 @router.expose(
     reraise_if_error_type=(
+        JobAbortedError,
         JobError,
+        JobMissingError,
         JobNotDoneError,
-        JobAbortedError,
         JobSchedulerError,
     )
 )
@@ -97,11 +104,11 @@ async def result(
             task_filter=task_filter,
             task_uuid=job_id,
         )
+    except TaskNotFoundError as exc:
+        raise JobMissingError(job_id=job_id) from exc
     except CeleryError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc
 
-    if _status.task_state == TaskState.ABORTED:
-        raise JobAbortedError(job_id=job_id)
     if _status.task_state == TaskState.FAILURE:
         # fallback exception to report
         exc_type = type(_result).__name__
@@ -127,9 +134,8 @@ async def result(
 
 @router.expose(reraise_if_error_type=(JobSchedulerError,))
 async def list_jobs(
-    task_manager: TaskManager, filter_: str, job_filter: AsyncJobFilter
+    task_manager: TaskManager, job_filter: AsyncJobFilter
 ) -> list[AsyncJobGet]:
-    _ = filter_
     assert task_manager  # nosec
     task_filter = TaskFilter.model_validate(job_filter.model_dump())
     try:

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

@@ -6,18 +6,15 @@
 from celery.worker.worker import WorkController  # type: ignore[import-untyped]
 from servicelib.celery.app_server import BaseAppServer
 from servicelib.logging_utils import log_context
-from settings_library.celery import CelerySettings
 
-from .common import create_task_manager
 from .utils import get_app_server, set_app_server
 
 _logger = logging.getLogger(__name__)
 
 
 def on_worker_init(
-    app_server: BaseAppServer,
-    celery_settings: CelerySettings,
     sender: WorkController,
+    app_server: BaseAppServer,
     **_kwargs,
 ) -> None:
     startup_complete_event = threading.Event()
@@ -26,21 +23,14 @@ def _init(startup_complete_event: threading.Event) -> None:
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
 
-        async def _setup_task_manager():
-            assert sender.app  # nosec
-            assert isinstance(sender.app, Celery)  # nosec
-
-            app_server.task_manager = await create_task_manager(
-                sender.app,
-                celery_settings,
-            )
+        assert sender.app  # nosec
+        assert isinstance(sender.app, Celery)  # nosec
 
-            set_app_server(sender.app, app_server)
+        set_app_server(sender.app, app_server)
 
         app_server.event_loop = loop
 
-        loop.run_until_complete(_setup_task_manager())
-        loop.run_until_complete(app_server.lifespan(startup_complete_event))
+        loop.run_until_complete(app_server.run_until_shutdown(startup_complete_event))
 
     thread = threading.Thread(
         group=None,