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