Freeze vLLM 0.20 profiles and capture trace routing
This commit is contained in:
@@ -0,0 +1,162 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Extract a prompt-bearing routing fixture while emitting a prompt-free manifest."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--trace", type=Path, required=True)
|
||||
parser.add_argument("--row-ids", type=int, nargs="+", required=True)
|
||||
parser.add_argument("--parent-of", type=int, nargs="*", default=[])
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--manifest", type=Path, required=True)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def common_prefix(left: list[Any], right: list[Any]) -> int:
|
||||
count = 0
|
||||
for lhs, rhs in zip(left, right):
|
||||
if lhs != rhs:
|
||||
break
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
target_ids = set(args.row_ids)
|
||||
parent_targets = set(args.parent_of)
|
||||
if not parent_targets.issubset(target_ids):
|
||||
raise SystemExit("--parent-of must be a subset of --row-ids")
|
||||
|
||||
trace_digest = hashlib.sha256()
|
||||
by_row: dict[int, dict[str, Any]] = {}
|
||||
by_chat_id: dict[str, dict[str, Any]] = {}
|
||||
with args.trace.open("rb") as handle:
|
||||
row_id = 0
|
||||
while True:
|
||||
offset = handle.tell()
|
||||
line = handle.readline()
|
||||
if not line:
|
||||
break
|
||||
trace_digest.update(line)
|
||||
row = json.loads(line)
|
||||
meta = {
|
||||
"row_id": row_id,
|
||||
"offset": offset,
|
||||
"chat_id": str(row.get("chat_id")),
|
||||
"parent_chat_id": str(row.get("parent_chat_id")),
|
||||
"turn": int(row.get("turn", 1)),
|
||||
"input_length": int(row["input_length"]),
|
||||
"output_length": int(row["output_length"]),
|
||||
"hash_ids": row.get("hash_ids") or [],
|
||||
}
|
||||
by_chat_id[meta["chat_id"]] = meta
|
||||
if row_id in target_ids:
|
||||
by_row[row_id] = meta
|
||||
row_id += 1
|
||||
|
||||
missing = target_ids - set(by_row)
|
||||
if missing:
|
||||
raise SystemExit(f"missing target row IDs: {sorted(missing)}")
|
||||
parent_rows: list[dict[str, Any]] = []
|
||||
for row_id in args.parent_of:
|
||||
parent_id = by_row[row_id]["parent_chat_id"]
|
||||
parent = by_chat_id.get(parent_id)
|
||||
if parent is None:
|
||||
raise SystemExit(f"row {row_id} parent chat {parent_id} is absent")
|
||||
parent_rows.append(parent)
|
||||
|
||||
# Put parents first so online prefix caching can materialize their shared
|
||||
# blocks before descendants are admitted.
|
||||
ordered_meta = parent_rows + [by_row[row_id] for row_id in args.row_ids]
|
||||
if len({row["row_id"] for row in ordered_meta}) != len(ordered_meta):
|
||||
raise SystemExit("fixture rows must be unique")
|
||||
|
||||
output_rows: list[dict[str, Any]] = []
|
||||
with args.trace.open("rb") as handle:
|
||||
for fixture_index, meta in enumerate(ordered_meta):
|
||||
handle.seek(meta["offset"])
|
||||
source = json.loads(handle.readline())
|
||||
prompt = source.get("prompt")
|
||||
if not isinstance(prompt, str) or not prompt:
|
||||
raise SystemExit(f"row {meta['row_id']} has no prompt")
|
||||
output_rows.append(
|
||||
{
|
||||
"fixture_index": fixture_index,
|
||||
"row_id": str(meta["row_id"]),
|
||||
"prompt": prompt,
|
||||
"prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(),
|
||||
"input_length": meta["input_length"],
|
||||
"output_length": meta["output_length"],
|
||||
"turn": meta["turn"],
|
||||
"chat_id": meta["chat_id"],
|
||||
"parent_chat_id": meta["parent_chat_id"],
|
||||
"hash_ids": meta["hash_ids"],
|
||||
}
|
||||
)
|
||||
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
with args.output.open("w") as handle:
|
||||
for row in output_rows:
|
||||
handle.write(json.dumps(row, sort_keys=True) + "\n")
|
||||
|
||||
pair_coverage = []
|
||||
by_chat = {row["chat_id"]: row for row in output_rows}
|
||||
for child in output_rows:
|
||||
parent = by_chat.get(child["parent_chat_id"])
|
||||
if parent is None:
|
||||
continue
|
||||
pair_coverage.append(
|
||||
{
|
||||
"parent_row_id": parent["row_id"],
|
||||
"child_row_id": child["row_id"],
|
||||
"parent_turn": parent["turn"],
|
||||
"child_turn": child["turn"],
|
||||
"trace_hash_common_prefix_blocks": common_prefix(
|
||||
parent["hash_ids"], child["hash_ids"]
|
||||
),
|
||||
}
|
||||
)
|
||||
fixture_digest = hashlib.sha256(args.output.read_bytes()).hexdigest()
|
||||
manifest = {
|
||||
"schema_version": "qwen30_trace_routing_fixture.v1",
|
||||
"source_trace": str(args.trace.resolve()),
|
||||
"source_trace_sha256": trace_digest.hexdigest(),
|
||||
"fixture_path": str(args.output.resolve()),
|
||||
"fixture_sha256": fixture_digest,
|
||||
"contains_prompt_text": False,
|
||||
"request_count": len(output_rows),
|
||||
"rows": [
|
||||
{
|
||||
key: row[key]
|
||||
for key in (
|
||||
"fixture_index",
|
||||
"row_id",
|
||||
"prompt_sha256",
|
||||
"input_length",
|
||||
"output_length",
|
||||
"turn",
|
||||
"chat_id",
|
||||
"parent_chat_id",
|
||||
)
|
||||
}
|
||||
| {"trace_hash_blocks": len(row["hash_ids"])}
|
||||
for row in output_rows
|
||||
],
|
||||
"prefix_pairs": pair_coverage,
|
||||
}
|
||||
args.manifest.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.manifest.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
||||
print(json.dumps({"requests": len(output_rows), "prefix_pairs": pair_coverage}))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user