diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 148925c..7e2a683 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,10 @@ 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 typecheck: name: mypy (advisory) 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..f00f40c --- /dev/null +++ b/src/gasbench/benchmarks/_checkpoint.py @@ -0,0 +1,266 @@ +"""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. + """ + + 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": 1, + "manifest": json.loads(_encode(dict(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. + 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) + try: + os.fsync(fd) + finally: + os.close(fd) + if persist and self._persist is not None: + self._persist(self.directory) + except BaseException: + self._usable = False + raise + finally: + if temp_path is not None: + temp_path.unlink(missing_ok=True) diff --git a/src/gasbench/benchmarks/common.py b/src/gasbench/benchmarks/common.py index 1719988..516371b 100644 --- a/src/gasbench/benchmarks/common.py +++ b/src/gasbench/benchmarks/common.py @@ -65,6 +65,7 @@ def run_batch_and_record( sample_index=sample_index, sample=sample, error_message=f"inference-failed: {str(e)[:160]}", + aug_pass=aug_pass, ) return diff --git a/src/gasbench/benchmarks/recording.py b/src/gasbench/benchmarks/recording.py index 9be7742..db41605 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: @@ -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, @@ -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, @@ -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/tests/unit/test_checkpoint.py b/tests/unit/test_checkpoint.py new file mode 100644 index 0000000..fae7687 --- /dev/null +++ b/tests/unit/test_checkpoint.py @@ -0,0 +1,217 @@ +"""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"), + ] + + +def test_failed_remote_commit_is_not_acknowledged(tmp_path, manifest): + calls = [] + + def persist(directory): + calls.append(directory) + if len(calls) > 2: + raise OSError("remote storage unavailable") + + store = RecorderCheckpoint(tmp_path, manifest, persist=persist) + with pytest.raises(OSError, match="remote storage"): + store.commit_batch([record("a")]) + 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 == [] diff --git a/tests/unit/test_recorder_checkpoint.py b/tests/unit/test_recorder_checkpoint.py new file mode 100644 index 0000000..700bd1c --- /dev/null +++ b/tests/unit/test_recorder_checkpoint.py @@ -0,0 +1,207 @@ +"""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