45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Extract simulator-visible graph/KV metadata from Qwen235 server starts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import re
|
|
from pathlib import Path
|
|
|
|
|
|
CONFIGS = ("tp4_ep1_mns64", "tp4_ep1_mns128", "tp8_ep8_mns64", "tp8_ep8_mns128")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--case-root", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
configs = {}
|
|
for config in CONFIGS:
|
|
log = args.case_root / "real" / config / "trial1" / "logs" / "server.log"
|
|
text = log.read_text(errors="replace")
|
|
token_matches = re.findall(r"GPU KV cache size:\s*([0-9,]+) tokens", text)
|
|
capture_matches = re.findall(r"cudagraph_capture_sizes': \[([^]]+)\]", text)
|
|
if not token_matches or not capture_matches:
|
|
raise ValueError(f"missing runtime metadata in {log}")
|
|
tokens = int(token_matches[-1].replace(",", ""))
|
|
if tokens % 16:
|
|
raise ValueError(f"KV token count is not block aligned: {tokens}")
|
|
capture = [int(value.strip()) for value in capture_matches[-1].split(",")]
|
|
configs[config] = {
|
|
"num_gpu_blocks": tokens // 16,
|
|
"capture_sizes": capture,
|
|
"source_log": str(log.resolve()),
|
|
}
|
|
payload = {"schema": "qwen235-v020-runtime-contract-v1", "configs": configs}
|
|
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, sort_keys=True))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|