fix(bench): wire _init_worker into pools, force RAYON_NUM_THREADS=1

Author: nikolay-e

benchmarks/adapters/calibrate.py

@@ -92,9 +92,16 @@ def evaluate_grid(
 
     from concurrent.futures.process import BrokenProcessPool
 
+    from benchmarks.common import _init_worker
+
     def _make_pool() -> ProcessPoolExecutor:
         ctx = mp.get_context("spawn")
-        p = ProcessPoolExecutor(max_workers=workers, mp_context=ctx, max_tasks_per_child=50)
+        p = ProcessPoolExecutor(
+            max_workers=workers,
+            mp_context=ctx,
+            max_tasks_per_child=50,
+            initializer=_init_worker,
+        )
         list(p.map(int, range(workers)))  # eager-spawn all workers
         return p
 
@@ -193,9 +200,16 @@ def evaluate_grid_cached(  # noqa: C901 — pool teardown + per-cell demux + ret
         if needed:
             pending.append((inst, needed))
 
+    from benchmarks.common import _init_worker
+
     def _make_pool() -> ProcessPoolExecutor:
         ctx = mp.get_context("spawn")
-        p = ProcessPoolExecutor(max_workers=workers, mp_context=ctx, max_tasks_per_child=40)
+        p = ProcessPoolExecutor(
+            max_workers=workers,
+            mp_context=ctx,
+            max_tasks_per_child=40,
+            initializer=_init_worker,
+        )
         list(p.map(int, range(workers)))
         return p
 

benchmarks/adapters/runner.py

@@ -297,7 +297,14 @@ def _drain(active_pool: ProcessPoolExecutor) -> None:  # noqa: C901
         _drain(pool)  # type: ignore[arg-type]
         return
 
+    from benchmarks.common import _init_worker
+
     ctx = mp.get_context("spawn")
-    with ProcessPoolExecutor(max_workers=workers, mp_context=ctx, max_tasks_per_child=50) as owned:
+    with ProcessPoolExecutor(
+        max_workers=workers,
+        mp_context=ctx,
+        max_tasks_per_child=50,
+        initializer=_init_worker,
+    ) as owned:
         _drain(owned)
         owned.shutdown(wait=False, cancel_futures=True)

benchmarks/common.py

@@ -302,11 +302,17 @@ def _init_worker() -> None:
 
     warnings.filterwarnings("ignore", category=SyntaxWarning)
 
-    os.environ.setdefault("RAYON_NUM_THREADS", os.environ.get("BENCH_RAYON_THREADS", "1"))
-    os.environ.setdefault("OMP_NUM_THREADS", "1")
-    os.environ.setdefault("MKL_NUM_THREADS", "1")
-    os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
-    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+    # Explicit assignment, not setdefault: parent may have inherited a
+    # different value (e.g. ambient RAYON_NUM_THREADS=24 from a shell
+    # profile) which would defeat the cap. With N workers each running
+    # Rayon at full core count, we get N x cores threads contending for
+    # the same cores — the calibration timeout breaks before the
+    # algorithm finishes.
+    os.environ["RAYON_NUM_THREADS"] = os.environ.get("BENCH_RAYON_THREADS", "1")
+    os.environ["OMP_NUM_THREADS"] = "1"
+    os.environ["MKL_NUM_THREADS"] = "1"
+    os.environ["OPENBLAS_NUM_THREADS"] = "1"
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
     for mod in (
         "_diffctx",

benchmarks/diffctx_eval_fn.py

@@ -85,6 +85,14 @@ def _arm_diffctx_kill_switch(timeout_s: float):
     import threading
 
     if timeout_s <= 0:
+        # Trace once per worker so we know whether the env propagated.
+        if not getattr(_arm_diffctx_kill_switch, "_warned_no_timeout", False):
+            print(
+                f"[worker pid={os.getpid()}] WARN: DIFFCTX_BENCH_TIMEOUT_SEC unset "
+                f"or <=0, timer disabled (raw value: {os.environ.get('DIFFCTX_BENCH_TIMEOUT_SEC', 'MISSING')!r})",
+                flush=True,
+            )
+            _arm_diffctx_kill_switch._warned_no_timeout = True  # type: ignore[attr-defined]
         return None
     timer = threading.Timer(timeout_s, lambda: os._exit(137))
     timer.daemon = True