From 14d1138c036f9a0f150d91a90606b6a861a5ee7a Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:20:42 -0700 Subject: [PATCH 01/54] Remove parameter sweep support The benchmark driver no longer expands list-valued config parameters into a cross product of runs: each `scaffold benchmark` invocation now runs exactly one worker in the resolved run directory, so the per-combination `param_set_i` subdirectories are gone and restart/checkpoint paths stay in one place. Config validation now always rejects list values for scalar keys, naming each offending key and its value ("problem_scale: parameter sweeps are no longer supported; got list [6, 7]") instead of the previous TypeErrors from `math.floor(list)` / `int - list` deep inside Config; load_config's config type "sweep" is renamed "benchmark". This supersedes findings R11, R12 and R16 of the round-2 review, which all stem from the half-finished sweep path. --- README.md | 6 +- ScaFFold/benchmark.py | 103 ++++++++------------------------- ScaFFold/cli.py | 11 ++-- ScaFFold/utils/config_utils.py | 37 +++++++----- tests/test_config.py | 91 +++++++++++++++++++++-------- 5 files changed, 119 insertions(+), 129 deletions(-) diff --git a/README.md b/README.md index f2827e8..3206b02 100644 --- a/README.md +++ b/README.md @@ -61,9 +61,9 @@ The model is trained from a random initialization until convergence, which is de ScaFFold benchmark training always uses PyTorch distributed execution with DistConv spatial parallelism. For a singleton run, launch one distributed rank rather than disabling distributed execution. -`benchmark` creates a folder for the benchmark run(s) at `base_run_dir` set in the config file. For reproducibility, we store a copy of the benchmark run config yml. Within each run subfolder, `benchmark` creates a yml config for that specific run. +Each `benchmark` invocation performs exactly one benchmark run, in a run folder created under `base_run_dir` set in the config file. Every run parameter must be single-valued; a list (e.g. `problem_scale: [6, 7]`) is rejected by name, since parameter sweeps are not supported. To compare parameter settings, launch one benchmark run per setting. For reproducibility, the run folder holds a copy of the benchmark config yml plus the fully merged `config.yaml` for that run. -After each run completes, statistics from the run are stored in `train_stats.csv`. Additionally, users can inspect plots of the training and validation losses over time in ` Date: Fri, 31 Jul 2026 15:21:01 -0700 Subject: [PATCH 02/54] Skip the final checkpoint save when a resume has nothing to train A resume whose checkpoint already covers config.epochs leaves the epoch loop before any epoch body runs, so val_loss_avg was never bound and the final-save block crashed with UnboundLocalError on every rank. train() now logs that there was nothing to resume and returns cleanly, leaving the existing checkpoint (which already covers those epochs) alone. Round-2 review: R01. --- ScaFFold/utils/trainer.py | 22 +++++++++- tests/test_resume.py | 84 ++++++++++++++++++++++++++++++++++++--- 2 files changed, 99 insertions(+), 7 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 108df71..29e7222 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -776,6 +776,11 @@ def train(self, profiler=None): # other ranks would make them disagree about whether to call # save_checkpoint on exit, deadlocking its internal collective. last_checkpoint_epoch = None + # Whether this invocation completed at least one NEW epoch. A resume + # whose checkpoint already covers config.epochs (or an --epochs lowered + # below the checkpointed epoch) leaves the loop at the max-epoch check + # before any epoch body runs, so none of the per-epoch metrics exist. + completed_new_epoch = False with open(self.outfile_path, "a", newline="") as outfile: start = time.time() while dice_score_train < self.config.target_dice: @@ -1030,6 +1035,7 @@ def train(self, profiler=None): dice_score_train = val_score epoch += 1 + completed_new_epoch = True # This check must exist otherwise the condition dice_score_train < self.config.target_dice will evaluate to False and incorrectly exit the training if math.isnan(dice_score_train): @@ -1039,12 +1045,26 @@ def train(self, profiler=None): completed_epochs = epoch - 1 + if not completed_new_epoch: + # The loop exited without running a single new epoch: the resumed + # checkpoint already covers every epoch this run was asked for. + # There is nothing new to save (the existing checkpoint already + # records epoch `completed_epochs`) and none of the per-epoch + # metrics the final save would write were ever computed, so skip + # it and return normally -- the caller's post-processing still has + # the CSV the original run left behind. + self.log.warning( + "Nothing to resume: the loaded checkpoint already covers epoch " + "%s, so no new epoch was trained and no checkpoint was written. " + "Increase 'epochs' (or lower 'target_dice') to train further.", + completed_epochs, + ) # Save a final checkpoint when the run exits (convergence or max epochs) # at an epoch that was not a checkpoint interval, so the converged # weights that produced the reported metrics are not lost. Skipped when # checkpointing is disabled, when no epoch completed, or when the last # completed epoch was already checkpointed inside the loop. - if ( + elif ( self.config.checkpoint_interval > 0 and completed_epochs >= 1 and last_checkpoint_epoch != completed_epochs diff --git a/tests/test_resume.py b/tests/test_resume.py index 209a6eb..93acdaf 100644 --- a/tests/test_resume.py +++ b/tests/test_resume.py @@ -28,6 +28,7 @@ from __future__ import annotations +import logging from pathlib import Path import numpy as np @@ -142,19 +143,27 @@ def cfg(): ) -def _stub_trainer(run_dir, *, train_from_scratch, log): +def _stub_trainer(run_dir, *, train_from_scratch, log, **config_overrides): """A PyTorchTrainer carrying only what cleanup_or_resume/train touch. Built via ``object.__new__`` (as the checkpointing tests do) so no dataset, model, or process group is needed. A real CheckpointManager over a tiny linear model backs the resume path so load/save round-trip faithfully. + + ``config_overrides`` set additional config fields (``epochs``, + ``target_dice``, ``checkpoint_interval``, ...) for tests that drive + ``train()``'s loop-entry and final-save control flow. """ t = object.__new__(PyTorchTrainer) - t.config = SimpleNamespace( - train_from_scratch=train_from_scratch, - run_dir=str(run_dir), - checkpoint_interval=-1, - ) + config_fields = { + "train_from_scratch": train_from_scratch, + "run_dir": str(run_dir), + "checkpoint_interval": -1, + "epochs": -1, + "target_dice": 0.95, + } + config_fields.update(config_overrides) + t.config = SimpleNamespace(**config_fields) t.world_rank = 0 t.global_step = 0 t.total_optimizer_steps = 0 @@ -346,6 +355,69 @@ def test_step_counters_roundtrip(tmp_path): assert t2.total_optimizer_steps == 37 +def test_restart_of_completed_run_trains_and_saves_nothing(tmp_path, caplog): + """Restarting a run whose checkpoint already covers ``epochs`` exits cleanly. + + The epoch loop breaks at the max-epoch check before any new epoch runs, so + no epoch metric exists to checkpoint: the final-save block must be skipped + (the last checkpoint already covers the completed epochs) rather than + saving with an unbound ``val_loss_avg``. ``train()`` has to return normally + so the worker's post-processing still runs off the CSV already on disk, and + the run must say plainly that there was nothing to resume. + """ + log = logging.getLogger("resume.r01") + run = tmp_path / "run" + run.mkdir() + + # Artifacts of a completed epochs=2 run: a checkpoint recording epoch 2 + # (written by the in-loop save) plus its CSV rows. + t1 = _stub_trainer( + run, + train_from_scratch=False, + log=log, + epochs=2, + checkpoint_interval=1, + ) + t1.checkpoint_manager.save_checkpoint( + epoch=2, + val_loss_avg=0.5, + extras={ + "train_mask_values": None, + "global_step": 8, + "total_optimizer_steps": 8, + }, + ) + csv = run / "train_stats.csv" + csv.write_text(_HEADER + "\n") + _write_rows(csv, [1, 2]) + + ckpt = t1.checkpoint_manager.last_ckpt_path + before = ckpt.read_bytes() + + # The user reruns the generated restart command against the same dir. + t2 = _stub_trainer( + run, + train_from_scratch=False, + log=log, + epochs=2, + checkpoint_interval=1, + ) + t2.config.restart = True + t2.cleanup_or_resume() + assert t2.start_epoch == 3 # past the last epoch: nothing left to train + + with caplog.at_level(logging.WARNING): + t2.train() # must not raise + + # Nothing new was written: the existing checkpoint is byte-identical (a + # fresh save would serialize this trainer's own randomly-initialized model) + # and the CSV still holds exactly the original run's epochs. + assert ckpt.read_bytes() == before + epochs = [ln.split(",")[0] for ln in csv.read_text().splitlines()[1:]] + assert epochs == ["1", "2"] + assert "nothing to resume" in caplog.text.lower() + + def test_total_steps_resume_predates_dedicated_key(tmp_path): """A checkpoint recording only global_step still resumes the step total. From 8ba619109c047f6574a09a79036d634c5eea14d7 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:23:08 -0700 Subject: [PATCH 03/54] Persist the achieved validation dice so a converged run resumes converged The training loop runs while the validation dice is below target_dice but that score was never checkpointed, so every restart of a converged epochs:-1 run re-trained one full epoch before rediscovering it had converged, inflating the epoch count and the FOM's total train time. The score now rides along in the checkpoint extras and seeds the loop variable on resume. Round-2 review: R06. --- ScaFFold/utils/trainer.py | 24 ++++++++++++++- tests/test_resume.py | 61 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 1 deletion(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 29e7222..cfd27ee 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -117,6 +117,10 @@ def __init__(self, model, config, device, log): self.global_step = 0 self.total_optimizer_steps = 0 self.start_epoch = -1 + # Validation dice already achieved by the epoch we resume from (0 for a + # fresh run). The training loop's exit condition is a threshold on this + # score, so it has to survive a restart: see cleanup_or_resume. + self.start_val_dice = 0.0 self.ps = getattr(self.config, "_parallel_strategy", None) self.spatial_mesh = None # Spatial mesh for use w/ DistConv self.data_num_replicas = self.world_size @@ -387,6 +391,7 @@ def cleanup_or_resume(self): pass self.start_epoch = 1 + self.start_val_dice = 0.0 else: # Load checkpoint via manager. An explicit restart must find a # checkpoint; a plain non-scratch launch may simply start fresh. @@ -399,6 +404,16 @@ def cleanup_or_resume(self): if "train_mask_values" in restored: self.train_set.mask_values = restored["train_mask_values"] + # Resume the convergence state, not just the weights. The loop runs + # while the validation dice is below target_dice, so a run that had + # already converged must re-enter the loop with the score it + # achieved -- starting from 0 would unconditionally re-train (and + # re-log, and re-checkpoint) one full epoch before the condition is + # re-tested, inflating the epoch count and the FOM's total train + # time. Checkpoints written before this key existed simply fall + # back to 0, i.e. the old behaviour. + self.start_val_dice = restored.get("val_dice", 0.0) + # Continue the optimizer-step counts from where the checkpoint # left off; otherwise a resumed run restarts them at 0 and # undercounts all pre-resume work in the reported step totals. @@ -766,7 +781,10 @@ def train(self, profiler=None): """ epoch = self.start_epoch - dice_score_train = 0 + # Seeded from the resumed checkpoint (0 for a fresh run) so an already + # converged run exits the loop immediately instead of training an + # extra epoch to rediscover that it converged. + dice_score_train = self.start_val_dice epoch_minibatch_times_s = [] # Track the last epoch checkpointed inside the loop so the final-save # decision below is identical on every rank. The in-loop checkpoint @@ -1027,6 +1045,9 @@ def train(self, profiler=None): "train_mask_values": self.train_set.mask_values, "global_step": self.global_step, "total_optimizer_steps": self.total_optimizer_steps, + # The convergence state: what a resume must restore to + # know this epoch already met (or missed) target_dice. + "val_dice": val_score, } self.checkpoint_manager.save_checkpoint(epoch, val_loss_avg, extras) last_checkpoint_epoch = epoch @@ -1073,6 +1094,7 @@ def train(self, profiler=None): "train_mask_values": self.train_set.mask_values, "global_step": self.global_step, "total_optimizer_steps": self.total_optimizer_steps, + "val_dice": val_score, } self.checkpoint_manager.save_checkpoint( completed_epochs, val_loss_avg, extras diff --git a/tests/test_resume.py b/tests/test_resume.py index 93acdaf..3e8aa4d 100644 --- a/tests/test_resume.py +++ b/tests/test_resume.py @@ -133,6 +133,7 @@ def cfg(): from types import SimpleNamespace # noqa: E402 +import ScaFFold.utils.trainer as trainer_mod # noqa: E402 from ScaFFold.utils.checkpointing import CheckpointManager # noqa: E402 from ScaFFold.utils.trainer import PyTorchTrainer # noqa: E402 @@ -418,6 +419,66 @@ def test_restart_of_completed_run_trains_and_saves_nothing(tmp_path, caplog): assert "nothing to resume" in caplog.text.lower() +def test_converged_resume_does_not_retrain(tiny_trainer, monkeypatch): + """A restart of an already-converged run must not train another epoch. + + The validation dice the checkpointed epoch achieved is persisted with the + checkpoint and restored on resume, so the ``while dice_score_train < + target_dice`` loop is never entered. Without it a converged ``epochs: -1`` + run re-trains, re-logs and re-checkpoints one full epoch on every restart, + inflating ``sum(epoch_duration)`` -- the FOM denominator -- and the + reported epoch count relative to the same run left un-restarted. + """ + overrides = { + "checkpoint_interval": 1, + "epochs": -1, + "target_dice": 0.95, + "train_from_scratch": 0, + } + + def converged_evaluate(*args, **kwargs): + # Two validation samples at hard dice 0.96, i.e. above target. + return (0.96 * 2, 0.1 * 2, 0.1, 2, 2) + + monkeypatch.setattr(trainer_mod, "evaluate", converged_evaluate) + + def stub_batches(trainer): + """Replace the DistConv forward path and count the batches it runs.""" + calls = {"n": 0} + + def _step(batch, **kwargs): + calls["n"] += 1 + return 1, torch.tensor(0.1), torch.tensor(0.96) + + monkeypatch.setattr(trainer, "_run_training_batch", _step) + return calls + + # The original run: converges in epoch 1 and checkpoints it. + first = tiny_trainer(config_overrides=overrides) + first_calls = stub_batches(first) + first.cleanup_or_resume() + first.train() + assert first_calls["n"] > 0 # it really did train + + ckpt_path = first.checkpoint_manager.last_ckpt_path + saved = torch.load(ckpt_path, map_location="cpu", weights_only=False) + assert saved["epoch"] == 1 + csv = Path(first.outfile_path) + assert [ln.split(",")[0] for ln in csv.read_text().splitlines()[1:]] == ["1"] + + # The user reruns the restart command: nothing is left to train. + second = tiny_trainer(config_overrides=overrides) + second_calls = stub_batches(second) + second.cleanup_or_resume() + assert second.start_epoch == 2 + second.train() + + assert second_calls["n"] == 0, "a converged run re-trained an extra epoch" + assert [ln.split(",")[0] for ln in csv.read_text().splitlines()[1:]] == ["1"] + reloaded = torch.load(ckpt_path, map_location="cpu", weights_only=False) + assert reloaded["epoch"] == 1 + + def test_total_steps_resume_predates_dedicated_key(tmp_path): """A checkpoint recording only global_step still resumes the step total. From 7bb34e9be5be44a1908aad8e03f49a0f03bc5599 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:24:27 -0700 Subject: [PATCH 04/54] Reset the cached best-loss state when cleaning up for a fresh run cleanup(train_from_scratch=True) deleted the checkpoint files but kept best_val_loss and last_saved_epoch, so the fresh run's is_best decisions stayed gated by the deleted run's best and it wrote no best checkpoint until it beat a score nothing backed any more. Round-2 review: R02. --- ScaFFold/utils/checkpointing.py | 9 ++++++++ tests/test_checkpointing.py | 39 +++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 823f154..932f6d1 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -107,6 +107,15 @@ def cleanup(self, train_from_scratch: bool) -> None: self._barrier() return + # Drop the cached state that described the run being deleted. Both + # fields are seeded from disk (or a previous save), so keeping them + # would let a deleted run's best gate this run's is_best decisions -- + # the fresh run would then never write a best checkpoint until it beat + # a score no file backs any more, leaving it with no best-checkpoint + # fallback. + self.best_val_loss = math.inf + self.last_saved_epoch = None + if self.world_rank == 0: for p in (self.last_ckpt_path, self.best_ckpt_path): if p.exists(): diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 4e1737d..d2d7183 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -26,6 +26,7 @@ from __future__ import annotations +import math import re import time from pathlib import Path @@ -244,6 +245,44 @@ def spy_load(path, *args, **kwargs): assert final_best["val_loss_avg"] == pytest.approx(0.2) +# --------------------------------------------------------------------------- +# R02 -- a from-scratch cleanup drops the deleted run's best, not just its files +# --------------------------------------------------------------------------- + + +def test_cleanup_from_scratch_resets_best(tmp_path): + """``cleanup(train_from_scratch=True)`` resets the cached best-loss state. + + The manager seeds ``best_val_loss`` from ``checkpoint_best.pth`` at + construction so a resumed run does not call its first epoch "best". When + the same directory is then wiped for a fresh run, that cached score + outlives the file it came from: every ``is_best`` decision of the new run + is gated by a deleted run's score, so no ``checkpoint_best.pth`` is written + until the retrain beats it -- leaving the run with no best-checkpoint + fallback at all. + """ + mgr, _ = _make_manager(tmp_path) + mgr.save_checkpoint(epoch=1, val_loss_avg=0.01) + assert mgr.best_ckpt_path.exists() + + # A driver reusing the run directory: the new manager seeds from disk. + mgr2, _ = _make_manager(tmp_path) + assert mgr2.best_val_loss == pytest.approx(0.01) + mgr2.save_checkpoint(epoch=2, val_loss_avg=0.9) + assert mgr2.last_saved_epoch == 2 + + mgr2.cleanup(train_from_scratch=True) + + assert not mgr2.last_ckpt_path.exists() + assert not mgr2.best_ckpt_path.exists() + assert mgr2.best_val_loss == math.inf + assert mgr2.last_saved_epoch is None + + # The fresh run's first epoch is its best, and a best checkpoint exists. + assert mgr2.save_checkpoint(epoch=1, val_loss_avg=0.5) is True + assert mgr2.best_ckpt_path.exists() + + # --------------------------------------------------------------------------- # F71 -- CPU tensors are cloned into the snapshot # --------------------------------------------------------------------------- From f9e5a8f1bc1cb065f550480a8a7e266377296ab4 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:25:57 -0700 Subject: [PATCH 05/54] Abort on non-finite epoch losses instead of checkpointing a diverged model The existing NaN check tests the hard-argmax dice, which stays finite even for an all-NaN model, so it could never fire; a diverged run kept looping below target while NaN weights overwrote checkpoint_last.pth. The reduced train and validation losses are now checked right after the data-parallel reductions -- identical on every rank, so all ranks raise together. Round-2 review: R03. --- ScaFFold/utils/trainer.py | 19 +++++++++++++++ tests/test_checkpointing.py | 48 +++++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index cfd27ee..07fdd2c 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -970,6 +970,25 @@ def train(self, profiler=None): # Reduced sample-weighted total and per-sample mean val loss. val_loss_epoch = val_info[1].item() val_loss_avg = val_loss_epoch / global_val_samples + + # Divergence check. The dice score below is computed from a + # hard argmax, so it stays finite even for an all-NaN model + # (argmax of NaN logits is 0 and the one-hots are finite) -- + # the loss is the only value that actually goes non-finite. Bail + # out before the CSV row and the checkpoint: continuing would + # keep overwriting checkpoint_last.pth with NaN weights (which + # then poison the next restart) while the loop's dice threshold + # can never be met. Both values come out of the data-parallel + # reductions above, so they are identical on every rank and + # every rank raises here together, leaving no unmatched + # collective behind. + if not (math.isfinite(overall_loss) and math.isfinite(val_loss_avg)): + raise ValueError( + f"Non-finite loss at epoch {epoch} " + f"(train_loss={overall_loss}, val_loss={val_loss_avg}): " + "training diverged, aborting before checkpointing." + ) + if not self.config.disable_scheduler: self.scheduler.step() else: diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index d2d7183..40151f6 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -346,6 +346,54 @@ def fake_evaluate(*args, **kwargs): assert saved["epoch"] == 2 +# --------------------------------------------------------------------------- +# R03 -- a diverged epoch aborts instead of checkpointing NaN weights +# --------------------------------------------------------------------------- + + +def test_divergence_aborts_before_poisoning_checkpoint(tiny_trainer, monkeypatch): + """Non-finite epoch losses abort the run before any checkpoint is written. + + The dice score is computed from a hard argmax, so an all-NaN model still + produces a *finite* dice (argmax of NaN logits is 0 and the one-hots are + finite): the dice check can never fire on divergence. Left unguarded, a + diverged ``epochs: -1`` run keeps looping on a finite plateau below target + while every checkpoint interval overwrites ``checkpoint_last.pth`` with NaN + weights, poisoning the next ``--restart``. The reduced losses are the + values that actually go non-finite, and being reductions they are identical + on every rank, so the check fires on all ranks together. + """ + trainer = tiny_trainer( + config_overrides={ + "checkpoint_interval": 1, + "epochs": 3, + "target_dice": 0.95, + } + ) + + # A diverged step: NaN loss, and a dice that stays finite. + monkeypatch.setattr( + trainer, + "_run_training_batch", + lambda batch, **kw: (1, torch.tensor(float("nan")), torch.tensor(0.0)), + ) + + def diverged_evaluate(*args, **kwargs): + # Exactly what evaluate() returns for an all-NaN model: a tiny but + # finite hard-argmax dice sum alongside a NaN validation loss. + return (7.4e-10, float("nan"), float("nan"), 2, 2) + + monkeypatch.setattr(trainer_mod, "evaluate", diverged_evaluate) + + trainer.cleanup_or_resume() + with pytest.raises(ValueError, match="[Nn]on-finite"): + trainer.train() + + # The run died before the diverged epoch could be checkpointed. + assert not trainer.checkpoint_manager.last_ckpt_path.exists() + assert not trainer.checkpoint_manager.best_ckpt_path.exists() + + # --------------------------------------------------------------------------- # F49 -- GradScaler-skipped steps do not advance the optimizer-step counter # --------------------------------------------------------------------------- From c2083a06427dd7cfbc16cf9b83898da06190f00a Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:31:42 -0700 Subject: [PATCH 06/54] Broadcast the checkpoint write outcome so all ranks fail together wait_for_save now re-raises the writer's exception instead of logging and dropping it, train() consumes the run's final save so a failed async write can no longer exit 0 with no checkpoint, and save_checkpoint broadcasts rank 0's outcome (success or error sentinel, including a deferred async failure) before anyone raises -- previously rank 0 raised ahead of its own broadcast and left the peers in an unmatched collective reporting a transport error instead of the disk error. Round-2 review: R04, R05. --- ScaFFold/utils/checkpointing.py | 230 ++++++++++++++++++++++++-------- ScaFFold/utils/trainer.py | 8 ++ tests/test_checkpointing.py | 181 +++++++++++++++++++++++++ 3 files changed, 364 insertions(+), 55 deletions(-) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 932f6d1..5fedc8c 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -25,6 +25,16 @@ import torch.distributed as dist +class CheckpointSaveError(RuntimeError): + """A checkpoint write failed. + + Raised identically on every rank. Only rank 0 writes, but its outcome is + broadcast, so the peers report the real disk error instead of the + unmatched-collective symptom (an opaque gloo transport error, or an NCCL + watchdog timeout minutes later) that a rank-0-only raise produces. + """ + + class CheckpointManager: """ Checkpoint Manager for DDP/Single-Process. @@ -90,6 +100,10 @@ def __init__( # Async handling self.executor = None self.future = None + # The exception behind the most recently reported save failure, kept + # only so rank 0 can chain it (and its traceback) onto the + # CheckpointSaveError every rank raises. + self._save_error_exc: Optional[BaseException] = None if self.async_save and self.world_rank == 0: # We only need 1 worker for serializing writes self.executor = ThreadPoolExecutor(max_workers=1) @@ -99,44 +113,106 @@ def __init__( self.base_dir.mkdir(parents=True, exist_ok=True) def cleanup(self, train_from_scratch: bool) -> None: - """Clear existing checkpoints if training from scratch.""" - # Ensure any pending async saves are finished before deleting - self.wait_for_save() + """Clear existing checkpoints if training from scratch. + + Rank-symmetric, like every other collective point here: any pending + async write is drained and its outcome broadcast, so a failure raises + on all ranks together (see ``save_checkpoint``). + """ + # Ensure any pending async save is finished before deleting. + error = self._drain_pending_save() + + if train_from_scratch: + # Drop the cached state that described the run being deleted. Both + # fields are seeded from disk (or a previous save), so keeping them + # would let a deleted run's best gate this run's is_best decisions + # -- the fresh run would then never write a best checkpoint until + # it beat a score no file backs any more, leaving it with no + # best-checkpoint fallback. + self.best_val_loss = math.inf + self.last_saved_epoch = None + + if self.world_rank == 0: + self._remove_checkpoint_files() + + error = self._broadcast_obj(error) + self._barrier() + if error is not None: + self._raise_save_error(error) + + def _remove_checkpoint_files(self) -> None: + """Delete this run's checkpoint files (rank 0 only).""" + for p in (self.last_ckpt_path, self.best_ckpt_path): + if p.exists(): + try: + p.unlink() + self._log(f"Removed existing checkpoint: {p}") + except Exception as e: + self._log(f"Failed to remove {p}: {e}") - if not train_from_scratch: - self._barrier() + def wait_for_save(self): + """Block until the background save (if any) is complete. + + The writer's exception is re-raised rather than logged and dropped: a + save that could not complete has to surface, otherwise the manager + state (and the process exit code) claims a checkpoint that is not on + disk. The future is consumed exactly once, so the failure is reported + at exactly one point. + + Callers that are inside a collective region must not let this + propagate directly -- use ``_drain_pending_save`` instead, per the + collective invariant documented on ``save_checkpoint``. + """ + if self.future is None: return + # Clear the handle *before* blocking on it so a failed write is + # reported once and does not re-raise at some later, arbitrary point. + future, self.future = self.future, None + if not future.done(): + self._log("Waiting for background checkpoint save to complete...") + future.result() # Blocks and re-raises whatever the writer raised + + def _drain_pending_save(self) -> Optional[str]: + """Consume the in-flight async save, reporting failure without raising. + + Returns a description of the writer's failure, or ``None``. Only rank 0 + ever has a pending write, so raising here directly would leave the + peers blocked in the next collective; callers broadcast this result and + raise on every rank together. + """ + try: + self.wait_for_save() + except Exception as e: + self._save_error_exc = e + return f"{type(e).__name__}: {e}" + return None - # Drop the cached state that described the run being deleted. Both - # fields are seeded from disk (or a previous save), so keeping them - # would let a deleted run's best gate this run's is_best decisions -- - # the fresh run would then never write a best checkpoint until it beat - # a score no file backs any more, leaving it with no best-checkpoint - # fallback. - self.best_val_loss = math.inf - self.last_saved_epoch = None + def _raise_save_error(self, description: str) -> None: + """Raise a broadcast save failure on this rank. - if self.world_rank == 0: - for p in (self.last_ckpt_path, self.best_ckpt_path): - if p.exists(): - try: - p.unlink() - self._log(f"Removed existing checkpoint: {p}") - except Exception as e: - self._log(f"Failed to remove {p}: {e}") + Rank 0 chains the original exception so its traceback survives; the + peers never saw it and raise the same message on its own. + """ + cause, self._save_error_exc = self._save_error_exc, None + raise CheckpointSaveError( + f"Checkpoint save failed on rank 0: {description}" + ) from cause + + def finalize_saves(self) -> None: + """Consume the outcome of the run's last save before the run ends. + + Nothing touches the manager after the training loop, so an + asynchronous write that failed there would never be observed: the + process would exit successfully having written no checkpoint (or left + a stale one), the benchmark would report success, and a later + ``--restart`` would resume from the wrong epoch or fail its pre-check. + Rank-symmetric, like ``save_checkpoint``. + """ + error = self._drain_pending_save() + error = self._broadcast_obj(error) self._barrier() - - def wait_for_save(self): - """Blocks until the background save (if any) is complete.""" - if self.future is not None: - # check if running - if not self.future.done(): - self._log("Waiting for background checkpoint save to complete...") - try: - self.future.result() # Blocks and raises exceptions if any occurred - except Exception as e: - self._log(f"Background save failed with error: {e}") - self.future = None + if error is not None: + self._raise_save_error(error) def snapshot_training_state(self) -> Dict[str, Any]: """Capture mutable in-memory training state without writing a checkpoint.""" @@ -184,7 +260,10 @@ def load_from_checkpoint(self, require_checkpoint: bool = False) -> int: With ``require_checkpoint`` (an explicit ``--restart``), a missing checkpoint raises instead of silently starting over. """ - self.wait_for_save() # Safety: don't load while writing + # Safety: don't load while writing. A failure from that write is + # folded into the decision broadcast below rather than raised here, so + # rank 0 never abandons its peers inside the broadcast. + error = self._drain_pending_save() # 1. Rank 0 is the sole reader: it selects the newest readable # checkpoint and deserializes it once, then broadcasts the loaded @@ -193,10 +272,16 @@ def load_from_checkpoint(self, require_checkpoint: bool = False) -> int: # read of one (multi-GB) file from the shared filesystem on restart -- # a restart I/O storm that serializes on the parallel FS. Peer ranks # therefore never open the checkpoint files at all. - result = self._select_and_load() if self.world_rank == 0 else None + result = None + if self.world_rank == 0: + result = ( + ("save_failed", error) if error is not None else self._select_and_load() + ) status, payload = self._broadcast_obj(result) # 2. Every rank acts on the same decision rank 0 reached. + if status == "save_failed": + self._raise_save_error(payload) if status == "empty": if require_checkpoint: # An explicit restart must resume real state; silently @@ -310,22 +395,55 @@ def save_checkpoint( """ Save checkpoint. If async_save is True, this returns immediately after CPU transfer. + + Collective invariant: every rank posts exactly one ``_broadcast_obj`` + followed by exactly one ``_barrier`` on every path through this method, + failures included. Rank 0 is the sole writer, but it never raises + before those collectives: what gets broadcast is the write's *outcome* + -- ``is_best`` on success, or an error description on failure, + including a *previous* asynchronous write whose failure is surfaced + here, at the next collective point. Every rank then raises the same + ``CheckpointSaveError`` together. Raising on rank 0 before the + broadcast would strand the peers in an unmatched collective, where a + plain disk error resurfaces as a gloo transport error or an NCCL + watchdog timeout that hides the real cause. """ - is_best = False + # Non-zero ranks contribute nothing; their placeholder is overwritten + # by rank 0's outcome in the broadcast below (and is a harmless no-op + # in the degenerate non-distributed case). + outcome = ("ok", False) if self.world_rank == 0: + outcome = self._rank0_save(epoch, val_loss_avg, extras) + + status, payload = self._broadcast_obj(outcome) + + # Barrier: ensure Rank 0 has finished the "Snapshot" phase before anyone continues. + # Even in async mode, we must wait for the CPU transfer to finish. + self._barrier() + + if status == "error": + self._raise_save_error(payload) + return payload + + def _rank0_save(self, epoch, val_loss_avg, extras): + """Perform rank 0's write and REPORT its outcome; never raises. + + Returns ``("ok", is_best)`` or ``("error", description)``. The caller + broadcasts that outcome so every rank fails together -- see the + collective invariant on ``save_checkpoint``. + """ + try: + # 1. Wait for previous async save to prevent OOM or race. If that + # write failed, this is where it surfaces. + if self.async_save: + self.wait_for_save() + # Decide is_best from the cached best loss (single source of truth), # not by re-reading checkpoint_best.pth from disk. The cache is # seeded once at construction and updated below, so the decision # never races the background writer that may still be replacing the # best checkpoint in async mode. - if val_loss_avg < self.best_val_loss: - is_best = True - self.best_val_loss = val_loss_avg - - if self.world_rank == 0: - # 1. Wait for previous async save to prevent OOM or race - if self.async_save: - self.wait_for_save() + is_best = val_loss_avg < self.best_val_loss model_to_save = ( self.model.module if hasattr(self.model, "module") else self.model @@ -350,10 +468,6 @@ def save_checkpoint( if extras: state_dict.update(extras) - # Record the epoch being written so callers can tell whether the - # last completed epoch has already been checkpointed. - self.last_saved_epoch = epoch - # 2. Save Trigger if self.async_save: # We must clone tensors to CPU now, because training will resume @@ -380,13 +494,19 @@ def save_checkpoint( self.log, ) - # Broadcast result (for logging elsewhere) - is_best = self._broadcast_obj(is_best) - - # Barrier: ensure Rank 0 has finished the "Snapshot" phase before anyone continues. - # Even in async mode, we must wait for the CPU transfer to finish. - self._barrier() - return is_best + # Only now claim the save: the bytes are on disk (sync) or handed + # to the writer (async). Recording the epoch lets callers tell + # whether the last completed epoch has already been checkpointed. + # An async write that fails later aborts the run at the next + # collective point, so this optimistic state is never observed by + # a run that keeps going. + if is_best: + self.best_val_loss = val_loss_avg + self.last_saved_epoch = epoch + return ("ok", is_best) + except Exception as e: + self._save_error_exc = e + return ("error", f"{type(e).__name__}: {e}") @staticmethod def _atomic_save(state_dict, path): diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 07fdd2c..da33c12 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -1119,6 +1119,14 @@ def train(self, profiler=None): completed_epochs, val_loss_avg, extras ) + # Nothing downstream of the training loop touches the checkpoint + # manager, so this is the last chance to observe the outcome of the + # run's final (possibly asynchronous) write. Without it a failed final + # save would let the process exit successfully with no checkpoint at + # all, and the next --restart would resume from a stale epoch or fail + # its pre-check. + self.checkpoint_manager.finalize_saves() + if epoch_minibatch_times_s: minibatch_time_s = statistics.median(epoch_minibatch_times_s) adiak_value("minibatch_time_s", minibatch_time_s) diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 40151f6..f102a6d 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -28,7 +28,9 @@ import math import re +import textwrap import time +from concurrent.futures import ThreadPoolExecutor from pathlib import Path import pytest @@ -187,6 +189,185 @@ def always_raise(obj, f, *args, **kwargs): mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) +# --------------------------------------------------------------------------- +# R04 -- async save failures are reported, and the run's LAST save is consumed +# --------------------------------------------------------------------------- + + +def _failing_torch_save(obj, f, *args, **kwargs): + raise RuntimeError("writer boom") + + +def test_async_save_failure_surfaces_at_next_save(tmp_path, monkeypatch): + """A background write that failed is reported at the next save, not dropped. + + In async mode ``save_checkpoint`` returns as soon as the CPU snapshot is + handed to the writer thread, so the failure can only be observed later. + Swallowing it leaves ``last_saved_epoch``/``best_val_loss`` claiming a + checkpoint that does not exist on disk. + """ + mgr, _ = _make_manager(tmp_path, async_save=True) + monkeypatch.setattr(torch, "save", _failing_torch_save) + + mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) # write fails in background + + with pytest.raises(RuntimeError, match="writer boom"): + mgr.save_checkpoint(epoch=2, val_loss_avg=0.4) + + assert not mgr.last_ckpt_path.exists() + + +def test_async_save_failure_surfaces_at_finalize(tmp_path, monkeypatch): + """The run's final save has its outcome consumed before the run ends.""" + mgr, _ = _make_manager(tmp_path, async_save=True) + monkeypatch.setattr(torch, "save", _failing_torch_save) + + mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) + + with pytest.raises(RuntimeError, match="writer boom"): + mgr.finalize_saves() + + assert not mgr.last_ckpt_path.exists() + + +def test_final_async_save_failure_fails_the_run(tiny_trainer, monkeypatch): + """``train()`` must not return successfully after a failed final save. + + Nothing downstream of the training loop touches the checkpoint manager, so + without an explicit wait the child process exits 0 with no checkpoint at + all: the benchmark reports success and a later ``--restart`` resumes from a + stale epoch or fails its pre-check. + """ + trainer = tiny_trainer( + config_overrides={ + "checkpoint_interval": 1, + "epochs": 1, + "target_dice": 0.95, + } + ) + # The config has no async knob at this scale; drive the manager directly. + mgr = trainer.checkpoint_manager + mgr.async_save = True + mgr.executor = ThreadPoolExecutor(max_workers=1) + + monkeypatch.setattr( + trainer, + "_run_training_batch", + lambda batch, **kw: (1, torch.tensor(0.3), torch.tensor(0.5)), + ) + monkeypatch.setattr( + trainer_mod, "evaluate", lambda *a, **k: (0.5 * 2, 0.4 * 2, 0.4, 2, 2) + ) + monkeypatch.setattr(torch, "save", _failing_torch_save) + + trainer.cleanup_or_resume() + try: + with pytest.raises(RuntimeError, match="writer boom"): + trainer.train() + finally: + mgr.executor.shutdown(wait=True) + + assert not mgr.last_ckpt_path.exists() + + +# --------------------------------------------------------------------------- +# R05 -- a rank-0 write failure fails every rank with the same error +# --------------------------------------------------------------------------- + +# Two-rank script: rank 0's torch.save fails. Both ranks must come out of +# save_checkpoint with the SAME real error rather than rank 0 raising the disk +# error while its peers die (gloo) or stall (NCCL) in an unmatched collective. +# Kept inline rather than in tests/helpers/rank_scripts/ because it is only +# meaningful together with the assertions below. +SAVE_FAIL_RANK_SCRIPT = textwrap.dedent( + '''\ + """Two-rank save-failure rank script (gloo, CPU).""" + + import os + import sys + + import torch + import torch.distributed as dist + + from ScaFFold.utils.checkpointing import CheckpointManager + + rank = int(os.environ["RANK"]) + dist.init_process_group(backend="gloo") + + torch.manual_seed(0) + mgr = CheckpointManager( + model=torch.nn.Linear(64, 64), + base_dir=os.environ["CKPT_DIR"], + world_rank=rank, + dist_enabled=True, + ) + + if rank == 0: + def _boom(obj, f, *args, **kwargs): + raise RuntimeError("injected disk failure on rank 0") + + torch.save = _boom + + # Markers are delimited (trailing '.', '<<...>>') because two ranks writing + # the same pipe can interleave without a newline between them. + try: + mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) + except BaseException as exc: # noqa: BLE001 - the point is what we caught + print(f"RANK {rank} RAISED {type(exc).__name__}.", flush=True) + message = str(exc).replace(chr(10), " ") + print(f"RANK {rank} MESSAGE <<{message}>>", flush=True) + else: + print(f"RANK {rank} NO_RAISE.", flush=True) + + print(f"RANK {rank} DONE.", flush=True) + sys.stdout.flush() + try: + dist.destroy_process_group() + except Exception: + pass + ''' +) + + +@_requires_gloo +def test_rank0_save_failure_fails_all_ranks(tmp_path): + """Under 2 gloo ranks, a rank-0 disk error reaches the peer as itself. + + Rank 0 is the only writer, so raising its error before the broadcast the + peers are already waiting in leaves them in an unmatched collective: gloo + reports an opaque "Connection closed by peer" and NCCL stalls until the + watchdog timeout, in both cases hiding the disk error that actually + happened. The write's OUTCOME must be broadcast instead, so both ranks + raise the same error together. + """ + ckpt_dir = tmp_path / "checkpoints" + ckpt_dir.mkdir(parents=True, exist_ok=True) + script = tmp_path / "save_fail_2rank.py" + script.write_text(SAVE_FAIL_RANK_SCRIPT) + + rc, out, err = mpi_runner.torchrun_gloo( + str(script), n=2, timeout=90, env={"CKPT_DIR": str(ckpt_dir)} + ) + + done = set(re.findall(r"RANK (\d+) DONE\.", out)) + assert rc == 0 and {"0", "1"} <= done, ( + f"expected both ranks to fail cleanly and finish, rc={rc}\n" + f"stdout:\n{out}\nstderr:\n{err[-3000:]}" + ) + + raised = dict(re.findall(r"RANK (\d+) RAISED (\w+)\.", out)) + messages = dict(re.findall(r"RANK (\d+) MESSAGE <<(.*?)>>", out)) + assert set(raised) == {"0", "1"}, f"both ranks must raise\nstdout:\n{out}" + assert raised["0"] == raised["1"], ( + f"ranks raised different error types: {raised}\nstdout:\n{out}" + ) + for rank in ("0", "1"): + assert "injected disk failure on rank 0" in messages.get(rank, ""), ( + f"rank {rank} did not see the real disk error: " + f"{messages.get(rank)!r}\nstdout:\n{out}" + ) + + # --------------------------------------------------------------------------- # F41 -- race-free best decision (cached best loss, no per-save probe) # --------------------------------------------------------------------------- From 6af26c21ed96ad8f03b5e2fb408266c65dc52526 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:34:51 -0700 Subject: [PATCH 07/54] Clean up checkpoint temp and quarantine debris A killed write strands a full-checkpoint-sized checkpoint_*.pth.tmp. and a quarantined checkpoint keeps its .corrupt copy forever; nothing ever removed either. The from-scratch cleanup now deletes both, and manager construction sweeps orphaned temp files (skipping this pid's) since the kill/restart cycle that creates them always takes the resume path. Round-2 review: R07. --- ScaFFold/utils/checkpointing.py | 47 ++++++++++++++++++++++++++-- tests/test_checkpointing.py | 54 +++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 2 deletions(-) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 5fedc8c..24f3fe6 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -111,6 +111,7 @@ def __init__( # Ensure base directory exists (Rank 0 only) if self.world_rank == 0: self.base_dir.mkdir(parents=True, exist_ok=True) + self._sweep_orphaned_tmp_files() def cleanup(self, train_from_scratch: bool) -> None: """Clear existing checkpoints if training from scratch. @@ -141,8 +142,20 @@ def cleanup(self, train_from_scratch: bool) -> None: self._raise_save_error(error) def _remove_checkpoint_files(self) -> None: - """Delete this run's checkpoint files (rank 0 only).""" - for p in (self.last_ckpt_path, self.best_ckpt_path): + """Delete this run's checkpoint files and debris (rank 0 only). + + Besides the two canonical files, a run directory can hold + ``checkpoint_*.pth.tmp.`` (an interrupted write whose Python-level + cleanup never ran) and ``checkpoint_*.pth.corrupt`` (a checkpoint + quarantined on resume). Both are full-checkpoint-sized and nothing else + removes them, so a "from scratch" cleanup that left them would claim to + have cleared the checkpoints while keeping their bytes on disk. + """ + debris = sorted( + set(self.base_dir.glob("checkpoint_*.pth.tmp.*")) + | set(self.base_dir.glob("checkpoint_*.pth.corrupt")) + ) + for p in (self.last_ckpt_path, self.best_ckpt_path, *debris): if p.exists(): try: p.unlink() @@ -150,6 +163,36 @@ def _remove_checkpoint_files(self) -> None: except Exception as e: self._log(f"Failed to remove {p}: {e}") + def _sweep_orphaned_tmp_files(self) -> None: + """Delete temp files stranded by checkpoint writes that were killed. + + ``_atomic_save`` unlinks its ``.tmp.`` file when the write + raises, but a SIGKILL (walltime, node failure) skips that Python-level + cleanup and strands a full-checkpoint-sized file. These accumulate one + per killed pid: the kill/restart cycle that produces them always takes + the *resume* path, so ``cleanup(train_from_scratch=True)`` never gets a + chance to clear them. + + Sweeping at construction is safe because run directories are per-run + and not shared between concurrently running jobs (F55), so any temp + file here belongs to a dead process -- except one from this pid, which + another manager in this process could still be writing. + + ``*.corrupt`` files are deliberately left alone here: + ``_quarantine_corrupt`` renames onto a fixed name, so at most two can + ever exist (they cannot accumulate) and they are the only evidence of + what a resume discarded. The from-scratch cleanup removes them. + """ + own_suffix = f".tmp.{os.getpid()}" + for path in sorted(self.base_dir.glob("checkpoint_*.pth.tmp.*")): + if path.name.endswith(own_suffix): + continue + try: + path.unlink() + self._log(f"Removed orphaned checkpoint temp file: {path}") + except OSError as e: + self._log(f"Failed to remove {path}: {e}") + def wait_for_save(self): """Block until the background save (if any) is complete. diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index f102a6d..2ef9cad 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -27,6 +27,7 @@ from __future__ import annotations import math +import os import re import textwrap import time @@ -464,6 +465,59 @@ def test_cleanup_from_scratch_resets_best(tmp_path): assert mgr2.best_ckpt_path.exists() +# --------------------------------------------------------------------------- +# R07 -- checkpoint debris (.tmp., .corrupt) does not accumulate +# --------------------------------------------------------------------------- + + +def test_cleanup_from_scratch_removes_stale_debris(tmp_path): + """A from-scratch cleanup clears checkpoint debris, not just the two files. + + ``_atomic_save``'s temp file survives a process kill (its unlink only runs + on a Python-level exception) and ``_quarantine_corrupt``'s ``.corrupt`` + rename is never undone. Both are full-checkpoint-sized, so a cleanup that + claims to have cleared the checkpoints while leaving them behind keeps + multi-GB files on the shared filesystem. + """ + mgr, _ = _make_manager(tmp_path) + mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) + + stale_tmp = tmp_path / "checkpoint_last.pth.tmp.999999" + stale_tmp.write_bytes(b"partial checkpoint") + quarantined = tmp_path / "checkpoint_best.pth.corrupt" + quarantined.write_bytes(b"truncated checkpoint") + + mgr.cleanup(train_from_scratch=True) + + assert not stale_tmp.exists() + assert not quarantined.exists() + assert list(tmp_path.iterdir()) == [] + + +def test_init_sweeps_orphaned_tmp_files(tmp_path): + """Constructing a manager sweeps temp files left by killed writes. + + Repeated walltime kills of a long run take the *resume* path, never the + from-scratch cleanup, so without this sweep one stranded temp file per + killed pid piles up in the run directory. A temp file belonging to this + process is left alone (another manager here may still be writing it), and + the quarantined ``.corrupt`` file is kept: at most two can ever exist and + they are the only evidence of what a resume discarded. + """ + orphan = tmp_path / "checkpoint_last.pth.tmp.999999" + orphan.write_bytes(b"partial checkpoint") + own = tmp_path / f"checkpoint_last.pth.tmp.{os.getpid()}" + own.write_bytes(b"possibly in flight") + quarantined = tmp_path / "checkpoint_last.pth.corrupt" + quarantined.write_bytes(b"truncated checkpoint") + + _make_manager(tmp_path) + + assert not orphan.exists() + assert own.exists() + assert quarantined.exists() + + # --------------------------------------------------------------------------- # F71 -- CPU tensors are cloned into the snapshot # --------------------------------------------------------------------------- From 3c4f5b4612e1bd5adf8c4286f70369e0f70982b6 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:35:13 -0700 Subject: [PATCH 08/54] Delete the unused checkpoint_validators module The module had no import sites anywhere in the package, tests or scripts, its usage comments reference a train.py flow that no longer exists, and compare_state_dicts2 crashes on any real optimizer state dict (it truth-tests an elementwise tensor comparison). Round-2 review: R10. --- ScaFFold/utils/checkpoint_validators.py | 143 ------------------------ 1 file changed, 143 deletions(-) delete mode 100644 ScaFFold/utils/checkpoint_validators.py diff --git a/ScaFFold/utils/checkpoint_validators.py b/ScaFFold/utils/checkpoint_validators.py deleted file mode 100644 index c2212b6..0000000 --- a/ScaFFold/utils/checkpoint_validators.py +++ /dev/null @@ -1,143 +0,0 @@ -# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. -# Produced at the Lawrence Livermore National Laboratory. -# Written by the LBANN Research Team (B. Van Essen, et al.) listed in -# the CONTRIBUTORS file. See the top-level LICENSE file for details. -# -# LLNL-CODE-697807. -# All rights reserved. -# -# This file is part of LBANN: Livermore Big Artificial Neural Network -# Toolkit. For details, see http://software.llnl.gov/LBANN or -# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. -# -# SPDX-License-Identifier: (Apache-2.0) - -import os - -# import pdb - -if hasattr(os, "sched_getaffinity"): - _orig_affinity = os.sched_getaffinity(0) -else: - _orig_affinity = None - -import torch - - -def compare_state_dicts(quantity: str, dict1, dict2): - def _compare(dict1, dict2, prefix=""): - equal = True - if dict1.keys() != dict2.keys(): - missing_in_dict2 = dict1.keys() - dict2.keys() - missing_in_dict1 = dict2.keys() - dict1.keys() - if missing_in_dict2: - print( - f"train.py: {quantity} missing in dict2: {', '.join(missing_in_dict2)} at {prefix}" - ) - equal = False - if missing_in_dict1: - print( - f"train.py: {quantity} missing in dict1: {', '.join(missing_in_dict1)} at {prefix}" - ) - equal = False - - for key in dict1.keys() & dict2.keys(): - full_key = f"{prefix}.{key}" if prefix else key - if isinstance(dict1[key], torch.Tensor) and isinstance( - dict2[key], torch.Tensor - ): - if not torch.equal(dict1[key], dict2[key]): - print( - f"train.py: {quantity} tensor discrepancy at {full_key}: dict1[{key}] != dict2[{key}]" - ) - equal = False - elif isinstance(dict1[key], dict) and isinstance(dict2[key], dict): - if not _compare(dict1[key], dict2[key], prefix=full_key): - equal = False - else: - if dict1[key] != dict2[key]: - print( - f"train.py: {quantity} value discrepancy at {full_key}: dict1[{key}]={dict1[key]}, dict2[{key}]={dict2[key]}" - ) - equal = False - return equal - - return _compare(dict1, dict2) - - -def compare_state_dicts2(*dicts): - keys = dicts[0].keys() - for key in keys: - values = [d[key] for d in dicts] - tensor_comparisons = [ - ( - torch.equal(values[i], values[i + 1]) - if torch.is_tensor(values[i]) - else values[i] == values[i + 1] - ) - for i in range(len(values) - 1) - ] - if not all(tensor_comparisons): - return False - return True - - -def compare_tensors(tensor1, tensor2): - return torch.all(torch.eq(tensor1, tensor2)) - - -def compare_items(item1, item2): - if isinstance(item1, torch.Tensor) and isinstance(item2, torch.Tensor): - return compare_tensors(item1, item2) - elif isinstance(item1, dict) and isinstance(item2, dict): - return compare_dicts3(item1, item2) - else: - return item1 == item2 - - -def compare_dicts3(dict1, dict2): - if dict1.keys() != dict2.keys(): - return False - for key in dict1.keys(): - if not compare_items(dict1[key], dict2[key]): - return False - return True - - -# -# Usage in `train.py` below: -# - -# For debugging, write the saved optimizer state to file to compare to loaded state on restart -# timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') -# optim_saved_path = f"{dir_checkpoint}/restarts/optim_saved_epoch{epoch}_{timestamp}.txt" -# with open(optim_saved_path, "w") as optim_f: -# optim_f.write(str(state_dict['optimizer_state_dict'])) - -# -# For debugging purposes, load that checkpoint into new model, optimizer, etc and compare to active -# -# prev_checkpoint = torch.load(checkpoint_path) -# newmodel = UNet(n_channels=3, n_classes=n_classes, trilinear=False, layers=unet_layers) -# newmodel = newmodel.to(memory_format=torch.channels_last_3d) -# newmodel.to(device=device) -# newmodel = torch.nn.parallel.DistributedDataParallel(newmodel, device_ids=[get_cuda_device()], output_device=get_cuda_device()) -# newmodel.module.load_state_dict(prev_checkpoint['model_state_dict']) -# newoptimizer = optim.RMSprop(newmodel.parameters(), -# lr=learning_rate, weight_decay=weight_decay, momentum=momentum, foreach=True) -# if optimizer_name == "ADAM": -# print(f"train.py(w{rank}|l{local_rank}): using ADAM optimizer .........") -# newoptimizer = optim.Adam(newmodel.parameters(), lr=learning_rate) -# elif optimizer_name == "SGD": -# print(f"train.py(w{rank}|l{local_rank}): using SGD optimizer .........") -# newoptimizer = optim.SGD(newmodel.parameters(), lr=learning_rate, momentum=0.9) -# newoptimizer.load_state_dict(prev_checkpoint['optimizer_state_dict']) -# newscheduler = optim.lr_scheduler.ReduceLROnPlateau(newoptimizer, 'max', patience=25) -# newscheduler.load_state_dict(prev_checkpoint['scheduler_state_dict']) - -# # Compare model state dicts -# model_compare = compare_state_dicts("model", model.state_dict(), newmodel.state_dict()) -# optimizer_compare = compare_state_dicts("optimizer", optimizer.state_dict(), newoptimizer.state_dict()) -# scheduler_compare = compare_state_dicts("scheduler", scheduler.state_dict(), newscheduler.state_dict()) -# print(f"train.py: model_compare={model_compare}, optimizer_compare={optimizer_compare}, scheduler_compare={scheduler_compare}") -# print(f"train.py: all equal? {all(compare_dicts3(sd, optimizer.state_dict()) for sd in [state_dict['optimizer_state_dict'], prev_checkpoint['optimizer_state_dict'], newoptimizer.state_dict()])}") From 638aac6b78c86c76d0f84480846fda5e70782f1d Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:39:10 -0700 Subject: [PATCH 09/54] State why no epoch was trained instead of guessing the remedy The nothing-to-resume warning suggested lowering target_dice, which is exactly backwards for the converged case, and read oddly for a fresh run that never entered the loop. Report the actual inputs (start epoch, epochs, starting val dice vs target) instead. Round-2 review: R01, R06 follow-up. --- ScaFFold/utils/trainer.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index da33c12..81dd5f2 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -1086,18 +1086,21 @@ def train(self, profiler=None): completed_epochs = epoch - 1 if not completed_new_epoch: - # The loop exited without running a single new epoch: the resumed - # checkpoint already covers every epoch this run was asked for. - # There is nothing new to save (the existing checkpoint already - # records epoch `completed_epochs`) and none of the per-epoch - # metrics the final save would write were ever computed, so skip - # it and return normally -- the caller's post-processing still has - # the CSV the original run left behind. + # The loop exited without running a single epoch: the state we + # resumed either already covers every epoch this run was asked for, + # or already met target_dice. There is nothing new to save (the + # checkpoint on disk already records epoch `completed_epochs`) and + # none of the per-epoch metrics the final save would write were + # ever computed, so skip it and return normally -- the caller's + # post-processing still has the CSV the original run left behind. self.log.warning( - "Nothing to resume: the loaded checkpoint already covers epoch " - "%s, so no new epoch was trained and no checkpoint was written. " - "Increase 'epochs' (or lower 'target_dice') to train further.", - completed_epochs, + "No new epoch was trained (start epoch %s, 'epochs' %s, " + "starting val dice %s vs target_dice %s): there was nothing to " + "resume, and no checkpoint was written.", + self.start_epoch, + self.config.epochs, + self.start_val_dice, + self.config.target_dice, ) # Save a final checkpoint when the run exits (convergence or max epochs) # at an epoch that was not a checkpoint interval, so the converged From 9df9a1bd0e1083a8576591866987f2e099df2cf3 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:21:18 -0700 Subject: [PATCH 10/54] Guard rank-0 windows in the get_dataset consensus Any failure in the rank-0 reuse/generate decision or the final meta-write and rename is now broadcast as an error sentinel, so peers raise the same error instead of hanging in bcast/Barrier. The reuse scan also skips .tmp_* staging dirs and tolerates unreadable metadata, and staging dir names carry pid+uuid so same-second jobs cannot collide. R27 --- ScaFFold/datagen/get_dataset.py | 108 ++++++++++++--- tests/datagen/test_mpi_consensus.py | 198 ++++++++++++++++++++++++++++ 2 files changed, 286 insertions(+), 20 deletions(-) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index c536b14..bda9572 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -16,9 +16,11 @@ import hashlib import json +import os import shutil import subprocess import time +import uuid from argparse import Namespace from pathlib import Path from typing import Any, Dict @@ -30,6 +32,11 @@ from ScaFFold.utils.utils import setup_mpi_logger META_FILENAME = "meta.yaml" +# Datasets are generated into a staging directory carrying this prefix and only +# renamed into their final ``__`` name once complete, so a +# reader never observes a half-written dataset. The prefix is also what the +# reuse scan skips and what the orphan cleanup collects. +TMP_PREFIX = ".tmp_" # Bumped from 2 to 3 when instance point clouds moved from float64 to float32: # the storage layout is unchanged, but float32 voxel binning shifts a handful of # boundary voxels, so a float64-era dataset must not be reused as if it were @@ -123,15 +130,44 @@ def _decide_reuse_or_generate( staging and final paths for a new generation. Making this decision in one place and broadcasting it prevents ranks from diverging when their views of the shared filesystem differ. + + The scan is deliberately forgiving: a candidate whose metadata is missing, + unreadable, or malformed is warned about and skipped rather than allowed to + raise. This function runs inside a window where every peer is already + waiting in the decision broadcast, so a crash here is a job-wide hang; a + poison directory (exactly what a killed job leaves behind) must never be + able to cause one. """ candidates = sorted( (p for p in base.iterdir() if p.is_dir()), key=lambda p: p.name, reverse=True ) for dataset_path in candidates: + # Staging dirs are not datasets: a job killed between the meta write and + # the rename leaves a complete meta.yaml inside one, and reusing it hands + # back a partially generated (and cleanup-eligible) directory. + if dataset_path.name.startswith(TMP_PREFIX): + continue meta_path = dataset_path / META_FILENAME if not meta_path.exists(): continue - meta = yaml.safe_load(meta_path.read_text()) + try: + meta = yaml.safe_load(meta_path.read_text()) + except Exception as exc: + log.warning( + "Skipping dataset candidate %s: unreadable %s (%s: %s)", + dataset_path, + META_FILENAME, + type(exc).__name__, + exc, + ) + continue + if not isinstance(meta, dict): + log.warning( + "Skipping dataset candidate %s: %s is empty or malformed", + dataset_path, + META_FILENAME, + ) + continue if meta.get("config_id") != config_id: continue if meta.get("dataset_format_version", 1) != DATASET_FORMAT_VERSION: @@ -145,7 +181,11 @@ def _decide_reuse_or_generate( log.info("No valid existing dataset found at %s. Generating new dataset.", base) ts = time.strftime("%Y%m%d-%H%M%S") dest = base / f"{ts}__{commit}" - tmp = base / f".tmp_{ts}" + # The staging name must be unique per job: a bare 1-second-granularity + # timestamp let two same-config jobs starting in the same second collide on + # ``mkdir(exist_ok=False)``, killing one of them mid-consensus. Adding the + # pid and a random suffix makes the name unique even across nodes. + tmp = base / f"{TMP_PREFIX}{ts}_{os.getpid()}_{uuid.uuid4().hex[:8]}" tmp.mkdir(parents=True, exist_ok=False) return ("generate", str(tmp), str(dest)) @@ -188,14 +228,28 @@ def get_dataset( # same branch. Scanning the shared filesystem independently per rank lets # divergent views (stale metadata caches, a racing job's rename) strand some # ranks in the generation collectives while others return early. + # Everything rank 0 does here happens while the peers are already blocked in + # the broadcast below, so a rank-0 exception would strand the whole job. + # Any failure is therefore turned into an error sentinel that travels + # through the same broadcast and makes every rank raise the same error. if rank == 0: - decision = _decide_reuse_or_generate( - base, config_id, commit, require_commit, log - ) + try: + decision = _decide_reuse_or_generate( + base, config_id, commit, require_commit, log + ) + except (Exception, SystemExit) as e: + decision = ( + "error", + f"rank 0 failed to select a dataset under {base}: " + f"{type(e).__name__}: {e}", + ) else: decision = None decision = comm.bcast(decision, root=0) + if decision[0] == "error": + raise RuntimeError(f"dataset selection failed: {decision[1]}") + if decision[0] == "reuse": return Path(decision[1]) @@ -230,21 +284,35 @@ def get_dataset( raise RuntimeError(f"dataset generation failed: {msgs or 'unknown error'}") # rank 0 writes metadata into the staging dir, then renames it into place so - # readers never observe a half-written dataset. + # readers never observe a half-written dataset. This is another rank-0-only + # window inside a collective sequence: the rename can fail (a racing job + # already published this name, quota, ...), so the outcome is broadcast + # rather than allowed to kill rank 0 while the peers wait for it. + finalize_err = "" if rank == 0: - meta = { - "config_id": config_id, - "dataset_format_version": DATASET_FORMAT_VERSION, - "config_subset": volume_config, - "include_keys": INCLUDE_KEYS, - "code_commit": commit, - "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), - } - (tmp / META_FILENAME).write_text( - yaml.safe_dump(meta, sort_keys=True, default_flow_style=False) - ) - tmp.rename(dest) + try: + meta = { + "config_id": config_id, + "dataset_format_version": DATASET_FORMAT_VERSION, + "config_subset": volume_config, + "include_keys": INCLUDE_KEYS, + "code_commit": commit, + "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + } + (tmp / META_FILENAME).write_text( + yaml.safe_dump(meta, sort_keys=True, default_flow_style=False) + ) + tmp.rename(dest) + except (Exception, SystemExit) as e: + finalize_err = ( + f"rank 0 failed to finalize dataset at {dest}: {type(e).__name__}: {e}" + ) + + # This broadcast doubles as the synchronization the old Barrier provided: no + # rank returns before rank 0 has published the rename (or reported that it + # could not), so nobody observes the staging path or a missing dataset. + finalize_err = comm.bcast(finalize_err, root=0) + if finalize_err: + raise RuntimeError(f"dataset generation failed: {finalize_err}") - # ensure the rename is visible everywhere before returning - comm.Barrier() return dest diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index 6a67e89..08bd0c9 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -38,6 +38,7 @@ from __future__ import annotations +import logging import re from argparse import Namespace from math import ceil @@ -340,6 +341,203 @@ def test_generation_success_finalizes_and_returns(tmp_path, monkeypatch): assert leftover == [] +# --------------------------------------------------------------------------- +# R27: rank 0 must never die between the collectives its peers have entered. +# Every rank-0-only step of the consensus (the reuse/generate decision and the +# final meta-write + rename) is wrapped so a failure travels to the peers as a +# broadcast sentinel instead of stranding them in ``bcast``/``Barrier``. +# --------------------------------------------------------------------------- + + +def _base_dir_for(config: Namespace) -> Path: + """The ``/`` directory ``get_dataset`` scans.""" + config_dict = vars(config).copy() + config_dict["dataset_format_version"] = gd.DATASET_FORMAT_VERSION + volume_config = gd._get_required_keys_dict(config_dict, gd.INCLUDE_KEYS) + return Path(config.dataset_dir) / gd._hash_volume_config(volume_config) + + +def test_decision_failure_is_broadcast_not_raised_before_bcast(tmp_path, monkeypatch): + """A rank-0 decision failure reaches peers through the broadcast. + + Any exception inside the rank-0-only decision (an unreadable base dir, a + staging ``mkdir`` hitting ENOSPC, ...) must be converted into an error + sentinel that is broadcast, so peers already waiting in ``bcast`` learn + about it and raise the same error. Before the fix rank 0 raised *before* + reaching the broadcast, leaving every peer blocked forever. + """ + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=2) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + + def explode(*_args, **_kwargs): + raise OSError("No space left on device") + + monkeypatch.setattr(gd, "_decide_reuse_or_generate", explode) + + with pytest.raises(RuntimeError) as excinfo: + gd.get_dataset(config) + + # Rank 0 reached the broadcast before raising, and the payload is the + # error sentinel every peer will see. + assert comm.calls == ["bcast"] + assert comm.bcast_payloads[0][0] == "error" + assert "No space left on device" in str(excinfo.value) + + +def test_non_root_raises_on_broadcast_decision_error(tmp_path, monkeypatch): + """A peer receiving the error sentinel raises instead of generating.""" + config = _reuse_config(tmp_path / "datasets") + sentinel = ("error", "rank 0: OSError: No space left on device") + comm = FakeComm(rank=1, size=2, bcast_returns=[sentinel]) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + + with pytest.raises(RuntimeError) as excinfo: + gd.get_dataset(config) + + assert "No space left on device" in str(excinfo.value) + # The peer stopped at the decision broadcast: no generation collectives. + assert comm.calls == ["bcast"] + + +def test_reuse_scan_skips_staging_dirs(tmp_path, monkeypatch): + """A complete ``meta.yaml`` stranded in a ``.tmp_*`` dir is never reused. + + A job killed between the meta write and the rename leaves a fully valid + meta inside its staging dir. Treating that as a publishable dataset hands + back a half-generated directory (and one that cleanup may delete). + """ + config = _reuse_config(tmp_path / "datasets") + base = _base_dir_for(config) + config_id = base.name + stranded = base / ".tmp_20260101-000000_1234" + stranded.mkdir(parents=True) + (stranded / gd.META_FILENAME).write_text( + yaml.safe_dump( + { + "config_id": config_id, + "dataset_format_version": gd.DATASET_FORMAT_VERSION, + } + ) + ) + + comm = FakeComm(rank=0, size=1) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + log = logging.getLogger("test_reuse_scan_skips_staging_dirs") + + decision = gd._decide_reuse_or_generate(base, config_id, "abc123", False, log) + + assert decision[0] == "generate", f"staging dir was reused: {decision}" + + +def test_reuse_scan_tolerates_corrupt_meta(tmp_path, monkeypatch): + """A corrupt/unreadable candidate meta is skipped, not fatal. + + A 0-byte ``meta.yaml`` (``yaml.safe_load`` -> ``None``) or an unparseable + one used to raise inside the rank-0-only scan. The scan must warn, skip the + directory, and keep looking -- here finding the good dataset next to it. + """ + config = _reuse_config(tmp_path / "datasets") + base = _base_dir_for(config) + config_id = base.name + base.mkdir(parents=True) + + # Sorted-descending scan order visits these two poison dirs first. + (base / "20260301-000000__zzz").mkdir() + (base / "20260301-000000__zzz" / gd.META_FILENAME).write_text("") + (base / "20260201-000000__yyy").mkdir() + (base / "20260201-000000__yyy" / gd.META_FILENAME).write_text("{[not yaml") + + good = _write_reusable_dataset(base, config_id) + log = logging.getLogger("test_reuse_scan_tolerates_corrupt_meta") + + decision = gd._decide_reuse_or_generate(base, config_id, "abc123", False, log) + + assert decision[0] == "reuse" + assert Path(decision[1]) == good + + +def test_staging_dir_names_are_collision_proof(tmp_path, monkeypatch): + """Two decisions in the same second stage into different directories. + + The old name was ``.tmp_%Y%m%d-%H%M%S`` with ``mkdir(exist_ok=False)``, so + two same-config jobs starting in the same second raced to a + ``FileExistsError`` on one of them -- inside the unguarded rank-0 window. + """ + config = _reuse_config(tmp_path / "datasets") + base = _base_dir_for(config) + base.mkdir(parents=True) + log = logging.getLogger("test_staging_dir_names_are_collision_proof") + + # Pin the clock so both decisions share a timestamp: only a non-time + # component can keep the names apart. + monkeypatch.setattr(gd.time, "strftime", lambda *_args: "20260101-000000") + + first = gd._decide_reuse_or_generate(base, base.name, "abc123", False, log) + second = gd._decide_reuse_or_generate(base, base.name, "abc123", False, log) + + assert first[0] == "generate" and second[0] == "generate" + assert first[1] != second[1], "same-second staging dirs collided" + assert Path(first[1]).is_dir() and Path(second[1]).is_dir() + + +def test_finalize_failure_is_broadcast_not_left_to_barrier(tmp_path, monkeypatch): + """A rank-0 rename failure is broadcast; peers raise instead of hanging. + + The rename happens *after* the generation consensus, so a failure there + (e.g. a racing job already created the destination) used to kill rank 0 + while every peer sat in the final ``Barrier``. The fix carries the failure + through one more collective and raises everywhere. + """ + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=2, allreduce_result=1) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + monkeypatch.setattr(volumegen, "main", lambda _config: None) + + # Pin the clock so the destination name is predictable, then have a + # "racing job" occupy it with a non-empty directory: the rename fails with + # ENOTEMPTY exactly as it did in the field. + monkeypatch.setattr(gd.time, "strftime", lambda *_args: "20260101-000000") + base = _base_dir_for(config) + dest = base / "20260101-000000__abc123" + dest.mkdir(parents=True) + (dest / "placeholder").write_text("created by a racing job") + + with pytest.raises(RuntimeError) as excinfo: + gd.get_dataset(config) + + message = str(excinfo.value) + assert "20260101-000000__abc123" in message + # The failure travelled through a collective *after* the generation + # consensus, so peers learn about it rather than waiting in the barrier. + assert "allgather" in comm.calls + assert comm.calls.index("allgather") < len(comm.calls) - 1 + + +def test_non_root_raises_on_broadcast_finalize_error(tmp_path, monkeypatch): + """A peer receiving the finalize error raises rather than returning dest.""" + config = _reuse_config(tmp_path / "datasets") + dest = tmp_path / "datasets" / "cid" / "20260101-000000__abc123" + comm, _tmp, _dest = _generate_decision_comm( + rank=1, size=2, dest=dest, allreduce_result=1 + ) + # Second bcast: root's finalize verdict (a failure message). + comm._bcast_returns.append("rank 0 failed to finalize: OSError: boom") + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + monkeypatch.setattr(volumegen, "main", lambda _config: None) + + with pytest.raises(RuntimeError) as excinfo: + gd.get_dataset(config) + + assert "boom" in str(excinfo.value) + assert not dest.exists() + + # --------------------------------------------------------------------------- # A missing instance file raises FileNotFoundError (a catchable Exception) # rather than calling sys.exit(1) (a BaseException that bypasses consensus), and From 3d26b7b3f83952359ec17c8aebb517e2be45c621 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:23:36 -0700 Subject: [PATCH 11/54] Publish meta.yaml atomically and reject damaged metadata meta.yaml is now written to a temp file, fsynced, and renamed into place, so a killed job cannot leave a truncated document. The loader treats only a missing meta.yaml as legacy v1; a present-but-unreadable or version-less one raises instead of silently reinterpreting a modern dataset. R28 --- ScaFFold/datagen/get_dataset.py | 31 ++++++++++++++-- ScaFFold/utils/data_loading.py | 33 ++++++++++++++--- tests/datagen/test_mpi_consensus.py | 48 ++++++++++++++++++++++++ tests/test_data_loading.py | 57 +++++++++++++++++++++++++++++ 4 files changed, 160 insertions(+), 9 deletions(-) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index bda9572..f200c61 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -115,6 +115,33 @@ def _git_commit_short(log) -> str: return "no-commit-id" +def _write_meta_atomic(meta_path: Path, meta: Dict[str, Any]) -> None: + """Write ``meta`` to ``meta_path`` atomically. + + ``meta.yaml`` is what the loader reads to decide how every sample in the + dataset is interpreted, so a partially written one is worse than none at + all: a truncated file parses as empty and silently reclassifies a modern + dataset as legacy v1. The document is therefore written to a temp file in + the same directory, flushed and fsynced, and only then ``os.replace``d onto + the final name -- an atomic rename within one filesystem. + """ + tmp_path = meta_path.parent / f".{meta_path.name}.tmp{os.getpid()}" + try: + with open(tmp_path, "w") as handle: + handle.write(yaml.safe_dump(meta, sort_keys=True, default_flow_style=False)) + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp_path, meta_path) + except BaseException: + # A failed write must not leave a temp file behind, and the final name + # must keep whatever complete document was already there. + try: + os.remove(tmp_path) + except OSError: + pass + raise + + def _decide_reuse_or_generate( base: Path, config_id: str, @@ -299,9 +326,7 @@ def get_dataset( "code_commit": commit, "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), } - (tmp / META_FILENAME).write_text( - yaml.safe_dump(meta, sort_keys=True, default_flow_style=False) - ) + _write_meta_atomic(tmp / META_FILENAME, meta) tmp.rename(dest) except (Exception, SystemExit) as e: finalize_err = ( diff --git a/ScaFFold/utils/data_loading.py b/ScaFFold/utils/data_loading.py index 564c380..dcc1a45 100644 --- a/ScaFFold/utils/data_loading.py +++ b/ScaFFold/utils/data_loading.py @@ -277,20 +277,41 @@ def _load_numpy_array(path, mmap_mode=None): return np.load(path, allow_pickle=False, mmap_mode=mmap_mode) def _load_dataset_format_version(self): + """Determine which on-disk layout this dataset uses. + + Only a *missing* ``meta.yaml`` means legacy v1: those datasets predate + the metadata file. A metadata file that exists but cannot be read or + does not carry a usable version is a damaged modern dataset, and + falling back to the legacy loader there silently transposes + channels-first volumes and remaps already-dense labels -- corrupt + training data with no error. Such a dataset is rejected instead, with a + message naming the file so it can be repaired or regenerated. + """ meta_path = self.dataset_root / META_FILENAME if not meta_path.exists(): return LEGACY_DATASET_FORMAT_VERSION try: with open(meta_path, "r") as meta_file: - meta = yaml.safe_load(meta_file) or {} + meta = yaml.safe_load(meta_file) except Exception as exc: - customlog( - f"Failed to read dataset metadata from {meta_path}: {exc}. Falling back to legacy loader." - ) - return LEGACY_DATASET_FORMAT_VERSION + raise ValueError( + f"Dataset metadata {meta_path} exists but could not be read " + f"({type(exc).__name__}: {exc}). A dataset carrying a " + f"{META_FILENAME} is not a legacy dataset; refusing to guess its " + "layout. Repair the file or regenerate the dataset." + ) from exc - return int(meta.get("dataset_format_version", LEGACY_DATASET_FORMAT_VERSION)) + version = meta.get("dataset_format_version") if isinstance(meta, dict) else None + try: + return int(version) + except (TypeError, ValueError): + raise ValueError( + f"Dataset metadata {meta_path} is missing a usable " + f"'dataset_format_version' (got {version!r}). A dataset carrying " + f"a {META_FILENAME} is not a legacy dataset; refusing to guess " + "its layout. Repair the file or regenerate the dataset." + ) from None @staticmethod def _prepare_legacy_image(img): diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index 08bd0c9..fa99a1b 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -538,6 +538,54 @@ def test_non_root_raises_on_broadcast_finalize_error(tmp_path, monkeypatch): assert not dest.exists() +# --------------------------------------------------------------------------- +# R28: meta.yaml is published atomically, so no reader ever sees a partial one. +# --------------------------------------------------------------------------- + + +def test_meta_write_is_atomic(tmp_path, monkeypatch): + """An interrupted meta write leaves the previous file intact and no temp. + + ``meta.yaml`` is the file that decides how every sample is interpreted (a + truncated one reclassifies the dataset as legacy v1), so it must appear at + its final name complete or not at all. + """ + target = tmp_path / gd.META_FILENAME + gd._write_meta_atomic(target, {"dataset_format_version": gd.DATASET_FORMAT_VERSION}) + good_bytes = target.read_bytes() + + # Interrupt the write after bytes have reached the temp file but before the + # rename -- the shape of a kill mid-write. + def boom(_fd): + raise OSError("simulated SIGKILL mid-write") + + monkeypatch.setattr(gd.os, "fsync", boom) + + with pytest.raises(OSError): + gd._write_meta_atomic(target, {"dataset_format_version": 99}) + + # The final name still holds the complete previous file, byte-for-byte, and + # no temp file is left behind for the reuse scan to trip over. + assert target.read_bytes() == good_bytes + assert [p.name for p in tmp_path.iterdir()] == [gd.META_FILENAME] + + +def test_published_dataset_has_no_partial_meta(tmp_path, monkeypatch): + """A successful generation publishes a parseable meta and no temp files.""" + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=1, allreduce_result=1) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + monkeypatch.setattr(volumegen, "main", lambda _config: None) + + result = Path(gd.get_dataset(config)) + + meta = yaml.safe_load((result / gd.META_FILENAME).read_text()) + assert meta["dataset_format_version"] == gd.DATASET_FORMAT_VERSION + # Nothing hidden alongside it (a temp meta would be a dotted sibling). + assert [p.name for p in result.iterdir() if p.name.startswith(".")] == [] + + # --------------------------------------------------------------------------- # A missing instance file raises FileNotFoundError (a catchable Exception) # rather than calling sys.exit(1) (a BaseException that bypasses consensus), and diff --git a/tests/test_data_loading.py b/tests/test_data_loading.py index 5f3a623..b29d212 100644 --- a/tests/test_data_loading.py +++ b/tests/test_data_loading.py @@ -594,3 +594,60 @@ def counting_load(path, mmap_mode=None): assert loaded["image"] == 0 assert mask_only.dtype == torch.int16 assert torch.equal(mask_only, expected) + + +# --------------------------------------------------------------------------- +# R28: a *present but broken* meta.yaml must not be mistaken for a v1 dataset +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "broken_meta", + [ + pytest.param("", id="zero-byte"), + pytest.param("{[not: valid: yaml", id="unparseable"), + pytest.param("- just\n- a\n- list\n", id="not-a-mapping"), + pytest.param("config_id: abc123\n", id="version-key-missing"), + pytest.param("dataset_format_version: two\n", id="version-not-an-int"), + ], +) +def test_broken_meta_raises_instead_of_silent_legacy(tmp_path, broken_meta): + """A corrupt ``meta.yaml`` is an error, never a silent legacy downgrade. + + Treating a broken meta as "no meta" reclassifies a modern dataset as legacy + v1: the loader then transposes channels-first volumes (a (3,N,N,N) sample + comes back (N,3,N,N)) and remaps already-dense labels. Training proceeds on + silently corrupted data. The dataset directory itself is intact here -- only + the metadata is damaged -- so the failure must be loud and actionable. + """ + root = _build_v2_constant_dataset(tmp_path / "ds", n_volumes=2) + (root / "meta.yaml").write_text(broken_meta) + + with pytest.raises(ValueError) as excinfo: + FractalDataset( + root / "volumes" / "training", + root / "masks" / "training", + data_dir=root / "train_unique_mask_vals", + ) + + message = str(excinfo.value) + assert "meta.yaml" in message + # The message must point at the offending file so it can be repaired. + assert str(root) in message + + +def test_absent_meta_is_still_legacy_v1(tiny_v1_dataset): + """The genuine legacy case (no ``meta.yaml`` at all) is unchanged. + + Control for the test above: v1 datasets predate the metadata file, so a + *missing* meta must keep selecting the legacy loader rather than raising. + """ + root = tiny_v1_dataset(n_categories=2, n_train=2, n_val=1, n=8) + assert not (root / "meta.yaml").exists() + + ds = FractalDataset( + root / "volumes" / "training", + root / "masks" / "training", + data_dir=root / "train_unique_mask_vals", + ) + assert ds.dataset_format_version == 1 From 956b18930b255f058ae88b10479d3fdac6951155 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:25:33 -0700 Subject: [PATCH 12/54] Reclaim orphaned dataset staging directories Killed generations left their .tmp_* staging trees under the config_id base forever, so every retry stacked another copy. Rank 0 now removes staging dirs that have sat untouched past an age threshold, which keeps a concurrent job's (far younger) staging dir safe. R37 --- ScaFFold/datagen/get_dataset.py | 56 ++++++++++++++++++++++++ tests/datagen/test_mpi_consensus.py | 68 +++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index f200c61..0b193f2 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -37,6 +37,10 @@ # reader never observes a half-written dataset. The prefix is also what the # reuse scan skips and what the orphan cleanup collects. TMP_PREFIX = ".tmp_" +# How long a staging directory must have sat untouched before it is treated as +# orphaned (left by a killed/OOM'd job) and reclaimed. See +# ``_cleanup_stale_staging_dirs`` for the safety argument behind the value. +STALE_STAGING_AGE_SECONDS = 24 * 60 * 60 # Bumped from 2 to 3 when instance point clouds moved from float64 to float32: # the storage layout is unchanged, but float32 voxel binning shifts a handful of # boundary voxels, so a float64-era dataset must not be reused as if it were @@ -115,6 +119,54 @@ def _git_commit_short(log) -> str: return "no-commit-id" +def _cleanup_stale_staging_dirs( + base: Path, log, max_age: float = STALE_STAGING_AGE_SECONDS +) -> None: + """Reclaim orphaned ``.tmp_*`` staging directories under one config base. + + A generation killed by a walltime limit, an OOM, or a node failure leaves + its whole staging tree behind, and nothing ever removed it: every retry + stacked another (potentially multi-terabyte) copy under the same config_id. + + Safety policy. Only directories that (a) live directly under *this* + config_id base, (b) carry the ``.tmp_`` prefix this module owns, and (c) + have been untouched for ``max_age`` are removed. The age gate is what keeps + a *concurrent* job's staging directory safe: unique staging names mean two + live jobs never share a directory, but they do share the base, so a live + peer's directory is visible here -- it is simply orders of magnitude younger + than the threshold (a day, against generations measured in minutes to + hours). Published datasets and anything outside ``base`` are never touched. + Failures are logged and ignored: cleanup is opportunistic and must never + break the decision it runs inside. + """ + now = time.time() + for path in base.iterdir(): + if not path.name.startswith(TMP_PREFIX) or not path.is_dir(): + continue + try: + # Newest mtime among the staging dir and its immediate children: a + # bounded, cheap probe (no recursive stat storm over a partially + # generated dataset) that still notices a job that has started + # laying down its split directories. + newest = path.stat().st_mtime + for child in path.iterdir(): + newest = max(newest, child.stat().st_mtime) + except OSError as exc: + log.warning("Could not stat staging dir %s: %s", path, exc) + continue + + age = now - newest + if age < max_age: + continue + + log.info( + "Removing orphaned dataset staging dir %s (untouched for %.1f hours)", + path, + age / 3600.0, + ) + shutil.rmtree(path, ignore_errors=True) + + def _write_meta_atomic(meta_path: Path, meta: Dict[str, Any]) -> None: """Write ``meta`` to ``meta_path`` atomically. @@ -165,6 +217,10 @@ def _decide_reuse_or_generate( poison directory (exactly what a killed job leaves behind) must never be able to cause one. """ + # Rank 0 is the only rank that touches this base, so this is also the one + # safe place to reclaim staging dirs orphaned by earlier killed jobs. + _cleanup_stale_staging_dirs(base, log) + candidates = sorted( (p for p in base.iterdir() if p.is_dir()), key=lambda p: p.name, reverse=True ) diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index fa99a1b..5086150 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -39,7 +39,9 @@ from __future__ import annotations import logging +import os import re +import time from argparse import Namespace from math import ceil from pathlib import Path @@ -538,6 +540,72 @@ def test_non_root_raises_on_broadcast_finalize_error(tmp_path, monkeypatch): assert not dest.exists() +# --------------------------------------------------------------------------- +# R37: orphaned staging dirs are reclaimed instead of accumulating forever. +# --------------------------------------------------------------------------- + + +def _age_tree(path: Path, seconds: float) -> None: + """Backdate ``path`` and everything under it by ``seconds``.""" + stamp = time.time() - seconds + for entry in sorted(path.rglob("*"), reverse=True): + os.utime(entry, (stamp, stamp)) + os.utime(path, (stamp, stamp)) + + +def test_stale_staging_dirs_are_cleaned(tmp_path, monkeypatch): + """A long-orphaned ``.tmp_*`` dir is reclaimed; a live one is not. + + Every killed generation leaves a full staging tree behind (potentially + terabytes) that nothing ever reclaims. Cleanup is age-gated so a *running* + job's staging dir -- which by construction is far younger than the + threshold -- is never deleted out from under it. + """ + config = _reuse_config(tmp_path / "datasets") + base = _base_dir_for(config) + base.mkdir(parents=True) + + orphan = base / f"{gd.TMP_PREFIX}20200101-000000_111_deadbeef" + (orphan / "volumes" / "training").mkdir(parents=True) + (orphan / "volumes" / "training" / "0.npy").write_bytes(b"stale payload") + _age_tree(orphan, 10 * gd.STALE_STAGING_AGE_SECONDS) + + live = base / f"{gd.TMP_PREFIX}20260101-000000_222_cafebabe" + (live / "volumes").mkdir(parents=True) + + log = logging.getLogger("test_stale_staging_dirs_are_cleaned") + gd._decide_reuse_or_generate(base, base.name, "abc123", False, log) + + assert not orphan.exists(), "orphaned staging dir was not reclaimed" + assert live.exists(), "a concurrent job's staging dir was deleted" + + +def test_cleanup_never_touches_published_datasets(tmp_path): + """Only ``.tmp_*`` dirs under this config_id base are ever removed. + + An old published dataset is precisely what reuse is for, and other configs' + (or other users') directories are none of this job's business. + """ + config = _reuse_config(tmp_path / "datasets") + base = _base_dir_for(config) + base.mkdir(parents=True) + + published = _write_reusable_dataset(base, "some-other-config") + _age_tree(published, 10 * gd.STALE_STAGING_AGE_SECONDS) + + # A staging dir belonging to a different config_id base entirely. + other_base = base.parent / "0123456789ab" + other_orphan = other_base / f"{gd.TMP_PREFIX}20200101-000000_333_f00d" + other_orphan.mkdir(parents=True) + _age_tree(other_orphan, 10 * gd.STALE_STAGING_AGE_SECONDS) + + log = logging.getLogger("test_cleanup_never_touches_published_datasets") + gd._decide_reuse_or_generate(base, base.name, "abc123", False, log) + + assert published.exists(), "an old published dataset was deleted" + assert other_orphan.exists(), "cleanup escaped this job's config_id base" + + # --------------------------------------------------------------------------- # R28: meta.yaml is published atomically, so no reader ever sees a partial one. # --------------------------------------------------------------------------- From 74b0fd490e2b72790148754ccea6bd43a94beedf Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:32:23 -0700 Subject: [PATCH 13/54] Key the fractal library directory by seed Categories and instances are derived from config.seed but resume is a pure file-existence check, so a run under a new seed silently adopted another seed's library and published a dataset stamped with the wrong seed. Library paths now carry seed, so cross-seed reuse is impossible; old-layout libraries are not found and are regenerated. R29 --- ScaFFold/datagen/category_search.py | 11 +- ScaFFold/datagen/instance.py | 13 +- ScaFFold/datagen/layout.py | 67 +++++ ScaFFold/datagen/volumegen.py | 11 +- tests/datagen/test_artifacts.py | 9 +- tests/datagen/test_library_layout.py | 248 ++++++++++++++++++ tests/datagen/test_mpi_consensus.py | 24 +- .../datagen_get_dataset_consensus.py | 14 +- .../datagen_instance_partition.py | 5 +- 9 files changed, 353 insertions(+), 49 deletions(-) create mode 100644 ScaFFold/datagen/layout.py create mode 100644 tests/datagen/test_library_layout.py diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py index 246666f..ae3b8fb 100644 --- a/ScaFFold/datagen/category_search.py +++ b/ScaFFold/datagen/category_search.py @@ -26,6 +26,7 @@ import numpy as np from mpi4py import MPI +from ScaFFold.datagen import layout from ScaFFold.datagen.generate_fractal_points import generate_fractal_points from ScaFFold.datagen.rng import SEED_MASK, derive_seed, seed_numba from ScaFFold.utils.config_utils import Config @@ -403,11 +404,11 @@ def main(config: Config) -> None: log.info("MPI size = %s", size) - # Setup directories - fracts_sub_dir = f"var{config.variance_threshold}" - fracts_write_dir = os.path.join( - config.fract_base_dir, fracts_sub_dir, "3DIFS_param" - ) + # Setup directories. The library is keyed by seed (see + # ScaFFold.datagen.layout): categories are drawn from a seed-derived + # candidate stream, so a run under a different seed must never resume onto + # another seed's parameter files. + fracts_write_dir = layout.category_param_dir(config) if rank == 0: log.info("Writing fractals to %s", fracts_write_dir) if os.path.exists(fracts_write_dir) and config.datagen_from_scratch: diff --git a/ScaFFold/datagen/instance.py b/ScaFFold/datagen/instance.py index e28d104..de8eb0d 100644 --- a/ScaFFold/datagen/instance.py +++ b/ScaFFold/datagen/instance.py @@ -28,6 +28,7 @@ import numpy as np from mpi4py import MPI +from ScaFFold.datagen import layout from ScaFFold.datagen.generate_fractal_points import generate_fractal_points from ScaFFold.datagen.rng import derive_seed, seed_numba from ScaFFold.utils.config_utils import Config @@ -239,12 +240,12 @@ def main(config: Config): log.info("MPI size = %s", size) - # Setup directories - fracts_sub_dir = f"var{config.variance_threshold}" - fracts_read_dir = os.path.join(config.fract_base_dir, fracts_sub_dir, "3DIFS_param") - instance_write_dir = os.path.join( - config.fract_base_dir, fracts_sub_dir, "instances", f"np{config.point_num}" - ) + # Setup directories. The library is keyed by seed (see + # ScaFFold.datagen.layout): every instance is generated from + # (seed, category, instance), so resuming onto another seed's files would + # silently mix data from two different seeds into one dataset. + fracts_read_dir = layout.category_param_dir(config) + instance_write_dir = layout.instance_dir(config) if rank == 0: log.info( "Generating instances for num_points=%s, writing to %s", diff --git a/ScaFFold/datagen/layout.py b/ScaFFold/datagen/layout.py new file mode 100644 index 0000000..e558c6d --- /dev/null +++ b/ScaFFold/datagen/layout.py @@ -0,0 +1,67 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""On-disk layout of the fractal library. + +Every artifact in the library is a deterministic function of the seed: category +IFS parameters come from a ``(seed, rank, attempt)`` candidate stream, and each +instance point cloud is generated from ``(seed, category, instance)``. Resume, +by contrast, is a pure file-existence test -- an instance is "already done" if +its file is on disk. + +Those two facts only compose safely if the path itself carries the seed. +Without it, a run under a new seed found the previous seed's files, generated +nothing, and produced a dataset whose metadata advertised the new seed while +its contents came from the old one. Keying the directory by seed makes the +existence question seed-specific, so data from one seed can never be mistaken +for another's: + + /var/seed/3DIFS_param/ + /var/seed/instances/np/ + +Libraries written under the older, seed-agnostic layout are simply not found +and are regenerated in the new location. + +These helpers are the single definition of that layout; every producer and +consumer (``category_search``, ``instance``, ``volumegen``) goes through them +so the two sides cannot drift apart. +""" + +from __future__ import annotations + +import os + + +def library_root(config) -> str: + """Return the root of the fractal library for this config's seed.""" + return os.path.join( + str(config.fract_base_dir), + f"var{config.variance_threshold}", + f"seed{int(config.seed)}", + ) + + +def category_param_dir(config) -> str: + """Return the directory holding this seed's category IFS parameter CSVs.""" + return os.path.join(library_root(config), "3DIFS_param") + + +def instance_dir(config) -> str: + """Return the directory holding this seed's instance point clouds. + + Instances are additionally keyed by point count, which is a property of the + cloud rather than of the category, so several point counts can coexist for + one seed. + """ + return os.path.join(library_root(config), "instances", f"np{config.point_num}") diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py index cfd6d1b..05567e0 100644 --- a/ScaFFold/datagen/volumegen.py +++ b/ScaFFold/datagen/volumegen.py @@ -22,6 +22,7 @@ import numpy as np from mpi4py import MPI +from ScaFFold.datagen import layout from ScaFFold.utils.config_utils import Config from ScaFFold.utils.data_types import MASK_DTYPE, VOLUME_DTYPE from ScaFFold.utils.utils import setup_mpi_logger @@ -240,7 +241,10 @@ def main(config: Dict): fractal_colors = np.random.rand(config.n_categories, 3) grid_size = resolve_grid_size(config) - fract_base_dir = str(config.fract_base_dir) + # The instance library is keyed by seed (see ScaFFold.datagen + # .layout), so a volume can only ever be built from point clouds + # this run's seed produced. Resolved once, outside the loop. + instances_dir = layout.instance_dir(config) # Generation loop start_time = time.time() @@ -269,12 +273,7 @@ def main(config: Dict): curr_instance = curr_vol[1 + 2 * curr_fract + 1] fractal_color = fractal_colors[curr_category] - instances_dir = ( - f"var{config.variance_threshold}/instances/np{config.point_num}" - ) - point_cloud_path = os.path.join( - fract_base_dir, instances_dir, f"{curr_category:06d}", f"{curr_category:06d}_{curr_instance:04d}.npy", diff --git a/tests/datagen/test_artifacts.py b/tests/datagen/test_artifacts.py index 224fefa..2f969a6 100644 --- a/tests/datagen/test_artifacts.py +++ b/tests/datagen/test_artifacts.py @@ -37,6 +37,7 @@ import pytest from ScaFFold.datagen import instance as inst +from ScaFFold.datagen import layout from ScaFFold.datagen import mask_detection as md from ScaFFold.datagen.volumegen import ( load_np_ptcloud, @@ -74,13 +75,15 @@ def _seed_category(fract_base: Path, *, point_num: int, keep: range) -> Path: Pre-seeding all but instance 0 means a ``main`` run only has to generate the single missing instance, keeping the test fast. Returns the instance dir. + The library lives under the seed-keyed layout, so the paths are derived from + the same config the run under test uses. """ - vt = 0.15 - param_dir = fract_base / f"var{vt}" / "3DIFS_param" + config = _make_config(fract_base, point_num=point_num) + param_dir = Path(layout.category_param_dir(config)) param_dir.mkdir(parents=True) np.savetxt(param_dir / "000000.csv", _contractive_params(), delimiter=",") - inst_dir = fract_base / f"var{vt}" / "instances" / f"np{point_num}" / "000000" + inst_dir = Path(layout.instance_dir(config)) / "000000" inst_dir.mkdir(parents=True) rng = np.random.default_rng(0) for i in keep: diff --git a/tests/datagen/test_library_layout.py b/tests/datagen/test_library_layout.py new file mode 100644 index 0000000..d09a524 --- /dev/null +++ b/tests/datagen/test_library_layout.py @@ -0,0 +1,248 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""The fractal library is keyed by the seed that produced it (R29). + +Categories and instances are *derived from* ``config.seed``: the IFS parameters +come from a seed-keyed candidate stream and every instance point cloud is +seeded from ``(seed, category, instance)``. Resume, however, is a pure +file-existence check, so a library laid out only by variance threshold and +point count let a run under one seed silently adopt another seed's data -- +and then publish a dataset whose metadata claimed the new seed. Two datasets +with identical provenance and different content. + +The fix puts the seed in the directory path, so the question "does this file +exist" is asked in a seed-specific place and can only ever be answered with +data that seed produced: + + /var/seed/3DIFS_param/ + /var/seed/instances/np/ + +Everything here runs single-process at tiny scale (one category, 60-point +clouds) so the real generators run in well under a second. +""" + +from __future__ import annotations + +import hashlib +from argparse import Namespace +from pathlib import Path + +import numpy as np +import pytest + +from ScaFFold.datagen import category_search as cs +from ScaFFold.datagen import instance as inst +from ScaFFold.datagen import layout, volumegen + +VT = 0.15 +POINT_NUM = 60 + + +def _contractive_params() -> np.ndarray: + """A 2-map IFS whose orbit stays bounded, so generation is fast and finite.""" + params = np.zeros((2, 13), dtype=np.float64) + params[:, 0] = params[:, 4] = params[:, 8] = 0.5 + params[1, 9] = params[1, 10] = params[1, 11] = 0.5 + params[0, 12] = 0.5 + return params + + +def _param_dir(fract_base: Path, seed: int) -> Path: + """The category directory the new layout mandates, spelled out literally.""" + return fract_base / f"var{VT}" / f"seed{seed}" / "3DIFS_param" + + +def _instance_dir(fract_base: Path, seed: int) -> Path: + """The instance directory the new layout mandates, spelled out literally.""" + return fract_base / f"var{VT}" / f"seed{seed}" / "instances" / f"np{POINT_NUM}" + + +def _seed_params(fract_base: Path, seed: int, n_categories: int = 1) -> Path: + param_dir = _param_dir(fract_base, seed) + param_dir.mkdir(parents=True, exist_ok=True) + for category in range(n_categories): + np.savetxt( + param_dir / f"{category:06d}.csv", _contractive_params(), delimiter="," + ) + return param_dir + + +def _inst_config(fract_base: Path, seed: int) -> Namespace: + return Namespace( + fract_base_dir=str(fract_base), + n_categories=1, + seed=seed, + variance_threshold=VT, + point_num=POINT_NUM, + datagen_from_scratch=False, + verbose=0, + ) + + +def _cs_config(fract_base: Path, seed: int) -> Namespace: + return Namespace( + fract_base_dir=str(fract_base), + n_categories=1, + seed=seed, + variance_threshold=VT, + point_num=POINT_NUM, + normalize=1, + datagen_from_scratch=False, + datagen_batch_size=4, + verbose=0, + ) + + +def _volumegen_config(dataset_dir: Path, fract_base: Path, seed: int) -> Namespace: + return Namespace( + dataset_dir=str(dataset_dir), + fract_base_dir=str(fract_base), + n_categories=1, + n_instances_used_per_fractal=1, + n_fracts_per_vol=1, + seed=seed, + variance_threshold=VT, + val_split=0, + vol_size=8, + point_num=POINT_NUM, + scale=1, + verbose=0, + ) + + +def _library_digest(instance_dir: Path) -> tuple[int, str]: + """(file count, content digest) for one instance directory.""" + digest = hashlib.sha256() + files = sorted(instance_dir.rglob("*.npy")) + for path in files: + digest.update(path.name.encode()) + digest.update(path.read_bytes()) + return len(files), digest.hexdigest() + + +# --------------------------------------------------------------------------- +# The layout helper is the single definition of the seed-keyed paths. +# --------------------------------------------------------------------------- + + +def test_layout_helpers_key_every_path_by_seed(tmp_path): + """Both library directories carry the seed, and differ across seeds.""" + fract_base = tmp_path / "fractals" + + for seed in (7, 999): + config = _inst_config(fract_base, seed) + assert Path(layout.category_param_dir(config)) == _param_dir(fract_base, seed) + assert Path(layout.instance_dir(config)) == _instance_dir(fract_base, seed) + + assert layout.category_param_dir(_inst_config(fract_base, 7)) != ( + layout.category_param_dir(_inst_config(fract_base, 999)) + ) + assert layout.instance_dir(_inst_config(fract_base, 7)) != ( + layout.instance_dir(_inst_config(fract_base, 999)) + ) + + +# --------------------------------------------------------------------------- +# Producers write under the seed; consumers read from under the seed. +# --------------------------------------------------------------------------- + + +def test_category_search_writes_under_the_seed_dir(tmp_path): + """A generated category CSV lands in this seed's parameter directory.""" + fract_base = tmp_path / "fractals" + + cs.main(_cs_config(fract_base, seed=42)) + + assert (_param_dir(fract_base, 42) / "000000.csv").exists() + # Nothing was written to a seed-agnostic location. + assert not (fract_base / f"var{VT}" / "3DIFS_param").exists() + + +def test_instances_are_written_under_the_seed_dir(tmp_path): + """Instance point clouds land under this seed's instance directory.""" + fract_base = tmp_path / "fractals" + _seed_params(fract_base, seed=7) + + inst.main(_inst_config(fract_base, seed=7)) + + count, _digest = _library_digest(_instance_dir(fract_base, 7)) + assert count == 145 + assert not (fract_base / f"var{VT}" / "instances").exists() + + +def test_same_seed_resume_generates_nothing_new(tmp_path): + """Re-running under the same seed reuses the library byte-for-byte. + + The resume path must stay cheap: the whole point of the library is that a + second run under the same configuration regenerates nothing. + """ + fract_base = tmp_path / "fractals" + _seed_params(fract_base, seed=7) + inst.main(_inst_config(fract_base, seed=7)) + + instance_dir = _instance_dir(fract_base, 7) + count_before, digest_before = _library_digest(instance_dir) + mtimes_before = {p: p.stat().st_mtime_ns for p in sorted(instance_dir.rglob("*"))} + + inst.main(_inst_config(fract_base, seed=7)) + + count_after, digest_after = _library_digest(instance_dir) + assert (count_after, digest_after) == (count_before, digest_before) + # No file was rewritten (0 new instances generated). + assert {p: p.stat().st_mtime_ns for p in sorted(instance_dir.rglob("*"))} == ( + mtimes_before + ) + + +def test_different_seed_cannot_reuse_another_seeds_instances(tmp_path): + """A second seed generates its own library instead of adopting the first. + + Before the fix this run reported "Generated 0 instances" and left the + first seed's bytes in place, so the dataset built on top of it carried the + wrong seed's data under the new seed's provenance. + """ + fract_base = tmp_path / "fractals" + _seed_params(fract_base, seed=7) + _seed_params(fract_base, seed=999) + + inst.main(_inst_config(fract_base, seed=7)) + count_7, digest_7 = _library_digest(_instance_dir(fract_base, 7)) + + inst.main(_inst_config(fract_base, seed=999)) + count_999, digest_999 = _library_digest(_instance_dir(fract_base, 999)) + + # Both libraries are complete and independent... + assert count_7 == count_999 == 145 + assert digest_999 != digest_7, "seed 999 reused seed 7's instances" + # ...and the first seed's data was left untouched. + assert _library_digest(_instance_dir(fract_base, 7)) == (count_7, digest_7) + + +def test_volumegen_reads_instances_for_its_own_seed(tmp_path): + """volumegen resolves point clouds under the seed it was configured with.""" + fract_base = tmp_path / "fractals" + _seed_params(fract_base, seed=7) + inst.main(_inst_config(fract_base, seed=7)) + + # Seed 7: the instances are where volumegen looks, so generation succeeds. + volumegen.main(_volumegen_config(tmp_path / "ds7", fract_base, seed=7)) + assert list((tmp_path / "ds7" / "volumes").rglob("*.npy")) + + # Seed 999: a different library entirely. Nothing has been generated for + # it, so volumegen must report the missing file rather than quietly + # rasterizing seed 7's clouds. + with pytest.raises(RuntimeError) as excinfo: + volumegen.main(_volumegen_config(tmp_path / "ds999", fract_base, seed=999)) + assert "seed999" in str(excinfo.value) diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index 5086150..c98fc5b 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -53,7 +53,7 @@ from ScaFFold.datagen import get_dataset as gd from ScaFFold.datagen import instance as inst -from ScaFFold.datagen import volumegen +from ScaFFold.datagen import layout, volumegen RANK_SCRIPTS = Path(__file__).resolve().parents[1] / "helpers" / "rank_scripts" @@ -684,15 +684,10 @@ def _seed_one_instance(fract_base: Path, config: Namespace, *, present: bool) -> volumegen selects instance indices with ``random.sample(range(145), ...)`` seeded by ``config.seed``; to be robust we populate every one of the 145 - instance slots for category 0 when ``present`` is True. + instance slots for category 0 when ``present`` is True. The library path is + seed-keyed, so it is derived from the same config the run under test uses. """ - inst_dir = ( - fract_base - / f"var{config.variance_threshold}" - / "instances" - / f"np{config.point_num}" - / "000000" - ) + inst_dir = Path(layout.instance_dir(config)) / "000000" inst_dir.mkdir(parents=True, exist_ok=True) if present: rng = np.random.default_rng(0) @@ -803,7 +798,7 @@ def _instance_config(fract_base: Path) -> Namespace: def _seed_ifs_params(fract_base: Path, config: Namespace, n_categories: int) -> None: """Write a contractive IFS param CSV per category so generation stays fast.""" - param_dir = fract_base / f"var{config.variance_threshold}" / "3DIFS_param" + param_dir = Path(layout.category_param_dir(config)) param_dir.mkdir(parents=True, exist_ok=True) params = np.zeros((2, 13), dtype=np.float64) params[:, 0] = params[:, 4] = params[:, 8] = 0.5 @@ -869,14 +864,7 @@ def no_scan(*_args, **_kwargs): assert rc == 0 # Rank 1 received the broadcast list and generated its share (pair [1, 0]). - generated = ( - fract_base - / f"var{config.variance_threshold}" - / "instances" - / f"np{config.point_num}" - / "000001" - / "000001_0000.npy" - ) + generated = Path(layout.instance_dir(config)) / "000001" / "000001_0000.npy" assert generated.exists() diff --git a/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py b/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py index 3799874..bc7e7a4 100644 --- a/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py +++ b/tests/helpers/rank_scripts/datagen_get_dataset_consensus.py @@ -45,7 +45,7 @@ from mpi4py import MPI import ScaFFold.datagen.get_dataset as gd -from ScaFFold.datagen import volumegen +from ScaFFold.datagen import layout, volumegen VT = 0.15 PN = 64 @@ -71,14 +71,10 @@ def _config(dataset_dir: Path, fract_base: Path) -> Namespace: def _instance_path(fract_base: Path, cat: int, inst: int) -> Path: - return ( - fract_base - / f"var{VT}" - / "instances" - / f"np{PN}" - / f"{cat:06d}" - / f"{cat:06d}_{inst:04d}.npy" - ) + # The instance library is keyed by seed; derive the path from the same + # config the run under test uses. + inst_root = Path(layout.instance_dir(_config(Path("unused"), fract_base))) + return inst_root / f"{cat:06d}" / f"{cat:06d}_{inst:04d}.npy" def _seed_instances(fract_base: Path) -> None: diff --git a/tests/helpers/rank_scripts/datagen_instance_partition.py b/tests/helpers/rank_scripts/datagen_instance_partition.py index 10df10f..37d0c98 100644 --- a/tests/helpers/rank_scripts/datagen_instance_partition.py +++ b/tests/helpers/rank_scripts/datagen_instance_partition.py @@ -44,6 +44,7 @@ from mpi4py import MPI import ScaFFold.datagen.instance as inst +from ScaFFold.datagen import layout VT = 0.15 PN = 64 @@ -62,7 +63,7 @@ def _config(fract_base: Path) -> Namespace: def _seed_ifs_params(fract_base: Path) -> None: - param_dir = fract_base / f"var{VT}" / "3DIFS_param" + param_dir = Path(layout.category_param_dir(_config(fract_base))) param_dir.mkdir(parents=True, exist_ok=True) params = np.zeros((2, 13), dtype=np.float64) params[:, 0] = params[:, 4] = params[:, 8] = 0.5 @@ -82,7 +83,7 @@ def main() -> None: _seed_ifs_params(fract_base) comm.Barrier() - inst_root = fract_base / f"var{VT}" / "instances" / f"np{PN}" + inst_root = Path(layout.instance_dir(_config(fract_base))) # Give each rank a divergent view of pre-existing instances. Only rank 0's # view should matter after the fix (its list is broadcast); rank 1's phantom From aa088c176bf8c6e79ca1185c1e4ddaffa2e1cec8 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:33:57 -0700 Subject: [PATCH 14/54] Scan for existing categories once on rank 0 and broadcast category_search derived its loop-gating remaining count from a per-rank filesystem scan, so ranks with divergent views could post mismatched collectives (bcast against reduce) and hang. Rank 0 now scans and broadcasts the existing-index list, mirroring the instance.py work-list fix. R30 --- ScaFFold/datagen/category_search.py | 17 +++- tests/datagen/test_category_search.py | 139 +++++++++++++++++++++++++- 2 files changed, 152 insertions(+), 4 deletions(-) diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py index ae3b8fb..98e2342 100644 --- a/ScaFFold/datagen/category_search.py +++ b/ScaFFold/datagen/category_search.py @@ -424,9 +424,20 @@ def main(config: Config) -> None: # the ones a fresh run produced. attempt_index = read_attempt_counter(fracts_write_dir, rank) - # Parse existing category files (rank 0 owns saving/dedup). Free indices are - # derived from these parsed names -- filling holes, never overwriting. - existing_indices = parse_category_indices(fracts_write_dir) + # Parse existing category files on rank 0 alone and broadcast the result. + # Free indices are derived from these parsed names -- filling holes, never + # overwriting -- and, critically, the count derived below gates a loop that + # contains collectives. Scanning the shared filesystem independently per + # rank lets divergent views (stale metadata caches, a concurrent job, a + # partially visible directory) put one rank inside the loop while another is + # past it, so the two post mismatched collectives on COMM_WORLD and the job + # hangs. One scan, one broadcast, one shared verdict. + if rank == 0: + existing_indices = parse_category_indices(fracts_write_dir) + else: + existing_indices = None + existing_indices = comm.bcast(existing_indices, root=0) + existing_params = [] if rank == 0: for idx in existing_indices: diff --git a/tests/datagen/test_category_search.py b/tests/datagen/test_category_search.py index 8608884..90236a8 100644 --- a/tests/datagen/test_category_search.py +++ b/tests/datagen/test_category_search.py @@ -12,8 +12,15 @@ # # SPDX-License-Identifier: (Apache-2.0) -"""Tests for category-search round sizing.""" +"""Tests for category-search round sizing and its work-scan consensus.""" +from argparse import Namespace +from pathlib import Path + +import numpy as np + +from ScaFFold.datagen import category_search as cs +from ScaFFold.datagen import layout from ScaFFold.datagen.category_search import compute_round_attempts @@ -53,3 +60,133 @@ def test_at_least_one_attempt_per_rank_when_work_remains(): # A positive remaining count always yields at least one attempt per rank so # the loop makes progress and can keep learning the acceptance rate. assert compute_round_attempts(1, 1024, 10000, 0.99) >= 1 + + +# --------------------------------------------------------------------------- +# R30: the initial work scan is made once on rank 0 and broadcast. +# +# ``categories_remaining`` gates a while loop that contains collectives, so it +# must be identical on every rank. Deriving it from a per-rank filesystem scan +# lets divergent views (a stale metadata cache, a racing job, a partially +# visible directory) put one rank inside the loop issuing ``bcast`` while +# another is past it issuing ``reduce`` -- mismatched collectives on +# COMM_WORLD, i.e. a hang. This mirrors the fix already applied in +# ``instance.py``: rank 0 scans, everyone else consumes the broadcast. +# --------------------------------------------------------------------------- + + +class CategorySearchComm: + """Single-process stand-in for COMM_WORLD recording the collective order.""" + + def __init__(self, rank=0, size=1, bcast_returns=None): + self.rank = rank + self.size = size + self.calls = [] + self.bcast_payloads = [] + self._bcast_returns = list(bcast_returns or []) + + def Get_rank(self): + return self.rank + + def Get_size(self): + return self.size + + def Barrier(self): + self.calls.append("Barrier") + + def bcast(self, obj, root=0): + self.calls.append("bcast") + self.bcast_payloads.append(obj) + if self.rank == root: + return obj + return self._bcast_returns.pop(0) + + def gather(self, obj, root=0): + self.calls.append("gather") + return [obj] if self.rank == root else None + + def reduce(self, value, op=None, root=0): + self.calls.append("reduce") + return value if self.rank == root else None + + +class FakeMPI: + """Namespace mimicking ``mpi4py.MPI`` for one ``CategorySearchComm``.""" + + def __init__(self, comm): + import mpi4py.MPI as real_mpi + + self.COMM_WORLD = comm + self.SUM = real_mpi.SUM + + +def _cs_config(fract_base: Path) -> Namespace: + return Namespace( + fract_base_dir=str(fract_base), + n_categories=1, + seed=42, + variance_threshold=0.15, + point_num=60, + normalize=1, + datagen_from_scratch=False, + datagen_batch_size=4, + verbose=0, + ) + + +def _seed_one_category(config: Namespace) -> None: + """Write the single category CSV this config asks for.""" + param_dir = Path(layout.category_param_dir(config)) + param_dir.mkdir(parents=True, exist_ok=True) + params = np.zeros((2, 13), dtype=np.float64) + params[:, 0] = params[:, 4] = params[:, 8] = 0.5 + params[1, 9] = params[1, 10] = params[1, 11] = 0.5 + params[0, 12] = 0.5 + np.savetxt(param_dir / "000000.csv", params, delimiter=",") + + +def test_work_scan_is_root_only_and_broadcast(tmp_path, monkeypatch): + """A non-root rank never scans; it consumes root's broadcast index list.""" + config = _cs_config(tmp_path / "fractals") + comm = CategorySearchComm(rank=1, size=2, bcast_returns=[[0]]) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + def no_scan(*_args, **_kwargs): + raise AssertionError("non-root rank must not scan the filesystem") + + monkeypatch.setattr(cs, "parse_category_indices", no_scan) + + cs.main(config) + + # Root said category 0 already exists, so this rank has nothing to do and + # went straight to the post-loop reductions. + assert comm.calls == ["Barrier", "bcast", "reduce", "reduce", "reduce", "reduce"] + # It contributed nothing to the scan broadcast (it is not the scanner). + assert comm.bcast_payloads[0] is None + + +def test_divergent_fs_views_take_the_same_collective_path(tmp_path, monkeypatch): + """Ranks disagreeing about the directory still issue identical collectives. + + Rank 0 sees the finished category; rank 1's own view is empty. Before the + fix rank 1 entered the work loop (``bcast``) while rank 0 was already past + it (``reduce``). With the scan broadcast, rank 1's view is irrelevant. + """ + # Rank 0: the category is on disk, so its scan finds it. + root_config = _cs_config(tmp_path / "root_view") + _seed_one_category(root_config) + root_comm = CategorySearchComm(rank=0, size=2) + monkeypatch.setattr(cs, "MPI", FakeMPI(root_comm)) + cs.main(root_config) + + # Rank 1: an empty directory (a divergent view), but root broadcast [0]. + peer_config = _cs_config(tmp_path / "peer_view") + peer_comm = CategorySearchComm(rank=1, size=2, bcast_returns=[[0]]) + monkeypatch.setattr(cs, "MPI", FakeMPI(peer_comm)) + cs.main(peer_config) + + assert root_comm.bcast_payloads[0] == [0] + assert peer_comm.calls == root_comm.calls, ( + "ranks with divergent filesystem views issued different collectives: " + f"rank 0 {root_comm.calls} vs rank 1 {peer_comm.calls}" + ) From 26bd51239b632c0558259132bd9e89f84114d689 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:35:32 -0700 Subject: [PATCH 15/54] Write category parameter CSVs atomically A category CSV truncated by a killed job kept its six-digit name, so the resume scan counted it as done forever while instance generation and the search's own resume both died parsing it. Categories are now staged under a temp name that no scan matches, fsynced, and renamed into place. R31 --- ScaFFold/datagen/category_search.py | 31 ++++++++++++- tests/datagen/test_category_search.py | 64 +++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 1 deletion(-) diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py index 98e2342..af53a9c 100644 --- a/ScaFFold/datagen/category_search.py +++ b/ScaFFold/datagen/category_search.py @@ -335,12 +335,41 @@ def save_valid_category( target = os.path.join(fracts_write_dir, "%06d.csv" % idx) if os.path.exists(target): raise FileExistsError(f"Refusing to overwrite existing category file: {target}") - np.savetxt(target, params, delimiter=",") + _savetxt_atomic(target, params) existing_indices.append(idx) existing_params.append(params) return idx +def _savetxt_atomic(target: str, params: np.array) -> None: + """Write one category's parameters to ``target`` atomically. + + A category CSV truncated by a killed job is poison: the six-digit name is + all the resume scan looks at, so the category counts as done forever, while + every consumer (instance generation, and the search's own resume) dies + parsing it. The file is therefore written to a temp name in the same + directory -- one that neither the resume glob (``NNNNNN.csv``) nor the + instance loader's ``*.csv`` filter can match -- flushed, fsynced, and only + then ``os.replace``d onto the final name. + """ + directory, name = os.path.split(target) + tmp_path = os.path.join(directory, f".{name}.tmp{os.getpid()}") + try: + with open(tmp_path, "w") as handle: + np.savetxt(handle, params, delimiter=",") + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp_path, target) + except BaseException: + # A failed write must leave nothing behind: no temp file, and no + # partial file under the name resume would accept. + try: + os.remove(tmp_path) + except OSError: + pass + raise + + def _attempt_state_path(fracts_write_dir: str, rank: int) -> str: return os.path.join(fracts_write_dir, f".rng_attempt_rank{rank}") diff --git a/tests/datagen/test_category_search.py b/tests/datagen/test_category_search.py index 90236a8..2cc43b1 100644 --- a/tests/datagen/test_category_search.py +++ b/tests/datagen/test_category_search.py @@ -18,6 +18,7 @@ from pathlib import Path import numpy as np +import pytest from ScaFFold.datagen import category_search as cs from ScaFFold.datagen import layout @@ -190,3 +191,66 @@ def test_divergent_fs_views_take_the_same_collective_path(tmp_path, monkeypatch) "ranks with divergent filesystem views issued different collectives: " f"rank 0 {root_comm.calls} vs rank 1 {peer_comm.calls}" ) + + +# --------------------------------------------------------------------------- +# R31: category CSVs appear complete or not at all. +# +# A category file truncated by a killed job is still counted as "done" by the +# resume scan, so nothing ever regenerates it: instance generation then dies +# parsing it, and the category search's own resume dies re-loading it. The +# pipeline cannot self-heal -- the file has to be deleted by hand. +# --------------------------------------------------------------------------- + + +def _params() -> np.ndarray: + params = np.zeros((2, 13), dtype=np.float64) + params[:, 0] = params[:, 4] = params[:, 8] = 0.5 + params[0, 12] = 0.5 + return params + + +def test_category_csv_write_is_atomic(tmp_path, monkeypatch): + """A killed mid-write leaves no category file under a name resume accepts.""" + param_dir = tmp_path / "3DIFS_param" + param_dir.mkdir() + + # One complete category, saved normally. + indices, saved = [], [] + first = _params() + assert cs.save_valid_category(str(param_dir), first, indices, saved) == 0 + assert np.loadtxt(param_dir / "000000.csv", delimiter=",").shape == (2, 13) + + # The next save is interrupted after some bytes have been written. + observed = {} + + def partial_then_raise(fname, arr, *args, **kwargs): + observed["listing"] = sorted(p.name for p in param_dir.iterdir()) + handle = fname if hasattr(fname, "write") else open(fname, "w") + handle.write("0.5,0.5,0.5\n") + handle.flush() + if handle is not fname: + handle.close() + raise OSError("simulated SIGKILL mid-write") + + monkeypatch.setattr(cs.np, "savetxt", partial_then_raise) + + second = _params() + second[0, 0] = 0.25 + with pytest.raises(OSError): + cs.save_valid_category(str(param_dir), second, indices, saved) + + # No truncated category is visible: the resume scan still sees exactly the + # one complete category, and every file it names parses. + assert cs.parse_category_indices(str(param_dir)) == [0] + assert not (param_dir / "000001.csv").exists() + for idx in cs.parse_category_indices(str(param_dir)): + assert np.loadtxt(param_dir / f"{idx:06d}.csv", delimiter=",").shape == (2, 13) + + # Mid-write, the partial data lived under a name neither the resume scan + # nor the instance loader (which takes every ``*.csv``) would pick up. + partial_names = [n for n in observed["listing"] if n != "000000.csv"] + assert all(not name.endswith(".csv") for name in partial_names), partial_names + + # And nothing was left behind afterwards. + assert sorted(p.name for p in param_dir.iterdir()) == ["000000.csv"] From 9b028fa856d8f27b85d67c8b448b83b636f6ac02 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:39:26 -0700 Subject: [PATCH 16/54] Correct the half-voxel shift in voxelization The centering offset subtracted an extra voxel, biasing every cloud toward the origin: the first half-voxel of each filled axis floored to -1 and was clipped onto plane 0 (1.5x the interior density, with the far plane at 0.5x). DATASET_FORMAT_VERSION is bumped so misregistered datasets are regenerated rather than reused. R33 --- ScaFFold/datagen/get_dataset.py | 14 +++-- ScaFFold/datagen/volumegen.py | 10 +++- tests/datagen/test_artifacts.py | 97 ++++++++++++++++++++++++++++++--- 3 files changed, 105 insertions(+), 16 deletions(-) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index 0b193f2..f76f208 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -44,11 +44,15 @@ # Bumped from 2 to 3 when instance point clouds moved from float64 to float32: # the storage layout is unchanged, but float32 voxel binning shifts a handful of # boundary voxels, so a float64-era dataset must not be reused as if it were -# float32. This version stamps new datasets, gates reuse below, and feeds the -# config_id hash, so an older dataset is neither matched nor scanned. The loader -# in data_loading.py keeps its own (lower) minimum-layout version and still reads -# v3 through the modern dense path. -DATASET_FORMAT_VERSION = 3 +# float32. Bumped from 3 to 4 when the voxel centering offset was corrected from +# (grid_size - 1 - span)/2 to (grid_size - span)/2: every volume and mask +# generated before that was misregistered by half a voxel (with the first +# half-voxel of each axis clipped onto plane 0), so those datasets must be +# regenerated rather than reused. This version stamps new datasets, gates reuse +# below, and feeds the config_id hash, so an older dataset is neither matched nor +# scanned. The loader in data_loading.py keeps its own (lower) minimum-layout +# version and still reads v4 through the modern dense path. +DATASET_FORMAT_VERSION = 4 INCLUDE_KEYS = [ "dataset_format_version", "n_categories", diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py index 05567e0..34e17f9 100644 --- a/ScaFFold/datagen/volumegen.py +++ b/ScaFFold/datagen/volumegen.py @@ -82,9 +82,15 @@ def points_to_voxel_indices( scaled = (points - mins) / voxel_size # 4) Center the occupied region: the largest axis fills the grid while the - # shorter axes are offset so their span sits in the middle. + # shorter axes are offset so their span sits in the middle. The free + # space to split between the two margins is (grid_size - span) voxels, + # measured in the same voxel units as ``scaled``; subtracting an extra 1 + # (as if the offset were an index rather than a length) shifted every + # cloud half a voxel toward the origin, floored the first half-voxel of + # each filled axis to -1, and let the clip below fold those points onto + # plane 0. span = scaled.max(axis=0) - offset = (grid_size - 1 - span) / 2.0 + offset = (grid_size - span) / 2.0 idx = np.floor(scaled + offset).astype(int) # 5) Clip to valid range (guards float rounding at the boundaries). diff --git a/tests/datagen/test_artifacts.py b/tests/datagen/test_artifacts.py index 2f969a6..c1c38e3 100644 --- a/tests/datagen/test_artifacts.py +++ b/tests/datagen/test_artifacts.py @@ -312,20 +312,22 @@ def test_scale_config_rejected(): # --------------------------------------------------------------------------- -def _dense_reference_indices(points: np.ndarray, grid_size: int, eps=1e-6): - """The pre-refactor index computation, kept verbatim as a reference. - - This is the exact math the old dense ``points_to_voxelgrid`` ran before - scattering ``True`` into a full ``grid_size**3`` boolean array; the scatter - API must reproduce the identical occupied-voxel set and painted values. +def _reference_indices(points: np.ndarray, grid_size: int, eps=1e-6, *, clip=True): + """The index computation of ``points_to_voxel_indices``, spelled out. + + Two tests need to see inside the function: the scatter-vs-dense equivalence + check (which needs the per-point indices the dense grid was built from) and + the centering check (which needs the indices *before* ``np.clip`` hides + out-of-range bins). Every test using this asserts the replica reproduces the + real function's output, so it cannot silently drift from it. """ mins = points.min(axis=0) maxs = points.max(axis=0) voxel_size = (float((maxs - mins).max()) + eps) / grid_size scaled = (points - mins) / voxel_size - offset = (grid_size - 1 - scaled.max(axis=0)) / 2.0 + offset = (grid_size - scaled.max(axis=0)) / 2.0 idx = np.floor(scaled + offset).astype(int) - return np.clip(idx, 0, grid_size - 1) + return np.clip(idx, 0, grid_size - 1) if clip else idx def test_voxel_indices_match_dense_grid(): @@ -337,7 +339,7 @@ def test_voxel_indices_match_dense_grid(): idx = points_to_voxel_indices(points, grid_size) # Reference dense grid built the old way, from the reference index math. - ref_idx = _dense_reference_indices(points, grid_size) + ref_idx = _reference_indices(points, grid_size) reference = np.zeros((grid_size,) * 3, dtype=bool) reference[ref_idx[:, 0], ref_idx[:, 1], ref_idx[:, 2]] = True @@ -421,3 +423,80 @@ def test_dataset_version_bumped_past_float64_era(): from ScaFFold.datagen import get_dataset as gd assert gd.DATASET_FORMAT_VERSION > 2 + + +# --------------------------------------------------------------------------- +# R33: voxel centering is a whole voxel, not half of one +# --------------------------------------------------------------------------- + + +def _assert_replica_tracks_real(grid_size: int = 16) -> None: + """Pin ``_reference_indices`` to the real function on a sparse cloud. + + A sparse cloud is essential here: a dense one occupies every voxel under + any offset, so the comparison would pass vacuously. + """ + sparse = np.random.default_rng(7).random((300, 3)).astype(np.float32) + assert np.array_equal( + np.unique(_reference_indices(sparse, grid_size), axis=0), + points_to_voxel_indices(sparse, grid_size), + ), "the replicated index arithmetic no longer matches points_to_voxel_indices" + + +def test_voxelization_never_bins_outside_the_grid(): + """No point lands outside ``[0, grid_size)`` before the safety clip. + + The centering offset positions a span of ``span`` voxels inside a grid of + ``grid_size`` voxels, so the free space to split between the two margins is + ``grid_size - span``. Using ``grid_size - 1 - span`` shifted every cloud + half a voxel toward the origin: points in the first half-voxel of each + filled axis floored to -1, and ``np.clip`` quietly folded them into bin 0. + """ + grid_size = 16 + _assert_replica_tracks_real(grid_size) + rng = np.random.default_rng(1234) + points = rng.random((200_000, 3)).astype(np.float32) + + pre_clip = _reference_indices(points, grid_size, clip=False) + assert pre_clip.min() >= 0, ( + f"{int((pre_clip < 0).any(axis=1).sum())} of {len(points)} points floored " + "below bin 0 and were clipped back in" + ) + assert pre_clip.max() <= grid_size - 1 + + +def test_voxelization_density_is_uniform_at_the_boundaries(): + """A uniform cloud fills the boundary planes like the interior ones. + + The half-voxel shift piled the clipped points onto plane 0 (1.5x the + interior density) and starved the far plane (0.5x), a systematic + misregistration in every generated volume and mask. + """ + grid_size = 16 + _assert_replica_tracks_real(grid_size) + rng = np.random.default_rng(1234) + points = rng.random((200_000, 3)).astype(np.float32) + + idx = _reference_indices(points, grid_size) + counts = np.bincount(idx[:, 0], minlength=grid_size) + interior = counts[2:-2].mean() + assert 0.9 <= counts[0] / interior <= 1.1, ( + f"boundary plane 0 holds {counts[0] / interior:.2f}x the interior density" + ) + assert 0.9 <= counts[-1] / interior <= 1.1, ( + f"boundary plane {grid_size - 1} holds {counts[-1] / interior:.2f}x the " + "interior density" + ) + + +def test_dataset_version_bumped_past_half_voxel_era(): + """The reuse marker advanced past 3: pre-fix datasets are misregistered. + + Correcting the offset changes the voxel contents of every generated volume + and mask, so a dataset built before the fix must not be handed to a run + after it. Bumping the version both stops the reuse scan from matching those + directories and changes the config_id they hash to. + """ + from ScaFFold.datagen import get_dataset as gd + + assert gd.DATASET_FORMAT_VERSION > 3 From d52dd07a9a976ba11cfd20d1610df712cb4376ee Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:41:00 -0700 Subject: [PATCH 17/54] Bound the int16 mask carrier by the largest class id The guard compared the class count against the int16 limit, but v2 masks ship raw category ids and a sparse split lists only the categories it contains, so a two-entry table holding id 40000 passed and then wrapped negative. The bound is now the largest id the carrier will hold, still the remapped count for legacy datasets. R34 --- ScaFFold/utils/data_loading.py | 32 ++++++++--- tests/test_data_loading.py | 97 ++++++++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+), 6 deletions(-) diff --git a/ScaFFold/utils/data_loading.py b/ScaFFold/utils/data_loading.py index dcc1a45..8a75b1d 100644 --- a/ScaFFold/utils/data_loading.py +++ b/ScaFFold/utils/data_loading.py @@ -165,16 +165,36 @@ def __init__( customlog(f"Dataset format version: {self.dataset_format_version}") # Masks are handed off in a signed 16-bit carrier (widened to long on - # the compute device), so every class id must fit that range. Legacy - # masks are remapped to 0..len(mask_values)-1; optimized masks store - # dense ids that stay within the same bound. - max_class_id = len(self.mask_values) - 1 + # the compute device), so the largest class id the carrier will hold + # must fit that range. + max_class_id = self._max_class_id() if max_class_id > np.iinfo(np.int16).max: raise ValueError( - f"{len(self.mask_values)} classes exceed the int16 mask carrier " - f"limit ({np.iinfo(np.int16).max})" + f"Mask class id {max_class_id} (from {len(self.mask_values)} " + f"classes) exceeds the int16 mask carrier limit " + f"({np.iinfo(np.int16).max}); it would wrap negative" ) + def _max_class_id(self): + """Return the largest class id ``_to_mask_carrier`` will have to carry. + + The bound differs by format, and using the wrong one is unsafe in one + direction and needlessly strict in the other. v2+ masks ship *raw* + ``category + 1`` ids, and the per-split table lists only the categories + present in that split -- so a sparse split can declare two classes while + holding an id in the tens of thousands, which the class *count* check + happily waved through. Legacy masks, by contrast, are remapped to + ``0..len(mask_values)-1``, so the count is exactly right there and their + (arbitrarily large) raw values are irrelevant. + """ + if self.dataset_format_version < DATASET_FORMAT_VERSION: + return len(self.mask_values) - 1 + + ids = np.asarray(self.mask_values) + if ids.size == 0: + return 0 + return int(ids.max()) + def _load_mask_values(self, data_dir): """Return the label-remap table for this split. diff --git a/tests/test_data_loading.py b/tests/test_data_loading.py index b29d212..38746bb 100644 --- a/tests/test_data_loading.py +++ b/tests/test_data_loading.py @@ -636,6 +636,103 @@ def test_broken_meta_raises_instead_of_silent_legacy(tmp_path, broken_meta): assert str(root) in message +def _build_v2_sparse_label_dataset(root: Path, label: int) -> Path: + """A v2 dataset whose only foreground label is the (large) ``label``. + + v2 masks store raw ``category + 1`` ids and the per-split pickle lists only + the categories actually present in that split, so a sparse split can hold a + handful of very large ids. + """ + vol_dir = root / "volumes" / "training" + mask_dir = root / "masks" / "training" + vol_dir.mkdir(parents=True) + mask_dir.mkdir(parents=True) + np.save(vol_dir / "0.npy", np.zeros((3, 4, 4, 4), dtype=VOLUME_DTYPE)) + mask = np.zeros((4, 4, 4), dtype=MASK_DTYPE) + mask[0, 0, 0] = label + np.save(mask_dir / "0_mask.npy", mask) + with open(root / "train_unique_mask_vals", "wb") as handle: + pickle.dump({"mask_values": [0, label]}, handle) + (root / "meta.yaml").write_text("dataset_format_version: 2\n") + return root + + +# --------------------------------------------------------------------------- +# R34: the int16 carrier guard must bound the largest class *id*, not the count +# --------------------------------------------------------------------------- + + +def test_int16_guard_checks_the_largest_class_id(tmp_path): + """A v2 split holding an id above the int16 range is rejected. + + The guard compared ``len(mask_values) - 1`` -- the class *count* -- against + the carrier limit, but v2 masks ship raw ids. A split listing only + ``[0, 40000]`` passed a two-class check and then wrapped 40000 to -25536 in + the int16 carrier, which survives the downstream ``.long()`` cast as a + negative label. + """ + root = _build_v2_sparse_label_dataset(tmp_path / "sparse", label=40000) + + with pytest.raises(ValueError) as excinfo: + FractalDataset( + root / "volumes" / "training", + root / "masks" / "training", + data_dir=root / "train_unique_mask_vals", + ) + + message = str(excinfo.value) + # The message names the offending id and how many classes the split has. + assert "40000" in message + assert "2" in message + assert str(np.iinfo(np.int16).max) in message + + +def test_int16_guard_accepts_ids_inside_the_range(tmp_path): + """A large-but-representable id still loads, and does not wrap negative.""" + label = int(np.iinfo(np.int16).max) + root = _build_v2_sparse_label_dataset(tmp_path / "edge", label=label) + + ds = FractalDataset( + root / "volumes" / "training", + root / "masks" / "training", + data_dir=root / "train_unique_mask_vals", + ) + carrier = ds[0]["mask"] + assert carrier.dtype == torch.int16 + assert int(carrier.min()) >= 0 + assert int(carrier.max()) == label + + +def test_int16_guard_uses_remapped_ids_for_legacy_datasets(tmp_path): + """v1 raw values are remapped to 0..n-1, so huge raw values are fine. + + Guarding on ``max(mask_values)`` alone would reject a legacy dataset that + the loader handles perfectly well: its carrier only ever holds the remapped + index, not the raw voxel value. + """ + raw_mask = np.zeros((4, 4, 4), dtype=MASK_DTYPE) + raw_mask[0, 0, 0] = 40000 + volume = np.zeros((4, 4, 4, 3), dtype=VOLUME_DTYPE) + root = _build_v1_split_dataset( + tmp_path / "legacy", + raw_mask, + volume, + train_vals=[0, 40000], + val_vals=[0, 40000], + ) + + ds = FractalDataset( + root / "volumes" / "training", + root / "masks" / "training", + data_dir=root / "train_unique_mask_vals", + ) + assert ds.dataset_format_version == 1 + carrier = ds[0]["mask"] + # 40000 was remapped to class index 1; nothing wrapped. + assert int(carrier.max()) == 1 + assert int(carrier.min()) == 0 + + def test_absent_meta_is_still_legacy_v1(tiny_v1_dataset): """The genuine legacy case (no ``meta.yaml`` at all) is unchanged. From 07723e619d01cbcb8742bde74a0805f544494ed7 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:42:19 -0700 Subject: [PATCH 18/54] Read dataset provenance from the ScaFFold source tree _git_commit_short ran git in the process working directory, so meta.yaml, the published directory name, and the commit-based reuse gate carried whatever repo the job was launched from. It now runs git in the package directory; a non-checkout install still degrades to no-commit-id. R35 --- ScaFFold/datagen/get_dataset.py | 17 ++++- tests/datagen/test_provenance.py | 109 +++++++++++++++++++++++++++++++ 2 files changed, 125 insertions(+), 1 deletion(-) create mode 100644 tests/datagen/test_provenance.py diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index f76f208..09e4b33 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -99,11 +99,26 @@ def _hash_volume_config(volume_config: Dict[str, Any]) -> str: return hashlib.sha256(s).hexdigest()[:12] -def _git_commit_short(log) -> str: +def _git_commit_short(log, source_dir: Path | None = None) -> str: + """Return the short commit of the ScaFFold checkout, or ``"no-commit-id"``. + + The commit identifies *the code that generated a dataset*: it is stamped + into ``meta.yaml``, into the published directory name, and is what + ``dataset_reuse_enforce_commit_id`` compares against. It must therefore be + read from the ScaFFold source tree rather than from the process working + directory, which is wherever the job was launched (a site workflow repo, a + scratch directory, ...) and has nothing to do with this code. + + ``source_dir`` overrides the directory git runs in; it defaults to this + module's own location and exists so the non-checkout case can be tested. + """ + if source_dir is None: + source_dir = Path(__file__).resolve().parent try: return ( subprocess.check_output( ["git", "rev-parse", "--short", "HEAD"], + cwd=str(source_dir), stderr=subprocess.DEVNULL, # Don't show console output to user ) .decode() diff --git a/tests/datagen/test_provenance.py b/tests/datagen/test_provenance.py new file mode 100644 index 0000000..d53bd33 --- /dev/null +++ b/tests/datagen/test_provenance.py @@ -0,0 +1,109 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Dataset provenance: the commit stamped on a dataset is ScaFFold's (R35). + +``meta.yaml``'s ``code_commit``, the published ``__`` +directory name, and the ``dataset_reuse_enforce_commit_id`` gate all key off +one string. Reading it from the *launch* directory made it a property of +wherever the job happened to start -- a site workflow repo, a scratch +directory -- instead of the code that generated the data. Reuse was then gated +on an unrelated repo's churn while real ScaFFold changes went undetected. +""" + +from __future__ import annotations + +import logging +import subprocess +from pathlib import Path + +from ScaFFold.datagen import get_dataset as gd + +LOG = logging.getLogger("test_provenance") + +# The ScaFFold source tree: the checkout whose commit must be stamped. +PACKAGE_DIR = Path(gd.__file__).resolve().parent + + +def _head_of(repo: Path) -> str: + return ( + subprocess.check_output(["git", "rev-parse", "--short", "HEAD"], cwd=repo) + .decode() + .strip() + ) + + +def _make_repo(path: Path) -> str: + """Create a throwaway git repo with one commit; return its short HEAD.""" + path.mkdir(parents=True, exist_ok=True) + subprocess.run(["git", "init", "-q"], cwd=path, check=True) + (path / "README").write_text("an unrelated project\n") + subprocess.run(["git", "add", "README"], cwd=path, check=True) + subprocess.run( + [ + "git", + "-c", + "user.name=test", + "-c", + "user.email=test@example.invalid", + "-c", + "commit.gpgsign=false", + "commit", + "-qm", + "init", + ], + cwd=path, + check=True, + ) + return _head_of(path) + + +def test_commit_is_read_from_the_scaffold_tree_not_the_cwd(tmp_path, monkeypatch): + """Running from an unrelated repo still stamps ScaFFold's commit.""" + expected = _head_of(PACKAGE_DIR) + + other = tmp_path / "workflow-repo" + other_head = _make_repo(other) + assert other_head != expected, "the throwaway repo must differ from ScaFFold" + + monkeypatch.chdir(other) + assert gd._git_commit_short(LOG) == expected + + +def test_commit_survives_a_non_repo_working_directory(tmp_path, monkeypatch): + """A scratch launch directory does not degrade provenance to no-commit-id.""" + expected = _head_of(PACKAGE_DIR) + + scratch = tmp_path / "scratch-cwd" + scratch.mkdir() + monkeypatch.chdir(scratch) + + assert gd._git_commit_short(LOG) == expected + + +def test_non_repo_install_reports_no_commit_id(tmp_path): + """An installed (non-git) ScaFFold still degrades gracefully. + + Provenance is best-effort: when the source tree is not a checkout there is + no commit to record, and reuse simply is not gated on one. + """ + not_a_repo = tmp_path / "site-packages" / "ScaFFold" / "datagen" + not_a_repo.mkdir(parents=True) + + assert gd._git_commit_short(LOG, source_dir=not_a_repo) == "no-commit-id" + + +def test_missing_source_dir_reports_no_commit_id(tmp_path): + """A source directory that does not exist is handled, not raised.""" + assert gd._git_commit_short(LOG, source_dir=tmp_path / "gone") == "no-commit-id" From 2f72060dd4c98c94a4ed2331ea7b68e3eb332a2e Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:43:13 -0700 Subject: [PATCH 19/54] Document the uneven spatial shard hazard with an xfail test Nothing validates vol_size against dc_num_shards, so 16 over 3 shards is accepted as 6/6/4 and per-shard pooling silently diverges from the global result. The fix belongs in DistConv, so this records the hazard as a strict xfail rather than working around it. R32 --- tests/test_data_loading.py | 66 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/tests/test_data_loading.py b/tests/test_data_loading.py index 38746bb..1f60055 100644 --- a/tests/test_data_loading.py +++ b/tests/test_data_loading.py @@ -636,6 +636,72 @@ def test_broken_meta_raises_instead_of_silent_legacy(tmp_path, broken_meta): assert str(root) in message +# --------------------------------------------------------------------------- +# R32: uneven spatial shards are accepted but computed wrong (known, unfixed) +# --------------------------------------------------------------------------- + + +@pytest.mark.xfail( + reason="uneven spatial shards mishandled; to be fixed upstream in DistConv", + strict=True, +) +def test_uneven_spatial_shards_pool_like_the_unsharded_volume(): + """Per-shard pooling must agree with pooling the whole volume. + + ``SpatialShardSpec`` slices with ``torch.chunk`` semantics and only rejects + an *empty* shard, so ``vol_size=16`` over 3 shards is accepted as 6/6/4 -- + nothing anywhere validates ``vol_size % num_shards`` or the per-U-Net-level + evenness that pooling needs. DistConv's ``DCTensor`` intercepts only + ``aten.convolution``, so ``max_pool3d`` runs independently on each local + shard; with unequal (or odd) shards the local pooling windows stop lining up + with the global ones, and by the second level whole planes are dropped and + values appear that are the max of no global window at all. The same happens + for equal-but-odd shards (``vol_size=10`` over 2 shards -> 5/5), and an odd + local shard entering a strided conv trips DistConv's own divisibility check + with a cryptic error deep in the first forward. + + A volume ramp is used so a pooled value names the plane it came from. + + XFAIL: the fix belongs in DistConv (its sharded ops must handle uneven + spatial decompositions), not in the loader, which is why this documents the + hazard instead of asserting a workaround. ``strict=True``: the comparison is + plain deterministic CPU pooling, so it cannot pass by chance -- if it ever + passes, DistConv/ScaFFold has changed and this test must be revisited. + """ + vol_size, num_shards = 16, 3 + ramp = torch.zeros(1, 1, vol_size, vol_size, vol_size) + for plane in range(vol_size): + ramp[0, 0, plane] = plane + + # What the dataset hands each rank: uneven shards, accepted without a word. + volume = np.arange(vol_size**3, dtype=np.float32).reshape((vol_size,) * 3) + shard_sizes = [ + dl.SpatialShardSpec( + shard_dims=(2,), num_shards=(num_shards,), shard_indices=(index,) + ) + .slice_array(volume, {2: 0, 3: 1, 4: 2}, "mask") + .shape[0] + for index in range(num_shards) + ] + assert shard_sizes == [6, 6, 4] + + # Two U-Net levels of pooling, globally versus per local shard. + pool = torch.nn.MaxPool3d(2) + global_pooled = pool(pool(ramp)) + shards = list(torch.split(ramp, shard_sizes, dim=2)) + shards = [pool(pool(shard)) for shard in shards] + sharded_pooled = torch.cat(shards, dim=2) + + assert sharded_pooled.shape == global_pooled.shape, ( + f"sharded pooling produced {list(sharded_pooled.shape[2:])} planes vs " + f"{list(global_pooled.shape[2:])} globally" + ) + assert torch.equal(sharded_pooled, global_pooled), ( + f"per-shard planes {sharded_pooled[0, 0, :, 0, 0].tolist()} vs global " + f"{global_pooled[0, 0, :, 0, 0].tolist()}" + ) + + def _build_v2_sparse_label_dataset(root: Path, label: int) -> Path: """A v2 dataset whose only foreground label is the (large) ``label``. From 2d77083a64d3f1ec8eaa1f203329ddb76b524ac9 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:56:43 -0700 Subject: [PATCH 20/54] Abort when the MPI world does not span the job cli.py decides run dirs and restart state on MPI rank 0 while training uses the launcher's rank environment; a plain torchrun makes mpi4py a singleton in every process, so each claims its own run dir and the job hangs. Cross-check the two world sizes at the CLI entry and abort with an explanation of the launcher wiring. R13 --- ScaFFold/cli.py | 44 +++++++++++ tests/test_cli.py | 187 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 231 insertions(+) create mode 100644 tests/test_cli.py diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 858e6ba..01119c4 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -24,9 +24,49 @@ from ScaFFold.utils import config_utils from ScaFFold.utils.collect_scheduler_info import collect_scheduler_metadata from ScaFFold.utils.create_restart_script import create_restart_script +from ScaFFold.utils.distributed import get_world_size from ScaFFold.utils.utils import setup_mpi_logger +def check_launcher_world_size(mpi_world_size): + """Verify that the MPI world spans the whole job. + + ScaFFold uses two different sources of truth for the job shape: the CLI and + the benchmark driver make their job-wide decisions on ``MPI.COMM_WORLD`` + rank 0, while the training path takes its rank and world size from the + launcher's environment (``get_world_rank`` / ``get_world_size``). Those + agree only when the job was started by an MPI-aware launcher. + + Under a plain ``torchrun`` (or a bare ``python`` invocation of several + processes) mpi4py initializes as an independent singleton in every process, + so every process believes it is MPI rank 0: each one runs the rank-0 block, + atomically claims its *own* timestamped run directory, and the job then + diverges -- non-zero launcher ranks crash on a broadcast that never + happened while rank 0 blocks in the first collective until it times out. + + Fail loudly here, before any run directory is created, instead of leaving + that mess behind. ``get_world_size`` falls back to the MPI communicator + when the environment reports nothing, so an unlauncher-ed single process + trivially agrees with itself. + """ + env_world_size = get_world_size() + if env_world_size != mpi_world_size: + raise RuntimeError( + f"Launcher/MPI world size mismatch: MPI.COMM_WORLD reports " + f"{mpi_world_size} rank(s) but the launcher environment reports " + f"{env_world_size}. ScaFFold decides run directories and restart " + "state on MPI rank 0 and broadcasts them, so an MPI world that " + "does not span the job is unrecoverable: every process acts as " + "rank 0, each claims a separate run directory, and the job hangs " + "in the first collective. This is what a plain 'torchrun' (or " + "launching the processes directly) produces, because mpi4py then " + "initializes as a singleton in every process. Launch ScaFFold " + "with an MPI-aware launcher (torchrun-hpc, flux run, srun, " + "mpirun) so that MPI spans all ranks, or run a single process " + "with no launcher environment set." + ) + + def _make_fresh_run_dir(base_run_dir, timestamp): """Create a fresh timestamped run directory without clobbering an existing one. @@ -299,6 +339,10 @@ def main(): comm = MPI.COMM_WORLD rank = comm.Get_rank() + # Every rank runs this identically, before any run directory is created, + # so a mis-launched job aborts uniformly instead of leaving per-rank run + # dirs behind and hanging. + check_launcher_world_size(comm.Get_size()) # Parse the command-line arguments. args = parser.parse_args() log = setup_mpi_logger(__file__, args.verbose) diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..9691499 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,187 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Tests for the ``scaffold`` CLI entry point (``ScaFFold.cli.main``). + +``cli.main`` makes every job-wide decision on MPI rank 0 and broadcasts it, so +these tests drive it with a *fake* communicator: the real ``MPI.COMM_WORLD`` in +this environment is always a one-rank singleton, which cannot express the +multi-rank shapes the CLI must get right (rank-0-decides-and-broadcasts, and +the mismatch between the MPI world and the launcher's environment). + +``_FakeComm`` records what rank 0 broadcasts and, for a non-zero rank, replays +a scripted sequence of values as if rank 0 had sent them. That makes it +possible to assert that a non-zero rank *uses the broadcast decision* instead +of consulting its own view of the filesystem. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +import yaml + +import ScaFFold.cli as cli + +REPO_ROOT = Path(__file__).resolve().parents[1] +DEFAULT_CONFIG = REPO_ROOT / "ScaFFold" / "configs" / "benchmark_default.yml" + + +# --------------------------------------------------------------------------- +# harness +# --------------------------------------------------------------------------- + + +class _FakeComm: + """A stand-in for ``MPI.COMM_WORLD`` with a settable rank and size. + + ``bcast`` returns the caller's object on rank 0 (recording it); on a + non-zero rank it pops the next scripted value from ``bcast_returns``, + falling back to the caller's object when the script is exhausted. + """ + + def __init__(self, rank: int = 0, size: int = 1, bcast_returns=None): + self._rank = rank + self._size = size + self._scripted = list(bcast_returns or []) + self.broadcast = [] + self.barriers = 0 + + def Get_rank(self) -> int: + return self._rank + + def Get_size(self) -> int: + return self._size + + def Barrier(self) -> None: + self.barriers += 1 + + def bcast(self, obj, root=0): + self.broadcast.append(obj) + if self._rank == root: + return obj + if self._scripted: + return self._scripted.pop(0) + return obj + + +class _FakeMPI: + """Minimal ``mpi4py.MPI`` stand-in exposing only ``COMM_WORLD``.""" + + def __init__(self, comm): + self.COMM_WORLD = comm + + +def write_config(tmp_path, updates=None, name="bench.yml"): + """Write a complete benchmark config into ``tmp_path`` and return its path.""" + config = yaml.safe_load(DEFAULT_CONFIG.read_text()) + config["base_run_dir"] = str(tmp_path / "runs") + config["dataset_dir"] = str(tmp_path / "datasets") + config["fract_base_dir"] = str(tmp_path / "fractals") + if updates: + config.update(updates) + path = tmp_path / name + path.write_text(yaml.dump(config)) + return path + + +def run_cli(monkeypatch, argv, *, comm=None): + """Run ``cli.main`` with a fake communicator and stubbed subcommand drivers. + + Returns ``(comm, calls)`` where ``calls`` maps the subcommand name to the + list of config dicts its driver was invoked with. + """ + import ScaFFold.benchmark as benchmark_mod + import ScaFFold.generate_fractals as generate_fractals_mod + + comm = comm if comm is not None else _FakeComm() + calls = {"benchmark": [], "generate_fractals": []} + + monkeypatch.setattr(sys, "argv", list(argv)) + monkeypatch.setattr(cli, "MPI", _FakeMPI(comm)) + monkeypatch.setattr( + benchmark_mod, + "main", + lambda kwargs_dict={}: calls["benchmark"].append(dict(kwargs_dict)), + ) + monkeypatch.setattr( + generate_fractals_mod, + "main", + lambda kwargs_dict={}: calls["generate_fractals"].append(dict(kwargs_dict)), + ) + cli.main() + return comm, calls + + +# --------------------------------------------------------------------------- +# R13: the MPI world must span the whole job +# --------------------------------------------------------------------------- + + +def test_mpi_singleton_under_multirank_launcher_aborts(monkeypatch, tmp_path): + """A 1-rank MPI world inside a 2-rank launcher job aborts with an explanation. + + This is the plain-``torchrun`` shape: every process is an mpi4py singleton + while the launcher says WORLD_SIZE=2. Left unchecked, every process runs + cli.py's rank-0 block, claims its own run directory, and the job then hangs + in the first real collective. + """ + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + cfg = write_config(tmp_path) + + with pytest.raises(RuntimeError) as excinfo: + run_cli( + monkeypatch, + ["scaffold", "benchmark", "-c", str(cfg)], + comm=_FakeComm(rank=0, size=1), + ) + + message = str(excinfo.value) + assert "1" in message and "2" in message + assert "torchrun" in message.lower() + # The abort happens before any run directory is claimed. + assert not (tmp_path / "runs").exists() + + +def test_matching_world_sizes_are_accepted(monkeypatch, tmp_path): + """An MPI world that matches the launcher environment runs normally.""" + monkeypatch.setenv("WORLD_SIZE", "4") + monkeypatch.setenv("RANK", "0") + cfg = write_config(tmp_path) + + _, calls = run_cli( + monkeypatch, + ["scaffold", "benchmark", "-c", str(cfg)], + comm=_FakeComm(rank=0, size=4), + ) + + assert len(calls["benchmark"]) == 1 + + +def test_no_launcher_env_is_not_a_mismatch(monkeypatch, tmp_path): + """With no launcher variables set, the MPI world alone defines the size.""" + for var in ("WORLD_SIZE", "RANK", "LOCAL_RANK", "SLURM_NTASKS", "FLUX_JOB_SIZE"): + monkeypatch.delenv(var, raising=False) + cfg = write_config(tmp_path) + + _, calls = run_cli( + monkeypatch, + ["scaffold", "benchmark", "-c", str(cfg)], + comm=_FakeComm(rank=0, size=1), + ) + + assert len(calls["benchmark"]) == 1 From c1e7a3dac80c3a4a9651deb64e1ea4194cd4550e Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:57:24 -0700 Subject: [PATCH 21/54] Substitute restart placeholders in combined flag tokens A run launched with --config=PATH emitted a literal --config=__CFG__ into restart.sh, because placeholder substitution only matched whole tokens; the restart then died with "Config file '__CFG__' not found". Substitute the value half of any --flag=PLACEHOLDER token as well. R14 --- ScaFFold/utils/create_restart_script.py | 24 ++++++-- tests/test_restart_script.py | 75 +++++++++++++++++++++++++ 2 files changed, 93 insertions(+), 6 deletions(-) diff --git a/ScaFFold/utils/create_restart_script.py b/ScaFFold/utils/create_restart_script.py index b43c0e3..18d6a9f 100644 --- a/ScaFFold/utils/create_restart_script.py +++ b/ScaFFold/utils/create_restart_script.py @@ -78,14 +78,26 @@ def _rewrite_config_and_add_restart(cli_args: List[str]) -> List[str]: return new_args +def _substitute_placeholder(tok: str, var_subs: dict[str, str]) -> str: + """Return ``tok`` with a placeholder replaced by its Bash expansion. + + A placeholder may be a whole token (``--config __CFG__``) or the value half + of a combined token (``--config=__CFG__``); argparse accepts both spellings + on the command line, so the rewriter can emit either. Anything else is + shell-quoted verbatim. + """ + if tok in var_subs: + return var_subs[tok] # e.g., "$RUN_DIR/config.yaml" + flag, sep, value = tok.partition("=") + if sep and value in var_subs: + # --config=__CFG__ -> --config="$RUN_DIR/config.yaml" + return shlex.quote(flag + sep) + var_subs[value] + return shlex.quote(tok) + + def _bash_array(var_name: str, argv: List[str], var_subs: dict[str, str]) -> str: """Render a Bash array declaration VAR=( ... ), safely quoted, with simple placeholder substitution.""" - parts = [] - for tok in argv: - if tok in var_subs: - parts.append(var_subs[tok]) # e.g., "$RUN_DIR/config.yaml" - else: - parts.append(shlex.quote(tok)) + parts = [_substitute_placeholder(tok, var_subs) for tok in argv] return f"{var_name}=( " + " ".join(parts) + " )" diff --git a/tests/test_restart_script.py b/tests/test_restart_script.py index 1059c7e..8e9f81c 100644 --- a/tests/test_restart_script.py +++ b/tests/test_restart_script.py @@ -182,3 +182,78 @@ def test_generated_script_is_valid_bash(monkeypatch, tmp_path): assert result.returncode == 0, ( f"bash -n failed for variant {i}:\n{result.stderr}" ) + + +# --------------------------------------------------------------------------- +# R14: combined ``--flag=value`` tokens +# --------------------------------------------------------------------------- + + +def test_config_equals_form_is_substituted(monkeypatch, tmp_path): + """``--config=PATH`` is repointed at the run dir, not left as a placeholder. + + The rewriter emits the placeholder as part of a combined token, so a + substitution that only matches whole tokens leaves ``--config=__CFG__`` in + the script and the restart dies with "Config file '__CFG__' not found". + """ + _isolate_env(monkeypatch) + argv = [ + "/usr/bin/scaffold", + "benchmark", + "--config=/some/where/config.yml", + "--epochs", + "10", + ] + + script = _generate(monkeypatch, tmp_path / "run", argv=argv) + + assert "__CFG__" not in script + assert '--config="$RUN_DIR/config.yaml"' in script + assert "/some/where/config.yml" not in script + + +@pytest.mark.parametrize( + "config_argv", + [ + ["-c", "/some/where/config.yml"], + ["--config", "/some/where/config.yml"], + ["--config=/some/where/config.yml"], + ], + ids=["short", "long-space", "long-equals"], +) +@pytest.mark.skipif(shutil.which("bash") is None, reason="bash not available") +def test_every_config_spelling_expands_to_the_run_dir_config( + monkeypatch, tmp_path, config_argv +): + """Bash expands every ``--config`` spelling to ``$RUN_DIR/config.yaml``. + + The generated PY array is sourced and expanded by a real shell so the test + asserts on the arguments the restarted CLI actually receives. + """ + _isolate_env(monkeypatch) + argv = ["/usr/bin/scaffold", "benchmark"] + config_argv + + script = _generate(monkeypatch, tmp_path / "run", argv=argv) + + py_decl = next(line for line in script.splitlines() if line.startswith("PY=(")) + probe = tmp_path / f"probe_{config_argv[0][-1]}.sh" + probe.write_text(f'RUN_DIR=/run/dir\n{py_decl}\nprintf "%s\\n" "${{PY[@]}}"\n') + result = subprocess.run( + ["bash", str(probe)], capture_output=True, text=True, check=True + ) + + tokens = result.stdout.split("\n") + assert "/run/dir/config.yaml" in tokens or ( + "--config=/run/dir/config.yaml" in tokens + ) + assert not any("__CFG__" in tok for tok in tokens) + + +def test_run_dir_placeholder_is_substituted(monkeypatch, tmp_path): + """The appended ``--run-dir`` placeholder still resolves (control).""" + _isolate_env(monkeypatch) + + script = _generate(monkeypatch, tmp_path / "run") + + assert "__RUN_DIR__" not in script + assert '--run-dir "$RUN_DIR"' in script From 8f9487cafe95ebf06465c2d701046cfbcb8d336c Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 15:58:53 -0700 Subject: [PATCH 22/54] Generate restart scripts at the true job scale The CLI holds MPI.COMM_WORLD but did not pass its size, and the generator's environment sniffing missed launcher variables the rank side honors (MV2, PALS, and bare SLURM/FLUX task counts), so e.g. a Cray PALS job got a single-process restart.sh. Pass the communicator size and share one world-size variable list with get_world_size. R17 --- ScaFFold/cli.py | 6 ++-- ScaFFold/utils/create_restart_script.py | 24 +++++++++++-- tests/test_cli.py | 33 +++++++++++++++++ tests/test_restart_script.py | 48 +++++++++++++++++++++++++ 4 files changed, 106 insertions(+), 5 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 01119c4..f3ea545 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -441,8 +441,10 @@ def main(): with open(benchmark_run_dir / "config.yaml", "w") as file: yaml.dump(combined_config, file) - # 4. Generate/Update the restart script in the directory - create_restart_script(benchmark_run_dir) + # 4. Generate/Update the restart script in the directory. The + # communicator size is ground truth for the job scale; environment + # sniffing is only the fallback for callers that lack it. + create_restart_script(benchmark_run_dir, world_size=comm.Get_size()) comm.Barrier() combined_config = comm.bcast(combined_config, root=0) diff --git a/ScaFFold/utils/create_restart_script.py b/ScaFFold/utils/create_restart_script.py index 18d6a9f..3f93ade 100644 --- a/ScaFFold/utils/create_restart_script.py +++ b/ScaFFold/utils/create_restart_script.py @@ -27,6 +27,21 @@ # were active in the generating run. Names mirror ScaFFold.utils.perf_measure. _PROFILING_ENV_VARS = ("PROFILE_TORCH", "CALI_CONFIG") +# Launcher variables carrying the total rank count, in the same priority order +# as ScaFFold.utils.distributed.get_world_size. The rank side and the restart +# generator must recognize the same set, or a job launched under a launcher +# only one of them knows about (e.g. Cray PALS) gets a restart script for the +# wrong number of ranks. +_WORLD_SIZE_ENV_VARS = ( + "WORLD_SIZE", + "MV2_COMM_WORLD_SIZE", + "OMPI_COMM_WORLD_SIZE", + "PMI_SIZE", + "PALS_NRANKS", + "SLURM_NTASKS", + "FLUX_JOB_SIZE", +) + def _rewrite_config_and_add_restart(cli_args: List[str]) -> List[str]: """ @@ -243,8 +258,11 @@ def _sniff_launch_shape(env: Mapping[str, str]) -> tuple[int | None, int, int]: Reads, in priority order: 1. Flux (FLUX_JOB_SIZE is total tasks, FLUX_JOB_NNODES is node count), 2. Slurm (SLURM_NTASKS / SLURM_NPROCS total tasks, SLURM_*NODES nodes), - 3. generic launcher hints for total rank count: torchrun's WORLD_SIZE, - Open MPI's OMPI_COMM_WORLD_SIZE, and PMI's PMI_SIZE. + 3. generic launcher hints for the total rank count, in the same order + and covering the same variables as + ``ScaFFold.utils.distributed.get_world_size``: keeping the two in + sync is what stops a restart script from relaunching the job at the + wrong scale. ``nodes`` is None when the environment does not report a node count. ``world_size`` is the best available total-rank estimate (>= 1). @@ -260,7 +278,7 @@ def _sniff_launch_shape(env: Mapping[str, str]) -> tuple[int | None, int, int]: total_tasks = int(env.get("SLURM_NTASKS") or env.get("SLURM_NPROCS") or 1) else: # No scheduler: fall back to generic launcher hints for the rank count. - for key in ("WORLD_SIZE", "OMPI_COMM_WORLD_SIZE", "PMI_SIZE"): + for key in _WORLD_SIZE_ENV_VARS: val = env.get(key) if val: total_tasks = int(val) diff --git a/tests/test_cli.py b/tests/test_cli.py index 9691499..ca47cdf 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -185,3 +185,36 @@ def test_no_launcher_env_is_not_a_mismatch(monkeypatch, tmp_path): ) assert len(calls["benchmark"]) == 1 + + +# --------------------------------------------------------------------------- +# R17: the restart script is generated at the true job scale +# --------------------------------------------------------------------------- + + +def test_restart_script_gets_the_mpi_world_size(monkeypatch, tmp_path): + """The CLI passes its communicator size to the restart-script generator. + + Without it the generator falls back to sniffing the environment, which + misses launcher variables the rank side honors (e.g. PALS_NRANKS) and + silently emits a single-process restart script for a multi-rank job. + """ + recorded = {} + + def _recorder(run_dir, world_size=None): + recorded["run_dir"] = run_dir + recorded["world_size"] = world_size + return Path(run_dir) / "restart.sh" + + monkeypatch.setattr(cli, "create_restart_script", _recorder) + monkeypatch.setenv("WORLD_SIZE", "4") + monkeypatch.setenv("RANK", "0") + cfg = write_config(tmp_path) + + run_cli( + monkeypatch, + ["scaffold", "benchmark", "-c", str(cfg)], + comm=_FakeComm(rank=0, size=4), + ) + + assert recorded["world_size"] == 4 diff --git a/tests/test_restart_script.py b/tests/test_restart_script.py index 8e9f81c..a8f3e08 100644 --- a/tests/test_restart_script.py +++ b/tests/test_restart_script.py @@ -25,6 +25,7 @@ from __future__ import annotations +import os import shutil import subprocess import sys @@ -45,8 +46,10 @@ "SLURM_JOB_NUM_NODES", "SLURM_NNODES", "WORLD_SIZE", + "MV2_COMM_WORLD_SIZE", "OMPI_COMM_WORLD_SIZE", "PMI_SIZE", + "PALS_NRANKS", ) # Profiling variables re-exported only when set in the generating run. @@ -257,3 +260,48 @@ def test_run_dir_placeholder_is_substituted(monkeypatch, tmp_path): assert "__RUN_DIR__" not in script assert '--run-dir "$RUN_DIR"' in script + + +# --------------------------------------------------------------------------- +# R17: launch-shape sniffing must match the rank side +# --------------------------------------------------------------------------- + +# Every variable ``ScaFFold.utils.distributed.get_world_size`` honors. The +# restart generator must derive the same world size from each of them, or a +# restart script silently relaunches the job at the wrong scale. +_WORLD_SIZE_ENV_VARS = ( + "WORLD_SIZE", + "MV2_COMM_WORLD_SIZE", + "OMPI_COMM_WORLD_SIZE", + "PMI_SIZE", + "PALS_NRANKS", + "SLURM_NTASKS", + "FLUX_JOB_SIZE", +) + + +@pytest.mark.parametrize("var", _WORLD_SIZE_ENV_VARS) +def test_sniffed_world_size_matches_rank_side(monkeypatch, var): + """The generator and ``get_world_size`` agree on every launcher variable.""" + from ScaFFold.utils.distributed import get_world_size + + _isolate_env(monkeypatch) + for other in _WORLD_SIZE_ENV_VARS: + monkeypatch.delenv(other, raising=False) + monkeypatch.setenv(var, "8") + + _, _, sniffed = crs._sniff_launch_shape(os.environ) + + assert sniffed == 8, f"{var} not recognized by the restart generator" + assert sniffed == get_world_size() + + +def test_pals_job_gets_a_multirank_restart_script(monkeypatch, tmp_path): + """A Cray PALS launch (PALS_NRANKS) emits the multi-rank template.""" + _isolate_env(monkeypatch) + monkeypatch.setenv("PALS_NRANKS", "8") + + script = _generate(monkeypatch, tmp_path / "run") + + assert "torchrun-hpc" in script + assert 'exec "${PY[@]}"' not in script From 7563905e2156bed06fed770bc9b8d27393202e60 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:00:57 -0700 Subject: [PATCH 23/54] Decide the restart pre-check on rank 0 and broadcast it Every rank stat-ed the shared filesystem for a checkpoint and raised on its own verdict, so a divergent view (stale attribute cache) either strands the peers in benchmark.py's timeout-less barrier or lets them run on after rank 0 aborted. Make it a rank-0 decision broadcast to all ranks, matching the other CLI decisions. R18 --- ScaFFold/cli.py | 48 ++++++++++++++++-------- tests/test_cli.py | 96 ++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 128 insertions(+), 16 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index f3ea545..4e354c4 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -90,6 +90,26 @@ def _make_fresh_run_dir(base_run_dir, timestamp): candidate = base_run_dir / f"{timestamp}-{suffix}" +def missing_checkpoint_error(combined_config): + """Return the "nothing to resume from" message, or None if a restart can run. + + Reports the first problem found rather than raising, so the caller can make + this a rank-0 decision and broadcast the verdict instead of letting every + rank stat the shared filesystem and possibly disagree. + """ + checkpoint_dir = Path(combined_config["run_dir"]) / combined_config.get( + "checkpoint_dir", "checkpoints" + ) + expected_checkpoints = ( + checkpoint_dir / "checkpoint_last.pth", + checkpoint_dir / "checkpoint_best.pth", + ) + if any(path.exists() for path in expected_checkpoints): + return None + expected = " or ".join(str(path) for path in expected_checkpoints) + return f"Restart requested but no checkpoint was found. Expected {expected}." + + def resolve_run_dir(args_dict, combined_config): """Decide the benchmark run directory and whether this launch resumes a run. @@ -449,23 +469,21 @@ def main(): comm.Barrier() combined_config = comm.bcast(combined_config, root=0) + # Restart pre-check. Like every other decision here it is made once, on + # rank 0, and broadcast: the check reads the filesystem, and ranks can see + # different views of a shared filesystem (stale NFS/Lustre attribute + # caches). A rank that decided for itself would either abort alone -- + # stranding its peers in benchmark.py's timeout-less barrier -- or keep + # running after rank 0 had already aborted. + restart_precheck_error = None if combined_config.get("restart", False): - run_dir = combined_config.get("run_dir") - if not run_dir: + if not combined_config.get("run_dir"): raise ValueError("--restart requires --run-dir") - - checkpoint_dir = Path(run_dir) / combined_config.get( - "checkpoint_dir", "checkpoints" - ) - expected_checkpoints = ( - checkpoint_dir / "checkpoint_last.pth", - checkpoint_dir / "checkpoint_best.pth", - ) - if not any(path.exists() for path in expected_checkpoints): - expected = " or ".join(str(path) for path in expected_checkpoints) - raise FileNotFoundError( - f"Restart requested but no checkpoint was found. Expected {expected}." - ) + if rank == 0: + restart_precheck_error = missing_checkpoint_error(combined_config) + restart_precheck_error = comm.bcast(restart_precheck_error, root=0) + if restart_precheck_error is not None: + raise FileNotFoundError(restart_precheck_error) if rank == 0: log.debug("combined_config = %s", combined_config) diff --git a/tests/test_cli.py b/tests/test_cli.py index ca47cdf..f5ddd56 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -98,9 +98,14 @@ def write_config(tmp_path, updates=None, name="bench.yml"): return path -def run_cli(monkeypatch, argv, *, comm=None): +def run_cli(monkeypatch, argv, *, comm=None, sync_env=True): """Run ``cli.main`` with a fake communicator and stubbed subcommand drivers. + ``sync_env`` makes the launcher environment agree with the fake + communicator, which is what a correctly launched job looks like; tests of + the mismatch check itself pass ``sync_env=False`` and set the environment + themselves. + Returns ``(comm, calls)`` where ``calls`` maps the subcommand name to the list of config dicts its driver was invoked with. """ @@ -110,6 +115,9 @@ def run_cli(monkeypatch, argv, *, comm=None): comm = comm if comm is not None else _FakeComm() calls = {"benchmark": [], "generate_fractals": []} + if sync_env: + monkeypatch.setenv("WORLD_SIZE", str(comm.Get_size())) + monkeypatch.setenv("RANK", str(comm.Get_rank())) monkeypatch.setattr(sys, "argv", list(argv)) monkeypatch.setattr(cli, "MPI", _FakeMPI(comm)) monkeypatch.setattr( @@ -148,6 +156,7 @@ def test_mpi_singleton_under_multirank_launcher_aborts(monkeypatch, tmp_path): monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)], comm=_FakeComm(rank=0, size=1), + sync_env=False, ) message = str(excinfo.value) @@ -182,6 +191,7 @@ def test_no_launcher_env_is_not_a_mismatch(monkeypatch, tmp_path): monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)], comm=_FakeComm(rank=0, size=1), + sync_env=False, ) assert len(calls["benchmark"]) == 1 @@ -218,3 +228,87 @@ def _recorder(run_dir, world_size=None): ) assert recorded["world_size"] == 4 + + +# --------------------------------------------------------------------------- +# R18: the restart pre-check is a rank-0 decision, broadcast to everyone +# --------------------------------------------------------------------------- + + +def _restart_argv(cfg, run_dir): + return [ + "scaffold", + "benchmark", + "-c", + str(cfg), + "--restart", + "--run-dir", + str(run_dir), + ] + + +def _make_checkpoint(run_dir): + ckpt_dir = run_dir / "checkpoints" + ckpt_dir.mkdir(parents=True, exist_ok=True) + (ckpt_dir / "checkpoint_last.pth").write_bytes(b"") + return ckpt_dir + + +def test_restart_without_checkpoint_is_rejected_on_rank0(monkeypatch, tmp_path): + """Rank 0 still rejects a restart with no checkpoint, naming the paths.""" + cfg = write_config(tmp_path) + run_dir = tmp_path / "prior" + run_dir.mkdir() + + with pytest.raises(FileNotFoundError) as excinfo: + run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=_FakeComm(rank=0)) + + assert "checkpoint_last.pth" in str(excinfo.value) + + +def test_restart_precheck_follows_the_broadcast_decision(monkeypatch, tmp_path): + """A non-zero rank trusts rank 0's verdict instead of stat-ing the FS itself. + + Simulates a stale attribute cache: rank 0 saw the checkpoint and broadcast + "go", while this rank's view of the shared filesystem shows nothing. A rank + that re-decides locally raises alone and strands its peers in the next + barrier. + """ + cfg = write_config(tmp_path) + run_dir = tmp_path / "prior" + run_dir.mkdir() # deliberately empty: this rank sees no checkpoint + rank0_config = { + "restart": True, + "run_dir": str(run_dir), + "checkpoint_dir": "checkpoints", + "verbose": 0, + } + + comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_config, None]) + _, calls = run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) + + assert len(calls["benchmark"]) == 1 + + +def test_restart_precheck_failure_raises_on_every_rank(monkeypatch, tmp_path): + """Rank 0's rejection is broadcast, so non-zero ranks raise too. + + Here the local filesystem view *does* show a checkpoint; the rank must + still fail, because rank 0 -- the only rank whose verdict counts -- did not + find one. Otherwise the job splits: rank 0 aborts and the rest run on. + """ + cfg = write_config(tmp_path) + run_dir = tmp_path / "prior" + run_dir.mkdir() + _make_checkpoint(run_dir) + rank0_config = { + "restart": True, + "run_dir": str(run_dir), + "checkpoint_dir": "checkpoints", + "verbose": 0, + } + rank0_error = "Restart requested but no checkpoint was found. Expected /nope." + + comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_config, rank0_error]) + with pytest.raises(FileNotFoundError, match="no checkpoint"): + run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) From e4aded529b1d0b378789c14a3528b56c1764da8d Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:01:59 -0700 Subject: [PATCH 24/54] Scope the run directory to the benchmark subcommand generate_fractals ran the CLI's rank-0 block too, littering base_run_dir with a timestamped benchmark directory whose restart.sh replayed generate_fractals with --restart/--run-dir, flags that subparser rejects (exit 2). Create the run dir, config dumps and restart script only for benchmark. R19 --- ScaFFold/cli.py | 54 ++++++++++++++++++++++++++--------------------- tests/test_cli.py | 51 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 24 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 4e354c4..3ed0a20 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -441,30 +441,36 @@ def main(): combined_config["vol_size"] = pow(2, combined_config["problem_scale"]) combined_config["point_num"] = int(combined_config["vol_size"] ** 3 / 256) - # Resolve the run directory and whether this launch resumes a run. - # This sets combined_config["benchmark_run_dir"] on every path and, - # when resuming, forces train_from_scratch off / restart on. - benchmark_run_dir, restarting = resolve_run_dir(vars(args), combined_config) - if restarting: - log.info("Resuming in existing directory: %s", benchmark_run_dir) - - # Add scheduler metadata and machine name to config.yaml - combined_config["scheduler_metadata"] = collect_scheduler_metadata() - combined_config["machine_name"] = socket.gethostname() - - # Dump configs (Overwrite is okay/desired on restart to capture new job IDs) - overrides = { - k: v for k, v in cli_args.items() if v is not None and k != "command" - } - with open(benchmark_run_dir / "overrides.yaml", "w") as file: - yaml.dump(overrides, file) - with open(benchmark_run_dir / "config.yaml", "w") as file: - yaml.dump(combined_config, file) - - # 4. Generate/Update the restart script in the directory. The - # communicator size is ground truth for the job scale; environment - # sniffing is only the fallback for callers that lack it. - create_restart_script(benchmark_run_dir, world_size=comm.Get_size()) + # The run directory, its config dumps and its restart script belong to + # the benchmark subcommand alone. Fractal generation writes nothing + # there, and the restart script it used to get replayed + # `generate_fractals --restart --run-dir ...` -- flags that subparser + # rejects, so the script could only ever exit 2. + if args.command == "benchmark": + # Resolve the run directory and whether this launch resumes a run. + # This sets combined_config["benchmark_run_dir"] on every path and, + # when resuming, forces train_from_scratch off / restart on. + benchmark_run_dir, restarting = resolve_run_dir(vars(args), combined_config) + if restarting: + log.info("Resuming in existing directory: %s", benchmark_run_dir) + + # Add scheduler metadata and machine name to config.yaml + combined_config["scheduler_metadata"] = collect_scheduler_metadata() + combined_config["machine_name"] = socket.gethostname() + + # Dump configs (Overwrite is okay/desired on restart to capture new job IDs) + overrides = { + k: v for k, v in cli_args.items() if v is not None and k != "command" + } + with open(benchmark_run_dir / "overrides.yaml", "w") as file: + yaml.dump(overrides, file) + with open(benchmark_run_dir / "config.yaml", "w") as file: + yaml.dump(combined_config, file) + + # 4. Generate/Update the restart script in the directory. The + # communicator size is ground truth for the job scale; environment + # sniffing is only the fallback for callers that lack it. + create_restart_script(benchmark_run_dir, world_size=comm.Get_size()) comm.Barrier() combined_config = comm.bcast(combined_config, root=0) diff --git a/tests/test_cli.py b/tests/test_cli.py index f5ddd56..f379aaa 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -312,3 +312,54 @@ def test_restart_precheck_failure_raises_on_every_rank(monkeypatch, tmp_path): comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_config, rank0_error]) with pytest.raises(FileNotFoundError, match="no checkpoint"): run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) + + +# --------------------------------------------------------------------------- +# R19: generate_fractals is not a benchmark run +# --------------------------------------------------------------------------- + + +def test_generate_fractals_creates_no_benchmark_run_dir(monkeypatch, tmp_path): + """Fractal generation leaves no benchmark run dir and no restart script. + + The rank-0 block used to run for every subcommand, so a generation job + littered base_run_dir with a timestamped benchmark directory holding a + restart.sh that replays ``generate_fractals --restart --run-dir ...`` -- + flags the generate_fractals subparser rejects, so the script exits 2. + """ + cfg = write_config(tmp_path) + + _, calls = run_cli(monkeypatch, ["scaffold", "generate_fractals", "-c", str(cfg)]) + + assert len(calls["generate_fractals"]) == 1 + assert not (tmp_path / "runs").exists(), "generation created a benchmark run dir" + assert list(tmp_path.rglob("restart.sh")) == [] + assert list(tmp_path.rglob("overrides.yaml")) == [] + + +def test_generate_fractals_config_reaches_the_driver(monkeypatch, tmp_path): + """The merged config still reaches the generation driver (control).""" + cfg = write_config(tmp_path) + + _, calls = run_cli( + monkeypatch, + ["scaffold", "generate_fractals", "-c", str(cfg), "--n-categories", "3"], + ) + + (config,) = calls["generate_fractals"] + assert config["n_categories"] == 3 + assert config["fract_base_dir"] == str(tmp_path / "fractals") + + +def test_benchmark_still_creates_its_run_dir(monkeypatch, tmp_path): + """The benchmark subcommand keeps its run dir, dumps and restart script.""" + cfg = write_config(tmp_path) + + _, calls = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)]) + + (config,) = calls["benchmark"] + run_dir = Path(config["benchmark_run_dir"]) + assert run_dir.is_dir() + assert (run_dir / "config.yaml").exists() + assert (run_dir / "overrides.yaml").exists() + assert (run_dir / "restart.sh").exists() From 13acd12398c3e7cdc80131ccb3fece82977d3d40 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:03:43 -0700 Subject: [PATCH 25/54] Honor auxiliary config keys set in YAML datagen_batch_size and verbose are accepted config keys, but Config never stores them, so the merge silently replaced them with argparse defaults. Carry the file's values into the merged config and let only options actually given on the command line override them: CLI flag > config file > argparse default. R20 --- ScaFFold/cli.py | 47 ++++++++++++++++++++++++++++++-- tests/test_cli.py | 68 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 2 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 3ed0a20..c951d02 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -90,6 +90,31 @@ def _make_fresh_run_dir(base_run_dir, timestamp): candidate = base_run_dir / f"{timestamp}-{suffix}" +def explicit_cli_keys(args, parsers): + """Return the names of the options actually given on the command line. + + argparse does not record which options were supplied, so a value counts as + explicit when it differs from the default of the first parser in + ``parsers`` that defines one (subcommand parser first, then the top-level + parser). Only these may outrank a config-file setting; everything else in + the namespace is an argparse default, which is the weakest source. + + The one ambiguity is a flag passed with exactly its default value: it looks + absent, so a config-file entry wins over it. Both spellings then agree on + the default, which is the only value the flag could have contributed. + """ + explicit = set() + for name, value in vars(args).items(): + default = None + for parser in parsers: + default = parser.get_default(name) + if default is not None: + break + if value != default: + explicit.add(name) + return explicit + + def missing_checkpoint_error(combined_config): """Return the "nothing to resume from" message, or None if a restart can run. @@ -365,6 +390,11 @@ def main(): check_launcher_world_size(comm.Get_size()) # Parse the command-line arguments. args = parser.parse_args() + subcommand_parsers = { + "benchmark": benchmark_parser, + "generate_fractals": generate_fractals_parser, + } + active_parser = subcommand_parsers[args.command] log = setup_mpi_logger(__file__, args.verbose) combined_config = None @@ -398,12 +428,23 @@ def main(): # into the run dir); keep the base config there. cli_args["config"] = config_paths[0] - # Combine configs: CLI args override config file values + # Combine configs, in increasing order of precedence: + # argparse default < config file < explicit command-line flag. combined_config = bench_config_dict.copy() + # Config only keeps the keys it consumes; the auxiliary keys it accepts + # (verbose, datagen_batch_size, ...) never become attributes, so put + # the file's values back first. Without this they are absent below and + # the argparse default overwrites what the user wrote in the config. + for key, value in merged_dict.items(): + combined_config.setdefault(key, value) + + explicit_cli = explicit_cli_keys(args, (active_parser, parser)) for key, value in cli_args.items(): + if key == "command": + continue if key not in combined_config: combined_config[key] = value - elif value is not None and key != "command": + elif key in explicit_cli and value is not None: log.info( "Overriding '%s=%s' with '%s=%s'", key, @@ -412,6 +453,8 @@ def main(): value, ) combined_config[key] = value + # The subcommand is always owned by the command line. + combined_config["command"] = cli_args["command"] # Recalculate unet_layers to capture any CLI overrides combined_config["unet_layers"] = ( diff --git a/tests/test_cli.py b/tests/test_cli.py index f379aaa..88782f5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -363,3 +363,71 @@ def test_benchmark_still_creates_its_run_dir(monkeypatch, tmp_path): assert (run_dir / "config.yaml").exists() assert (run_dir / "overrides.yaml").exists() assert (run_dir / "restart.sh").exists() + + +# --------------------------------------------------------------------------- +# R20: auxiliary keys set in YAML must survive; CLI > YAML > argparse default +# --------------------------------------------------------------------------- + + +def test_yaml_aux_keys_reach_the_driver(monkeypatch, tmp_path): + """Auxiliary keys set in the config file are not replaced by defaults. + + ``Config`` accepts ``datagen_batch_size``/``verbose`` but does not store + them, so they used to vanish from the merged config and the argparse + default was installed instead -- making them settable only on the command + line despite being documented, validated config keys. + """ + cfg = write_config(tmp_path, {"datagen_batch_size": 500, "verbose": 1}) + + _, calls = run_cli(monkeypatch, ["scaffold", "generate_fractals", "-c", str(cfg)]) + + (config,) = calls["generate_fractals"] + assert config["datagen_batch_size"] == 500 + assert config["verbose"] == 1 + + +def test_cli_flag_outranks_yaml_aux_key(monkeypatch, tmp_path): + """An explicit command-line flag still wins over the config file.""" + cfg = write_config(tmp_path, {"datagen_batch_size": 500, "verbose": 0}) + + _, calls = run_cli( + monkeypatch, + [ + "scaffold", + "-v", + "generate_fractals", + "-c", + str(cfg), + "--datagen-batch-size", + "250", + ], + ) + + (config,) = calls["generate_fractals"] + assert config["datagen_batch_size"] == 250 + assert config["verbose"] == 1 + + +def test_argparse_default_used_when_yaml_is_silent(monkeypatch, tmp_path): + """With neither a flag nor a config entry, the argparse default applies.""" + cfg = write_config(tmp_path) + + _, calls = run_cli(monkeypatch, ["scaffold", "generate_fractals", "-c", str(cfg)]) + + (config,) = calls["generate_fractals"] + assert config["datagen_batch_size"] == 10000 + assert config["verbose"] == 0 + + +def test_run_config_records_the_effective_aux_values(monkeypatch, tmp_path): + """The run dir's config.yaml records what the run actually used.""" + cfg = write_config(tmp_path, {"verbose": 1}) + + _, calls = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)]) + + (config,) = calls["benchmark"] + dumped = yaml.safe_load( + (Path(config["benchmark_run_dir"]) / "config.yaml").read_text() + ) + assert dumped["verbose"] == 1 From f494ec8b7bf71bab3e78989c974984062092f157 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:05:05 -0700 Subject: [PATCH 26/54] Keep memory diagnostics from crashing CPU-only runs mem_stats called torch.cuda.current_device() unguarded, so gather_and_print_mem -- which BaseTrainer.__init__ invokes unconditionally -- killed any CPU/gloo run launched with -v. Report the missing device and log a fallback instead; the early return is uniform across ranks, so no collective is skipped on one rank only. R21 --- ScaFFold/utils/utils.py | 23 +++++++++++++++++- tests/conftest.py | 8 +++--- tests/test_infra.py | 54 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 81 insertions(+), 4 deletions(-) diff --git a/ScaFFold/utils/utils.py b/ScaFFold/utils/utils.py index 5d9e334..12c196b 100644 --- a/ScaFFold/utils/utils.py +++ b/ScaFFold/utils/utils.py @@ -113,12 +113,24 @@ def setup_mpi_logger( def mem_stats(): + """Return this rank's GPU memory counters. + + On a host with no visible GPU (a CPU/gloo run, or a job launched with the + devices masked off) there are no counters to read: report that instead of + raising, so a diagnostic call cannot take down a run that is otherwise + perfectly able to proceed. + """ + rank = dist.get_rank() if dist.is_initialized() else 0 + if not torch.cuda.is_available(): + return {"rank": rank, "device": "cpu", "cuda_available": False} + dev = torch.cuda.current_device() free, total = torch.cuda.mem_get_info() # device-level (driver) view stats = torch.cuda.memory_stats(dev) # allocator internals return { - "rank": dist.get_rank() if dist.is_initialized() else 0, + "rank": rank, "device": dev, + "cuda_available": True, "allocated": torch.cuda.memory_allocated( dev ), # bytes currently used by tensors @@ -135,6 +147,15 @@ def mem_stats(): def gather_and_print_mem(log, tag=""): if log.getEffectiveLevel() > 10: # 10 -> DEBUG return + if not torch.cuda.is_available(): + # Uniform across ranks, so returning here skips the all_gather on every + # rank rather than desynchronizing them. + log.debug( + "=== %s === no CUDA device visible on this rank; " + "GPU memory statistics unavailable", + tag, + ) + return stats = mem_stats() if dist.is_initialized(): world = dist.get_world_size() diff --git a/tests/conftest.py b/tests/conftest.py index 376afbc..e35f739 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -467,6 +467,7 @@ def make( n_train: int = 4, n_val: int = 2, n: int = 16, + log_level: int = logging.INFO, config_overrides: Optional[dict] = None, ) -> "PyTorchTrainer": dataset_root = tiny_dataset( @@ -495,9 +496,10 @@ def make( device = torch.device("cpu") log = logging.getLogger(f"tiny_trainer.{id(config)}") - # INFO (20) > DEBUG (10) => gather_and_print_mem short-circuits and - # never touches CUDA / torch.distributed. - log.setLevel(logging.INFO) + # At the INFO default (20 > DEBUG's 10) gather_and_print_mem + # short-circuits and never touches CUDA / torch.distributed; pass + # log_level=logging.DEBUG to exercise the memory diagnostics. + log.setLevel(log_level) return PyTorchTrainer(model, config, device, log) diff --git a/tests/test_infra.py b/tests/test_infra.py index bf0fbeb..c279f34 100644 --- a/tests/test_infra.py +++ b/tests/test_infra.py @@ -22,6 +22,7 @@ from __future__ import annotations +import logging import os import numpy as np @@ -29,6 +30,7 @@ from ScaFFold.utils.data_loading import FractalDataset from ScaFFold.utils.data_types import MASK_DTYPE, VOLUME_DTYPE +from ScaFFold.utils.utils import gather_and_print_mem, mem_stats from tests.helpers import mpi_runner # --------------------------------------------------------------------------- @@ -230,3 +232,55 @@ def test_torchrun_gloo_two_ranks(tmp_path): # Both ranks reported; all_reduce of ranks {0,1} sums to 1. assert "RANK 0/2 sum=1.0" in out assert "RANK 1/2 sum=1.0" in out + + +# --------------------------------------------------------------------------- +# R21: memory diagnostics on a CPU-only run +# --------------------------------------------------------------------------- + + +def _debug_logger(name): + log = logging.getLogger(name) + log.setLevel(logging.DEBUG) + return log + + +def test_mem_stats_without_cuda(caplog): + """``mem_stats`` reports "no GPU" instead of raising on a CPU-only host.""" + if torch.cuda.is_available(): + import pytest + + pytest.skip("test covers the CPU-only path") + + stats = mem_stats() + + assert stats["cuda_available"] is False + assert "rank" in stats + + +def test_gather_and_print_mem_without_cuda(caplog): + """A DEBUG-level CPU run logs a fallback instead of crashing. + + ``BaseTrainer.__init__`` calls this unconditionally, so a CPU/gloo run with + ``-v`` used to die in trainer construction with "No CUDA GPUs are + available". + """ + if torch.cuda.is_available(): + import pytest + + pytest.skip("test covers the CPU-only path") + + log = _debug_logger("test_gather_and_print_mem_without_cuda") + with caplog.at_level(logging.DEBUG, logger=log.name): + gather_and_print_mem(log, "after_trainer_setup") + + messages = " ".join(record.getMessage() for record in caplog.records) + assert "after_trainer_setup" in messages + assert "cuda" in messages.lower() or "gpu" in messages.lower() + + +def test_trainer_constructs_with_debug_logging(tiny_trainer): + """The real call site survives: a trainer builds with a DEBUG logger.""" + trainer = tiny_trainer(log_level=logging.DEBUG) + + assert trainer.log.getEffectiveLevel() == logging.DEBUG From ffe7c0b49b6027d41a3fcff77316ffe2f4325b39 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:06:03 -0700 Subject: [PATCH 27/54] Preserve the base config under its own name in the run dir Copying the base config into the run dir kept its original filename, so a base config named config.yaml overwrote the merged config.yaml the CLI had just written -- and restart.sh points -c at that file. Always copy it to base_config.yaml. R22 --- README.md | 2 +- ScaFFold/benchmark.py | 9 +++++-- tests/test_config.py | 57 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 3206b02..86acb7b 100644 --- a/README.md +++ b/README.md @@ -61,7 +61,7 @@ The model is trained from a random initialization until convergence, which is de ScaFFold benchmark training always uses PyTorch distributed execution with DistConv spatial parallelism. For a singleton run, launch one distributed rank rather than disabling distributed execution. -Each `benchmark` invocation performs exactly one benchmark run, in a run folder created under `base_run_dir` set in the config file. Every run parameter must be single-valued; a list (e.g. `problem_scale: [6, 7]`) is rejected by name, since parameter sweeps are not supported. To compare parameter settings, launch one benchmark run per setting. For reproducibility, the run folder holds a copy of the benchmark config yml plus the fully merged `config.yaml` for that run. +Each `benchmark` invocation performs exactly one benchmark run, in a run folder created under `base_run_dir` set in the config file. Every run parameter must be single-valued; a list (e.g. `problem_scale: [6, 7]`) is rejected by name, since parameter sweeps are not supported. To compare parameter settings, launch one benchmark run per setting. For reproducibility, the run folder holds a copy of the benchmark config yml as `base_config.yaml` plus the fully merged `config.yaml` for that run. After the run completes, statistics from the run are stored in `train_stats.csv`. Additionally, users can inspect plots of the training and validation losses over time in ` Date: Fri, 31 Jul 2026 16:07:18 -0700 Subject: [PATCH 28/54] Validate the U-Net bottleneck against the problem scale Config accepted any unet_bottleneck_dim, so a value outside 0..problem_scale-1 built a U-Net with zero or too many pooling levels and failed much later with an opaque max_pool3d size error naming no config key. Reject it at config time, in Config and again after the CLI applies overrides. R25 --- ScaFFold/cli.py | 6 +++++- ScaFFold/utils/config_utils.py | 38 ++++++++++++++++++++++++++++++++- tests/test_cli.py | 33 ++++++++++++++++++++++++++++ tests/test_config.py | 39 ++++++++++++++++++++++++++++++++++ 4 files changed, 114 insertions(+), 2 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index c951d02..623f868 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -456,7 +456,11 @@ def main(): # The subcommand is always owned by the command line. combined_config["command"] = cli_args["command"] - # Recalculate unet_layers to capture any CLI overrides + # Recalculate unet_layers to capture any CLI overrides. The overridden + # pair has to be re-validated: Config only saw the config-file values. + config_utils.validate_unet_dims( + combined_config["problem_scale"], combined_config["unet_bottleneck_dim"] + ) combined_config["unet_layers"] = ( combined_config["problem_scale"] - combined_config["unet_bottleneck_dim"] ) diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 3289b42..92af20c 100644 --- a/ScaFFold/utils/config_utils.py +++ b/ScaFFold/utils/config_utils.py @@ -27,6 +27,40 @@ def require_positive_int(name: str, value: int) -> int: return value +def validate_unet_dims(problem_scale, unet_bottleneck_dim) -> int: + """Check that ``problem_scale``/``unet_bottleneck_dim`` describe a real U-Net. + + The U-Net has ``unet_layers = problem_scale - unet_bottleneck_dim`` + down/up levels over a ``2**problem_scale`` volume, so the bottleneck + exponent must satisfy ``0 <= unet_bottleneck_dim <= problem_scale - 1``: + a larger value asks for a bottleneck no smaller than the input (zero or + negative layers) and a negative one asks for more pooling levels than the + volume has. Both are only discovered later as an opaque + ``max_pool3d`` size error -- in production, after the whole dataset has + been generated -- so reject them here, at config time, naming the two keys + that have to change. + + Returns the validated bottleneck dimension. + """ + if isinstance(unet_bottleneck_dim, bool) or not isinstance( + unet_bottleneck_dim, int + ): + raise ValueError( + f"unet_bottleneck_dim must be an integer; got {unet_bottleneck_dim!r}" + ) + unet_layers = problem_scale - unet_bottleneck_dim + if unet_bottleneck_dim < 0 or unet_layers < 1: + raise ValueError( + f"unet_bottleneck_dim={unet_bottleneck_dim} is out of range for " + f"problem_scale={problem_scale}: it must satisfy " + f"0 <= unet_bottleneck_dim <= problem_scale - 1 " + f"(i.e. <= {problem_scale - 1}) so that the U-Net has at least one " + f"layer, but unet_layers = problem_scale - unet_bottleneck_dim = " + f"{unet_layers}. Raise problem_scale or lower unet_bottleneck_dim." + ) + return unet_bottleneck_dim + + class Config: """ A class for storing configuration settings for a specific run. @@ -182,7 +216,9 @@ def __init__(self, config_dict, strict=True): "WARNING: problem_scale found to be non-integer. Truncating to nearest int." ) self.problem_scale = math.floor(self.problem_scale) - self.unet_bottleneck_dim = config_dict["unet_bottleneck_dim"] + self.unet_bottleneck_dim = validate_unet_dims( + self.problem_scale, config_dict["unet_bottleneck_dim"] + ) self.unet_layers = self.problem_scale - self.unet_bottleneck_dim self.n_fracts_per_vol = config_dict["n_fracts_per_vol"] self.n_instances_used_per_fractal = config_dict["n_instances_used_per_fractal"] diff --git a/tests/test_cli.py b/tests/test_cli.py index 88782f5..b215e44 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -431,3 +431,36 @@ def test_run_config_records_the_effective_aux_values(monkeypatch, tmp_path): (Path(config["benchmark_run_dir"]) / "config.yaml").read_text() ) assert dumped["verbose"] == 1 + + +# --------------------------------------------------------------------------- +# R25: an out-of-range bottleneck is rejected before any work starts +# --------------------------------------------------------------------------- + + +def test_cli_override_bottleneck_out_of_range_rejected(monkeypatch, tmp_path): + """A command-line override that empties the U-Net is caught at config time. + + The CLI recomputes unet_layers after applying overrides, so the check has + to run there too -- not only inside Config. + """ + cfg = write_config(tmp_path) + + with pytest.raises(ValueError) as excinfo: + run_cli( + monkeypatch, + [ + "scaffold", + "benchmark", + "-c", + str(cfg), + "--problem-scale", + "4", + "--unet-bottleneck-dim", + "4", + ], + ) + + message = str(excinfo.value) + assert "unet_bottleneck_dim" in message + assert "problem_scale" in message diff --git a/tests/test_config.py b/tests/test_config.py index 0c8976b..28cc15b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -329,3 +329,42 @@ def test_base_config_copy_never_clobbers_merged_config( preserved = yaml.safe_load((run_dir / "base_config.yaml").read_text()) assert preserved["local_batch_size"] == BASE["local_batch_size"] assert "machine_name" not in preserved + + +# --------------------------------------------------------------------------- +# unet_bottleneck_dim range (R25) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "problem_scale, bottleneck", + [(5, -1), (5, 5), (5, 6)], + ids=["negative", "zero-layers", "negative-layers"], +) +def test_bottleneck_out_of_range_rejected(problem_scale, bottleneck): + """An out-of-range bottleneck fails at config time, naming both keys. + + Left unvalidated it produced a U-Net with more pooling levels than the + volume has, and the run died hours later inside max_pool3d with "Given + input size: (2048x1x1x1)" -- naming no config key at all. + """ + bad = {**BASE, "problem_scale": problem_scale, "unet_bottleneck_dim": bottleneck} + + with pytest.raises(ValueError) as excinfo: + config_utils.Config(bad) + + message = str(excinfo.value) + assert "unet_bottleneck_dim" in message + assert "problem_scale" in message + assert str(bottleneck) in message + assert str(problem_scale) in message + + +@pytest.mark.parametrize("bottleneck", [0, 3, 4]) +def test_bottleneck_in_range_accepted(bottleneck): + """The full valid range (at least one U-Net layer) is accepted.""" + cfg = config_utils.Config( + {**BASE, "problem_scale": 5, "unet_bottleneck_dim": bottleneck} + ) + assert cfg.unet_layers == 5 - bottleneck + assert cfg.unet_layers >= 1 From 2ac96437f1d0b6e6289635ff7b56d018b9bdcf48 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:08:57 -0700 Subject: [PATCH 29/54] Never let a trace export strand the other ranks export_chrome_trace runs before the barrier that precedes rank-0 post-processing and raises when the run had zero profiled steps (or the write fails), killing the profiling rank and blocking every other rank in that barrier until timeout. Move it into a helper that logs the failure and returns. R15 --- ScaFFold/worker.py | 40 +++++++++++++++++++--- tests/test_reporting.py | 76 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 4 deletions(-) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 459a8cf..8eafa51 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -96,6 +96,41 @@ def wrap_model_ddp(model, device, ps): ) +def export_profiler_trace(prof, config, log, rank, world_size, ranks_per_node): + """Write this rank's chrome trace, reporting failures instead of raising. + + A profiling rank reaches this while the others are already heading for the + barrier that precedes rank-0 post-processing, so an exception here does not + just lose a trace: it kills this rank and leaves every other rank blocked + in that barrier until the collective times out. Failures are real (a run + with zero training batches never starts the profiler, and export_chrome_trace + then raises; the trace can also fill the filesystem), so log them and let + the job finish. + + Returns the path written, or None if the trace could not be written. + """ + tracename = ( + f"torch-{socket.gethostname()}-r{rank}" + f"-N{world_size // ranks_per_node}-n{world_size}" + f"-ps{config.problem_scale}-e{config.epochs}" + f"-nipf{config.n_instances_used_per_fractal}-{int(time.time())}.json" + ) + try: + prof.export_chrome_trace(tracename) + except Exception as e: + log.error( + "Could not write PyTorch trace '%s': %s: %s. Continuing so the " + "run can finish; a run with zero profiled steps never starts the " + "profiler and has no trace to export.", + tracename, + type(e).__name__, + e, + ) + return None + log.info("Wrote PyTorch trace '%s'", tracename) + return tracename + + @annotate() def main(kwargs_dict: dict = {}): # @@ -261,10 +296,7 @@ def main(kwargs_dict: dict = {}): trainer.train(profiler=prof if TORCH_PERF_LOCAL else None) end_code_region("train") if TORCH_PERF_LOCAL: - hostname = socket.gethostname() - tracename = f"torch-{hostname}-r{rank}-N{world_size // ranks_per_node}-n{world_size}-ps{config.problem_scale}-e{config.epochs}-nipf{config.n_instances_used_per_fractal}-{int(time.time())}.json" - prof.export_chrome_trace(tracename) - log.info("Wrote PyTorch trace '%s'", tracename) + export_profiler_trace(prof, config, log, rank, world_size, ranks_per_node) # Results are final here; synchronize before rank-0 post-processing so a # post-processing failure on rank 0 cannot strand the other ranks in a diff --git a/tests/test_reporting.py b/tests/test_reporting.py index a6f0106..f77fb19 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -270,3 +270,79 @@ def test_torch_profiler_independent_of_caliper(self, monkeypatch): ) finally: importlib.reload(perf_measure) + + +class TestProfilerTraceExport: + """R15: a failed trace export must not strand the other ranks.""" + + @staticmethod + def _unstepped_profiler(): + """A profiler whose window never opened (a run with zero batches).""" + from torch.profiler import ProfilerActivity, profile, schedule + + prof = profile( + activities=[ProfilerActivity.CPU], + schedule=schedule(wait=1, warmup=1, active=3, repeat=1), + ) + with prof: + pass # no prof.step(): the schedule never leaves its wait phase + return prof + + @staticmethod + def _config(run_dir): + return SimpleNamespace( + problem_scale=4, + epochs=1, + n_instances_used_per_fractal=2, + run_dir=str(run_dir), + ) + + def test_zero_step_export_is_reported_not_raised(self, tmp_path, caplog): + """Exporting an unstepped profiler logs an error instead of raising. + + The export runs before the ``dist.barrier()`` that precedes rank-0 + post-processing, so a raise here kills the profiling rank and leaves + every other rank blocked in that barrier until the collective timeout. + """ + import logging + + import ScaFFold.worker as worker + + log = logging.getLogger("test_zero_step_export") + with caplog.at_level(logging.DEBUG, logger=log.name): + result = worker.export_profiler_trace( + self._unstepped_profiler(), + self._config(tmp_path), + log, + rank=0, + world_size=1, + ranks_per_node=1, + ) + + assert result is None + messages = " ".join(record.getMessage() for record in caplog.records) + assert "trace" in messages.lower() + + def test_successful_export_writes_a_trace(self, tmp_path, caplog): + """A profiler with a completed window still writes its trace (control).""" + import logging + + from torch.profiler import ProfilerActivity, profile, schedule + + import ScaFFold.worker as worker + + prof = profile( + activities=[ProfilerActivity.CPU], + schedule=schedule(wait=1, warmup=1, active=1, repeat=1), + ) + with prof: + for _ in range(4): + prof.step() + + log = logging.getLogger("test_successful_export") + path = worker.export_profiler_trace( + prof, self._config(tmp_path), log, rank=0, world_size=1, ranks_per_node=1 + ) + + assert path is not None + assert Path(path).exists() From f722a9afdc7cc6e1fd97caf4b7af9d8c336c7592 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:11:11 -0700 Subject: [PATCH 30/54] Detect the local rank count and place traces in the run dir get_local_size ignored LOCAL_WORLD_SIZE/PMI_LOCAL_SIZE/PALS_LOCAL_SIZE that get_local_rank honors, so it returned 1 and the per-node profiler gate selected every rank while the trace name claimed one node per rank. Add the missing variables, round the node count up, and write the trace into config.run_dir instead of the working directory. R23 --- ScaFFold/utils/distributed.py | 13 +++++++- ScaFFold/worker.py | 19 +++++++++--- tests/test_reporting.py | 58 +++++++++++++++++++++++++++++++++++ tests/test_worker_dist.py | 50 ++++++++++++++++++++++++++++++ 4 files changed, 134 insertions(+), 6 deletions(-) diff --git a/ScaFFold/utils/distributed.py b/ScaFFold/utils/distributed.py index c46a5bd..afad5ab 100644 --- a/ScaFFold/utils/distributed.py +++ b/ScaFFold/utils/distributed.py @@ -62,11 +62,22 @@ def get_local_rank(required: bool = False) -> int: def get_local_size(required: bool = False) -> int: - """Return the number of local MPI ranks.""" + """Return the number of local MPI ranks. + + Recognizes the same launchers as ``get_local_rank``: a variable honored + there but not here silently yields 1, which makes per-node logic (e.g. the + profiler's one-rank-per-node gate) treat every rank as node-local. + """ + if "LOCAL_WORLD_SIZE" in os.environ: + return int(os.environ["LOCAL_WORLD_SIZE"]) if "MV2_COMM_WORLD_LOCAL_SIZE" in os.environ: return int(os.environ["MV2_COMM_WORLD_LOCAL_SIZE"]) if "OMPI_COMM_WORLD_LOCAL_SIZE" in os.environ: return int(os.environ["OMPI_COMM_WORLD_LOCAL_SIZE"]) + if "PMI_LOCAL_SIZE" in os.environ: + return int(os.environ["PMI_LOCAL_SIZE"]) + if "PALS_LOCAL_SIZE" in os.environ: + return int(os.environ["PALS_LOCAL_SIZE"]) if "SLURM_NNODES" in os.environ and "SLURM_NTASKS" in os.environ: return int(os.environ["SLURM_NTASKS"]) // int(os.environ["SLURM_NNODES"]) # Flux does not have an env variable for this, so we assume an diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 8eafa51..2f36b9a 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -17,6 +17,7 @@ import socket import time from argparse import Namespace +from pathlib import Path import numpy as np import psutil @@ -107,28 +108,36 @@ def export_profiler_trace(prof, config, log, rank, world_size, ranks_per_node): then raises; the trace can also fill the filesystem), so log them and let the job finish. + The trace is written into ``config.run_dir`` so it lands with the rest of + the run's artifacts instead of wherever the job happened to be launched + from. + Returns the path written, or None if the trace could not be written. """ + # Round up: with a partly-filled last node, flooring would report one node + # too few (and a ranks_per_node larger than the job would report none). + nodes = max(1, math.ceil(world_size / max(1, ranks_per_node))) tracename = ( f"torch-{socket.gethostname()}-r{rank}" - f"-N{world_size // ranks_per_node}-n{world_size}" + f"-N{nodes}-n{world_size}" f"-ps{config.problem_scale}-e{config.epochs}" f"-nipf{config.n_instances_used_per_fractal}-{int(time.time())}.json" ) + tracepath = Path(getattr(config, "run_dir", None) or os.getcwd()) / tracename try: - prof.export_chrome_trace(tracename) + prof.export_chrome_trace(str(tracepath)) except Exception as e: log.error( "Could not write PyTorch trace '%s': %s: %s. Continuing so the " "run can finish; a run with zero profiled steps never starts the " "profiler and has no trace to export.", - tracename, + tracepath, type(e).__name__, e, ) return None - log.info("Wrote PyTorch trace '%s'", tracename) - return tracename + log.info("Wrote PyTorch trace '%s'", tracepath) + return tracepath @annotate() diff --git a/tests/test_reporting.py b/tests/test_reporting.py index f77fb19..a4948c4 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -18,6 +18,7 @@ import matplotlib import numpy as np +import pytest matplotlib.use("Agg") @@ -288,6 +289,20 @@ def _unstepped_profiler(): pass # no prof.step(): the schedule never leaves its wait phase return prof + @staticmethod + def _stepped_profiler(): + """A profiler with a completed capture window.""" + from torch.profiler import ProfilerActivity, profile, schedule + + prof = profile( + activities=[ProfilerActivity.CPU], + schedule=schedule(wait=1, warmup=1, active=1, repeat=1), + ) + with prof: + for _ in range(4): + prof.step() + return prof + @staticmethod def _config(run_dir): return SimpleNamespace( @@ -346,3 +361,46 @@ def test_successful_export_writes_a_trace(self, tmp_path, caplog): assert path is not None assert Path(path).exists() + + def test_trace_lands_in_the_run_dir(self, tmp_path, caplog): + """R23: the trace goes to the run dir, not whatever CWD happens to be.""" + import logging + + import ScaFFold.worker as worker + + prof = self._stepped_profiler() + log = logging.getLogger("test_trace_lands_in_the_run_dir") + + path = worker.export_profiler_trace( + prof, self._config(tmp_path), log, rank=0, world_size=1, ranks_per_node=1 + ) + + assert Path(path).parent == tmp_path + assert list(tmp_path.glob("torch-*.json")) == [Path(path)] + + @pytest.mark.parametrize( + "world_size, ranks_per_node, expected", + [(8, 4, "-N2-n8-"), (6, 4, "-N2-n6-"), (1, 1, "-N1-n1-")], + ids=["even", "ragged-last-node", "singleton"], + ) + def test_trace_name_counts_nodes_not_ranks( + self, tmp_path, world_size, ranks_per_node, expected + ): + """R23: the N field is a node count, and never rounds a node away.""" + import logging + + import ScaFFold.worker as worker + + prof = self._stepped_profiler() + log = logging.getLogger("test_trace_name_counts_nodes") + + path = worker.export_profiler_trace( + prof, + self._config(tmp_path), + log, + rank=0, + world_size=world_size, + ranks_per_node=ranks_per_node, + ) + + assert expected in Path(path).name diff --git a/tests/test_worker_dist.py b/tests/test_worker_dist.py index 40be05b..b1ae524 100644 --- a/tests/test_worker_dist.py +++ b/tests/test_worker_dist.py @@ -226,3 +226,53 @@ def fake_train(self, profiler=None): assert trainer.config.global_batch_size == trainer.config.local_batch_size # The worker destroyed the process group before rank-0 post-processing. assert not torch.distributed.is_initialized() + + +# --------------------------------------------------------------------------- +# Local size detection (R23) +# --------------------------------------------------------------------------- + +_LOCAL_SIZE_CASES = [ + # torchrun exports LOCAL_WORLD_SIZE alongside LOCAL_RANK. + ({"LOCAL_WORLD_SIZE": "4"}, 4), + ({"MV2_COMM_WORLD_LOCAL_SIZE": "4"}, 4), + ({"OMPI_COMM_WORLD_LOCAL_SIZE": "4"}, 4), + ({"PMI_LOCAL_SIZE": "4"}, 4), + ({"PALS_LOCAL_SIZE": "4"}, 4), + ({"SLURM_NTASKS": "8", "SLURM_NNODES": "2"}, 4), + ({"FLUX_JOB_SIZE": "8", "FLUX_JOB_NNODES": "2"}, 4), +] + +_LOCAL_SIZE_VARS = [ + "LOCAL_WORLD_SIZE", + "MV2_COMM_WORLD_LOCAL_SIZE", + "OMPI_COMM_WORLD_LOCAL_SIZE", + "PMI_LOCAL_SIZE", + "PALS_LOCAL_SIZE", + "SLURM_NTASKS", + "SLURM_NNODES", + "FLUX_JOB_SIZE", + "FLUX_JOB_NNODES", +] + + +def test_local_size_detection_matrix(monkeypatch): + """Every launcher that reports a local rank has its local size honored too. + + An unrecognized variable silently yields 1, which makes the per-node + profiler gate ``rank % ranks_per_node == 0`` select *every* rank and + mislabels the trace's node count. + """ + for env, want_local_size in _LOCAL_SIZE_CASES: + for var in _LOCAL_SIZE_VARS: + monkeypatch.delenv(var, raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert distributed_mod.get_local_size() == want_local_size, env + + +def test_local_size_defaults_to_one(monkeypatch): + """With nothing to go on, one rank per node is still the assumption.""" + for var in _LOCAL_SIZE_VARS: + monkeypatch.delenv(var, raising=False) + assert distributed_mod.get_local_size() == 1 From 9e87457850562379a1309df9481c750f50ad6b8c Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:12:21 -0700 Subject: [PATCH 31/54] Parse PROFILE_TORCH like the other profiler flags The gate enabled profiling for any value except the literal "off", so PROFILE_TORCH=0, =false, =no and ="" all turned the profiler ON -- the opposite of what they say, and inconsistent with the sub-option flags in the same module. BEHAVIOR CHANGE: profiling is now enabled only by 1/true/on/yes (any case); every other value, including previously-enabling ones such as "enabled", leaves it off. R24 --- README.md | 2 +- ScaFFold/utils/perf_measure.py | 20 ++++++----- tests/test_reporting.py | 65 ++++++++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 86acb7b..2442494 100644 --- a/README.md +++ b/README.md @@ -172,7 +172,7 @@ For n  in n_volumes: #### 1. Profiling with the PyTorch Profiler -Set `PROFILE_TORCH=ON` to generate a PyTorch profiling trace that can be read into [Perfetto](https://ui.perfetto.dev/). +Set `PROFILE_TORCH=ON` to generate a PyTorch profiling trace that can be read into [Perfetto](https://ui.perfetto.dev/). The trace is written into the run directory. `1`, `true`, `on` and `yes` (any case) enable profiling; every other value, including `0`, `false`, `no` and `off`, leaves it disabled. #### 2. Profiling with Caliper & Adiak diff --git a/ScaFFold/utils/perf_measure.py b/ScaFFold/utils/perf_measure.py index c69b4b6..2f17cc9 100644 --- a/ScaFFold/utils/perf_measure.py +++ b/ScaFFold/utils/perf_measure.py @@ -18,6 +18,17 @@ CALI_PERF_ENV_VAR = "CALI_CONFIG" TORCH_PERF_ENV_VAR = "PROFILE_TORCH" + +def _profiler_env_flag(name): + """Return True only for an affirmative value of the environment variable. + + Every profiler toggle -- the master switch and its sub-options alike -- + goes through this, so "0"/"false"/"no"/"off"/"" all mean off and there is + no spelling that means the opposite of what it says. + """ + return os.environ.get(name, "").lower() in ("1", "true", "on", "yes") + + _CALI_PERF_ENABLED = False TORCH_PERF_ENABLED = False if CALI_PERF_ENV_VAR in os.environ: @@ -33,10 +44,7 @@ # The torch profiler is gated purely on its own environment variable: Caliper # and the torch profiler may both be enabled at once. -if ( - TORCH_PERF_ENV_VAR in os.environ - and os.environ.get(TORCH_PERF_ENV_VAR).lower() != "off" -): +if _profiler_env_flag(TORCH_PERF_ENV_VAR): try: from torch.profiler import ProfilerActivity from torch.profiler import profile as torchprofile @@ -100,10 +108,6 @@ def _profiler_env_int(name, default): return default -def _profiler_env_flag(name): - return os.environ.get(name, "").lower() in ("1", "true", "on", "yes") - - def get_torch_context(ranks_per_node, rank): if TORCH_PERF_ENABLED: TORCH_PERF_LOCAL = TORCH_PERF_ENABLED and (rank % ranks_per_node == 0) diff --git a/tests/test_reporting.py b/tests/test_reporting.py index a4948c4..e9cead1 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -404,3 +404,68 @@ def test_trace_name_counts_nodes_not_ranks( ) assert expected in Path(path).name + + +class TestProfileTorchGate: + """R24: PROFILE_TORCH is parsed like every other profiler flag.""" + + @staticmethod + def _reload_with(monkeypatch_context, value): + import importlib + + import ScaFFold.utils.perf_measure as perf_measure + + if value is None: + monkeypatch_context.delenv("PROFILE_TORCH", raising=False) + else: + monkeypatch_context.setenv("PROFILE_TORCH", value) + monkeypatch_context.delenv("CALI_CONFIG", raising=False) + importlib.reload(perf_measure) + return perf_measure + + @pytest.mark.parametrize("value", [None, "", "0", "false", "no", "off", "OFF"]) + def test_disabled_values(self, monkeypatch, value): + """Anything that is not an affirmative value leaves profiling off. + + ``PROFILE_TORCH=0`` used to *enable* the profiler: the gate only + rejected the literal "off", so every conventional way of saying "no" + silently turned profiling on. + """ + import importlib + + import ScaFFold.utils.perf_measure as perf_measure + + try: + with monkeypatch.context() as m: + assert not self._reload_with(m, value).TORCH_PERF_ENABLED + finally: + importlib.reload(perf_measure) + + @pytest.mark.parametrize("value", ["1", "true", "on", "ON", "yes", "TRUE"]) + def test_enabled_values(self, monkeypatch, value): + """The affirmative spellings still enable the profiler.""" + import importlib + + import ScaFFold.utils.perf_measure as perf_measure + + try: + with monkeypatch.context() as m: + assert self._reload_with(m, value).TORCH_PERF_ENABLED + finally: + importlib.reload(perf_measure) + + def test_gate_matches_the_sub_option_parser(self, monkeypatch): + """The master switch and the sub-option flags agree on every spelling.""" + import importlib + + import ScaFFold.utils.perf_measure as perf_measure + + try: + for value in ("1", "true", "on", "yes", "0", "false", "no", "off", ""): + with monkeypatch.context() as m: + module = self._reload_with(m, value) + assert module.TORCH_PERF_ENABLED == module._profiler_env_flag( + "PROFILE_TORCH" + ), value + finally: + importlib.reload(perf_measure) From b9a42aa2529559e818b521c61c57c057a06e6521 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:13:20 -0700 Subject: [PATCH 32/54] Require the profiler window to skip at least one step The profiler context wraps checkpoint cleanup and every warmup batch while prof.step() advances only per training batch, so PROFILE_TORCH_WAIT=0 buffered that entire prologue in host memory as one unbounded step. Clamp wait to 1 and say so, rather than reordering the context or stepping it from warmup. R26 --- ScaFFold/utils/perf_measure.py | 15 +++++++ tests/test_reporting.py | 80 ++++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/ScaFFold/utils/perf_measure.py b/ScaFFold/utils/perf_measure.py index 2f17cc9..7236e69 100644 --- a/ScaFFold/utils/perf_measure.py +++ b/ScaFFold/utils/perf_measure.py @@ -122,7 +122,22 @@ def get_torch_context(ranks_per_node, rank): # trace. The window (skip `wait`, prime `warmup`, capture `active`, once) # is tunable via the environment. Callers must drive it with # ``prof.step()`` once per training step for the schedule to advance. + # The context is entered around checkpoint cleanup and the warmup + # batches, but prof.step() only advances once per *training* batch, so + # everything before the first training batch lands in step 0. The + # window must therefore skip at least one step: with wait=0 that whole + # prologue -- warmup_batches forward+backward passes per rank -- is + # buffered in host memory as a single unbounded step, which is the very + # thing the bounded window exists to prevent. wait = _profiler_env_int("PROFILE_TORCH_WAIT", 1) + if wait < 1: + print( + "PROFILE_TORCH_WAIT must be at least 1: the profiler window " + "opens before the warmup batches, whose work would otherwise " + "accumulate in host memory as one unbounded step. Using " + "PROFILE_TORCH_WAIT=1." + ) + wait = 1 warmup = _profiler_env_int("PROFILE_TORCH_WARMUP", 1) active = _profiler_env_int("PROFILE_TORCH_ACTIVE", 3) or 1 diff --git a/tests/test_reporting.py b/tests/test_reporting.py index e9cead1..3b658fd 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -469,3 +469,83 @@ def test_gate_matches_the_sub_option_parser(self, monkeypatch): ), value finally: importlib.reload(perf_measure) + + +class TestProfilerSchedule: + """R26: the schedule must not record everything before the first step.""" + + @staticmethod + def _context_with(monkeypatch_context, env): + """Reload perf_measure with ``env`` applied and build a profiler context.""" + import importlib + + import ScaFFold.utils.perf_measure as perf_measure + + monkeypatch_context.setenv("PROFILE_TORCH", "1") + monkeypatch_context.delenv("CALI_CONFIG", raising=False) + for name in ("PROFILE_TORCH_WAIT", "PROFILE_TORCH_WARMUP"): + monkeypatch_context.delenv(name, raising=False) + for key, value in env.items(): + monkeypatch_context.setenv(key, value) + importlib.reload(perf_measure) + assert perf_measure.TORCH_PERF_ENABLED + ctx, is_local = perf_measure.get_torch_context(1, 0) + assert is_local + return ctx + + def test_wait_zero_does_not_record_step_zero(self, monkeypatch, capsys): + """PROFILE_TORCH_WAIT=0 is clamped so step 0 records nothing. + + worker.main enters the profiler context around checkpoint cleanup and + every warmup batch, and ``prof.step()`` only advances once per training + batch -- so a schedule that is already active at step 0 buffers all of + that as a single unbounded step, which is exactly what the bounded + window exists to prevent. + """ + import importlib + + from torch.profiler import ProfilerAction + + import ScaFFold.utils.perf_measure as perf_measure + + try: + with monkeypatch.context() as m: + ctx = self._context_with(m, {"PROFILE_TORCH_WAIT": "0"}) + assert ctx.schedule(0) == ProfilerAction.NONE + output = capsys.readouterr().out + assert "PROFILE_TORCH_WAIT" in output + finally: + importlib.reload(perf_measure) + + def test_default_schedule_skips_step_zero(self, monkeypatch): + """The default window already skips step 0 (control).""" + import importlib + + from torch.profiler import ProfilerAction + + import ScaFFold.utils.perf_measure as perf_measure + + try: + with monkeypatch.context() as m: + ctx = self._context_with(m, {}) + assert ctx.schedule(0) == ProfilerAction.NONE + finally: + importlib.reload(perf_measure) + + def test_larger_wait_is_preserved(self, monkeypatch): + """A wait longer than the minimum is left alone.""" + import importlib + + from torch.profiler import ProfilerAction + + import ScaFFold.utils.perf_measure as perf_measure + + try: + with monkeypatch.context() as m: + ctx = self._context_with( + m, {"PROFILE_TORCH_WAIT": "3", "PROFILE_TORCH_WARMUP": "1"} + ) + assert ctx.schedule(2) == ProfilerAction.NONE + assert ctx.schedule(3) == ProfilerAction.WARMUP + finally: + importlib.reload(perf_measure) From 6a0ae3a9c46f9efa896f5f5fd83fc7b405abe6de Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:16:44 -0700 Subject: [PATCH 33/54] Cover the config round-trip across a restart Regression guard for the merge order: a restart driven by the run dir's config.yaml (what restart.sh emits) must reproduce the first run's CLI overrides and auxiliary config keys and still take the resume path. R20 R22 --- tests/test_cli.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tests/test_cli.py b/tests/test_cli.py index b215e44..381744f 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -464,3 +464,36 @@ def test_cli_override_bottleneck_out_of_range_rejected(monkeypatch, tmp_path): message = str(excinfo.value) assert "unet_bottleneck_dim" in message assert "problem_scale" in message + + +# --------------------------------------------------------------------------- +# The whole config path survives a restart (R20/R22 together) +# --------------------------------------------------------------------------- + + +def test_config_round_trips_through_a_restart(monkeypatch, tmp_path): + """A restart driven by the run dir's config.yaml reproduces the run. + + This is exactly what the generated restart.sh does: ``-c + $RUN_DIR/config.yaml --restart --run-dir $RUN_DIR``. It exercises the merge + in both directions -- CLI overrides and auxiliary config keys have to come + back out of the dumped config, and the resume flags have to win. + """ + cfg = write_config(tmp_path, {"verbose": 1, "datagen_batch_size": 500}) + + _, calls = run_cli( + monkeypatch, + ["scaffold", "benchmark", "-c", str(cfg), "--local-batch-size", "2"], + ) + run_dir = Path(calls["benchmark"][0]["benchmark_run_dir"]) + _make_checkpoint(run_dir) + + _, resumed = run_cli(monkeypatch, _restart_argv(run_dir / "config.yaml", run_dir)) + + (config,) = resumed["benchmark"] + assert config["local_batch_size"] == 2 # CLI override from the first run + assert config["verbose"] == 1 # auxiliary key set in YAML + assert config["datagen_batch_size"] == 500 + assert config["restart"] is True + assert config["train_from_scratch"] is False + assert config["benchmark_run_dir"] == str(run_dir) From d6c7836eff30b62ebf45041ef4e824a0ff7a7d72 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:47:25 -0700 Subject: [PATCH 34/54] Fence the remaining rank-0 windows in checkpointing Rank 0 is the only rank that touches the run directory, and it does so while its peers are already committed to the next collective, so every one of those windows has to report failure *through* that collective. Three were left open. cleanup() ran _remove_checkpoint_files between the drain and the broadcast: individual unlinks were tolerated, but the glob and stat around them can raise on a shared filesystem (ESTALE, EACCES), re-creating the R05 hazard in a narrow window. load_from_checkpoint() folded only the drain error into its decision broadcast, leaving _select_and_load -- which stats, deserializes and renames -- able to kill rank 0 before the broadcast the peers were waiting in. Both now travel as decisions; cleanup's payload carries the phase so the peers name the operation that actually failed rather than reporting a "save". __init__'s orphan sweep is rank-0-only outside any collective, so a directory it cannot list would abort rank 0 alone and strand the peers at cleanup's first collective. It only reclaims space, so it is now best-effort with a warning. VA-1, VA-2, VA-3 --- ScaFFold/utils/checkpointing.py | 95 +++++++++++++++++++++++++------- tests/test_checkpointing.py | 98 ++++++++++++++++++++++++++++++++- 2 files changed, 171 insertions(+), 22 deletions(-) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 24f3fe6..22c6051 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -26,12 +26,13 @@ class CheckpointSaveError(RuntimeError): - """A checkpoint write failed. + """A rank-0 checkpoint operation failed (a write, a cleanup, a load). - Raised identically on every rank. Only rank 0 writes, but its outcome is - broadcast, so the peers report the real disk error instead of the - unmatched-collective symptom (an opaque gloo transport error, or an NCCL - watchdog timeout minutes later) that a rank-0-only raise produces. + Raised identically on every rank. Only rank 0 touches the run directory, + but its outcome is broadcast, so the peers report the real disk error + instead of the unmatched-collective symptom (an opaque gloo transport + error, or an NCCL watchdog timeout minutes later) that a rank-0-only raise + produces. """ @@ -111,17 +112,33 @@ def __init__( # Ensure base directory exists (Rank 0 only) if self.world_rank == 0: self.base_dir.mkdir(parents=True, exist_ok=True) - self._sweep_orphaned_tmp_files() + try: + self._sweep_orphaned_tmp_files() + except Exception as e: + # Construction is rank-0-only work outside any collective, so a + # raise here aborts rank 0 alone and leaves the peers waiting in + # the manager's first collective (``cleanup``). The sweep only + # reclaims space, so a directory that cannot be listed (a stale + # NFS/Lustre handle, a permissions oddity) degrades to a warning + # rather than taking the job down asymmetrically. + self._log( + f"Could not sweep orphaned checkpoint temp files in " + f"{self.base_dir}: {type(e).__name__}: {e}" + ) def cleanup(self, train_from_scratch: bool) -> None: """Clear existing checkpoints if training from scratch. Rank-symmetric, like every other collective point here: any pending - async write is drained and its outcome broadcast, so a failure raises - on all ranks together (see ``save_checkpoint``). + async write is drained, the rank-0 deletion is fenced, and whichever + failed is broadcast, so a failure raises on all ranks together (see + ``save_checkpoint``). The broadcast payload is ``(phase, description)`` + so the peers -- which saw neither the write nor the deletion -- report + the same operation rank 0 did. """ # Ensure any pending async save is finished before deleting. error = self._drain_pending_save() + failure = None if error is None else ("save", error) if train_from_scratch: # Drop the cached state that described the run being deleted. Both @@ -134,12 +151,29 @@ def cleanup(self, train_from_scratch: bool) -> None: self.last_saved_epoch = None if self.world_rank == 0: - self._remove_checkpoint_files() - - error = self._broadcast_obj(error) + # Rank 0 alone touches the filesystem here, while every peer is + # already committed to the broadcast below. Individual unlinks + # are tolerated inside, but the glob/stat around them can still + # raise on a shared filesystem (ESTALE, EACCES), and raising in + # this window strands the peers in an unmatched collective -- + # the R05 hazard. Report the failure through the broadcast, like + # a failed write. + try: + self._remove_checkpoint_files() + except Exception as e: + if failure is None: + self._save_error_exc = e + failure = ("cleanup", f"{type(e).__name__}: {e}") + else: + # A drained write already failed; that outcome is the + # one being reported, so this is only logged. + self._log(f"Clearing existing checkpoints also failed: {e}") + + failure = self._broadcast_obj(failure) self._barrier() - if error is not None: - self._raise_save_error(error) + if failure is not None: + phase, description = failure + self._raise_save_error(description, phase=phase) def _remove_checkpoint_files(self) -> None: """Delete this run's checkpoint files and debris (rank 0 only). @@ -230,15 +264,17 @@ def _drain_pending_save(self) -> Optional[str]: return f"{type(e).__name__}: {e}" return None - def _raise_save_error(self, description: str) -> None: - """Raise a broadcast save failure on this rank. + def _raise_save_error(self, description: str, phase: str = "save") -> None: + """Raise a broadcast rank-0 failure on this rank. - Rank 0 chains the original exception so its traceback survives; the - peers never saw it and raise the same message on its own. + ``phase`` names the operation that failed (``save``, ``cleanup``, + ``load``) so the message describes what actually went wrong. Rank 0 + chains the original exception so its traceback survives; the peers never + saw it and raise the same message on its own. """ cause, self._save_error_exc = self._save_error_exc, None raise CheckpointSaveError( - f"Checkpoint save failed on rank 0: {description}" + f"Checkpoint {phase} failed on rank 0: {description}" ) from cause def finalize_saves(self) -> None: @@ -317,14 +353,26 @@ def load_from_checkpoint(self, require_checkpoint: bool = False) -> int: # therefore never open the checkpoint files at all. result = None if self.world_rank == 0: - result = ( - ("save_failed", error) if error is not None else self._select_and_load() - ) + if error is not None: + result = ("save_failed", error) + else: + # The selection itself stats, deserializes and renames files + # while the peers are already blocked in the broadcast below, so + # anything it raises (a stale handle on ``exists``, a rename + # denied, an unpickling MemoryError) has to become a decision + # rather than a rank-0-only death. + try: + result = self._select_and_load() + except Exception as e: + self._save_error_exc = e + result = ("load_failed", f"{type(e).__name__}: {e}") status, payload = self._broadcast_obj(result) # 2. Every rank acts on the same decision rank 0 reached. if status == "save_failed": self._raise_save_error(payload) + if status == "load_failed": + self._raise_save_error(payload, phase="load") if status == "empty": if require_checkpoint: # An explicit restart must resume real state; silently @@ -410,6 +458,11 @@ def _select_and_load(self): * ``("empty", None)`` -- no checkpoint files exist; * ``("ok", checkpoint_dict)`` -- a candidate deserialized cleanly; * ``("unreadable", [paths])`` -- every candidate was corrupt. + + The filesystem calls around those decisions (``exists``, the quarantine + rename) can still fail outright; the caller runs this inside a guard + that turns such a failure into a broadcast ``("load_failed", ...)`` + decision, because raising here would strand the peers. """ candidates = [] if self.last_ckpt_path.exists(): diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 2ef9cad..acca106 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -39,7 +39,7 @@ import torch.distributed as dist import ScaFFold.utils.trainer as trainer_mod -from ScaFFold.utils.checkpointing import CheckpointManager +from ScaFFold.utils.checkpointing import CheckpointManager, CheckpointSaveError from ScaFFold.utils.trainer import PyTorchTrainer from tests.helpers import mpi_runner @@ -518,6 +518,102 @@ def test_init_sweeps_orphaned_tmp_files(tmp_path): assert quarantined.exists() +# --------------------------------------------------------------------------- +# VA-1/VA-2/VA-3 -- the remaining rank-0 filesystem windows are fenced +# +# Rank 0 is the only rank that touches the run directory, and it does so while +# its peers are already committed to the next collective. Every such window must +# therefore report its failure *through* that collective; raising inside it +# leaves the peers in an unmatched collective, where a plain disk error +# resurfaces as an opaque gloo transport error or an NCCL watchdog timeout. +# --------------------------------------------------------------------------- + + +def _record_collectives(monkeypatch): + """Stub the process-group collectives, recording what a rank posts.""" + posted = {"broadcasts": [], "barriers": 0} + + def fake_broadcast(objs, src=0): + posted["broadcasts"].append(objs[0]) + + def fake_barrier(*args, **kwargs): + posted["barriers"] += 1 + + monkeypatch.setattr(dist, "broadcast_object_list", fake_broadcast) + monkeypatch.setattr(dist, "barrier", fake_barrier) + return posted + + +def _raise_stale(*args, **kwargs): + raise OSError("[Errno 116] Stale file handle") + + +def test_cleanup_rank0_fs_error_travels_through_the_broadcast(tmp_path, monkeypatch): + """A failure while clearing checkpoints fails every rank, not just rank 0. + + ``_remove_checkpoint_files`` globs and stats the run directory; on a shared + filesystem those can raise (ESTALE, EACCES) even though each individual + unlink is already tolerated. That happens between the drain and the + broadcast, so an unfenced raise re-creates exactly the hazard R05 closed. + """ + mgr, _ = _make_manager(tmp_path) + mgr.dist_enabled = True + posted = _record_collectives(monkeypatch) + monkeypatch.setattr(CheckpointManager, "_remove_checkpoint_files", _raise_stale) + + with pytest.raises(CheckpointSaveError) as excinfo: + mgr.cleanup(train_from_scratch=True) + + assert "Stale file handle" in str(excinfo.value) + # The peers' collectives were posted before the raise, so they fail with the + # same error instead of hanging. + assert len(posted["broadcasts"]) == 1 + assert "Stale file handle" in str(posted["broadcasts"][0]) + assert posted["barriers"] == 1 + + +def test_load_rank0_selection_error_travels_through_the_broadcast( + tmp_path, monkeypatch +): + """A failure inside ``_select_and_load`` is a broadcast decision too. + + The load path folded only the *drain* error into its decision broadcast, so + rank 0 stat-ing or renaming a checkpoint candidate could still die before + the broadcast the peers were already waiting in. + """ + mgr, _ = _make_manager(tmp_path) + mgr.dist_enabled = True + posted = _record_collectives(monkeypatch) + monkeypatch.setattr(CheckpointManager, "_select_and_load", _raise_stale) + + with pytest.raises(CheckpointSaveError) as excinfo: + mgr.load_from_checkpoint() + + assert "Stale file handle" in str(excinfo.value) + assert len(posted["broadcasts"]) == 1 + assert posted["broadcasts"][0][0] == "load_failed" + + +def test_init_survives_an_unlistable_run_dir(tmp_path, monkeypatch, capsys): + """A run directory that cannot be listed warns instead of killing __init__. + + The orphan sweep runs on rank 0 only and outside any collective, so a raise + here aborts rank 0 alone and strands the peers at the manager's first + collective (``cleanup``). ``pathlib`` already swallows PermissionError, but + a stale NFS/Lustre handle propagates; the sweep is opportunistic, so it must + degrade to a warning. + """ + base = tmp_path / "checkpoints" + base.mkdir() + (base / "checkpoint_last.pth.tmp.999999").write_bytes(b"partial checkpoint") + monkeypatch.setattr(Path, "glob", _raise_stale) + + mgr, _ = _make_manager(base) + + assert mgr.base_dir == base + assert "sweep" in capsys.readouterr().out.lower() + + # --------------------------------------------------------------------------- # F71 -- CPU tensors are cloned into the snapshot # --------------------------------------------------------------------------- From 2fb1bf53661670bc1cfa2e734799ecf6ae02c3f6 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:51:15 -0700 Subject: [PATCH 35/54] Fence the remaining rank-0 windows in datagen The category search broadcasts its scan of the existing category files, but rank 0 then loaded those files' parameters -- for its duplicate guard -- just *after* that broadcast, unfenced. A CSV that will not parse (ragged, or hand-edited) killed rank 0 while every peer had already taken the broadcast and entered the work loop, stranding them in its next collective: the same window class the broadcast was introduced to close. Scan and load are now one guarded decision reported through one broadcast, as get_dataset reports its selection. The datagen consensus guards caught (Exception, SystemExit), which is not the same as "everything": a KeyboardInterrupt delivered to rank 0 alone (Ctrl-C on the launching terminal, a watchdog SIGINT) unwound straight past the broadcast and hung the peers. They now catch BaseException. SystemExit keeps being converted -- it must never escape get_dataset, since a peer would read the silent unwind as success -- while the rank that was interrupted re-raises the interrupt after posting the sentinel, so it keeps the operator's exit status and its peers still learn to stop. volumegen's guard, the same pattern feeding the same allreduce, is fixed with it. VB-2, VB-3 --- ScaFFold/datagen/category_search.py | 45 +++++++++++++++++------ ScaFFold/datagen/get_dataset.py | 42 ++++++++++++++++++--- ScaFFold/datagen/volumegen.py | 12 +++++- tests/datagen/test_category_search.py | 53 +++++++++++++++++++++++++-- tests/datagen/test_mpi_consensus.py | 51 ++++++++++++++++++++++++++ 5 files changed, 182 insertions(+), 21 deletions(-) diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py index af53a9c..f3b3496 100644 --- a/ScaFFold/datagen/category_search.py +++ b/ScaFFold/datagen/category_search.py @@ -461,21 +461,44 @@ def main(config: Config) -> None: # partially visible directory) put one rank inside the loop while another is # past it, so the two post mismatched collectives on COMM_WORLD and the job # hangs. One scan, one broadcast, one shared verdict. - if rank == 0: - existing_indices = parse_category_indices(fracts_write_dir) - else: - existing_indices = None - existing_indices = comm.bcast(existing_indices, root=0) - + # + # Rank 0 also loads the parameters of the categories already on disk (only + # it writes, so only it needs them for the duplicate guard). That read is + # part of the same rank-0-only window: a category CSV that will not parse -- + # ragged, or hand-edited -- would otherwise kill rank 0 *after* the peers + # had already taken the broadcast and moved on to the next collective. Scan + # and load are therefore one guarded decision, reported through one + # broadcast, exactly as ``get_dataset`` reports its selection. existing_params = [] + interrupt = None if rank == 0: - for idx in existing_indices: - existing_params.append( - np.loadtxt( - os.path.join(fracts_write_dir, "%06d.csv" % idx), - delimiter=",", + try: + existing_indices = parse_category_indices(fracts_write_dir) + for idx in existing_indices: + existing_params.append( + np.loadtxt( + os.path.join(fracts_write_dir, "%06d.csv" % idx), + delimiter=",", + ) ) + scan = ("ok", existing_indices) + except BaseException as e: + existing_params = [] + scan = ( + "error", + f"rank 0 failed to scan existing categories in " + f"{fracts_write_dir}: {type(e).__name__}: {e}", ) + interrupt = e if isinstance(e, KeyboardInterrupt) else None + else: + scan = None + status, payload = comm.bcast(scan, root=0) + if status == "error": + # Rank 0 keeps an operator's interrupt; every rank aborts either way. + if interrupt is not None: + raise interrupt + raise RuntimeError(f"category search failed: {payload}") + existing_indices = payload # Calculate number of remaining fractal categories to generate existing_categories = len(existing_indices) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index 09e4b33..e54b29f 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -292,6 +292,21 @@ def _decide_reuse_or_generate( return ("generate", str(tmp), str(dest)) +def _reraisable(exc: BaseException) -> BaseException | None: + """Return ``exc`` when the rank that caught it should re-raise it verbatim. + + The consensus guards below turn any failure into a sentinel so every rank + aborts together, and every rank then raises a ``RuntimeError`` carrying the + collected messages. That is the right report for a genuine error -- and for + ``SystemExit``, which must never escape ``get_dataset`` (a peer would treat + the silent unwind as success). A ``KeyboardInterrupt`` is different: it is + not an error but an operator abort, so the rank that received it re-raises + it after posting the sentinel, keeping the interrupt's own semantics (and + exit status) while its peers still learn to stop. + """ + return exc if isinstance(exc, KeyboardInterrupt) else None + + def get_dataset( config: Namespace, require_commit: bool = False, # default: ignore commit mismatches for reuse @@ -334,22 +349,32 @@ def get_dataset( # the broadcast below, so a rank-0 exception would strand the whole job. # Any failure is therefore turned into an error sentinel that travels # through the same broadcast and makes every rank raise the same error. + interrupt = None if rank == 0: try: decision = _decide_reuse_or_generate( base, config_id, commit, require_commit, log ) - except (Exception, SystemExit) as e: + except BaseException as e: + # BaseException, not (Exception, SystemExit): a KeyboardInterrupt + # delivered to rank 0 alone (Ctrl-C on the launching terminal, a + # site watchdog SIGINT) would otherwise skip the broadcast and hang + # every peer -- the exact failure this guard exists to prevent. decision = ( "error", f"rank 0 failed to select a dataset under {base}: " f"{type(e).__name__}: {e}", ) + interrupt = _reraisable(e) else: decision = None decision = comm.bcast(decision, root=0) if decision[0] == "error": + # Rank 0 keeps the abort signal it was actually given; the peers, which + # only ever saw the sentinel, report it as a generation failure. + if interrupt is not None: + raise interrupt raise RuntimeError(f"dataset selection failed: {decision[1]}") if decision[0] == "reuse": @@ -364,13 +389,15 @@ def get_dataset( err = "" # A worker failure must not skip any collective below: catch everything - # (including SystemExit, which is a BaseException and would otherwise bypass - # the consensus) so every rank always reaches the allreduce and gather. + # (BaseException, so neither SystemExit nor a KeyboardInterrupt delivered to + # one rank can bypass the consensus) so every rank always reaches the + # allreduce and gather. try: volumegen.main(config) - except (Exception, SystemExit) as e: + except BaseException as e: ok = False err = f"volumegen attempt failed: rank {rank}: {type(e).__name__}: {e}" + interrupt = _reraisable(e) # Reach a global verdict, then have every rank participate in the error # gather so no rank is left in a mismatched collective on the failure path. @@ -382,6 +409,8 @@ def get_dataset( shutil.rmtree(tmp, ignore_errors=True) # Every rank raises with the collected messages, so a non-root rank # never returns an unfinalized dataset path. + if interrupt is not None: + raise interrupt msgs = "; ".join(e for e in errs if e) raise RuntimeError(f"dataset generation failed: {msgs or 'unknown error'}") @@ -403,16 +432,19 @@ def get_dataset( } _write_meta_atomic(tmp / META_FILENAME, meta) tmp.rename(dest) - except (Exception, SystemExit) as e: + except BaseException as e: finalize_err = ( f"rank 0 failed to finalize dataset at {dest}: {type(e).__name__}: {e}" ) + interrupt = _reraisable(e) # This broadcast doubles as the synchronization the old Barrier provided: no # rank returns before rank 0 has published the rename (or reported that it # could not), so nobody observes the staging path or a missing dataset. finalize_err = comm.bcast(finalize_err, root=0) if finalize_err: + if interrupt is not None: + raise interrupt raise RuntimeError(f"dataset generation failed: {finalize_err}") return dest diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py index 34e17f9..abb05fe 100644 --- a/ScaFFold/datagen/volumegen.py +++ b/ScaFFold/datagen/volumegen.py @@ -229,6 +229,7 @@ def main(config: Dict): # the failure is then propagated to all ranks via an allreduce. ok = True err = "" + interrupt = None try: if start_idx >= end_idx: @@ -337,11 +338,14 @@ def main(config: Dict): total_time, len(volumes_contents_subset) / total_time, ) - except (Exception, SystemExit) as e: + except BaseException as e: # Capture the failure locally instead of letting it unwind past the - # collective below, which would desynchronize the ranks. + # collective below, which would desynchronize the ranks. BaseException, + # not (Exception, SystemExit): a KeyboardInterrupt delivered to one rank + # would otherwise skip the consensus and hang the others. ok = False err = f"rank {rank}: {type(e).__name__}: {e}" + interrupt = e if isinstance(e, KeyboardInterrupt) else None # Consensus on the generation status. This replaces a bare Barrier: every # rank always executes exactly this collective (regardless of success or @@ -350,6 +354,10 @@ def main(config: Dict): all_ok = comm.allreduce(1 if ok else 0, op=MPI.MIN) == 1 errs = comm.allgather(err) if not all_ok: + # The interrupted rank re-raises the operator's abort verbatim; the + # others report the gathered failure. + if interrupt is not None: + raise interrupt msgs = "; ".join(e for e in errs if e) raise RuntimeError(f"volume generation failed: {msgs or 'unknown error'}") diff --git a/tests/datagen/test_category_search.py b/tests/datagen/test_category_search.py index 2cc43b1..1b861c3 100644 --- a/tests/datagen/test_category_search.py +++ b/tests/datagen/test_category_search.py @@ -149,7 +149,7 @@ def _seed_one_category(config: Namespace) -> None: def test_work_scan_is_root_only_and_broadcast(tmp_path, monkeypatch): """A non-root rank never scans; it consumes root's broadcast index list.""" config = _cs_config(tmp_path / "fractals") - comm = CategorySearchComm(rank=1, size=2, bcast_returns=[[0]]) + comm = CategorySearchComm(rank=1, size=2, bcast_returns=[("ok", [0])]) monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) def no_scan(*_args, **_kwargs): @@ -182,17 +182,64 @@ def test_divergent_fs_views_take_the_same_collective_path(tmp_path, monkeypatch) # Rank 1: an empty directory (a divergent view), but root broadcast [0]. peer_config = _cs_config(tmp_path / "peer_view") - peer_comm = CategorySearchComm(rank=1, size=2, bcast_returns=[[0]]) + peer_comm = CategorySearchComm(rank=1, size=2, bcast_returns=[("ok", [0])]) monkeypatch.setattr(cs, "MPI", FakeMPI(peer_comm)) cs.main(peer_config) - assert root_comm.bcast_payloads[0] == [0] + assert root_comm.bcast_payloads[0] == ("ok", [0]) assert peer_comm.calls == root_comm.calls, ( "ranks with divergent filesystem views issued different collectives: " f"rank 0 {root_comm.calls} vs rank 1 {peer_comm.calls}" ) +# --------------------------------------------------------------------------- +# VB-2: the whole rank-0 scan window is fenced, not just the index parse. +# +# Rank 0 also reads the parameters of every category already on disk, right +# after the scan broadcast. A CSV that will not parse therefore killed rank 0 +# while its peers had already consumed the broadcast and moved on -- the same +# stranding the scan broadcast was introduced to prevent. +# --------------------------------------------------------------------------- + + +def test_unparseable_existing_category_is_broadcast_not_raised_on_root( + tmp_path, monkeypatch +): + """A ragged category CSV becomes a broadcast error, not a rank-0-only death.""" + config = _cs_config(tmp_path / "fractals") + param_dir = Path(layout.category_param_dir(config)) + param_dir.mkdir(parents=True, exist_ok=True) + # Six-digit name, so the scan accepts it; ragged rows, so loadtxt raises. + (param_dir / "000000.csv").write_text("0.5,0.5,0.5\n0.5,0.5\n") + + comm = CategorySearchComm(rank=0, size=2) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + with pytest.raises(RuntimeError) as excinfo: + cs.main(config) + + # Rank 0 stopped at the scan broadcast, and what it broadcast is the error + # sentinel its peers need in order to abort with it. + assert comm.calls == ["Barrier", "bcast"] + assert comm.bcast_payloads[-1][0] == "error" + assert "000000.csv" in str(excinfo.value) or "3DIFS_param" in str(excinfo.value) + + +def test_peer_raises_on_broadcast_scan_error(tmp_path, monkeypatch): + """A peer receiving the scan sentinel raises instead of entering the loop.""" + config = _cs_config(tmp_path / "fractals") + sentinel = ("error", "rank 0 failed to scan existing categories: ValueError: boom") + comm = CategorySearchComm(rank=1, size=2, bcast_returns=[sentinel]) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + with pytest.raises(RuntimeError, match="boom"): + cs.main(config) + + # It never reached the work loop's collectives. + assert comm.calls == ["Barrier", "bcast"] + + # --------------------------------------------------------------------------- # R31: category CSVs appear complete or not at all. # diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index c98fc5b..a4c954a 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -540,6 +540,57 @@ def test_non_root_raises_on_broadcast_finalize_error(tmp_path, monkeypatch): assert not dest.exists() +# --------------------------------------------------------------------------- +# VB-3: the consensus guards catch BaseException, not (Exception, SystemExit). +# +# ``KeyboardInterrupt`` is neither, so an interrupt delivered to rank 0 alone +# (Ctrl-C on the launching terminal, a site watchdog's SIGINT) unwound straight +# past the broadcast and hung every peer -- the failure mode the guard exists to +# prevent, arriving through the one exception class it did not cover. +# --------------------------------------------------------------------------- + + +def test_rank0_interrupt_still_posts_the_decision_sentinel(tmp_path, monkeypatch): + """A KeyboardInterrupt on rank 0 reaches the peers as an error sentinel.""" + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=2) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + + def interrupted(*_args, **_kwargs): + raise KeyboardInterrupt + + monkeypatch.setattr(gd, "_decide_reuse_or_generate", interrupted) + + # Rank 0 keeps the operator's abort ... + with pytest.raises(KeyboardInterrupt): + gd.get_dataset(config) + + # ... but only after telling the peers to stop. + assert comm.calls == ["bcast"] + assert comm.bcast_payloads[0][0] == "error" + assert "KeyboardInterrupt" in comm.bcast_payloads[0][1] + + +def test_interrupt_during_generation_reaches_the_consensus(tmp_path, monkeypatch): + """An interrupt inside volumegen still drives the allreduce and allgather.""" + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=2, allreduce_result=0, allgather_peers=[""]) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + + def interrupted(_config): + raise KeyboardInterrupt + + monkeypatch.setattr(volumegen, "main", interrupted) + + with pytest.raises(KeyboardInterrupt): + gd.get_dataset(config) + + assert "allreduce" in comm.calls and "allgather" in comm.calls + assert "KeyboardInterrupt" in comm.allgather_payloads[0] + + # --------------------------------------------------------------------------- # R37: orphaned staging dirs are reclaimed instead of accumulating forever. # --------------------------------------------------------------------------- From 1cce9bb10c5101c039f64d8cb9b3634d810a834d Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:55:36 -0700 Subject: [PATCH 36/54] Age staging dirs by their deepest recent write The orphan cleanup judged a staging directory by the newest mtime among the directory and its immediate children. Volume writing lands at depth >= 2 (volumes//N.npy), and generation is not bounded by the staleness threshold -- at the larger scales it runs for days -- so a perfectly healthy generation stopped touching anything the probe could see and read as "untouched for 48 hours". A concurrent same-config start then rmtree'd it out from under its peers, which died on FileNotFoundError. Liveness is now reported rather than inferred: every rank writing volumes refreshes /.heartbeat every five minutes, and cleanup keeps any directory whose heartbeat is younger than the threshold. The marker is removed before the staging dir is published, so a dataset does not carry it. Backing the heartbeat up (for staging dirs written before it existed, or killed before the first beat) the mtime probe now walks a bounded number of directory levels instead of one. It stats directories, whose mtimes change when entries are created in them, so it notices a writer without a stat storm over the volume files, and it stops at the first recent entry -- the live case, the one that must not be misjudged, is the cheap one. A tree too wide to walk within the bound is called live: failing to reclaim disk is recoverable, deleting a running job's dataset is not. VB-1 --- ScaFFold/datagen/get_dataset.py | 119 ++++++++++++++++++++++------ ScaFFold/datagen/volumegen.py | 58 +++++++++++++- tests/datagen/test_mpi_consensus.py | 114 ++++++++++++++++++++++++++ 3 files changed, 266 insertions(+), 25 deletions(-) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index e54b29f..168f002 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -37,10 +37,17 @@ # reader never observes a half-written dataset. The prefix is also what the # reuse scan skips and what the orphan cleanup collects. TMP_PREFIX = ".tmp_" -# How long a staging directory must have sat untouched before it is treated as -# orphaned (left by a killed/OOM'd job) and reclaimed. See +# How long a staging directory must have shown no sign of life before it is +# treated as orphaned (left by a killed/OOM'd job) and reclaimed. See # ``_cleanup_stale_staging_dirs`` for the safety argument behind the value. STALE_STAGING_AGE_SECONDS = 24 * 60 * 60 +# Bounds on the liveness probe that backs up the heartbeat: how many directory +# levels below a staging dir it looks at, and how many directories it is willing +# to visit before it gives up and calls the tree live. Directory mtimes change +# when entries are created in them, so a couple of levels is enough to notice a +# writer without stat-ing every volume file. +_STALE_PROBE_MAX_DEPTH = 3 +_STALE_PROBE_MAX_DIRS = 10000 # Bumped from 2 to 3 when instance point clouds moved from float64 to float32: # the storage layout is unchanged, but float32 voxel binning shifts a handful of # boundary voxels, so a float64-era dataset must not be reused as if it were @@ -138,6 +145,67 @@ def _git_commit_short(log, source_dir: Path | None = None) -> str: return "no-commit-id" +def _has_recent_write(path: Path, cutoff: float) -> bool: + """Return True if the top levels of ``path`` were written after ``cutoff``. + + A directory's mtime changes whenever an entry is created in it, so the + directory mtimes near the top of a staging tree are a cheap proxy for "a + writer is active down there" -- no stat of the (possibly hundreds of + thousands of) volume files is needed. The walk therefore stats the staging + directory, its immediate children, and directories down to + ``_STALE_PROBE_MAX_DEPTH``, and stops at the first recent entry, which makes + the live case (the one that must not be misjudged) the cheap one. + + A tree wide enough to exceed ``_STALE_PROBE_MAX_DIRS`` is reported as recent + rather than walked further: failing to reclaim disk is recoverable, deleting + a running job's dataset is not. + """ + stack = [(path, 0)] + dirs_seen = 0 + while stack: + current, depth = stack.pop() + if current.stat().st_mtime > cutoff: + return True + if depth >= _STALE_PROBE_MAX_DEPTH: + continue + with os.scandir(current) as entries: + for entry in entries: + if entry.is_dir(follow_symlinks=False): + dirs_seen += 1 + if dirs_seen > _STALE_PROBE_MAX_DIRS: + return True + stack.append((Path(entry.path), depth + 1)) + elif depth == 0 and entry.stat().st_mtime > cutoff: + # Files directly in the staging dir (the heartbeat, + # volumes_contents.csv, meta.yaml) are few and cheap. + return True + return False + + +def _staging_dir_is_live(path: Path, cutoff: float) -> bool: + """Return True if ``path`` shows any sign of a generation still running. + + Two signals, in order of authority: + + 1. ``/.heartbeat``, refreshed by every writing rank every few + minutes for as long as volumes are being written (see + ``volumegen.StagingHeartbeat``). This is the reliable one, because it + does not depend on where in the tree the writers currently are. + 2. a bounded-depth mtime probe, which covers staging directories written + before the heartbeat existed, or killed before the first beat. + + Raises ``OSError`` if the directory cannot be examined; the caller treats + that as "cannot tell" and leaves the directory alone. + """ + heartbeat = path / volumegen.STAGING_HEARTBEAT_NAME + try: + if heartbeat.stat().st_mtime > cutoff: + return True + except OSError: + pass # No heartbeat: fall back to the mtime probe. + return _has_recent_write(path, cutoff) + + def _cleanup_stale_staging_dirs( base: Path, log, max_age: float = STALE_STAGING_AGE_SECONDS ) -> None: @@ -149,39 +217,39 @@ def _cleanup_stale_staging_dirs( Safety policy. Only directories that (a) live directly under *this* config_id base, (b) carry the ``.tmp_`` prefix this module owns, and (c) - have been untouched for ``max_age`` are removed. The age gate is what keeps - a *concurrent* job's staging directory safe: unique staging names mean two - live jobs never share a directory, but they do share the base, so a live - peer's directory is visible here -- it is simply orders of magnitude younger - than the threshold (a day, against generations measured in minutes to - hours). Published datasets and anything outside ``base`` are never touched. - Failures are logged and ignored: cleanup is opportunistic and must never - break the decision it runs inside. + show no sign of life for ``max_age`` are removed. Published datasets and + anything outside ``base`` are never touched. Failures are logged and + ignored: cleanup is opportunistic and must never break the decision it runs + inside. + + "No sign of life" is the delicate part, because the staging directory of a + *running* job is visible here (unique staging names mean two live jobs never + share a directory, but they do share the base). Age alone is not enough: + generation is not bounded by a day -- at the larger scales it is measured in + days -- and after the first minutes it writes only at depth >= 2, so the top + of the tree stops changing while the job is perfectly healthy. Judging by + the top-level mtimes alone therefore let a concurrent same-config start + rmtree a live generation out from under its peers. ``_staging_dir_is_live`` + is the answer: an explicit heartbeat maintained by the writers, backed by a + bounded-depth mtime probe for directories that predate it. """ now = time.time() + cutoff = now - max_age for path in base.iterdir(): if not path.name.startswith(TMP_PREFIX) or not path.is_dir(): continue try: - # Newest mtime among the staging dir and its immediate children: a - # bounded, cheap probe (no recursive stat storm over a partially - # generated dataset) that still notices a job that has started - # laying down its split directories. - newest = path.stat().st_mtime - for child in path.iterdir(): - newest = max(newest, child.stat().st_mtime) + if _staging_dir_is_live(path, cutoff): + continue except OSError as exc: - log.warning("Could not stat staging dir %s: %s", path, exc) - continue - - age = now - newest - if age < max_age: + log.warning("Could not examine staging dir %s: %s", path, exc) continue log.info( - "Removing orphaned dataset staging dir %s (untouched for %.1f hours)", + "Removing orphaned dataset staging dir %s (no write in the last " + "%.1f hours, and no live generation heartbeat)", path, - age / 3600.0, + max_age / 3600.0, ) shutil.rmtree(path, ignore_errors=True) @@ -430,6 +498,9 @@ def get_dataset( "code_commit": commit, "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), } + # The liveness marker described this directory while it was being + # written; it has no meaning in a published dataset. + (tmp / volumegen.STAGING_HEARTBEAT_NAME).unlink(missing_ok=True) _write_meta_atomic(tmp / META_FILENAME, meta) tmp.rename(dest) except BaseException as e: diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py index abb05fe..997a73f 100644 --- a/ScaFFold/datagen/volumegen.py +++ b/ScaFFold/datagen/volumegen.py @@ -27,6 +27,55 @@ from ScaFFold.utils.data_types import MASK_DTYPE, VOLUME_DTYPE from ScaFFold.utils.utils import setup_mpi_logger +# Liveness marker for the directory being generated into. Volume writing is the +# long phase of a generation and it happens deep inside the tree +# (``volumes//N.npy``), so the top of the staging directory can look +# untouched for many hours while the job is perfectly healthy. Every writing +# rank therefore refreshes this file periodically, and +# ``get_dataset._staging_dir_is_live`` reads it instead of trying to infer +# liveness from mtimes it cannot cheaply see. The name is owned here, next to +# the writer; ``get_dataset`` (which already imports this module) reads it from +# here so the two sides cannot drift. +STAGING_HEARTBEAT_NAME = ".heartbeat" +# Refresh interval. Small enough to be negligible against the staleness +# threshold (a day), large enough that it is one utime per rank per few minutes +# no matter how fast volumes are written. +STAGING_HEARTBEAT_INTERVAL_SECONDS = 5 * 60 + + +class StagingHeartbeat: + """Periodically touch a staging directory's heartbeat file. + + ``beat()`` is called from the volume loop and is a no-op until the interval + has elapsed, so it costs one comparison per volume. Every rank writing into + the directory beats the same file: the marker means "somebody is still + working here", and a last-writer-wins utime is exactly the semantics wanted. + Failures are swallowed -- a heartbeat that cannot be written must never take + down a generation that is otherwise fine (the cleanup's second signal, the + bounded mtime probe, still applies). + """ + + def __init__( + self, staging_dir, interval: float = STAGING_HEARTBEAT_INTERVAL_SECONDS + ) -> None: + self.path = os.path.join(str(staging_dir), STAGING_HEARTBEAT_NAME) + self.interval = interval + self._last_beat = float("-inf") + + def beat(self, now: float | None = None) -> bool: + """Touch the marker if the interval has elapsed; return whether it did.""" + now = time.time() if now is None else now + if now - self._last_beat < self.interval: + return False + self._last_beat = now + try: + with open(self.path, "a"): + pass + os.utime(self.path, (now, now)) + except OSError: + return False + return True + def load_np_ptcloud(path: str) -> np.ndarray: """ @@ -253,9 +302,16 @@ def main(config: Dict): # this run's seed produced. Resolved once, outside the loop. instances_dir = layout.instance_dir(config) - # Generation loop + # Generation loop. Every rank reports that this staging directory + # is still being written to, so a concurrent job's orphan cleanup + # can tell a live multi-hour generation from one killed a day ago + # (the volumes themselves land two levels down, where a cheap + # top-level mtime probe cannot see them). + heartbeat = StagingHeartbeat(dataset_dir) + heartbeat.beat() start_time = time.time() for i, curr_vol in enumerate(volumes_contents_subset): + heartbeat.beat() if i % 10 == 0: log.debug("Rank %s processing local volume %s", rank, i) diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index a4c954a..5185eb3 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -657,6 +657,120 @@ def test_cleanup_never_touches_published_datasets(tmp_path): assert other_orphan.exists(), "cleanup escaped this job's config_id base" +# --------------------------------------------------------------------------- +# VB-1: a staging dir is aged by its DEEPEST recent write, not its top level. +# +# Generation is not bounded by the staleness threshold -- at the larger scales +# it runs for days -- and after the first minutes it writes only at depth >= 2 +# (``volumes//N.npy``). Judging liveness from the staging dir and its +# immediate children alone therefore reported a healthy multi-day generation as +# "untouched for 48 hours", and a concurrent same-config start rmtree'd it out +# from under its peers (which then died on FileNotFoundError). +# --------------------------------------------------------------------------- + + +def _cleanup(base: Path, name: str) -> None: + gd._cleanup_stale_staging_dirs(base, logging.getLogger(name)) + + +def test_live_generation_survives_a_deep_write(tmp_path): + """A >24h-old staging dir with a fresh deep write is NOT reclaimed.""" + base = tmp_path / "cid" + base.mkdir(parents=True) + + live = base / f"{gd.TMP_PREFIX}20260101-000000_222_cafebabe" + split = live / "volumes" / "training" + split.mkdir(parents=True) + _age_tree(live, 2 * gd.STALE_STAGING_AGE_SECONDS) + + # The job is alive and still laying down volumes: the write lands two levels + # down, so only that directory's mtime is current. + (split / "0.npy").write_bytes(b"payload from a running job") + + _cleanup(base, "test_live_generation_survives_a_deep_write") + + assert live.exists(), "a live generation's staging dir was reclaimed" + + +def test_live_generation_survives_on_its_heartbeat_alone(tmp_path): + """A fresh heartbeat keeps a staging dir whose whole tree looks ancient. + + The mtime probe is bounded, so a generation writing deeper than it looks + (or on a filesystem with coarse directory mtimes) still has to be safe. The + heartbeat is the signal that does not depend on the shape of the tree. + """ + base = tmp_path / "cid" + base.mkdir(parents=True) + + live = base / f"{gd.TMP_PREFIX}20260101-000000_333_f00d" + (live / "volumes" / "training" / "deep" / "deeper").mkdir(parents=True) + heartbeat = live / volumegen.STAGING_HEARTBEAT_NAME + heartbeat.write_text("") + _age_tree(live, 2 * gd.STALE_STAGING_AGE_SECONDS) + now = time.time() + os.utime(heartbeat, (now, now)) + + _cleanup(base, "test_live_generation_survives_on_its_heartbeat_alone") + + assert live.exists(), "a heartbeating generation's staging dir was reclaimed" + + +def test_dead_staging_dir_with_stale_heartbeat_is_reclaimed(tmp_path): + """The heartbeat must not turn cleanup into a no-op (R37 still holds).""" + base = tmp_path / "cid" + base.mkdir(parents=True) + + orphan = base / f"{gd.TMP_PREFIX}20200101-000000_111_deadbeef" + split = orphan / "volumes" / "training" + split.mkdir(parents=True) + (split / "0.npy").write_bytes(b"stale payload") + (orphan / volumegen.STAGING_HEARTBEAT_NAME).write_text("") + _age_tree(orphan, 2 * gd.STALE_STAGING_AGE_SECONDS) + + _cleanup(base, "test_dead_staging_dir_with_stale_heartbeat_is_reclaimed") + + assert not orphan.exists(), "an orphaned staging dir was not reclaimed" + + +def test_generation_writes_and_then_drops_the_heartbeat(tmp_path, monkeypatch): + """volumegen marks the staging dir live; publishing removes the marker.""" + config = _reuse_config(tmp_path / "datasets") + comm = FakeComm(rank=0, size=1, allreduce_result=1) + monkeypatch.setattr(gd, "MPI", FakeMPI(comm)) + monkeypatch.setattr(gd, "_git_commit_short", lambda log: "abc123") + + beating = {} + + def fake_volumegen(cfg): + # Stand in for the write loop: report the staging dir as live. + volumegen.StagingHeartbeat(cfg.dataset_dir).beat() + beating["path"] = Path(cfg.dataset_dir) / volumegen.STAGING_HEARTBEAT_NAME + beating["existed_during_generation"] = beating["path"].exists() + + monkeypatch.setattr(volumegen, "main", fake_volumegen) + + published = Path(gd.get_dataset(config)) + + assert beating["existed_during_generation"], "no heartbeat during generation" + assert not (published / volumegen.STAGING_HEARTBEAT_NAME).exists(), ( + "the staging heartbeat was published with the dataset" + ) + + +def test_heartbeat_respects_its_interval(tmp_path): + """``beat`` is a no-op until the interval elapses, then refreshes.""" + staging = tmp_path / "staging" + staging.mkdir() + heartbeat = volumegen.StagingHeartbeat(staging, interval=60) + + assert heartbeat.beat(now=1000.0) is True + first = Path(heartbeat.path).stat().st_mtime + assert heartbeat.beat(now=1030.0) is False # inside the interval + assert Path(heartbeat.path).stat().st_mtime == first + assert heartbeat.beat(now=1090.0) is True + assert Path(heartbeat.path).stat().st_mtime > first + + # --------------------------------------------------------------------------- # R28: meta.yaml is published atomically, so no reader ever sees a partial one. # --------------------------------------------------------------------------- From 5cfa82f977257b7f433843707521626f28755edb Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 16:57:19 -0700 Subject: [PATCH 37/54] Sweep category-search temp files and warn on old-layout libraries Both atomic writers in the category search name their temp file after the writing pid and unlink it when the write raises -- but a SIGKILL (walltime, OOM, node failure) skips that Python-level cleanup. Nothing ever looked at the strays again, so .NNNNNN.csv.tmp and .rng_attempt_rank*.tmp piled up one per killed process. Rank 0 now sweeps them at startup, best-effort, the way instance.py already sweeps its equivalents. The seed-keyed relayout is silent about libraries in the old location: the existence check simply does not find them. From outside that looks like a library that was there yesterday being regenerated for no reason -- hours of work at scale. One warning naming the old directory and the new one explains it. VB-4, VB-6 --- ScaFFold/datagen/category_search.py | 34 +++++++++++++ ScaFFold/datagen/layout.py | 35 +++++++++++++ tests/datagen/test_category_search.py | 72 +++++++++++++++++++++++++++ 3 files changed, 141 insertions(+) diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py index f3b3496..3cc1f81 100644 --- a/ScaFFold/datagen/category_search.py +++ b/ScaFFold/datagen/category_search.py @@ -403,6 +403,35 @@ def write_attempt_counter(fracts_write_dir: str, rank: int, attempt_index: int) os.replace(tmp, path) +def _sweep_stale_temp_files(fracts_write_dir: str, log) -> None: + """Remove temp files stranded by killed writes in the category directory. + + Both atomic writers here (``_savetxt_atomic`` and ``write_attempt_counter``) + unlink their temp file when the write raises, but a SIGKILL -- walltime, an + OOM, a node failure -- skips that Python-level cleanup and strands it. The + names carry the writer's pid, so they accumulate one per killed process and + nothing else ever removes them; ``instance.py`` sweeps its equivalents for + exactly this reason. + + Called on rank 0 before any rank has written anything this run, and + best-effort: this is housekeeping, and it runs just before a Barrier the + peers are heading into, so it must not raise. + """ + patterns = ( + # .NNNNNN.csv.tmp -- a partially written category CSV. + f"{fracts_write_dir}/.*.csv.tmp*", + # .rng_attempt_rank.tmp -- a partially written attempt counter. + f"{fracts_write_dir}/.rng_attempt_rank*.tmp*", + ) + for pattern in patterns: + for stale in glob.glob(pattern): + try: + os.remove(stale) + log.info("Removed stale category-search temp file %s", stale) + except OSError as exc: + log.warning("Could not remove stale temp file %s: %s", stale, exc) + + def main(config: Config) -> None: """ Generate fractal categories. @@ -440,10 +469,15 @@ def main(config: Config) -> None: fracts_write_dir = layout.category_param_dir(config) if rank == 0: log.info("Writing fractals to %s", fracts_write_dir) + # A library in the pre-seed layout is invisible to everything below, so + # say why it is being ignored rather than appearing to regenerate work + # that is plainly still on disk. + layout.warn_if_legacy_library(config, log) if os.path.exists(fracts_write_dir) and config.datagen_from_scratch: log.info("Removing existing fractals directory") shutil.rmtree(fracts_write_dir) os.makedirs(fracts_write_dir, exist_ok=True) + _sweep_stale_temp_files(fracts_write_dir, log) # Wait until dir setup completes comm.Barrier() diff --git a/ScaFFold/datagen/layout.py b/ScaFFold/datagen/layout.py index e558c6d..eff02f8 100644 --- a/ScaFFold/datagen/layout.py +++ b/ScaFFold/datagen/layout.py @@ -57,6 +57,41 @@ def category_param_dir(config) -> str: return os.path.join(library_root(config), "3DIFS_param") +def legacy_category_param_dir(config) -> str: + """Return where the category CSVs lived before the layout was seed-keyed.""" + return os.path.join( + str(config.fract_base_dir), + f"var{config.variance_threshold}", + "3DIFS_param", + ) + + +def warn_if_legacy_library(config, log) -> bool: + """Warn when a library in the old, seed-agnostic layout is being ignored. + + The relayout is deliberately silent about old data -- an existence check in + the seed-keyed location simply does not find it -- which from the outside + looks like a library that was there yesterday being regenerated for no + reason (at large scales, hours of work). One line naming both directories + turns that into an explained, expected event. Returns whether the old + layout was present, so callers can test the condition directly. + """ + legacy = legacy_category_param_dir(config) + if not os.path.isdir(legacy): + return False + log.warning( + "Found a fractal library in the old, seed-agnostic layout at %s. " + "Libraries are now keyed by seed, so this one cannot be reused (a run " + "under a different seed would silently adopt another seed's data) and " + "the categories for seed %s will be generated at %s. Delete the old " + "directory once you no longer need it.", + legacy, + int(config.seed), + category_param_dir(config), + ) + return True + + def instance_dir(config) -> str: """Return the directory holding this seed's instance point clouds. diff --git a/tests/datagen/test_category_search.py b/tests/datagen/test_category_search.py index 1b861c3..3ff19a2 100644 --- a/tests/datagen/test_category_search.py +++ b/tests/datagen/test_category_search.py @@ -193,6 +193,78 @@ def test_divergent_fs_views_take_the_same_collective_path(tmp_path, monkeypatch) ) +# --------------------------------------------------------------------------- +# VB-4: temp files stranded by killed writes are swept, as in instance.py. +# +# Both atomic writers name their temp file after the writing pid, so a job +# killed mid-write leaves one behind per killed process, forever: nothing in +# this module ever looked at them again. +# --------------------------------------------------------------------------- + + +def test_stale_temp_files_are_swept(tmp_path, monkeypatch): + """Category and attempt-counter temps from dead pids are removed.""" + config = _cs_config(tmp_path / "fractals") + param_dir = Path(layout.category_param_dir(config)) + param_dir.mkdir(parents=True, exist_ok=True) + _seed_one_category(config) # n_categories=1, so the search has no work + + stale_csv = param_dir / ".000001.csv.tmp999999" + stale_csv.write_text("0.5,0.5\n") + stale_counter = param_dir / ".rng_attempt_rank3.tmp999999" + stale_counter.write_text("17") + + comm = CategorySearchComm(rank=0, size=1) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + cs.main(config) + + assert not stale_csv.exists(), "a stranded category temp file was kept" + assert not stale_counter.exists(), "a stranded attempt-counter temp was kept" + # The real artifact is untouched: only the temp names are swept. + assert (param_dir / "000000.csv").exists() + + +# --------------------------------------------------------------------------- +# VB-6: a library in the old, seed-agnostic layout is explained, not ignored. +# --------------------------------------------------------------------------- + + +def test_old_layout_library_is_reported(tmp_path, monkeypatch, caplog): + """A pre-relayout library produces one warning naming both directories.""" + config = _cs_config(tmp_path / "fractals") + legacy = Path(layout.legacy_category_param_dir(config)) + legacy.mkdir(parents=True) + (legacy / "000000.csv").write_text("") + _seed_one_category(config) # nothing to generate under the new layout + + comm = CategorySearchComm(rank=0, size=1) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + with caplog.at_level("WARNING"): + cs.main(config) + + messages = " ".join(record.getMessage() for record in caplog.records) + assert str(legacy) in messages + assert layout.category_param_dir(config) in messages + + +def test_no_warning_without_an_old_layout(tmp_path, monkeypatch, caplog): + """The warning does not fire for a fresh library (control).""" + config = _cs_config(tmp_path / "fractals") + _seed_one_category(config) + + comm = CategorySearchComm(rank=0, size=1) + monkeypatch.setattr(cs, "MPI", FakeMPI(comm)) + + with caplog.at_level("WARNING"): + cs.main(config) + + assert "seed-agnostic" not in " ".join( + record.getMessage() for record in caplog.records + ) + + # --------------------------------------------------------------------------- # VB-2: the whole rank-0 scan window is fenced, not just the index parse. # From fcdcee40ccfa72d48e47064830a6baa832425fc1 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:00:43 -0700 Subject: [PATCH 38/54] Resolve restart state before the pre-check restart and run_dir can come from the config file as well as the command line -- the generated restart.sh replays a dumped config.yaml, which carries both -- and an absent --restart cannot outrank a file that sets it, since an unset store_true flag is indistinguishable from its default. resolve_run_dir read only the command line while the pre-check read the merged config, so the two disagreed in both directions. A config-file `restart: true` with no run directory therefore created a fresh timestamped directory, wrote its config dumps and restart script into it, and only then died for want of a run dir; it now fails before anything is claimed. And a run_dir inherited from a reused config.yaml stayed in the config while a fresh directory was created to train in, so the pre-check passed on another run's checkpoints and the job died hours later with its dataset generated. resolve_run_dir now reads the merged values, and writes its answer back to benchmark_run_dir, restart and run_dir (cleared for a fresh run), so nothing downstream can re-derive a different one; missing_checkpoint_error keys off the resolved benchmark run dir. The pre-check is also gated on the benchmark subcommand, which is the only one that has run directories at all. VC-1 --- ScaFFold/cli.py | 73 +++++++++++++++++++++++++++++++++------------ tests/test_cli.py | 75 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 19 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 623f868..6a3923b 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -121,8 +121,22 @@ def missing_checkpoint_error(combined_config): Reports the first problem found rather than raising, so the caller can make this a rank-0 decision and broadcast the verdict instead of letting every rank stat the shared filesystem and possibly disagree. + + The directory checked is the *resolved* benchmark run dir -- the one this + launch will actually train in -- and not the raw ``run_dir`` key. The two + can differ: a config.yaml dumped by a restarted run carries that run's + ``run_dir``, so reusing it as a base config had this check stat another + run's checkpoints, pass, and let the job die hours later with its dataset + already generated. ``resolve_run_dir`` keeps the two in agreement; this + prefers the resolved value so they cannot drift apart again. """ - checkpoint_dir = Path(combined_config["run_dir"]) / combined_config.get( + run_dir = combined_config.get("benchmark_run_dir") or combined_config.get("run_dir") + if not run_dir: + return ( + "Restart requested but no run directory was resolved. Pass " + "'--run-dir ' (or set run_dir in the config file)." + ) + checkpoint_dir = Path(run_dir) / combined_config.get( "checkpoint_dir", "checkpoints" ) expected_checkpoints = ( @@ -140,20 +154,34 @@ def resolve_run_dir(args_dict, combined_config): The semantics are fixed and unambiguous: - * ``--run-dir DIR`` (with or without ``--restart``): resume in that exact - directory. ``train_from_scratch`` is forced off and ``restart`` on so the - downstream benchmark driver takes its restart path. - * ``--restart`` without ``--run-dir``: rejected with a clear error. The run - directory to resume must be named explicitly; the most recent directory - is never guessed. - * neither flag: create a fresh timestamped directory under ``base_run_dir``, + * a run directory (``--run-dir DIR``, or ``run_dir`` in the config file), + with or without a restart flag: resume in that exact directory. + ``train_from_scratch`` is forced off and ``restart`` on so the downstream + benchmark driver takes its restart path. + * a restart requested with no run directory anywhere: rejected with a clear + error, before any directory is created. The directory to resume must be + named explicitly; the most recent one is never guessed. + * neither: create a fresh timestamped directory under ``base_run_dir``, retrying with a numeric suffix on a same-second name collision. - ``combined_config['benchmark_run_dir']`` is set in every path so the driver - can always read it. Returns ``(benchmark_run_dir: Path, restarting: bool)``. + Both keys are read from the *merged* config, not from the command line + alone. A config file is a first-class source for them -- the generated + ``restart.sh`` replays a dumped ``config.yaml``, which carries both -- and + an absent ``--restart`` cannot outrank a file that sets it, because an + unset ``store_true`` flag is indistinguishable from its default. Reading + only the command line here let the two disagree: a config-file restart + created a *fresh* run directory and only then failed for want of a run dir, + and a stale ``run_dir`` inherited from a reused ``config.yaml`` was left in + the config for the restart pre-check to stat while training happened + somewhere else entirely. + + The resolved answer is written back to ``benchmark_run_dir``, ``restart`` + and ``run_dir``, so every later reader -- the pre-check, the dumped + ``config.yaml``, the benchmark driver -- sees exactly what was decided here. + Returns ``(benchmark_run_dir: Path, restarting: bool)``. """ - restart_flag = bool(args_dict.get("restart")) - run_dir_arg = args_dict.get("run_dir") + restart_flag = bool(combined_config.get("restart") or args_dict.get("restart")) + run_dir_arg = combined_config.get("run_dir") or args_dict.get("run_dir") if run_dir_arg is not None: benchmark_run_dir = Path(run_dir_arg) @@ -163,9 +191,10 @@ def resolve_run_dir(args_dict, combined_config): restarting = True elif restart_flag: raise ValueError( - "--restart requires --run-dir: pass the directory of the run to " - "resume (e.g. '--restart --run-dir '). The most recent run " - "directory is not resolved automatically." + "A restart was requested (--restart, or restart: true in the " + "config file) but no run directory was given: pass the directory " + "of the run to resume (e.g. '--restart --run-dir '). The " + "most recent run directory is not resolved automatically." ) else: base_run_dir = Path(combined_config["base_run_dir"]) @@ -179,10 +208,15 @@ def resolve_run_dir(args_dict, combined_config): ) restarting = False + # Write the resolution back, so nothing downstream can re-derive a + # different answer from the raw inputs. ``run_dir`` is cleared for a fresh + # run: left set, a value inherited from a reused config.yaml names a + # directory this run has nothing to do with. combined_config["benchmark_run_dir"] = str(benchmark_run_dir) + combined_config["restart"] = restarting + combined_config["run_dir"] = str(benchmark_run_dir) if restarting else None if restarting: combined_config["train_from_scratch"] = False - combined_config["restart"] = True return benchmark_run_dir, restarting @@ -528,10 +562,11 @@ def main(): # caches). A rank that decided for itself would either abort alone -- # stranding its peers in benchmark.py's timeout-less barrier -- or keep # running after rank 0 had already aborted. + # Only the benchmark subcommand has run directories or checkpoints, so it + # is the only one this applies to: fractal generation reading a benchmark's + # config.yaml must not be judged on that run's restart state. restart_precheck_error = None - if combined_config.get("restart", False): - if not combined_config.get("run_dir"): - raise ValueError("--restart requires --run-dir") + if args.command == "benchmark" and combined_config.get("restart", False): if rank == 0: restart_precheck_error = missing_checkpoint_error(combined_config) restart_precheck_error = comm.bcast(restart_precheck_error, root=0) diff --git a/tests/test_cli.py b/tests/test_cli.py index 381744f..1db35af 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -314,6 +314,81 @@ def test_restart_precheck_failure_raises_on_every_rank(monkeypatch, tmp_path): run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) +# --------------------------------------------------------------------------- +# VC-1: the restart state is resolved once, before anything acts on it. +# +# ``restart`` and ``run_dir`` can come from the config file as well as the +# command line (the generated restart.sh replays a dumped config.yaml, which +# carries both), and an absent ``--restart`` cannot outrank a file that sets +# it. Resolving the run directory from the command line alone while the +# pre-check read the merged config made the two disagree. +# --------------------------------------------------------------------------- + + +def test_yaml_restart_without_run_dir_fails_before_creating_a_run_dir( + monkeypatch, tmp_path +): + """``restart: true`` with no run dir aborts without claiming a directory. + + The run directory was resolved from the command line, which said nothing + about a restart, so a fresh timestamped directory was created and populated + -- and only then did the merged config's ``restart`` trip the pre-check. + """ + cfg = write_config(tmp_path, {"restart": True}) + + with pytest.raises(ValueError, match="run directory"): + run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)]) + + assert not (tmp_path / "runs").exists(), "a run dir was created before the abort" + + +def test_reused_restart_config_resolves_to_one_run_dir(monkeypatch, tmp_path): + """A config.yaml from a restarted run cannot split the run across two dirs. + + Reusing such a file as a base config left ``run_dir`` pointing at the run + it was dumped by while a *fresh* directory was created to train in. The + pre-check then passed on the old run's checkpoints, and the job died hours + later with its dataset already generated. Whatever the file resolves to, + the directory the pre-check judges and the directory the run uses must be + the same one. + """ + cfg = write_config(tmp_path) + _, first = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)]) + first_run = Path(first["benchmark"][0]["benchmark_run_dir"]) + _make_checkpoint(first_run) + + # Restart it once, so its config.yaml records the restart state. + run_cli(monkeypatch, _restart_argv(first_run / "config.yaml", first_run)) + runs_before = sorted(p.name for p in (tmp_path / "runs").iterdir()) + + # Now reuse that config.yaml as a plain base config, with no flags at all. + _, reused = run_cli( + monkeypatch, ["scaffold", "benchmark", "-c", str(first_run / "config.yaml")] + ) + + (config,) = reused["benchmark"] + assert config["run_dir"] == config["benchmark_run_dir"], ( + "the pre-check's run_dir and the run's directory disagree" + ) + assert Path(config["benchmark_run_dir"]) == first_run + assert sorted(p.name for p in (tmp_path / "runs").iterdir()) == runs_before + + +def test_fresh_run_records_no_restart_state(monkeypatch, tmp_path): + """A fresh run's config.yaml carries no restart state to inherit.""" + cfg = write_config(tmp_path) + + _, calls = run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)]) + + (config,) = calls["benchmark"] + assert config["restart"] is False + assert config["run_dir"] is None + dumped = yaml.safe_load( + (Path(config["benchmark_run_dir"]) / "config.yaml").read_text() + ) + assert dumped["restart"] is False and dumped["run_dir"] is None + + # --------------------------------------------------------------------------- # R19: generate_fractals is not a benchmark run # --------------------------------------------------------------------------- From e338e6ae617fd93c2f112f30eecfbfe2c5ff74e7 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:03:28 -0700 Subject: [PATCH 39/54] Broadcast config validation outcomes Rank 0 builds the whole job's config while every peer waits in a barrier, so anything raising in there stranded them: the recomputed bottleneck check and the n_categories check added with R25, the Config() validation that has always been there, and now the run-dir resolution. The user saw a hang where an error message belonged. The rank-0 block moves into build_run_config and runs inside a guard whose outcome is broadcast, exactly like the restart pre-check below it. Rank 0 re-raises the original exception, keeping its traceback; the peers rebuild it from the type name and message that crossed the wire, so a builtin type comes back as itself and anything else degrades to a RuntimeError naming the original. The exception object itself is deliberately not pickled across: a failure to unpickle on the receiving side would turn the error being reported back into the hang it was reported to avoid. VC-2 --- ScaFFold/cli.py | 262 ++++++++++++++++++++++++++++------------------ tests/test_cli.py | 66 +++++++++++- 2 files changed, 222 insertions(+), 106 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 6a3923b..88bf1f9 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -13,6 +13,7 @@ # SPDX-License-Identifier: (Apache-2.0) import argparse +import builtins import socket import sys from datetime import datetime @@ -220,6 +221,142 @@ def resolve_run_dir(args_dict, combined_config): return benchmark_run_dir, restarting +def rebuild_error(type_name, message): + """Rebuild rank 0's configuration error on a peer that never saw it. + + Only the type name and message cross the wire: an arbitrary exception + object may not survive a pickle round trip, and a failure to unpickle on + the receiving side would turn the error being reported back into the hang + it was reported to avoid. Builtin exception types are rebuilt as + themselves, so a caller's ``except ValueError`` still catches what rank 0 + raised; anything else degrades to ``RuntimeError`` naming the original + type. + """ + cls = getattr(builtins, type_name, None) + if isinstance(cls, type) and issubclass(cls, Exception): + return cls(message) + return RuntimeError(f"{type_name}: {message}") + + +def build_run_config(args, parsers, log, world_size): + """Build the job-wide config, and lay down the benchmark's run directory. + + Rank-0-only work: it reads and merges the config files, applies the + command-line overrides, validates the result, and -- for the benchmark + subcommand -- resolves the run directory, dumps the configs into it and + writes its restart script. Returns the merged config the caller broadcasts. + + Raises whatever the validation or the filesystem raises; the caller runs + this inside a guard and broadcasts the outcome, since every peer is already + waiting in the barrier that follows. + """ + log.debug("args = %s", args) + + # --config may be a single path (generate_fractals) or a list of + # paths (benchmark, action="append"): base config plus overrides. + config_paths = args.config if isinstance(args.config, list) else [args.config] + merged_dict = config_utils.load_config_files(config_paths) + # Validate the merged result and derive dependent settings. Every run + # parameter must be single-valued; a list is rejected here by name. + bench_config = config_utils.Config(merged_dict) + bench_config_dict = vars(bench_config) + cli_args = vars(args) + # Downstream consumers expect a single config path (e.g. to copy it + # into the run dir); keep the base config there. + cli_args["config"] = config_paths[0] + + # Combine configs, in increasing order of precedence: + # argparse default < config file < explicit command-line flag. + combined_config = bench_config_dict.copy() + # Config only keeps the keys it consumes; the auxiliary keys it accepts + # (verbose, datagen_batch_size, ...) never become attributes, so put + # the file's values back first. Without this they are absent below and + # the argparse default overwrites what the user wrote in the config. + for key, value in merged_dict.items(): + combined_config.setdefault(key, value) + + explicit_cli = explicit_cli_keys(args, parsers) + for key, value in cli_args.items(): + if key == "command": + continue + if key not in combined_config: + combined_config[key] = value + elif key in explicit_cli and value is not None: + log.info( + "Overriding '%s=%s' with '%s=%s'", + key, + combined_config[key], + key, + value, + ) + combined_config[key] = value + # The subcommand is always owned by the command line. + combined_config["command"] = cli_args["command"] + + # Recalculate unet_layers to capture any CLI overrides. The overridden + # pair has to be re-validated: Config only saw the config-file values. + config_utils.validate_unet_dims( + combined_config["problem_scale"], combined_config["unet_bottleneck_dim"] + ) + combined_config["unet_layers"] = ( + combined_config["problem_scale"] - combined_config["unet_bottleneck_dim"] + ) + config_utils.require_positive_int("n_categories", combined_config["n_categories"]) + + # Resolve paths to absolute, matching Config() behavior + if "base_run_dir" in combined_config and combined_config["base_run_dir"]: + combined_config["base_run_dir"] = str( + Path(combined_config["base_run_dir"]).resolve() + ) + + if "dataset_dir" in combined_config and combined_config["dataset_dir"]: + combined_config["dataset_dir"] = str( + Path(combined_config["dataset_dir"]).resolve() + ) + + if "fract_base_dir" in combined_config and combined_config["fract_base_dir"]: + combined_config["fract_base_dir"] = str( + Path(combined_config["fract_base_dir"]).resolve() + ) + + # Calculate these variables after override + combined_config["vol_size"] = pow(2, combined_config["problem_scale"]) + combined_config["point_num"] = int(combined_config["vol_size"] ** 3 / 256) + + # The run directory, its config dumps and its restart script belong to + # the benchmark subcommand alone. Fractal generation writes nothing + # there, and the restart script it used to get replayed + # `generate_fractals --restart --run-dir ...` -- flags that subparser + # rejects, so the script could only ever exit 2. + if args.command == "benchmark": + # Resolve the run directory and whether this launch resumes a run. + # This sets combined_config["benchmark_run_dir"] on every path, and + # writes the resolved restart/run_dir back into the config. + benchmark_run_dir, restarting = resolve_run_dir(cli_args, combined_config) + if restarting: + log.info("Resuming in existing directory: %s", benchmark_run_dir) + + # Add scheduler metadata and machine name to config.yaml + combined_config["scheduler_metadata"] = collect_scheduler_metadata() + combined_config["machine_name"] = socket.gethostname() + + # Dump configs (Overwrite is okay/desired on restart to capture new job IDs) + overrides = { + k: v for k, v in cli_args.items() if v is not None and k != "command" + } + with open(benchmark_run_dir / "overrides.yaml", "w") as file: + yaml.dump(overrides, file) + with open(benchmark_run_dir / "config.yaml", "w") as file: + yaml.dump(combined_config, file) + + # 4. Generate/Update the restart script in the directory. The + # communicator size is ground truth for the job scale; environment + # sniffing is only the fallback for callers that lack it. + create_restart_script(benchmark_run_dir, world_size=world_size) + + return combined_config + + def main(): """ Command line interface for ScaFFold. @@ -446,114 +583,31 @@ def main(): "resume (e.g. '--restart --run-dir ')." ) + # Rank 0 builds the config for the whole job while every other rank waits + # in the barrier below. Everything in there can fail on user input (an + # unknown config key, an out-of-range bottleneck, a restart with no run + # dir) or on the filesystem, and a rank-0-only raise leaves the peers + # blocked in that barrier -- a hang instead of the error message the user + # needs. The outcome is therefore broadcast, exactly like the restart + # pre-check below, and every rank raises the same error together. + config_error = None + rank0_error = None if rank == 0: - log.debug("args = %s", args) - - # --config may be a single path (generate_fractals) or a list of - # paths (benchmark, action="append"): base config plus overrides. - config_paths = args.config if isinstance(args.config, list) else [args.config] - merged_dict = config_utils.load_config_files(config_paths) - # Validate the merged result and derive dependent settings. Every run - # parameter must be single-valued; a list is rejected here by name. - bench_config = config_utils.Config(merged_dict) - bench_config_dict = vars(bench_config) - cli_args = vars(args) - # Downstream consumers expect a single config path (e.g. to copy it - # into the run dir); keep the base config there. - cli_args["config"] = config_paths[0] - - # Combine configs, in increasing order of precedence: - # argparse default < config file < explicit command-line flag. - combined_config = bench_config_dict.copy() - # Config only keeps the keys it consumes; the auxiliary keys it accepts - # (verbose, datagen_batch_size, ...) never become attributes, so put - # the file's values back first. Without this they are absent below and - # the argparse default overwrites what the user wrote in the config. - for key, value in merged_dict.items(): - combined_config.setdefault(key, value) - - explicit_cli = explicit_cli_keys(args, (active_parser, parser)) - for key, value in cli_args.items(): - if key == "command": - continue - if key not in combined_config: - combined_config[key] = value - elif key in explicit_cli and value is not None: - log.info( - "Overriding '%s=%s' with '%s=%s'", - key, - combined_config[key], - key, - value, - ) - combined_config[key] = value - # The subcommand is always owned by the command line. - combined_config["command"] = cli_args["command"] - - # Recalculate unet_layers to capture any CLI overrides. The overridden - # pair has to be re-validated: Config only saw the config-file values. - config_utils.validate_unet_dims( - combined_config["problem_scale"], combined_config["unet_bottleneck_dim"] - ) - combined_config["unet_layers"] = ( - combined_config["problem_scale"] - combined_config["unet_bottleneck_dim"] - ) - config_utils.require_positive_int( - "n_categories", combined_config["n_categories"] - ) - - # Resolve paths to absolute, matching Config() behavior - if "base_run_dir" in combined_config and combined_config["base_run_dir"]: - combined_config["base_run_dir"] = str( - Path(combined_config["base_run_dir"]).resolve() - ) - - if "dataset_dir" in combined_config and combined_config["dataset_dir"]: - combined_config["dataset_dir"] = str( - Path(combined_config["dataset_dir"]).resolve() - ) - - if "fract_base_dir" in combined_config and combined_config["fract_base_dir"]: - combined_config["fract_base_dir"] = str( - Path(combined_config["fract_base_dir"]).resolve() + try: + combined_config = build_run_config( + args, (active_parser, parser), log, comm.Get_size() ) - - # Calculate these variables after override - combined_config["vol_size"] = pow(2, combined_config["problem_scale"]) - combined_config["point_num"] = int(combined_config["vol_size"] ** 3 / 256) - - # The run directory, its config dumps and its restart script belong to - # the benchmark subcommand alone. Fractal generation writes nothing - # there, and the restart script it used to get replayed - # `generate_fractals --restart --run-dir ...` -- flags that subparser - # rejects, so the script could only ever exit 2. - if args.command == "benchmark": - # Resolve the run directory and whether this launch resumes a run. - # This sets combined_config["benchmark_run_dir"] on every path and, - # when resuming, forces train_from_scratch off / restart on. - benchmark_run_dir, restarting = resolve_run_dir(vars(args), combined_config) - if restarting: - log.info("Resuming in existing directory: %s", benchmark_run_dir) - - # Add scheduler metadata and machine name to config.yaml - combined_config["scheduler_metadata"] = collect_scheduler_metadata() - combined_config["machine_name"] = socket.gethostname() - - # Dump configs (Overwrite is okay/desired on restart to capture new job IDs) - overrides = { - k: v for k, v in cli_args.items() if v is not None and k != "command" - } - with open(benchmark_run_dir / "overrides.yaml", "w") as file: - yaml.dump(overrides, file) - with open(benchmark_run_dir / "config.yaml", "w") as file: - yaml.dump(combined_config, file) - - # 4. Generate/Update the restart script in the directory. The - # communicator size is ground truth for the job scale; environment - # sniffing is only the fallback for callers that lack it. - create_restart_script(benchmark_run_dir, world_size=comm.Get_size()) + except Exception as e: + combined_config = None + rank0_error = e + config_error = (type(e).__name__, str(e)) comm.Barrier() + config_error = comm.bcast(config_error, root=0) + if config_error is not None: + # Rank 0 re-raises the original (keeping its traceback); the peers + # rebuild it from what crossed the wire. + raise rank0_error if rank0_error is not None else rebuild_error(*config_error) combined_config = comm.bcast(combined_config, root=0) # Restart pre-check. Like every other decision here it is made once, on diff --git a/tests/test_cli.py b/tests/test_cli.py index 1db35af..e91955a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -284,7 +284,9 @@ def test_restart_precheck_follows_the_broadcast_decision(monkeypatch, tmp_path): "verbose": 0, } - comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_config, None]) + # Rank 0 broadcasts, in order: no config error, the config, no pre-check + # error. + comm = _FakeComm(rank=1, size=2, bcast_returns=[None, rank0_config, None]) _, calls = run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) assert len(calls["benchmark"]) == 1 @@ -309,7 +311,7 @@ def test_restart_precheck_failure_raises_on_every_rank(monkeypatch, tmp_path): } rank0_error = "Restart requested but no checkpoint was found. Expected /nope." - comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_config, rank0_error]) + comm = _FakeComm(rank=1, size=2, bcast_returns=[None, rank0_config, rank0_error]) with pytest.raises(FileNotFoundError, match="no checkpoint"): run_cli(monkeypatch, _restart_argv(cfg, run_dir), comm=comm) @@ -541,6 +543,66 @@ def test_cli_override_bottleneck_out_of_range_rejected(monkeypatch, tmp_path): assert "problem_scale" in message +# --------------------------------------------------------------------------- +# VC-2: a rank-0 config failure is broadcast, not left to the barrier. +# +# Rank 0 builds the whole job's config while every peer waits in a barrier, so +# anything that raises in there -- the config-file validation, the recomputed +# bottleneck check, the run-dir resolution -- has to travel to the peers as a +# decision. Otherwise the user gets a hang instead of the error message. +# --------------------------------------------------------------------------- + + +def test_config_failure_is_broadcast_before_the_barrier(monkeypatch, tmp_path): + """Rank 0 posts its collectives and broadcasts the error it hit.""" + cfg = write_config(tmp_path) + comm = _FakeComm(rank=0, size=2) + + with pytest.raises(ValueError): + run_cli( + monkeypatch, + [ + "scaffold", + "benchmark", + "-c", + str(cfg), + "--problem-scale", + "4", + "--unet-bottleneck-dim", + "4", + ], + comm=comm, + ) + + # The peers' barrier was matched, and what they receive names the failure. + assert comm.barriers == 1 + errors = [ + payload + for payload in comm.broadcast + if isinstance(payload, tuple) and payload[0] == "ValueError" + ] + assert errors, f"no error sentinel was broadcast: {comm.broadcast}" + assert "unet_bottleneck_dim" in errors[0][1] + + +def test_peer_raises_the_broadcast_config_error(monkeypatch, tmp_path): + """A peer rebuilds rank 0's error instead of running with no config.""" + cfg = write_config(tmp_path) + rank0_error = ("ValueError", "unet_bottleneck_dim (4) must be < problem_scale (4)") + + comm = _FakeComm(rank=1, size=2, bcast_returns=[rank0_error]) + with pytest.raises(ValueError, match="unet_bottleneck_dim"): + run_cli(monkeypatch, ["scaffold", "benchmark", "-c", str(cfg)], comm=comm) + + +def test_unknown_error_types_degrade_to_runtime_error(): + """A non-builtin exception type is still reported, as a RuntimeError.""" + assert type(cli.rebuild_error("FileNotFoundError", "gone")) is FileNotFoundError + rebuilt = cli.rebuild_error("SomeSiteSpecificError", "boom") + assert isinstance(rebuilt, RuntimeError) + assert "SomeSiteSpecificError" in str(rebuilt) and "boom" in str(rebuilt) + + # --------------------------------------------------------------------------- # The whole config path survives a restart (R20/R22 together) # --------------------------------------------------------------------------- From f27aa2d23aae6911c038445b22baeef4fecbb55a Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:06:20 -0700 Subject: [PATCH 40/54] Harden launcher env parsing and CLI entry ordering The launcher-environment helpers called int() on whatever a variable held. A site wrapper that exports WORLD_SIZE= (an unset shell variable) or a placeholder like "auto" therefore killed every scaffold invocation with a bare ValueError naming neither the variable nor a remedy -- and it fired from the top of the entry point, so even `scaffold --help` died. Each lookup now goes through a helper that treats an unusable value as absent and consults the next source in the priority order, warning when the value looked deliberate (as _sniff_launch_shape already did with `if val:`). The Slurm and Flux tasks-per-node divisions no longer trust the node count either: zero is as unusable as a word. The world-size cross-check also moved after parse_args. Asking what the flags are is not a job launch, and answering it with a launcher mismatch -- exactly the environment someone debugging one is sitting in -- helps nobody. It still runs before any run directory is created, which is the property R13 needs. VC-3, VC-4 --- ScaFFold/cli.py | 12 ++- ScaFFold/utils/distributed.py | 165 ++++++++++++++++++++++------------ tests/test_cli.py | 22 +++++ tests/test_worker_dist.py | 53 +++++++++++ 4 files changed, 190 insertions(+), 62 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 88bf1f9..650a126 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -555,12 +555,16 @@ def main(): comm = MPI.COMM_WORLD rank = comm.Get_rank() - # Every rank runs this identically, before any run directory is created, - # so a mis-launched job aborts uniformly instead of leaving per-rank run - # dirs behind and hanging. - check_launcher_world_size(comm.Get_size()) # Parse the command-line arguments. args = parser.parse_args() + # Every rank runs this identically, before any run directory is created, so + # a mis-launched job aborts uniformly instead of leaving per-rank run dirs + # behind and hanging. It runs *after* parsing so that the arguments argparse + # handles by itself -- ``--help``, a usage error -- still behave: asking + # what the flags are is not a job launch, and answering it with a launcher + # mismatch (which is exactly the environment someone debugging one is + # sitting in) helps nobody. + check_launcher_world_size(comm.Get_size()) subcommand_parsers = { "benchmark": benchmark_parser, "generate_fractals": generate_fractals_parser, diff --git a/ScaFFold/utils/distributed.py b/ScaFFold/utils/distributed.py index afad5ab..5cb0700 100644 --- a/ScaFFold/utils/distributed.py +++ b/ScaFFold/utils/distributed.py @@ -12,6 +12,7 @@ # # SPDX-License-Identifier: (Apache-2.0) +import logging import os import os.path import socket @@ -21,6 +22,52 @@ import torch import torch.distributed +logger = logging.getLogger(__name__) + + +def _env_int(name: str) -> Optional[int]: + """Return the launcher variable ``name`` as an int, or None if unusable. + + Launcher variables are not always what they claim to be. A site wrapper + that exports ``WORLD_SIZE=`` (empty, e.g. from an unset shell variable) or + a placeholder like ``auto`` is common enough, and a bare ``int()`` turned + it into a ``ValueError`` raised from the first of these helpers anything + called -- killing every ``scaffold`` invocation, ``--help`` included, with + a traceback that named neither the variable nor a remedy. + + An unusable value is treated as absent so the next source in the priority + order is consulted (ultimately the MPI communicator, or the documented + default). A non-empty value that is not an integer is warned about, because + unlike an empty one it looks deliberate and the fallback may not be what + its author intended. ``create_restart_script._sniff_launch_shape`` skips + empty values for the same reason. + """ + value = os.environ.get(name) + if value is None: + return None + value = value.strip() + if not value: + return None + try: + return int(value) + except ValueError: + logger.warning( + "Ignoring launcher variable %s=%r: not an integer. Falling back to " + "the next source for the job shape.", + name, + value, + ) + return None + + +def _first_env_int(names) -> Optional[int]: + """Return the first usable integer among ``names``, in priority order.""" + for name in names: + value = _env_int(name) + if value is not None: + return value + return None + def get_num_gpus() -> int: """Return the number of GPUs on this node.""" @@ -40,22 +87,50 @@ def _mpi_comm_world(): return None +_LOCAL_RANK_VARS = ( + "LOCAL_RANK", + "MV2_COMM_WORLD_LOCAL_RANK", + "OMPI_COMM_WORLD_LOCAL_RANK", + "PMI_LOCAL_RANK", + "PALS_LOCAL_RANKID", + "SLURM_LOCALID", + "FLUX_TASK_LOCAL_ID", +) + +_LOCAL_SIZE_VARS = ( + "LOCAL_WORLD_SIZE", + "MV2_COMM_WORLD_LOCAL_SIZE", + "OMPI_COMM_WORLD_LOCAL_SIZE", + "PMI_LOCAL_SIZE", + "PALS_LOCAL_SIZE", +) + +_WORLD_RANK_VARS = ( + "RANK", + "MV2_COMM_WORLD_RANK", + "OMPI_COMM_WORLD_RANK", + "PMI_RANK", + "PALS_RANKID", + "SLURM_PROCID", + "FLUX_TASK_RANK", +) + +_WORLD_SIZE_VARS = ( + "WORLD_SIZE", + "MV2_COMM_WORLD_SIZE", + "OMPI_COMM_WORLD_SIZE", + "PMI_SIZE", + "PALS_NRANKS", + "SLURM_NTASKS", + "FLUX_JOB_SIZE", +) + + def get_local_rank(required: bool = False) -> int: """Return the local MPI rank.""" - if "LOCAL_RANK" in os.environ: - return int(os.environ["LOCAL_RANK"]) - if "MV2_COMM_WORLD_LOCAL_RANK" in os.environ: - return int(os.environ["MV2_COMM_WORLD_LOCAL_RANK"]) - if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ: - return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]) - if "PMI_LOCAL_RANK" in os.environ: - return int(os.environ["PMI_LOCAL_RANK"]) - if "PALS_LOCAL_RANKID" in os.environ: - return int(os.environ["PALS_LOCAL_RANKID"]) - if "SLURM_LOCALID" in os.environ: - return int(os.environ["SLURM_LOCALID"]) - if "FLUX_TASK_LOCAL_ID" in os.environ: - return int(os.environ["FLUX_TASK_LOCAL_ID"]) + value = _first_env_int(_LOCAL_RANK_VARS) + if value is not None: + return value if required: raise RuntimeError("Could not get local rank") return 0 @@ -68,22 +143,18 @@ def get_local_size(required: bool = False) -> int: there but not here silently yields 1, which makes per-node logic (e.g. the profiler's one-rank-per-node gate) treat every rank as node-local. """ - if "LOCAL_WORLD_SIZE" in os.environ: - return int(os.environ["LOCAL_WORLD_SIZE"]) - if "MV2_COMM_WORLD_LOCAL_SIZE" in os.environ: - return int(os.environ["MV2_COMM_WORLD_LOCAL_SIZE"]) - if "OMPI_COMM_WORLD_LOCAL_SIZE" in os.environ: - return int(os.environ["OMPI_COMM_WORLD_LOCAL_SIZE"]) - if "PMI_LOCAL_SIZE" in os.environ: - return int(os.environ["PMI_LOCAL_SIZE"]) - if "PALS_LOCAL_SIZE" in os.environ: - return int(os.environ["PALS_LOCAL_SIZE"]) - if "SLURM_NNODES" in os.environ and "SLURM_NTASKS" in os.environ: - return int(os.environ["SLURM_NTASKS"]) // int(os.environ["SLURM_NNODES"]) - # Flux does not have an env variable for this, so we assume an - # even distribution. - if "FLUX_JOB_SIZE" in os.environ and "FLUX_JOB_NNODES" in os.environ: - return int(os.environ["FLUX_JOB_SIZE"]) // int(os.environ["FLUX_JOB_NNODES"]) + value = _first_env_int(_LOCAL_SIZE_VARS) + if value is not None: + return value + # Slurm and Flux report only totals; assume an even distribution. A zero + # node count is as unusable as a non-numeric one, so it falls through + # rather than raising ZeroDivisionError. + ntasks, nnodes = _env_int("SLURM_NTASKS"), _env_int("SLURM_NNODES") + if ntasks is not None and nnodes: + return ntasks // nnodes + job_size, job_nodes = _env_int("FLUX_JOB_SIZE"), _env_int("FLUX_JOB_NNODES") + if job_size is not None and job_nodes: + return job_size // job_nodes if required: raise RuntimeError("Could not get local size") return 1 @@ -91,20 +162,9 @@ def get_local_size(required: bool = False) -> int: def get_world_rank(required: bool = False) -> int: """Return the global MPI rank.""" - if "RANK" in os.environ: - return int(os.environ["RANK"]) - if "MV2_COMM_WORLD_RANK" in os.environ: - return int(os.environ["MV2_COMM_WORLD_RANK"]) - if "OMPI_COMM_WORLD_RANK" in os.environ: - return int(os.environ["OMPI_COMM_WORLD_RANK"]) - if "PMI_RANK" in os.environ: - return int(os.environ["PMI_RANK"]) - if "PALS_RANKID" in os.environ: - return int(os.environ["PALS_RANKID"]) - if "SLURM_PROCID" in os.environ: - return int(os.environ["SLURM_PROCID"]) - if "FLUX_TASK_RANK" in os.environ: - return int(os.environ["FLUX_TASK_RANK"]) + value = _first_env_int(_WORLD_RANK_VARS) + if value is not None: + return value comm = _mpi_comm_world() if comm is not None: return comm.Get_rank() @@ -115,20 +175,9 @@ def get_world_rank(required: bool = False) -> int: def get_world_size(required: bool = False) -> int: """Return the number of MPI ranks.""" - if "WORLD_SIZE" in os.environ: - return int(os.environ["WORLD_SIZE"]) - if "MV2_COMM_WORLD_SIZE" in os.environ: - return int(os.environ["MV2_COMM_WORLD_SIZE"]) - if "OMPI_COMM_WORLD_SIZE" in os.environ: - return int(os.environ["OMPI_COMM_WORLD_SIZE"]) - if "PMI_SIZE" in os.environ: - return int(os.environ["PMI_SIZE"]) - if "PALS_NRANKS" in os.environ: - return int(os.environ["PALS_NRANKS"]) - if "SLURM_NTASKS" in os.environ: - return int(os.environ["SLURM_NTASKS"]) - if "FLUX_JOB_SIZE" in os.environ: - return int(os.environ["FLUX_JOB_SIZE"]) + value = _first_env_int(_WORLD_SIZE_VARS) + if value is not None: + return value comm = _mpi_comm_world() if comm is not None: return comm.Get_size() diff --git a/tests/test_cli.py b/tests/test_cli.py index e91955a..f7d83ca 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -181,6 +181,28 @@ def test_matching_world_sizes_are_accepted(monkeypatch, tmp_path): assert len(calls["benchmark"]) == 1 +def test_help_works_under_a_mismatched_launcher_env(monkeypatch, tmp_path, capsys): + """``--help`` is answered even when the launcher environment disagrees. + + The cross-check ran before ``parse_args``, so asking what the flags are + raised the launcher-mismatch error -- in exactly the environment (a + half-configured shell) where someone is most likely to be asking. + """ + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + + with pytest.raises(SystemExit) as excinfo: + run_cli( + monkeypatch, + ["scaffold", "--help"], + comm=_FakeComm(rank=0, size=1), + sync_env=False, + ) + + assert excinfo.value.code == 0 + assert "usage" in capsys.readouterr().out.lower() + + def test_no_launcher_env_is_not_a_mismatch(monkeypatch, tmp_path): """With no launcher variables set, the MPI world alone defines the size.""" for var in ("WORLD_SIZE", "RANK", "LOCAL_RANK", "SLURM_NTASKS", "FLUX_JOB_SIZE"): diff --git a/tests/test_worker_dist.py b/tests/test_worker_dist.py index b1ae524..e3441ea 100644 --- a/tests/test_worker_dist.py +++ b/tests/test_worker_dist.py @@ -21,6 +21,7 @@ import logging +import pytest import torch import ScaFFold.utils.distributed as distributed_mod @@ -276,3 +277,55 @@ def test_local_size_defaults_to_one(monkeypatch): for var in _LOCAL_SIZE_VARS: monkeypatch.delenv(var, raising=False) assert distributed_mod.get_local_size() == 1 + + +# --------------------------------------------------------------------------- +# VC-3: an unusable launcher variable is ignored, not fatal +# +# These helpers run at the top of every entry point, so a bare int() on a +# variable a site wrapper exported empty ("WORLD_SIZE=") or as a placeholder +# ("auto") killed the invocation -- ``scaffold --help`` included -- with a +# ValueError naming neither the variable nor a remedy. +# --------------------------------------------------------------------------- + + +def _clear_launcher_env(monkeypatch): + for var in set(_LAUNCHER_VARS) | set(_LOCAL_SIZE_VARS): + monkeypatch.delenv(var, raising=False) + + +@pytest.mark.parametrize("value", ["", " ", "auto"]) +def test_unusable_launcher_values_fall_through(monkeypatch, value): + """Empty and non-numeric values are treated as absent, never raise.""" + _clear_launcher_env(monkeypatch) + for var in ("WORLD_SIZE", "RANK", "LOCAL_RANK", "LOCAL_WORLD_SIZE"): + monkeypatch.setenv(var, value) + + # Falls through to the next source -- here the (singleton) communicator and + # the documented defaults. + assert distributed_mod.get_world_size() == 1 + assert distributed_mod.get_world_rank() == 0 + assert distributed_mod.get_local_rank() == 0 + assert distributed_mod.get_local_size() == 1 + + +def test_unusable_value_defers_to_the_next_launcher_variable(monkeypatch, caplog): + """A garbage value does not mask a usable variable further down the order.""" + _clear_launcher_env(monkeypatch) + monkeypatch.setenv("WORLD_SIZE", "auto") + monkeypatch.setenv("PALS_NRANKS", "8") + + with caplog.at_level(logging.WARNING, logger=distributed_mod.logger.name): + assert distributed_mod.get_world_size() == 8 + + messages = " ".join(record.getMessage() for record in caplog.records) + assert "WORLD_SIZE" in messages, "the ignored value was not reported" + + +def test_zero_node_count_does_not_divide_by_zero(monkeypatch): + """A nonsense node count falls through instead of raising.""" + _clear_launcher_env(monkeypatch) + monkeypatch.setenv("SLURM_NTASKS", "8") + monkeypatch.setenv("SLURM_NNODES", "0") + + assert distributed_mod.get_local_size() == 1 From 9492d9ad636d5891e147194cc21114b5f1ab7de0 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:08:01 -0700 Subject: [PATCH 41/54] Derive restart node shape from local size Flux and Slurm state their node count, so the restart script reproduced their shape correctly. Everything else -- Cray PALS, plain torchrun -- fell back to NODES=1, which relaunched an 8-rank job spread over 2 nodes as NODES=1 TASKS_PER_NODE=8: an oversubscribed node, or a job the scheduler rejects. get_local_size already reads those launchers' per-node variables (PALS_LOCAL_ SIZE, LOCAL_WORLD_SIZE, ...), so the node count is derived from it, ceil-ing the division as the rank side does. required=True keeps "one rank per node" distinct from "nothing reported a per-node count"; only the latter falls back to the historical single-node assumption. VC-5 --- ScaFFold/utils/create_restart_script.py | 28 +++++++++++++-- tests/test_restart_script.py | 48 +++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/ScaFFold/utils/create_restart_script.py b/ScaFFold/utils/create_restart_script.py index 3f93ade..1d5caf1 100644 --- a/ScaFFold/utils/create_restart_script.py +++ b/ScaFFold/utils/create_restart_script.py @@ -15,6 +15,7 @@ # restart_script.py from __future__ import annotations +import math import os import shlex import stat @@ -23,6 +24,8 @@ from pathlib import Path from typing import List, Union +from ScaFFold.utils.distributed import get_local_size + # Profiling toggles that must be reproduced on restart -- but only when they # were active in the generating run. Names mirror ScaFFold.utils.perf_measure. _PROFILING_ENV_VARS = ("PROFILE_TORCH", "CALI_CONFIG") @@ -299,6 +302,11 @@ def create_restart_script(run_dir: str | Path, world_size: int | None = None) -> torchrun, Open MPI, PMI). The multi-rank torchrun-hpc template is emitted whenever the resulting world size is greater than one; the local single-process template is used only for a world size of one. + + The node count comes from the scheduler when it reports one (Flux, Slurm); + otherwise it is derived from the per-node rank count + ``ScaFFold.utils.distributed.get_local_size`` reads, so a PALS or torchrun + job is not relaunched with every rank crammed onto one node. """ run_dir = Path(run_dir) run_dir.mkdir(parents=True, exist_ok=True) @@ -339,8 +347,24 @@ def create_restart_script(run_dir: str | Path, world_size: int | None = None) -> if use_torchrun: # Calculate tasks per node for torchrun (-n arg). if nodes is None: - nodes = 1 - tasks_per_node = max(1, total_tasks // nodes) + # No scheduler reported a node count: this is a PALS or plain + # torchrun launch. Assuming one node put the job's whole rank count + # on a single node (an 8-rank job across 2 nodes came back as + # NODES=1 TASKS_PER_NODE=8), which either oversubscribes one node or + # is rejected outright. The rank side's ``get_local_size`` reads the + # same launchers' per-node variables, so ask it how many ranks share + # this node and derive the node count from that. ``required=True`` + # distinguishes "one rank per node" from "nothing reported a + # per-node count", where the historical single-node assumption is + # still the best guess available. + try: + local_size = get_local_size(required=True) + except RuntimeError: + local_size = total_tasks + tasks_per_node = max(1, min(local_size, total_tasks)) + nodes = math.ceil(total_tasks / tasks_per_node) + else: + tasks_per_node = max(1, total_tasks // nodes) script = _render_torchrun_hpc_restart( py_array_decl, nodes, tasks_per_node, env_setup diff --git a/tests/test_restart_script.py b/tests/test_restart_script.py index a8f3e08..291c57a 100644 --- a/tests/test_restart_script.py +++ b/tests/test_restart_script.py @@ -305,3 +305,51 @@ def test_pals_job_gets_a_multirank_restart_script(monkeypatch, tmp_path): assert "torchrun-hpc" in script assert 'exec "${PY[@]}"' not in script + + +# --------------------------------------------------------------------------- +# VC-5: the node shape comes from the local rank count when no scheduler +# reports one. Flux and Slurm state their node count; PALS and plain torchrun +# do not, and assuming one node relaunched an 8-rank/2-node job as +# NODES=1 TASKS_PER_NODE=8 -- an oversubscribed node, or a rejected job. +# --------------------------------------------------------------------------- + + +def test_pals_multinode_shape_uses_the_local_rank_count(monkeypatch, tmp_path): + """8 PALS ranks, 4 per node -> NODES=2, TASKS_PER_NODE=4.""" + _isolate_env(monkeypatch) + monkeypatch.delenv("PALS_LOCAL_SIZE", raising=False) + monkeypatch.delenv("LOCAL_WORLD_SIZE", raising=False) + monkeypatch.setenv("PALS_NRANKS", "8") + monkeypatch.setenv("PALS_LOCAL_SIZE", "4") + + script = _generate(monkeypatch, tmp_path / "run") + + assert 'NODES="2"' in script + assert 'TASKS_PER_NODE="4"' in script + + +def test_single_node_torchrun_shape_is_unchanged(monkeypatch, tmp_path): + """8 torchrun ranks all on one node stay NODES=1, TASKS_PER_NODE=8.""" + _isolate_env(monkeypatch) + monkeypatch.delenv("PALS_LOCAL_SIZE", raising=False) + monkeypatch.setenv("WORLD_SIZE", "8") + monkeypatch.setenv("LOCAL_WORLD_SIZE", "8") + + script = _generate(monkeypatch, tmp_path / "run") + + assert 'NODES="1"' in script + assert 'TASKS_PER_NODE="8"' in script + + +def test_unknown_local_size_keeps_the_single_node_assumption(monkeypatch, tmp_path): + """With nothing reporting a per-node count, the old assumption stands.""" + _isolate_env(monkeypatch) + for var in ("PALS_LOCAL_SIZE", "LOCAL_WORLD_SIZE", "PMI_LOCAL_SIZE"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("PALS_NRANKS", "8") + + script = _generate(monkeypatch, tmp_path / "run") + + assert 'NODES="1"' in script + assert 'TASKS_PER_NODE="8"' in script From 7cdef0302b702b77e8063689666d6a80516956d6 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:12:10 -0700 Subject: [PATCH 42/54] Wording and cosmetics The "no new epoch was trained" warning covers two ways of entering train() with nothing to do, and described one of them as the other: a fresh epochs:0 run was told "there was nothing to resume", sending its user to look for a checkpoint that was never part of the story. It now says the run had no epoch left to run, which is true of a completed resume and a zero-epoch fresh start alike. (The deeper problem in that scenario -- worker.py's rank-0-only genfromtxt raising after destroy_process_group, so rank 0 exits non-zero while its peers exit 0 -- is pre-existing structure, unchanged here.) explicit_cli_keys' docstring claimed that a flag passed with exactly its default value costs nothing because "both spellings then agree on the default". They do not: --datagen-batch-size 10000 next to datagen_batch_size: 500 in the config yields 500. The rationale is corrected to name the real trade-off rather than deny it; the defaults stay where they are, which is what the R20 tests pin. perf_measure now has a logger instead of printing. Its messages are all "your profiling request was not honored as written", which belongs on a channel a caller can filter or capture, not in the middle of the run's stdout. The R26 test follows it to caplog. And the R21 tests' function-body `import pytest` moves to module level as a skipif marker. The config-error re-raise added a commit ago is spelled as a statement rather than a conditional expression. VA-4, VC-6, VC-cosmetics --- ScaFFold/cli.py | 17 +++++++++++++---- ScaFFold/utils/perf_measure.py | 22 ++++++++++++++++----- ScaFFold/utils/trainer.py | 20 ++++++++++--------- tests/test_infra.py | 20 +++++++++---------- tests/test_reporting.py | 22 ++++++++------------- tests/test_resume.py | 35 ++++++++++++++++++++++++++++++++-- 6 files changed, 92 insertions(+), 44 deletions(-) diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 650a126..105e87a 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -100,9 +100,16 @@ def explicit_cli_keys(args, parsers): parser). Only these may outrank a config-file setting; everything else in the namespace is an argparse default, which is the weakest source. - The one ambiguity is a flag passed with exactly its default value: it looks - absent, so a config-file entry wins over it. Both spellings then agree on - the default, which is the only value the flag could have contributed. + The one ambiguity is a flag passed with exactly its default value: it is + indistinguishable from an absent flag, so a config-file entry outranks it. + Where the flag has no default (``None``) that is harmless -- passing a + value always makes it explicit -- but the two flags that do have one, + ``--datagen-batch-size`` (10000) and ``-v`` (0), lose the argument in that + one case: ``--datagen-batch-size 10000`` next to ``datagen_batch_size: 500`` + in the config file yields 500. The alternative is to give every flag a + ``None`` default and re-derive the real defaults elsewhere, which buys a + narrow correctness win by scattering the defaults; the ambiguity is + documented instead. """ explicit = set() for name, value in vars(args).items(): @@ -611,7 +618,9 @@ def main(): if config_error is not None: # Rank 0 re-raises the original (keeping its traceback); the peers # rebuild it from what crossed the wire. - raise rank0_error if rank0_error is not None else rebuild_error(*config_error) + if rank0_error is not None: + raise rank0_error + raise rebuild_error(*config_error) combined_config = comm.bcast(combined_config, root=0) # Restart pre-check. Like every other decision here it is made once, on diff --git a/ScaFFold/utils/perf_measure.py b/ScaFFold/utils/perf_measure.py index 7236e69..9d87a2b 100644 --- a/ScaFFold/utils/perf_measure.py +++ b/ScaFFold/utils/perf_measure.py @@ -12,12 +12,20 @@ # # SPDX-License-Identifier: (Apache-2.0) +import logging import os from contextlib import nullcontext CALI_PERF_ENV_VAR = "CALI_CONFIG" TORCH_PERF_ENV_VAR = "PROFILE_TORCH" +# This module is imported before (and independently of) the run's MPI logger, +# so it keeps its own. Everything it has to say is about the user's profiling +# request not being honored as written, which belongs on a diagnostic channel +# that a caller can filter or capture -- not on stdout, where it lands in the +# middle of whatever the run is printing. +logger = logging.getLogger(__name__) + def _profiler_env_flag(name): """Return True only for an affirmative value of the environment variable. @@ -39,8 +47,11 @@ def _profiler_env_flag(name): _CALI_PERF_ENABLED = True except Exception as e: - print("User requested Caliper annotations, but could not import Caliper") - print(f"Exception: {e}") + logger.warning( + "User requested Caliper annotations, but could not import Caliper: %s: %s", + type(e).__name__, + e, + ) # The torch profiler is gated purely on its own environment variable: Caliper # and the torch profiler may both be enabled at once. @@ -51,8 +62,9 @@ def _profiler_env_flag(name): TORCH_PERF_ENABLED = True except Exception: - print( - "User requested PyTorch profiling, but could not import the PyTorch profiler" + logger.warning( + "User requested PyTorch profiling, but could not import the " + "PyTorch profiler" ) @@ -131,7 +143,7 @@ def get_torch_context(ranks_per_node, rank): # thing the bounded window exists to prevent. wait = _profiler_env_int("PROFILE_TORCH_WAIT", 1) if wait < 1: - print( + logger.warning( "PROFILE_TORCH_WAIT must be at least 1: the profiler window " "opens before the warmup batches, whose work would otherwise " "accumulate in host memory as one unbounded step. Using " diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 81dd5f2..3498cc1 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -1086,17 +1086,19 @@ def train(self, profiler=None): completed_epochs = epoch - 1 if not completed_new_epoch: - # The loop exited without running a single epoch: the state we - # resumed either already covers every epoch this run was asked for, - # or already met target_dice. There is nothing new to save (the - # checkpoint on disk already records epoch `completed_epochs`) and - # none of the per-epoch metrics the final save would write were - # ever computed, so skip it and return normally -- the caller's - # post-processing still has the CSV the original run left behind. + # The loop exited without running a single epoch: the epoch budget + # was already exhausted (a resume whose checkpoint covers every + # epoch asked for, or a fresh run configured with none) or the + # starting state already met target_dice. There is nothing new to + # save -- any checkpoint on disk already records epoch + # `completed_epochs`, and none of the per-epoch metrics the final + # save would write were ever computed -- so skip it and return + # normally; the caller's post-processing still has whatever CSV is + # there. self.log.warning( "No new epoch was trained (start epoch %s, 'epochs' %s, " - "starting val dice %s vs target_dice %s): there was nothing to " - "resume, and no checkpoint was written.", + "starting val dice %s vs target_dice %s): this run had no epoch " + "left to run, and no checkpoint was written.", self.start_epoch, self.config.epochs, self.start_val_dice, diff --git a/tests/test_infra.py b/tests/test_infra.py index c279f34..aaf3a73 100644 --- a/tests/test_infra.py +++ b/tests/test_infra.py @@ -26,6 +26,7 @@ import os import numpy as np +import pytest import torch from ScaFFold.utils.data_loading import FractalDataset @@ -245,19 +246,23 @@ def _debug_logger(name): return log -def test_mem_stats_without_cuda(caplog): - """``mem_stats`` reports "no GPU" instead of raising on a CPU-only host.""" - if torch.cuda.is_available(): - import pytest +# The bug was a CUDA-free host taking the CUDA path, so these only mean +# something where CUDA is genuinely unavailable. +_requires_no_cuda = pytest.mark.skipif( + torch.cuda.is_available(), reason="covers the CPU-only path" +) - pytest.skip("test covers the CPU-only path") +@_requires_no_cuda +def test_mem_stats_without_cuda(caplog): + """``mem_stats`` reports "no GPU" instead of raising on a CPU-only host.""" stats = mem_stats() assert stats["cuda_available"] is False assert "rank" in stats +@_requires_no_cuda def test_gather_and_print_mem_without_cuda(caplog): """A DEBUG-level CPU run logs a fallback instead of crashing. @@ -265,11 +270,6 @@ def test_gather_and_print_mem_without_cuda(caplog): ``-v`` used to die in trainer construction with "No CUDA GPUs are available". """ - if torch.cuda.is_available(): - import pytest - - pytest.skip("test covers the CPU-only path") - log = _debug_logger("test_gather_and_print_mem_without_cuda") with caplog.at_level(logging.DEBUG, logger=log.name): gather_and_print_mem(log, "after_trainer_setup") diff --git a/tests/test_reporting.py b/tests/test_reporting.py index 3b658fd..4a462ca 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -13,6 +13,7 @@ # SPDX-License-Identifier: (Apache-2.0) import csv +import logging from pathlib import Path from types import SimpleNamespace @@ -319,8 +320,6 @@ def test_zero_step_export_is_reported_not_raised(self, tmp_path, caplog): post-processing, so a raise here kills the profiling rank and leaves every other rank blocked in that barrier until the collective timeout. """ - import logging - import ScaFFold.worker as worker log = logging.getLogger("test_zero_step_export") @@ -340,8 +339,6 @@ def test_zero_step_export_is_reported_not_raised(self, tmp_path, caplog): def test_successful_export_writes_a_trace(self, tmp_path, caplog): """A profiler with a completed window still writes its trace (control).""" - import logging - from torch.profiler import ProfilerActivity, profile, schedule import ScaFFold.worker as worker @@ -364,8 +361,6 @@ def test_successful_export_writes_a_trace(self, tmp_path, caplog): def test_trace_lands_in_the_run_dir(self, tmp_path, caplog): """R23: the trace goes to the run dir, not whatever CWD happens to be.""" - import logging - import ScaFFold.worker as worker prof = self._stepped_profiler() @@ -387,8 +382,6 @@ def test_trace_name_counts_nodes_not_ranks( self, tmp_path, world_size, ranks_per_node, expected ): """R23: the N field is a node count, and never rounds a node away.""" - import logging - import ScaFFold.worker as worker prof = self._stepped_profiler() @@ -493,7 +486,7 @@ def _context_with(monkeypatch_context, env): assert is_local return ctx - def test_wait_zero_does_not_record_step_zero(self, monkeypatch, capsys): + def test_wait_zero_does_not_record_step_zero(self, monkeypatch, caplog): """PROFILE_TORCH_WAIT=0 is clamped so step 0 records nothing. worker.main enters the profiler context around checkpoint cleanup and @@ -509,11 +502,12 @@ def test_wait_zero_does_not_record_step_zero(self, monkeypatch, capsys): import ScaFFold.utils.perf_measure as perf_measure try: - with monkeypatch.context() as m: - ctx = self._context_with(m, {"PROFILE_TORCH_WAIT": "0"}) - assert ctx.schedule(0) == ProfilerAction.NONE - output = capsys.readouterr().out - assert "PROFILE_TORCH_WAIT" in output + with caplog.at_level(logging.WARNING, logger=perf_measure.logger.name): + with monkeypatch.context() as m: + ctx = self._context_with(m, {"PROFILE_TORCH_WAIT": "0"}) + assert ctx.schedule(0) == ProfilerAction.NONE + messages = " ".join(record.getMessage() for record in caplog.records) + assert "PROFILE_TORCH_WAIT" in messages finally: importlib.reload(perf_measure) diff --git a/tests/test_resume.py b/tests/test_resume.py index 3e8aa4d..dfe7b2c 100644 --- a/tests/test_resume.py +++ b/tests/test_resume.py @@ -364,7 +364,7 @@ def test_restart_of_completed_run_trains_and_saves_nothing(tmp_path, caplog): (the last checkpoint already covers the completed epochs) rather than saving with an unbound ``val_loss_avg``. ``train()`` has to return normally so the worker's post-processing still runs off the CSV already on disk, and - the run must say plainly that there was nothing to resume. + the run must say plainly that it had no epoch left to run. """ log = logging.getLogger("resume.r01") run = tmp_path / "run" @@ -416,7 +416,38 @@ def test_restart_of_completed_run_trains_and_saves_nothing(tmp_path, caplog): assert ckpt.read_bytes() == before epochs = [ln.split(",")[0] for ln in csv.read_text().splitlines()[1:]] assert epochs == ["1", "2"] - assert "nothing to resume" in caplog.text.lower() + assert "no epoch left to run" in caplog.text.lower() + + +def test_fresh_run_with_no_epochs_does_not_claim_a_resume(tmp_path, caplog): + """A fresh ``epochs: 0`` run reports the truth: nothing was resumed. + + The same message covers both ways of entering ``train()`` with no epoch to + run, so it must not describe one of them as the other. This run has no + checkpoint and never asked for one -- telling its user "there was nothing + to resume" sends them looking for a checkpoint that was never part of the + story. + """ + log = logging.getLogger("resume.va4") + run = tmp_path / "run" + run.mkdir() + + trainer = _stub_trainer( + run, + train_from_scratch=True, + log=log, + epochs=0, + checkpoint_interval=1, + ) + trainer.cleanup_or_resume() + assert trainer.start_epoch == 1 # a fresh run: nothing was resumed + + with caplog.at_level(logging.WARNING): + trainer.train() + + assert "no epoch left to run" in caplog.text.lower() + assert "nothing to resume" not in caplog.text.lower() + assert not trainer.checkpoint_manager.last_ckpt_path.exists() def test_converged_resume_does_not_retrain(tiny_trainer, monkeypatch): From bd380271d5c7ebe0880a93b98e362ef4c5589668 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:18:59 -0700 Subject: [PATCH 43/54] Pin the non-repo provenance test to its own directory test_non_repo_install_reports_no_commit_id builds a "not a checkout" directory under tmp_path, but git walks upwards until it finds a repository: run with --basetemp inside a ScaFFold checkout, that directory inherits the checkout's HEAD and the test fails on where it was run rather than on what it tests. GIT_CEILING_DIRECTORIES stops the walk at tmp_path. VB-7 --- tests/datagen/test_provenance.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/datagen/test_provenance.py b/tests/datagen/test_provenance.py index d53bd33..d907c39 100644 --- a/tests/datagen/test_provenance.py +++ b/tests/datagen/test_provenance.py @@ -92,12 +92,19 @@ def test_commit_survives_a_non_repo_working_directory(tmp_path, monkeypatch): assert gd._git_commit_short(LOG) == expected -def test_non_repo_install_reports_no_commit_id(tmp_path): +def test_non_repo_install_reports_no_commit_id(tmp_path, monkeypatch): """An installed (non-git) ScaFFold still degrades gracefully. Provenance is best-effort: when the source tree is not a checkout there is no commit to record, and reuse simply is not gated on one. + + "Not a checkout" has to be made true of the *whole path*, not just the leaf: + git walks upwards until it finds a repository, so with ``--basetemp`` inside + a ScaFFold checkout this directory inherits that checkout's HEAD and the + test fails on where it was run rather than on what it tests. The ceiling + stops the walk at ``tmp_path``. """ + monkeypatch.setenv("GIT_CEILING_DIRECTORIES", str(tmp_path)) not_a_repo = tmp_path / "site-packages" / "ScaFFold" / "datagen" not_a_repo.mkdir(parents=True) From b11569d1d6c378708f0c541a91b8ff94925dbd0c Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:26:24 -0700 Subject: [PATCH 44/54] Copy the committed checkpoint instead of re-serializing it for best An improving epoch pickled and fsynced the identical state dict twice, once for checkpoint_last.pth and again for checkpoint_best.pth. The best file is now copied (tmp + fsync + os.replace) from the last file the same writer just committed, halving the serialization CPU and checkpoint bytes. R08 --- ScaFFold/utils/checkpointing.py | 43 ++++++++++++++++++-- tests/test_checkpointing.py | 71 +++++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 3 deletions(-) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 22c6051..84ed96a 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -15,6 +15,7 @@ import math import os import random +import shutil import traceback from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -630,6 +631,32 @@ def _atomic_save(state_dict, path): pass raise + @staticmethod + def _atomic_copy(src, dst): + """Copy an already-committed checkpoint onto another name atomically. + + Same discipline as ``_atomic_save`` -- copy into a temp file in the + same directory, fsync it, then ``os.replace`` -- so ``dst`` is never + observed half-written and a crash mid-copy cannot damage the previous + good file there. + """ + src = Path(src) + dst = Path(dst) + tmp_path = dst.with_name(f"{dst.name}.tmp.{os.getpid()}") + try: + with open(src, "rb") as fsrc, open(tmp_path, "wb") as fdst: + shutil.copyfileobj(fsrc, fdst) + fdst.flush() + os.fsync(fdst.fileno()) + os.replace(tmp_path, dst) + except Exception: + try: + if tmp_path.exists(): + tmp_path.unlink() + except OSError: + pass + raise + @classmethod def _write_to_disk(cls, state_dict, last_path, best_path, is_best, log): """Worker function to perform actual disk I/O. @@ -642,10 +669,20 @@ def _write_to_disk(cls, state_dict, last_path, best_path, is_best, log): try: # Save 'last' atomically. cls._atomic_save(state_dict, last_path) - # Save 'best' atomically (re-serialize rather than copy a file that - # a concurrent writer might still be replacing). + # 'best' is byte-identical to the 'last' just committed, so copy + # that file instead of pickling and fsyncing the same state a + # second time (double the serialization CPU and double the bytes + # pushed at the shared filesystem on every improving epoch). + # + # There is no concurrent writer to race: checkpoint writes are + # serialized through a single writer -- the caller's thread in sync + # mode, or the one-worker ThreadPoolExecutor in async mode, whose + # previous write ``_rank0_save`` drains before submitting the next + # -- and only rank 0 ever writes. So this very thread performed the + # ``os.replace`` onto ``last_path`` a moment ago and nothing else + # can be replacing it now. if is_best: - cls._atomic_save(state_dict, best_path) + cls._atomic_copy(last_path, best_path) except Exception: if log is not None: log.error("Saving checkpoint failed:\n%s", traceback.format_exc()) diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index acca106..82ab353 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -518,6 +518,77 @@ def test_init_sweeps_orphaned_tmp_files(tmp_path): assert quarantined.exists() +# --------------------------------------------------------------------------- +# R08 -- an improving epoch serializes the state dict once, not twice +# --------------------------------------------------------------------------- + + +def _count_torch_saves(monkeypatch): + """Count ``torch.save`` calls made by the checkpoint writer.""" + calls = [] + real_save = torch.save + + def counting_save(obj, f, *args, **kwargs): + calls.append(str(getattr(f, "name", f))) + return real_save(obj, f, *args, **kwargs) + + monkeypatch.setattr(torch, "save", counting_save) + return calls + + +def test_best_checkpoint_copied_not_reserialized(tmp_path, monkeypatch): + """The best checkpoint reuses the bytes just written to 'last'. + + An improving epoch used to pickle *and fsync* the identical state dict a + second time, doubling both the serialization CPU and the checkpoint bytes + pushed at the shared filesystem. Copying the file that was just committed + costs one read (usually from page cache) and one write instead. + """ + mgr, model = _make_manager(tmp_path) + saves = _count_torch_saves(monkeypatch) + + assert mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) is True + assert len(saves) == 1 + assert saves == [str(mgr.last_ckpt_path) + f".tmp.{os.getpid()}"] + + # A second improving epoch: still exactly one serialization. + saves.clear() + assert mgr.save_checkpoint(epoch=2, val_loss_avg=0.25) is True + assert len(saves) == 1 + + # A non-improving epoch writes 'last' only and leaves 'best' alone. + best_bytes = mgr.best_ckpt_path.read_bytes() + saves.clear() + assert mgr.save_checkpoint(epoch=3, val_loss_avg=0.9) is False + assert len(saves) == 1 + assert mgr.best_ckpt_path.read_bytes() == best_bytes + + # And the best checkpoint is still a real, loadable checkpoint holding the + # epoch-2 state -- byte-identical to the 'last' file it was copied from. + monkeypatch.undo() + best = torch.load(mgr.best_ckpt_path, map_location="cpu", weights_only=False) + assert best["epoch"] == 2 + assert best["val_loss_avg"] == pytest.approx(0.25) + for name, tensor in model.state_dict().items(): + assert torch.equal(best["model_state_dict"][name], tensor) + + +def test_best_copy_failure_is_not_silent(tmp_path, monkeypatch): + """A failed best-checkpoint copy surfaces and leaves no partial file.""" + mgr, _ = _make_manager(tmp_path) + + def boom(*args, **kwargs): + raise OSError("[Errno 28] No space left on device") + + monkeypatch.setattr("ScaFFold.utils.checkpointing.shutil.copyfileobj", boom) + + with pytest.raises(CheckpointSaveError, match="No space left on device"): + mgr.save_checkpoint(epoch=1, val_loss_avg=0.5) + + assert not mgr.best_ckpt_path.exists() + assert list(tmp_path.glob("checkpoint_*.tmp.*")) == [] + + # --------------------------------------------------------------------------- # VA-1/VA-2/VA-3 -- the remaining rank-0 filesystem windows are fenced # From 913c92c57b518348d0986c358a12c92ad0be863a Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:29:52 -0700 Subject: [PATCH 45/54] Take the warmup training-state snapshot on the host snapshot_training_state cloned model and optimizer state device-to-device, holding ~3x parameter bytes of accelerator memory across the whole warmup phase (and warmup's own memory peak). It now copies to CPU; load_state_dict puts the state back on each parameter's device on restore. R09 --- ScaFFold/utils/checkpointing.py | 47 ++++++++------ tests/test_checkpointing.py | 109 ++++++++++++++++++++++++++++++++ 2 files changed, 135 insertions(+), 21 deletions(-) diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 84ed96a..5a04ecd 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -295,17 +295,29 @@ def finalize_saves(self) -> None: self._raise_save_error(error) def snapshot_training_state(self) -> Dict[str, Any]: - """Capture mutable in-memory training state without writing a checkpoint.""" + """Capture mutable in-memory training state without writing a checkpoint. + + The copy is taken on the *host*. Cloning device-to-device instead would + hold roughly 3x parameter bytes of accelerator memory (the weights plus + Adam's two moment buffers) for as long as the snapshot lives -- which is + the whole of warmup, precisely where the run establishes its peak + memory. ``restore_training_state`` puts the state back on the model's + own device, so the detour is invisible to callers. + """ model_ref = self.model.module if hasattr(self.model, "module") else self.model return { - "model_state_dict": self._clone_state_dict(model_ref.state_dict()), - "optimizer_state_dict": self._clone_state_dict(self.optimizer.state_dict()) + "model_state_dict": self._transfer_dict_to_cpu(model_ref.state_dict()), + "optimizer_state_dict": self._transfer_dict_to_cpu( + self.optimizer.state_dict() + ) if self.optimizer else None, - "scheduler_state_dict": self._clone_state_dict(self.scheduler.state_dict()) + "scheduler_state_dict": self._transfer_dict_to_cpu( + self.scheduler.state_dict() + ) if self.scheduler else None, - "grad_scaler_state_dict": self._clone_state_dict( + "grad_scaler_state_dict": self._transfer_dict_to_cpu( self.grad_scaler.state_dict() ) if self.grad_scaler @@ -315,7 +327,13 @@ def snapshot_training_state(self) -> Dict[str, Any]: } def restore_training_state(self, snapshot: Dict[str, Any]) -> None: - """Restore an in-memory training snapshot.""" + """Restore an in-memory training snapshot. + + The snapshot is host-resident (see ``snapshot_training_state``); the + ``load_state_dict`` calls below copy into the live parameters and move + optimizer state onto each parameter's device, so the restored state + ends up exactly where it started. + """ model_ref = self.model.module if hasattr(self.model, "module") else self.model model_ref.load_state_dict(snapshot["model_state_dict"]) @@ -696,8 +714,8 @@ def _transfer_dict_to_cpu(self, obj): ``Tensor.cpu()`` is a no-op for tensors already on CPU (it returns the same object), so CPU-resident state must be cloned explicitly; - otherwise the async writer would serialize tensors the training loop - keeps mutating in place. + otherwise the async writer -- or the warmup snapshot -- would keep an + alias of tensors the training loop mutates in place. """ if torch.is_tensor(obj): t = obj.detach() @@ -711,19 +729,6 @@ def _transfer_dict_to_cpu(self, obj): else: return obj - def _clone_state_dict(self, obj): - """Recursively clone tensors so in-memory snapshots are isolated.""" - if torch.is_tensor(obj): - return obj.detach().clone() - elif isinstance(obj, dict): - return {k: self._clone_state_dict(v) for k, v in obj.items()} - elif isinstance(obj, list): - return [self._clone_state_dict(v) for v in obj] - elif isinstance(obj, tuple): - return tuple(self._clone_state_dict(v) for v in obj) - else: - return obj - def _quarantine_corrupt(self, path): """Rename an unreadable checkpoint aside so it is not retried on the next restart (a persistent corrupt 'last' would otherwise break every diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 82ab353..5d7afa7 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -705,6 +705,115 @@ def test_cpu_tensors_cloned(tmp_path): assert torch.equal(snapshot["nested"][0], torch.ones(3)) +# --------------------------------------------------------------------------- +# R09 -- the warmup snapshot is host-resident, not a device-side copy +# --------------------------------------------------------------------------- + + +def _all_tensors(obj): + """Yield every tensor reachable from a snapshot payload.""" + if torch.is_tensor(obj): + yield obj + elif isinstance(obj, dict): + for value in obj.values(): + yield from _all_tensors(value) + elif isinstance(obj, (list, tuple)): + for value in obj: + yield from _all_tensors(value) + + +class _StubOptimizer: + """Stand-in exposing only the state_dict the snapshot reads.""" + + def __init__(self, state): + self._state = state + + def state_dict(self): + return self._state + + +def test_snapshot_holds_no_device_resident_tensors(tmp_path): + """``snapshot_training_state`` copies to the host, not device-to-device. + + The snapshot used to clone model and optimizer state on whatever device + they lived on, pinning ~3x parameter bytes of accelerator memory (model + clone + Adam's two moment buffers) from before the first warmup batch + until the restore -- i.e. straight across warmup's own peak, which is + where a memory-marginal configuration OOMs. Device tensors are simulated + with FakeTensorMode so this runs without a GPU. + """ + from torch._subclasses.fake_tensor import FakeTensorMode + + mgr, _ = _make_manager(tmp_path) + + with FakeTensorMode(): + mgr.model = torch.nn.Linear(8, 4, device="cuda") + mgr.optimizer = _StubOptimizer( + { + "state": { + 0: { + "step": torch.zeros(1, device="cuda"), + "exp_avg": torch.zeros(4, 8, device="cuda"), + "exp_avg_sq": torch.zeros(4, 8, device="cuda"), + } + }, + "param_groups": [{"params": [0], "lr": 0.1}], + } + ) + snapshot = mgr.snapshot_training_state() + + devices = {t.device.type for t in _all_tensors(snapshot)} + + assert devices == {"cpu"}, f"snapshot kept device-resident tensors: {devices}" + + +def test_snapshot_restore_round_trips_values_and_devices(tmp_path): + """A host-resident snapshot still restores exact state on its own device. + + ``load_state_dict`` copies into the live parameters and moves optimizer + state back to each parameter's device, so nothing downstream has to know + the snapshot took a detour through the host. + """ + mgr, model = _make_manager(tmp_path) + optimizer = mgr.optimizer + + # One applied step so the optimizer carries real per-parameter state. + model(torch.randn(2, 8)).sum().backward() + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + reference_params = {k: v.detach().clone() for k, v in model.state_dict().items()} + reference_state = { + pid: {k: v.detach().clone() for k, v in state.items() if torch.is_tensor(v)} + for pid, state in optimizer.state_dict()["state"].items() + } + param_devices = {k: v.device for k, v in model.state_dict().items()} + + snapshot = mgr.snapshot_training_state() + + # Warmup-shaped mutation: more steps, then roll back. + for _ in range(3): + model(torch.randn(2, 8)).sum().backward() + optimizer.step() + optimizer.zero_grad(set_to_none=True) + assert not torch.equal(model.state_dict()["weight"], reference_params["weight"]) + + mgr.restore_training_state(snapshot) + + for name, tensor in model.state_dict().items(): + assert torch.equal(tensor, reference_params[name]), name + assert tensor.device == param_devices[name], name + restored_state = optimizer.state_dict()["state"] + for pid, state in reference_state.items(): + for key, value in state.items(): + assert torch.equal(restored_state[pid][key], value), (pid, key) + for group, param in zip(optimizer.param_groups, model.parameters()): + del group + for value in optimizer.state[param].values(): + if torch.is_tensor(value) and value.dim() > 0: + assert value.device == param.device + + # --------------------------------------------------------------------------- # F50 -- a final checkpoint is written when the run exits between intervals # --------------------------------------------------------------------------- From a4e9e047222490e8c00f233a996857b71d792b0b Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:31:33 -0700 Subject: [PATCH 46/54] Cache the parsed category IFS parameters across a category's instances instance.main re-ran np.genfromtxt on the same category CSV for every (category, instance) work item -- 145 identical shared-filesystem reads per category. Work items for a category are contiguous in the block partition, so a one-entry cache gives one parse per category per rank. R36 --- ScaFFold/datagen/instance.py | 24 ++++++++++---- tests/datagen/test_artifacts.py | 57 +++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 6 deletions(-) diff --git a/ScaFFold/datagen/instance.py b/ScaFFold/datagen/instance.py index de8eb0d..ea84a91 100644 --- a/ScaFFold/datagen/instance.py +++ b/ScaFFold/datagen/instance.py @@ -346,14 +346,26 @@ def main(config: Config): start_time = time.time() + # One-entry cache of the most recently parsed category CSV. The work list + # is built category-major and block-sliced, so every rank's items for a + # given category are contiguous: a single entry is enough to turn the + # per-item re-parse (145 identical reads of the same small file off the + # shared filesystem, per category) into one parse per category per rank. + # ``generate_instance_points`` copies before scaling, so sharing the parsed + # array across instances cannot leak weights from one item into the next. + cached_category = None + params = None + for i, category_instance_pair in enumerate(instances_to_generate_for_this_rank): category, instance = category_instance_pair - category_IFS_params = IFS_param_csv_names[category] - params = np.genfromtxt( - f"{fracts_read_dir}/{category_IFS_params}", - dtype=DEFAULT_NP_DTYPE, - delimiter=",", - ) + if category != cached_category: + category_IFS_params = IFS_param_csv_names[category] + params = np.genfromtxt( + f"{fracts_read_dir}/{category_IFS_params}", + dtype=DEFAULT_NP_DTYPE, + delimiter=",", + ) + cached_category = category weights = weights_all[instance] # Generate a validated, weighted point cloud. Weighting can turn a diff --git a/tests/datagen/test_artifacts.py b/tests/datagen/test_artifacts.py index c1c38e3..f07bd3a 100644 --- a/tests/datagen/test_artifacts.py +++ b/tests/datagen/test_artifacts.py @@ -234,6 +234,63 @@ def test_resume_rejects_truncated(tmp_path): assert np.isfinite(np.load(victim)).all() +# --------------------------------------------------------------------------- +# R36: a category's IFS parameters are parsed once, not once per instance +# --------------------------------------------------------------------------- + + +def test_category_params_parsed_once_per_category(tmp_path, monkeypatch): + """``main`` parses each category CSV once per rank, not once per work item. + + The parse used to sit inside the per-item loop, so a full generation read + and re-parsed the same small CSV 145 times per category off the shared + filesystem. Work items for a category are contiguous in the block + partition, so a one-entry cache collapses that to one parse per category. + """ + fract_base = tmp_path / "fractals" + point_num = 60 + n_categories = 2 + missing_per_category = 3 + + config = _make_config(fract_base, point_num=point_num) + config.n_categories = n_categories + + param_dir = Path(layout.category_param_dir(config)) + param_dir.mkdir(parents=True) + instance_root = Path(layout.instance_dir(config)) + rng = np.random.default_rng(0) + for category in range(n_categories): + np.savetxt( + param_dir / f"{category:06d}.csv", _contractive_params(), delimiter="," + ) + # Pre-seed all but a few instances so the run stays fast; the ones left + # missing are what the loop (and the parse) actually iterates over. + inst_dir = instance_root / f"{category:06d}" + inst_dir.mkdir(parents=True) + for i in range(missing_per_category, 145): + np.save(inst_dir / f"{category:06d}_{i:04d}.npy", rng.random((10, 3))) + + parses = [] + real_genfromtxt = np.genfromtxt + + def counting_genfromtxt(fname, *args, **kwargs): + parses.append(Path(str(fname)).name) + return real_genfromtxt(fname, *args, **kwargs) + + monkeypatch.setattr(inst.np, "genfromtxt", counting_genfromtxt) + inst.main(config) + + category_parses = [name for name in parses if name[0].isdigit()] + # Every missing instance was generated ... + for category in range(n_categories): + for i in range(missing_per_category): + assert ( + instance_root / f"{category:06d}" / f"{category:06d}_{i:04d}.npy" + ).exists() + # ... from n_categories parses, not one per (category, instance) item. + assert sorted(category_parses) == ["000000.csv", "000001.csv"] + + # --------------------------------------------------------------------------- # F62: mask scanner requires exactly one file per id # --------------------------------------------------------------------------- From be4a01b891351bf031830893b5f728b8efe62674 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:43:40 -0700 Subject: [PATCH 47/54] Warm the ragged final batch shapes during warmup Warmup ran only leading, full-size batches, so with local_batch_size>1 and an indivisible shard the epoch's partial batch met cuDNN/MIOpen for the first time inside the first timed epoch -- a measured 95 s stall in epoch_duration, the FOM denominator. Warmup now runs one extra iteration per distinct ragged size (agreed across ranks, since validation shards are unpadded) inside the existing snapshot/restore envelope. R41 --- ScaFFold/utils/trainer.py | 69 ++++++++++++++++++++ tests/test_perf_hotpath.py | 127 +++++++++++++++++++++++++++++++++++++ 2 files changed, 196 insertions(+) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 3498cc1..9e53ecc 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -713,6 +713,71 @@ def _sync_gather_minibatch_timer(self, minibatch_events): minibatch_time_s = statistics.median(minibatch_times.cpu().tolist()) return minibatch_time_s + def _warmup_ragged_batches(self, batch): + """Warm the narrower final batch each loader ends its epoch with. + + Warmup runs only the *leading* batches of the train loader, which are + all ``local_batch_size`` wide, and neither loader drops its last batch. + So when a rank's sample count is not a multiple of the batch size, the + final batch of every epoch presents a set of convolution problems + nothing has ever run: with ``cudnn.benchmark`` on, its algorithm search + (MIOpen find) then happens inside the *first timed epoch* -- a one-off + stall measured at 95 s that lands in ``epoch_duration``, i.e. straight + in the FOM denominator, while staying invisible to the per-minibatch + timer (which excludes partial batches). + + One extra iteration per distinct ragged size fixes that, cut from a + batch warmup already fetched so no additional I/O is needed. The + validation shapes are covered by the same (training) step: the forward + convolutions are what validation shares, and this runs inside warmup's + snapshot/restore envelope, so the extra step cannot affect training. + ``local_batch_size = 1`` never has a ragged batch and does no extra work. + + The set of sizes is agreed across ranks first. The training shards are + padded to equal length, so their remainder is already identical + everywhere, but validation is sharded *unpadded* on purpose (F-series: + an unbiased SUM-reduced metric), so per-rank counts -- and their + remainders -- differ. A rank that decided on its own would run a step + its peers did not, and the collectives inside that step (the gradient + all-reduce, the sharded loss reductions) would deadlock. + """ + if batch is None: + return + local_batch_size = self.config.local_batch_size + available = batch["image"].shape[0] + + ragged_sizes = {len(self.train_sampler) % local_batch_size} + local_val_ragged = torch.tensor( + [len(self.val_sampler) % local_batch_size], device=self.device + ) + gathered_val_ragged = [ + torch.empty_like(local_val_ragged) for _ in range(self.world_size) + ] + torch.distributed.all_gather(gathered_val_ragged, local_val_ragged) + ragged_sizes.update(int(size.item()) for size in gathered_val_ragged) + + # The batches already run are all ``available`` wide, and ``available`` + # is itself rank-invariant (the padded training shards give every rank + # the same leading batch size), so this loop is identical on all ranks. + warmed = {0, available} + for ragged in sorted(ragged_sizes): + if ragged in warmed: + continue + if ragged > available: + self.log.debug( + f" warmup: cannot build a {ragged}-sample batch from a " + f"{available}-sample batch; skipping that shape" + ) + continue + warmed.add(ragged) + self.log.debug( + f" warmup: running the ragged batch shape ({ragged} samples)" + ) + self._run_training_batch( + {key: value[:ragged] for key, value in batch.items()}, + log_prefix=f"warmup ragged ({ragged}): ", + ) + def warmup(self): """Run warmup iterations before the main training loop.""" warmup_batches = self.config.warmup_batches @@ -735,10 +800,12 @@ def warmup(self): self.optimizer.zero_grad(set_to_none=True) try: + last_batch = None for batch_idx, batch in enumerate(self.train_loader): if batch_idx >= max_batches: break + last_batch = batch self._run_training_batch( batch, log_prefix="warmup: ", @@ -749,6 +816,8 @@ def warmup(self): f" warmup: batch {batch_idx} completed in {batch_t_end - start_warmup} seconds" ) + self._warmup_ragged_batches(last_batch) + self.val_loader.sampler.set_epoch(0) if max_val_batches > 0: diff --git a/tests/test_perf_hotpath.py b/tests/test_perf_hotpath.py index 840bc3d..0eb35d0 100644 --- a/tests/test_perf_hotpath.py +++ b/tests/test_perf_hotpath.py @@ -188,3 +188,130 @@ def test_training_batch_gathers_mem_only_first_batch(): # first-batch predicate. assert "gather_mem_stats=True" not in src assert "first_batch" in src + + +# --------------------------------------------------------------------------- +# R41: warmup covers the ragged final batch +# +# Warmup only ever runs the *leading* batches of the train loader, which are +# all local_batch_size wide, and neither loader drops its last batch. When a +# rank's sample count is not a multiple of the batch size, the narrower final +# batch is therefore a set of convolution shapes nothing has warmed, and with +# cudnn.benchmark on the algorithm search for it runs inside the first *timed* +# epoch (measured: 95 s) -- straight into epoch_duration, the FOM denominator. +# --------------------------------------------------------------------------- + + +def _stub_warmup_steps(trainer, monkeypatch, *, mutate=False): + """Record the batch sizes warmup runs; optionally mutate model state. + + The real training step is hardwired through DistConv and cannot run on the + CPU ``ps=None`` fixture, so the step itself is stubbed; what is under test + is which batches warmup feeds it. + """ + import ScaFFold.utils.trainer as tr + + sizes = [] + + def fake_training_batch(batch, **kwargs): + sizes.append(int(batch["image"].shape[0])) + if mutate: + with torch.no_grad(): + for param in trainer.model.parameters(): + param.add_(1.0) + return int(batch["image"].shape[0]), torch.tensor(0.0), torch.tensor(0.0) + + monkeypatch.setattr(trainer, "_run_training_batch", fake_training_batch) + monkeypatch.setattr(tr, "evaluate", lambda *args, **kwargs: (0.0, 0.0)) + return sizes + + +def test_warmup_covers_the_ragged_train_batch(tiny_trainer, monkeypatch): + # 5 local training samples at local_batch_size 2: every epoch ends with a + # 1-sample batch that the leading warmup batches never present. + trainer = tiny_trainer( + n_train=5, + n_val=4, + config_overrides={"local_batch_size": 2, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch) + + trainer.warmup() + + assert sizes == [2, 2, 1] + + +def test_warmup_adds_no_extra_batch_when_shards_divide_evenly( + tiny_trainer, monkeypatch +): + # local_batch_size 1 can never produce a ragged batch: no extra work. + trainer = tiny_trainer( + n_train=4, + n_val=2, + config_overrides={"local_batch_size": 1, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch) + + trainer.warmup() + + assert sizes == [1, 1] + + +def test_warmup_covers_a_ragged_validation_batch(tiny_trainer, monkeypatch): + # Training divides evenly (4 / 2) but validation does not (3 / 2), so the + # 1-sample shape still has to be warmed. + trainer = tiny_trainer( + n_train=4, + n_val=3, + config_overrides={"local_batch_size": 2, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch) + + trainer.warmup() + + assert sizes == [2, 2, 1] + + +def test_warmup_ragged_sizes_are_agreed_across_ranks(tiny_trainer, monkeypatch): + # Validation is sharded unpadded, so peers can end their epoch with a + # different partial size. Every rank must run the same extra steps or the + # collectives inside them diverge; the peer's remainder (2) is gathered and + # warmed here even though this rank's own is 1. + import torch.distributed as dist + + trainer = tiny_trainer( + n_train=6, + n_val=4, + config_overrides={"local_batch_size": 3, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch) + trainer.world_size = 2 + + def fake_all_gather(tensor_list, tensor, *args, **kwargs): + tensor_list[0].copy_(tensor) + tensor_list[1].fill_(2) + + monkeypatch.setattr(dist, "all_gather", fake_all_gather) + + trainer.warmup() + + assert sizes == [3, 3, 1, 2] + + +def test_warmup_rolls_back_state_including_the_ragged_batch(tiny_trainer, monkeypatch): + # The extra ragged iteration stays inside warmup's snapshot/restore + # envelope, so nothing it touches survives into training. + trainer = tiny_trainer( + n_train=5, + n_val=4, + config_overrides={"local_batch_size": 2, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch, mutate=True) + before = {k: v.detach().clone() for k, v in trainer.model.state_dict().items()} + + trainer.warmup() + + assert sizes == [2, 2, 1] + after = trainer.model.state_dict() + for name, tensor in before.items(): + assert torch.equal(after[name], tensor), name From e6bb5baea0339c832526c38f06a9b0ac149d0406 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:43:40 -0700 Subject: [PATCH 48/54] Close the figures standard_viz opens pyplot retained every figure the run-summary plots created, and a sweep calls standard_viz.main in-process once per combination, so figures (and their canvases) piled up for the whole sweep. Each is now closed after its savefig, and unconditionally in a finally since plotting errors are swallowed. R43 --- ScaFFold/viz/standard_viz.py | 19 +++++++++++++--- tests/test_reporting.py | 44 ++++++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/ScaFFold/viz/standard_viz.py b/ScaFFold/viz/standard_viz.py index 50d5beb..d9f2dda 100644 --- a/ScaFFold/viz/standard_viz.py +++ b/ScaFFold/viz/standard_viz.py @@ -27,6 +27,10 @@ def main(config: RunConfig): figures_path = Path(config.run_dir) / "figures" figures_path.mkdir(parents=True, exist_ok=True) + # pyplot keeps a strong reference to every figure until it is closed, and a + # sweep calls this once per combination in the same process, so an unclosed + # figure is retained (canvas included) for the rest of the run. + figures = [] try: epochs = [] train_loss = [] @@ -53,7 +57,7 @@ def main(config: RunConfig): legend_loc = (0, -0.17) # Plot training loss - plt.figure() + figures.append(plt.figure()) plt.plot(epochs, train_loss, label="Train Loss", linewidth=line_thickness) plt.xlabel("Epoch", fontsize=fontsize) plt.ylabel("Train loss", fontsize=fontsize) @@ -63,9 +67,10 @@ def main(config: RunConfig): plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize) plt.grid(True, axis="y") plt.savefig(figures_path / "train_loss.png", dpi=300, bbox_inches="tight") + plt.close(figures[-1]) # Plot validation dice - plt.figure() + figures.append(plt.figure()) plt.plot(epochs, val_dice, label="Val Dice Score", linewidth=line_thickness) plt.xlabel("Epoch", fontsize=fontsize) plt.ylabel("Val dice score", fontsize=fontsize) @@ -74,10 +79,11 @@ def main(config: RunConfig): plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize) plt.grid(True, axis="y") plt.savefig(figures_path / "val_dice.png", dpi=300, bbox_inches="tight") + plt.close(figures[-1]) # Plot validation loss if available if val_loss: - plt.figure() + figures.append(plt.figure()) plt.plot(epochs, val_loss, label="Val Loss", linewidth=line_thickness) plt.xlabel("Epoch", fontsize=fontsize) plt.ylabel("Val loss", fontsize=fontsize) @@ -86,5 +92,12 @@ def main(config: RunConfig): plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize) plt.grid(True, axis="y") plt.savefig(figures_path / "val_loss.png", dpi=300, bbox_inches="tight") + plt.close(figures[-1]) except Exception as e: logger.error(f"Failed to generate figures: {e}") + finally: + # Errors here are logged and swallowed, so the close has to be + # unconditional: a failure between figure() and savefig() would + # otherwise leak exactly the figure nobody goes looking for. + for figure in figures: + plt.close(figure) diff --git a/tests/test_reporting.py b/tests/test_reporting.py index 4a462ca..af50a04 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -57,6 +57,50 @@ def test_figures_dir_idempotent(self, tmp_path): assert (run_dir / "figures" / "train_loss.png").exists() +class TestFigureLifetime: + """R43: standard_viz closes every figure it opens. + + ``worker.main`` calls ``standard_viz.main`` in-process once per sweep + combination, and pyplot keeps a strong reference to every unclosed figure, + so the canvases (and their Figure/Axes/Line objects) accumulate for the + whole sweep -- 36 live figures / 42 MiB after 12 combinations, plus + matplotlib's max_open_warning from the seventh on. + """ + + def _config(self, tmp_path, name): + run_dir = tmp_path / name + run_dir.mkdir() + (run_dir / "train_stats.csv").write_text( + "epoch,overall_loss,val_dice,val_loss_avg\n1,0.9,0.40,0.8\n2,0.5,0.70,0.4\n" + ) + return SimpleNamespace( + run_dir=str(run_dir), vol_size=32, n_categories=5, unet_layers=2 + ) + + def test_no_figures_left_open(self, tmp_path): + """Repeated calls (a sweep) leave no figure behind.""" + plt.close("all") + for i in range(3): + standard_viz.main(self._config(tmp_path, f"run{i}")) + assert plt.get_fignums() == [], f"figures leaked after call {i}" + + def test_no_figures_left_open_when_plotting_fails(self, tmp_path, monkeypatch): + """A failure between figure() and savefig() does not strand a figure. + + ``main`` logs and swallows plotting errors, so without an unconditional + close the leak survives exactly the case it is hardest to notice. + """ + plt.close("all") + + def boom(*args, **kwargs): + raise OSError("[Errno 28] No space left on device") + + monkeypatch.setattr(plt, "savefig", boom) + standard_viz.main(self._config(tmp_path, "failing_run")) + + assert plt.get_fignums() == [] + + class TestDiceFigure: """F70: Validation Dice figure saved as val_dice.png, not val_loss.png.""" From cb1a22c8ebed031e17279812835485dc223a3453 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:43:41 -0700 Subject: [PATCH 49/54] Make activation checkpointing reachable from the config The U-Net's use_checkpointing (made correct by F46) had no config key and no production caller, so the memory/compute trade it offers could not be taken. Per the user's decision to wire the feature rather than delete it, an activation_checkpointing 0/1 flag now enables it in worker.py, before the DDP wrap hides the method behind .module. R44 --- ScaFFold/configs/benchmark_default.yml | 1 + ScaFFold/utils/config_utils.py | 19 +++++++++ ScaFFold/worker.py | 9 ++++ tests/test_config.py | 33 ++++++++++++++ tests/test_worker_dist.py | 59 ++++++++++++++++++++++---- 5 files changed, 112 insertions(+), 9 deletions(-) diff --git a/ScaFFold/configs/benchmark_default.yml b/ScaFFold/configs/benchmark_default.yml index c41f4ed..29be50f 100644 --- a/ScaFFold/configs/benchmark_default.yml +++ b/ScaFFold/configs/benchmark_default.yml @@ -39,5 +39,6 @@ loss_freq: 1 # Number of epochs between logging the overal normalize: 1 # Cateogry search normalization parameter group_norm_groups: 8 # Number of groups used by GroupNorm in the UNet blocks. warmup_batches: 64 # How many warmup batches per rank to run before training. +activation_checkpointing: 0 # If 1, recompute UNet block activations during the backward pass instead of storing them: less memory, more compute. ce_weight_sample_fraction: 0.1 # Fraction of training masks to sample when estimating background vs foreground CE weights. dataset_reuse_enforce_commit_id: 0 # Enforce matching commit IDs for dataset reuse. diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 92af20c..60731fe 100644 --- a/ScaFFold/utils/config_utils.py +++ b/ScaFFold/utils/config_utils.py @@ -27,6 +27,19 @@ def require_positive_int(name: str, value: int) -> int: return value +def require_flag(name: str, value) -> bool: + """Validate an on/off config toggle written as 0/1 (or a YAML boolean). + + ``bool(value)`` would quietly accept ``2``, ``-1`` or ``"no"`` (all true), + so a mistyped toggle would enable the feature it was meant to disable. + """ + if isinstance(value, bool): + return value + if isinstance(value, int) and value in (0, 1): + return bool(value) + raise ValueError(f"{name} must be 0, 1, or a boolean; got {value!r}") + + def validate_unet_dims(problem_scale, unet_bottleneck_dim) -> int: """Check that ``problem_scale``/``unet_bottleneck_dim`` describe a real U-Net. @@ -105,6 +118,7 @@ class Config: "normalize", "group_norm_groups", "warmup_batches", + "activation_checkpointing", "ce_weight_sample_fraction", "dataset_reuse_enforce_commit_id", "target_dice", @@ -163,6 +177,7 @@ class Config: "loss_freq", "group_norm_groups", "warmup_batches", + "activation_checkpointing", "ce_weight_sample_fraction", "target_dice", "checkpoint_interval", @@ -252,6 +267,10 @@ def __init__(self, config_dict, strict=True): self.normalize = config_dict["normalize"] self.group_norm_groups = config_dict.get("group_norm_groups", 8) self.warmup_batches = config_dict.get("warmup_batches") + self.activation_checkpointing = require_flag( + "activation_checkpointing", + config_dict.get("activation_checkpointing", 0), + ) self.ce_weight_sample_fraction = config_dict.get( "ce_weight_sample_fraction", 0.1 ) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 2f36b9a..40458ba 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -227,6 +227,15 @@ def main(kwargs_dict: dict = {}): ) model = model.to(device, memory_format=torch.channels_last_3d) + if config.activation_checkpointing: + # Has to happen before the DDP wrap: afterwards the model is only + # reachable as ``model.module``, and the wrapper does not forward the + # method. + log.info( + "activation_checkpointing TRUE -- recomputing block activations in " + "the backward pass instead of storing them" + ) + model.use_checkpointing() # Wrap with DistConvDDP that corrects gradient scaling for dc submesh model = wrap_model_ddp(model, device, ps) # Store ps for use in the training loop diff --git a/tests/test_config.py b/tests/test_config.py index 28cc15b..086a4f5 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -98,6 +98,39 @@ def test_invalid_type_message_names_type(tmp_path): config_utils.load_config(str(path), "bogus") +def test_activation_checkpointing_is_a_real_option(): + """R44: activation checkpointing is reachable from a config, defaulting off. + + The U-Net has always had ``use_checkpointing``, but with no config key and + no caller it could not be turned on: any attempt was rejected as an unknown + key. + """ + cfg = config_utils.Config(dict(BASE)) + assert cfg.activation_checkpointing is False + assert ( + config_utils.Config({**BASE, "activation_checkpointing": 1}) + ).activation_checkpointing is True + assert ( + config_utils.Config({**BASE, "activation_checkpointing": True}) + ).activation_checkpointing is True + assert ( + config_utils.Config({**BASE, "activation_checkpointing": 0}) + ).activation_checkpointing is False + + +@pytest.mark.parametrize("value", [2, -1, "yes", 0.0]) +def test_activation_checkpointing_rejects_non_flag_values(value): + """Anything that is not a 0/1 (or bool) toggle is rejected by name.""" + with pytest.raises(ValueError, match="activation_checkpointing"): + config_utils.Config({**BASE, "activation_checkpointing": value}) + + +def test_activation_checkpointing_documented_in_the_default_config(): + """The shipped config is the parameter reference, so the key lives there.""" + text = (CONFIG_DIR / "benchmark_default.yml").read_text() + assert "activation_checkpointing:" in text + + def test_async_save_is_real_option(): """async_save is an accepted, defaulted option (consumed by the trainer).""" cfg = config_utils.Config(dict(BASE)) diff --git a/tests/test_worker_dist.py b/tests/test_worker_dist.py index e3441ea..b5c6b32 100644 --- a/tests/test_worker_dist.py +++ b/tests/test_worker_dist.py @@ -172,21 +172,21 @@ def fake_ddp(model, parallel_strategy=None, **kwargs): # --------------------------------------------------------------------------- -def test_worker_singleton_smoke(monkeypatch, tiny_config, tiny_dataset): - """worker.main completes end to end as a one-rank gloo job on CPU. +def _run_singleton_worker( + monkeypatch, tiny_config, tiny_dataset, *, port, config_overrides=None +): + """Run ``worker.main`` as a one-rank gloo job on CPU; return (rc, trainer). - ScaFFold always runs distributed; the supported singleton case is a - one-rank launch. The worker initializes the (gloo) process group itself, - builds a real unsharded ParallelStrategy, and tears the group down before - rank-0 post-processing. + Training itself is stubbed (one synthetic epoch row so post-processing has + data), so what this exercises is the worker's own setup path. """ monkeypatch.setenv("MASTER_ADDR", "127.0.0.1") - monkeypatch.setenv("MASTER_PORT", "29513") + monkeypatch.setenv("MASTER_PORT", str(port)) # Force the CPU path so initialize_dist selects gloo: this test must not # depend on a working GPU/NCCL stack. monkeypatch.setattr(torch.cuda, "is_available", lambda: False) - cfg = tiny_config() + cfg = tiny_config(**(config_overrides or {})) kwargs = dict(vars(cfg)) kwargs.update( { @@ -213,7 +213,20 @@ def fake_train(self, profiler=None): monkeypatch.setattr(worker_mod.PyTorchTrainer, "train", fake_train) result = worker_mod.main(kwargs_dict=kwargs) - trainer = seen.get("trainer") + return result, seen.get("trainer") + + +def test_worker_singleton_smoke(monkeypatch, tiny_config, tiny_dataset): + """worker.main completes end to end as a one-rank gloo job on CPU. + + ScaFFold always runs distributed; the supported singleton case is a + one-rank launch. The worker initializes the (gloo) process group itself, + builds a real unsharded ParallelStrategy, and tears the group down before + rank-0 post-processing. + """ + result, trainer = _run_singleton_worker( + monkeypatch, tiny_config, tiny_dataset, port=29513 + ) assert result == 0 assert trainer is not None @@ -229,6 +242,34 @@ def fake_train(self, profiler=None): assert not torch.distributed.is_initialized() +# --------------------------------------------------------------------------- +# R44: the activation-checkpointing config flag reaches the model +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("flag, expected", [(0, False), (1, True)]) +def test_activation_checkpointing_flag_reaches_the_model( + monkeypatch, tiny_config, tiny_dataset, flag, expected +): + """``activation_checkpointing: 1`` turns the U-Net's flag on. + + ``use_checkpointing`` had no caller at all, and it has to be invoked + *before* the DDP wrap: afterwards the model is only reachable through + ``.module``, which is exactly why this asserts on the wrapped model's + inner module. + """ + _result, trainer = _run_singleton_worker( + monkeypatch, + tiny_config, + tiny_dataset, + port=29520 + flag, + config_overrides={"activation_checkpointing": flag}, + ) + + model = getattr(trainer.model, "module", trainer.model) + assert model.checkpointing is expected + + # --------------------------------------------------------------------------- # Local size detection (R23) # --------------------------------------------------------------------------- From c7e3f5d57a156fe35b2be21413ec644eda18f392 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 17:51:48 -0700 Subject: [PATCH 50/54] Compile the UNet GroupNorm on GPU ATen's GroupNorm computes its per-group statistics with a kernel that launches one workgroup per (batch, group) row. At the benchmark defaults (local_batch_size=1, group_norm_groups=8) that is 8 of an MI300A's 228 CUs, and GroupNorm was 86.98 ms of the 186.97 ms scale-7 training step -- 46.5% of wall time, the single largest cost in the step (R38). FastGroupNorm subclasses nn.GroupNorm and routes its forward through a lazily built, process-wide torch.compile(F.group_norm, dynamic=False, fullgraph=True), which hands the reduction to Inductor to tile across the whole device. The module holds the same weight/bias under the same names with no new buffers, so checkpoints round-trip in both directions against a plain nn.GroupNorm build. The compiled path is used only where it is safe, and every rejection falls back to the stock kernel: CPU tensors (the CPU suite pays no compile latency), tensor subclasses such as DistConv's DCTensor (Dynamo cannot trace __torch_dispatch__ wrappers), an already-compiled enclosing region, an explicit SCAFFOLD_GROUPNORM_COMPILE=0, and -- permanently, with one warning -- any exception out of torch.compile. Dynamo's per-function recompile limit is raised from its stock 8: one UNet needs 10 cache entries (5 distinct GroupNorm shapes, each again under no_grad for evaluation), and past the limit Dynamo gives up and silently reverts every GroupNorm to the slow kernel. Measured on one MI300A with review/round2/repros/perf/step_bench.py --layout cl (1x3x128^3, layers=4, bf16 autocast, GradScaler disabled, warm MIOpen db), median of 20 steps, eager numbers from the same build with SCAFFOLD_GROUPNORM_COMPILE=0: step 184.69 ms -> 100.71 ms (1.83x) forward 104.04 ms -> 30.64 ms backward 75.59 ms -> 64.66 ms peak alloc 9.80 GiB -> 8.22 GiB --batch 2 279.58 ms -> 187.61 ms (1.49x; no regression at B>1) torch.profiler over the same step: GroupNorm 86.98 ms/step (46.5% of wall) -> 6.56 ms/step (6.7% of a 97.91 ms step). Compilation is one-time: the first compiled step takes 15.15 s with a cold Inductor cache and produces 5 graphs, after which 12 steps at 97.4 ms produce none; the first no_grad forward adds the other 5. The default 64 warmup batches absorb it outside every timed epoch. Determinism: no gate needed. Two separate processes running three fwd+bwd+Adam steps of the scale-7 UNet under the more_determinism settings (use_deterministic_algorithms(True, warn_only=True), cudnn.benchmark=False, fixed seeds) hash bitwise identically with the compiled path, exactly as they do with the eager one, so no config flag is plumbed through. Not yet visible in the default configuration: worker.py wraps activations in DCTensor even at dc_num_shards=[1,1,1], and GroupNorm then keeps the eager kernel. Verified through a worker.py-shaped DistConvDDP harness that all three paths (DCTensor, plain tensors, plain + use_checkpointing) train without error and agree on their losses; the speedup lands as soon as the unsharded wrap is skipped (R39). --- ScaFFold/unet/group_norm.py | 203 +++++++++++++++++ ScaFFold/unet/unet_parts.py | 6 +- tests/test_groupnorm.py | 421 ++++++++++++++++++++++++++++++++++++ 3 files changed, 629 insertions(+), 1 deletion(-) create mode 100644 ScaFFold/unet/group_norm.py create mode 100644 tests/test_groupnorm.py diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py new file mode 100644 index 0000000..1134924 --- /dev/null +++ b/ScaFFold/unet/group_norm.py @@ -0,0 +1,203 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""GroupNorm with a ``torch.compile``d fast path on GPU. + +ATen's GroupNorm computes its per-group statistics with a kernel that launches +one workgroup per ``(batch, group)`` row. At this benchmark's defaults +(``local_batch_size=1``, ``group_norm_groups=8``) that is 8 workgroups, so on a +228-CU MI300A the normalization runs at a small fraction of achievable +bandwidth and dominates the step: measured 87 ms of a 187 ms step (47%) at +scale 7. Compiling the same functional GroupNorm hands the reduction to +Inductor, which tiles it across the whole device; the same measurement then +gives a 184.7 ms step at 100.7 ms, with GroupNorm down to ~7% of it. + +``FastGroupNorm`` is a drop-in ``nn.GroupNorm``: same parameters, same names, +same shapes, same numerics -- only the kernel differs, so checkpoints are +interchangeable in both directions with any other GroupNorm-based build. The +compiled path is used only when it is safe and worthwhile, and every rejection +falls back to stock eager ``F.group_norm``: + +* non-CUDA tensors (the CPU test suite never pays compile latency), +* tensor subclasses such as DistConv's ``DCTensor``, whose ``__torch_dispatch__`` + wrapper Dynamo cannot trace, +* an already-compiled enclosing region (the functional call inlines instead), +* an explicit opt-out via ``SCAFFOLD_GROUPNORM_COMPILE=0``, +* any failure inside ``torch.compile`` -- logged once, then eager forever after. + +Determinism: the compiled kernels are bitwise reproducible. Two separate +processes running three fwd+bwd+Adam steps of the scale-7 UNet under +``more_determinism`` (``use_deterministic_algorithms(True, warn_only=True)``, +``cudnn.benchmark=False``, fixed seeds) hash identically with the compiled path, +exactly as they do with the eager one, so no determinism gate is needed. +""" + +import logging +import os + +import torch +import torch.nn as nn +import torch.nn.functional as F + +logger = logging.getLogger(__name__) + +#: Opt-out (``0``/``false``/``off``/``no``) or explicit opt-in (``1``/``true``/ +#: ``on``/``yes``) for the compiled GroupNorm path. Unset means "on wherever it +#: is safe", which is what every production run wants. +COMPILE_ENV_VAR = "SCAFFOLD_GROUPNORM_COMPILE" + +#: Dynamo caches one entry per distinct guard set on the traced function. A +#: UNet presents one entry per distinct activation shape (5 at scale 7) times +#: grad-enabled/no-grad (training vs. evaluation), i.e. 10 -- above the stock +#: limit of 8, which would silently drop the whole model back to eager mid-run. +#: The traced function is a single ``F.group_norm`` call, so the extra entries +#: cost only their one-time compilation. +_MIN_RECOMPILE_LIMIT = 64 + +# Lazily built on the first eligible forward: importing ScaFFold must not drag +# in Dynamo, and a run that never reaches the GPU must not pay for it. +_compiled_group_norm = None + +# Set once if torch.compile raises; the eager path is then used everywhere. +_compile_failed = False + +# None = decide per tensor; True/False = forced by SCAFFOLD_GROUPNORM_COMPILE or +# by set_compile_enabled(). +_compile_override = None + + +def _env_override(): + """Read ``SCAFFOLD_GROUPNORM_COMPILE``; ``None`` when unset or unparsable.""" + raw = os.environ.get(COMPILE_ENV_VAR) + if raw is None: + return None + value = raw.strip().lower() + if value in ("1", "true", "on", "yes"): + return True + if value in ("0", "false", "off", "no"): + return False + logger.warning( + f"Ignoring unrecognized {COMPILE_ENV_VAR}={raw!r}; " + "expected one of 1/0/true/false/on/off/yes/no" + ) + return None + + +_compile_override = _env_override() + + +def set_compile_enabled(enabled): + """Force the compiled path on (``True``) or off (``False``). + + ``None`` restores the default, which is the environment variable if set and + otherwise "compile wherever it is safe". Forcing it on does not override + the device and tensor-subclass checks -- those are correctness conditions, + not preferences. Returns the previous setting so callers (tests) can + restore it. + """ + global _compile_override + previous = _compile_override + _compile_override = _env_override() if enabled is None else bool(enabled) + return previous + + +def _group_norm(input, num_groups, weight, bias, eps): + """The function Dynamo traces: plain functional GroupNorm, nothing else.""" + return F.group_norm(input, num_groups, weight, bias, eps) + + +def _raise_recompile_limit(): + """Lift Dynamo's per-function recompile cap to cover every UNet GN shape. + + Only ever raises it, so a caller that deliberately set a larger limit keeps + theirs. ``cache_size_limit`` is the older spelling of ``recompile_limit``; + set whichever exists. + """ + config = torch._dynamo.config + for name in ("recompile_limit", "cache_size_limit"): + current = getattr(config, name, None) + if isinstance(current, int) and current < _MIN_RECOMPILE_LIMIT: + setattr(config, name, _MIN_RECOMPILE_LIMIT) + + +def _get_compiled_group_norm(): + """Build (once) the compiled functional GroupNorm shared by every module. + + One compiled callable for the whole model, not one per module: the shapes, + not the instances, are what Dynamo specializes on, and sharing keeps the + 18 GroupNorms of a scale-7 UNet down to 5 compilations. ``dynamic=False`` + keeps the specialized kernels (this benchmark runs fixed shapes); + ``fullgraph=True`` turns anything Dynamo cannot handle into an exception we + catch, rather than a silent graph break that reintroduces the slow kernel. + """ + global _compiled_group_norm + if _compiled_group_norm is None: + _raise_recompile_limit() + _compiled_group_norm = torch.compile(_group_norm, dynamic=False, fullgraph=True) + return _compiled_group_norm + + +def _use_compiled(input): + """Whether this particular input should take the compiled path.""" + if _compile_failed or _compile_override is False: + return False + # Tensor subclasses (DistConv's DCTensor) route their ops through + # __torch_dispatch__, which Dynamo cannot trace; eager keeps the wrapper's + # semantics -- including which of its outputs come back wrapped -- exactly + # as they are today. worker.py wraps activations in DCTensor even at + # dc_num_shards=[1,1,1], so this fast path engages once that wrap is + # skipped for the unsharded case (or whenever the model is driven with + # plain tensors, as the tests and the standalone benchmarks do). + if type(input) is not torch.Tensor: + return False + # CPU GroupNorm is not the bottleneck and compiling it would put a + # multi-second C++ build in front of every unit test. + if not input.is_cuda: + return False + # Already inside a compiled region: let the functional call be inlined. + if torch.compiler.is_compiling(): + return False + return True + + +class FastGroupNorm(nn.GroupNorm): + """``nn.GroupNorm`` that runs its GPU forward through ``torch.compile``. + + Identical state: ``weight``/``bias`` of shape ``(num_channels,)``, no + buffers, so state dicts are interchangeable with plain ``nn.GroupNorm`` + in both directions. + """ + + def forward(self, input): + # super().forward() is the stock kernel; deferring to it keeps the eager + # path identical to nn.GroupNorm's by construction. + if not _use_compiled(input): + return super().forward(input) + global _compile_failed + try: + return _get_compiled_group_norm()( + input, self.num_groups, self.weight, self.bias, self.eps + ) + except Exception as e: + # Compilation is an optimization, never a correctness requirement: + # a broken Inductor/Triton install, an unwritable cache directory or + # an untraceable input must degrade to the stock kernel, not kill a + # multi-node run. GroupNorm is pure, so retrying eagerly is safe. + _compile_failed = True + logger.warning( + f"torch.compile of GroupNorm failed ({type(e).__name__}: {e}); " + "falling back to the eager kernel for the rest of this run. " + f"Set {COMPILE_ENV_VAR}=0 to skip this attempt entirely." + ) + return super().forward(input) diff --git a/ScaFFold/unet/unet_parts.py b/ScaFFold/unet/unet_parts.py index 681e44b..c9e6cb0 100644 --- a/ScaFFold/unet/unet_parts.py +++ b/ScaFFold/unet/unet_parts.py @@ -20,6 +20,8 @@ from ScaFFold.utils.perf_measure import annotate +from .group_norm import FastGroupNorm + _doubleconv_annotate = annotate(fmt="DoubleConv.{}") _down_annotate = annotate(fmt="Down.{}") _up_annotate = annotate(fmt="Up.{}") @@ -31,7 +33,9 @@ def _group_norm(num_groups, num_channels): raise ValueError( f"group_norm_groups={num_groups} must evenly divide num_channels={num_channels}" ) - return nn.GroupNorm(num_groups, num_channels) + # FastGroupNorm is nn.GroupNorm plus a compiled GPU kernel; it holds the + # same parameters under the same names, so checkpoints are unaffected. + return FastGroupNorm(num_groups, num_channels) class DoubleConv(nn.Module): diff --git a/tests/test_groupnorm.py b/tests/test_groupnorm.py new file mode 100644 index 0000000..b1e8e6a --- /dev/null +++ b/tests/test_groupnorm.py @@ -0,0 +1,421 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Tests for the compiled GroupNorm fast path (``ScaFFold.unet.group_norm``). + +The optimization must be invisible everywhere except in the profile: the same +state dict as a stock ``nn.GroupNorm`` model (checkpoints stay interchangeable +in both directions), the same numbers within reduction-order noise, and an +eager fallback for every input the compiled kernel cannot or should not take +(CPU, tensor subclasses such as DistConv's ``DCTensor``, a broken compiler). +""" + +from __future__ import annotations + +import logging + +import pytest +import torch +import torch.nn as nn + +from ScaFFold.unet import group_norm as gn_mod +from ScaFFold.unet.group_norm import FastGroupNorm +from ScaFFold.unet.unet_model import UNet + +_N = 16 +_N_CHANNELS = 3 +_N_CLASSES = 2 +_GROUPS = 8 + + +@pytest.fixture(autouse=True) +def _restore_compile_state(): + """Keep per-test overrides of the module-level compile state contained.""" + previous = gn_mod.set_compile_enabled(None) + failed = gn_mod._compile_failed + yield + gn_mod._compile_override = previous + gn_mod._compile_failed = failed + + +def _make_unet(seed: int, group_norm_cls=None): + """Build the worker.py-shaped UNet, optionally with a different norm class.""" + torch.manual_seed(seed) + if group_norm_cls is None: + return UNet( + n_channels=_N_CHANNELS, + n_classes=_N_CLASSES, + trilinear=False, + layers=2, + group_norm_groups=_GROUPS, + ) + import ScaFFold.unet.unet_parts as parts + + original = parts.FastGroupNorm + parts.FastGroupNorm = group_norm_cls + try: + return UNet( + n_channels=_N_CHANNELS, + n_classes=_N_CLASSES, + trilinear=False, + layers=2, + group_norm_groups=_GROUPS, + ) + finally: + parts.FastGroupNorm = original + + +def _make_input(seed: int = 0, channels: int = _N_CHANNELS, size: int = _N): + generator = torch.Generator().manual_seed(seed) + return torch.randn(1, channels, size, size, size, generator=generator) + + +# --------------------------------------------------------------------------- +# state dict compatibility +# --------------------------------------------------------------------------- + + +def test_state_dict_matches_plain_groupnorm_model(): + """Names, shapes and dtypes must be unchanged from the nn.GroupNorm build. + + A checkpoint written before this optimization has to keep loading, so the + parameter inventory of the model may not shift by even one key. + """ + new_model = _make_unet(seed=0) + old_model = _make_unet(seed=0, group_norm_cls=nn.GroupNorm) + + new_sd = new_model.state_dict() + old_sd = old_model.state_dict() + assert list(new_sd.keys()) == list(old_sd.keys()) + for key in old_sd: + assert new_sd[key].shape == old_sd[key].shape, key + assert new_sd[key].dtype == old_sd[key].dtype, key + # The optimization must not have introduced buffers either. + assert [name for name, _ in new_model.named_buffers()] == [ + name for name, _ in old_model.named_buffers() + ] + + +def test_checkpoint_round_trip_both_directions(tmp_path): + """An old checkpoint loads into the new model and vice versa, strict=True. + + Both directions matter: runs resumed onto the new code must accept old + checkpoints, and checkpoints written by the new code must stay readable by + anything still building plain ``nn.GroupNorm`` (e.g. an older analysis + script). After each load the two models must agree bit for bit. + """ + new_model = _make_unet(seed=0) + old_model = _make_unet(seed=1, group_norm_cls=nn.GroupNorm) + x = _make_input(seed=3) + + old_path = tmp_path / "old.pth" + torch.save({"model_state_dict": old_model.state_dict()}, old_path) + loaded = torch.load(old_path, weights_only=True) + missing = new_model.load_state_dict(loaded["model_state_dict"], strict=True) + assert not missing.missing_keys and not missing.unexpected_keys + + new_model.eval() + old_model.eval() + with torch.no_grad(): + assert torch.equal(new_model(x), old_model(x)) + + # ... and the reverse: new checkpoint into the plain-GroupNorm model. + fresh_new = _make_unet(seed=2) + new_path = tmp_path / "new.pth" + torch.save({"model_state_dict": fresh_new.state_dict()}, new_path) + reloaded = torch.load(new_path, weights_only=True) + result = old_model.load_state_dict(reloaded["model_state_dict"], strict=True) + assert not result.missing_keys and not result.unexpected_keys + + fresh_new.eval() + with torch.no_grad(): + assert torch.equal(old_model(x), fresh_new(x)) + + +def test_unet_uses_fast_group_norm(): + """Every norm layer in the model is the fast one -- no half-converted build.""" + model = _make_unet(seed=0) + norms = [m for m in model.modules() if isinstance(m, nn.GroupNorm)] + assert norms, "UNet should contain GroupNorm layers" + assert all(isinstance(m, FastGroupNorm) for m in norms) + + +# --------------------------------------------------------------------------- +# CPU behavior: identical numerics, and no compilation at all +# --------------------------------------------------------------------------- + + +def test_cpu_output_bit_identical_to_eager(): + """On CPU the fast module is literally the stock kernel, so bits must match.""" + fast = FastGroupNorm(_GROUPS, 64) + plain = nn.GroupNorm(_GROUPS, 64) + with torch.no_grad(): + plain.weight.copy_(fast.weight) + plain.bias.copy_(fast.bias) + x = _make_input(seed=5, channels=64, size=8) + assert torch.equal(fast(x), plain(x)) + + +def test_cpu_never_invokes_torch_compile(monkeypatch): + """The CPU unit suite must not pay Inductor's compile latency. + + Guards the ``input.is_cuda`` check: if it ever regresses, a CPU-only test + run would start building C++ kernels for every GroupNorm shape. + """ + calls = [] + + def _boom(*a, **kw): + calls.append(a) + raise AssertionError("torch.compile must not be called for CPU tensors") + + monkeypatch.setattr(torch, "compile", _boom) + monkeypatch.setattr(gn_mod, "_compiled_group_norm", None) + gn_mod.set_compile_enabled(True) # even when explicitly forced on + + model = _make_unet(seed=0) + with torch.no_grad(): + model(_make_input(seed=6)) + assert not calls + + +def test_tensor_subclass_input_stays_eager(): + """DistConv wraps activations in a ``__torch_dispatch__`` tensor subclass. + + Dynamo cannot trace those wrappers, so the predicate must reject anything + that is not exactly ``torch.Tensor`` before a compile is attempted. + """ + + class _Wrapper(torch.Tensor): + pass + + plain = torch.randn(1, 8, 4, 4, 4) + assert gn_mod._use_compiled(plain) is False # CPU + assert gn_mod._use_compiled(plain.as_subclass(_Wrapper)) is False + + +def test_compile_failure_falls_back_to_eager(monkeypatch, caplog): + """A broken compiler degrades to the stock kernel instead of killing the run. + + Simulated by making the compiled callable raise; the module must return the + eager result, warn once, and stop trying for the rest of the process. + """ + + def _raises(*args, **kwargs): + raise RuntimeError("simulated Inductor failure") + + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) + gn_mod._compile_failed = False + + fast = FastGroupNorm(_GROUPS, 64) + x = _make_input(seed=7, channels=64, size=8) + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + out = fast(x) + assert torch.equal( + out, nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias) + ) + assert any("falling back to the eager kernel" in r.message for r in caplog.records) + assert gn_mod._compile_failed is True + # Latched off: the predicate now refuses even a would-be eligible tensor. + monkeypatch.undo() + assert gn_mod._use_compiled(torch.randn(1, 8, 4, 4, 4)) is False + + +@pytest.mark.parametrize( + "value,expected", + [ + ("0", False), + ("false", False), + ("OFF", False), + ("no", False), + ("1", True), + ("true", True), + ("On", True), + ("yes", True), + ("maybe", None), + ], +) +def test_env_var_controls_the_fast_path(monkeypatch, value, expected): + """``SCAFFOLD_GROUPNORM_COMPILE`` is the documented run-time opt-out.""" + monkeypatch.setenv(gn_mod.COMPILE_ENV_VAR, value) + gn_mod.set_compile_enabled(None) + assert gn_mod._compile_override is expected + + +def test_env_var_unset_means_auto(monkeypatch): + monkeypatch.delenv(gn_mod.COMPILE_ENV_VAR, raising=False) + gn_mod.set_compile_enabled(None) + assert gn_mod._compile_override is None + + +def test_recompile_limit_is_raised_never_lowered(): + """Dynamo's stock cap of 8 is below what one UNet needs. + + A scale-7 UNet presents 5 distinct GroupNorm shapes, and evaluation runs the + same 5 again under ``no_grad`` -- 10 cache entries (measured). Past the cap + Dynamo gives up and every GroupNorm silently reverts to the slow kernel, so + the module raises the limit; it must never lower one a caller chose. + """ + import torch._dynamo + + config = torch._dynamo.config + name = ( + "recompile_limit" if hasattr(config, "recompile_limit") else "cache_size_limit" + ) + original = getattr(config, name) + try: + setattr(config, name, 8) + gn_mod._raise_recompile_limit() + assert getattr(config, name) >= 10 + setattr(config, name, 4096) + gn_mod._raise_recompile_limit() + assert getattr(config, name) == 4096 + finally: + setattr(config, name, original) + + +# --------------------------------------------------------------------------- +# GPU behavior: numerics, single compile, checkpointing +# --------------------------------------------------------------------------- + + +def _assert_close(actual, expected, tol, what): + diff = (actual.float() - expected.float()).abs().max().item() + scale = expected.float().abs().max().item() + assert diff <= tol * max(scale, 1.0), f"{what}: max|diff|={diff:.3e}" + return diff + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape", [(1, 64, 32, 32, 32), (2, 128, 16, 16, 16)]) +@pytest.mark.parametrize("autocast", [False, True]) +def test_gpu_compiled_matches_eager(shape, autocast): + """Compiled forward and gradients match eager, fp32 and under bf16 autocast. + + ``(1, 64, 32^3)`` is the hot production shape (``[1, 64, 128^3]``) at + reduced size -- same channel count and group count, same reduction + structure, small enough for a unit test. Tolerances are loose enough for + reduction-order differences and tight enough to catch a real numerics bug; + observed maxima on MI300A are ~1e-6 relative. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(11) + x = torch.randn(*shape, device=device, generator=generator) + grad_out = torch.randn(*shape, device=device, generator=generator) + + fast = FastGroupNorm(_GROUPS, shape[1]).to(device) + with torch.no_grad(): + fast.weight.normal_(1.0, 0.1, generator=generator) + fast.bias.normal_(0.0, 0.1, generator=generator) + + def run(compiled): + gn_mod.set_compile_enabled(compiled) + inp = x.clone().requires_grad_(True) + fast.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast): + out = fast(inp) + out.backward(grad_out.to(out.dtype)) + return ( + out.detach(), + inp.grad.detach(), + fast.weight.grad.detach().clone(), + fast.bias.grad.detach().clone(), + ) + + eager = run(False) + compiled = run(True) + assert gn_mod._compiled_group_norm is not None, "compiled path was not taken" + assert not gn_mod._compile_failed + + _assert_close(compiled[0], eager[0], 1e-5, "output") + _assert_close(compiled[1], eager[1], 1e-4, "d_input") + _assert_close(compiled[2], eager[2], 1e-4, "d_weight") + _assert_close(compiled[3], eager[3], 1e-4, "d_bias") + # Autocast policy must be preserved: GroupNorm is an fp32 op, so the + # compiled path may not quietly hand back bf16 activations. + assert compiled[0].dtype == eager[0].dtype + + +@pytest.mark.gpu +def test_gpu_steady_state_does_not_recompile(): + """Fixed shapes must compile once and then never again. + + A recompile inside a timed epoch would show up as a multi-second outlier in + ``epoch_duration`` (and therefore the FOM), so the guard set has to be + stable across steps. + """ + from torch._dynamo.utils import counters + + gn_mod.set_compile_enabled(True) + device = torch.device("cuda") + fast = FastGroupNorm(_GROUPS, 64).to(device) + x = torch.randn(1, 64, 16, 16, 16, device=device, requires_grad=True) + + def step(): + fast.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = fast(x) + out.sum().backward() + + step() # first call: compiles + before = counters["stats"]["unique_graphs"] + for _ in range(5): + step() + assert counters["stats"]["unique_graphs"] == before, "recompiled in steady state" + + +@pytest.mark.gpu +def test_gpu_activation_checkpointing_matches_eager(): + """The compiled kernel must survive recompute under use_checkpointing(). + + Non-reentrant checkpointing replays the block's forward inside the backward + pass; a compiled region has to produce the same activations both times or + the gradients silently change. + + Compared as relative L2 error per gradient tensor, because whole-network + agreement is not bitwise even without this change: with cudnn.benchmark on + and bf16 autocast, two eager runs of this model differ by ~4e-3 relative + (measured), and checkpointing on vs. off differs by the same amount. + Measured here: compiled vs. eager 6.4e-3, i.e. the same order as that noise + floor -- while a genuinely wrong kernel would be O(1). + """ + device = torch.device("cuda") + x = _make_input(seed=9).to(device).requires_grad_(True) + tolerance = 5e-2 + + def grads(compiled, checkpointing): + gn_mod.set_compile_enabled(compiled) + model = _make_unet(seed=0).to(device) + if checkpointing: + model.use_checkpointing() + model.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = model(x) + out.float().pow(2).mean().backward() + return {n: p.grad.detach().clone() for n, p in model.named_parameters()} + + def assert_agrees(actual, expected, label): + for name in expected: + reference = expected[name].float() + error = (actual[name].float() - reference).norm().item() + relative = error / max(reference.norm().item(), 1e-12) + assert relative < tolerance, f"{label} {name}: rel L2 {relative:.3e}" + + eager = grads(False, True) + compiled = grads(True, True) + compiled_nockpt = grads(True, False) + assert gn_mod._compiled_group_norm is not None, "compiled path was not taken" + assert not gn_mod._compile_failed + assert_agrees(compiled, eager, "checkpointed grad") + assert_agrees(compiled_nockpt, compiled, "grad") From c669b3f747ecc75e2e30a8e01442f40c2a025cc9 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 18:37:39 -0700 Subject: [PATCH 51/54] Post the ragged-size all_gather before the empty-loader return MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A rank whose warmup fetched no batch returned before the ragged-size all_gather, skipping a collective its peers post — the divergence class this round closes elsewhere. Latent today (padded training shards make loader lengths rank-invariant), but the collective pattern must not depend on local loader state. Found by the final verification pass. --- ScaFFold/utils/trainer.py | 14 +++++++++++--- tests/test_perf_hotpath.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 9e53ecc..b284708 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -741,11 +741,15 @@ def _warmup_ragged_batches(self, batch): its peers did not, and the collectives inside that step (the gradient all-reduce, the sharded loss reductions) would deadlock. """ - if batch is None: - return local_batch_size = self.config.local_batch_size - available = batch["image"].shape[0] + # Agree on the ragged sizes FIRST. The all_gather below is a + # collective: every rank must post it even when its own warmup fetched + # no batch at all (``batch is None``), or its peers block in the + # gather. Batch availability is rank-invariant in practice (padded + # training shards give every rank the same loader length), but the + # collective pattern must not depend on local loader state, so the + # no-batch early return comes after the gather. ragged_sizes = {len(self.train_sampler) % local_batch_size} local_val_ragged = torch.tensor( [len(self.val_sampler) % local_batch_size], device=self.device @@ -756,6 +760,10 @@ def _warmup_ragged_batches(self, batch): torch.distributed.all_gather(gathered_val_ragged, local_val_ragged) ragged_sizes.update(int(size.item()) for size in gathered_val_ragged) + if batch is None: + return + available = batch["image"].shape[0] + # The batches already run are all ``available`` wide, and ``available`` # is itself rank-invariant (the padded training shards give every rank # the same leading batch size), so this loop is identical on all ranks. diff --git a/tests/test_perf_hotpath.py b/tests/test_perf_hotpath.py index 0eb35d0..a683c59 100644 --- a/tests/test_perf_hotpath.py +++ b/tests/test_perf_hotpath.py @@ -298,6 +298,36 @@ def fake_all_gather(tensor_list, tensor, *args, **kwargs): assert sizes == [3, 3, 1, 2] +def test_warmup_ragged_all_gather_is_posted_even_without_a_batch( + tiny_trainer, monkeypatch +): + # The ragged-size agreement is a collective: a rank whose warmup fetched + # no batch at all (empty train loader) must still post the all_gather, or + # its peers block in it. The no-batch early return has to come after the + # gather, even though such a rank then runs no extra step itself. + import torch.distributed as dist + + trainer = tiny_trainer( + n_train=4, + n_val=3, + config_overrides={"local_batch_size": 2, "warmup_batches": 2}, + ) + sizes = _stub_warmup_steps(trainer, monkeypatch) + gathers = [] + + def fake_all_gather(tensor_list, tensor, *args, **kwargs): + gathers.append(int(tensor.item())) + for out in tensor_list: + out.copy_(tensor) + + monkeypatch.setattr(dist, "all_gather", fake_all_gather) + + trainer._warmup_ragged_batches(None) + + assert gathers, "rank skipped the ragged-size all_gather when batch was None" + assert sizes == [] + + def test_warmup_rolls_back_state_including_the_ragged_batch(tiny_trainer, monkeypatch): # The extra ragged iteration stays inside warmup's snapshot/restore # envelope, so nothing it touches survives into training. From c825cc846a20531df438b6c22aab52e219eef750 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 18:37:39 -0700 Subject: [PATCH 52/54] Correct the best-copy perf claim and widen the copy buffer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The R08 comment claimed the copy halves the bytes pushed at the shared filesystem; measured on Lustre the copy is ~1.02x — the real saving is the serialization CPU. Note the redundancy trade-off (best is now a byte copy of last), use a 16 MiB copy buffer to cut syscall count on parallel filesystems, and document that the Dynamo recompile-limit raise clobbers a deliberately smaller limit. Found by the final verification pass. --- ScaFFold/unet/group_norm.py | 5 +++-- ScaFFold/utils/checkpointing.py | 15 +++++++++++---- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py index 1134924..06d7407 100644 --- a/ScaFFold/unet/group_norm.py +++ b/ScaFFold/unet/group_norm.py @@ -121,8 +121,9 @@ def _raise_recompile_limit(): """Lift Dynamo's per-function recompile cap to cover every UNet GN shape. Only ever raises it, so a caller that deliberately set a larger limit keeps - theirs. ``cache_size_limit`` is the older spelling of ``recompile_limit``; - set whichever exists. + theirs -- but note the converse: a limit deliberately set *smaller* than + ours is clobbered up to ``_MIN_RECOMPILE_LIMIT``. ``cache_size_limit`` is + the older spelling of ``recompile_limit``; set whichever exists. """ config = torch._dynamo.config for name in ("recompile_limit", "cache_size_limit"): diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 5a04ecd..00dfc14 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -663,7 +663,10 @@ def _atomic_copy(src, dst): tmp_path = dst.with_name(f"{dst.name}.tmp.{os.getpid()}") try: with open(src, "rb") as fsrc, open(tmp_path, "wb") as fdst: - shutil.copyfileobj(fsrc, fdst) + # A large buffer keeps the syscall count low on parallel + # filesystems (the default 64 KiB means ~1k read/write pairs + # per 64 MiB checkpoint). + shutil.copyfileobj(fsrc, fdst, length=16 * 1024 * 1024) fdst.flush() os.fsync(fdst.fileno()) os.replace(tmp_path, dst) @@ -688,9 +691,13 @@ def _write_to_disk(cls, state_dict, last_path, best_path, is_best, log): # Save 'last' atomically. cls._atomic_save(state_dict, last_path) # 'best' is byte-identical to the 'last' just committed, so copy - # that file instead of pickling and fsyncing the same state a - # second time (double the serialization CPU and double the bytes - # pushed at the shared filesystem on every improving epoch). + # that file instead of pickling the same state a second time. The + # saving is the serialization CPU (pickle + zip of the full state + # dict); the filesystem traffic is roughly a wash -- the copy + # writes the same bytes and adds a read (measured ~1.02x faster on + # Lustre). Trade-off: 'best' is now a byte copy of 'last', so a + # silently corrupted 'last' write would propagate into 'best' + # rather than being an independent serialization. # # There is no concurrent writer to race: checkpoint writes are # serialized through a single writer -- the caller's thread in sync From 3b8ead857fbd4e376e10c579dbcc8009fe5dae8f Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 31 Jul 2026 18:52:46 -0700 Subject: [PATCH 53/54] Drop the unreachable packaged fractal library The seed-keyed relayout (R29) made the shipped seed-unknown CSVs at ScaFFold/fractals/var0.15/3DIFS_param unreachable under any configuration. Remove them, their package-data glob, the dead Config.library_root (zero readers), and the README claim; libraries regenerate deterministically from the configured seed. Resolves verification item VB-5 per user decision. --- README.md | 4 +++- ScaFFold/fractals/var0.15/3DIFS_param/000000.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000001.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000002.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000003.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000004.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000005.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000006.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000007.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000008.csv | 3 --- ScaFFold/fractals/var0.15/3DIFS_param/000009.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000010.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000011.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000012.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000013.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000014.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000015.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000016.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000017.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000018.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000019.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000020.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000021.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000022.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000023.csv | 3 --- ScaFFold/fractals/var0.15/3DIFS_param/000024.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000025.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000026.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000027.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000028.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000029.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000030.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000031.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000032.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000033.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000034.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000035.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000036.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000037.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000038.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000039.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000040.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000041.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000042.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000043.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000044.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000045.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000046.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000047.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000048.csv | 2 -- ScaFFold/fractals/var0.15/3DIFS_param/000049.csv | 2 -- ScaFFold/utils/config_utils.py | 4 ---- pyproject.toml | 1 - 53 files changed, 3 insertions(+), 108 deletions(-) delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000000.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000001.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000002.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000003.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000004.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000005.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000006.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000007.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000008.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000009.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000010.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000011.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000012.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000013.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000014.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000015.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000016.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000017.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000018.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000019.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000020.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000021.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000022.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000023.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000024.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000025.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000026.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000027.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000028.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000029.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000030.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000031.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000032.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000033.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000034.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000035.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000036.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000037.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000038.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000039.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000040.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000041.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000042.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000043.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000044.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000045.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000046.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000047.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000048.csv delete mode 100644 ScaFFold/fractals/var0.15/3DIFS_param/000049.csv diff --git a/README.md b/README.md index 2442494..53d449c 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,9 @@ The model is trained from a random initialization until convergence, which is de 1. If running the benchmark for the first time, or running with different fractal parameters (`n_categories`, `variance_threshold`) than previously, generate fractal classes and instances: `scaffold generate_fractals -c ScaFFold/configs/benchmark_default.yml` - Note that the benchmark ships with an initial set of 50 fractal classes. + Fractal category libraries are generated deterministically from the + configured seed (under `fract_base_dir/var<...>/seed<...>/`) and reused by + later runs with the same seed. 1. Once fractal generation completes, run the benchmark: `torchrun-hpc -N 1 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark -c ScaFFold/configs/benchmark_default.yml` diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000000.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000000.csv deleted file mode 100644 index 824252c..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000000.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.113880658597081297e-01,1.525788435216699490e-01,-1.677513246653659085e-01,9.450327943163072675e-01,1.437558218378187647e-01,-7.781853097899118499e-01,-7.520371219608632529e-01,4.049719188683833515e-01,3.599506321979530910e-01,-9.209462080360120151e-01,9.797064915552313735e-01,2.046581090479335785e-01,5.952758981070249700e-01 --3.273838635851094025e-01,5.796182968623042608e-01,2.204038860996404559e-02,7.254736836236364006e-02,-3.624484794849551772e-01,3.976647887245436941e-01,-6.272009791693362590e-01,2.300328757717320372e-01,1.439593364137665699e-01,-1.758502593004573900e-01,7.577588977727744979e-01,-4.555583614631009137e-01,4.047241018929750300e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000001.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000001.csv deleted file mode 100644 index fca2ed3..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000001.csv +++ /dev/null @@ -1,2 +0,0 @@ -2.473679574324576524e-01,1.597856635526939684e-02,-1.261210131422938474e-01,2.860947005456315750e-01,8.355771386484407426e-01,-1.930246242945548030e-01,-4.801705885775757743e-01,1.904275728954856195e-01,7.806079367730189844e-01,7.953170851152218113e-01,-5.816986628081313171e-01,6.209542113667643193e-01,1.988928060480311955e-01 --7.041010339743891677e-02,-2.454306473738818717e-01,9.405874516413117448e-01,7.268593945019246050e-01,6.070031387640555387e-01,-9.102525915665182765e-02,6.339356801972018118e-01,-1.532413735258475462e-01,3.654495384649214529e-02,-4.532769270980656628e-01,1.964159205402833397e-01,-5.359870555644301593e-01,8.011071939519688323e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000002.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000002.csv deleted file mode 100644 index c3e5385..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000002.csv +++ /dev/null @@ -1,2 +0,0 @@ --9.413772084350502389e-02,5.254718465280536766e-01,-4.808322979877577286e-01,5.811542374432154823e-01,-4.454009637949891687e-01,-4.249149866229446904e-01,-3.000182472995254201e-01,4.022670562303001240e-01,-2.934854435005793682e-01,-6.605442110513501941e-01,-6.253032129726847632e-01,-5.719885961970654353e-01,4.224941885002395647e-01 --4.407723055180619021e-01,1.325088026302621014e-01,1.840254341100278079e-01,-5.508321197309193895e-01,8.464500652955120330e-01,5.041753392202226181e-01,5.651656783277005935e-01,-8.325090493066080732e-01,-1.381507388162375172e-01,5.571341416171868843e-01,-1.025781698153782617e-01,-8.838988798440101657e-01,5.775058114997604353e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000003.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000003.csv deleted file mode 100644 index 18274fd..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000003.csv +++ /dev/null @@ -1,2 +0,0 @@ --1.294466006455852192e-01,1.974941747383867074e-01,4.427934740152306148e-01,2.336055684661912935e-01,1.300312996679606758e-01,8.725504616000674396e-01,-1.902059536103457571e-01,-3.532649574684778582e-01,-8.430890597933748953e-01,-9.596612485558557726e-01,6.321097711736041180e-01,8.135413458612708038e-01,2.489526692860659640e-01 --1.399185307891788188e-01,-1.602045901301751840e-01,4.335897375838684287e-01,-9.384258169419115170e-01,-1.069737283930594085e-01,2.989241164297768982e-01,-6.355754104429354179e-01,2.746690226736152596e-01,2.732139814068013095e-01,9.765182437472603727e-01,-4.228379364440870702e-02,-5.170594086561124403e-01,7.510473307139340360e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000004.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000004.csv deleted file mode 100644 index ac107cb..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000004.csv +++ /dev/null @@ -1,2 +0,0 @@ --7.074987909298180444e-01,-1.415561527827760013e-01,6.453558418406579733e-02,9.717892764870739164e-01,4.692315433286387005e-02,-9.087912295403599572e-02,-7.306123549338172651e-01,5.199024492355210914e-01,3.242137268086431323e-01,-7.685133956354921470e-01,-2.890512330200536439e-01,-4.414459117103588515e-01,1.767437117859010087e-01 --5.783609543921015561e-01,-6.261288719190640784e-01,6.083280289156212106e-01,5.708229102420392387e-01,6.839066129793194282e-01,-3.730802596414783956e-01,4.320233943214566441e-01,-2.520497672481589735e-01,2.188189615840767654e-01,-5.142628557025761271e-01,9.167783772484443539e-01,-9.404699392000455127e-01,8.232562882140990190e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000005.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000005.csv deleted file mode 100644 index 898a388..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000005.csv +++ /dev/null @@ -1,2 +0,0 @@ --6.016222179722698904e-01,-8.568626391040801149e-01,4.339904630487174675e-01,8.566190336460910437e-02,5.955551394131912701e-01,-7.233426894750305536e-02,-3.439913328223502820e-01,-1.295598008765910247e-01,5.054214022018066466e-01,5.403179943497680160e-01,9.887300162494017108e-01,5.065235980037172681e-01,2.036756775286962806e-01 -4.255811510519635910e-02,-5.861700389763935259e-01,1.539188938105566784e-01,7.464882854730132689e-01,9.985987594821947866e-01,-3.852978585836623893e-01,5.848150331480403974e-01,-8.218568908737176049e-01,-4.786801609318307449e-01,-7.997867917133967275e-01,-3.226694643578802424e-01,-9.055189731557820032e-01,7.963243224713036916e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000006.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000006.csv deleted file mode 100644 index 345f2a4..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000006.csv +++ /dev/null @@ -1,2 +0,0 @@ -4.449557690511229957e-02,1.859345105226184458e-02,-5.492709805712474580e-01,5.002339770143411357e-01,6.022059536539468017e-01,-2.581455709557001210e-01,5.384026026512653829e-01,3.525173161797230392e-01,4.535819387452528773e-01,5.054013100483643051e-01,5.534237431467781132e-01,6.486536701245786407e-01,5.893263018191835512e-01 --4.935786839220561717e-01,-3.483405932567602559e-01,-4.920557235894853498e-01,-3.225857219554804090e-01,-5.369237708506735540e-01,-4.285587807484951828e-01,-2.187430959790792606e-02,8.181072307028691704e-02,4.255774677072714507e-01,-8.957260803321442921e-01,-6.285123937957712847e-01,-7.274955330956049959e-01,4.106736981808163378e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000007.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000007.csv deleted file mode 100644 index fa04c8c..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000007.csv +++ /dev/null @@ -1,2 +0,0 @@ -4.567482715684858530e-01,-2.162812897008825619e-01,8.688925426344686898e-01,9.403035945928239769e-01,-4.923377969407702892e-01,9.039413044894315519e-01,-7.007586459777630505e-01,9.958966466641376858e-01,-2.148840978447630334e-02,-8.820192395688601916e-01,-6.962417306340771272e-01,6.989108727983952551e-01,6.175732443919085268e-01 --1.311082673391426034e-01,2.651764139972054846e-01,2.416531271291553207e-01,2.658826889156997719e-01,5.321835857677645887e-01,8.850640913149898648e-01,-4.146324561488827776e-01,6.016104083259397051e-01,-6.006142415160782289e-01,6.902838974364782221e-01,5.080744400285828188e-01,-5.543689403825666773e-01,3.824267556080915287e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000008.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000008.csv deleted file mode 100644 index b017040..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000008.csv +++ /dev/null @@ -1,3 +0,0 @@ -1.955211818468614027e-01,-1.860012909088504252e-02,1.679370276531577666e-01,-2.474658865690229081e-01,2.002748188638228122e-01,-2.296462069177727106e-01,4.272668474314074150e-01,5.518250090890584048e-01,-1.814099098142296640e-01,-7.759223213118924267e-01,2.396552032441381375e-01,-9.984320127889498853e-01,3.158600947042149998e-01 -5.269496943889822038e-01,5.090567464178141766e-01,3.638063575297016961e-01,-9.409697282110314198e-01,-8.838575707448308449e-01,-9.681971946347656122e-02,-6.464087288109372498e-01,-4.699168427008522109e-01,4.879093108230523335e-01,7.225552560446799610e-01,5.544701020084910059e-01,6.687273760328038552e-01,6.074177711483321751e-01 -3.702243567552572223e-01,2.665339551565091281e-01,-7.302570963874586152e-01,7.128226896410159164e-01,-1.412872752126193010e-01,-5.914873767070649713e-01,4.858707230431502655e-01,2.870662392522644879e-01,-8.632888349796041805e-01,8.908557359947888443e-02,-1.967254792508705830e-01,-6.141181327493672182e-01,7.672213414745278348e-02 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000009.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000009.csv deleted file mode 100644 index f56b11b..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000009.csv +++ /dev/null @@ -1,2 +0,0 @@ -2.948060069505098468e-01,6.887599776806021534e-01,7.611229734423119453e-02,4.932702124290413437e-01,-4.882699298136421451e-01,-3.811733224774245254e-01,7.278710433957895631e-01,6.852985410340497463e-01,-7.967518041925278904e-01,6.365195144826076845e-01,7.940485085563173673e-01,-7.048888688479171272e-02,6.538648345727335887e-01 -2.849341248173462571e-01,-5.989665908310131126e-01,-8.440934355375118159e-01,-5.676180605895972953e-01,2.243341626806083511e-02,-4.393691291202828086e-01,-6.561728033554823369e-01,-1.447073895277080080e-01,-3.027716266953806024e-01,-1.488449066054653436e-01,-7.989435645608391479e-01,6.936025509345438156e-01,3.461351654272663003e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000010.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000010.csv deleted file mode 100644 index af8b628..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000010.csv +++ /dev/null @@ -1,2 +0,0 @@ --2.352521358159322951e-01,-8.804830225066977434e-01,2.445692685925464627e-01,-3.804268259383307704e-01,7.002478757426469080e-01,4.949375238071500593e-01,6.340824702532006363e-01,-1.661009284036518707e-01,6.403120148905094844e-02,-4.544636202226781663e-03,-5.614632388117013484e-01,-8.225067911333789894e-01,7.543884465970274178e-01 --6.205195826406055826e-01,2.082688392648848197e-01,7.051714514820610624e-01,6.615150945575909436e-01,-5.430135870340266901e-01,-7.590298999219038389e-01,-2.484973171764268685e-01,4.507062852245327100e-01,-3.971212226677103274e-01,3.603919868974059249e-01,-6.455363862654992513e-01,6.976679002200263380e-01,2.456115534029726655e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000011.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000011.csv deleted file mode 100644 index 60713bc..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000011.csv +++ /dev/null @@ -1,2 +0,0 @@ --2.647713297675924338e-01,4.532537039449799909e-01,9.126597053670870707e-02,6.297942569420966752e-01,-7.970366903049754814e-01,-5.805081819440571778e-01,1.847828454797417752e-01,1.632357583131875955e-01,1.112052944356816120e-01,5.581182508420681199e-01,-1.235235401936412014e-01,5.525957004344514978e-01,3.917940489256083736e-01 --3.328722698643555855e-01,-6.568958305252938779e-01,9.548663670929014025e-02,2.135325385282085264e-01,-7.994872824743104456e-01,3.120617004916548254e-01,-8.149604985800276147e-01,3.155236614340959367e-01,-5.804561978693489888e-01,-3.465348059373873912e-01,8.779303525751596116e-01,-4.964138499509700431e-01,6.082059510743916819e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000012.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000012.csv deleted file mode 100644 index d75a5f8..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000012.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.732047938669280640e-01,7.457174836228177561e-01,-1.134959963046562326e-01,-2.524155789085107404e-01,6.491113757559063835e-01,5.923873502246719269e-01,-4.163840681869901417e-01,-9.955960643681671662e-02,-4.399149502577506254e-01,-3.298124569267795181e-01,-4.480024355377016931e-01,4.831814457889009873e-01,7.022590685295611035e-01 -3.640511591326804908e-01,-8.138787463613117446e-01,6.945993218897195121e-01,1.561110571982835538e-01,1.420189997595979747e-01,8.921447905553927527e-01,8.983211135903812483e-02,-7.873694636782158085e-02,-1.924595656350303052e-01,3.444934594172788245e-01,1.245560132691707622e-01,-2.441271125304933509e-01,2.977409314704388410e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000013.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000013.csv deleted file mode 100644 index bf897a7..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000013.csv +++ /dev/null @@ -1,2 +0,0 @@ --9.477179726636597579e-02,1.411777244118233021e-01,-3.210166902811444345e-01,8.538383221528667022e-02,5.145040185660725296e-01,5.449889235641092178e-01,1.961008681633682471e-01,5.250846593606550705e-01,2.072609251252632845e-01,8.556312962216905404e-01,7.079983983577677886e-01,-5.970939704307753892e-01,1.986178895167198533e-01 -2.340421417737004184e-02,-2.331838675905175684e-01,8.789698166911528165e-01,4.337374001870073492e-01,-6.744896384287986102e-01,6.995573996808193140e-01,-9.371564747554361752e-02,7.122203806164681961e-01,-3.226197608587253463e-01,5.228286702908606642e-01,6.219275677253508494e-01,6.982539942183398907e-01,8.013821104832802300e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000014.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000014.csv deleted file mode 100644 index 6fd2a78..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000014.csv +++ /dev/null @@ -1,2 +0,0 @@ --5.079619751044526232e-01,-8.448208339950418200e-01,3.379361696842797524e-01,-7.424757815754379209e-01,-2.914860368679528246e-01,7.863175574507479393e-01,-3.913127807203062858e-01,-3.395110295018073376e-01,1.729246237110684259e-01,1.155790784528218929e-01,-1.778336752766045414e-01,-8.237750280597011532e-01,8.128837685573842009e-01 --5.053085469842428790e-01,3.413777808104254685e-01,5.072078013440961541e-02,-2.297771208541388166e-01,1.179085091368414773e-01,-5.353419781682577927e-02,3.033551256644697602e-01,-8.399081644307380135e-01,-2.580876839431038849e-01,-2.915499965121404191e-01,-8.274711651504977894e-01,-3.831424144957116251e-01,1.871162314426157436e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000015.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000015.csv deleted file mode 100644 index 040ab5e..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000015.csv +++ /dev/null @@ -1,2 +0,0 @@ -1.959234370454023821e-01,7.195879279286905295e-01,-8.847434444070787496e-01,-2.025052870905292846e-01,-6.099125369957605347e-01,4.364854998130400787e-01,2.605566328535515730e-01,6.345444901137731186e-02,2.845477814208994261e-01,3.682037052176181380e-01,-5.482600452804842206e-01,8.228596647488930493e-01,2.413865640087159981e-01 --3.105662416993948405e-02,1.541091773938909615e-01,9.832144790547923119e-01,2.617717287788325908e-01,2.964779787430120717e-01,2.867701916196363499e-01,-1.595092390377108593e-01,5.184743556024293820e-01,7.055853027879592787e-01,-2.253650882070545869e-02,3.223158737418712061e-01,-3.130584359547923246e-01,7.586134359912839464e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000016.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000016.csv deleted file mode 100644 index 1262706..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000016.csv +++ /dev/null @@ -1,2 +0,0 @@ --1.300116634675734240e-01,1.252286096526100678e-01,2.499311708856801761e-01,7.830409313847326302e-01,-3.485355694096761159e-01,-2.040759407213834642e-01,3.636277174597402073e-01,-9.857803444259161108e-01,4.953546011072127442e-01,-9.229964481663586184e-01,8.261855069165144894e-01,7.432769447647864514e-01,5.045106210869472196e-01 --9.419389020480599672e-01,-2.909144464059969515e-01,2.371749371359630487e-01,-9.140320234271737121e-02,2.578548696646267846e-01,-6.247044474540388581e-01,5.120075657145519710e-01,7.003880575966952016e-01,-7.351851386341190508e-01,-3.580478535291309328e-02,-3.983862417971921754e-01,-8.667134200938420019e-01,4.954893789130526693e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000017.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000017.csv deleted file mode 100644 index 0a04245..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000017.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.510560689158552794e-01,-2.754404723500034624e-01,-2.684881793088436108e-01,-1.948162421580679204e-01,1.792154998944521793e-01,-3.924494838945669084e-01,-5.433151157278728327e-01,3.050183437518030338e-01,-5.297884663705332287e-01,-4.402434005450059917e-01,3.033450134513224761e-01,5.080005533136775497e-01,1.765071572885704987e-01 --3.296381408504345245e-02,8.213734474661573692e-01,2.884182165116890850e-01,-5.803866935447996589e-01,3.041849216693381930e-01,5.667920446859808781e-01,2.250925284276168448e-01,8.236273827714153395e-01,-4.114514278581935525e-01,2.460381915048910351e-01,-1.767376131099762659e-01,-3.058606496513562867e-01,8.234928427114294180e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000018.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000018.csv deleted file mode 100644 index bca2acf..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000018.csv +++ /dev/null @@ -1,2 +0,0 @@ --5.517066569583584101e-01,6.457030079067409556e-01,1.312452431616923931e-01,-9.654657530459793691e-01,1.001632917413284307e-01,4.808171457391481329e-01,2.887487385246092497e-01,9.356047048047881898e-01,1.042754743685525565e-01,-5.230136305881964986e-01,7.032043120523423507e-01,8.023204324713959501e-01,7.872166104195181813e-01 --5.753978519462528141e-01,7.086092336367590949e-01,8.862996973827064195e-01,-4.508828284143380216e-01,-1.498953015148070111e-01,-1.479318170320651493e-01,-3.642550894021263641e-01,1.915439725996275211e-01,7.097594434528575746e-02,4.751618882704935487e-01,3.779955724380190674e-02,2.814613891706765347e-01,2.127833895804817355e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000019.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000019.csv deleted file mode 100644 index a21ea50..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000019.csv +++ /dev/null @@ -1,2 +0,0 @@ -1.016398917604210972e-01,5.367241283642387728e-01,-6.818635487683848417e-01,-6.492151914242854094e-01,2.951107976835281033e-01,1.210983103608658240e-01,-1.914729327636166545e-01,2.084900726109191194e-01,-9.913674641194303305e-03,-1.318795603716438336e-01,-4.591725358362441778e-01,3.168885180251357347e-02,2.613304272286978147e-01 --2.518998374757852599e-01,-1.545791580965771850e-01,2.591599328760181287e-01,-9.227238644730679784e-01,-6.523143964998687760e-01,1.325715926486024099e-01,-2.634595670714392490e-01,-6.084513786591814188e-01,5.945406152909971098e-01,-7.413295649974527279e-01,-9.836113191580013737e-01,-6.613550934121019687e-01,7.386695727713020743e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000020.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000020.csv deleted file mode 100644 index a24562e..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000020.csv +++ /dev/null @@ -1,2 +0,0 @@ -6.921866875646716100e-01,5.108296431449499408e-03,-8.297196043425478784e-01,-8.782207198751741384e-01,4.881693383540661735e-01,7.602399262614862874e-01,7.929628517423199519e-01,2.747828187373047015e-01,-8.565292578863379358e-01,-8.028470318278093654e-01,-7.666729973480264082e-01,-3.618054675021058486e-01,6.474809588861711873e-01 --4.195656375439105190e-02,3.829378217234820081e-01,-4.738148839593532280e-01,-1.013350142767268647e-01,6.518052790254404982e-01,-2.989039824402659473e-01,3.810369849047898771e-01,-4.761398937049154956e-01,-8.736647575845664093e-01,8.671975808320215862e-01,4.518295413402169114e-01,3.044248270515332866e-01,3.525190411138288682e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000021.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000021.csv deleted file mode 100644 index d861c0d..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000021.csv +++ /dev/null @@ -1,2 +0,0 @@ -1.667868049973333378e-01,-9.628664509718032338e-01,6.014400922608953426e-01,2.972001483931350219e-01,4.941659040388075574e-01,-6.774253220554804500e-01,-1.993299704297657460e-01,6.394927405908403806e-01,-1.095498513905281968e-01,1.797181927692330650e-01,8.966277273316463070e-01,2.586404039871941229e-01,5.552709887608781036e-01 -8.904906062660500332e-01,3.456353981853357293e-01,-5.089078336971060157e-01,5.564474698994368307e-01,1.613474062509465679e-01,-4.536022710624829646e-01,7.605093590363174449e-01,-3.094273061226622268e-01,-6.937981682429052999e-01,8.031386423573720901e-01,-7.976443746091934628e-01,6.529445161208398130e-01,4.447290112391218964e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000022.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000022.csv deleted file mode 100644 index d9cf9f6..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000022.csv +++ /dev/null @@ -1,2 +0,0 @@ --5.460898717556510462e-01,5.858062690874137335e-01,2.194226187473129475e-01,9.188630804699871035e-02,-1.707137636423472493e-01,-3.178585699238225537e-01,3.633095509488837305e-01,-6.887805288505008949e-01,-1.468860672171268256e-02,-2.862987915279111562e-01,-7.723918312056159419e-01,8.126650586523440634e-01,4.458798622363638331e-01 -5.547486983713716402e-01,-5.675784088131794469e-01,-6.187415038981081139e-02,6.835868811760836827e-02,7.617313312413553916e-03,5.521100702462113929e-01,5.867352194847055280e-01,-8.129437845772602422e-01,-1.254568018619570680e-01,8.382074387221385425e-01,9.974232510580696154e-01,1.902631261807281593e-01,5.541201377636362224e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000023.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000023.csv deleted file mode 100644 index c1d0e5a..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000023.csv +++ /dev/null @@ -1,3 +0,0 @@ -2.812878692303713013e-01,-5.380780165501937162e-01,-1.995746226435606285e-01,3.556609488625348536e-01,-3.570058976622436653e-02,3.226211706468127272e-01,7.604367324861782684e-02,-6.726022941928089249e-01,-1.244510164772598682e-01,-5.242674835142189238e-01,6.111367647898253708e-01,5.446881171705224567e-01,1.870729455165833777e-01 --1.804572792799459258e-01,-5.615234318506530098e-01,-5.678667666618428811e-01,-1.973882851907129421e-01,-9.342817705045243226e-01,-2.696256933877165807e-01,-4.754476132734983818e-01,8.379911380637130591e-01,-1.430754625396699620e-01,6.703727340713743210e-01,-7.950803881856465249e-01,1.646446229560936114e-01,5.812129149706145581e-01 --7.303265877923637017e-01,5.163190035467037919e-01,6.970165582377609859e-01,-2.277609099967679018e-01,5.773156478708756367e-01,7.154948998404351279e-01,-4.402130399012762485e-01,-5.815862404831599886e-01,-9.434039458871834594e-01,-3.147493553756923745e-01,-6.215697199671859074e-01,1.823798551890962738e-01,2.317141395128020642e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000024.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000024.csv deleted file mode 100644 index 3eb8c33..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000024.csv +++ /dev/null @@ -1,2 +0,0 @@ -2.166769833429884606e-01,-2.331594788641218052e-01,-8.006870701326018747e-01,-5.266093172746755258e-01,-8.322914134913979023e-01,7.446673507975531958e-01,-2.886575388105372397e-02,-2.990935453031380309e-01,3.936163556018033027e-01,-3.542536932837532238e-01,2.834018169147607402e-01,-9.586592464730825380e-01,4.259584301232757220e-01 -8.416839786286203218e-01,-5.562313049662526154e-02,-3.647897665777326548e-01,-6.475376571556650251e-01,-3.655643797640151238e-01,-3.196437920194410420e-01,5.442135586745493470e-01,5.508198826067678411e-03,-8.528986112361351957e-01,-6.533886080407793617e-01,-8.277524479526696677e-01,-6.784479736983972664e-01,5.740415698767242780e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000025.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000025.csv deleted file mode 100644 index 4f7835f..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000025.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.328441505555097546e-01,-1.387237543865100786e-01,-8.385785443603737122e-01,-6.460892635786552596e-01,-6.206316056012477489e-01,-2.482763796893627806e-02,8.854727682161622759e-02,-3.627820436506015156e-01,-2.149069615655596621e-02,-8.638772932614158240e-01,-9.940866772762460002e-01,7.740573663063279319e-01,4.901549186583971096e-01 --5.533853946343669783e-02,6.783593497875153311e-01,-3.470947229869234540e-01,-8.081860602696091522e-01,-2.234790784515090500e-01,5.246894583705352666e-02,4.184989961414606885e-01,5.085439272585972059e-01,2.229705198326268345e-01,-2.452705816636990832e-02,-3.598764438670341015e-01,-4.548214145779025941e-01,5.098450813416028904e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000026.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000026.csv deleted file mode 100644 index cf18c74..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000026.csv +++ /dev/null @@ -1,2 +0,0 @@ -9.215997820884049840e-01,-8.886305956631088687e-01,7.918854213071213621e-02,8.549195867901333568e-01,-7.539338040902159310e-01,2.127998942252871117e-01,9.334583063583679063e-02,2.117336123070212572e-01,-3.905267778142720303e-01,4.922987444042437044e-01,3.538552884251142672e-01,-4.995810348035092385e-01,3.171922691986440168e-01 --2.710991105815863111e-01,6.963869144978018788e-02,-2.614443785391928898e-01,-1.745302413226517135e-01,2.582270480554498260e-01,-7.700432645988855018e-01,-9.312241172350603780e-01,5.326596688321914019e-01,6.785270353740011640e-01,-6.480593296020313865e-01,-5.862245859094883382e-01,-8.798502128158780522e-02,6.828077308013559277e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000027.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000027.csv deleted file mode 100644 index 5645631..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000027.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.363859180192507736e-01,-2.005356922417877996e-02,4.105299683991425752e-03,4.520536430696797670e-01,-2.397748696508128496e-01,-9.149591317297329773e-01,-2.687516636255817826e-01,2.502854147579025579e-02,-5.658580150798908637e-01,-3.444933603162623204e-01,3.980110762078155062e-01,6.019507953845388837e-01,2.306736239861526538e-01 --9.924270668037067367e-01,1.880015741363572079e-01,7.482014409903954277e-01,-8.717841019029326510e-01,3.178806298249037265e-01,4.369798231563537527e-02,-4.363922340879544670e-01,-3.301960725800534568e-01,5.876506713161555595e-01,-6.176395243356269660e-02,-3.698941006853062596e-01,-6.461432233004753556e-01,7.693263760138474572e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000028.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000028.csv deleted file mode 100644 index 3901f6d..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000028.csv +++ /dev/null @@ -1,2 +0,0 @@ --9.488858097444059503e-01,4.079256280033416449e-01,-7.273191618538961123e-01,-2.164528040135893505e-02,1.801964698567355416e-01,2.110643357786530228e-01,-6.481561009612546442e-02,-2.415201336393544285e-01,8.209213192744986287e-01,1.342520688379920113e-01,6.654275517850556376e-01,4.674644534794492046e-02,2.727174515041750347e-01 -5.889383911906385105e-01,-4.716661362499952048e-01,8.398891628230311657e-01,9.444296451726641450e-01,6.852416895301041144e-01,-2.718291944059443299e-01,9.759285577328178363e-01,4.221736051403957024e-01,-5.860661824356676597e-01,6.789765242522334265e-01,-6.864599914320155261e-01,8.076484878207652596e-01,7.272825484958250764e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000029.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000029.csv deleted file mode 100644 index 7f0e048..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000029.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.156004044367652916e-01,4.539226489835086475e-01,3.386522834986691599e-01,-6.897275360788317489e-01,-6.389720151715840846e-01,-3.773349166157233814e-01,1.813606744935740700e-01,-6.205184082106425247e-01,-3.558681169774229325e-01,3.970835926210039002e-01,-8.915932500008973971e-01,-7.113662644290583703e-01,2.064773053549207871e-01 --2.222304390465441593e-01,2.761428589612251461e-01,-4.556405012454298742e-01,1.006421622526911808e-01,1.163138470867763896e-01,4.306829086925703098e-01,7.112480073546407766e-01,8.255028592978328472e-01,-9.373835503173011396e-02,-7.714024572181317208e-01,9.021782591265152806e-01,-1.140412270850570398e-01,7.935226946450792962e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000030.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000030.csv deleted file mode 100644 index 76e9192..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000030.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.054766017168615289e-01,3.285035685711270581e-01,3.339186103540698891e-01,8.138242613802961767e-01,-5.205905000696835483e-01,1.529228105404809579e-01,8.466677815415111219e-01,-7.061057292640156025e-01,-1.216890214920127722e-01,-1.716434030659630405e-01,-3.234074144154277519e-01,-5.208430114476063633e-01,5.154903134260301334e-01 -8.269501238092979989e-02,9.228403821803805585e-01,-9.202930276956799993e-01,-1.265069840482921926e-01,2.493516975241008016e-01,-3.998067562089782090e-01,-4.105033927816157391e-01,2.800337569818500683e-02,-5.976660197456356016e-01,4.214440523544955575e-01,4.705808164701692498e-01,4.943614828825371177e-01,4.845096865739698111e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000031.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000031.csv deleted file mode 100644 index 297ec32..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000031.csv +++ /dev/null @@ -1,2 +0,0 @@ -5.931682635522148583e-01,8.699251727039116755e-01,-8.264376444320893356e-01,6.777316299673374900e-03,1.259812425166662031e-02,7.526947568173611991e-01,-6.456449467182574509e-01,-1.867494885666358684e-01,-5.751692444575797758e-01,1.762422487332968579e-01,-4.008003102356285652e-01,-8.546965786953328870e-01,6.116071270069189936e-01 -9.061805329925218810e-01,5.725102286958729803e-01,6.505301386644584127e-01,-5.634820490409915283e-01,1.322069399068586115e-01,-5.320988916994693341e-01,-8.018914495696449762e-02,-1.672922951958961679e-01,-5.236940947370711807e-01,2.913933672990698387e-01,-7.670309932568115663e-01,2.266688501822122781e-01,3.883928729930810619e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000032.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000032.csv deleted file mode 100644 index 7b3c85b..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000032.csv +++ /dev/null @@ -1,2 +0,0 @@ --4.029727807509135840e-01,2.221234395356765923e-01,1.653311729326147894e-01,2.854384277332262521e-01,-5.299668975252573855e-01,2.277161429752445621e-01,-8.605748587143935424e-01,2.679358058329555092e-01,-1.753822281114920667e-01,-4.080365861113419701e-01,-7.990197638852070128e-02,6.451913142929026623e-01,2.646000967988433872e-01 --5.825951636626720553e-01,9.933789450712708913e-01,6.017632717440635215e-01,1.285660707793123692e-01,9.225675787278662110e-01,7.629667460674474100e-01,-1.269325586916481008e-01,-2.399724811514538647e-01,2.244383569074661633e-01,3.864334357836947120e-01,1.407009021060205978e-01,-4.470113212371948919e-01,7.353999032011566683e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000033.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000033.csv deleted file mode 100644 index a6af9d6..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000033.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.343344504186707145e-01,-1.788377734180790490e-02,5.238121179021124618e-01,-1.801466014877619592e-01,2.444788960837480651e-01,5.831808836713883171e-01,-8.022784620348435425e-01,-4.295701503377884478e-01,-3.972178837057684930e-01,7.118887696693811939e-01,-8.639924687935303105e-01,4.474523280618651899e-01,4.561726115084100419e-01 --3.159104803788210791e-01,-8.391006409409682565e-01,-3.390316252365024319e-01,1.553039413054588813e-01,8.170194759350031255e-01,7.527728133817603862e-01,6.978584447967217663e-01,1.128938492491866619e-01,-8.250572721858402403e-01,-5.512768790771354066e-01,3.630337106170709038e-01,-2.222138950067416019e-01,5.438273884915899581e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000034.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000034.csv deleted file mode 100644 index d18f1b5..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000034.csv +++ /dev/null @@ -1,2 +0,0 @@ -5.697263627900581717e-01,9.724449544795059630e-01,-1.674757567980476036e-01,-1.372771123535918569e-01,-7.764922940367462445e-01,5.410911337820014655e-01,3.028425898014133200e-01,4.025386950980982537e-01,-1.476033977383572893e-01,3.294617484232353899e-01,-5.024261635732010234e-01,-4.464667963982544840e-01,6.101539033152794111e-01 --7.758966548598598134e-01,6.989084956917288594e-01,7.492372266276172699e-01,1.570862193960709252e-01,-2.601786362671802966e-01,2.727571671873869619e-01,-1.349546389088733811e-01,1.024561812582318598e-01,5.506075183810597018e-01,2.805526623632508265e-01,-9.364615058079310828e-01,5.565384469523662059e-01,3.898460966847205333e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000035.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000035.csv deleted file mode 100644 index f61ede0..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000035.csv +++ /dev/null @@ -1,2 +0,0 @@ --7.898235741040993130e-01,5.427817216628392227e-01,6.300274118070616769e-01,-5.211653741616515401e-01,-6.981622668479592342e-02,9.578781805862026655e-01,1.666377264587282081e-01,1.637757194655065085e-01,7.571393872230314237e-02,9.020579858164654574e-01,8.210889142502095783e-01,9.377974534116928496e-01,8.615315349321857052e-01 -2.679807526537674178e-01,-1.111544236701951238e-01,-5.024699421507621278e-01,7.457314846733458236e-01,-4.885664083857288453e-02,-7.673089298447215434e-01,-8.883324858657668521e-01,-2.464007149940794505e-01,6.129545902906914367e-01,2.592414698046310306e-01,-3.794121113471968787e-01,1.341081142088484945e-02,1.384684650678143225e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000036.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000036.csv deleted file mode 100644 index bcb6f88..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000036.csv +++ /dev/null @@ -1,2 +0,0 @@ --7.654273192322054165e-01,-4.094450069738897469e-02,-7.595383743431962653e-01,-8.557505909970739566e-01,9.064524921829164583e-01,1.517470713028805651e-01,9.075466163718142187e-01,-4.142547321400837923e-01,4.595383450523824465e-01,2.242844526762681756e-01,6.565226571017350743e-01,-6.406292457846078925e-01,2.164925796749257170e-01 -4.134224368111192316e-01,-1.279193807354295220e-01,1.395348624965011552e-01,-7.113928155897839556e-01,-2.037598375484022117e-01,-1.397999201487876153e-01,-8.100399318887028244e-01,-3.540527934775026253e-01,-8.142623919357250273e-01,-7.886095433794413356e-01,-3.674068192061692439e-01,-8.141558656567047247e-01,7.835074203250743663e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000037.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000037.csv deleted file mode 100644 index d666294..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000037.csv +++ /dev/null @@ -1,2 +0,0 @@ -6.234462065020718313e-01,1.300245200362042386e-01,-2.842597130938973038e-01,1.595438141453187075e-01,-4.713993691785938189e-01,1.486436548511567146e-01,-6.366613328949060069e-01,7.675794196907781419e-02,-9.937538509582277690e-02,5.037179007533050257e-01,-6.402726075739686440e-01,-3.221377519083947760e-01,4.313968476057338797e-01 --2.673181139834390763e-03,-2.286479913327021940e-02,6.059211548471905573e-01,-3.648636787800783043e-01,-4.414015900975392093e-01,2.960060241178898988e-01,-3.576760461529826518e-01,-9.644156958072134245e-01,-5.959032460738362680e-01,-7.214674024314973177e-01,5.674627493986834637e-01,3.261734992328386706e-01,5.686031523942661758e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000038.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000038.csv deleted file mode 100644 index df4a712..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000038.csv +++ /dev/null @@ -1,2 +0,0 @@ --7.420103884864390764e-01,-6.381524612883886505e-01,-1.613639618245801266e-01,7.199764450024330742e-01,-1.006210277087187244e-01,5.736088536664056825e-01,-2.078207707668398019e-01,3.313608111485661922e-01,-1.567612276534600113e-01,3.527694725218057936e-01,-1.767897433287017872e-01,-3.553205709558269199e-01,5.623580600962118092e-01 -1.803325284453796140e-01,1.231843219533304001e-01,5.739878271111997776e-03,-8.697875645645707365e-01,4.277980606243028117e-01,1.010399067634688564e-01,9.441666464685340987e-01,-2.209014237470598996e-01,3.359090300130838092e-01,5.940980517447476128e-01,2.482757629902709873e-01,-1.834052354989281763e-01,4.376419399037881908e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000039.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000039.csv deleted file mode 100644 index a71f835..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000039.csv +++ /dev/null @@ -1,2 +0,0 @@ -7.369773596082351830e-02,-2.482646733587001719e-01,5.057638234463792681e-01,-6.242771609481161388e-01,-7.429818940347492351e-01,-4.707267465509237248e-01,-4.626361468138358024e-01,-6.467146033806137062e-02,2.131345959536605772e-01,7.725903761561427885e-01,-4.240286243640691843e-01,-4.218771823875742122e-01,5.823322236392094453e-01 --6.554490572494988676e-01,-7.815506609888773770e-01,-7.332898854928275867e-01,6.612085576900739170e-02,2.178726193278017753e-01,5.569210460230202830e-01,-2.319012715747397202e-01,-9.022663389078762197e-01,-4.304867213238519064e-01,-5.467178901348748177e-01,5.657685803999605856e-01,7.731699809585836913e-01,4.176677763607904992e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000040.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000040.csv deleted file mode 100644 index 2012dad..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000040.csv +++ /dev/null @@ -1,2 +0,0 @@ -3.730533296626739048e-01,1.433183679154126366e-01,3.637273449949744997e-01,2.796358796157163429e-01,9.971687232829262726e-03,7.480311521939866370e-01,3.937355475477484212e-01,4.560000728934507919e-01,-4.984654853379433259e-01,-3.455550319734554954e-01,-9.835584381128927856e-01,-2.633451200187155727e-01,1.651205815987711323e-01 -8.112653261883815414e-01,-3.950454378848380355e-01,3.176817652040930806e-01,-5.239337818191858176e-03,4.250105481704846699e-01,3.604098546798490954e-01,5.490989594952853103e-01,-7.353484188252294995e-01,1.375370159687769878e-01,3.739887895489748537e-01,7.194665959008399447e-01,-2.052330554944428176e-01,8.348794184012288122e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000041.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000041.csv deleted file mode 100644 index 9697871..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000041.csv +++ /dev/null @@ -1,2 +0,0 @@ -1.985151568679008882e-01,-4.118089649714407052e-02,-1.173690980919479543e-01,-6.147768674282323431e-01,5.384205158174693029e-01,-8.092057756629023046e-02,2.779284151732872576e-01,2.862361042030712177e-01,3.840843723060081150e-01,9.898121724234101304e-01,9.713830508604506253e-01,7.822167263660317893e-01,7.111855449058056555e-01 -5.892885167675354641e-01,-1.963827275432317165e-03,-7.606511330012744043e-01,5.586488854728632880e-01,7.846991525938085132e-02,-6.184691865438018965e-01,2.366672318317677437e-01,-4.646883212217289838e-02,2.801530524023061464e-01,-6.722047100717012391e-01,-7.126605435777322306e-01,-6.314786757874539802e-01,2.888144550941944000e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000042.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000042.csv deleted file mode 100644 index 9760c67..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000042.csv +++ /dev/null @@ -1,2 +0,0 @@ --5.511430473636560468e-01,2.075233910672833471e-01,9.276713288167097726e-01,-1.822552515865676348e-01,4.964535495465318693e-01,4.198110192190658285e-01,-7.110023362885946607e-01,3.517890700491155265e-02,1.069361824193393318e-02,-7.884513208221513025e-01,-8.560972160329762826e-01,-9.047373606248456657e-01,8.266997873488395321e-01 --5.807999610356651132e-02,-1.039348906028856323e-01,3.141343076567044701e-01,-2.056624804791262751e-01,5.792891943323019710e-02,-9.509068945320611199e-01,-3.764789001983093186e-01,2.488096127203558439e-01,-1.853125052175255139e-01,8.301118240605835918e-01,-1.216569088549470656e-01,8.391547337468705514e-01,1.733002126511604402e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000043.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000043.csv deleted file mode 100644 index a0119e8..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000043.csv +++ /dev/null @@ -1,2 +0,0 @@ --5.633361199129609531e-01,-1.242868737984672567e-01,-8.018904487385585256e-01,-5.010751601456076010e-02,-5.054966895260570858e-02,4.270084989655040797e-01,5.613658733821358382e-02,1.067847504111867352e-01,9.715247219631861775e-02,-5.884270629564716248e-01,2.869301460815778526e-01,-8.916215668672899941e-01,1.127259639768470878e-01 -5.813484676405529239e-01,-7.608577159641982668e-02,-1.726841049851548515e-01,1.495963314716997061e-01,9.282798648258450136e-02,-9.208945702831410340e-01,-3.570382956573512345e-01,-3.065376051478874153e-01,-3.776733142035426649e-01,-9.376120443254172265e-01,-1.646485576867517953e-01,4.424834047857026942e-01,8.872740360231529122e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000044.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000044.csv deleted file mode 100644 index 088d30f..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000044.csv +++ /dev/null @@ -1,2 +0,0 @@ -1.245129250727008419e-01,-6.636582935874648648e-01,-7.095021996374384354e-01,-5.708048087285482186e-01,-2.613952915421429157e-01,4.360332099195141087e-01,4.386695896700476549e-01,2.591668816930605690e-01,2.710449407828230406e-01,-1.994249360309625629e-01,-9.884196752400828956e-01,-8.948557108504355817e-01,4.008892167512856930e-01 --7.977101440753309181e-01,-5.891213181154428824e-01,-5.394841756954262824e-01,6.845658522320718919e-01,5.996325324640972010e-01,-7.950515319880713250e-02,1.314403032344408917e-01,8.974954145234985692e-01,3.185168035203989056e-02,-8.571054947316427697e-01,8.942330986453337349e-01,4.800334580863008238e-01,5.991107832487143625e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000045.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000045.csv deleted file mode 100644 index 379affb..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000045.csv +++ /dev/null @@ -1,2 +0,0 @@ -4.156697549963126459e-01,1.938812896637998051e-02,3.762020023825551895e-02,-4.543409658210944002e-01,-4.160769028245498991e-01,-4.404768450621054932e-02,-1.226606756646233531e-01,2.246644739370420307e-01,-1.930408370461424994e-01,4.195049700021282746e-01,-7.252898605016289135e-01,-1.186045724865631978e-01,2.376483080486445632e-01 --5.048426707707873717e-01,2.678319927728600724e-01,-2.276997135948761741e-01,5.838793463819900165e-01,-5.592597689481533241e-01,5.818759639934625305e-01,-5.955380441213122822e-01,-6.677174966534085154e-01,2.195016551675315064e-01,-5.308958776483587716e-01,8.413556421996504220e-02,6.724915832844733377e-01,7.623516919513554368e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000046.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000046.csv deleted file mode 100644 index 674f6e6..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000046.csv +++ /dev/null @@ -1,2 +0,0 @@ --1.237880297177609279e-01,-5.389232504861669604e-01,8.351615604828057648e-01,-5.197299868564686509e-01,-6.477376740352802642e-02,6.235481432700507032e-01,8.648963276371262054e-01,-4.513602462196104614e-01,1.121472463841981515e-01,9.872965050461273151e-01,9.679965448742415823e-01,-1.147346194826741606e-01,3.912812329938487044e-01 --5.999286838715911507e-02,-2.097115124782031881e-01,3.409859382911564207e-01,-1.638644314851909201e-01,-5.392371544653988824e-01,1.453514884790074735e-01,-9.158051602213537201e-01,7.348874336021522513e-01,5.552381537384321053e-01,-2.804504387202546578e-01,7.872454148265923823e-01,2.773799831190322251e-01,6.087187670061512401e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000047.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000047.csv deleted file mode 100644 index 4580b5a..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000047.csv +++ /dev/null @@ -1,2 +0,0 @@ --4.781135907766911330e-01,-1.276623027908179164e-01,-5.244677607605363612e-02,7.793736542701723558e-02,-4.965349935613236898e-01,-7.450674808538539917e-01,2.215427033187817862e-01,-2.835553192890256646e-01,-5.879019088656525227e-01,-8.655409501489235158e-01,1.484038914439216317e-01,2.340787049439099210e-01,4.441034483063370786e-01 -8.918907925254648816e-02,9.583522210758066429e-02,3.380023039545712038e-01,-2.757160640620373027e-01,9.491265801938739699e-01,-2.949240243853041843e-01,2.368000271881480767e-01,-3.274535743006037336e-01,2.320714265769052709e-01,6.681697133141641931e-01,-8.151860272732713852e-02,-5.223786728288988268e-01,5.558965516936629214e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000048.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000048.csv deleted file mode 100644 index 4b22744..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000048.csv +++ /dev/null @@ -1,2 +0,0 @@ --2.430657731263985433e-01,-5.883176849573654721e-01,-1.051727327384854860e-02,-5.613318623862584289e-01,-3.654238252241008844e-01,-2.903699040052056812e-01,-3.376162129993560690e-01,-1.937379177653244522e-01,1.447762165448007732e-01,9.132031396935695877e-01,2.309821948964190241e-01,4.969710849235846606e-01,3.556727315444812021e-01 --5.700025402755233284e-01,3.962163276669257161e-01,4.106860990482503748e-01,-4.445226161030975121e-01,2.216226823069946672e-01,5.337581703209348660e-01,-5.759837441605704100e-01,-3.921802841679722373e-01,-5.148726436855111110e-01,-4.398568392874997457e-01,-8.504659419589648550e-01,6.585902602061193267e-01,6.443272684555187979e-01 diff --git a/ScaFFold/fractals/var0.15/3DIFS_param/000049.csv b/ScaFFold/fractals/var0.15/3DIFS_param/000049.csv deleted file mode 100644 index e66cf87..0000000 --- a/ScaFFold/fractals/var0.15/3DIFS_param/000049.csv +++ /dev/null @@ -1,2 +0,0 @@ --3.200427851155407399e-01,3.576359420223551577e-01,-2.513016740058249265e-01,-7.546934815443087086e-01,-2.786933152965342941e-01,2.327058657559735178e-01,7.113022836300710861e-01,6.680266543764468157e-01,2.330623308024013518e-01,-6.132416434718550580e-01,-1.253524191738015769e-01,4.190854187019641408e-01,6.428379861042887722e-01 --2.051650609723685292e-01,-4.086715839711074771e-02,1.337484780156206199e-02,-3.217966923911506072e-01,2.645280631801716353e-01,-5.515626691943973370e-01,-7.300788769231552067e-01,-9.944065418179377502e-01,-6.942904879158731113e-01,3.461188780593014158e-01,-2.042033665456883806e-01,-6.116735230775620646e-01,3.571620138957111168e-01 diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 60731fe..3d7ebe3 100644 --- a/ScaFFold/utils/config_utils.py +++ b/ScaFFold/utils/config_utils.py @@ -18,8 +18,6 @@ import yaml -import ScaFFold.paths - def require_positive_int(name: str, value: int) -> int: if not isinstance(value, int) or isinstance(value, bool) or value < 1: @@ -89,7 +87,6 @@ class Config: "dataset_dir", "fract_base_dir", "job_name", - "library_root", "n_categories", "problem_scale", "unet_bottleneck_dim", @@ -209,7 +206,6 @@ def _validate_keys(cls, config_dict, strict): def __init__(self, config_dict, strict=True): self._validate_keys(config_dict, strict) - self.library_root = str(ScaFFold.paths.scaffold_root).rstrip("/") + "/ScaFFold/" self.base_run_dir = str(Path(config_dict["base_run_dir"]).resolve()) self.dataset_dir = str( Path(config_dict.get("dataset_dir", "datasets/")).resolve() diff --git a/pyproject.toml b/pyproject.toml index c0e6c3e..80b5aa7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,6 @@ build-backend = "setuptools.build_meta" [tool.setuptools] package-data = { "ScaFFold" = [ "package_data/weights_ins145.csv", - "fractals/var0.15/3DIFS_param/*", "configs/*", ] } include-package-data = true From a6a22a1807c162d2363ed0057d6c38a305e49c4d Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Wed, 5 Aug 2026 10:48:55 -0700 Subject: [PATCH 54/54] Say what a comment means instead of citing a review issue Ninety-one comments and docstrings pointed at issue numbers from the two review rounds -- 47 F-codes from the first, which were never cleaned up, and 44 R-codes from the second. A reader outside those reviews cannot resolve any of them, and the reviews are not part of the repository. Most were section banners that already carried their description beside the code, so they only lose the prefix. The dozen that had the code embedded in prose are reworded to stand on their own: the checkpoint cleanup now names the hazard it guards (peers stranded in an unmatched collective) rather than the issue that found it. Also generalise the parsed-CSV cache comment in the instance generator, which quoted the literal 145 instances per category. --- ScaFFold/datagen/instance.py | 5 +++-- ScaFFold/utils/checkpointing.py | 7 +++---- tests/conftest.py | 3 ++- tests/datagen/test_artifacts.py | 18 ++++++++--------- tests/datagen/test_category_search.py | 4 ++-- tests/datagen/test_library_layout.py | 2 +- tests/datagen/test_mpi_consensus.py | 8 ++++---- tests/datagen/test_provenance.py | 2 +- tests/datagen/test_rng_determinism.py | 12 +++++------ tests/test_checkpointing.py | 29 ++++++++++++++------------- tests/test_cli.py | 14 ++++++------- tests/test_config.py | 6 +++--- tests/test_data_loading.py | 18 ++++++++--------- tests/test_evaluate.py | 6 +++--- tests/test_infra.py | 2 +- tests/test_perf_hotpath.py | 20 +++++++++--------- tests/test_reporting.py | 22 ++++++++++---------- tests/test_restart_script.py | 4 ++-- tests/test_resume.py | 4 ++-- tests/test_worker_dist.py | 4 ++-- 20 files changed, 96 insertions(+), 94 deletions(-) diff --git a/ScaFFold/datagen/instance.py b/ScaFFold/datagen/instance.py index ea84a91..12d9c02 100644 --- a/ScaFFold/datagen/instance.py +++ b/ScaFFold/datagen/instance.py @@ -349,8 +349,9 @@ def main(config: Config): # One-entry cache of the most recently parsed category CSV. The work list # is built category-major and block-sliced, so every rank's items for a # given category are contiguous: a single entry is enough to turn the - # per-item re-parse (145 identical reads of the same small file off the - # shared filesystem, per category) into one parse per category per rank. + # per-item re-parse -- one redundant read of the same small file off the + # shared filesystem for every instance generated in the category -- into + # one parse per category per rank. # ``generate_instance_points`` copies before scaling, so sharing the parsed # array across instances cannot leak weights from one item into the next. cached_category = None diff --git a/ScaFFold/utils/checkpointing.py b/ScaFFold/utils/checkpointing.py index 00dfc14..c8d8b9e 100644 --- a/ScaFFold/utils/checkpointing.py +++ b/ScaFFold/utils/checkpointing.py @@ -156,9 +156,8 @@ def cleanup(self, train_from_scratch: bool) -> None: # already committed to the broadcast below. Individual unlinks # are tolerated inside, but the glob/stat around them can still # raise on a shared filesystem (ESTALE, EACCES), and raising in - # this window strands the peers in an unmatched collective -- - # the R05 hazard. Report the failure through the broadcast, like - # a failed write. + # this window strands the peers in an unmatched collective. + # Report the failure through the broadcast, like a failed write. try: self._remove_checkpoint_files() except Exception as e: @@ -209,7 +208,7 @@ def _sweep_orphaned_tmp_files(self) -> None: chance to clear them. Sweeping at construction is safe because run directories are per-run - and not shared between concurrently running jobs (F55), so any temp + and not shared between concurrently running jobs, so any temp file here belongs to a dead process -- except one from this pid, which another manager in this process could still be writing. diff --git a/tests/conftest.py b/tests/conftest.py index e35f739..b390e37 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -326,7 +326,8 @@ def tiny_v1_dataset(tmp_path): Same shape as ``tiny_dataset`` but: no ``meta.yaml`` (so the loader falls back to the legacy path), channels-last volumes ``(N, N, N, 3)``, and masks holding raw values that get remapped through the per-split ``mask_values`` - pickle. Needed by the F42 dataset-loading tests. + pickle. Needed by the dataset-loading tests that cover legacy label + remapping. """ def make( diff --git a/tests/datagen/test_artifacts.py b/tests/datagen/test_artifacts.py index f07bd3a..77d16dc 100644 --- a/tests/datagen/test_artifacts.py +++ b/tests/datagen/test_artifacts.py @@ -92,7 +92,7 @@ def _seed_category(fract_base: Path, *, point_num: int, keep: range) -> Path: # --------------------------------------------------------------------------- -# F09: non-finite weighted instances are retried / fall back, never saved +# Non-finite weighted instances are retried / fall back, never saved # --------------------------------------------------------------------------- @@ -142,7 +142,7 @@ def _patched_genfromtxt(fname, *args, **kwargs): # --------------------------------------------------------------------------- -# F09 defense: points_to_voxelgrid rejects non-finite input +# Defense in depth: points_to_voxelgrid rejects non-finite input # --------------------------------------------------------------------------- @@ -157,7 +157,7 @@ def test_voxelgrid_rejects_nonfinite(bad): # --------------------------------------------------------------------------- -# F32: instance save is atomic and resume handles temp / truncated files +# Instance save is atomic and resume handles temp / truncated files # --------------------------------------------------------------------------- @@ -235,7 +235,7 @@ def test_resume_rejects_truncated(tmp_path): # --------------------------------------------------------------------------- -# R36: a category's IFS parameters are parsed once, not once per instance +# A category's IFS parameters are parsed once, not once per instance # --------------------------------------------------------------------------- @@ -292,7 +292,7 @@ def counting_genfromtxt(fname, *args, **kwargs): # --------------------------------------------------------------------------- -# F62: mask scanner requires exactly one file per id +# Mask scanner requires exactly one file per id # --------------------------------------------------------------------------- @@ -322,7 +322,7 @@ def test_mask_stem_index_rejects_ambiguous(tmp_path): # --------------------------------------------------------------------------- -# F64: isotropic + centered normalization; the scale knob is rejected +# Isotropic + centered normalization; the scale knob is rejected # --------------------------------------------------------------------------- @@ -365,7 +365,7 @@ def test_scale_config_rejected(): # --------------------------------------------------------------------------- -# F34: rasterization scatters point indices instead of traversing a dense grid +# Rasterization scatters point indices instead of traversing a dense grid # --------------------------------------------------------------------------- @@ -442,7 +442,7 @@ def test_scatter_paint_matches_boolean_mask(): # --------------------------------------------------------------------------- -# F66: instance point clouds are stored and loaded as float32 +# Instance point clouds are stored and loaded as float32 # --------------------------------------------------------------------------- @@ -483,7 +483,7 @@ def test_dataset_version_bumped_past_float64_era(): # --------------------------------------------------------------------------- -# R33: voxel centering is a whole voxel, not half of one +# Voxel centering is a whole voxel, not half of one # --------------------------------------------------------------------------- diff --git a/tests/datagen/test_category_search.py b/tests/datagen/test_category_search.py index 3ff19a2..353edbd 100644 --- a/tests/datagen/test_category_search.py +++ b/tests/datagen/test_category_search.py @@ -64,7 +64,7 @@ def test_at_least_one_attempt_per_rank_when_work_remains(): # --------------------------------------------------------------------------- -# R30: the initial work scan is made once on rank 0 and broadcast. +# The initial work scan is made once on rank 0 and broadcast. # # ``categories_remaining`` gates a while loop that contains collectives, so it # must be identical on every rank. Deriving it from a per-rank filesystem scan @@ -313,7 +313,7 @@ def test_peer_raises_on_broadcast_scan_error(tmp_path, monkeypatch): # --------------------------------------------------------------------------- -# R31: category CSVs appear complete or not at all. +# Category CSVs appear complete or not at all. # # A category file truncated by a killed job is still counted as "done" by the # resume scan, so nothing ever regenerates it: instance generation then dies diff --git a/tests/datagen/test_library_layout.py b/tests/datagen/test_library_layout.py index d09a524..a5346da 100644 --- a/tests/datagen/test_library_layout.py +++ b/tests/datagen/test_library_layout.py @@ -12,7 +12,7 @@ # # SPDX-License-Identifier: (Apache-2.0) -"""The fractal library is keyed by the seed that produced it (R29). +"""The fractal library is keyed by the seed that produced it. Categories and instances are *derived from* ``config.seed``: the IFS parameters come from a seed-keyed candidate stream and every instance point cloud is diff --git a/tests/datagen/test_mpi_consensus.py b/tests/datagen/test_mpi_consensus.py index 5185eb3..72b462c 100644 --- a/tests/datagen/test_mpi_consensus.py +++ b/tests/datagen/test_mpi_consensus.py @@ -344,7 +344,7 @@ def test_generation_success_finalizes_and_returns(tmp_path, monkeypatch): # --------------------------------------------------------------------------- -# R27: rank 0 must never die between the collectives its peers have entered. +# Rank 0 must never die between the collectives its peers have entered. # Every rank-0-only step of the consensus (the reuse/generate decision and the # final meta-write + rename) is wrapped so a failure travels to the peers as a # broadcast sentinel instead of stranding them in ``bcast``/``Barrier``. @@ -592,7 +592,7 @@ def interrupted(_config): # --------------------------------------------------------------------------- -# R37: orphaned staging dirs are reclaimed instead of accumulating forever. +# Orphaned staging dirs are reclaimed instead of accumulating forever. # --------------------------------------------------------------------------- @@ -716,7 +716,7 @@ def test_live_generation_survives_on_its_heartbeat_alone(tmp_path): def test_dead_staging_dir_with_stale_heartbeat_is_reclaimed(tmp_path): - """The heartbeat must not turn cleanup into a no-op (R37 still holds).""" + """The heartbeat must not stop orphaned staging dirs being reclaimed.""" base = tmp_path / "cid" base.mkdir(parents=True) @@ -772,7 +772,7 @@ def test_heartbeat_respects_its_interval(tmp_path): # --------------------------------------------------------------------------- -# R28: meta.yaml is published atomically, so no reader ever sees a partial one. +# meta.yaml is published atomically, so no reader ever sees a partial one. # --------------------------------------------------------------------------- diff --git a/tests/datagen/test_provenance.py b/tests/datagen/test_provenance.py index d907c39..90f74b1 100644 --- a/tests/datagen/test_provenance.py +++ b/tests/datagen/test_provenance.py @@ -12,7 +12,7 @@ # # SPDX-License-Identifier: (Apache-2.0) -"""Dataset provenance: the commit stamped on a dataset is ScaFFold's (R35). +"""Dataset provenance: the commit stamped on a dataset is ScaFFold's. ``meta.yaml``'s ``code_commit``, the published ``__`` directory name, and the ``dataset_reuse_enforce_commit_id`` gate all key off diff --git a/tests/datagen/test_rng_determinism.py b/tests/datagen/test_rng_determinism.py index 9c1bec3..2feafbe 100644 --- a/tests/datagen/test_rng_determinism.py +++ b/tests/datagen/test_rng_determinism.py @@ -34,7 +34,7 @@ from ScaFFold.datagen import category_search as cs # --------------------------------------------------------------------------- -# F06: seed_numba controls the njit RNG stream +# seed_numba controls the njit RNG stream # --------------------------------------------------------------------------- @@ -64,7 +64,7 @@ def draw_with_seed(seed): # --------------------------------------------------------------------------- -# F06: whole-instance generation is deterministic across fresh processes +# Whole-instance generation is deterministic across fresh processes # --------------------------------------------------------------------------- @@ -106,7 +106,7 @@ def test_instance_generation_deterministic(fresh_python): # --------------------------------------------------------------------------- -# F06: per-item seed derivation is independent of rank/world-size layout +# Per-item seed derivation is independent of rank/world-size layout # --------------------------------------------------------------------------- @@ -165,7 +165,7 @@ def _run_search_loop(write_dir, base_seed, rank, attempt_start, n_wanted, accept # --------------------------------------------------------------------------- -# F07: resume continues the candidate stream instead of replaying it +# Resume continues the candidate stream instead of replaying it # --------------------------------------------------------------------------- @@ -211,7 +211,7 @@ def accept_all(_params): # --------------------------------------------------------------------------- -# F07: duplicate candidates are rejected by the dedup guard +# Duplicate candidates are rejected by the dedup guard # --------------------------------------------------------------------------- @@ -240,7 +240,7 @@ def test_duplicate_candidate_skipped(tmp_path): # --------------------------------------------------------------------------- -# F63: index allocation fills gaps and never overwrites an existing file +# Index allocation fills gaps and never overwrites an existing file # --------------------------------------------------------------------------- diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 5d7afa7..659d432 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -94,7 +94,7 @@ def elapsed_time(self, other): # --------------------------------------------------------------------------- -# F10 -- atomic writes + fallback + no swallowing +# Atomic writes + fallback + no swallowing # --------------------------------------------------------------------------- @@ -191,7 +191,7 @@ def always_raise(obj, f, *args, **kwargs): # --------------------------------------------------------------------------- -# R04 -- async save failures are reported, and the run's LAST save is consumed +# Async save failures are reported, and the run's LAST save is consumed # --------------------------------------------------------------------------- @@ -272,7 +272,7 @@ def test_final_async_save_failure_fails_the_run(tiny_trainer, monkeypatch): # --------------------------------------------------------------------------- -# R05 -- a rank-0 write failure fails every rank with the same error +# A rank-0 write failure fails every rank with the same error # --------------------------------------------------------------------------- # Two-rank script: rank 0's torch.save fails. Both ranks must come out of @@ -370,7 +370,7 @@ def test_rank0_save_failure_fails_all_ranks(tmp_path): # --------------------------------------------------------------------------- -# F41 -- race-free best decision (cached best loss, no per-save probe) +# Race-free best decision (cached best loss, no per-save probe) # --------------------------------------------------------------------------- @@ -428,7 +428,7 @@ def spy_load(path, *args, **kwargs): # --------------------------------------------------------------------------- -# R02 -- a from-scratch cleanup drops the deleted run's best, not just its files +# A from-scratch cleanup drops the deleted run's best, not just its files # --------------------------------------------------------------------------- @@ -466,7 +466,7 @@ def test_cleanup_from_scratch_resets_best(tmp_path): # --------------------------------------------------------------------------- -# R07 -- checkpoint debris (.tmp., .corrupt) does not accumulate +# Checkpoint debris (.tmp., .corrupt) does not accumulate # --------------------------------------------------------------------------- @@ -519,7 +519,7 @@ def test_init_sweeps_orphaned_tmp_files(tmp_path): # --------------------------------------------------------------------------- -# R08 -- an improving epoch serializes the state dict once, not twice +# An improving epoch serializes the state dict once, not twice # --------------------------------------------------------------------------- @@ -625,7 +625,8 @@ def test_cleanup_rank0_fs_error_travels_through_the_broadcast(tmp_path, monkeypa ``_remove_checkpoint_files`` globs and stats the run directory; on a shared filesystem those can raise (ESTALE, EACCES) even though each individual unlink is already tolerated. That happens between the drain and the - broadcast, so an unfenced raise re-creates exactly the hazard R05 closed. + broadcast, so an unfenced raise would strand the peers in an unmatched + collective -- the hazard that reporting through the broadcast closes. """ mgr, _ = _make_manager(tmp_path) mgr.dist_enabled = True @@ -686,7 +687,7 @@ def test_init_survives_an_unlistable_run_dir(tmp_path, monkeypatch, capsys): # --------------------------------------------------------------------------- -# F71 -- CPU tensors are cloned into the snapshot +# CPU tensors are cloned into the snapshot # --------------------------------------------------------------------------- @@ -706,7 +707,7 @@ def test_cpu_tensors_cloned(tmp_path): # --------------------------------------------------------------------------- -# R09 -- the warmup snapshot is host-resident, not a device-side copy +# The warmup snapshot is host-resident, not a device-side copy # --------------------------------------------------------------------------- @@ -815,7 +816,7 @@ def test_snapshot_restore_round_trips_values_and_devices(tmp_path): # --------------------------------------------------------------------------- -# F50 -- a final checkpoint is written when the run exits between intervals +# A final checkpoint is written when the run exits between intervals # --------------------------------------------------------------------------- @@ -858,7 +859,7 @@ def fake_evaluate(*args, **kwargs): # --------------------------------------------------------------------------- -# R03 -- a diverged epoch aborts instead of checkpointing NaN weights +# A diverged epoch aborts instead of checkpointing NaN weights # --------------------------------------------------------------------------- @@ -906,7 +907,7 @@ def diverged_evaluate(*args, **kwargs): # --------------------------------------------------------------------------- -# F49 -- GradScaler-skipped steps do not advance the optimizer-step counter +# GradScaler-skipped steps do not advance the optimizer-step counter # --------------------------------------------------------------------------- @@ -956,7 +957,7 @@ def test_skipped_step_not_counted(): # --------------------------------------------------------------------------- -# F72 -- on resume only rank 0 reads the checkpoint file; peers get the +# On resume only rank 0 reads the checkpoint file; peers get the # deserialized state over the process group (no N-way filesystem read storm) # --------------------------------------------------------------------------- diff --git a/tests/test_cli.py b/tests/test_cli.py index f7d83ca..18952d7 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -135,7 +135,7 @@ def run_cli(monkeypatch, argv, *, comm=None, sync_env=True): # --------------------------------------------------------------------------- -# R13: the MPI world must span the whole job +# The MPI world must span the whole job # --------------------------------------------------------------------------- @@ -220,7 +220,7 @@ def test_no_launcher_env_is_not_a_mismatch(monkeypatch, tmp_path): # --------------------------------------------------------------------------- -# R17: the restart script is generated at the true job scale +# The restart script is generated at the true job scale # --------------------------------------------------------------------------- @@ -253,7 +253,7 @@ def _recorder(run_dir, world_size=None): # --------------------------------------------------------------------------- -# R18: the restart pre-check is a rank-0 decision, broadcast to everyone +# The restart pre-check is a rank-0 decision, broadcast to everyone # --------------------------------------------------------------------------- @@ -414,7 +414,7 @@ def test_fresh_run_records_no_restart_state(monkeypatch, tmp_path): # --------------------------------------------------------------------------- -# R19: generate_fractals is not a benchmark run +# generate_fractals is not a benchmark run # --------------------------------------------------------------------------- @@ -465,7 +465,7 @@ def test_benchmark_still_creates_its_run_dir(monkeypatch, tmp_path): # --------------------------------------------------------------------------- -# R20: auxiliary keys set in YAML must survive; CLI > YAML > argparse default +# Auxiliary keys set in YAML must survive; CLI > YAML > argparse default # --------------------------------------------------------------------------- @@ -533,7 +533,7 @@ def test_run_config_records_the_effective_aux_values(monkeypatch, tmp_path): # --------------------------------------------------------------------------- -# R25: an out-of-range bottleneck is rejected before any work starts +# An out-of-range bottleneck is rejected before any work starts # --------------------------------------------------------------------------- @@ -626,7 +626,7 @@ def test_unknown_error_types_degrade_to_runtime_error(): # --------------------------------------------------------------------------- -# The whole config path survives a restart (R20/R22 together) +# The whole config path survives a restart # --------------------------------------------------------------------------- diff --git a/tests/test_config.py b/tests/test_config.py index 086a4f5..89375ee 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -99,7 +99,7 @@ def test_invalid_type_message_names_type(tmp_path): def test_activation_checkpointing_is_a_real_option(): - """R44: activation checkpointing is reachable from a config, defaulting off. + """Activation checkpointing is reachable from a config, defaulting off. The U-Net has always had ``use_checkpointing``, but with no config key and no caller it could not be turned on: any attempt was rejected as an unknown @@ -308,7 +308,7 @@ def test_list_valued_key_rejected_in_run_config(tmp_path): # --------------------------------------------------------------------------- -# The base config is preserved under a name of its own (R22) +# The base config is preserved under a name of its own # --------------------------------------------------------------------------- @@ -365,7 +365,7 @@ def test_base_config_copy_never_clobbers_merged_config( # --------------------------------------------------------------------------- -# unet_bottleneck_dim range (R25) +# unet_bottleneck_dim range # --------------------------------------------------------------------------- diff --git a/tests/test_data_loading.py b/tests/test_data_loading.py index 1f60055..8a6b1fb 100644 --- a/tests/test_data_loading.py +++ b/tests/test_data_loading.py @@ -112,7 +112,7 @@ def _build_v1_split_dataset( # --------------------------------------------------------------------------- -# F11: index -> file mapping must not depend on os.listdir order +# Index -> file mapping must not depend on os.listdir order # --------------------------------------------------------------------------- @@ -256,7 +256,7 @@ def fake_cuda(self, *args, **kwargs): # --------------------------------------------------------------------------- -# F42: legacy label remapping must use one global (union) table +# Legacy label remapping must use one global (union) table # --------------------------------------------------------------------------- @@ -394,7 +394,7 @@ def test_v2_datasets_unaffected(tiny_dataset): # --------------------------------------------------------------------------- -# F43: id -> path resolved once at init; no per-item directory scans +# id -> path resolved once at init; no per-item directory scans # --------------------------------------------------------------------------- @@ -446,7 +446,7 @@ def test_duplicate_stem_raises(tmp_path): # --------------------------------------------------------------------------- -# F13: narrow (int16) mask carrier, widened to long on the compute device +# Narrow (int16) mask carrier, widened to long on the compute device # --------------------------------------------------------------------------- @@ -475,7 +475,7 @@ def test_mask_carrier_is_narrow_int16(tiny_dataset): # --------------------------------------------------------------------------- -# F21: non-sharded volume prep is zero-copy (no redundant full-volume copy) +# Non-sharded volume prep is zero-copy (no redundant full-volume copy) # --------------------------------------------------------------------------- @@ -525,7 +525,7 @@ def test_nonsharded_getitem_tensors_bit_identical(tiny_dataset): # --------------------------------------------------------------------------- -# F24: mask-only accessor loads the mask without touching the image volume +# Mask-only accessor loads the mask without touching the image volume # --------------------------------------------------------------------------- @@ -597,7 +597,7 @@ def counting_load(path, mmap_mode=None): # --------------------------------------------------------------------------- -# R28: a *present but broken* meta.yaml must not be mistaken for a v1 dataset +# A *present but broken* meta.yaml must not be mistaken for a v1 dataset # --------------------------------------------------------------------------- @@ -637,7 +637,7 @@ def test_broken_meta_raises_instead_of_silent_legacy(tmp_path, broken_meta): # --------------------------------------------------------------------------- -# R32: uneven spatial shards are accepted but computed wrong (known, unfixed) +# Uneven spatial shards are accepted but computed wrong (known, unfixed) # --------------------------------------------------------------------------- @@ -724,7 +724,7 @@ def _build_v2_sparse_label_dataset(root: Path, label: int) -> Path: # --------------------------------------------------------------------------- -# R34: the int16 carrier guard must bound the largest class *id*, not the count +# The int16 carrier guard must bound the largest class *id*, not the count # --------------------------------------------------------------------------- diff --git a/tests/test_evaluate.py b/tests/test_evaluate.py index b55df77..8f2500e 100644 --- a/tests/test_evaluate.py +++ b/tests/test_evaluate.py @@ -165,7 +165,7 @@ def _run(evaluate, ps, net, batches, n_categories): # --------------------------------------------------------------------------- -# F05: reported Dice must be the hard (argmax) segmentation Dice +# Reported Dice must be the hard (argmax) segmentation Dice # --------------------------------------------------------------------------- @@ -241,7 +241,7 @@ def test_imperfect_model_dice_below_one(eval_env): # --------------------------------------------------------------------------- -# F23: the degenerate single-class path must be rejected, not faked +# The degenerate single-class path must be rejected, not faked # --------------------------------------------------------------------------- @@ -255,7 +255,7 @@ def test_n_categories_zero_rejected(eval_env): # --------------------------------------------------------------------------- -# F53: val_loss_avg must be sample-weighted, not equal-weighted per batch +# val_loss_avg must be sample-weighted, not equal-weighted per batch # --------------------------------------------------------------------------- diff --git a/tests/test_infra.py b/tests/test_infra.py index aaf3a73..a3e8f93 100644 --- a/tests/test_infra.py +++ b/tests/test_infra.py @@ -236,7 +236,7 @@ def test_torchrun_gloo_two_ranks(tmp_path): # --------------------------------------------------------------------------- -# R21: memory diagnostics on a CPU-only run +# Memory diagnostics on a CPU-only run # --------------------------------------------------------------------------- diff --git a/tests/test_perf_hotpath.py b/tests/test_perf_hotpath.py index a683c59..eac6a5f 100644 --- a/tests/test_perf_hotpath.py +++ b/tests/test_perf_hotpath.py @@ -33,7 +33,7 @@ def _reference_onehot(labels, num_classes): def test_labels_to_onehot_matches_one_hot_and_is_float32_contiguous(): - # F12: the scatter-based one-hot must be bit-identical to the + # The scatter-based one-hot must be bit-identical to the # F.one_hot().permute().float() chain, but float32 (not an int64 # intermediate) and contiguous in channel-first layout. torch.manual_seed(0) @@ -49,7 +49,7 @@ def test_labels_to_onehot_matches_one_hot_and_is_float32_contiguous(): def test_ce_log_probs_path_matches_cross_entropy(): - # F20: feeding a precomputed log_softmax via NLL must equal computing CE + # Feeding a precomputed log_softmax via NLL must equal computing CE # from raw logits, for both weighted and unweighted cases, with no mesh. torch.manual_seed(1) b, c = 2, 5 @@ -68,7 +68,7 @@ def test_ce_log_probs_path_matches_cross_entropy(): def test_ce_uses_single_spatial_collective(monkeypatch): - # F19: the CE numerator and its normalizer are reduced together in one + # The CE numerator and its normalizer are reduced together in one # SpatialAllReduce, not two. Count applications; the packed path issues # exactly one per CE call (down from two). import ScaFFold.utils.losses as losses @@ -95,7 +95,7 @@ def counting_apply(tensor, mesh): def test_torch_profiler_context_is_bounded(monkeypatch): - # F47: the profiler must be built with a bounded schedule (not RECORD for + # The profiler must be built with a bounded schedule (not RECORD for # every step) and with the expensive record_shapes/with_stack options off # by default, so a long run cannot grow an unbounded trace. import ScaFFold.utils.perf_measure as pm @@ -129,7 +129,7 @@ def _module_source(module): def test_training_batch_returns_detached_dice(tiny_trainer): - # F51: the dice score returned from a batch must be detached so the epoch + # The dice score returned from a batch must be detached so the epoch # accumulator does not retain each batch's autograd graph. Run one real # CPU batch (ps=None path) and inspect the returned tensor. trainer = tiny_trainer() @@ -142,7 +142,7 @@ def test_training_batch_returns_detached_dice(tiny_trainer): def test_zero_grad_uses_set_to_none(tiny_trainer): - # F52: after a batch, grads should be released (None) rather than zeroed + # After a batch, grads should be released (None) rather than zeroed # buffers, avoiding a per-step memset over all gradient memory. trainer = tiny_trainer() batch = next(iter(trainer.train_loader)) @@ -153,7 +153,7 @@ def test_zero_grad_uses_set_to_none(tiny_trainer): def test_evaluate_defers_item_sync_to_end(): - # F22: the validation loop must not call .item() per batch. Verify the + # The validation loop must not call .item() per batch. Verify the # foreground stats helper returns a tensor sum (no host sync) and that the # source has no per-batch .item() inside the loop. import ScaFFold.utils.evaluate as ev @@ -169,7 +169,7 @@ def test_evaluate_defers_item_sync_to_end(): def test_evaluate_copies_are_non_blocking(): - # F54: both eval H2D copies must pass non_blocking=True to exploit the + # Both eval H2D copies must pass non_blocking=True to exploit the # pinned-memory loaders. import ScaFFold.utils.evaluate as ev @@ -178,7 +178,7 @@ def test_evaluate_copies_are_non_blocking(): def test_training_batch_gathers_mem_only_first_batch(): - # F18: the epoch loop must not request mem-stat gathering for every batch + # The epoch loop must not request mem-stat gathering for every batch # (which resets peak counters and issues collectives inside the timed # region). It should gate on the first batch of the run. import ScaFFold.utils.trainer as tr @@ -191,7 +191,7 @@ def test_training_batch_gathers_mem_only_first_batch(): # --------------------------------------------------------------------------- -# R41: warmup covers the ragged final batch +# Warmup covers the ragged final batch # # Warmup only ever runs the *leading* batches of the train loader, which are # all local_batch_size wide, and neither loader drops its last batch. When a diff --git a/tests/test_reporting.py b/tests/test_reporting.py index af50a04..9740bf9 100644 --- a/tests/test_reporting.py +++ b/tests/test_reporting.py @@ -30,7 +30,7 @@ class TestFiguresDir: - """F45: standard_viz.main() creates figures_path with idempotence.""" + """standard_viz.main() creates figures_path with idempotence.""" def test_figures_dir_idempotent(self, tmp_path): """Generate figures twice into same run_dir; second call should succeed.""" @@ -58,7 +58,7 @@ def test_figures_dir_idempotent(self, tmp_path): class TestFigureLifetime: - """R43: standard_viz closes every figure it opens. + """standard_viz closes every figure it opens. ``worker.main`` calls ``standard_viz.main`` in-process once per sweep combination, and pyplot keeps a strong reference to every unclosed figure, @@ -102,7 +102,7 @@ def boom(*args, **kwargs): class TestDiceFigure: - """F70: Validation Dice figure saved as val_dice.png, not val_loss.png.""" + """Validation Dice figure saved as val_dice.png, not val_loss.png.""" def test_dice_figure_filename(self, tmp_path): """Save Dice figure as val_dice.png; val_loss.png (if present) contains loss series.""" @@ -179,7 +179,7 @@ def recording_savefig(fname, *a, **kw): class TestMaskPanelLabels: - """F69: plot_img_and_mask labels panels with correct class index, not off-by-one.""" + """plot_img_and_mask labels panels with correct class index, not off-by-one.""" def test_mask_panel_labels(self): """Render a 3-class mask; check that each panel title matches the class shown.""" @@ -214,7 +214,7 @@ def test_mask_panel_labels(self): class TestVisualizerVolume: - """F65: data_visualizer renders 4D channels-first volumes from volumegen.""" + """data_visualizer renders 4D channels-first volumes from volumegen.""" def test_visualizer_accepts_4d_volume(self, tmp_path): """4D float volume (3, N, N, N) renders without exception.""" @@ -268,7 +268,7 @@ def test_visualizer_still_accepts_3d_mask(self, tmp_path): class TestTorchProfiler: - """F68: Torch profiler enabled independently of Caliper even when CALI_CONFIG set.""" + """Torch profiler enabled independently of Caliper even when CALI_CONFIG set.""" def test_torch_profiler_independent_of_caliper(self, monkeypatch): """Both profilers come up together when both are requested. @@ -319,7 +319,7 @@ def test_torch_profiler_independent_of_caliper(self, monkeypatch): class TestProfilerTraceExport: - """R15: a failed trace export must not strand the other ranks.""" + """A failed trace export must not strand the other ranks.""" @staticmethod def _unstepped_profiler(): @@ -404,7 +404,7 @@ def test_successful_export_writes_a_trace(self, tmp_path, caplog): assert Path(path).exists() def test_trace_lands_in_the_run_dir(self, tmp_path, caplog): - """R23: the trace goes to the run dir, not whatever CWD happens to be.""" + """The trace goes to the run dir, not whatever CWD happens to be.""" import ScaFFold.worker as worker prof = self._stepped_profiler() @@ -425,7 +425,7 @@ def test_trace_lands_in_the_run_dir(self, tmp_path, caplog): def test_trace_name_counts_nodes_not_ranks( self, tmp_path, world_size, ranks_per_node, expected ): - """R23: the N field is a node count, and never rounds a node away.""" + """The N field is a node count, and never rounds a node away.""" import ScaFFold.worker as worker prof = self._stepped_profiler() @@ -444,7 +444,7 @@ def test_trace_name_counts_nodes_not_ranks( class TestProfileTorchGate: - """R24: PROFILE_TORCH is parsed like every other profiler flag.""" + """PROFILE_TORCH is parsed like every other profiler flag.""" @staticmethod def _reload_with(monkeypatch_context, value): @@ -509,7 +509,7 @@ def test_gate_matches_the_sub_option_parser(self, monkeypatch): class TestProfilerSchedule: - """R26: the schedule must not record everything before the first step.""" + """The schedule must not record everything before the first step.""" @staticmethod def _context_with(monkeypatch_context, env): diff --git a/tests/test_restart_script.py b/tests/test_restart_script.py index 291c57a..328cfdb 100644 --- a/tests/test_restart_script.py +++ b/tests/test_restart_script.py @@ -188,7 +188,7 @@ def test_generated_script_is_valid_bash(monkeypatch, tmp_path): # --------------------------------------------------------------------------- -# R14: combined ``--flag=value`` tokens +# Combined ``--flag=value`` tokens # --------------------------------------------------------------------------- @@ -263,7 +263,7 @@ def test_run_dir_placeholder_is_substituted(monkeypatch, tmp_path): # --------------------------------------------------------------------------- -# R17: launch-shape sniffing must match the rank side +# Launch-shape sniffing must match the rank side # --------------------------------------------------------------------------- # Every variable ``ScaFFold.utils.distributed.get_world_size`` honors. The diff --git a/tests/test_resume.py b/tests/test_resume.py index dfe7b2c..a1238b8 100644 --- a/tests/test_resume.py +++ b/tests/test_resume.py @@ -38,7 +38,7 @@ import ScaFFold.cli as cli # --------------------------------------------------------------------------- -# resolve_run_dir matrix (F02) +# resolve_run_dir matrix # --------------------------------------------------------------------------- @@ -128,7 +128,7 @@ def cfg(): # --------------------------------------------------------------------------- -# stats-file header handling on resume (F01) + step counters (F14) +# stats-file header handling on resume + step counters # --------------------------------------------------------------------------- from types import SimpleNamespace # noqa: E402 diff --git a/tests/test_worker_dist.py b/tests/test_worker_dist.py index b5c6b32..a418462 100644 --- a/tests/test_worker_dist.py +++ b/tests/test_worker_dist.py @@ -243,7 +243,7 @@ def test_worker_singleton_smoke(monkeypatch, tiny_config, tiny_dataset): # --------------------------------------------------------------------------- -# R44: the activation-checkpointing config flag reaches the model +# The activation-checkpointing config flag reaches the model # --------------------------------------------------------------------------- @@ -271,7 +271,7 @@ def test_activation_checkpointing_flag_reaches_the_model( # --------------------------------------------------------------------------- -# Local size detection (R23) +# Local size detection # --------------------------------------------------------------------------- _LOCAL_SIZE_CASES = [