feat(multipool): clean up orphaned .snapshot_scratch partials

Scratch partials under runs/<id>/.snapshot_scratch/step_<S>/ were only ever
removed by consolidate_step (after it assembles a step's checkpoint). When the
async consolidation job stalls, is preempted, or is never submitted, nothing
removes them and they accumulate — in practice runs that never consolidated a
single step each held multiple TB of dead partials (one sweep reclaimed ~6.3 TB).

Two small, complementary cleanups, both safe because resume reads only the
consolidated training_<S>.pth, never raw partials:

  * _prune_scratch_through(scratch_dir, S) in consolidate_step (off the loop):
    drops every step_<=S>, not just S. The final save is the highest step, so its
    consolidation sweeps the whole run — a completed run ends with zero scratch,
    even if some intermediate consolidations were skipped.
  * prune_old_scratch(scratch_dir, keep_last_n=2) in snapshot() (rank 0, off the
    collective): keeps only the newest KEEP_LAST_N_SCRATCH snapshots — the backstop
    for when consolidation never runs at all, so a stuck run can't grow without
    bound (residue at most keep_last_n step dirs).

CPU-only unit tests for both helpers; three_pool/CLAUDE.md updated.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: danbraunai-goodfire

param_decomp_lab/tests/test_consolidate_scratch.py

@@ -0,0 +1,68 @@
+"""Unit tests for scratch cleanup — the two complementary mechanisms.
+
+CPU-only, no torch: pure filesystem logic.
+
+  * ``_prune_scratch_through``: consolidation drops everything at/below the step it
+    just finished — so a completed run's scratch goes fully to zero.
+  * ``prune_old_scratch``: the train-loop backstop keeps only the newest N snapshots,
+    bounding scratch when consolidation never runs at all.
+"""
+
+from pathlib import Path
+
+import pytest
+
+from param_decomp_lab.three_pool.consolidate import (
+    SNAPSHOT_SCRATCH_DIRNAME,
+    _prune_scratch_through,
+    prune_old_scratch,
+)
+
+
+def _make_scratch(tmp_path: Path, steps: list[int]) -> Path:
+    scratch = tmp_path / SNAPSHOT_SCRATCH_DIRNAME
+    for s in steps:
+        (scratch / f"step_{s}").mkdir(parents=True)
+    return scratch
+
+
+def _steps(scratch: Path) -> set[int]:
+    return {int(d.name.removeprefix("step_")) for d in scratch.glob("step_*")}
+
+
+def test_through_drops_at_or_below_step_keeps_newer(tmp_path: Path) -> None:
+    scratch = _make_scratch(tmp_path, [20, 40, 60])
+    _prune_scratch_through(scratch, step=40)
+    assert _steps(scratch) == {60}  # 20, 40 consolidated/superseded; 60 still pending
+
+
+def test_through_final_step_clears_a_finished_run(tmp_path: Path) -> None:
+    # The highest step's consolidation sweeps everything — even a backlog of earlier,
+    # never-consolidated steps — so a completed run leaves zero scratch (not N).
+    scratch = _make_scratch(tmp_path, [20, 40, 60])
+    _prune_scratch_through(scratch, step=60)
+    assert _steps(scratch) == set()
+
+
+def test_old_keeps_newest_n_drops_the_rest(tmp_path: Path) -> None:
+    scratch = _make_scratch(tmp_path, [20, 40, 60])
+    prune_old_scratch(scratch, keep_last_n=2)
+    assert _steps(scratch) == {40, 60}  # ordered by step number, not name/mtime
+
+
+def test_old_noop_when_at_or_under_limit(tmp_path: Path) -> None:
+    scratch = _make_scratch(tmp_path, [10])
+    prune_old_scratch(scratch, keep_last_n=2)
+    assert _steps(scratch) == {10}
+
+
+def test_noop_when_scratch_absent(tmp_path: Path) -> None:
+    absent = tmp_path / SNAPSHOT_SCRATCH_DIRNAME
+    prune_old_scratch(absent)  # glob on a missing dir is empty — must not raise
+    _prune_scratch_through(absent, step=10)
+
+
+def test_old_keep_last_n_must_be_at_least_one(tmp_path: Path) -> None:
+    # keep_last_n=0 would delete the just-written step out from under its consolidation.
+    with pytest.raises(AssertionError):
+        prune_old_scratch(_make_scratch(tmp_path, [10]), keep_last_n=0)

param_decomp_lab/three_pool/CLAUDE.md

@@ -19,7 +19,7 @@ layer's output independently — is a separate, stable concept the chunkwise poo
 | `config.py` | `ThreePoolTopology` (`ci` / `ppgd` / `chunkwise` `PoolSpec`s of per-rank batch + `sites_per_chunk`) + `resolve(ordered_sites, batch_size) -> ResolvedLayout`. Authors per-rank batch, NOT rank ids; the resolver derives ranks/chunks/world_size in canonical order. Parse-time validation = cross-divisibility of the three per-rank batches |
 | `layout.py` | `World` topology + the runtime `Chunk` (`.ranks` / `.sites`); `build_world` constructs every process group (threading `pg_timeout` into each); `BatchEdge` — symmetric per-edge batch-slice geometry (CI↔chunk, CI↔PPGD) answering routing for both fan directions |
 | `checkpoint.py` | offline state_dict assembly from on-disk partials (`assemble_model_state_dict_from_partials`) + the leader key-partition helpers (`owned_model_state_keys` / `ci_fn_state_keys`) |
-| `consolidate.py` | `consolidate_step` — async, off-train-loop, **streaming** assembly of `model_<step>.pth` + `training_<step>.pth` from a step's small (parameter-shaped) scratch partials; prunes old `training_*.pth` + `ppgd_*/`; deletes the scratch dir. The data-shaped PPGD sources are NOT here — each adversary rank writes its own `ppgd_<step>/rank_<r>.pth` at snapshot time, in parallel (see "Checkpoint save" below). `load_ppgd_shard` reads a rank's shard on resume; `unconsolidated_steps` lists recoverable steps |
+| `consolidate.py` | `consolidate_step` — async, off-train-loop, **streaming** assembly of `model_<step>.pth` + `training_<step>.pth` from a step's small (parameter-shaped) scratch partials; prunes old `training_*.pth` + `ppgd_*/`; drops scratch at/below the step (`_prune_scratch_through`). The data-shaped PPGD sources are NOT here — each adversary rank writes its own `ppgd_<step>/rank_<r>.pth` at snapshot time, in parallel (see "Checkpoint save" below). `load_ppgd_shard` reads a rank's shard on resume; `unconsolidated_steps` lists recoverable steps. `prune_old_scratch` (called on-loop from `snapshot()`) bounds scratch when consolidation stalls/never runs |
 | `consolidate_cli.py` | `python -m …consolidate_cli <run> [--step N]` — manual CPU-only recovery for a failed/preempted async consolidation (separate module to avoid an import cycle with `experiments.lm.run`) |
 | `role.py` | `PoolRole = CIRole \| ChunkRole \| PPGDRole` — this rank's pool role; per-pool fields are union variants, not optional attrs |
 | `context.py` | `PoolContext = CIContext \| ChunkContext \| PPGDContext` — `world` + `role` + this pool's portals; the trainer matches on it to dispatch step fns |
@@ -186,7 +186,9 @@ the assembled CPU model, never the full partial set) → assembles the full
 `ComponentModel` state_dict + `ThreePoolTrainingState` → writes `model_<S>.pth` +
 `training_<S>.pth` → prunes old `training_*.pth` to the last
 `DEFAULT_KEEP_LAST_N_TRAINING` (=3; **all `model_*.pth` are kept**) → prunes old
