From d348a446ae2dfd0407cedd9ce332170bd5c57287 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Sat, 17 Oct 2026 11:47:50 +0300 Subject: [PATCH] fix: keep the root container off app.conf so spawn children start setup_di now registers the container in a weak registry keyed by app, so pickling the app no longer pickles the container. In a spawn child the worker_process_init handler binds its container to the unpickled current app when the main names match. Closes #41 --- CONTEXT.md | 10 +++--- README.md | 8 +++-- modern_di_celery/main.py | 31 +++++++++++++---- tests/test_spawn.py | 73 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 108 insertions(+), 14 deletions(-) create mode 100644 tests/test_spawn.py diff --git a/CONTEXT.md b/CONTEXT.md index 41a7171..48eb15d 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -1,7 +1,7 @@ # modern-di-celery -A Celery adapter over [`modern-di`](https://github.com/modern-python/modern-di): it stores a root -container on the Celery app, opens and closes it on the worker lifecycle signals, and resolves a +A Celery adapter over [`modern-di`](https://github.com/modern-python/modern-di): it registers a root +container for the Celery app, opens and closes it on the worker lifecycle signals, and resolves a task's marked parameters from a child container built per task invocation. ## Language @@ -16,9 +16,9 @@ one. Celery's vocabulary is Celery's — task, worker, broker, beat, and the poo [its docs](https://docs.celeryq.dev) say they mean. The three below are this package's own. **Root container**: -The `Scope.APP` `Container` handed to `setup_di` and stored on the Celery app's `conf`; the worker -lifecycle signals open and close it, and every per-task child descends from it. `fetch_di_container` -returns this one. +The `Scope.APP` `Container` handed to `setup_di` and registered for the Celery app, outside its +`conf`; the worker lifecycle signals open and close it, and every per-task child descends from it. +`fetch_di_container` returns this one. _Avoid_: APP container, APP-scope container — both have been used for it, and neither says the thing that matters at a call site, which is that this is the container a per-task child is built *from*. diff --git a/README.md b/README.md index 0304f30..1bd0ce4 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,7 @@ uv add modern-di-celery # or: pip install modern-di-celery ## Usage -`setup_di` stores the container on `app.conf` and registers handlers that open and close it on `worker_process_init`/`worker_process_shutdown` (once per forked process under the prefork and solo pools) and on `worker_init`/`worker_shutdown` (once in the main worker process, the only signals the threads, gevent, and eventlet pools send). `@inject` builds a `Scope.REQUEST` child container per task call and resolves the `FromDI`-marked parameters from it; alternatively set `task_cls=DITask` to inject every task without a per-task decorator. +`setup_di` registers the container for the app and connects handlers that open and close it on `worker_process_init`/`worker_process_shutdown` (once per forked process under the prefork and solo pools) and on `worker_init`/`worker_shutdown` (once in the main worker process, the only signals the threads, gevent, and eventlet pools send). `@inject` builds a `Scope.REQUEST` child container per task call and resolves the `FromDI`-marked parameters from it; alternatively set `task_cls=DITask` to inject every task without a per-task decorator. ```python import typing @@ -76,11 +76,15 @@ def greet( These worker signals fire only when a real `celery worker` starts and stops, so a script or test that runs tasks eagerly (`task_always_eager = True`) must drive the container lifecycle itself, either by sending those signals or by opening and closing the container directly. Celery tasks are sync, so REQUEST-scoped finalizers run via `close_sync()`; there is no framework request/message object, so no context provider is registered. +### Prefork with spawn + +Celery runs its prefork pool on billiard, its own copy of multiprocessing, which starts children with the fork method by default on macOS and Linux. With spawn (Windows, or `billiard.set_start_method("spawn", force=True)`), each child unpickles the app and re-imports your task modules, so `setup_di` runs again in the child. The `worker_process_init` handler then registers the child's container for the app the child runs. It matches apps by `main` name, so give each app in a process its own name. Under spawn, Celery's prefork pool also needs `FORKED_BY_MULTIPROCESSING=1` in the environment. Without it tasks fail inside Celery's tracer, with or without modern-di. + ## API | Symbol | Description | |---|---| -| `setup_di(app, container)` | Stores the root container (`Scope.APP`) on `app.conf` and opens/closes it on `worker_process_init`/`worker_process_shutdown` and `worker_init`/`worker_shutdown`. Returns the container | +| `setup_di(app, container)` | Registers the root container (`Scope.APP`) for the app, outside `app.conf`, and opens/closes it on `worker_process_init`/`worker_process_shutdown` and `worker_init`/`worker_shutdown`. Returns the container | | `FromDI(dependency)` | Inert marker for `Annotated[T, FromDI(...)]` in task signatures; accepts a provider instance or a type | | `inject(task)` | Decorator that builds a `Scope.REQUEST` child per call, resolves the `FromDI`-annotated parameters, and closes the child with `close_sync()` afterwards. Raises `TypeError` at decoration if the task also declares `*args`/`**kwargs`. Raises `RuntimeError` naming `setup_di` when a task reaches it without `setup_di` called | | `DITask` | `Task` subclass that applies `@inject` to a task's `run` automatically; pass `task_cls=DITask` to `Celery(...)` or `base=DITask` to `@app.task(...)` | diff --git a/modern_di_celery/main.py b/modern_di_celery/main.py index 7419c0b..0dca1b0 100644 --- a/modern_di_celery/main.py +++ b/modern_di_celery/main.py @@ -1,23 +1,35 @@ import inspect import typing +import weakref from celery import Celery, Task, current_app, signals +from celery.local import Proxy from celery.utils.functional import head_from_fun from modern_di import Container, Scope, integrations -# The root container lives on ``app.conf`` under this named key; read back with -# ``current_app.conf`` inside ``inject``. A named constant keeps writer and -# reader in provable agreement. -_ROOT_CONTAINER_KEY = "modern_di_container" +_ROOT_CONTAINERS: "weakref.WeakKeyDictionary[Celery, Container]" = weakref.WeakKeyDictionary() + + +def _real_app(app: Celery) -> Celery: + if isinstance(app, Proxy): + return typing.cast(Celery, app._get_current_object()) # noqa: SLF001 + return app def setup_di(app: Celery, container: Container) -> Container: - app.conf[_ROOT_CONTAINER_KEY] = container + _ROOT_CONTAINERS[app] = container + main = app.main def _open_container(**_: typing.Any) -> None: # noqa: ANN401 container.open() + def _open_and_bind_container(**_: typing.Any) -> None: # noqa: ANN401 + container.open() + process_app = _real_app(typing.cast(Celery, current_app)) + if process_app.main == main: + _ROOT_CONTAINERS.setdefault(process_app, container) + def _close_container(**_: typing.Any) -> None: # noqa: ANN401 container.close_sync() @@ -34,7 +46,12 @@ def _close_container(**_: typing.Any) -> None: # noqa: ANN401 # ContainerClosedError. The overlap between the two pairs is harmless: # Container.open() is a no-op when already open, and close_sync() is a no-op # when nothing was cached. - signals.worker_process_init.connect(_open_container, weak=False) + # + # A spawn child unpickles the app it is sent and makes that copy current, then + # re-imports the task modules, so setup_di runs again on a new app object. The + # process-init handler binds its container to that current copy by main name, + # unless the copy already has one. + signals.worker_process_init.connect(_open_and_bind_container, weak=False) signals.worker_process_shutdown.connect(_close_container, weak=False) signals.worker_init.connect(_open_container, weak=False) signals.worker_shutdown.connect(_close_container, weak=False) @@ -43,7 +60,7 @@ def _close_container(**_: typing.Any) -> None: # noqa: ANN401 def fetch_di_container(app: Celery) -> Container: try: - return typing.cast(Container, app.conf[_ROOT_CONTAINER_KEY]) + return _ROOT_CONTAINERS[_real_app(app)] except KeyError: msg = ( "No modern-di container found on the app. " diff --git a/tests/test_spawn.py b/tests/test_spawn.py new file mode 100644 index 0000000..ece0dfa --- /dev/null +++ b/tests/test_spawn.py @@ -0,0 +1,73 @@ +import pickle +import typing + +import pytest +from celery import Celery, current_app, signals +from modern_di import Container + +from modern_di_celery import fetch_di_container, setup_di +from tests.dependencies import Dependencies + + +def _make_app(main: str, set_as_current: bool = True) -> Celery: + return Celery(main, broker="memory://", backend="cache+memory://", set_as_current=set_as_current) + + +@pytest.fixture(autouse=True) +def _restore_current_app() -> typing.Iterator[None]: + previous = current_app._get_current_object() # noqa: SLF001 + yield + previous.set_current() + + +def test_a_configured_app_survives_a_pickle_round_trip() -> None: + app = _make_app("pickled") + setup_di(app, Container(groups=[Dependencies])) + + assert isinstance(pickle.loads(pickle.dumps(app)), Celery) # noqa: S301 + + +def test_main_survives_a_pickle_round_trip() -> None: + app = _make_app("pickled-main") + setup_di(app, Container(groups=[Dependencies])) + + assert pickle.loads(pickle.dumps(app)).main == "pickled-main" # noqa: S301 + + +def test_a_spawn_child_binds_the_recreated_container_to_its_current_app() -> None: + payload = pickle.dumps(_make_app("spawn-child", set_as_current=False)) + unpickled = pickle.loads(payload) # noqa: S301 + recreated = _make_app("spawn-child", set_as_current=False) + container = setup_di(recreated, Container(groups=[Dependencies])) + container.close_sync() + unpickled.set_current() + + signals.worker_process_init.send(sender=None) + + assert fetch_di_container(typing.cast(Celery, current_app)) is container + assert container.closed is False + + +def test_apps_with_different_main_names_keep_their_own_containers() -> None: + alpha = _make_app("alpha", set_as_current=False) + beta = _make_app("beta", set_as_current=False) + alpha_container = setup_di(alpha, Container(groups=[Dependencies])) + beta_container = setup_di(beta, Container(groups=[Dependencies])) + pickle.loads(pickle.dumps(alpha)).set_current() # noqa: S301 + + signals.worker_process_init.send(sender=None) + + assert fetch_di_container(typing.cast(Celery, current_app)) is alpha_container + assert fetch_di_container(alpha) is alpha_container + assert fetch_di_container(beta) is beta_container + + +def test_an_app_with_a_container_keeps_it_when_a_later_handler_fires() -> None: + current = _make_app("shared-main") + later = _make_app("shared-main", set_as_current=False) + current_container = setup_di(current, Container(groups=[Dependencies])) + setup_di(later, Container(groups=[Dependencies])) + + signals.worker_process_init.send(sender=None) + + assert fetch_di_container(current) is current_container