Resume benchmark runs from checkpoints by default

.github/workflows/ci.yml

@@ -59,6 +59,28 @@ jobs:
       - run: pip install numpy pandas pyarrow
       - run: pytest tests/unit/test_recorder_checkpoint.py
 
+  benchmark-resume:
+    name: benchmark interruption recovery (CPU)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+      # Use the project's versions, with CPU wheels for the inference test doubles.
+      - run: |
+          python - <<'EOF'
+          import subprocess, sys, tomllib
+          with open("pyproject.toml", "rb") as f:
+              requirements = tomllib.load(f)["project"]["optional-dependencies"]["gpu"]
+          cpu = [r for r in requirements if r.split("==")[0] in {"torch", "torchvision", "torchaudio"}]
+          subprocess.check_call([sys.executable, "-m", "pip", "install", *cpu,
+                                 "--index-url", "https://download.pytorch.org/whl/cpu"])
+          EOF
+      - run: pip install '.[gpu]' pytest
+      - run: pytest tests/unit/test_benchmark_resume.py tests/unit/test_cached_sample_labels.py
+
   typecheck:
     name: mypy (advisory)
     runs-on: ubuntu-latest

README.md

@@ -77,6 +77,24 @@ Model directory must contain: `model_config.yaml`, `model.py`, `*.safetensors`
 
 Results are automatically saved to a timestamped JSON file.
 
+Every benchmark run checkpoints through `BenchmarkRunRecorder` after each inference
+batch. Checkpoints default to `<cache-dir>/runs/<run-id>/checkpoint`; use
+`--checkpoint-dir` to choose another location. To resume, rerun the same command
+with the same `--run-id` and storage directory. Completed predictions are restored
+before decoding, and scores and parquet output are rebuilt through the normal
+recorder path. A new run ID starts a new evaluation.
+
+The checkpoint freezes sample selection (including the robustness pass), source
+content hashes, model/evaluator identity, seed (42 when omitted), and scoring
+settings. Changed inputs fail closed. Preparing a new run reads the selected files
+to fingerprint them; resumed runs validate only pending samples. Model setup and
+uncommitted work may repeat after interruption.
+
+Checkpoint storage must survive the process or container being replaced. The
+Python API accepts `checkpoint_persist(directory)` for filesystems requiring an
+explicit remote commit. One coordinator must own each run directory; automatic
+container replacement and distributed ownership are the caller's responsibility.
+
 ---
 
 ## Python API

src/gasbench/benchmark.py

@@ -41,9 +41,16 @@ async def run_benchmark(
     aug_weight: float = 0.2,
     aug_cache_dir: Optional[str] = None,
     aug_cache_readonly: bool = False,
+    checkpoint_dir: Optional[str] = None,
+    checkpoint_persist=None,
 ) -> Dict:
     """
     Args:
+        checkpoint_dir: Storage directory; defaults to <cache_dir>/runs/<run_id>/checkpoint.
+            Every run checkpoints. Reuse run_id and the same inputs to resume.
+        checkpoint_persist: Optional filesystem commit callback, called after each
+            durable write. Required when a mounted filesystem needs an explicit
+            remote commit; failures abort the run.
         model_path: Path to a custom PyTorch model directory
         modality: Type of modality to test ("image" or "video")
         mode: Benchmark mode - "debug", "small", or "full" (default: "full")
@@ -127,6 +134,8 @@ async def run_benchmark(
             content_category,
             score_composition,
             multiclass_scoring,
+            checkpoint_dir=checkpoint_dir,
+            checkpoint_persist=checkpoint_persist,
             n_aug_per_dataset=n_aug_per_dataset,
             aug_weight=aug_weight,
             aug_cache_dir=aug_cache_dir,
@@ -232,6 +241,8 @@ async def execute_benchmark(
     aug_weight: float = 0.2,
     aug_cache_dir: Optional[str] = None,
     aug_cache_readonly: bool = False,
+    checkpoint_dir: Optional[str] = None,
+    checkpoint_persist=None,
 ) -> float:
     """Execute the actual benchmark evaluation."""
 
@@ -260,6 +271,8 @@ async def execute_benchmark(
             content_category=content_category,
             score_composition=score_composition,
             multiclass_scoring=multiclass_scoring,
+            checkpoint_dir=checkpoint_dir,
+            checkpoint_persist=checkpoint_persist,
             n_aug_per_dataset=n_aug_per_dataset,
             aug_weight=aug_weight,
             aug_cache_dir=aug_cache_dir,
@@ -288,6 +301,8 @@ async def execute_benchmark(
             content_category=content_category,
             score_composition=score_composition,
             multiclass_scoring=multiclass_scoring,
+            checkpoint_dir=checkpoint_dir,
+            checkpoint_persist=checkpoint_persist,
             n_aug_per_dataset=n_aug_per_dataset,
             aug_weight=aug_weight,
             aug_cache_dir=aug_cache_dir,
@@ -316,6 +331,8 @@ async def execute_benchmark(
             content_category=content_category,
             score_composition=score_composition,
             multiclass_scoring=multiclass_scoring,
+            checkpoint_dir=checkpoint_dir,
+            checkpoint_persist=checkpoint_persist,
         )
         benchmark_score = benchmark_results.get("audio_results", {}).get("benchmark_score", 0.0)
     else:

src/gasbench/benchmarks/_checkpoint.py

@@ -65,6 +65,21 @@ class RecorderCheckpoint:
     is rejected.
     """
 
