kiln_ai.sandbox.spawn

Shared multiprocessing spawn helpers.

Stdlib only — no Pydantic / Kiln-model / DB / UI imports.

 1"""Shared multiprocessing spawn helpers.
 2
 3Stdlib only — no Pydantic / Kiln-model / DB / UI imports.
 4"""
 5
 6import multiprocessing.process
 7import sys
 8import threading
 9import types
10
11_spawn_lock = threading.Lock()
12"""Process-global lock shared by code-evals and code tools.
13
14Both paths spawn through :func:`kiln_ai.tools.sandbox_bridge.run_bridged_child`,
15which calls :func:`start_process_with_light_main` below. Serializes the
16``__main__`` stub-swap window around ``p.start()``, preventing the PyInstaller
17concurrent-spawn bug #7410.
18"""
19
20
21def start_process_with_light_main(
22    p: multiprocessing.process.BaseProcess,
23) -> None:
24    """Start *p* under ``_spawn_lock`` with ``sys.modules['__main__']`` swapped for a stub.
25
26    Prevents the spawn child from re-importing the parent's heavy
27    ``__main__`` module (e.g. dev_server.py which transitively pulls
28    in litellm, google.cloud.aiplatform, etc.).
29
30    ``multiprocessing.spawn.get_preparation_data`` reads
31    ``sys.modules['__main__']`` to decide what to re-execute in the
32    child.  By swapping in a lightweight stub named ``"__main__"``
33    (with no ``__file__`` or ``__spec__``) before ``p.start()``
34    (which pickles the prep data), neither ``init_main_from_name``
35    nor ``init_main_from_path`` is set, so the child skips the
36    heavy re-import entirely.
37
38    The stub is named ``"__main__"`` (not ``"__mp_main__"``, which
39    would collide with multiprocessing's internal main-module
40    name).  Third-party code that inspects
41    ``sys.modules["__main__"].__name__`` during the brief swap
42    window sees the expected ``"__main__"`` value.
43
44    This stays within ``multiprocessing.spawn``'s existing
45    bootstrap, so ``freeze_support()`` and PyInstaller frozen
46    builds continue to work.
47
48    Thread safety: the swap is process-global, but the window is
49    sub-millisecond (only spans ``p.start()``), is serialized by
50    ``_spawn_lock``, and ``__main__`` is not read by request-
51    handling code, so cross-thread visibility is benign.
52    """
53    with _spawn_lock:
54        _real_main = sys.modules.get("__main__")
55        if _real_main is not None:
56            _light_main = types.ModuleType("__main__")
57            sys.modules["__main__"] = _light_main
58            try:
59                p.start()
60            finally:
61                sys.modules["__main__"] = _real_main
62        else:
63            p.start()
def start_process_with_light_main(p: multiprocessing.process.BaseProcess) -> None:
22def start_process_with_light_main(
23    p: multiprocessing.process.BaseProcess,
24) -> None:
25    """Start *p* under ``_spawn_lock`` with ``sys.modules['__main__']`` swapped for a stub.
26
27    Prevents the spawn child from re-importing the parent's heavy
28    ``__main__`` module (e.g. dev_server.py which transitively pulls
29    in litellm, google.cloud.aiplatform, etc.).
30
31    ``multiprocessing.spawn.get_preparation_data`` reads
32    ``sys.modules['__main__']`` to decide what to re-execute in the
33    child.  By swapping in a lightweight stub named ``"__main__"``
34    (with no ``__file__`` or ``__spec__``) before ``p.start()``
35    (which pickles the prep data), neither ``init_main_from_name``
36    nor ``init_main_from_path`` is set, so the child skips the
37    heavy re-import entirely.
38
39    The stub is named ``"__main__"`` (not ``"__mp_main__"``, which
40    would collide with multiprocessing's internal main-module
41    name).  Third-party code that inspects
42    ``sys.modules["__main__"].__name__`` during the brief swap
43    window sees the expected ``"__main__"`` value.
44
45    This stays within ``multiprocessing.spawn``'s existing
46    bootstrap, so ``freeze_support()`` and PyInstaller frozen
47    builds continue to work.
48
49    Thread safety: the swap is process-global, but the window is
50    sub-millisecond (only spans ``p.start()``), is serialized by
51    ``_spawn_lock``, and ``__main__`` is not read by request-
52    handling code, so cross-thread visibility is benign.
53    """
54    with _spawn_lock:
55        _real_main = sys.modules.get("__main__")
56        if _real_main is not None:
57            _light_main = types.ModuleType("__main__")
58            sys.modules["__main__"] = _light_main
59            try:
60                p.start()
61            finally:
62                sys.modules["__main__"] = _real_main
63        else:
64            p.start()

Start p under _spawn_lock with sys.modules['__main__'] swapped for a stub.

Prevents the spawn child from re-importing the parent's heavy __main__ module (e.g. dev_server.py which transitively pulls in litellm, google.cloud.aiplatform, etc.).

multiprocessing.spawn.get_preparation_data reads sys.modules['__main__'] to decide what to re-execute in the child. By swapping in a lightweight stub named "__main__" (with no __file__ or __spec__) before p.start() (which pickles the prep data), neither init_main_from_name nor init_main_from_path is set, so the child skips the heavy re-import entirely.

The stub is named "__main__" (not "__mp_main__", which would collide with multiprocessing's internal main-module name). Third-party code that inspects sys.modules["__main__"].__name__ during the brief swap window sees the expected "__main__" value.

This stays within multiprocessing.spawn's existing bootstrap, so freeze_support() and PyInstaller frozen builds continue to work.

Thread safety: the swap is process-global, but the window is sub-millisecond (only spans p.start()), is serialized by _spawn_lock, and __main__ is not read by request- handling code, so cross-thread visibility is benign.