diff --git a/runs/frontier-qwen30-vllm020-profile-v1/extract_trace_profile_support.py b/runs/frontier-qwen30-vllm020-profile-v1/extract_trace_profile_support.py new file mode 100644 index 0000000..e3e37c6 --- /dev/null +++ b/runs/frontier-qwen30-vllm020-profile-v1/extract_trace_profile_support.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +"""Freeze the exact request cohort and its operator-profile support. + +This calls AITuner's production trace loader, including its input-length +filter, uniform max-request downsampling, output override, and sampling-u +threshold semantics. Prefix reuse below is a no-eviction upper bound; the +actual cache state remains scheduler/config dependent. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +from pathlib import Path +from typing import Any, Iterable + +from aituner.spec import load_study_spec +from aituner.trace import load_trace_requests, select_requests_for_threshold + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--study", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--cohort-output", type=Path, required=True) + parser.add_argument( + "--thresholds", + type=float, + nargs="+", + default=[0.125, 0.25, 0.5, 1.0], + ) + return parser.parse_args() + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def nearest_rank(values: Iterable[float], percentile: float) -> float: + ordered = sorted(values) + if not ordered: + return 0.0 + index = min( + len(ordered) - 1, + max(0, math.ceil(percentile / 100.0 * len(ordered)) - 1), + ) + return float(ordered[index]) + + +def distribution(values: list[float]) -> dict[str, float]: + return { + "min": min(values, default=0.0), + "p10": nearest_rank(values, 10), + "p25": nearest_rank(values, 25), + "p50": nearest_rank(values, 50), + "p75": nearest_rank(values, 75), + "p90": nearest_rank(values, 90), + "p95": nearest_rank(values, 95), + "p99": nearest_rank(values, 99), + "max": max(values, default=0.0), + } + + +def prefix_upper_bound(requests: list[Any], block_size: int) -> dict[str, Any]: + seen: set[Any] = set() + total_blocks = 0 + reusable_blocks = 0 + leading_reusable_blocks: list[float] = [] + reusable_tokens: list[float] = [] + rows_with_hashes = 0 + for request in requests: + hashes = request.metadata.get("hash_ids") + if not isinstance(hashes, list): + continue + rows_with_hashes += 1 + total_blocks += len(hashes) + reusable = sum(hash_id in seen for hash_id in hashes) + reusable_blocks += reusable + leading = 0 + for hash_id in hashes: + if hash_id not in seen: + break + leading += 1 + leading_reusable_blocks.append(float(leading)) + reusable_tokens.append( + float(min(request.prompt_tokens_hint or 0, leading * block_size)) + ) + seen.update(hashes) + return { + "semantics": ( + "arrival-ordered, infinite-capacity/no-eviction upper bound; " + "not an observed KV-cache hit rate" + ), + "rows_with_hash_ids": rows_with_hashes, + "total_blocks": total_blocks, + "unique_blocks": len(seen), + "any_position_reusable_block_ratio": ( + reusable_blocks / total_blocks if total_blocks else 0.0 + ), + "leading_reusable_blocks_per_request": distribution( + leading_reusable_blocks + ), + "leading_reusable_tokens_per_request": distribution(reusable_tokens), + } + + +def summarize(requests: list[Any], block_size: int) -> dict[str, Any]: + input_lengths = [float(request.prompt_tokens_hint or 0) for request in requests] + output_lengths = [ + float(request.completion_tokens_hint or 0) for request in requests + ] + hash_counts = [ + float(len(request.metadata.get("hash_ids") or [])) for request in requests + ] + arrivals = [request.arrival_s for request in requests] + interarrivals = [ + max(0.0, arrivals[index] - arrivals[index - 1]) + for index in range(1, len(arrivals)) + ] + return { + "request_count": len(requests), + "input_tokens": distribution(input_lengths), + "output_tokens": distribution(output_lengths), + "hash_blocks_per_request": distribution(hash_counts), + "interarrival_s": distribution(interarrivals), + "sampling_u": distribution([request.sampling_u for request in requests]), + "multi_turn_fraction": ( + sum( + isinstance(request.metadata.get("turn"), (int, float)) + and request.metadata["turn"] > 1 + for request in requests + ) + / len(requests) + if requests + else 0.0 + ), + "prefix_reuse_upper_bound": prefix_upper_bound(requests, block_size), + } + + +def cohort_row(request: Any) -> dict[str, Any]: + return { + "row_id": request.row_id, + "arrival_s": request.arrival_s, + "sampling_u": request.sampling_u, + "input_tokens": request.prompt_tokens_hint, + "output_tokens": request.completion_tokens_hint, + "hash_ids": request.metadata.get("hash_ids"), + "turn": request.metadata.get("turn"), + "parent_chat_id": request.metadata.get("parent_chat_id"), + "type": request.metadata.get("type"), + } + + +def main() -> None: + args = parse_args() + study = load_study_spec(args.study) + window, cohort = load_trace_requests(study, study_spec_path=args.study) + block_size = int(window.source_payload.get("block_size") or 1) + + args.cohort_output.parent.mkdir(parents=True, exist_ok=True) + with args.cohort_output.open("w", encoding="utf-8") as handle: + for request in cohort: + handle.write(json.dumps(cohort_row(request), sort_keys=True) + "\n") + + payload = { + "schema_version": "qwen30_trace_profile_support.v1", + "study": str(args.study.resolve()), + "study_sha256": sha256_file(args.study), + "trace": str(window.trace_path), + "trace_sha256": sha256_file(window.trace_path), + "window_id": window.window_id, + "block_size": block_size, + "loader_contract": { + "input_length_filter": { + "min": study.trace.input_length_filter.min_input_tokens, + "max": study.trace.input_length_filter.max_input_tokens, + } + if study.trace.input_length_filter is not None + else None, + "completion_tokens_override": study.trace.completion_tokens_override, + "max_requests_per_probe": study.trace.max_requests_per_probe, + "replay_time_scale": study.trace.replay_time_scale, + "ordering": "arrival_s", + "downsampling": "AITuner _downsample_requests before sampling_u threshold", + }, + "full_downsampled_cohort": summarize(cohort, block_size), + "threshold_cohorts": { + str(threshold): summarize( + select_requests_for_threshold(cohort, threshold=threshold), + block_size, + ) + for threshold in args.thresholds + }, + "limits": [ + "Trace length/hash support does not determine dynamic decode or mixed batch shapes.", + "Those shapes depend jointly on arrival history, SLO pressure, TP execution time, MNS, chunking, and KV eviction.", + "MoE expert routing is not present in the trace and must be measured from model execution.", + ], + } + payload["cohort_output"] = str(args.cohort_output.resolve()) + payload["cohort_sha256"] = sha256_file(args.cohort_output) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + print(json.dumps(payload, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main()