aituner/scripts/collectivespec/create_controlled_greedy_window.py

#!/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())