fix(infra): persistent XLA compile cache in harvest/clustering workers; one policy definition

[claude-spd1]

Description

Harvest and clustering workers never enabled JAX's persistent compilation cache, so a fanned-out harvest recompiled the identical full 8B open_jax_run forward once per worker, per invocation — while the trainer's cache dir already sat on shared FS.

Commit 1 — the cache fix:

• Hoists the cache policy into core as param_decomp.compile_cache.enable_persistent_compilation_cache — now the ONE definition. Called at startup by the LM composition root, pretrain.train, and both workers (harvest/scripts/run_worker.py, clustering/scripts/run_worker.py). The pretrain copy had already drifted (missing min_entry_size_bytes, int threshold, dead out_dir is None guard) — deduped.
• Policy decision (the audit's "confirm intent" item, decided here since no record explains the old value): jax_persistent_cache_min_compile_time_secs 60 → 5 (jax default: 1). Under 60s, eval-tier and init compiles (1–60s) recompiled on every requeue, and the M1 fix's value would silently hinge on whether the worker forward happens to compile in over a minute. 5s still keeps trivial slice/utility jits from churning the dir. Happy to bump back if someone knows why 60 was chosen.
• Documents at the hoisted site: rank-0's xla_dump_* XLA_FLAGS don't fork its cache key (jax strips debug options when keying); jax_persistent_cache_enable_xla_caches='all' must stay unset (runtime CUDA_ERROR_NOT_FOUND on jax 0.10.1/B200; the default already persists the per-fusion autotune cache on jax ≥ 0.10).

Commit 2 — sanctioned rider (audit-classed consistency LOW, folded in per triage since compiler_options enter the compile-cache key — threading them later would invalidate the worker cache entries this PR creates): the run's pinned runtime.compiler_options now reach open_jax_run's consumer forward and make_arithmetic_grid_step, the two step factories that compiled with bare defaults. _take_columns (trivial gather) deliberately stays without options, noted at the site.

Related Issue

Bridge task jax-worker-persistent-cache (2026-07-10 JAX compile audit: M1 + cache-policy LOW + compiler_options LOW; ledger in lore 2026-07-10--jax-compile-audit-remaining-findings).

Motivation and Context

A fanned-out harvest pays the multi-minute forward compile N times for identical HLO; requeued trainer runs additionally re-paid every 1–60s compile. Known open item NOT addressed here (belongs to task jax-compile-observability): the 2026-06-19 observation of the trainer cache not writing new entries.

How Has This Been Tested?

• make check clean (ruff + basedpyright, whole workspace).
• pytest over the touched surfaces: clustering ensemble pipeline + membership bridge, harvest batch, pretrain smoke, arithmetic eval — 20 passed.
• Functional wiring check on CPU: enable_persistent_compilation_cache on a temp dir → jax auto-creates the dir and writes entries on compile (verified with threshold dropped to 0 so a toy jit qualifies).

Does this PR introduce a breaking change?

No. Behavioral deltas: workers now populate/read the shared xla_compilation_cache sibling of runs/; compiles in the 5–60s band now cache; the consumer forward and arithmetic grid step compile under the run's pinned compiler_options.

🤖 Generated with Claude Code

Crew-Address: task/jax-worker-persistent-cache