✨ Add a Distributed Task Queue (using Celery) (🏗️ ⚠️DEVOPS)

.env-devel

@@ -49,6 +49,8 @@ CATALOG_SERVICES_DEFAULT_RESOURCES='{"CPU": {"limit": 0.1, "reservation": 0.1},
 CATALOG_SERVICES_DEFAULT_SPECIFICATIONS='{}'
 CATALOG_TRACING=null
 
+CELERY_RESULT_EXPIRES=P7D
+
 CLUSTERS_KEEPER_COMPUTATIONAL_BACKEND_DEFAULT_CLUSTER_AUTH='{"type":"tls","tls_ca_file":"/home/scu/.dask/dask-crt.pem","tls_client_cert":"/home/scu/.dask/dask-crt.pem","tls_client_key":"/home/scu/.dask/dask-key.pem"}'
 CLUSTERS_KEEPER_COMPUTATIONAL_BACKEND_DOCKER_IMAGE_TAG=master-github-latest
 CLUSTERS_KEEPER_DASK_NTHREADS=0

api/specs/web-server/_storage.py

@@ -210,15 +210,15 @@ async def export_data(data_export: DataExportPost, location_id: LocationID):
     response_model=Envelope[StorageAsyncJobStatus],
     name="get_async_job_status",
 )
-async def get_async_job_status(storage_async_job_get: StorageAsyncJobGet, job_id: UUID):
+async def get_async_job_status(job_id: UUID):
     """Get async job status"""
 
 
 @router.post(
     "/storage/async-jobs/{job_id}:abort",
     name="abort_async_job",
 )
-async def abort_async_job(storage_async_job_get: StorageAsyncJobGet, job_id: UUID):
+async def abort_async_job(job_id: UUID):
     """aborts execution of an async job"""
 
 
@@ -227,7 +227,7 @@ async def abort_async_job(storage_async_job_get: StorageAsyncJobGet, job_id: UUI
     response_model=Envelope[StorageAsyncJobResult],
     name="get_async_job_result",
 )
-async def get_async_job_result(storage_async_job_get: StorageAsyncJobGet, job_id: UUID):
+async def get_async_job_result(job_id: UUID):
     """Get the result of the async job"""
 
 

packages/common-library/src/common_library/async_tools.py

@@ -0,0 +1,24 @@
+import asyncio
+import functools
+from collections.abc import Awaitable, Callable
+from concurrent.futures import Executor
+from typing import ParamSpec, TypeVar
+
+R = TypeVar("R")
+P = ParamSpec("P")
+
+
+def make_async(
+    executor: Executor | None = None,
+) -> Callable[[Callable[P, R]], Callable[P, Awaitable[R]]]:
+    def decorator(func: Callable[P, R]) -> Callable[P, Awaitable[R]]:
+        @functools.wraps(func)
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            loop = asyncio.get_running_loop()
+            return await loop.run_in_executor(
+                executor, functools.partial(func, *args, **kwargs)
+            )
+
+        return wrapper
+
+    return decorator

packages/common-library/tests/test_async_tools.py

@@ -0,0 +1,45 @@
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+from common_library.async_tools import make_async
+
+
+@make_async()
+def sync_function(x: int, y: int) -> int:
+    return x + y
+
+
+@make_async()
+def sync_function_with_exception() -> None:
+    raise ValueError("This is an error!")
+
+
+@pytest.mark.asyncio
+async def test_make_async_returns_coroutine():
+    result = sync_function(2, 3)
+    assert asyncio.iscoroutine(result), "Function should return a coroutine"
+
+
+@pytest.mark.asyncio
+async def test_make_async_execution():
+    result = await sync_function(2, 3)
+    assert result == 5, "Function should return 5"
+
+
+@pytest.mark.asyncio
+async def test_make_async_exception():
+    with pytest.raises(ValueError, match="This is an error!"):
+        await sync_function_with_exception()
+
+
+@pytest.mark.asyncio
+async def test_make_async_with_executor():
+    executor = ThreadPoolExecutor()
+
+    @make_async(executor)
+    def heavy_computation(x: int) -> int:
+        return x * x
+
+    result = await heavy_computation(4)
+    assert result == 16, "Function should return 16"

packages/models-library/src/models_library/api_schemas_rpc_async_jobs/async_jobs.py

@@ -1,10 +1,8 @@
-from datetime import datetime
 from typing import Any, TypeAlias
 from uuid import UUID
 
 from models_library.users import UserID
-from pydantic import BaseModel, model_validator
-from typing_extensions import Self
+from pydantic import BaseModel
 
 from ..progress_bar import ProgressReport
 
@@ -15,18 +13,6 @@ class AsyncJobStatus(BaseModel):
     job_id: AsyncJobId
     progress: ProgressReport
     done: bool
-    started: datetime
-    stopped: datetime | None
-
-    @model_validator(mode="after")
-    def _check_consistency(self) -> Self:
-        is_done = self.done
-        is_stopped = self.stopped is not None
-
-        if is_done != is_stopped:
-            msg = f"Inconsistent data: {self.done=}, {self.stopped=}"
-            raise ValueError(msg)
-        return self
 
 
 class AsyncJobResult(BaseModel):

packages/models-library/src/models_library/api_schemas_webserver/storage.py

@@ -1,4 +1,3 @@
-from datetime import datetime
 from pathlib import Path
 from typing import Annotated, Any
 
@@ -62,8 +61,6 @@ class StorageAsyncJobStatus(OutputSchema):
     job_id: AsyncJobId
     progress: ProgressReport
     done: bool
-    started: datetime
-    stopped: datetime | None
 
     @classmethod
     def from_rpc_schema(
@@ -73,8 +70,6 @@ def from_rpc_schema(
             job_id=async_job_rpc_status.job_id,
             progress=async_job_rpc_status.progress,
             done=async_job_rpc_status.done,
-            started=async_job_rpc_status.started,
-            stopped=async_job_rpc_status.stopped,
         )
 
 

packages/settings-library/src/settings_library/celery.py

@@ -0,0 +1,52 @@
+from datetime import timedelta
+from typing import Annotated
+
+from pydantic import Field
+from pydantic_settings import SettingsConfigDict
+from settings_library.rabbit import RabbitSettings
+from settings_library.redis import RedisSettings
+
+from .base import BaseCustomSettings
+
+
+class CelerySettings(BaseCustomSettings):
+    CELERY_RABBIT_BROKER: Annotated[
+        RabbitSettings, Field(json_schema_extra={"auto_default_from_env": True})
+    ]
+    CELERY_REDIS_RESULT_BACKEND: Annotated[
+        RedisSettings, Field(json_schema_extra={"auto_default_from_env": True})
+    ]
+    CELERY_RESULT_EXPIRES: Annotated[
+        timedelta,
+        Field(
+            description="Time after which task results will be deleted (default to seconds, or see https://pydantic-docs.helpmanual.io/usage/types/#datetime-types for string formating)."
+        ),
+    ] = timedelta(days=7)
+    CELERY_RESULT_PERSISTENT: Annotated[
+        bool,
+        Field(
+            description="If set to True, result messages will be persistent (after a broker restart)."
+        ),
+    ] = True
+
+    model_config = SettingsConfigDict(
+        json_schema_extra={
+            "examples": [
+                {
+                    "CELERY_RABBIT_BROKER": {
+                        "RABBIT_USER": "guest",
+                        "RABBIT_SECURE": False,
+                        "RABBIT_PASSWORD": "guest",
+                        "RABBIT_HOST": "localhost",
+                        "RABBIT_PORT": 5672,
+                    },
+                    "CELERY_REDIS_RESULT_BACKEND": {
+                        "REDIS_HOST": "localhost",
+                        "REDIS_PORT": 6379,
+                    },
+                    "CELERY_RESULT_EXPIRES": timedelta(days=1),  # type: ignore[dict-item]
+                    "CELERY_RESULT_PERSISTENT": True,
+                }
+            ],
+        }
+    )

packages/settings-library/src/settings_library/redis.py

@@ -1,6 +1,5 @@
 from enum import IntEnum
 
-from pydantic import TypeAdapter
 from pydantic.networks import RedisDsn
 from pydantic.types import SecretStr
 
@@ -18,13 +17,14 @@ class RedisDatabase(IntEnum):
     DISTRIBUTED_IDENTIFIERS = 6
     DEFERRED_TASKS = 7
     DYNAMIC_SERVICES = 8
+    CELERY_TASKS = 9
 
 
 class RedisSettings(BaseCustomSettings):
     # host
     REDIS_SECURE: bool = False
     REDIS_HOST: str = "redis"
-    REDIS_PORT: PortInt = TypeAdapter(PortInt).validate_python(6789)
+    REDIS_PORT: PortInt = 6789
 
     # auth
     REDIS_USER: str | None = None

services/docker-compose-ops.yml

@@ -95,6 +95,7 @@ services:
         distributed_identifiers:${REDIS_HOST}:${REDIS_PORT}:6:${REDIS_PASSWORD},
         deferred_tasks:${REDIS_HOST}:${REDIS_PORT}:7:${REDIS_PASSWORD},
         dynamic_services:${REDIS_HOST}:${REDIS_PORT}:8:${REDIS_PASSWORD}
+        celery_tasks:${REDIS_HOST}:${REDIS_PORT}:9:${REDIS_PASSWORD}
       # If you add/remove a db, do not forget to update the --databases entry in the docker-compose.yml
     ports:
       - "18081:8081"

services/docker-compose.yml

@@ -1177,17 +1177,28 @@ services:
       S3_ENDPOINT: ${S3_ENDPOINT}
       S3_REGION: ${S3_REGION}
       S3_SECRET_KEY: ${S3_SECRET_KEY}
+      STORAGE_WORKER_MODE: "false"
       STORAGE_LOGLEVEL: ${STORAGE_LOGLEVEL}
       STORAGE_MONITORING_ENABLED: 1
       STORAGE_PROFILING: ${STORAGE_PROFILING}
       STORAGE_PORT: ${STORAGE_PORT}
       TRACING_OPENTELEMETRY_COLLECTOR_PORT: ${TRACING_OPENTELEMETRY_COLLECTOR_PORT}
       TRACING_OPENTELEMETRY_COLLECTOR_ENDPOINT: ${TRACING_OPENTELEMETRY_COLLECTOR_ENDPOINT}
-    networks:
+    networks: &storage_networks
       - default
       - interactive_services_subnet
       - storage_subnet
 
+  sto-worker:
+    image: ${DOCKER_REGISTRY:-itisfoundation}/storage:${DOCKER_IMAGE_TAG:-master-github-latest}
+    init: true
+    hostname: "sto-worker-{{.Node.Hostname}}-{{.Task.Slot}}"
+    environment:
+      <<: *storage_environment
+      STORAGE_WORKER_MODE: "true"
+      CELERY_CONCURRENCY: 1
+    networks: *storage_networks
+
   rabbit:
     image: itisfoundation/rabbitmq:3.11.2-management
     init: true
@@ -1272,7 +1283,7 @@ services:
         "--loglevel",
         "verbose",
         "--databases",
-        "9",
+        "10",
         "--appendonly",
         "yes",
         "--requirepass",

services/storage/Makefile

@@ -15,7 +15,7 @@ openapi.json: .env-ignore
 	@set -o allexport; \
 	source $<; \
 	set +o allexport; \
-	python3 -c "import json; from $(APP_PACKAGE_NAME).main import *; print( json.dumps(the_app.openapi(), indent=2) )" > $@
+	python3 -c "import json; from $(APP_PACKAGE_NAME).main import *; print( json.dumps(app.openapi(), indent=2) )" > $@
 
 	# validates OAS file: $@
 	$(call validate_openapi_specs,$@)

services/storage/docker/boot.sh

@@ -56,8 +56,17 @@ if [ "${SC_BOOT_MODE}" = "debug" ]; then
       --log-level \"${SERVER_LOG_LEVEL}\"
   "
 else
-  exec uvicorn simcore_service_storage.main:the_app \
-    --host 0.0.0.0 \
-    --port ${STORAGE_PORT} \
-    --log-level "${SERVER_LOG_LEVEL}"
+  if [ "${STORAGE_WORKER_MODE}" = "true" ]; then
+    exec celery \
+      --app=simcore_service_storage.modules.celery.worker_main:app \
+      worker --pool=threads \
+      --loglevel="${SERVER_LOG_LEVEL}" \
+      --hostname="${HOSTNAME}" \
+      --concurrency="${CELERY_CONCURRENCY}"
+  else
+    exec uvicorn simcore_service_storage.main:app \
+      --host 0.0.0.0 \
+      --port ${STORAGE_PORT} \
+      --log-level "${SERVER_LOG_LEVEL}"
+  fi
 fi

services/storage/docker/healthcheck.py

@@ -20,18 +20,51 @@
 
 
 import os
+import socket
+import subprocess
 import sys
 from urllib.request import urlopen
 
+from simcore_service_storage.main import ApplicationSettings
+
 SUCCESS, UNHEALTHY = 0, 1
 
 # Disabled if boots with debugger
-ok = os.environ.get("SC_BOOT_MODE", "").lower() == "debug"
+ok = os.getenv("SC_BOOT_MODE", "").lower() == "debug"
 
 # Queries host
 # pylint: disable=consider-using-with
+
+
+def _is_celery_worker_healthy():
+    app_settings = ApplicationSettings.create_from_envs()
+
+    assert app_settings.STORAGE_CELERY
+    broker_url = app_settings.STORAGE_CELERY.CELERY_RABBIT_BROKER.dsn
+
+    try:
+        result = subprocess.run(
+            [
+                "celery",
+                "--broker",
+                broker_url,
+                "inspect",
+                "ping",
+                "--destination",
+                "celery@" + socket.gethostname(),
+            ],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        return "pong" in result.stdout
+    except subprocess.CalledProcessError:
+        return False
+
+
 ok = (
     ok
+    or (bool(os.getenv("STORAGE_WORKER_MODE", "") and _is_celery_worker_healthy()))
     or urlopen(
         "{host}{baseurl}".format(
             host=sys.argv[1], baseurl=os.environ.get("SIMCORE_NODE_BASEPATH", "")

services/storage/requirements/_base.in

@@ -16,7 +16,9 @@
 
 aioboto3 # s3 storage
 aiofiles # i/o
+asgi_lifespan
 asyncpg # database
+celery[redis]
 httpx
 opentelemetry-instrumentation-botocore
 packaging