fix(slurm): scrub SLURM_*/SBATCH_* from the sbatch submission env

[danbraunai-goodfire]

Description

sbatch propagates the submitting shell's env into the new job, and SLURM reads SLURM_* vars as directives — so submitting from inside an allocation (dev pod, launcher job, agent session, a tmux server created in an srun step) carries the caller's step/allocation state into the new job. Two observed failure modes of the same leak:

• Loud: an inherited SLURM_CPUS_PER_TASK conflicts with the new job's own allocation and its inner srun dies in ~2s with "cpus-per-task set by two different environment variables".
• Silent (the expensive one): an inherited SLURM_CPU_BIND pins every rank to the originating pod's few cores. This was the root cause of the June 2026 fleet-wide ~3× slowdown — jobs sbatch'd from tmux windows (the tmux server was created inside an srun --overlap step) ran 8 ranks + dataloaders on 4 cores, cost days of debugging and nearly a vendor escalation before being root-caused on our side.

One line of substance in infra/slurm.py::_submit_script: the sbatch subprocess runs with os.environ minus SLURM_*/SBATCH_*/PMIX_*, so the new job's SLURM state comes only from its own allocation. Covers every pd-* submitter (they share _submit_script). The June incident's fix scrubbed one source (the tmux server env); this scrubs at the sink, so pd-* submissions are immune regardless of what launched them. PMIX_* is included to match what the June root-cause fix stripped.

Nothing the ranks rely on travels via inheritance: the run env is exported explicitly in the generated sbatch script from runtime.launch_env (the config-captured single source of truth); ambient inheritance of everything else (PATH, HOME, tokens) is untouched. One deliberate behavior change: exporting SBATCH_* option-vars (e.g. SBATCH_PARTITION) in your shell no longer steers pd-* submissions — submissions are config-driven by design.

Extracted from #941 (same base commit); #941/#954 dedupe on rebase once this merges.

Related Issue

Extracted from #941. Sink-side counterpart of the June 2026 CPU-affinity incident (Slack thread).

Motivation and Context

Launching from anywhere other than a bare login shell currently fails in ~2s with a confusing srun error — and the same leak class already caused a silent fleet-wide 3× slowdown once.

How Has This Been Tested?

All of this week's compile-probe launches (~15 pd-lm submissions, incl. an 8-node dp64 run) were submitted from inside 1-CPU launcher jobs on and-btdr with the SLURM_*/SBATCH_* scrub — every submission's inner srun started cleanly, which fails without it.

Does this PR introduce a breaking change?

No (aside from the deliberate SBATCH_* note above).

🤖 Generated with Claude Code

https://claude.ai/code/session_01E34cK5BaVr6ZFAwRidNp6n

[danbraunai-goodfire]

Tested with a 6-arm matrix across clusters, submission methods, and dp (full 8L-slice llama8b runs, 30 steps each):

arm	cluster	method	dp	result
negative control (no scrub, plain feature/jax)	and-btdr	inside 4-CPU allocation	8	❌ child's srun step OOM-killed at 55s during model load — inherited SLURM_MEM_* caps from the launcher allocation (a third symptom of the same leak)
scrubbed, same context as neg control	and-btdr	inside allocation	8	✅ COMPLETED, 30/30 steps
scrubbed	and-gmi (H200)	inside allocation	8	✅ COMPLETED, 30/30 steps — step-30 losses match btdr to 4 digits
scrubbed	and-gmi	inside allocation	16 (2 nodes)	✅ COMPLETED, 30/30 steps
scrubbed + poisoned June-incident env (SLURM_CPU_BIND=map_cpu:0 exported at submit)	and-btdr	inside allocation	16	✅ live rank probe: Cpus_allowed_list: 0-191 (full node, not the poison), SLURM_CPU_BIND absent from rank /proc/environ
unpoisoned differential control for the above	and-btdr	inside allocation	16	reproduced a btdr-side NCCL failure identically without any poison (ncclAlltoAll "Too many open files" / ncclOsSocketTryAccept hang) — orthogonal to this PR, btdr multi-node issue on this config

Two side findings: (1) the and-gmi login node's CPU lacks AVX so jax cannot import there — on gmi every pd-lm submission necessarily comes from inside an allocation, i.e. the path this PR fixes is the only path; (2) the btdr dp16 NCCL socket failure above is a real, separate cluster issue (also hit a dp64 probe arm yesterday) and deserves its own investigation.

Logs archived: scrub_test_2026-07-06_logs.tgz in each cluster's artifacts root.

🤖 Generated with Claude Code