#!/usr/bin/env python3 """Materialize a fixed-load, greedy subset of an immutable replay window. The output preserves prompt/arrival/sampling provenance while setting every selected request's temperature to zero. It is intentionally a controlled mechanism workload, not a replacement for a production sampling trace. """ from __future__ import annotations import argparse import hashlib import json import math from pathlib import Path from typing import Any def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def resolve_trace_path(windows_path: Path, trace_file: str) -> Path: candidate = Path(trace_file) if candidate.is_absolute(): return candidate direct = (windows_path.parent / candidate).resolve() if direct.exists(): return direct if candidate.parts and candidate.parts[0] == "trace_windows": return (windows_path.parent / Path(*candidate.parts[1:])).resolve() return direct def numeric(value: Any, *, field: str, row_index: int) -> float: if isinstance(value, bool) or not isinstance(value, (int, float)): raise SystemExit(f"row {row_index}: {field} must be numeric") value = float(value) if not math.isfinite(value): raise SystemExit(f"row {row_index}: {field} must be finite") return value def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--windows-path", type=Path, required=True) parser.add_argument("--window-id", required=True) parser.add_argument("--sampling-u-threshold", type=float, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--output-window-id", required=True) args = parser.parse_args() if not args.windows_path.is_file(): raise SystemExit(f"windows file does not exist: {args.windows_path}") threshold = float(args.sampling_u_threshold) if not math.isfinite(threshold) or not 0.0 <= threshold <= 1.0: raise SystemExit("--sampling-u-threshold must be finite and in [0, 1]") source_windows = json.loads(args.windows_path.read_text(encoding="utf-8")) windows = source_windows.get("windows") if isinstance(source_windows, dict) else source_windows if not isinstance(windows, list): raise SystemExit("windows payload must contain a list") source_window = next( ( item for item in windows if isinstance(item, dict) and item.get("window_id") == args.window_id ), None, ) if source_window is None: raise SystemExit(f"window id not found: {args.window_id}") trace_file = source_window.get("trace_file") if not isinstance(trace_file, str) or not trace_file: raise SystemExit(f"window {args.window_id} has no trace_file") source_trace = resolve_trace_path(args.windows_path, trace_file) if not source_trace.is_file(): raise SystemExit(f"source trace does not exist: {source_trace}") args.output_dir.mkdir(parents=True, exist_ok=True) output_trace = args.output_dir / f"{args.output_window_id}.jsonl" selected = 0 sampling_values: list[float] = [] timestamps: list[float] = [] with source_trace.open("r", encoding="utf-8") as src, output_trace.open( "w", encoding="utf-8" ) as dest: for row_index, line in enumerate(src): if not line.strip(): continue row = json.loads(line) if not isinstance(row, dict): raise SystemExit(f"row {row_index}: expected an object") sampling_u = numeric(row.get("sampling_u"), field="sampling_u", row_index=row_index) if sampling_u > threshold: continue timestamp = numeric(row.get("timestamp"), field="timestamp", row_index=row_index) row["temperature"] = 0.0 dest.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n") selected += 1 sampling_values.append(sampling_u) timestamps.append(timestamp) if selected == 0: output_trace.unlink(missing_ok=True) raise SystemExit("threshold selected zero requests") output_window = dict(source_window) output_window.update( { "window_id": args.output_window_id, "trace_file": str(output_trace.resolve()), "num_requests": selected, "sampling_strategy": "fixed_uniform_score_then_greedy_temperature", "collectivespec_source_window_id": args.window_id, "collectivespec_sampling_u_threshold": threshold, "collectivespec_temperature_override": 0.0, "collectivespec_source_trace_path": str(source_trace), "collectivespec_source_trace_sha256": sha256(source_trace), } ) output_windows = { "kind": "collectivespec_controlled_greedy_window", "source_windows_path": str(args.windows_path.resolve()), "source_window_id": args.window_id, "sampling_u_threshold": threshold, "temperature_override": 0.0, "windows": [output_window], } output_manifest = args.output_dir / "windows.json" output_manifest.write_text( json.dumps(output_windows, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8", ) sanity = { "n": selected, "sampling_u": { "min": min(sampling_values), "max": max(sampling_values), "distinct_value_count": len(set(sampling_values)), "all_at_or_below_threshold": all(value <= threshold for value in sampling_values), }, "timestamp": { "min": min(timestamps), "max": max(timestamps), "non_negative": all(value >= 0 for value in timestamps), "nondecreasing": all( timestamps[index] >= timestamps[index - 1] for index in range(1, len(timestamps)) ), }, "temperature": {"distinct_value_count": 1, "all_zero": True}, "source_trace_sha256": sha256(source_trace), "output_trace_sha256": sha256(output_trace), } print( json.dumps( { "windows_path": str(output_manifest), "trace_path": str(output_trace), "data_sanity": sanity, }, ensure_ascii=False, indent=2, sort_keys=True, ) ) return 0 if __name__ == "__main__": raise SystemExit(main())