diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7e2a683..d7e034e 100644 --- a/.github/workflows/ci.yml +++ b/.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 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/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/_checkpoint.py b/src/gasbench/benchmarks/_checkpoint.py index f00f40c..c6657e1 100644 --- a/src/gasbench/benchmarks/_checkpoint.py +++ b/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) 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 516371b..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: @@ -67,6 +79,7 @@ def run_batch_and_record( error_message=f"inference-failed: {str(e)[:160]}", aug_pass=aug_pass, ) + tracker.checkpoint() return batch_inference_time = (time.time() - start) * 1000 @@ -75,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, @@ -91,6 +104,8 @@ def run_batch_and_record( aug_pass=aug_pass, ) + tracker.checkpoint() + @dataclass class BenchmarkRunConfig: @@ -106,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) @@ -113,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) @@ -135,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( @@ -239,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, @@ -250,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( @@ -263,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 db41605..c8357a3 100644 --- a/src/gasbench/benchmarks/recording.py +++ b/src/gasbench/benchmarks/recording.py @@ -172,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"), @@ -233,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"), @@ -291,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"), 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 index fae7687..ff4f1ec 100644 --- a/tests/unit/test_checkpoint.py +++ b/tests/unit/test_checkpoint.py @@ -170,17 +170,27 @@ def crash_replace(source, target): ] -def test_failed_remote_commit_is_not_acknowledged(tmp_path, manifest): +@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 len(calls) > 2: - raise OSError("remote storage unavailable") + if failure == "callback" and len(calls) > 2: + fail() store = RecorderCheckpoint(tmp_path, manifest, persist=persist) - with pytest.raises(OSError, match="remote storage"): + 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")]) @@ -215,3 +225,14 @@ def test_invalid_batch_never_changes_progress(tmp_path, manifest, rows): 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 index 700bd1c..81001ee 100644 --- a/tests/unit/test_recorder_checkpoint.py +++ b/tests/unit/test_recorder_checkpoint.py @@ -205,3 +205,23 @@ def test_plain_recorder_still_records_without_checkpoint_configuration(): 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)