237 lines
9.3 KiB
Python
237 lines
9.3 KiB
Python
#!/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 stable_session_key(value: Any) -> str:
|
|
"""Render a session identifier without conflating JSON types."""
|
|
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
|
|
|
|
|
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}")
|
|
source_sampling_strategy = source_window.get("sampling_strategy")
|
|
session_coherent_source = source_sampling_strategy == "session_coherent_uniform_score"
|
|
|
|
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] = []
|
|
session_counts: dict[str, dict[str, int]] = {}
|
|
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)
|
|
session_entry: dict[str, int] | None = None
|
|
if session_coherent_source:
|
|
if "session_root" not in row:
|
|
raise SystemExit(
|
|
f"row {row_index}: session-coherent source has no session_root"
|
|
)
|
|
session_key = stable_session_key(row["session_root"])
|
|
session_entry = session_counts.setdefault(
|
|
session_key, {"total_turns": 0, "selected_turns": 0}
|
|
)
|
|
session_entry["total_turns"] += 1
|
|
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 session_entry is not None:
|
|
session_entry["selected_turns"] += 1
|
|
if selected == 0:
|
|
output_trace.unlink(missing_ok=True)
|
|
raise SystemExit("threshold selected zero requests")
|
|
split_sessions = [
|
|
entry
|
|
for entry in session_counts.values()
|
|
if 0 < entry["selected_turns"] < entry["total_turns"]
|
|
]
|
|
if split_sessions:
|
|
output_trace.unlink(missing_ok=True)
|
|
raise SystemExit(
|
|
"session-coherent source produced a partial selected session; "
|
|
f"split_session_count={len(split_sessions)}"
|
|
)
|
|
selected_session_turn_counts = [
|
|
entry["selected_turns"]
|
|
for entry in session_counts.values()
|
|
if entry["selected_turns"] > 0
|
|
]
|
|
|
|
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": (
|
|
"session_coherent_uniform_score_then_greedy_temperature"
|
|
if session_coherent_source
|
|
else "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),
|
|
"collectivespec_source_sampling_strategy": source_sampling_strategy,
|
|
"collectivespec_session_closed_within_window": session_coherent_source,
|
|
}
|
|
)
|
|
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},
|
|
"session": {
|
|
"source_is_session_coherent": session_coherent_source,
|
|
"session_root_count": len(session_counts) if session_coherent_source else None,
|
|
"selected_session_count": (
|
|
len(selected_session_turn_counts) if session_coherent_source else None
|
|
),
|
|
"split_session_count": len(split_sessions) if session_coherent_source else None,
|
|
"all_selected_or_dropped": not split_sessions if session_coherent_source else None,
|
|
"selected_turns": {
|
|
"min": min(selected_session_turn_counts)
|
|
if selected_session_turn_counts
|
|
else None,
|
|
"max": max(selected_session_turn_counts)
|
|
if selected_session_turn_counts
|
|
else None,
|
|
"distinct_value_count": len(set(selected_session_turn_counts)),
|
|
},
|
|
},
|
|
"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())
|