216 lines
7.7 KiB
Python
216 lines
7.7 KiB
Python
#!/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()
|