fix: keep the root container off app.conf so spawn children start

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*.
 

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(...)` |

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. "

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