-`ppgd_*/` → deletes `step_<S>/`. It runs inside the async slow-eval job
+`ppgd_*/` → drops scratch at/below `S` (`_prune_scratch_through`: `step_<S>/` plus any
+older step, consolidated or not — resume uses the newest checkpoint, so they're dead). It
+runs inside the async slow-eval job
 (`experiments/lm/async_eval.py`), as a CPU-only phase BEFORE the eval pass, so the
 assembled `model_<S>.pth` exists before the eval loads it. It reads **only** the
 small partials — the PPGD shards were already written at snapshot time, so
@@ -210,6 +212,23 @@ resume uses the newest checkpoint. Idempotent: a no-op if `training_<S>.pth`
 already exists (and it cleans any leftover scratch in that case); the scratch dir
 is deleted only on success.
 
+**Scratch cleanup (two mechanisms).** `consolidate_step` is the only thing that
+removes scratch, so when its async job stalls, is preempted, or is never submitted,
+partials pile up — runs that never consolidated a single step were observed holding
+**multiple TB** of dead `step_<S>/` partials. Two complementary cleanups bound it,
+both safe because resume reads only the consolidated `training_<S>.pth`, never raw
+partials:
+
+  * **`_prune_scratch_through(scratch_dir, S)`** (in `consolidate_step`, off the loop):
+    drops every `step_<=S>`, not just `S`. Because the final save is the highest step,
+    its consolidation sweeps the whole run — a *completed* run ends with **zero**
+    scratch, even if some intermediate consolidations were skipped.
+  * **`prune_old_scratch(scratch_dir, keep_last_n=2)`** (in `snapshot()`, rank 0, off
+    the collective): keeps only the newest `KEEP_LAST_N_SCRATCH` snapshots. The backstop
+    for when consolidation *never* runs — without it a stuck run grows without bound;
+    with it the residue is at most `keep_last_n` step dirs. (Two, not one, so the
+    previous save's still-in-flight consolidation keeps its partials.)
+
 This is the fix for how run 34446 (p-a5b667e9) died: the old synchronous rank-0
 read held the other ranks at a barrier past the NCCL watchdog. There is no
 on-loop read to outrun the watchdog now.
@@ -393,7 +412,8 @@ that whole packet stays bf16 rather than splitting the buffer by dtype (off the
 consolidation finishes, resume from the previous consolidated step — at most one
 save-interval of lost progress (the scratch partials for the unconsolidated step
 are left on disk; they can be consolidated manually via `consolidate_step` if
-that interval matters).
+that interval matters — but only the newest `KEEP_LAST_N_SCRATCH` snapshots survive
+the on-loop `prune_old_scratch`, so don't count on partials from many saves back).
 
 `from_snapshot` validates the saved topology against the current one, but the
 comparison runs on EVERY rank, so it compares only the **rank-invariant** core

param_decomp_lab/three_pool/consolidate.py

@@ -55,6 +55,12 @@
 resume uses the newest checkpoint, and an older checkpoint with its shards pruned
 still resumes (the adversary re-warms via ``n_warmup``)."""
 
+KEEP_LAST_N_SCRATCH = 2
+"""How many snapshots' scratch to retain on the train loop. ``consolidate_step``
+clears each step's scratch as it lands, so this only bites when that async job stalls
+or never runs — the backstop that stops partials growing without bound. Two, not one,
+leaves the previous save's still-in-flight consolidation its partials."""
+
 
 def step_scratch_dir(scratch_dir: Path, step: int) -> Path:
     return scratch_dir / f"step_{step}"
@@ -109,7 +115,7 @@ def consolidate_step(
         # that wrote the checkpoint but died before removing it (e.g. crashed in
         # prune), so this step stops looking unconsolidated.
         logger.info(f"consolidate: {training_path.name} already exists; skipping")
-        shutil.rmtree(step_dir, ignore_errors=True)
+        _prune_scratch_through(scratch_dir, step)
         return
 
     if not step_dir.is_dir():
@@ -188,8 +194,8 @@ def consolidate_step(
     _prune_old_training(out_dir, keep_last_n=keep_last_n_training)
     _prune_old_ppgd(out_dir, keep_last_n=DEFAULT_KEEP_LAST_N_PPGD)
 
-    shutil.rmtree(step_dir, ignore_errors=True)
-    logger.info(f"consolidate: removed scratch {step_dir}")
+    _prune_scratch_through(scratch_dir, step)
+    logger.info(f"consolidate: removed scratch through step {step}")
 
 
 def _prune_old_training(out_dir: Path, *, keep_last_n: int) -> None:
@@ -245,6 +251,38 @@ def _prune_old_ppgd(out_dir: Path, *, keep_last_n: int) -> None:
         logger.info(f"consolidate: pruned old {ppgd_shard_dirname(step)}/")
 
 
+def _prune_scratch_through(scratch_dir: Path, step: int) -> None:
+    """Remove scratch for every snapshot at or before ``step`` (now consolidated).
+
+    Called by ``consolidate_step`` once ``training_<step>.pth`` exists. Resume loads
+    the newest checkpoint, never raw partials, so a step ``<=`` a consolidated one is
+    dead — and since the final save is the highest step, its consolidation sweeps a
+    finished run's scratch to nothing (even if earlier consolidations were skipped).
+    """
+    for d in scratch_dir.glob("step_*"):
+        if d.is_dir() and int(d.name.removeprefix("step_")) <= step:
+            shutil.rmtree(d, ignore_errors=True)
+
+
+def prune_old_scratch(scratch_dir: Path, *, keep_last_n: int = KEEP_LAST_N_SCRATCH) -> None:
+    """Delete all but the newest ``keep_last_n`` ``step_<S>/`` scratch dirs.
+
+    The train-loop backstop to ``_prune_scratch_through``: when the async
+    consolidation that normally clears scratch never runs, this still bounds it. Safe
+    on the loop for the same reason — resume loads only the consolidated
+    ``training_<S>.pth``, never raw partials, so dropping an older un-consolidated step
+    costs at most a manual-recovery convenience, never a resume target.
+    """
+    assert keep_last_n >= 1, "keep_last_n=0 would delete the just-written step"
+    steps = sorted(
+        (d for d in scratch_dir.glob("step_*") if d.is_dir()),
+        key=lambda d: int(d.name.removeprefix("step_")),
+    )
+    for d in steps[:-keep_last_n]:
+        shutil.rmtree(d, ignore_errors=True)
+        logger.info(f"scratch GC: dropped stale {d.name}/")
+
+
 def unconsolidated_steps(out_dir: Path) -> list[int]:
     """Steps that have scratch partials on disk but no `training_<step>.pth` yet.
 

param_decomp_lab/three_pool/optimize.py

@@ -84,6 +84,7 @@ class boundary.
     CONSOLIDATE_META_FILENAME,
     load_ppgd_shard,
     ppgd_shard_dirname,
+    prune_old_scratch,
     step_scratch_dir,
 )
 from param_decomp_lab.three_pool.context import (
@@ -484,6 +485,11 @@ def snapshot(self, scratch_dir: Path) -> None:
         dist.barrier(group=p2p_group)
         trace("snapshot: barrier (post-write rejoin) done")
 
+        # Backstop the async consolidation that normally clears scratch: if its job
+        # stalls or never runs, this keeps partials from growing without bound.
+        if self.ctx.role.rank == 0:
+            prune_old_scratch(scratch_dir)
+
     def _build_my_partial(self) -> dict[str, Any]:
         my_named_params = self._named_params_for_my_optimizer()
         my_optimizer_by_name: dict[str, dict[str, Any]] = (

param_decomp_lab/three_pool/two_pool_optimize.py

@@ -73,6 +73,7 @@
     CONSOLIDATE_META_FILENAME,
     load_ppgd_shard,
     ppgd_shard_dirname,
+    prune_old_scratch,
     step_scratch_dir,
 )
 from param_decomp_lab.three_pool.context import ChunkContext
@@ -392,6 +393,11 @@ def snapshot(self, scratch_dir: Path) -> None:
         dist.barrier(group=p2p_group)
         trace("snapshot: barrier (post-write rejoin) done")
 
+        # Backstop the async consolidation that normally clears scratch: if its job
+        # stalls or never runs, this keeps partials from growing without bound.
+        if self.ctx.role.rank == 0:
+            prune_old_scratch(scratch_dir)
+
     @classmethod
     def from_snapshot(
         cls,