diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 148925c..d7e034e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -30,7 +30,9 @@ jobs: # registry importable without the benchmark stack — bitmind-subnet # depends on this contract. This job fails if anyone adds an eager # heavy import to the registry path or a heavy dep to base deps. - - run: pip install . + - run: pip install . pytest + # Checkpoint recovery is CPU-only and must not require the ML extras. + - run: pytest tests/unit/test_checkpoint.py - run: | python - <<'EOF' from gasbench.dataset.config import load_benchmark_datasets_from_yaml @@ -52,6 +54,32 @@ jobs: assert not leaked, f"heavy deps in base install: {leaked}" print("base install contract holds:", dists) EOF + # Recorder recovery shares the existing metrics and parquet paths, without + # loading torch or the inference backends (same lazy-import policy). + - 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) diff --git a/README.md b/README.md index 40c0122..0a88d73 100644 --- a/README.md +++ b/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 `/runs//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 diff --git a/pyproject.toml b/pyproject.toml index 7ba02d0..386f3cd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "gasbench" -version = "0.9.10" +version = "0.10.0" description = "GASBench - ML model benchmark evaluation package" readme = "README.md" requires-python = ">=3.10" diff --git a/src/gasbench/benchmark.py b/src/gasbench/benchmark.py index a08be5a..4a439c2 100644 --- a/src/gasbench/benchmark.py +++ b/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 /runs//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: diff --git a/src/gasbench/benchmarks/__init__.py b/src/gasbench/benchmarks/__init__.py index 9211411..dce5822 100644 --- a/src/gasbench/benchmarks/__init__.py +++ b/src/gasbench/benchmarks/__init__.py @@ -1,18 +1,23 @@ -"""Benchmark execution and metrics.""" +"""Lazy exports, matching gasbench: recorder recovery needs no inference backend.""" -from .image_bench import run_image_benchmark -from .video_bench import run_video_benchmark -from .utils import ( - Metrics, - update_generator_stats, - calculate_per_source_accuracy, -) +from importlib import import_module -__all__ = [ - "run_image_benchmark", - "run_video_benchmark", - "Metrics", - "update_generator_stats", - "calculate_per_source_accuracy", -] +_LAZY_EXPORTS = { + "run_image_benchmark": (".image_bench", "run_image_benchmark"), + "run_video_benchmark": (".video_bench", "run_video_benchmark"), + "Metrics": (".utils.metrics", "Metrics"), + "update_generator_stats": (".utils.metrics", "update_generator_stats"), + "calculate_per_source_accuracy": ( + ".utils.metrics", + "calculate_per_source_accuracy", + ), +} +__all__ = list(_LAZY_EXPORTS) + + +def __getattr__(name): + if name not in _LAZY_EXPORTS: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module_name, attr = _LAZY_EXPORTS[name] + return getattr(import_module(module_name, __name__), attr) diff --git a/src/gasbench/benchmarks/_checkpoint.py b/src/gasbench/benchmarks/_checkpoint.py new file mode 100644 index 0000000..c6657e1 --- /dev/null +++ b/src/gasbench/benchmarks/_checkpoint.py @@ -0,0 +1,289 @@ +"""Append-only inference checkpoints, independent of Modal and the ML stack. + +One coordinator must own a run directory at a time. This class does not provide +distributed locking. The manifest must identify the run, model, evaluator, +complete sample plan (including source revisions), seeds, and scoring settings. +Records use BenchmarkRunRecorder fields: run_id, dataset_name, iteration_index, +sample_id, and aug_pass. No parallel prediction schema or identifier is introduced. + +Local durability uses fsync and atomic rename. Distributed filesystems require a +``persist(directory)`` callback that commits the mounted filesystem before a +batch is acknowledged. A callback failure makes this instance unusable; recover +by reopening the directory from the durable filesystem, never by deleting it. +""" + +import hashlib +import json +import os +import tempfile +from copy import deepcopy +from pathlib import Path +from typing import Callable, Iterable, Mapping, Optional + + +class CheckpointError(RuntimeError): + """Checkpoint is incompatible, corrupt, or cannot safely be persisted.""" + + +def _encode(value, *, sort_keys: bool = True) -> bytes: + return json.dumps( + value, sort_keys=sort_keys, separators=(",", ":"), allow_nan=False + ).encode("utf-8") + + +def _digest(value) -> str: + return hashlib.sha256(_encode(value)).hexdigest() + + +def _read(path: Path): + try: + return json.loads(path.read_bytes()) + except (OSError, ValueError) as exc: + raise CheckpointError(f"Cannot read checkpoint {path.name}") from exc + + +def prediction_key(row: Mapping) -> tuple: + """Identity of one planned prediction, using the existing recorder schema.""" + fields = ("run_id", "dataset_name", "sample_id") + if any(not isinstance(row.get(field), str) or not row[field] for field in fields): + raise CheckpointError("Prediction requires run_id, dataset_name, and sample_id") + index = row.get("iteration_index") + aug_pass = row.get("aug_pass", False) + if type(index) is not int or index < 1 or type(aug_pass) is not bool: + raise CheckpointError( + "Prediction requires a positive sample index and boolean aug_pass" + ) + return (row["run_id"], row["dataset_name"], index, row["sample_id"], aug_pass) + + +class RecorderCheckpoint: + """Internal persistence for BenchmarkRunRecorder rows. + + Reopening validates every committed batch. Temporary writes are ignored; + corrupt or missing committed batches fail closed, including a missing tail. + Replaying an identical record is a no-op; changing a completed prediction + 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, + manifest: Mapping, + *, + persist: Optional[Callable[[Path], None]] = None, + ): + self.directory = Path(directory) + self._persist = persist + self._usable = True + self._records = {} + self._next_batch = 0 + if not isinstance(manifest, Mapping) or not manifest: + raise ValueError("A nonempty run manifest is required") + # Normalize JSON types and detach mutable caller-owned objects. + expected = { + "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") + self._manifest_digest = _digest(expected) + # Provision the run directory before execution. Sync its parent as well + # so a local crash cannot lose a newly created directory entry. + self.directory.mkdir(parents=True, exist_ok=True) + parent_fd = os.open(self.directory.parent, os.O_RDONLY) + try: + os.fsync(parent_fd) + finally: + os.close(parent_fd) + manifest_path = self.directory / "manifest.json" + batches = sorted(self.directory.glob("batch-*.json")) + if manifest_path.exists(): + if _encode(_read(manifest_path)) != _encode(expected): + raise CheckpointError("Checkpoint manifest differs from this run") + else: + if batches: + raise CheckpointError("Checkpoint batches exist without a manifest") + self._write(manifest_path, expected) + + head_path = self.directory / "head.json" + if not head_path.exists(): + if batches: + raise CheckpointError("Checkpoint batches exist without a commit head") + self._write_head(-1, None) + head_envelope = _read(head_path) + if not isinstance(head_envelope, dict) or not isinstance( + head_envelope.get("payload"), dict + ): + raise CheckpointError("Invalid checkpoint commit head") + head = head_envelope["payload"] + try: + head_valid = head_envelope.get("sha256") == _digest(head) + except (TypeError, ValueError): + head_valid = False + if not head_valid: + raise CheckpointError("Invalid checkpoint commit head checksum") + if ( + not isinstance(head, dict) + or type(head.get("last_batch")) is not int + or head["last_batch"] < -1 + or head["last_batch"] >= len(batches) + or (head["last_batch"] == -1 and head.get("sha256") is not None) + ): + raise CheckpointError("Invalid or incomplete checkpoint commit head") + # Files newer than the head were never acknowledged. They can be safely + # replaced when the interrupted batch is replayed. + for path in batches[: head["last_batch"] + 1]: + if path.name != self._batch_name(self._next_batch): + raise CheckpointError("Checkpoint batch sequence is incomplete") + envelope = _read(path) + if not isinstance(envelope, dict): + raise CheckpointError(f"Invalid checkpoint batch {path.name}") + payload = envelope.get("payload") + try: + valid = ( + isinstance(payload, dict) + and envelope.get("sha256") == _digest(payload) + and payload.get("manifest_sha256") == self._manifest_digest + and payload.get("batch_index") == self._next_batch + ) + except (TypeError, ValueError): + valid = False + if not valid: + raise CheckpointError(f"Invalid checkpoint batch {path.name}") + records = self._validate_records(payload.get("records")) + if any(prediction_key(row) in self._records for row in records): + raise CheckpointError("Duplicate work in committed checkpoint batches") + self._records.update((prediction_key(row), row) for row in records) + self._next_batch += 1 + if ( + batches + and head["last_batch"] >= 0 + and envelope["sha256"] != head.get("sha256") + ): + raise CheckpointError("Checkpoint tail differs from commit head") + + @staticmethod + def _batch_name(index: int) -> str: + return f"batch-{index:012d}.json" + + def _validate_records(self, records): + if not isinstance(records, list) or not records: + raise CheckpointError("Checkpoint batch must contain records") + seen = set() + for row in records: + if not isinstance(row, dict): + raise CheckpointError("Checkpoint record must be an object") + key = prediction_key(row) + if row.get("run_id") != self.run_id: + raise CheckpointError("Prediction belongs to a different benchmark run") + if key in seen: + raise CheckpointError("Duplicate prediction in checkpoint batch") + seen.add(key) + return records + + @property + def completed_predictions(self) -> frozenset: + return frozenset(self._records) + + def contains_prediction(self, row: Mapping) -> bool: + return prediction_key(row) in self._records + + @property + def records(self) -> list: + return deepcopy(list(self._records.values())) + + def commit_batch(self, records: Iterable[Mapping]) -> int: + """Commit new records; return their count after persistence succeeds. + + The caller may replay an overlapping batch on recovery. All previously + completed records must match exactly, and only new records are appended. + """ + if not self._usable: + raise CheckpointError("Persistence failed; reopen the durable checkpoint") + rows = self._validate_records( + json.loads(_encode(list(records), sort_keys=False)) + ) + new_rows = [] + for row in rows: + old = self._records.get(prediction_key(row)) + if old is None: + new_rows.append(row) + elif _encode(old) != _encode(row): + raise CheckpointError( + f"Conflicting completed work: {prediction_key(row)}" + ) + if not new_rows: + return 0 + payload = { + "manifest_sha256": self._manifest_digest, + "batch_index": self._next_batch, + "records": new_rows, + } + self._write( + self.directory / self._batch_name(self._next_batch), + {"payload": payload, "sha256": _digest(payload)}, + persist=False, + ) + self._write_head(self._next_batch, _digest(payload)) + self._records.update((prediction_key(row), row) for row in new_rows) + self._next_batch += 1 + return len(new_rows) + + def _write_head(self, index: int, checksum: Optional[str]) -> None: + payload = {"last_batch": index, "sha256": checksum} + self._write( + self.directory / "head.json", + {"payload": payload, "sha256": _digest(payload)}, + ) + + 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. + temp_path = None + try: + content = _encode(value, sort_keys=False) + 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) + try: + os.fsync(fd) + finally: + os.close(fd) + if persist and self._persist is not None: + self._persist(self.directory) + finally: + 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 diff --git a/src/gasbench/benchmarks/audio_bench.py b/src/gasbench/benchmarks/audio_bench.py index aa548b0..68d1e9d 100644 --- a/src/gasbench/benchmarks/audio_bench.py +++ b/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)} diff --git a/src/gasbench/benchmarks/common.py b/src/gasbench/benchmarks/common.py index 1719988..e25f355 100644 --- a/src/gasbench/benchmarks/common.py +++ b/src/gasbench/benchmarks/common.py @@ -1,6 +1,12 @@ import json import time -from dataclasses import dataclass +import hashlib +import uuid +import platform +from importlib.metadata import version, PackageNotFoundError +from pathlib import Path +from typing import Callable +from dataclasses import asdict, dataclass, field, fields from typing import Dict, List, Optional, Tuple import numpy as np @@ -15,8 +21,11 @@ apply_mode_to_datasets, ) from ..processing.transforms import extract_target_size_from_input_specs +from ..dataset.iterator import DatasetIterator +from ._checkpoint import CheckpointError, RecorderCheckpoint from .recording import ( BenchmarkRunRecorder, + build_sample_id, compute_metrics_from_df, compute_per_dataset_from_df, compute_generator_stats_from_df, @@ -57,6 +66,9 @@ def run_batch_and_record( start = time.time() try: outputs = session.run(None, {input_specs[0].name: batch_array}) + if len(outputs[0]) != len(batch_metadata): + raise ValueError("Model output batch size differs from input") + predictions = [process_model_output(output) for output in outputs[0]] except Exception as e: logger.error(f"Inference failed: {e} (batch shape: {batch_array.shape})") for label, sample, sample_index, dataset_name, sample_seed in batch_metadata: @@ -65,7 +77,9 @@ def run_batch_and_record( sample_index=sample_index, sample=sample, error_message=f"inference-failed: {str(e)[:160]}", + aug_pass=aug_pass, ) + tracker.checkpoint() return batch_inference_time = (time.time() - start) * 1000 @@ -74,7 +88,7 @@ def run_batch_and_record( for i, (label, sample, sample_index, dataset_name, sample_seed) in enumerate( batch_metadata ): - predicted, pred_probs = process_model_output(outputs[0][i]) + predicted, pred_probs = predictions[i] tracker.add_ok( dataset_name=dataset_name, sample_index=sample_index, @@ -90,6 +104,8 @@ def run_batch_and_record( aug_pass=aug_pass, ) + tracker.checkpoint() + @dataclass class BenchmarkRunConfig: @@ -105,6 +121,8 @@ class BenchmarkRunConfig: crop_prob: float records_parquet_path: Optional[str] run_id: Optional[str] = None + checkpoint_dir: Optional[str] = None + checkpoint_persist: Optional[Callable[[Path], None]] = None dataset_filters: Optional[List[str]] = None holdout_weight: float = 1.0 # (Legacy) weight multiplier for holdout datasets in benchmark_score only holdouts_only: bool = False # If True, only run holdout datasets (requires holdout_config_path) @@ -112,6 +130,8 @@ class BenchmarkRunConfig: score_composition: Optional[Dict[str, float]] = None # Target score weight share per provenance class, e.g. {"public": 0.5, "holdout": 0.3, "gasstation": 0.2}; weights all metrics incl. sn34_score multiclass_scoring: bool = False # Derive sn34_score from Gorodkin multiclass MCC + multiclass Brier instead of the binary real-vs-not-real collapse. No-op for audio (2 classes). n_aug_per_dataset: int = 0 # Number of samples per dataset to re-evaluate with robustness augmentations (0 = disabled) + aug_cache_dir: Optional[str] = None + aug_cache_readonly: bool = False aug_weight: float = 0.2 # Weight of aug_sn34_score in blended final score (when n_aug_per_dataset > 0) @@ -134,6 +154,7 @@ class BenchmarkPlan: target_size: Tuple[int, int] dataset_info: Dict sampling_summary: SamplingSummary + samples: Dict = field(default_factory=dict) def build_plan( @@ -238,10 +259,228 @@ def build_plan( ) +def fingerprint_files(paths, root): + """Bind model/evaluator inputs by content, excluding transient Python bytecode.""" + digest = hashlib.sha256() + for path in sorted(paths): + if not path.is_file() or "__pycache__" in path.parts or path.suffix == ".pyc": + continue + digest.update(str(path.relative_to(root)).encode()) + digest.update(b"\0") + with path.open("rb") as stream: + while chunk := stream.read(1024 * 1024): + digest.update(chunk) + digest.update(b"\0") + return digest.hexdigest() + + +def runtime_versions(): + versions = { + "python": platform.python_version(), + "platform": platform.system(), + "machine": platform.machine(), + } + for package in ( + "gasbench", + "numpy", + "torch", + "torchvision", + "scipy", + "pillow", + "opencv-python-headless", + "decord", + "torchcodec", + ): + try: + versions[package] = version(package) + except PackageNotFoundError: + versions[package] = None + return versions + + +def sample_files(sample): + if "video_frames" in sample: + return [Path(p) for p in sample["video_frames"]] + return [ + Path(sample[f"{modality}_path"]) + for modality in ("image", "video", "audio") + if f"{modality}_path" in sample + ] + + +def sample_digest(sample): + paths = sample_files(sample) + if not paths or any(not p.is_file() for p in paths): + raise CheckpointError("A selected sample is missing from the cache") + return fingerprint_files(paths, paths[0].parent) + + +def verify_sample(sample): + """Reject changed pending input before preprocessing; completed rows need no I/O.""" + if "content_sha256" in sample: + try: + if sample_digest(sample) != sample["content_sha256"]: + raise CheckpointError( + "A selected sample changed since this run started" + ) + except OSError as exc: + raise CheckpointError("Cannot read a selected sample") from exc + + +def use_augmentation_cache(sample, path): + """Use only the derived artifact selected for this run, or regenerate it.""" + path = Path(path) + if "augmentation_cache_sha256" not in sample: + return path.is_file() + expected = sample["augmentation_cache_sha256"] + if expected is None: + return False + try: + if not path.is_file() or fingerprint_files([path], path.parent) != expected: + raise CheckpointError("Selected augmentation cache changed or disappeared") + except OSError as exc: + raise CheckpointError("Cannot read selected augmentation cache") from exc + return True + + +def create_dataset_iterator(config, plan, dataset_config, *, aug_pass=False): + """Use the iterator's frozen selection for both normal and resumed execution.""" + samples = plan.samples[dataset_config.name]["aug" if aug_pass else "base"] + return DatasetIterator( + dataset_config, + max_samples=max(len(samples), 1), + cache_dir=config.cache_dir, + download=False, + lazy_read=True, + frozen_samples=samples, + ) + + def create_tracker( - config: BenchmarkRunConfig, plan: BenchmarkPlan, input_specs + config: BenchmarkRunConfig, + plan: BenchmarkPlan, + input_specs, + *, + session, + seed, + skip_missing=False, + download_latest_gasstation_data=False, ) -> BenchmarkRunRecorder: - return BenchmarkRunRecorder( + """Freeze inputs once, then open the single recorder for this run.""" + config.run_id = config.run_id or str(uuid.uuid4()) + # run_id is an external identifier, not a path supplied by the caller. + if Path(config.run_id).name != config.run_id or config.run_id in (".", ".."): + raise ValueError("run_id must be a single path component") + directory = ( + Path(config.checkpoint_dir) + if config.checkpoint_dir + else Path(config.cache_dir) / "runs" / config.run_id / "checkpoint" + ) + config.checkpoint_dir = str(directory) + saved = RecorderCheckpoint.read_manifest(directory) + model_dir = getattr(session, "model_dir", None) + if model_dir is None: + raise ValueError( + "Benchmark inference sessions must expose model_dir for checkpoint identity" + ) + model_dir = Path(model_dir) + if not model_dir.is_dir(): + raise CheckpointError("Model directory is unavailable") + evaluator_dir = Path(__file__).resolve().parents[1] + settings = { + item.name: getattr(config, item.name) + for item in fields(config) + if item.name + not in ( + "hf_token", + "records_parquet_path", + "checkpoint_dir", + "checkpoint_persist", + ) + } + context = { + "settings": settings, + "seed": seed, + "runtime": runtime_versions(), + "skip_missing": skip_missing, + "model_sha256": fingerprint_files(model_dir.rglob("*"), model_dir), + "evaluator_sha256": fingerprint_files( + evaluator_dir.rglob("*.py"), evaluator_dir + ), + "input_specs": [ + {"name": spec.name, "shape": list(spec.shape), "type": spec.type} + for spec in input_specs + ], + "preprocessing": session.get_preprocessing_config() + if hasattr(session, "get_preprocessing_config") + else {}, + "datasets": [asdict(dataset) for dataset in plan.available_datasets], + "sampling_plan": plan.sampling_plan, + "target_size": list(plan.target_size), + } + # Normalize tuples/paths consistently with the recorder's JSON manifest. + context = json.loads(json.dumps(context)) + if saved is not None: + previous = saved.get("benchmark", {}) + if {k: v for k, v in previous.items() if k != "samples"} != context: + raise CheckpointError( + "Checkpoint configuration or model differs from this run" + ) + if "samples" not in previous: + raise CheckpointError("Checkpoint has no frozen sample selection") + plan.samples = previous["samples"] + else: + for dataset in plan.available_datasets: + selections = {} + for pass_name, cap in ( + ("base", plan.sampling_plan[dataset.name]), + ("aug", config.n_aug_per_dataset), + ): + if cap <= 0: + selections[pass_name] = [] + continue + iterator = DatasetIterator( + dataset, + max_samples=cap, + cache_dir=config.cache_dir, + download=not skip_missing + and ( + download_latest_gasstation_data + if "gasstation" in dataset.name.lower() + else True + ), + hf_token=config.hf_token, + seed=seed, + metadata_only=True, + ) + selections[pass_name] = list(iterator) + for sample in selections[pass_name]: + sample["content_sha256"] = sample_digest(sample) + if pass_name == "aug" and config.aug_cache_dir: + from .aug_cache import img_aug_cache_path, vid_aug_cache_path + + cache_path = ( + img_aug_cache_path + if config.modality == "image" + else vid_aug_cache_path + ) + path = Path( + cache_path( + config.aug_cache_dir, + build_sample_id(sample), + plan.target_size, + ) + ) + # Newly generated cache files are outputs of this attempt; + # only artifacts present in the frozen plan may be inputs. + sample["augmentation_cache_sha256"] = ( + fingerprint_files([path], path.parent) + if path.is_file() + else None + ) + plan.samples[dataset.name] = selections + context["samples"] = plan.samples + tracker = BenchmarkRunRecorder( run_id=config.run_id, mode=config.mode, modality=config.modality, @@ -249,7 +488,14 @@ def create_tracker( model_input_name=input_specs[0].name if input_specs else None, augment_level=config.augment_level or 0, crop_prob=config.crop_prob or 0.0, + checkpoint_dir=directory, + checkpoint_context=context, + checkpoint_persist=config.checkpoint_persist, + ) + get_logger(__name__).info( + f"Run {config.run_id}: restored {tracker.count} rows from {directory}" ) + return tracker def finalize_run( @@ -262,6 +508,7 @@ def finalize_run( extra_fields: Optional[Dict] = None, ): logger = get_logger(__name__) + tracker.checkpoint() df = tracker.to_dataframe() metric_pack = compute_metrics_from_df( df, diff --git a/src/gasbench/benchmarks/image_bench.py b/src/gasbench/benchmarks/image_bench.py index e051716..9d1254d 100644 --- a/src/gasbench/benchmarks/image_bench.py +++ b/src/gasbench/benchmarks/image_bench.py @@ -16,13 +16,16 @@ DEFAULT_IMAGE_BATCH_SIZE, ) -from ..dataset.iterator import DatasetIterator +from ._checkpoint import CheckpointError from .recording import BenchmarkRunRecorder, log_dataset_summary, build_sample_id from .common import ( BenchmarkRunConfig, build_plan, create_tracker, + create_dataset_iterator, + verify_sample, + use_augmentation_cache, finalize_run, stack_uniform_batch, run_batch_and_record, @@ -56,7 +59,9 @@ def __init__( robustness_pass=False, aug_cache_dir=None, aug_cache_readonly=False, + tracker=None, ): + self.tracker = tracker self.dataset_iterator = dataset_iterator self.target_size = target_size self.batch_size = batch_size @@ -79,6 +84,7 @@ def __init__( def _read_and_preprocess(self, sample, sample_index, dataset_name): """Read file from disk (if lazy), decode, and augment. Runs in worker thread.""" + verify_sample(sample) try: image_path = sample.get("image_path") if image_path: @@ -95,7 +101,7 @@ def _read_and_preprocess(self, sample, sample_index, dataset_name): if self.aug_cache_dir: sid = build_sample_id(sample) cache_path = img_aug_cache_path(self.aug_cache_dir, sid, self.target_size) - if os.path.exists(cache_path): + if use_augmentation_cache(sample, cache_path): aug_hwc = np.load(cache_path) else: aug_hwc, _, _, _ = apply_robustness_augmentations( @@ -129,6 +135,8 @@ def _read_and_preprocess(self, sample, sample_index, dataset_name): "dataset_name": dataset_name, "sample_seed": sample_seed, } + except CheckpointError: + raise except Exception as e: logger.warning(f"Failed to preprocess sample {sample_index}: {e}") return None @@ -148,6 +156,11 @@ def _producer_loop(self): while len(pending) < max_in_flight and not exhausted: try: idx, sample = next(sample_iter) + if self.tracker is not None and self.tracker.is_checkpointed( + dataset_name=dataset_name, sample_index=idx, + sample=sample, aug_pass=self.robustness_pass, + ): + continue future = self.executor.submit( self._read_and_preprocess, sample, idx, dataset_name ) @@ -166,6 +179,8 @@ def _producer_loop(self): break try: result = future.result() + except CheckpointError: + raise except Exception: continue if result is not None: @@ -194,6 +209,8 @@ def __next__(self): try: batch = self.batch_queue.get(timeout=300) if batch is None: + if self.error: + raise self.error raise StopIteration return batch except Empty: @@ -257,9 +274,13 @@ async def run_image_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, ) -> pd.DataFrame: """Test model on benchmark image datasets for AI-generated content detection.""" + seed = 42 if seed is None else seed + if batch_size is None: batch_size = DEFAULT_IMAGE_BATCH_SIZE @@ -284,6 +305,8 @@ async def run_image_benchmark( crop_prob=crop_prob or 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, @@ -292,6 +315,8 @@ async def run_image_benchmark( multiclass_scoring=multiclass_scoring, n_aug_per_dataset=n_aug_per_dataset, aug_weight=aug_weight, + aug_cache_dir=aug_cache_dir, + aug_cache_readonly=aug_cache_readonly, ) plan = build_plan(logger, run_config, input_specs) @@ -300,7 +325,13 @@ async def run_image_benchmark( benchmark_results["image_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 benchmark_results.setdefault("errors", []) for dataset_idx, dataset_config in enumerate(plan.available_datasets): @@ -311,22 +342,8 @@ async def run_image_benchmark( ) try: - is_gasstation = "gasstation" in dataset_config.name.lower() - if skip_missing: - should_download = False - else: - should_download = ( - download_latest_gasstation_data if is_gasstation else True - ) - - dataset_iterator = DatasetIterator( - dataset_config, - max_samples=dataset_cap, - cache_dir=cache_dir, - download=should_download, - hf_token=hf_token, - seed=seed, - lazy_read=True, + dataset_iterator = create_dataset_iterator( + run_config, plan, dataset_config, aug_pass=False, ) if skip_missing and dataset_iterator.get_total_cached_count() == 0: @@ -335,6 +352,7 @@ async def run_image_benchmark( pipeline = PrefetchPipeline( dataset_iterator=dataset_iterator, + tracker=tracker, target_size=plan.target_size, batch_size=batch_size, seed=seed, @@ -378,6 +396,8 @@ async def run_image_benchmark( logger, tracker, dataset_config.name, include_skipped=False ) + except CheckpointError: + raise except Exception as e: logger.error(f"Failed to process dataset {dataset_config.name}: {e}") benchmark_results["errors"].append( @@ -396,19 +416,8 @@ async def run_image_benchmark( f"{dataset_config.name}" ) try: - is_gasstation = "gasstation" in dataset_config.name.lower() - should_download = ( - download_latest_gasstation_data if is_gasstation else True - ) if not skip_missing else False - - aug_iterator = DatasetIterator( - dataset_config, - max_samples=n_aug_per_dataset, - cache_dir=cache_dir, - download=should_download, - hf_token=hf_token, - seed=seed, - lazy_read=True, + aug_iterator = create_dataset_iterator( + run_config, plan, dataset_config, aug_pass=True, ) if skip_missing and aug_iterator.get_total_cached_count() == 0: @@ -416,6 +425,7 @@ async def run_image_benchmark( aug_pipeline = PrefetchPipeline( dataset_iterator=aug_iterator, + tracker=tracker, target_size=plan.target_size, batch_size=batch_size, seed=aug_seed, @@ -453,6 +463,8 @@ async def run_image_benchmark( finally: aug_pipeline.close() + except CheckpointError: + raise except Exception as e: logger.error( f"Robustness pass failed for {dataset_config.name}: {e}" @@ -468,6 +480,8 @@ async def run_image_benchmark( ) return df + except CheckpointError: + raise except Exception as e: logger.error(f"Benchmark image testing failed: {e}") benchmark_results["image_results"] = {"error": str(e)} diff --git a/src/gasbench/benchmarks/recording.py b/src/gasbench/benchmarks/recording.py index 9be7742..c8357a3 100644 --- a/src/gasbench/benchmarks/recording.py +++ b/src/gasbench/benchmarks/recording.py @@ -2,12 +2,14 @@ import time import os import hashlib -from typing import Any, Dict, List, Optional, Tuple +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple import pandas as pd import numpy as np -from .utils import Metrics +from .utils.metrics import Metrics +from ._checkpoint import RecorderCheckpoint from ..constants import MODALITY_NUM_CLASSES @@ -22,7 +24,15 @@ def __init__( model_name: Optional[str] = None, augment_level: Optional[int] = 0, crop_prob: float = 0.0, + checkpoint_dir: Optional[Path] = None, + checkpoint_context: Optional[Dict[str, Any]] = None, + checkpoint_persist: Optional[Callable[[Path], None]] = None, ): + # Identity must be supplied by the coordinator, never regenerated on resume. + if checkpoint_dir is not None and (not run_id or not checkpoint_context): + raise ValueError("Checkpointing requires run_id and benchmark context") + if checkpoint_dir is None and (checkpoint_context is not None or checkpoint_persist is not None): + raise ValueError("Checkpoint context/persistence requires checkpoint_dir") self.run_id = run_id or str(uuid.uuid4()) self.run_started_at = int(time.time()) self.mode = mode @@ -39,6 +49,65 @@ def __init__( # log_dataset_summary never has to materialise a full DataFrame. # Structure: {dataset_name: {"ok": int, "correct": int, "skipped": int}} self._dataset_counts: Dict[str, Dict[str, int]] = {} + self._checkpoint = None + self._checkpointed_count = 0 + if checkpoint_dir is not None: + self._checkpoint = RecorderCheckpoint( + checkpoint_dir, + { + "run_id": self.run_id, + "recorder": { + "mode": self.mode, "modality": self.modality, + "target_height": self.target_height, "target_width": self.target_width, + "input_name": self.model_input_name, "model_name": self.model_name, + "augment_level": self.augment_level, "crop_prob": self.crop_prob, + }, + "benchmark": checkpoint_context, + }, + persist=checkpoint_persist, + ) + for row in self._checkpoint.records: + self._append_row(row) + if self.rows: + self.run_started_at = self.rows[0]["run_started_at"] + self._checkpointed_count = len(self.rows) + + def checkpoint(self) -> int: + """Commit pending recorder rows; distributed storage must supply persist. + + checkpoint_context must include the model/evaluator identity, immutable + sample plan with source revisions, seed, and scoring settings. Recorder + settings and run_id are included automatically. One coordinator owns + this run directory; this method is not a distributed writer lock. + """ + if self._checkpoint is None or self._checkpointed_count == len(self.rows): + return 0 + count = self._checkpoint.commit_batch(self.rows[self._checkpointed_count:]) + self._checkpointed_count = len(self.rows) + return count + + def is_checkpointed(self, *, dataset_name: str, sample_index: int, + sample: Dict[str, Any], aug_pass: bool = False) -> bool: + """Check durable progress before decoding or inferring a planned sample.""" + if self._checkpoint is None: + return False + return self._checkpoint.contains_prediction({ + "run_id": self.run_id, "dataset_name": dataset_name, + "iteration_index": sample_index, "sample_id": build_sample_id(sample), + "aug_pass": aug_pass, + }) + + def _append_row(self, row: Dict[str, Any]) -> None: + """Apply the same counters for live records and checkpoint restoration.""" + self.rows.append(row) + if row.get("aug_pass", False) or row["status"] == "error": + return + ds = self._dataset_counts.setdefault(row["dataset_name"], {"ok": 0, "correct": 0, "skipped": 0}) + if row["status"] == "ok": + ds["ok"] += 1 + ds["correct"] += int(bool(row["correct"])) + elif row["status"] == "skipped": + ds["skipped"] += 1 @property def count(self) -> int: @@ -103,7 +172,7 @@ def add_ok( "dataset_path": sample.get("dataset_path"), "hf_resolved_revision": sample.get("hf_resolved_revision"), "archive_filename": sample.get("archive_filename"), - "path_in_archive": sample.get("member_path"), + "path_in_archive": sample.get("path_in_archive") or sample.get("member_path"), "source_file": sample.get("source_file"), "iso_week": sample.get("iso_week"), "cache_relpath": sample.get("cache_relpath"), @@ -119,15 +188,7 @@ def add_ok( row["sample_compound_id"] = build_compound_id(row) row["sample_display_uri"] = build_display_uri(row) - self.rows.append(row) - - # Maintain incremental counters so per-dataset logging is O(1). - # Aug-pass rows are not counted here — they are reported separately. - if not aug_pass: - ds = self._dataset_counts.setdefault(dataset_name, {"ok": 0, "correct": 0, "skipped": 0}) - ds["ok"] += 1 - if bool(predicted == label): - ds["correct"] += 1 + self._append_row(row) def add_skip( self, @@ -136,6 +197,7 @@ def add_skip( sample_index: int, sample: Dict[str, Any], reason: str, + aug_pass: bool = False, ): row = { "run_id": self.run_id, @@ -152,6 +214,7 @@ def add_skip( "iteration_index": int(sample_index), "media_type": sample.get("media_type"), "status": "skipped", + "aug_pass": bool(aug_pass), "label": None, "predicted": None, "probs": None, @@ -170,7 +233,7 @@ def add_skip( "dataset_path": sample.get("dataset_path"), "hf_resolved_revision": sample.get("hf_resolved_revision"), "archive_filename": sample.get("archive_filename"), - "path_in_archive": sample.get("member_path"), + "path_in_archive": sample.get("path_in_archive") or sample.get("member_path"), "source_file": sample.get("source_file"), "iso_week": sample.get("iso_week"), "cache_relpath": sample.get("cache_relpath"), @@ -183,10 +246,7 @@ def add_skip( row["sample_id"] = build_sample_id(row) row["sample_compound_id"] = build_compound_id(row) row["sample_display_uri"] = build_display_uri(row) - self.rows.append(row) - - ds = self._dataset_counts.setdefault(dataset_name, {"ok": 0, "correct": 0, "skipped": 0}) - ds["skipped"] += 1 + self._append_row(row) def add_error( self, @@ -195,6 +255,7 @@ def add_error( sample_index: int, sample: Dict[str, Any], error_message: str, + aug_pass: bool = False, ): row = { "run_id": self.run_id, @@ -211,6 +272,7 @@ def add_error( "iteration_index": int(sample_index), "media_type": sample.get("media_type"), "status": "error", + "aug_pass": bool(aug_pass), "label": None, "predicted": None, "probs": None, @@ -229,7 +291,7 @@ def add_error( "dataset_path": sample.get("dataset_path"), "hf_resolved_revision": sample.get("hf_resolved_revision"), "archive_filename": sample.get("archive_filename"), - "path_in_archive": sample.get("member_path"), + "path_in_archive": sample.get("path_in_archive") or sample.get("member_path"), "source_file": sample.get("source_file"), "iso_week": sample.get("iso_week"), "cache_relpath": sample.get("cache_relpath"), @@ -242,7 +304,7 @@ def add_error( row["sample_id"] = build_sample_id(row) row["sample_compound_id"] = build_compound_id(row) row["sample_display_uri"] = build_display_uri(row) - self.rows.append(row) + self._append_row(row) def get_dataset_summary(self, dataset_name: str, include_skipped: bool = False) -> Dict[str, Any]: """Return per-dataset accuracy from incremental counters — O(1), no DataFrame.""" diff --git a/src/gasbench/benchmarks/utils/__init__.py b/src/gasbench/benchmarks/utils/__init__.py index 5a9caeb..871caad 100644 --- a/src/gasbench/benchmarks/utils/__init__.py +++ b/src/gasbench/benchmarks/utils/__init__.py @@ -1,20 +1,23 @@ -from .metrics import ( - Metrics, - update_generator_stats, - calculate_per_source_accuracy, -) -from .inference import create_inference_session, process_model_output -from .pytorch_session import PyTorchInferenceSession -from .custom_model_loader import load_custom_model, validate_model_directory +"""Lazy exports, matching gasbench: recorder recovery needs no inference backend.""" -__all__ = [ - "Metrics", - "update_generator_stats", - "calculate_per_source_accuracy", - "create_inference_session", - "process_model_output", - "PyTorchInferenceSession", - "load_custom_model", - "validate_model_directory", -] +from importlib import import_module +_LAZY_EXPORTS = { + "Metrics": (".metrics", "Metrics"), + "update_generator_stats": (".metrics", "update_generator_stats"), + "calculate_per_source_accuracy": (".metrics", "calculate_per_source_accuracy"), + "create_inference_session": (".inference", "create_inference_session"), + "process_model_output": (".inference", "process_model_output"), + "PyTorchInferenceSession": (".pytorch_session", "PyTorchInferenceSession"), + "load_custom_model": (".custom_model_loader", "load_custom_model"), + "validate_model_directory": (".custom_model_loader", "validate_model_directory"), +} + +__all__ = list(_LAZY_EXPORTS) + + +def __getattr__(name): + if name not in _LAZY_EXPORTS: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module_name, attr = _LAZY_EXPORTS[name] + return getattr(import_module(module_name, __name__), attr) diff --git a/src/gasbench/benchmarks/video_bench.py b/src/gasbench/benchmarks/video_bench.py index a8ddb70..357f5a0 100644 --- a/src/gasbench/benchmarks/video_bench.py +++ b/src/gasbench/benchmarks/video_bench.py @@ -16,13 +16,16 @@ ) from ..config import DEFAULT_VIDEO_BATCH_SIZE from ..constants import MAX_VIDEO_NUM_FRAMES -from ..dataset.iterator import DatasetIterator +from ._checkpoint import CheckpointError from .recording import BenchmarkRunRecorder, log_dataset_summary, build_sample_id from .common import ( BenchmarkRunConfig, build_plan, create_tracker, + create_dataset_iterator, + verify_sample, + use_augmentation_cache, finalize_run, stack_uniform_batch, run_batch_and_record, @@ -58,7 +61,9 @@ def __init__( robustness_pass=False, aug_cache_dir=None, aug_cache_readonly=False, + tracker=None, ): + self.tracker = tracker self.dataset_iterator = dataset_iterator self.target_size = target_size self.batch_size = batch_size @@ -83,6 +88,7 @@ def __init__( def _read_and_preprocess(self, sample, sample_index, dataset_name): """Read file from disk (if lazy), decode, and augment. Runs in worker thread.""" + verify_sample(sample) try: video_path = sample.get("video_path") if video_path: @@ -108,7 +114,7 @@ def _read_and_preprocess(self, sample, sample_index, dataset_name): if self.aug_cache_dir: sid = build_sample_id(sample) cache_path = vid_aug_cache_path(self.aug_cache_dir, sid, self.target_size) - if os.path.exists(cache_path): + if use_augmentation_cache(sample, cache_path): aug_thwc = np.load(cache_path) else: aug_thwc, _, _, _ = apply_video_robustness_augmentations( @@ -130,6 +136,8 @@ def _read_and_preprocess(self, sample, sample_index, dataset_name): level=self.augment_level, crop_prob=self.crop_prob, ) + except CheckpointError: + raise except Exception as e: logger.error(f"Video augmentation failed: {e}") return None @@ -146,6 +154,8 @@ def _read_and_preprocess(self, sample, sample_index, dataset_name): "dataset_name": dataset_name, "sample_seed": sample_seed, } + except CheckpointError: + raise except Exception as e: logger.warning(f"Failed to preprocess video sample {sample_index}: {e}") return None @@ -165,6 +175,11 @@ def _producer_loop(self): while len(pending) < max_in_flight and not exhausted: try: idx, sample = next(sample_iter) + if self.tracker is not None and self.tracker.is_checkpointed( + dataset_name=dataset_name, sample_index=idx, + sample=sample, aug_pass=self.robustness_pass, + ): + continue future = self.executor.submit( self._read_and_preprocess, sample, idx, dataset_name ) @@ -183,6 +198,8 @@ def _producer_loop(self): break try: result = future.result() + except CheckpointError: + raise except Exception: continue if result is not None: @@ -213,6 +230,8 @@ def __next__(self): try: batch = self.batch_queue.get(timeout=300) if batch is None: + if self.error: + raise self.error raise StopIteration return batch except Empty: @@ -276,9 +295,13 @@ async def run_video_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, ) -> pd.DataFrame: """Test model on benchmark video datasets for AI-generated content detection.""" + seed = 42 if seed is None else seed + if batch_size is None: batch_size = DEFAULT_VIDEO_BATCH_SIZE @@ -322,6 +345,8 @@ async def run_video_benchmark( crop_prob=crop_prob or 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, @@ -330,6 +355,8 @@ async def run_video_benchmark( multiclass_scoring=multiclass_scoring, n_aug_per_dataset=n_aug_per_dataset, aug_weight=aug_weight, + aug_cache_dir=aug_cache_dir, + aug_cache_readonly=aug_cache_readonly, ) plan = build_plan(logger, run_config, input_specs) if not plan: @@ -337,7 +364,13 @@ async def run_video_benchmark( benchmark_results["video_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 with video_archive_manager(cache_dir=cache_dir) as video_cache: skipped_samples = 0 @@ -354,22 +387,8 @@ async def run_video_benchmark( ) try: - is_gasstation = "gasstation" in dataset_config.name.lower() - if skip_missing: - should_download = False - else: - should_download = ( - download_latest_gasstation_data if is_gasstation else True - ) - - dataset_iterator = DatasetIterator( - dataset_config, - max_samples=dataset_cap, - cache_dir=cache_dir, - download=should_download, - hf_token=hf_token, - seed=seed, - lazy_read=True, + dataset_iterator = create_dataset_iterator( + run_config, plan, dataset_config, aug_pass=False, ) if skip_missing and dataset_iterator.get_total_cached_count() == 0: @@ -378,6 +397,7 @@ async def run_video_benchmark( pipeline = VideoPrefetchPipeline( dataset_iterator=dataset_iterator, + tracker=tracker, target_size=plan.target_size, batch_size=batch_size, seed=seed, @@ -423,6 +443,8 @@ async def run_video_benchmark( logger, tracker, dataset_config.name, include_skipped=True ) + except CheckpointError: + raise except Exception as e: logger.error( f"Failed to process dataset {dataset_config.name}: {e}" @@ -444,19 +466,8 @@ async def run_video_benchmark( f"{dataset_config.name}" ) try: - is_gasstation = "gasstation" in dataset_config.name.lower() - should_download = ( - download_latest_gasstation_data if is_gasstation else True - ) if not skip_missing else False - - aug_iterator = DatasetIterator( - dataset_config, - max_samples=n_aug_per_dataset, - cache_dir=cache_dir, - download=should_download, - hf_token=hf_token, - seed=seed, - lazy_read=True, + aug_iterator = create_dataset_iterator( + run_config, plan, dataset_config, aug_pass=True, ) if skip_missing and aug_iterator.get_total_cached_count() == 0: @@ -464,6 +475,7 @@ async def run_video_benchmark( aug_pipeline = VideoPrefetchPipeline( dataset_iterator=aug_iterator, + tracker=tracker, target_size=plan.target_size, batch_size=batch_size, seed=aug_seed, @@ -503,6 +515,8 @@ async def run_video_benchmark( finally: aug_pipeline.close() + except CheckpointError: + raise except Exception as e: logger.error( f"Video robustness pass failed for {dataset_config.name}: {e}" @@ -528,6 +542,8 @@ async def run_video_benchmark( return df + except CheckpointError: + raise except Exception as e: logger.error(f"Benchmark video testing failed: {e}") benchmark_results["video_results"] = {"error": str(e)} diff --git a/src/gasbench/cli.py b/src/gasbench/cli.py index 3d0d1bb..fb1077b 100644 --- a/src/gasbench/cli.py +++ b/src/gasbench/cli.py @@ -159,6 +159,7 @@ def command_run(args): records_parquet_path=parquet_path, skip_missing=getattr(args, "skip_missing", False), run_id=getattr(args, "run_id", None), + checkpoint_dir=getattr(args, "checkpoint_dir", None), holdout_weight=getattr(args, "holdout_weight", 1.0), holdouts_only=holdouts_only, score_composition=score_composition, @@ -526,6 +527,10 @@ def main(): action="store_true", help="Skip datasets that are not already cached (do not download missing datasets)", ) + run_parser.add_argument( + "--checkpoint-dir", + help="Checkpoint storage directory (default: /runs//checkpoint). Reuse with --run-id to resume.", + ) run_parser.add_argument( "--run-id", type=str, diff --git a/src/gasbench/dataset/iterator.py b/src/gasbench/dataset/iterator.py index 7468d5c..89fd72a 100644 --- a/src/gasbench/dataset/iterator.py +++ b/src/gasbench/dataset/iterator.py @@ -36,6 +36,8 @@ def __init__( hf_token: Optional[str] = None, seed: Optional[int] = None, lazy_read: bool = False, + metadata_only: bool = False, + frozen_samples: Optional[list] = None, ): self.config = dataset_config self.max_samples = max_samples or DEFAULT_MAX_SAMPLES @@ -46,10 +48,12 @@ def __init__( self.hf_token = hf_token self.seed = seed self.lazy_read = lazy_read + self.metadata_only = metadata_only + self.frozen_samples = frozen_samples self.is_gasstation = "gasstation" in dataset_config.name.lower() - if self.is_gasstation: + if self.is_gasstation and frozen_samples is None: self.target_weeks = gasstation_utils.calculate_target_weeks( num_weeks, dataset_path=dataset_config.path, @@ -72,7 +76,7 @@ def __init__( self.source_kind = getattr(self.config, "source", "huggingface") self.hf_resolved_revision = None - if download: + if download and frozen_samples is None: self.ensure_cached() def __iter__(self): @@ -118,6 +122,10 @@ def get_samples(self): Note: This generator yields ALL available cached samples. The __next__() method is responsible for enforcing max_samples limit. This avoids double-counting issues. """ + if self.frozen_samples is not None: + yield from self.frozen_samples + return + try: if self.is_gasstation: # Load from all week directories @@ -148,6 +156,8 @@ def get_samples(self): yield sample except Exception as e: + if self.metadata_only: + raise logger.error(f"Failed to get samples from {self.config.name}: {e}") import traceback @@ -564,6 +574,8 @@ def get_total_cached_count(self) -> int: This reads metadata files directly without loading actual samples. """ + if self.frozen_samples is not None: + return len(self.frozen_samples) try: if self.is_gasstation: return gasstation_utils.get_total_cached_samples(self.week_dirs) @@ -623,6 +635,29 @@ def extract_index(filename): # The current registry owns labels; caches may predate a taxonomy change. metadata = {**metadata, "media_type": self.config.media_type} + if self.metadata_only: + sample = { + **metadata, + "dataset_name": self.config.name, + "dataset_path": self.config.path, + "source_kind": self.source_kind, + "hf_resolved_revision": dataset_info.get( + "hf_resolved_revision", self.hf_resolved_revision + ), + "cache_relpath": filename, + f"{self.config.modality}_path": file_path, + } + if self.config.modality == "video" and os.path.isdir(file_path): + from ..dataset.download import IMAGE_FILE_EXTENSIONS + + sample["video_frames"] = sorted( + str(p) for p in Path(file_path).iterdir() + if p.suffix.lower() in IMAGE_FILE_EXTENSIONS + ) + del sample["video_path"] + yield sample + continue + if self.config.modality == "image": try: if self.lazy_read: @@ -703,36 +738,15 @@ def extract_index(filename): elif self.config.modality == "audio": try: - # Check if this is a preprocessed tensor (.pt file) - if filename.endswith(".pt"): - # Load preprocessed tensor directly (much faster!) - import torch - - data = torch.load(file_path, map_location="cpu") - sample = { - "preprocessed_waveform": data["waveform"], - "label": data["label"], - "dataset_name": self.config.name, - "media_type": self.config.media_type, - "cached_filename": filename, - "is_preprocessed": True, - **data.get("metadata", {}), - **metadata, - } - yield sample - else: - # Load raw audio bytes (legacy/fallback) - with open(file_path, "rb") as f: - audio_bytes = f.read() - sample = { - "audio_bytes": audio_bytes, - "dataset_name": self.config.name, - "media_type": self.config.media_type, - "cached_filename": filename, # Include filename for format detection - "is_preprocessed": False, - **metadata, - } - yield sample + yield load_audio_sample({ + **metadata, + "audio_path": file_path, + "dataset_name": self.config.name, + "dataset_path": self.config.path, + "source_kind": self.source_kind, + "hf_resolved_revision": dataset_info.get("hf_resolved_revision", self.hf_resolved_revision), + "cache_relpath": filename, + }) except Exception as e: logger.warning(f"Failed to load cached audio {filename}: {e}") continue @@ -740,3 +754,21 @@ def extract_index(filename): except Exception as e: logger.error(f"Failed to load cached dataset from {cache_dir}: {e}") raise + + +def load_audio_sample(sample): + """Materialize cached audio after resume filtering, preserving cache provenance.""" + path = Path(sample["audio_path"]) + if path.suffix == ".pt": + import torch + + data = torch.load(path, map_location="cpu") + return { + **data.get("metadata", {}), **sample, + "preprocessed_waveform": data["waveform"], "label": data["label"], + "cached_filename": path.name, "is_preprocessed": True, + } + return { + **sample, "audio_bytes": path.read_bytes(), + "cached_filename": path.name, "is_preprocessed": False, + } diff --git a/tests/unit/test_benchmark_resume.py b/tests/unit/test_benchmark_resume.py new file mode 100644 index 0000000..e61f0b7 --- /dev/null +++ b/tests/unit/test_benchmark_resume.py @@ -0,0 +1,323 @@ +"""Exercise interruption recovery through real iterators, pipelines and scoring.""" + +import asyncio +import importlib +import json +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest + +from gasbench.benchmarks import common +from gasbench.benchmarks._checkpoint import CheckpointError, RecorderCheckpoint +from gasbench.dataset.config import BenchmarkDatasetConfig + + +class Interrupted(BaseException): + pass + + +class Session: + def __init__(self, model_dir, stop_after=None, error=False): + self.model_dir = model_dir + self.stop_after = stop_after + self.calls = [] + self.error = error + + def run(self, _, inputs): + if self.stop_after is not None and len(self.calls) >= self.stop_after: + raise Interrupted() + values = next(iter(inputs.values())) + self.calls.extend(int(value.flat[0]) for value in values) + if self.error: + raise RuntimeError("model failed") + return [np.tile([2.0, 0.0], (len(values), 1))] + + +@pytest.fixture(params=["image", "video", "audio", "audio-tensor"]) +def benchmark(request, tmp_path, monkeypatch): + modality = request.param.split("-")[0] + module = importlib.import_module(f"gasbench.benchmarks.{modality}_bench") + dataset = BenchmarkDatasetConfig("tiny", "test/repo", modality, "real") + cache = tmp_path / "cache" + dataset_dir = cache / "datasets" / dataset.name + samples_dir = dataset_dir / "samples" + samples_dir.mkdir(parents=True) + metadata = {} + for i in range(5): + name = f"sample_{i}.bin" + if request.param == "audio-tensor": + import torch + + name = f"sample_{i}.pt" + torch.save({ + "waveform": torch.full((8,), float(i)), "label": 0, + # Embedded legacy metadata must not change the frozen sample ID. + "metadata": {"path_in_archive": "different-embedded-path"}, + }, samples_dir / name) + else: + (samples_dir / name).write_bytes(bytes([i])) + metadata[name] = {"source_file": name, "member_path": name} + (dataset_dir / "sample_metadata.json").write_text(json.dumps(metadata)) + (dataset_dir / "dataset_info.json").write_text( + json.dumps({"hf_resolved_revision": "fixed-revision"}) + ) + model_dir = tmp_path / "model" + model_dir.mkdir() + (model_dir / "model.py").write_text("fixture model") + monkeypatch.setattr(common, "discover_benchmark_datasets", lambda **_: [dataset]) + monkeypatch.setattr( + common, "calculate_weighted_dataset_sampling", lambda *_: {dataset.name: 5} + ) + decoded = [] + + def decode(sample, **kwargs): + decoded.append(sample["source_file"]) + data = sample.get("image", sample.get("video_bytes", sample.get("audio_bytes"))) + shape = {"image": (2, 2, 3), "video": (2, 2, 2, 3), "audio": (8,)}[modality] + return np.full(shape, data[0], dtype=np.float32), 0 + + def augment(array, *args, **kwargs): + return array, None, None, None + + if modality == "image": + monkeypatch.setattr(module, "process_image_sample", decode) + monkeypatch.setattr(module, "apply_random_augmentations", augment) + monkeypatch.setattr(module, "apply_robustness_augmentations", augment) + elif modality == "video": + monkeypatch.setattr(module, "process_video_bytes_sample", decode) + monkeypatch.setattr(module, "apply_random_augmentations", augment) + monkeypatch.setattr(module, "apply_video_robustness_augmentations", augment) + else: + monkeypatch.setattr(module, "process_audio_sample", decode) + specs = [SimpleNamespace(name="input", shape=[None, 3, 2, 2], type="tensor(float)")] + + def run(session, run_id="run", **extra): + results = {"errors": []} + asyncio.run( + getattr(module, f"run_{modality}_benchmark")( + session, + specs, + results, + run_id=run_id, + cache_dir=str(cache), + batch_size=extra.pop("batch_size", 1), + skip_missing=True, + seed=7, + **extra, + ) + ) + return results + + return SimpleNamespace( + run=run, + model_dir=model_dir, + cache=cache, + samples_dir=samples_dir, + dataset=dataset, + decoded=decoded, + modality=modality, + module=module, + checkpoint=cache / "runs" / "run" / "checkpoint", + ) + + +def checkpoint_rows(directory): + manifest = RecorderCheckpoint.read_manifest(directory) + return RecorderCheckpoint(directory, manifest).records + + +@pytest.mark.parametrize("batch_size", [1, 2]) +def test_resume_skips_committed_inputs_and_preserves_results( + benchmark, tmp_path, batch_size +): + b = benchmark + uninterrupted = b.run( + Session(b.model_dir), + batch_size=batch_size, + run_id="control", + records_parquet_path=str(tmp_path / "control.parquet"), + ) + with pytest.raises(Interrupted): + b.run(Session(b.model_dir, stop_after=2), batch_size=batch_size) + committed = checkpoint_rows(b.checkpoint) + assert len(committed) == 2 + # Cached discovery changes and completed source files disappear. Resume must + # consume the saved plan, without loading the completed inputs again. + for row in committed: + (b.samples_dir / row["source_file"]).unlink() + (b.samples_dir / "new.bin").write_bytes(b"\xff") + b.decoded.clear() + resumed_session = Session(b.model_dir) + resumed = b.run( + resumed_session, + batch_size=batch_size, + records_parquet_path=str(tmp_path / "resumed.parquet"), + ) + assert len(resumed_session.calls) == 3 + assert not {row["source_file"] for row in committed}.intersection(b.decoded) + rows = checkpoint_rows(b.checkpoint) + assert len(rows) == 5 + keys = ["dataset_name", "iteration_index", "sample_id", "aug_pass"] + assert len({tuple(row[k] for k in keys) for row in rows}) == 5 + assert ( + resumed[f"{b.modality}_results"]["benchmark_score"] + == uninterrupted[f"{b.modality}_results"]["benchmark_score"] + ) + left, right = ( + pd.read_parquet(tmp_path / f"{name}.parquet") for name in ("control", "resumed") + ) + columns = keys + ["label", "predicted", "correct"] + pd.testing.assert_frame_equal( + left[columns].sort_values(keys).reset_index(drop=True), + right[columns].sort_values(keys).reset_index(drop=True), + ) + complete_session = Session(b.model_dir) + b.run(complete_session, batch_size=batch_size) + assert complete_session.calls == [] + + +def test_changed_pending_sample_aborts_instead_of_returning_partial_score(benchmark): + b = benchmark + with pytest.raises(Interrupted): + b.run(Session(b.model_dir, stop_after=0)) + next(b.samples_dir.iterdir()).write_bytes(b"changed") + with pytest.raises(CheckpointError, match="changed"): + b.run(Session(b.model_dir)) + + +def test_changed_model_rejected_before_inference(benchmark): + b = benchmark + with pytest.raises(Interrupted): + b.run(Session(b.model_dir, stop_after=1)) + (b.model_dir / "model.py").write_text("different model") + session = Session(b.model_dir) + with pytest.raises(CheckpointError, match="configuration or model"): + b.run(session) + assert session.calls == [] + + +@pytest.mark.parametrize("model_error", [False, True]) +def test_persistence_failure_aborts_every_modality(benchmark, model_error): + b = benchmark + session = Session(b.model_dir, error=model_error) + storage_error = OSError("storage unavailable") + + def persist(directory): + if list(Path(directory).glob("batch-*.json")): + raise storage_error + + with pytest.raises(CheckpointError) as error: + b.run(session, checkpoint_persist=persist) + assert error.value.__cause__ is storage_error + assert len(session.calls) == 1 + + +@pytest.mark.parametrize("benchmark", ["audio-tensor"], indirect=True) +def test_unreadable_audio_preserves_pending_batch_and_remaining_samples(benchmark): + b = benchmark + selected = list(common.DatasetIterator( + b.dataset, max_samples=5, cache_dir=str(b.cache), download=False, + seed=7, metadata_only=True, + )) + # Corrupt an input before freezing the plan, with a valid sample ahead of it + # waiting for its batch and more valid samples after it. + bad_sample = selected[1] + Path(bad_sample["audio_path"]).write_bytes(b"not a torch file") + expected = {sample["source_file"] for sample in selected} - {bad_sample["source_file"]} + session = Session(b.model_dir) + result = b.run(session, batch_size=2) + assert result["errors"] + assert len(session.calls) == len(expected) + assert {row["source_file"] for row in checkpoint_rows(b.checkpoint)} == expected + resumed_session = Session(b.model_dir) + b.run(resumed_session, batch_size=2) + assert resumed_session.calls == [] + + +def test_failed_inference_is_committed_and_not_retried(benchmark): + b = benchmark + b.run(Session(b.model_dir, error=True)) + assert all(row["status"] == "error" for row in checkpoint_rows(b.checkpoint)) + session = Session(b.model_dir) + b.run(session) + assert session.calls == [] + + +def test_augmentation_resumes_separately_from_base_pass(benchmark): + b = benchmark + if b.modality == "audio": + pytest.skip("Audio has no robustness pass") + with pytest.raises(Interrupted): + b.run(Session(b.model_dir, stop_after=6), n_aug_per_dataset=3) + rows = checkpoint_rows(b.checkpoint) + assert sum(row["aug_pass"] for row in rows) == 1 + session = Session(b.model_dir) + b.run(session, n_aug_per_dataset=3) + rows = checkpoint_rows(b.checkpoint) + assert len(session.calls) == 2 + assert sum(row["aug_pass"] for row in rows) == 3 + assert sum(not row["aug_pass"] for row in rows) == 5 + + +def test_top_level_api_checkpoints_by_default(benchmark, monkeypatch): + import gasbench.benchmark as driver + + b = benchmark + session = Session(b.model_dir) + + async def load_model(*args): + return session, [ + SimpleNamespace(name="input", shape=[None, 3, 2, 2], type="tensor(float)") + ] + + monkeypatch.setattr(driver, "load_model_for_benchmark", load_model) + for _ in range(2): + result = asyncio.run( + driver.run_benchmark( + str(b.model_dir), + b.modality, + cache_dir=str(b.cache), + run_id="api", + skip_missing=True, + batch_size=2, + ) + ) + assert result["benchmark_completed"] + assert len(session.calls) == 5 + assert len(checkpoint_rows(Path(result["checkpoint_dir"]))) == 5 + + +def test_augmentation_cache_cannot_change_across_attempts(benchmark, tmp_path): + from gasbench.benchmarks.aug_cache import img_aug_cache_path, vid_aug_cache_path + from gasbench.benchmarks.recording import build_sample_id + + b = benchmark + if b.modality == "audio": + pytest.skip("Audio has no robustness pass") + aug_dir = tmp_path / "augmentations" + cache_path = img_aug_cache_path if b.modality == "image" else vid_aug_cache_path + for path in b.samples_dir.iterdir(): + sample = { + "source_kind": "huggingface", + "dataset_path": "test/repo", + "source_file": path.name, + "member_path": path.name, + } + artifact = Path(cache_path(str(aug_dir), build_sample_id(sample), (2, 2))) + artifact.parent.mkdir(parents=True, exist_ok=True) + np.save( + artifact, np.zeros((2, 2, 3) if b.modality == "image" else (2, 2, 2, 3)) + ) + with pytest.raises(Interrupted): + b.run( + Session(b.model_dir, stop_after=5), + n_aug_per_dataset=3, + aug_cache_dir=str(aug_dir), + ) + for path in aug_dir.rglob("*.npy"): + path.write_bytes(b"changed") + with pytest.raises(CheckpointError, match="augmentation cache"): + b.run(Session(b.model_dir), n_aug_per_dataset=3, aug_cache_dir=str(aug_dir)) diff --git a/tests/unit/test_checkpoint.py b/tests/unit/test_checkpoint.py new file mode 100644 index 0000000..ff4f1ec --- /dev/null +++ b/tests/unit/test_checkpoint.py @@ -0,0 +1,238 @@ +"""Exercise crash recovery, replay safety, and persistence failures.""" + +import json +import os +import subprocess +import sys + +import pytest + +from gasbench.benchmarks._checkpoint import ( + CheckpointError, + RecorderCheckpoint, + prediction_key, +) + + +@pytest.fixture +def manifest(): + return { + "run_id": "run-a", + "model_hash": "model-a", + "benchmark_version": "test", + "seed": 7, + "sample_plan": ["source@revision/sample-a", "source@revision/sample-b"], + "settings": {"augmentation": True, "weight": 0.2}, + } + + +def record(sample_id, prediction=1): + sample_id, _, pass_name = sample_id.partition("/") + return { + "run_id": "run-a", + "dataset_name": "dataset", + "sample_id": sample_id, + "iteration_index": 1, + "aug_pass": pass_name == "aug", + "predicted": prediction, + "probs": [0.1, 0.9], + } + + +def test_resume_and_overlapping_replay_preserve_exact_records(tmp_path, manifest): + first, augmented, second = [record(key) for key in ("a/base", "a/aug", "b/base")] + store = RecorderCheckpoint(tmp_path, manifest) + store.commit_batch([first, augmented]) + resumed = RecorderCheckpoint(tmp_path, manifest) + assert resumed.completed_predictions == { + prediction_key(first), + prediction_key(augmented), + } + assert resumed.commit_batch([first, second]) == 1 + assert resumed.commit_batch([first, augmented, second]) == 0 + assert RecorderCheckpoint(tmp_path, manifest).records == [first, augmented, second] + + +def test_conflicting_replay_rejects_entire_batch(tmp_path, manifest): + store = RecorderCheckpoint(tmp_path, manifest) + store.commit_batch([record("a")]) + with pytest.raises(CheckpointError, match="Conflicting"): + store.commit_batch([record("b"), record("a", prediction=0)]) + assert RecorderCheckpoint(tmp_path, manifest).records == [record("a")] + + +@pytest.mark.parametrize( + "field", + ["run_id", "model_hash", "benchmark_version", "seed", "sample_plan", "settings"], +) +def test_changed_run_identity_fails_closed(tmp_path, manifest, field): + RecorderCheckpoint(tmp_path, manifest).commit_batch([record("a")]) + changed = {**manifest, field: "changed"} + with pytest.raises(CheckpointError, match="manifest differs"): + RecorderCheckpoint(tmp_path, changed) + assert RecorderCheckpoint(tmp_path, manifest).records == [record("a")] + + +def test_callers_cannot_mutate_checkpoint_state(tmp_path, manifest): + store = RecorderCheckpoint(tmp_path, manifest) + row = record("a") + store.commit_batch([row]) + row["probs"][0] = 100 + exposed = store.records + exposed[0]["probs"][0] = 200 + assert ( + store.records == RecorderCheckpoint(tmp_path, manifest).records == [record("a")] + ) + + +@pytest.mark.parametrize( + "damage", + [ + "truncated", + "checksum", + "missing_batch", + "missing_tail", + "head_checksum", + "missing_manifest", + ], +) +def test_committed_corruption_never_silently_restarts(tmp_path, manifest, damage): + store = RecorderCheckpoint(tmp_path, manifest) + store.commit_batch([record("a")]) + store.commit_batch([record("b")]) + first = sorted(tmp_path.glob("batch-*.json"))[0] + if damage == "truncated": + first.write_text('{"payload":') + elif damage == "checksum": + envelope = json.loads(first.read_text()) + envelope["payload"]["records"][0]["predicted"] = 0 + first.write_text(json.dumps(envelope)) + elif damage == "missing_batch": + first.unlink() + elif damage == "missing_tail": + sorted(tmp_path.glob("batch-*.json"))[-1].unlink() + elif damage == "head_checksum": + head_path = tmp_path / "head.json" + head = json.loads(head_path.read_text()) + head["payload"]["last_batch"] = -1 + head_path.write_text(json.dumps(head)) + else: + (tmp_path / "manifest.json").unlink() + with pytest.raises(CheckpointError): + RecorderCheckpoint(tmp_path, manifest) + + +def test_batch_from_another_manifest_is_rejected(tmp_path, manifest): + source = RecorderCheckpoint(tmp_path / "source", manifest) + source.commit_batch([record("a")]) + RecorderCheckpoint(tmp_path / "target", {**manifest, "seed": 8}) + batch = next(source.directory.glob("batch-*.json")) + (tmp_path / "target" / batch.name).write_bytes(batch.read_bytes()) + (tmp_path / "target" / "head.json").write_bytes( + (source.directory / "head.json").read_bytes() + ) + with pytest.raises(CheckpointError, match="Invalid checkpoint batch"): + RecorderCheckpoint(tmp_path / "target", {**manifest, "seed": 8}) + + +@pytest.mark.parametrize("crash_at", ["batch_rename", "head_rename", "after_commit"]) +def test_killed_process_recovers_only_complete_batches(tmp_path, manifest, crash_at): + RecorderCheckpoint(tmp_path, manifest).commit_batch([record("a")]) + program = """ +import json, os, sys +from pathlib import Path +from gasbench.benchmarks._checkpoint import RecorderCheckpoint +store = RecorderCheckpoint(Path(sys.argv[1]), json.loads(sys.argv[2])) +replace = os.replace +def crash_replace(source, target): + if sys.argv[3] == 'batch_rename' or (sys.argv[3] == 'head_rename' and target.name == 'head.json'): + os._exit(71) + replace(source, target) +os.replace = crash_replace +store.commit_batch([{'run_id': 'run-a', 'dataset_name': 'dataset', 'sample_id': 'b', 'iteration_index': 1, 'aug_pass': False, 'predicted': 1, 'probs': [0.1, 0.9]}]) +os._exit(71) +""" + result = subprocess.run( + [sys.executable, "-c", program, str(tmp_path), json.dumps(manifest), crash_at], + env=os.environ.copy(), + check=False, + ) + assert result.returncode == 71 + resumed = RecorderCheckpoint(tmp_path, manifest) + assert resumed.records == ( + [record("a")] if crash_at != "after_commit" else [record("a"), record("b")] + ) + resumed.commit_batch([record("b"), record("c")]) + assert RecorderCheckpoint(tmp_path, manifest).records == [ + record("a"), + record("b"), + record("c"), + ] + + +@pytest.mark.parametrize("failure", ["callback", "rename", "fsync"]) +def test_failed_commit_is_not_acknowledged(tmp_path, manifest, monkeypatch, failure): + calls = [] + storage_error = RuntimeError("remote storage unavailable") if failure == "callback" else OSError("local storage unavailable") + + def fail(*args): + raise storage_error + + def persist(directory): + calls.append(directory) + if failure == "callback" and len(calls) > 2: + fail() + + store = RecorderCheckpoint(tmp_path, manifest, persist=persist) + if failure == "rename": + monkeypatch.setattr(os, "replace", fail) + elif failure == "fsync": + monkeypatch.setattr(os, "fsync", fail) + with pytest.raises(CheckpointError, match="storage unavailable") as error: + store.commit_batch([record("a")]) + assert error.value.__cause__ is storage_error + assert store.completed_predictions == set() + with pytest.raises(CheckpointError, match="reopen"): + store.commit_batch([record("b")]) + + +def test_persistence_callback_sees_complete_checkpoint_before_ack(tmp_path, manifest): + snapshots = [] + + def persist(directory): + if (directory / "head.json").exists(): + snapshots.append(RecorderCheckpoint(directory, manifest).records) + + store = RecorderCheckpoint(tmp_path, manifest, persist=persist) + store.commit_batch([record("a")]) + assert snapshots == [[], [record("a")]] + + +@pytest.mark.parametrize( + "rows", + [ + [record("a"), record("a")], + [{**record("a"), "sample_id": ""}], + [{**record("a"), "iteration_index": 0}], + [{**record("a"), "run_id": "other"}], + [{**record("a"), "aug_pass": "false"}], + [None], + [], + ], +) +def test_invalid_batch_never_changes_progress(tmp_path, manifest, rows): + store = RecorderCheckpoint(tmp_path, manifest) + with pytest.raises(CheckpointError): + store.commit_batch(rows) + assert RecorderCheckpoint(tmp_path, manifest).records == [] + + +def test_saved_manifest_is_checked_before_any_predictions_exist(tmp_path, manifest): + RecorderCheckpoint(tmp_path, manifest) + assert RecorderCheckpoint.read_manifest(tmp_path) == manifest + path = tmp_path / "manifest.json" + saved = json.loads(path.read_text()) + saved["manifest"]["seed"] = "corrupted" + path.write_text(json.dumps(saved)) + with pytest.raises(CheckpointError, match="manifest checksum"): + RecorderCheckpoint.read_manifest(tmp_path) diff --git a/tests/unit/test_recorder_checkpoint.py b/tests/unit/test_recorder_checkpoint.py new file mode 100644 index 0000000..81001ee --- /dev/null +++ b/tests/unit/test_recorder_checkpoint.py @@ -0,0 +1,227 @@ +"""Recovery uses real recorder rows, summaries, scoring, and parquet output.""" + +import pandas as pd +import pytest + +from gasbench.benchmarks.recording import ( + BenchmarkRunRecorder, + compute_metrics_from_df, + compute_per_dataset_from_df, +) +from gasbench.benchmarks._checkpoint import CheckpointError + + +@pytest.fixture +def context(): + return { + "model_hash": "model", + "benchmark_version": "test", + "seed": 42, + "sample_plan": ["dataset@rev/a", "dataset@rev/b"], + "score_composition": {"public": 1.0}, + } + + +def sample(name): + return { + "source_kind": "hf", + "dataset_path": "repo", + "source_file": name, + "hf_resolved_revision": "revision", + "media_type": "real", + } + + +def add_prediction( + recorder, name, *, index=1, dataset="dataset", aug_pass=False, label=0 +): + recorder.add_ok( + dataset_name=dataset, + sample_index=index, + sample=sample(name), + label=label, + predicted=label, + probs=[0.8, 0.1, 0.1] if label == 0 else [0.1, 0.8, 0.1], + inference_time_ms=2.0, + batch_inference_time_ms=2.0, + batch_id=1, + batch_size=1, + sample_seed=42 + index, + aug_pass=aug_pass, + ) + + +def test_resumed_recorder_uses_existing_metrics_summaries_and_parquet( + tmp_path, context, monkeypatch +): + monkeypatch.setattr("gasbench.benchmarks.recording.time.time", lambda: 100) + original = BenchmarkRunRecorder(run_id="run", modality="image") + interrupted = BenchmarkRunRecorder( + run_id="run", + modality="image", + checkpoint_dir=tmp_path / "checkpoint", + checkpoint_context=context, + ) + for recorder in (original, interrupted): + add_prediction(recorder, "a") + add_prediction(recorder, "a", aug_pass=True) + recorder.add_skip( + dataset_name="dataset", sample_index=3, sample=sample("c"), reason="decode" + ) + recorder.add_error( + dataset_name="dataset", + sample_index=4, + sample=sample("d"), + error_message="inference", + ) + assert interrupted.checkpoint() == 4 + # Later wall-clock time must not create a second run/start timestamp. + monkeypatch.setattr("gasbench.benchmarks.recording.time.time", lambda: 200) + resumed = BenchmarkRunRecorder( + run_id="run", + modality="image", + checkpoint_dir=tmp_path / "checkpoint", + checkpoint_context=context, + ) + assert resumed.run_started_at == original.run_started_at + for recorder in (original, resumed): + add_prediction(recorder, "b", index=2, label=1) + assert resumed.checkpoint() == 1 + assert resumed.checkpoint() == 0 + pd.testing.assert_frame_equal(original.to_dataframe(), resumed.to_dataframe()) + assert resumed.get_dataset_summary("dataset", True) == original.get_dataset_summary( + "dataset", True + ) + assert compute_metrics_from_df(resumed.to_dataframe()) == compute_metrics_from_df( + original.to_dataframe() + ) + assert compute_per_dataset_from_df( + resumed.to_dataframe() + ) == compute_per_dataset_from_df(original.to_dataframe()) + resumed.write_parquet(str(tmp_path / "resumed.parquet")) + original.write_parquet(str(tmp_path / "original.parquet")) + pd.testing.assert_frame_equal( + pd.read_parquet(tmp_path / "resumed.parquet"), + pd.read_parquet(tmp_path / "original.parquet"), + ) + + +def test_existing_prediction_fields_distinguish_datasets_indices_and_passes( + tmp_path, context +): + recorder = BenchmarkRunRecorder( + run_id="run", checkpoint_dir=tmp_path, checkpoint_context=context + ) + variants = [dict(), {"dataset": "other"}, {"index": 2}, {"aug_pass": True}] + for variant in variants: + add_prediction(recorder, "same", **variant) + assert not recorder.is_checkpointed( + dataset_name="dataset", sample_index=1, sample=sample("same") + ) + assert recorder.checkpoint() == 4 + resumed = BenchmarkRunRecorder( + run_id="run", checkpoint_dir=tmp_path, checkpoint_context=context + ) + for variant in variants: + assert resumed.is_checkpointed( + dataset_name=variant.get("dataset", "dataset"), + sample_index=variant.get("index", 1), + sample=sample("same"), + aug_pass=variant.get("aug_pass", False), + ) + assert not resumed.is_checkpointed( + dataset_name="dataset", sample_index=3, sample=sample("same") + ) + + +def test_restore_uses_same_counters_for_augmented_skips_and_errors(tmp_path, context): + recorder = BenchmarkRunRecorder( + run_id="run", checkpoint_dir=tmp_path, checkpoint_context=context + ) + add_prediction(recorder, "a") + recorder.add_error( + dataset_name="dataset", + sample_index=1, + sample=sample("a"), + error_message="aug-failed", + aug_pass=True, + ) + recorder.add_skip( + dataset_name="dataset", + sample_index=2, + sample=sample("b"), + reason="aug-decode", + aug_pass=True, + ) + recorder.checkpoint() + resumed = BenchmarkRunRecorder( + run_id="run", checkpoint_dir=tmp_path, checkpoint_context=context + ) + assert resumed.get_dataset_summary("dataset", True) == { + "accuracy": 1.0, + "correct": 1, + "total": 1, + "skipped": 0, + } + assert resumed.rows == recorder.rows + assert resumed.is_checkpointed( + dataset_name="dataset", sample_index=1, sample=sample("a"), aug_pass=True + ) + + +@pytest.mark.parametrize( + "changes", + [ + {"run_id": "other"}, + {"modality": "video"}, + {"target_size": (48, 48)}, + {"crop_prob": 0.5}, + ], +) +def test_recorder_identity_cannot_drift_on_resume(tmp_path, context, changes): + kwargs = dict( + run_id="run", + modality="image", + checkpoint_dir=tmp_path, + checkpoint_context=context, + ) + recorder = BenchmarkRunRecorder(**kwargs) + add_prediction(recorder, "a") + recorder.checkpoint() + with pytest.raises(CheckpointError, match="manifest differs"): + BenchmarkRunRecorder(**{**kwargs, **changes}) + + +def test_checkpoint_requires_explicit_run_and_benchmark_context(tmp_path, context): + with pytest.raises(ValueError, match="requires run_id"): + BenchmarkRunRecorder(checkpoint_dir=tmp_path, checkpoint_context=context) + with pytest.raises(ValueError, match="requires run_id"): + BenchmarkRunRecorder(run_id="run", checkpoint_dir=tmp_path) + + +def test_plain_recorder_still_records_without_checkpoint_configuration(): + recorder = BenchmarkRunRecorder() + add_prediction(recorder, "a") + assert recorder.count == 1 + assert recorder.checkpoint() == 0 + assert recorder.get_dataset_summary("dataset")["total"] == 1 + + +def test_archive_path_identity_matches_resume_lookup_for_every_status(tmp_path, context): + recorder = BenchmarkRunRecorder( + run_id="run", checkpoint_dir=tmp_path, checkpoint_context=context, + ) + source = {**sample("a"), "path_in_archive": "canonical", "member_path": "legacy"} + recorder.add_ok( + dataset_name="dataset", sample_index=1, sample=source, label=0, + predicted=0, probs=[0.9, 0.1], inference_time_ms=1, + batch_inference_time_ms=1, batch_id=1, batch_size=1, sample_seed=42, + ) + recorder.add_skip(dataset_name="dataset", sample_index=2, sample=source, reason="decode") + recorder.add_error(dataset_name="dataset", sample_index=3, sample=source, error_message="model") + recorder.checkpoint() + resumed = BenchmarkRunRecorder( + run_id="run", checkpoint_dir=tmp_path, checkpoint_context=context, + ) + for index in (1, 2, 3): + assert resumed.is_checkpointed(dataset_name="dataset", sample_index=index, sample=source)