From 058e583d730b3edb7bf76a6f2d51955450371dc2 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 26 Sep 2026 06:24:32 +0000 Subject: [PATCH 1/3] Fix tokenization process-pool readiness under startup skew --- src/art/trajectories/_parallel.py | 55 +++- .../trajectories/test_process_readiness.py | 248 ++++++++++++++++++ 2 files changed, 292 insertions(+), 11 deletions(-) create mode 100644 tests/unit/trajectories/test_process_readiness.py diff --git a/src/art/trajectories/_parallel.py b/src/art/trajectories/_parallel.py index 433aa5b5e..d33951dd5 100644 --- a/src/art/trajectories/_parallel.py +++ b/src/art/trajectories/_parallel.py @@ -10,6 +10,7 @@ import math import multiprocessing from multiprocessing.process import BaseProcess +from multiprocessing.synchronize import Barrier import os from pathlib import Path import pickle @@ -38,6 +39,7 @@ _PROCESS_MIN_ITEMS = 4 _PROCESS_MIN_THREAD_SECONDS = 1.0 _PROCESS_EXIT_GRACE_SECONDS = 5.0 +_PROCESS_STARTUP_TIMEOUT_SECONDS = 60.0 def _cgroup_cpu_limit() -> int | None: @@ -94,15 +96,24 @@ def _executor(capacity: int) -> ThreadPoolExecutor: _PROCESS_EXECUTOR_CAPACITY = 0 _PROCESS_STARTUP: tuple[Future[int], ...] = () _PROCESS_BACKEND_DISABLED = False +_PROCESS_READY_BARRIER: Barrier | None = None def _process_context() -> multiprocessing.context.BaseContext: return multiprocessing.get_context("spawn") +def _initialize_process_worker(barrier: Barrier) -> None: + global _PROCESS_READY_BARRIER + _PROCESS_READY_BARRIER = barrier + + def _process_identity() -> int: - # Keep every submitted warmup occupied until the bounded pool is started. - time.sleep(0.25) + # A task cannot finish and be reused by the same worker until every worker + # has entered this bounded handshake, regardless of import/startup skew. + if _PROCESS_READY_BARRIER is None: + raise RuntimeError("process worker readiness barrier is missing") + _PROCESS_READY_BARRIER.wait() return os.getpid() @@ -123,7 +134,11 @@ def _submit_process_warmup( def _finish_process_warmup(futures: tuple[Future[int], ...], capacity: int) -> None: if not futures: return - worker_pids = {future.result() for future in futures} + deadline = time.monotonic() + _PROCESS_STARTUP_TIMEOUT_SECONDS + worker_pids = { + future.result(timeout=max(0.0, deadline - time.monotonic())) + for future in futures + } if len(worker_pids) != capacity: raise RuntimeError( f"started {len(worker_pids)} process workers, expected {capacity}" @@ -144,14 +159,20 @@ def _start_process_executor( or _PROCESS_EXECUTOR_CAPACITY < process_capacity ): previous = _PROCESS_EXECUTOR if _PROCESS_EXECUTOR_PID == pid else None + context = _process_context() + barrier = context.Barrier( + process_capacity, timeout=_PROCESS_STARTUP_TIMEOUT_SECONDS + ) executor = ProcessPoolExecutor( max_workers=process_capacity, - mp_context=_process_context(), + mp_context=context, + initializer=_initialize_process_worker, + initargs=(barrier,), ) try: startup = _submit_process_warmup(executor, process_capacity) except BaseException: - executor.shutdown(wait=False, cancel_futures=True) + _shutdown_process_executor(0, executor) raise _PROCESS_EXECUTOR = executor _PROCESS_EXECUTOR_PID = pid @@ -177,7 +198,11 @@ def _complete_process_warmup( def _process_executor(capacity: int) -> ProcessPoolExecutor: executor, process_capacity, startup = _start_process_executor(capacity) - _finish_process_warmup(startup, process_capacity) + try: + _finish_process_warmup(startup, process_capacity) + except (OSError, RuntimeError): + _shutdown_process_executor(0, executor) + raise _complete_process_warmup(executor, startup) return executor @@ -206,7 +231,9 @@ def _process_executor_workers(executor: ProcessPoolExecutor) -> list[BaseProcess return list(processes.values()) if processes else [] -def _shutdown_process_executor(grace: float | None = None) -> None: +def _shutdown_process_executor( + grace: float | None = None, executor: ProcessPoolExecutor | None = None +) -> None: """Stop the shared process pool within a bounded time. Runs before concurrent.futures joins its workers at interpreter exit. Idle @@ -214,7 +241,8 @@ def _shutdown_process_executor(grace: float | None = None) -> None: with tensorization nobody can consume anymore are terminated after ``grace`` seconds so the interpreter never waits on them indefinitely. """ - executor = _release_process_executor() + if executor is None: + executor = _release_process_executor() if executor is None: return if grace is None: @@ -598,6 +626,7 @@ async def _ordered_process_map( ) -> list[TokenizedTrajectory | TokenizedMultiHistoryTrajectory]: loop = asyncio.get_running_loop() thread_executor = _executor(capacity) + executor: ProcessPoolExecutor | None = None try: executor, process_capacity, startup = _start_process_executor(capacity) await loop.run_in_executor( @@ -607,9 +636,13 @@ async def _ordered_process_map( process_capacity, ) _complete_process_warmup(executor, startup) - except BrokenProcessPool: - raise - except (OSError, RuntimeError) as error: + except (BrokenProcessPool, OSError, RuntimeError) as error: + if executor is not None: + await loop.run_in_executor( + thread_executor, _shutdown_process_executor, 0, executor + ) + if isinstance(error, BrokenProcessPool): + raise raise _ProcessBackendError( f"could not start process workers: {type(error).__name__}: {error}" ) from None diff --git a/tests/unit/trajectories/test_process_readiness.py b/tests/unit/trajectories/test_process_readiness.py new file mode 100644 index 000000000..263d0535c --- /dev/null +++ b/tests/unit/trajectories/test_process_readiness.py @@ -0,0 +1,248 @@ +from __future__ import annotations + +import asyncio +from concurrent.futures import Future, ProcessPoolExecutor +from concurrent.futures.process import BrokenProcessPool +import os +import threading +import time +from types import SimpleNamespace +from typing import Any + +import pytest +from test_parallel_tokenize import _exchange_trajectory + +import art +from art.trajectories import _parallel + + +def _held_initializer(barrier: Any, counter: Any, release: Any) -> None: + with counter.get_lock(): + index = counter.value + counter.value += 1 + if index and not release.wait(15): + raise RuntimeError("public delayed initializer timed out") + if barrier is not None: + _parallel._initialize_process_worker(barrier) + + +def _old_identity() -> int: + time.sleep(0.25) + return os.getpid() + + +def _failed_initializer() -> None: + raise RuntimeError("public initializer failure") + + +def test_short_tasks_are_not_a_worker_census() -> None: + context = _parallel._process_context() + release = context.Event() + counter = context.Value("i", 0) + with ProcessPoolExecutor( + max_workers=2, + mp_context=context, + initializer=_held_initializer, + initargs=(None, counter, release), + ) as pool: + try: + futures = [pool.submit(_old_identity) for _ in range(2)] + assert len({f.result(timeout=20) for f in futures}) == 1 + finally: + release.set() + + +async def test_skewed_startup_and_cancelled_waiter_keep_true_readiness( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _parallel._shutdown_process_executor(grace=0) + context = _parallel._process_context() + release = context.Event() + counter = context.Value("i", 0) + created: list[Any] = [] + + def factory(**kwargs: Any) -> ProcessPoolExecutor: + assert kwargs["initializer"] is _parallel._initialize_process_worker + kwargs["initializer"] = _held_initializer + kwargs["initargs"] = (*kwargs["initargs"], counter, release) + pool = ProcessPoolExecutor(**kwargs) + created.append(pool) + return pool + + monkeypatch.setattr(_parallel, "ProcessPoolExecutor", factory) + try: + pool, capacity, futures = _parallel._start_process_executor(2) + task = asyncio.create_task( + asyncio.to_thread(_parallel._finish_process_warmup, futures, capacity) + ) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + deadline = time.monotonic() + 20 + while counter.value < 2: + assert time.monotonic() < deadline + await asyncio.sleep(0.02) + await asyncio.sleep(0.6) # Longer than two old 0.25-second identity tasks. + assert not any(f.done() for f in futures) + release.set() + await asyncio.wait_for( + asyncio.to_thread(_parallel._finish_process_warmup, futures, capacity), 30 + ) + assert len({f.result() for f in futures}) == capacity == 2 + _parallel._complete_process_warmup(pool, futures) + assert _parallel._process_executor(2) is pool + assert list(pool.map(abs, [-2, -1])) == [2, 1] + finally: + release.set() + workers = [p for pool in created for p in (pool._processes or {}).values()] + _parallel._shutdown_process_executor(grace=2) + assert all(not p.is_alive() for p in workers) + + +def test_initializer_failure_propagates_and_pool_is_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _parallel._shutdown_process_executor(grace=0) + + def factory(**kwargs: Any) -> ProcessPoolExecutor: + kwargs["initializer"] = _failed_initializer + kwargs["initargs"] = () + return ProcessPoolExecutor(**kwargs) + + monkeypatch.setattr(_parallel, "ProcessPoolExecutor", factory) + try: + _, capacity, futures = _parallel._start_process_executor(2) + with pytest.raises(BrokenProcessPool): + _parallel._finish_process_warmup(futures, capacity) + finally: + _parallel._shutdown_process_executor(grace=0) + + +def test_missing_initializer_and_barrier_timeout_refuse( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(_parallel, "_PROCESS_READY_BARRIER", None) + with pytest.raises(RuntimeError, match="barrier is missing"): + _parallel._process_identity() + monkeypatch.setattr( + _parallel, "_PROCESS_READY_BARRIER", threading.Barrier(2, timeout=0.02) + ) + with pytest.raises(threading.BrokenBarrierError): + _parallel._process_identity() + + +def test_duplicate_pid_still_refuses() -> None: + futures: tuple[Future[int], ...] = (Future(), Future()) + for future in futures: + future.set_result(123) + with pytest.raises(RuntimeError, match="started 1 process workers, expected 2"): + _parallel._finish_process_warmup(futures, 2) + + +def test_parent_wait_has_one_total_deadline(monkeypatch: pytest.MonkeyPatch) -> None: + observed: list[float] = [] + clock = iter([100.0, 101.0, 104.0]) + monkeypatch.setattr( + _parallel, "time", SimpleNamespace(monotonic=lambda: next(clock)) + ) + monkeypatch.setattr(_parallel, "_PROCESS_STARTUP_TIMEOUT_SECONDS", 5.0) + + class Ready: + def result(self, *, timeout: float) -> int: + observed.append(timeout) + return len(observed) + + _parallel._finish_process_warmup((Ready(), Ready()), 2) # type: ignore[arg-type] + assert observed == [4.0, 1.0] + + +def test_parent_timeout_when_no_worker_reaches_barrier( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(_parallel, "_PROCESS_STARTUP_TIMEOUT_SECONDS", 0.02) + start = time.monotonic() + with pytest.raises(TimeoutError): + _parallel._finish_process_warmup((Future(),), 1) + assert time.monotonic() - start < 1.0 + + +async def test_public_timeout_falls_back_without_changing_results( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(_parallel, "_PROCESS_BACKEND_DISABLED", False) + monkeypatch.setattr(_parallel, "_supports_processes", lambda **_: True) + monkeypatch.setattr(_parallel, "_processes_enabled", lambda *_, **__: True) + monkeypatch.setattr(_parallel, "_cpu_capacity", lambda: 2) + + def unavailable(_: int) -> Any: + raise TimeoutError("public bounded startup timeout") + + monkeypatch.setattr(_parallel, "_start_process_executor", unavailable) + trajectories = [_exchange_trajectory(i) for i in range(4)] + with pytest.warns(RuntimeWarning, match="using threads"): + result = await art.tokenize(trajectories, model="test/model") + assert [value.tokens for value in result] == [[1, 2], [1, 3], [1, 4], [1, 5]] + assert [value.trajectory for value in result] == trajectories + assert _parallel._PROCESS_BACKEND_DISABLED + + +async def test_timeout_stops_exact_spawned_family( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _parallel._shutdown_process_executor(grace=0) + monkeypatch.setattr(_parallel, "_PROCESS_STARTUP_TIMEOUT_SECONDS", 0.05) + original = _parallel._submit_process_warmup + workers: list[Any] = [] + + def submit(pool: Any, capacity: int) -> Any: + result = original(pool, capacity) + workers.extend(pool._processes.values()) + return result + + monkeypatch.setattr(_parallel, "_submit_process_warmup", submit) + try: + with pytest.raises(_parallel._ProcessBackendError, match="TimeoutError"): + await _parallel._ordered_process_map([], [], workers=2, capacity=2) + assert len(workers) == 2 and all(not p.is_alive() for p in workers) + finally: + _parallel._shutdown_process_executor(grace=0) + + +def test_partial_submit_failure_stops_unpublished_pool( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _parallel._shutdown_process_executor(grace=0) + original = _parallel._submit_process_warmup + workers: list[Any] = [] + failure = RuntimeError("public submission failure") + + def submit(pool: Any, capacity: int) -> Any: + original(pool, capacity) + workers.extend(pool._processes.values()) + raise failure + + monkeypatch.setattr(_parallel, "_submit_process_warmup", submit) + with pytest.raises(RuntimeError) as caught: + _parallel._start_process_executor(2) + assert caught.value is failure + assert len(workers) == 2 and all(not p.is_alive() for p in workers) + assert _parallel._PROCESS_EXECUTOR is None + + +def test_failed_pool_cleanup_does_not_release_replacement( + monkeypatch: pytest.MonkeyPatch, +) -> None: + replacement = object() + monkeypatch.setattr(_parallel, "_PROCESS_EXECUTOR", replacement) + closed: list[dict[str, bool]] = [] + + class Failed: + _processes: dict[int, Any] = {} + + def shutdown(self, **kwargs: bool) -> None: + closed.append(kwargs) + + _parallel._shutdown_process_executor(0, Failed()) # type: ignore[arg-type] + assert _parallel._PROCESS_EXECUTOR is replacement + assert closed == [{"wait": False, "cancel_futures": True}] From fb59da75f7bd997964c681b05717ef9645a5127c Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 26 Sep 2026 06:32:37 +0000 Subject: [PATCH 2/3] Type readiness test doubles explicitly --- tests/unit/trajectories/test_process_readiness.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/unit/trajectories/test_process_readiness.py b/tests/unit/trajectories/test_process_readiness.py index 263d0535c..db299359a 100644 --- a/tests/unit/trajectories/test_process_readiness.py +++ b/tests/unit/trajectories/test_process_readiness.py @@ -7,7 +7,7 @@ import threading import time from types import SimpleNamespace -from typing import Any +from typing import Any, cast import pytest from test_parallel_tokenize import _exchange_trajectory @@ -153,7 +153,9 @@ def result(self, *, timeout: float) -> int: observed.append(timeout) return len(observed) - _parallel._finish_process_warmup((Ready(), Ready()), 2) # type: ignore[arg-type] + _parallel._finish_process_warmup( + cast(tuple[Future[int], ...], (Ready(), Ready())), 2 + ) assert observed == [4.0, 1.0] @@ -243,6 +245,6 @@ class Failed: def shutdown(self, **kwargs: bool) -> None: closed.append(kwargs) - _parallel._shutdown_process_executor(0, Failed()) # type: ignore[arg-type] + _parallel._shutdown_process_executor(0, cast(ProcessPoolExecutor, Failed())) assert _parallel._PROCESS_EXECUTOR is replacement assert closed == [{"wait": False, "cancel_futures": True}] From b052ce9f9b547e4dc17f39fbac26fb194593f678 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 26 Sep 2026 07:05:41 +0000 Subject: [PATCH 3/3] Coordinate worker cleanup with the executor manager --- src/art/trajectories/_parallel.py | 8 +- .../trajectories/test_process_readiness.py | 118 ++++++++++++++++++ 2 files changed, 125 insertions(+), 1 deletion(-) diff --git a/src/art/trajectories/_parallel.py b/src/art/trajectories/_parallel.py index d33951dd5..c3f35fa5d 100644 --- a/src/art/trajectories/_parallel.py +++ b/src/art/trajectories/_parallel.py @@ -248,6 +248,7 @@ def _shutdown_process_executor( if grace is None: grace = _PROCESS_EXIT_GRACE_SECONDS workers = _process_executor_workers(executor) + manager = getattr(executor, "_executor_manager_thread", None) executor.shutdown(wait=False, cancel_futures=True) deadline = time.monotonic() + max(0.0, grace) for worker in workers: @@ -257,10 +258,15 @@ def _shutdown_process_executor( worker.terminate() for worker in workers: worker.join(1.0) + deadline = time.monotonic() + 1.0 for worker in workers: if worker.is_alive(): worker.kill() - worker.join(1.0) + worker.join(max(0.0, deadline - time.monotonic())) + # The manager may reap a child concurrently with the joins above. Wait for + # it to publish the exit status before returning ownership to the caller. + if manager is not None and manager is not threading.current_thread(): + manager.join(max(0.0, deadline - time.monotonic())) def _register_process_exit_hook() -> None: diff --git a/tests/unit/trajectories/test_process_readiness.py b/tests/unit/trajectories/test_process_readiness.py index db299359a..df8c9b059 100644 --- a/tests/unit/trajectories/test_process_readiness.py +++ b/tests/unit/trajectories/test_process_readiness.py @@ -3,6 +3,7 @@ import asyncio from concurrent.futures import Future, ProcessPoolExecutor from concurrent.futures.process import BrokenProcessPool +from multiprocessing.process import BaseProcess import os import threading import time @@ -232,6 +233,70 @@ def submit(pool: Any, capacity: int) -> Any: assert _parallel._PROCESS_EXECUTOR is None +def test_cleanup_joins_manager_before_reporting_reaped_worker( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _parallel._shutdown_process_executor(grace=0) + original_submit = _parallel._submit_process_warmup + original_waitpid = os.waitpid + original_join = BaseProcess.join + original_thread_join = threading.Thread.join + workers: list[Any] = [] + managers: list[Any] = [] + reaped = threading.Event() + publish = threading.Event() + failure = RuntimeError("public partial submission during manager reaping") + + def waitpid(pid: int, flags: int) -> tuple[int, int]: + result = original_waitpid(pid, flags) + if threading.current_thread() in managers and result[0] and not reaped.is_set(): + reaped.set() + assert publish.wait(10) + return result + + def join(process: Any, timeout: float | None = None) -> None: + if ( + threading.current_thread() is threading.main_thread() + and process in workers + and timeout == 1.0 + ): + assert reaped.wait(5) + original_join(process, timeout) + + def thread_join(thread: Any, timeout: float | None = None) -> None: + if thread in managers and threading.current_thread() is threading.main_thread(): + assert timeout is not None and 0 <= timeout <= 1.0 + publish.set() + original_thread_join(thread, timeout) + + def submit(pool: Any, capacity: int) -> Any: + original_submit(pool, capacity) + workers.extend(pool._processes.values()) + managers.append(pool._executor_manager_thread) + raise failure + + monkeypatch.setattr(os, "waitpid", waitpid) + monkeypatch.setattr(BaseProcess, "join", join) + monkeypatch.setattr(threading.Thread, "join", thread_join) + monkeypatch.setattr(_parallel, "_submit_process_warmup", submit) + try: + with pytest.raises(RuntimeError) as caught: + _parallel._start_process_executor(2) + assert caught.value is failure + assert reaped.is_set() and publish.is_set() + assert len(workers) == 2 and all(not p.is_alive() for p in workers) + assert all(not manager.is_alive() for manager in managers) + finally: + publish.set() + monkeypatch.undo() + for manager in managers: + manager.join(10) + for worker in workers: + if worker.is_alive(): + worker.kill() + worker.join(2) + + def test_failed_pool_cleanup_does_not_release_replacement( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -248,3 +313,56 @@ def shutdown(self, **kwargs: bool) -> None: _parallel._shutdown_process_executor(0, cast(ProcessPoolExecutor, Failed())) assert _parallel._PROCESS_EXECUTOR is replacement assert closed == [{"wait": False, "cancel_futures": True}] + + +def test_manager_wait_shares_final_reap_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + clock = iter([100.0, 100.0, 101.0, 101.8, 102.0]) + monkeypatch.setattr( + _parallel, "time", SimpleNamespace(monotonic=lambda: next(clock)) + ) + joins: list[tuple[str, float | None]] = [] + + class Worker: + def join(self, timeout: float | None = None) -> None: + joins.append(("worker", timeout)) + + def is_alive(self) -> bool: + return True + + def terminate(self) -> None: + pass + + def kill(self) -> None: + pass + + class Manager: + def join(self, timeout: float | None = None) -> None: + joins.append(("manager", timeout)) + + class Pool: + _processes = {1: Worker()} + _executor_manager_thread = Manager() + + def shutdown(self, **kwargs: bool) -> None: + pass + + _parallel._shutdown_process_executor(0, cast(ProcessPoolExecutor, Pool())) + assert joins == [ + ("worker", 0), + ("worker", 1.0), + ("worker", pytest.approx(0.2)), + ("manager", 0), + ] + + +def test_cleanup_does_not_join_current_manager() -> None: + class Pool: + _processes: dict[int, Any] = {} + _executor_manager_thread = threading.current_thread() + + def shutdown(self, **kwargs: bool) -> None: + pass + + _parallel._shutdown_process_executor(0, cast(ProcessPoolExecutor, Pool()))