Freeze Qwen30 trace profile support

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()