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/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_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") 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: