#!/usr/bin/env python3 """Phase-5 private-manifest transforms and timestamp-scheduled P3 client wrapper.""" from __future__ import annotations import argparse import asyncio import json import math import sys from pathlib import Path from typing import Any import aiohttp import opprof_phase3_client as p3 def _numeric(values: list[float | int]) -> dict[str, Any]: finite = [float(value) for value in values if math.isfinite(float(value))] return { "n": len(values), "finite_n": len(finite), "missing_n": len(values) - len(finite), "min": min(finite) if finite else None, "max": max(finite) if finite else None, "distinct_n": len(set(finite)), "sum": sum(finite), } def _r16(rows: list[dict[str, Any]]) -> float: groups = [rows[index : index + 16] for index in range(0, len(rows) - 15, 16)] useful = sum(sum(int(row["input_tokens"]) for row in group) for group in groups) rectangular = sum(16 * max(int(row["input_tokens"]) for row in group) for group in groups) return 1.0 - useful / rectangular def _source_timestamps(path: Path, indices: set[int], field: str) -> dict[int, float]: result: dict[int, float] = {} maximum = max(indices) with path.open(encoding="utf-8") as source: for index, line in enumerate(source): if index in indices: value = float(json.loads(line)[field]) if not math.isfinite(value): raise ValueError(f"non-finite source timestamp at {index}") result[index] = value if index >= maximum: break if set(result) != indices: raise ValueError("timestamp source did not cover all source_index values") return result def transform(args: argparse.Namespace) -> dict[str, Any]: rows = p3.load_manifest(Path(args.input))[: args.take_first] if len(rows) != args.take_first: raise ValueError(f"requested {args.take_first} rows, found {len(rows)}") indices = [int(row[args.join_key]) for row in rows] timestamps = _source_timestamps( Path(args.timestamp_source), set(indices), args.timestamp_field ) source_times = [timestamps[index] for index in indices] if any(right < left for left, right in zip(source_times, source_times[1:])): raise ValueError("selected timestamps are not nondecreasing") if source_times[-1] <= source_times[0]: raise ValueError("selected timestamp span is not positive") end_s = (len(rows) - 1) / args.target_rate scale = end_s / (source_times[-1] - source_times[0]) recorded_slots = [(value - source_times[0]) * scale for value in source_times] uniform_slots = [index / args.target_rate for index in range(len(rows))] slots = recorded_slots if args.arrival == "recorded-scaled" else uniform_slots for index, row in enumerate(rows): row["arrival"] = args.arrival row["arrival_s"] = slots[index] row["original_index"] = index row["source_timestamp"] = source_times[index] original = list(rows) max_added_delay = 0.0 if args.service_order == "length-binned": edges = [int(item) for item in args.length_bin_edges.split(",")] def bin_id(row: dict[str, Any]) -> int: length = int(row["input_tokens"]) for index, edge in enumerate(edges): if length <= edge: return index raise ValueError(f"input length {length} exceeds final edge") reordered: list[dict[str, Any]] = [] for offset in range(0, len(rows), args.reorder_block_size): block = rows[offset : offset + args.reorder_block_size] ordered = sorted( block, key=lambda row: ( bin_id(row), int(row["input_tokens"]), int(row["original_index"]), ), ) block_slots = slots[offset : offset + len(block)] for position, row in enumerate(ordered): added = max(0.0, block_slots[position] - float(row["arrival_s"])) max_added_delay = max(max_added_delay, added) row["arrival_s"] = block_slots[position] reordered.extend(ordered) rows = reordered if max_added_delay > args.max_added_delay_seconds + 1e-9: raise ValueError( f"fairness cap exceeded: {max_added_delay} > {args.max_added_delay_seconds}" ) elif args.service_order != "original": raise ValueError(f"unsupported service order: {args.service_order}") if sorted(row["request_id"] for row in rows) != sorted( row["request_id"] for row in original ): raise AssertionError("request identity changed") for key in ("input_tokens", "output_tokens"): if sum(int(row[key]) for row in rows) != sum(int(row[key]) for row in original): raise AssertionError(f"{key} total changed") arrival_values = [float(row["arrival_s"]) for row in rows] if any(right < left for left, right in zip(arrival_values, arrival_values[1:])): raise AssertionError("assigned arrival slots are not nondecreasing") output = Path(args.out) p3.atomic_jsonl(output, rows, mode=0o600) summary = { "schema": 1, "path": str(output), "sha256": p3.sha256_file(output), "rows": len(rows), "arrival": args.arrival, "service_order": args.service_order, "target_rate": args.target_rate, "input_tokens": _numeric([int(row["input_tokens"]) for row in rows]), "output_tokens": _numeric([int(row["output_tokens"]) for row in rows]), "arrival_s": _numeric(arrival_values), "source_timestamp": _numeric(source_times), "r16": _r16(rows), "max_added_delay_seconds": max_added_delay, "request_id_set_sha256": p3.hashlib.sha256( "\n".join(sorted(str(row["request_id"]) for row in rows)).encode() ).hexdigest(), "invariants": { "same_request_ids": True, "same_input_tokens": True, "same_output_tokens": True, "arrival_nondecreasing": True, "fairness_cap": max_added_delay <= args.max_added_delay_seconds + 1e-9, "no_prompt_in_summary": True, }, } p3.atomic_json(output.with_suffix(output.suffix + ".summary.json"), summary, mode=0o600) print(json.dumps(summary, sort_keys=True)) return summary async def finite_timestamp_load( ctx: p3.RunContext, session: aiohttp.ClientSession, rate: float ) -> list[dict[str, Any]]: if "arrival_s" not in ctx.rows[0]: return await _ORIGINAL_FINITE_LOAD(ctx, session, rate) sem = asyncio.Semaphore(ctx.args.max_concurrency) tasks: list[asyncio.Task[dict[str, Any]]] = [] async def limited(row: dict[str, Any], scheduled: float) -> dict[str, Any]: async with sem: return await p3.request_one(ctx, session, row, scheduled) for expected in ctx.rows: scheduled = ctx.t0 + float(expected["arrival_s"]) delay = scheduled - asyncio.get_running_loop().time() if delay > 0: try: await asyncio.wait_for(ctx.stop_event.wait(), timeout=delay) break except asyncio.TimeoutError: pass if ctx.stop_event.is_set(): break row = await ctx.next_row() if row["request_id"] != expected["request_id"]: raise AssertionError("timestamp scheduler row drift") tasks.append(asyncio.create_task(limited(row, scheduled))) return await asyncio.gather(*tasks) if tasks else [] _ORIGINAL_FINITE_LOAD = p3.finite_load p3.finite_load = finite_timestamp_load def build_transform_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser() parser.add_argument("--in", dest="input", required=True) parser.add_argument("--take-first", type=int, required=True) parser.add_argument("--timestamp-source", required=True) parser.add_argument("--join-key", default="source_index") parser.add_argument("--timestamp-field", default="timestamp") parser.add_argument("--arrival", choices=("recorded-scaled", "uniform"), required=True) parser.add_argument("--target-rate", type=float, required=True) parser.add_argument("--service-order", choices=("original", "length-binned"), required=True) parser.add_argument("--reorder-block-size", type=int, default=32) parser.add_argument("--analysis-cohort-size", type=int, default=16) parser.add_argument("--length-bin-edges", default="512,1024,2048,4096,8192,16384,32768") parser.add_argument("--max-added-delay-seconds", type=float, default=64) parser.add_argument("--out", required=True) return parser def main() -> None: if len(sys.argv) > 1 and sys.argv[1] == "transform": transform(build_transform_parser().parse_args(sys.argv[2:])) return fixed_rate = None if "--fixed-request-rate" in sys.argv: index = sys.argv.index("--fixed-request-rate") fixed_rate = float(sys.argv[index + 1]) del sys.argv[index : index + 2] args = p3.build_parser().parse_args() if args.command != "run": p3.main() return if fixed_rate is not None: if args.load_point != "moderate" or fixed_rate <= 0 or not math.isfinite(fixed_rate): raise ValueError("--fixed-request-rate requires positive finite moderate rate") result_dir = Path(args.result_dir) result_dir.mkdir(parents=True, exist_ok=True) source = result_dir / "fixed-rate-source.json" p3.atomic_json(source, {"clean": {"completed_throughput_rps": fixed_rate}}) args.saturation_result = str(source) args.rate_fraction = 1.0 if args.profile_after_clean and not args.profile_trace_dir: raise ValueError("--profile-after-clean requires --profile-trace-dir") print(json.dumps(asyncio.run(p3.run_load(args)), sort_keys=True)) if __name__ == "__main__": main()