[DOP-21799] Refactor celery initialization (#165)

Author: IlyasDevelopment

docker/entrypoint_worker.sh

@@ -6,4 +6,3 @@ set -e
 
 # exec is required to forward all signals to the main process
 exec python -m celery -A syncmaster.worker.celery worker --max-tasks-per-child=1 "$@"
-

docs/conf.py

@@ -147,11 +147,6 @@
 # If true, `todo` and `todoList` produce output, else they produce nothing.
 todo_include_todos = False
 
-# set the environment variable before imports
-# TODO: remove after global init of WorkerAppSettings in worker/__init__.py
-os.environ["SYNCMASTER__ENCRYPTION__CRYPTO_KEY"] = "crypto_key"
-os.environ["SYNCMASTER__DATABASE__URL"] = "postgresql+asyncpg://syncmaster:changeme@db:5432/syncmaster"
-os.environ["SYNCMASTER__BROKER__URL"] = "amqp://guest:guest@localhost:5672/"
 
 # -- Options for HTMLHelp output ------------------------------------------
 

syncmaster/backend/__init__.py

@@ -1,5 +1,6 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
+from celery import Celery
 from fastapi import FastAPI, HTTPException
 from fastapi.exceptions import RequestValidationError
 from pydantic import ValidationError
@@ -20,6 +21,15 @@
 from syncmaster.exceptions import SyncmasterError
 
 
+def celery_factory(settings: Settings) -> Celery:
+    app = Celery(
+        __name__,
+        broker=settings.broker.url,
+        backend="db+" + settings.database.sync_url,
+    )
+    return app
+
+
 def application_factory(settings: Settings) -> FastAPI:
     application = FastAPI(
         title="Syncmaster",
@@ -30,6 +40,7 @@ def application_factory(settings: Settings) -> FastAPI:
         redoc_url=None,
     )
     application.state.settings = settings
+    application.state.celery = celery_factory(settings)
     application.include_router(api_router)
     application.exception_handler(RequestValidationError)(validation_exception_handler)
     application.exception_handler(ValidationError)(validation_exception_handler)
@@ -44,6 +55,7 @@ def application_factory(settings: Settings) -> FastAPI:
         {
             Settings: lambda: settings,
             UnitOfWork: get_uow(session_factory, settings=settings),
+            Celery: lambda: application.state.celery,
         },
     )
 

syncmaster/backend/api/v1/runs.py

@@ -5,6 +5,7 @@
 from typing import Annotated
 
 from asgi_correlation_id import correlation_id
+from celery import Celery
 from fastapi import APIRouter, Depends, Query
 from jinja2 import Template
 from kombu.exceptions import KombuError
@@ -24,7 +25,6 @@
     ReadRunSchema,
     RunPageSchema,
 )
-from syncmaster.worker import celery
 
 router = APIRouter(tags=["Runs"], responses=get_error_responses())
 
@@ -84,6 +84,7 @@ async def read_run(
 async def start_run(
     create_run_data: CreateRunSchema,
     settings: Annotated[Settings, Depends(Stub(Settings))],
+    celery: Annotated[Celery, Depends(Stub(Celery))],
     unit_of_work: UnitOfWork = Depends(UnitOfWork),
     current_user: User = Depends(get_user(is_active=True)),
 ) -> ReadRunSchema:

syncmaster/scheduler/__init__.py

@@ -1,4 +1,14 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
-from syncmaster.scheduler.transfer_fetcher import TransferFetcher
-from syncmaster.scheduler.transfer_job_manager import TransferJobManager
+from celery import Celery
+
+from syncmaster.scheduler.settings import SchedulerAppSettings
+
+
+def celery_factory(settings: SchedulerAppSettings) -> Celery:
+    app = Celery(
+        __name__,
+        broker=settings.broker.url,
+        backend="db+" + settings.database.sync_url,
+    )
+    return app

syncmaster/scheduler/celery.py

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from syncmaster.scheduler import celery_factory
+from syncmaster.scheduler.settings import SchedulerAppSettings
+
+# Global object, since the TransferJobManager.send_job_to_celery method is static
+app = celery_factory(SchedulerAppSettings())

syncmaster/scheduler/transfer_job_manager.py

@@ -9,10 +9,10 @@
 from syncmaster.backend.services.unit_of_work import UnitOfWork
 from syncmaster.db.models import RunType, Status, Transfer
 from syncmaster.exceptions.run import CannotConnectToTaskQueueError
+from syncmaster.scheduler.celery import app as celery
 from syncmaster.scheduler.settings import SchedulerAppSettings as Settings
 from syncmaster.scheduler.utils import get_async_session
 from syncmaster.schemas.v1.connections.connection import ReadAuthDataSchema
-from syncmaster.worker import celery
 
 
 class TransferJobManager:
@@ -50,8 +50,11 @@ def update_jobs(self, transfers: list[Transfer]) -> None:
     @staticmethod
     async def send_job_to_celery(transfer_id: int) -> None:
         """
-        Do not pass additional arguments like settings,
+        1. Do not pass additional arguments like settings,
         otherwise they will be serialized in jobs table.
+        2. Instance methods are bound to specific objects and cannot be reliably serialized
+        due to the weak reference problem. Use a static method instead, as it is not
+        object-specific and can be serialized.
         """
         settings = Settings()
 

syncmaster/worker/__init__.py

@@ -6,7 +6,7 @@
 from syncmaster.worker.settings import WorkerAppSettings
 
 
-def create_celery_app(settings) -> Celery:
+def celery_factory(settings: WorkerAppSettings) -> Celery:
     app = Celery(
         __name__,
         broker=settings.broker.url,
@@ -17,8 +17,3 @@ def create_celery_app(settings) -> Celery:
         ],
     )
     return app
-
-
-# TODO: initialize celery app in __name__ == "__main__"
-#       then initialize celery app in backend via dependency injection and initialize in scheduler
-celery = create_celery_app(WorkerAppSettings())

syncmaster/worker/celery.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from syncmaster.worker import celery_factory
+from syncmaster.worker.settings import WorkerAppSettings
+
+app = celery_factory(WorkerAppSettings())

syncmaster/worker/transfer.py

@@ -15,7 +15,7 @@
 from syncmaster.db.repositories.utils import decrypt_auth_data
 from syncmaster.exceptions.run import RunNotFoundError
 from syncmaster.settings.log import setup_logging
-from syncmaster.worker import celery
+from syncmaster.worker.celery import app as celery
 from syncmaster.worker.controller import TransferController
 from syncmaster.worker.settings import WorkerAppSettings