Bound tensorization process-pool shutdown at interpreter exit

[bradhilton]

Summary

• shut the shared tensorization ProcessPoolExecutor down explicitly at interpreter exit, before concurrent.futures joins its workers, so the join is bounded
• idle workers still leave via the normal shutdown sentinel; workers busy with work nobody can consume anymore are terminated after a 5 s grace (then killed), so cancellation and interpreter shutdown never orphan children or wait on them indefinitely
• resolve a wandb-artifact:/// checkpoint's tokenizer configuration once per process instead of making a W&B API request on every tokenization
• keep the adaptive thread/process selection, worker reuse across calls, ordered results, source rebinding, and the thread fallback unchanged; experiment scripts still need no pool or cleanup code

Root cause of the shutdown hang

concurrent.futures.process._python_exit runs inside threading._shutdown() and joins the pool's manager thread. The manager only sends shutdown sentinels once pending_work_items is empty, so any worker still executing an item (for example a tensorization whose awaiting task was cancelled by SIGINT-driven cleanup) blocks interpreter exit for as long as that item runs. ART never shut the global pool down itself, so the stdlib owned the final join.

Reproduced with a worker holding one in-flight item at exit: on main the driver never exits (faulthandler: main thread in process._python_exit → join, manager thread in wait_result_broken_or_wakeup); with this change the same script exits in 5.8 s with no surviving children. The idle-pool path (normal exit, cancellation, worker exception, SIGINT) already exited in ~2 s on main and still does. The exact item that kept the 046 workers busy was not captured (py-spy cannot attach on that host), so the fix bounds the join regardless of cause.

Checkpoint tokenizer configs

_artifact_config fetched the artifact from W&B on every call; only the base model name was cached, and only alias-independently. Experiment 046 passes model=<wandb-artifact checkpoint> to art.tensorize, so every trajectory tokenized in every worker paid a 0.2–0.5 s API round trip (measured against the 046 checkpoint artifact). The resolved _TokenizerConfig is now cached per model string (bounded, single-flight, exceptions not cached), and callers receive a copy because _tokenizer_config mutates the config when base_model is given. The metadata describes the checkpoint collection, so it is stable across versions and aliases, the same assumption the existing base-model cache already made.

Validation

• new test_interpreter_exit_is_bounded_after_process_tokenization: real spawn pool, two art.tokenize calls reuse the same worker PIDs, then one never-finishing item plus a cancelled tokenization at exit; asserts exit code 0, exit within 30 s of the last result (observed ~2 s beyond the 1 s test grace), and no surviving worker PIDs. Against main the same test hits a 150 s timeout.
• new test_process_exit_hook_runs_before_stdlib_joins_workers
• new test_checkpoint_tokenizer_configs_are_resolved_once_per_process: one API call across repeated and concurrent resolutions, independent copies, base_model override does not leak into the cache
• pytest tests/unit/trajectories/test_tokenize.py tests/unit/trajectories/test_parallel_tokenize.py (253 passed)
• Ruff and ty clean for the changed files; full ty check src tests unchanged from main

🤖 Generated with Claude Code