Add reproducible CollectiveSpec opportunity screen
This commit is contained in:
175
scripts/collectivespec/create_controlled_greedy_window.py
Normal file
175
scripts/collectivespec/create_controlled_greedy_window.py
Normal file
@@ -0,0 +1,175 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user