+    @staticmethod
+    def read_manifest(directory: Path):
+        """Read a saved plan before opening the recorder, which validates its batches."""
+        path = Path(directory) / "manifest.json"
+        if not path.exists():
+            if any(Path(directory).glob("batch-*.json")) or (Path(directory) / "head.json").exists():
+                raise CheckpointError("Checkpoint progress exists without a manifest")
+            return None
+        value = _read(path)
+        if not isinstance(value, dict) or value.get("schema_version") != 2 or not isinstance(value.get("manifest"), dict):
+            raise CheckpointError("Invalid checkpoint manifest")
+        if value.get("sha256") != _digest(value["manifest"]):
+            raise CheckpointError("Invalid checkpoint manifest checksum")
+        return value["manifest"]
+
     def __init__(
         self,
         directory: Path,
@@ -81,9 +96,10 @@ def __init__(
             raise ValueError("A nonempty run manifest is required")
         # Normalize JSON types and detach mutable caller-owned objects.
         expected = {
-            "schema_version": 1,
+            "schema_version": 2,
             "manifest": json.loads(_encode(dict(manifest))),
         }
+        expected["sha256"] = _digest(expected["manifest"])
         self.run_id = expected["manifest"].get("run_id")
         if not isinstance(self.run_id, str) or not self.run_id:
             raise ValueError("Checkpoint requires the benchmark run_id")
@@ -240,27 +256,34 @@ def _write_head(self, index: int, checksum: Optional[str]) -> None:
     def _write(self, destination: Path, value, *, persist: bool = True) -> None:
         # Preserve recorder column order in stored rows; checksums still use
         # canonical key ordering so integrity does not depend on JSON layout.
-        content = _encode(value, sort_keys=False)
         temp_path = None
         try:
-            with tempfile.NamedTemporaryFile(
-                dir=self.directory, prefix=".pending-", delete=False
-            ) as handle:
-                temp_path = Path(handle.name)
-                handle.write(content)
-                handle.flush()
-                os.fsync(handle.fileno())
-            os.replace(temp_path, destination)
-            fd = os.open(self.directory, os.O_RDONLY)
+            content = _encode(value, sort_keys=False)
             try:
-                os.fsync(fd)
+                with tempfile.NamedTemporaryFile(
+                    dir=self.directory, prefix=".pending-", delete=False
+                ) as handle:
+                    temp_path = Path(handle.name)
+                    handle.write(content)
+                    handle.flush()
+                    os.fsync(handle.fileno())
+                os.replace(temp_path, destination)
+                fd = os.open(self.directory, os.O_RDONLY)
+                try:
+                    os.fsync(fd)
+                finally:
+                    os.close(fd)
+                if persist and self._persist is not None:
+                    self._persist(self.directory)
             finally:
-                os.close(fd)
-            if persist and self._persist is not None:
-                self._persist(self.directory)
-        except BaseException:
+                if temp_path is not None:
+                    temp_path.unlink(missing_ok=True)
+        except BaseException as exc:
             self._usable = False
+            # Sample/dataset handlers must never swallow a failed durable write.
+            # Preserve process interruptions while normalizing storage errors.
+            if isinstance(exc, Exception) and not isinstance(exc, CheckpointError):
+                raise CheckpointError(
+                    f"Cannot persist checkpoint {destination.name}: {exc}"
+                ) from exc
             raise
-        finally:
-            if temp_path is not None:
-                temp_path.unlink(missing_ok=True)

src/gasbench/benchmarks/audio_bench.py

@@ -4,13 +4,16 @@
 
 from ..logger import get_logger
 from ..processing.media import process_audio_sample
-from ..dataset.iterator import DatasetIterator
+from ..dataset.iterator import load_audio_sample
 
+from ._checkpoint import CheckpointError
 from .recording import BenchmarkRunRecorder, log_dataset_summary
 from .common import (
     BenchmarkRunConfig,
     build_plan,
     create_tracker,
+    create_dataset_iterator,
+    verify_sample,
     finalize_run,
     run_batch_and_record,
 )
@@ -45,6 +48,8 @@ def process_batch(
         batch_audio_np = [b.squeeze() if b.ndim > 1 else b for b in batch_audio_np]
 
         batch_array = np.stack(batch_audio_np)
+    except CheckpointError:
+        raise
     except Exception as e:
         logger.error(f"Failed to stack audio batch: {e}")
         for label, sample, sample_index, dataset_name, sample_seed in batch_metadata:
@@ -54,6 +59,7 @@ def process_batch(
                 sample=sample,
                 error_message=f"stack-failed: {str(e)[:160]}",
             )
+        tracker.checkpoint()
         return
 
     run_batch_and_record(
@@ -82,12 +88,16 @@ async def run_audio_benchmark(
     content_category: Optional[str] = None,
     score_composition: dict = None,
     multiclass_scoring: bool = False,
+    checkpoint_dir: Optional[str] = None,
+    checkpoint_persist=None,
 ) -> pd.DataFrame:
     """Test model on benchmark audio datasets for AI-generated content detection.
 
     Uses binary classification: 0=real, 1=synthetic (semisynthetic treated as synthetic).
     """
 
+    seed = 42 if seed is None else seed
+
     if batch_size is None:
         batch_size = DEFAULT_AUDIO_BATCH_SIZE
 
@@ -113,6 +123,8 @@ async def run_audio_benchmark(
             crop_prob=0.0,
             records_parquet_path=records_parquet_path,
             run_id=run_id,
+            checkpoint_dir=checkpoint_dir,
+            checkpoint_persist=checkpoint_persist,
             dataset_filters=dataset_filters,
             holdout_weight=holdout_weight,
             holdouts_only=holdouts_only,
@@ -127,7 +139,13 @@ async def run_audio_benchmark(
             benchmark_results["audio_results"] = {"error": "No datasets available"}
             return 0.0
 
-        tracker = create_tracker(run_config, plan, input_specs)
+        tracker = create_tracker(
+            run_config, plan, input_specs, session=session, seed=seed,
+            skip_missing=skip_missing,
+            download_latest_gasstation_data=download_latest_gasstation_data,
+        )
+        benchmark_results["run_id"] = run_config.run_id
+        benchmark_results["checkpoint_dir"] = run_config.checkpoint_dir
 
         # Target sample rate for audio processing
         target_sr = 16000
@@ -144,21 +162,8 @@ async def run_audio_benchmark(
             )
 
             try:
-                is_gasstation = "gasstation" in dataset_cfg.name.lower()
-                if skip_missing:
-                    should_download = False
-                else:
-                    should_download = (
-                        download_latest_gasstation_data if is_gasstation else True
-                    )
-
-                dataset_iterator = DatasetIterator(
-                    dataset_cfg,
-                    max_samples=dataset_cap,
-                    cache_dir=cache_dir,
-                    download=should_download,
-                    hf_token=hf_token,
-                    seed=seed,
+                dataset_iterator = create_dataset_iterator(
+                    run_config, plan, dataset_cfg, aug_pass=False,
                 )
 
                 if skip_missing and dataset_iterator.get_total_cached_count() == 0:
@@ -172,11 +177,17 @@ async def run_audio_benchmark(
 
                 for sample in dataset_iterator:
                     sample_index += 1
+                    if tracker.is_checkpointed(
+                        dataset_name=dataset_cfg.name, sample_index=sample_index, sample=sample,
+                    ):
+                        continue
+                    verify_sample(sample)
                     try:
+                        audio_sample = load_audio_sample(sample)
                         # Check if sample is already preprocessed
-                        if sample.get("is_preprocessed", False):
-                            audio_array = sample.get("preprocessed_waveform")
-                            label = sample.get("label")
+                        if audio_sample.get("is_preprocessed", False):
+                            audio_array = audio_sample.get("preprocessed_waveform")
+                            label = audio_sample.get("label")
 
                             if audio_array is None or label is None:
                                 continue
@@ -187,7 +198,7 @@ async def run_audio_benchmark(
                             # Process raw audio bytes
                             sample_seed_val = None if seed is None else (seed + sample_index)
                             audio_array, label = process_audio_sample(
-                                sample, 
+                                audio_sample,
                                 target_sr=target_sr,
                                 seed=sample_seed_val,
                             )
@@ -220,6 +231,8 @@ async def run_audio_benchmark(
                             if tracker.count % 500 == 0:
                                 logger.info(f"Progress: {tracker.count} samples")
 
+                    except CheckpointError:
+                        raise
                     except Exception as e:
                         logger.warning(
                             f"Failed to process audio sample from {dataset_cfg.name}: {e}"
@@ -244,6 +257,8 @@ async def run_audio_benchmark(
                     logger, tracker, dataset_cfg.name, include_skipped=False
                 )
 
+            except CheckpointError:
+                raise
             except Exception as e:
                 logger.error(f"Failed to process dataset {dataset_cfg.name}: {e}")
                 benchmark_results["errors"].append(
@@ -260,6 +275,8 @@ async def run_audio_benchmark(
         )
         return df
 
+    except CheckpointError:
+        raise
     except Exception as e:
         logger.error(f"Benchmark audio testing failed: {e}")
         benchmark_results["audio_results"] = {"error": str(e)}