#!/usr/bin/env python3 """Validate ReplayServe fixture directories.""" from __future__ import annotations import argparse import csv import json import math import sys from pathlib import Path from typing import Any def positive_int(value: str) -> int: parsed = int(value) if parsed <= 0: raise argparse.ArgumentTypeError("must be positive") return parsed def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Validate ReplayServe fixtures.") parser.add_argument("fixture_dirs", nargs="+", type=Path) parser.add_argument("--max-tokens", type=positive_int, default=32768) parser.add_argument("--block-size", type=positive_int, default=16) return parser.parse_args() def parse_block_hash_ids(value: str) -> list[int]: stripped = value.strip() if not stripped: return [] return [int(part) for part in stripped.split("|") if part] def expected_block_counts(input_length: int, block_size: int) -> list[int]: hash_count = math.ceil(input_length / block_size) if hash_count == 0: return [] last_count = input_length % block_size if last_count == 0: last_count = block_size return [block_size] * (hash_count - 1) + [last_count] def load_jsonl(path: Path) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] with path.open("r", encoding="utf-8") as handle: for line_number, line in enumerate(handle, start=1): stripped = line.strip() if not stripped: continue try: row = json.loads(stripped) except json.JSONDecodeError as exc: raise ValueError(f"{path}: line {line_number}: invalid JSON") from exc if not isinstance(row, dict): raise ValueError(f"{path}: line {line_number}: JSON value must be object") rows.append(row) return rows def load_csv(path: Path) -> list[dict[str, str]]: with path.open("r", encoding="utf-8", newline="") as handle: reader = csv.DictReader(handle) required = { "arrived_at", "num_prefill_tokens", "num_decode_tokens", "session_id", "block_hash_ids", } missing = required - set(reader.fieldnames or []) if missing: raise ValueError(f"{path}: missing CSV columns: {sorted(missing)}") return list(reader) def require_paths(fixture_dir: Path) -> tuple[Path, Path, Path, Path]: source_path = fixture_dir / "source.jsonl" csv_path = fixture_dir / "frontier.csv" sidecar_path = fixture_dir / "sidecar.jsonl" manifest_path = fixture_dir / "manifest.json" for path in (source_path, csv_path, sidecar_path, manifest_path): if not path.exists(): raise ValueError(f"{fixture_dir}: missing {path.name}") return source_path, csv_path, sidecar_path, manifest_path def validate_fixture(fixture_dir: Path, block_size: int, max_tokens: int) -> str: source_path, csv_path, sidecar_path, manifest_path = require_paths(fixture_dir) source_rows = load_jsonl(source_path) csv_rows = load_csv(csv_path) sidecar_rows = load_jsonl(sidecar_path) with manifest_path.open("r", encoding="utf-8") as handle: manifest = json.load(handle) row_count = len(csv_rows) if len(source_rows) != row_count or len(sidecar_rows) != row_count: raise ValueError( f"{fixture_dir}: row count mismatch source={len(source_rows)} " f"csv={row_count} sidecar={len(sidecar_rows)}" ) if manifest.get("row_count") != row_count: raise ValueError( f"{fixture_dir}: manifest row_count={manifest.get('row_count')} " f"does not match csv rows={row_count}" ) if manifest.get("block_size") != block_size: raise ValueError( f"{fixture_dir}: manifest block_size={manifest.get('block_size')} " f"does not match expected {block_size}" ) if manifest.get("max_tokens") != max_tokens: raise ValueError( f"{fixture_dir}: manifest max_tokens={manifest.get('max_tokens')} " f"does not match expected {max_tokens}" ) previous_timestamp: float | None = None max_total_tokens = 0 partial_final_block_rows = 0 for index, (source, csv_row, sidecar) in enumerate( zip(source_rows, csv_rows, sidecar_rows) ): prefix = f"{fixture_dir}: row {index}" input_length = int(csv_row["num_prefill_tokens"]) output_length = int(csv_row["num_decode_tokens"]) total_tokens = input_length + output_length if total_tokens > max_tokens: raise ValueError( f"{prefix}: total_tokens={total_tokens} exceeds max_tokens={max_tokens}" ) max_total_tokens = max(max_total_tokens, total_tokens) timestamp = float(csv_row["arrived_at"]) if previous_timestamp is not None and timestamp < previous_timestamp: raise ValueError(f"{prefix}: timestamp is not monotonic") previous_timestamp = timestamp hash_ids = parse_block_hash_ids(csv_row["block_hash_ids"]) expected_hash_count = math.ceil(input_length / block_size) if len(hash_ids) != expected_hash_count: raise ValueError( f"{prefix}: hash count {len(hash_ids)} != {expected_hash_count}" ) counts = expected_block_counts(input_length, block_size) if sum(counts) != input_length: raise ValueError(f"{prefix}: expected block counts do not sum to input") partial_final_block_rows += int(input_length % block_size != 0) if int(csv_row["session_id"]) != int(source["chat_id"]): raise ValueError(f"{prefix}: session_id does not match source chat_id") if timestamp != float(source["timestamp"]): raise ValueError(f"{prefix}: arrived_at does not match source timestamp") if input_length != int(source["input_length"]): raise ValueError(f"{prefix}: num_prefill_tokens does not match source") if output_length != int(source["output_length"]): raise ValueError(f"{prefix}: num_decode_tokens does not match source") if hash_ids != source["hash_ids"]: raise ValueError(f"{prefix}: block_hash_ids do not match source hash_ids") required_sidecar_keys = { "request_id", "chat_id", "parent_chat_id", "turn", "type", "timestamp", "input_length", "output_length", "hash_ids", "block_token_counts", } missing = required_sidecar_keys - set(sidecar) if missing: raise ValueError(f"{prefix}: missing sidecar keys {sorted(missing)}") if int(sidecar["request_id"]) != index: raise ValueError(f"{prefix}: sidecar request_id mismatch") if int(sidecar["chat_id"]) != int(source["chat_id"]): raise ValueError(f"{prefix}: sidecar chat_id mismatch") if int(sidecar["parent_chat_id"]) != int(source["parent_chat_id"]): raise ValueError(f"{prefix}: sidecar parent_chat_id mismatch") if int(sidecar["turn"]) != int(source["turn"]): raise ValueError(f"{prefix}: sidecar turn mismatch") if sidecar["type"] != source["type"]: raise ValueError(f"{prefix}: sidecar type mismatch") if float(sidecar["timestamp"]) != float(source["timestamp"]): raise ValueError(f"{prefix}: sidecar timestamp mismatch") if int(sidecar["input_length"]) != input_length: raise ValueError(f"{prefix}: sidecar input_length mismatch") if int(sidecar["output_length"]) != output_length: raise ValueError(f"{prefix}: sidecar output_length mismatch") if sidecar["hash_ids"] != hash_ids: raise ValueError(f"{prefix}: sidecar hash_ids mismatch") if sidecar["block_token_counts"] != counts: raise ValueError(f"{prefix}: sidecar block_token_counts mismatch") if manifest.get("max_total_tokens") != max_total_tokens: raise ValueError( f"{fixture_dir}: manifest max_total_tokens=" f"{manifest.get('max_total_tokens')} does not match {max_total_tokens}" ) if manifest.get("partial_final_block_rows") != partial_final_block_rows: raise ValueError( f"{fixture_dir}: manifest partial_final_block_rows=" f"{manifest.get('partial_final_block_rows')} does not match " f"{partial_final_block_rows}" ) if manifest.get("overflow_count") != 0: raise ValueError(f"{fixture_dir}: manifest overflow_count is not zero") if manifest.get("timestamp_monotonic") is not True: raise ValueError(f"{fixture_dir}: manifest timestamp_monotonic is not true") return ( f"{fixture_dir.name}: rows={row_count} max_total_tokens={max_total_tokens} " f"partial_final_block_rows={partial_final_block_rows}" ) def main() -> int: args = parse_args() try: for fixture_dir in args.fixture_dirs: print( validate_fixture( fixture_dir=fixture_dir, block_size=args.block_size, max_tokens=args.max_tokens, ) ) except Exception as exc: print(f"validate_fixtures.py: error: {exc}", file=sys.stderr) return 1 return 0 if __name__ == "__main__": raise SystemExit(main())