Bound tensorization process-pool shutdown at interpreter exit

src/art/trajectories/_parallel.py

@@ -1,13 +1,15 @@
 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
 from dataclasses import dataclass, field
 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

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)

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")

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: