From 46c176f1c253334d959132825895242ca3b7c7f3 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 4 Sep 2026 18:50:48 +0000 Subject: [PATCH 1/2] Bound tensorization process-pool shutdown at interpreter exit The shared spawn-based ProcessPoolExecutor behind art.tokenize/tensorize was never shut down by ART, so concurrent.futures owned the final join during threading._shutdown(). Its manager thread only sends the worker shutdown sentinels once no work item is pending, so a worker still busy with an item nobody can consume anymore (for example after SIGINT-driven cleanup cancelled the awaiting task) blocked interpreter exit for as long as that item ran, leaving drivers alive with idle-looking children. Register a threading exit hook that runs before the stdlib join: release the global pool, shut it down without waiting, give workers a short grace period to leave via the sentinel, then terminate and finally kill any survivor. Idle pools still exit through the normal sentinel path, worker reuse during the process lifetime is unchanged, and the thread fallback and broken-pool handling are untouched. Co-Authored-By: Claude Fable 5.1 --- src/art/trajectories/_parallel.py | 66 +++++++++- .../trajectories/test_parallel_tokenize.py | 116 ++++++++++++++++++ 2 files changed, 179 insertions(+), 3 deletions(-) diff --git a/src/art/trajectories/_parallel.py b/src/art/trajectories/_parallel.py index 0940bba0e..433aa5b5e 100644 --- a/src/art/trajectories/_parallel.py +++ b/src/art/trajectories/_parallel.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import atexit from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from concurrent.futures import Future, ProcessPoolExecutor, ThreadPoolExecutor from concurrent.futures.process import BrokenProcessPool @@ -8,6 +9,7 @@ from functools import lru_cache import math import multiprocessing +from multiprocessing.process import BaseProcess import os from pathlib import Path import pickle @@ -35,6 +37,7 @@ _PROCESS_MAX_WORKERS = 4 _PROCESS_MIN_ITEMS = 4 _PROCESS_MIN_THREAD_SECONDS = 1.0 +_PROCESS_EXIT_GRACE_SECONDS = 5.0 def _cgroup_cpu_limit() -> int | None: @@ -179,17 +182,74 @@ def _process_executor(capacity: int) -> ProcessPoolExecutor: return executor -def _discard_process_executor() -> None: +def _release_process_executor() -> ProcessPoolExecutor | None: global _PROCESS_EXECUTOR, _PROCESS_EXECUTOR_CAPACITY, _PROCESS_EXECUTOR_PID global _PROCESS_STARTUP with _PROCESS_EXECUTOR_LOCK: previous = _PROCESS_EXECUTOR + owned = _PROCESS_EXECUTOR_PID == os.getpid() _PROCESS_EXECUTOR = None _PROCESS_EXECUTOR_PID = None _PROCESS_EXECUTOR_CAPACITY = 0 _PROCESS_STARTUP = () - if previous is not None: - previous.shutdown(wait=False, cancel_futures=True) + return previous if owned else None + + +def _discard_process_executor() -> None: + previous = _release_process_executor() + if previous is not None: + previous.shutdown(wait=False, cancel_futures=True) + + +def _process_executor_workers(executor: ProcessPoolExecutor) -> list[BaseProcess]: + processes = getattr(executor, "_processes", None) + return list(processes.values()) if processes else [] + + +def _shutdown_process_executor(grace: float | None = None) -> None: + """Stop the shared process pool within a bounded time. + + Runs before concurrent.futures joins its workers at interpreter exit. Idle + workers leave as soon as they read the shutdown sentinel; workers still busy + 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: + return + if grace is None: + grace = _PROCESS_EXIT_GRACE_SECONDS + workers = _process_executor_workers(executor) + executor.shutdown(wait=False, cancel_futures=True) + deadline = time.monotonic() + max(0.0, grace) + for worker in workers: + worker.join(max(0.0, deadline - time.monotonic())) + for worker in workers: + if worker.is_alive(): + worker.terminate() + for worker in workers: + worker.join(1.0) + for worker in workers: + if worker.is_alive(): + worker.kill() + worker.join(1.0) + + +def _register_process_exit_hook() -> None: + # threading's private atexit list runs before concurrent.futures joins its + # worker threads and processes, which is the only point early enough to + # bound that join. Fall back to atexit where the hook is unavailable. + register = getattr(threading, "_register_atexit", None) + if register is not None: + try: + register(_shutdown_process_executor) + return + except RuntimeError: + return + atexit.register(_shutdown_process_executor) + + +_register_process_exit_hook() @dataclass diff --git a/tests/unit/trajectories/test_parallel_tokenize.py b/tests/unit/trajectories/test_parallel_tokenize.py index 2fb22a778..046a02994 100644 --- a/tests/unit/trajectories/test_parallel_tokenize.py +++ b/tests/unit/trajectories/test_parallel_tokenize.py @@ -426,6 +426,122 @@ def test_process_context_does_not_reexecute_unguarded_script( assert marker.read_text() == "run\n" +def test_process_exit_hook_runs_before_stdlib_joins_workers() -> None: + from concurrent.futures import process as stdlib_process + + registered: list[Any] = getattr(threading, "_threading_atexits") + hooks = [getattr(hook, "func", hook) for hook in registered] + + assert _parallel._shutdown_process_executor in hooks + # threading runs its exit hooks in reverse registration order. + assert hooks.index(_parallel._shutdown_process_executor) > hooks.index( + stdlib_process._python_exit + ) + _parallel._shutdown_process_executor() # no pool: nothing to do + + +def test_interpreter_exit_is_bounded_after_process_tokenization( + tmp_path: Path, +) -> None: + script = tmp_path / "exit_after_process_tokenization.py" + script.write_text( + textwrap.dedent( + """ + import asyncio + import multiprocessing + import time + + from openai.types.chat import ChatCompletion + + import art + from art.trajectories import _parallel + + _parallel._supports_processes = lambda **_: True + _parallel._processes_enabled = lambda *_: True + _parallel._PROCESS_EXIT_GRACE_SECONDS = 1.0 + + + def trajectory(index): + token = index + 2 + completion = ChatCompletion.model_validate( + { + "id": f"chat-{index}", + "object": "chat.completion", + "created": 0, + "model": "test/model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "answer"}, + "prompt_token_ids": [1], + "token_ids": [token], + "logprobs": { + "content": [ + { + "token": f"token_id:{token}", + "logprob": -0.2, + "bytes": [], + "top_logprobs": [], + } + ] + }, + } + ], + } + ) + return art.Trajectory( + messages_and_choices=[completion.choices[0]], + metadata={"index": index}, + ) + + + async def main(): + trajectories = [trajectory(index) for index in range(4)] + first = await art.tokenize(trajectories, model="test/model") + workers = sorted(p.pid for p in multiprocessing.active_children()) + second = await art.tokenize(trajectories, model="test/model") + assert [value.tokens for value in first] == [[1, 2], [1, 3], [1, 4], [1, 5]] + assert [value.trajectory for value in second] == trajectories + assert workers == sorted(p.pid for p in multiprocessing.active_children()) + print("WORKERS", *workers, flush=True) + # Leave one worker busy with work nobody will ever consume, and one + # tokenization cancelled mid-flight, as an interrupted driver would. + _parallel._PROCESS_EXECUTOR.submit(time.sleep, 600) + task = asyncio.ensure_future(art.tokenize(trajectories, model="test/model")) + await asyncio.sleep(0) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + + asyncio.run(main()) + print("MAIN_DONE", time.time(), flush=True) + """ + ) + ) + + completed = subprocess.run( + [sys.executable, str(script)], + capture_output=True, + text=True, + timeout=600, + cwd=Path.cwd(), + ) + finished_at = time.time() + + assert completed.returncode == 0, completed.stderr + lines = dict(line.split(" ", 1) for line in completed.stdout.splitlines()) + workers = [int(pid) for pid in lines["WORKERS"].split()] + assert len(workers) >= 2 + assert finished_at - float(lines["MAIN_DONE"]) < 30 + for pid in workers: + with pytest.raises(ProcessLookupError): + os.kill(pid, 0) + + def test_child_deserialization_failure_is_a_transfer_error() -> None: with pytest.raises(_parallel._ProcessTransferError, match="process input"): _parallel._tokenize_process_payload(b"not a pickle") From 414d7c997da0a9fe01b27afef847c194af95ee76 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 4 Sep 2026 19:05:27 +0000 Subject: [PATCH 2/2] Resolve checkpoint tokenizer configs once per process _artifact_config fetched the W&B artifact on every call, so tokenizing with a wandb-artifact:/// model paid an API round trip per trajectory in every tensorization worker; only the base model name was cached. Cache the resolved _TokenizerConfig per model string (bounded, single-flight, exceptions not cached) and hand callers a copy, since _tokenizer_config adjusts the config when base_model is given. The metadata describes the checkpoint collection and does not change between versions, the same assumption the base-model cache already makes. Co-Authored-By: Claude Fable 5.1 --- src/art/trajectories/_tokenize.py | 18 +++++++- tests/unit/trajectories/test_tokenize.py | 56 ++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 8daecb7fc..e9bbe3ab9 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -4,7 +4,7 @@ import codecs from collections.abc import Mapping, Sequence from copy import deepcopy -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime from functools import lru_cache from hashlib import sha256 @@ -1503,7 +1503,23 @@ def _artifact_base_model(identity: str) -> str: return base_model +_ARTIFACT_CONFIG_LOCK = threading.Lock() + + def _artifact_config(model: str) -> _TokenizerConfig: + """Resolve a checkpoint's tokenizer configuration, once per process. + + The metadata describes the checkpoint collection, so it is stable across + versions and aliases; re-fetching it would cost a W&B API round trip on every + tokenization. Callers get their own copy because _tokenizer_config adjusts it. + """ + with _ARTIFACT_CONFIG_LOCK: + config = _cached_artifact_config(model) + return replace(config) + + +@lru_cache(maxsize=1024) +def _cached_artifact_config(model: str) -> _TokenizerConfig: from wandb.apis.public import Api artifact_path = _artifact_name(model) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 6b2faf73b..6666d5510 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -3114,6 +3114,9 @@ def artifact(self, name: str) -> SimpleNamespace: monkeypatch.setitem(sys.modules, "wandb", wandb) monkeypatch.setitem(sys.modules, "wandb.apis", apis) monkeypatch.setitem(sys.modules, "wandb.apis.public", public) + from art.trajectories._tokenize import _cached_artifact_config + + _cached_artifact_config.cache_clear() exchange = _chat_exchange([], [], model=model) extra = exchange.response.choices[0].model_extra assert extra is not None @@ -3141,6 +3144,59 @@ def artifact(self, name: str) -> SimpleNamespace: assert tokenizer.calls[0]["thinking"] is True +def test_checkpoint_tokenizer_configs_are_resolved_once_per_process( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from concurrent.futures import ThreadPoolExecutor + + from art.trajectories._tokenize import ( + _ARTIFACT_BASE_MODELS, + _cached_artifact_config, + _tokenizer_config, + ) + + artifact_names: list[str] = [] + + class Api: + def artifact(self, name: str) -> SimpleNamespace: + artifact_names.append(name) + return SimpleNamespace( + metadata={ + "base_model": "base/model", + "renderer": {"chat_template_kwargs": {"thinking": True}}, + } + ) + + wandb = ModuleType("wandb") + apis = ModuleType("wandb.apis") + public = ModuleType("wandb.apis.public") + setattr(public, "Api", Api) + setattr(apis, "public", public) + setattr(wandb, "apis", apis) + monkeypatch.setitem(sys.modules, "wandb", wandb) + monkeypatch.setitem(sys.modules, "wandb.apis", apis) + monkeypatch.setitem(sys.modules, "wandb.apis.public", public) + model = "wandb-artifact:///entity/project/cached-run" + _cached_artifact_config.cache_clear() + _ARTIFACT_BASE_MODELS.pop("entity/project/cached-run", None) + try: + first = _tokenizer_config(model, None) + overridden = _tokenizer_config(model, "other/model") + second = _tokenizer_config(model, None) + with ThreadPoolExecutor(max_workers=4) as executor: + parallel = list(executor.map(_tokenizer_config, [model] * 4, [None] * 4)) + finally: + _cached_artifact_config.cache_clear() + _ARTIFACT_BASE_MODELS.pop("entity/project/cached-run", None) + + assert artifact_names == ["entity/project/cached-run:latest"] + assert first == second and first is not second + assert first.base_model == "base/model" + assert overridden.base_model == "other/model" + assert second.base_model == "base/model" + assert all(config == first for config in parallel) + + def test_loaded_tokenizers_are_cached_by_model_and_revision( monkeypatch: pytest.MonkeyPatch, ) -> None: