264 lines
8.6 KiB
Python
Executable File
264 lines
8.6 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Low-overhead collectors and result summarizer for Kimi-K3 PD attribution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import re
|
|
import statistics
|
|
import subprocess
|
|
import time
|
|
from pathlib import Path
|
|
|
|
|
|
HCA_COUNTERS = (
|
|
"port_xmit_data",
|
|
"port_rcv_data",
|
|
"port_xmit_wait",
|
|
"port_rcv_errors",
|
|
"port_xmit_discards",
|
|
)
|
|
|
|
|
|
def read_int(path: Path) -> int | None:
|
|
try:
|
|
return int(path.read_text().strip())
|
|
except (FileNotFoundError, PermissionError, ValueError):
|
|
return None
|
|
|
|
|
|
def read_hcas() -> dict[str, dict[str, int | None]]:
|
|
result = {}
|
|
for index in range(4):
|
|
hca = f"mlx5_{index}"
|
|
root = Path(f"/sys/class/infiniband/{hca}/ports/1/counters")
|
|
result[hca] = {name: read_int(root / name) for name in HCA_COUNTERS}
|
|
return result
|
|
|
|
|
|
def read_gpus() -> list[dict[str, float | int | str | None]]:
|
|
fields = (
|
|
"index,timestamp,utilization.gpu,memory.used,power.draw,clocks.sm,"
|
|
"pcie.link.gen.current,pcie.link.width.current"
|
|
)
|
|
proc = subprocess.run(
|
|
[
|
|
"nvidia-smi",
|
|
f"--query-gpu={fields}",
|
|
"--format=csv,noheader,nounits",
|
|
],
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
rows = []
|
|
for line in proc.stdout.splitlines():
|
|
values = [value.strip() for value in line.split(",")]
|
|
if len(values) != 8:
|
|
continue
|
|
parsed: list[float | int | str | None] = [int(values[0]), values[1]]
|
|
for value in values[2:]:
|
|
try:
|
|
parsed.append(float(value))
|
|
except ValueError:
|
|
parsed.append(None)
|
|
rows.append(
|
|
dict(
|
|
zip(
|
|
(
|
|
"index",
|
|
"timestamp",
|
|
"gpu_util_pct",
|
|
"memory_used_mib",
|
|
"power_w",
|
|
"sm_clock_mhz",
|
|
"pcie_gen",
|
|
"pcie_width",
|
|
),
|
|
parsed,
|
|
)
|
|
)
|
|
)
|
|
return rows
|
|
|
|
|
|
def sample(args: argparse.Namespace) -> None:
|
|
out = Path(args.out)
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
deadline = time.monotonic() + args.duration
|
|
with out.open("w") as handle:
|
|
while time.monotonic() < deadline:
|
|
started = time.monotonic()
|
|
record = {
|
|
"wall_time": time.time(),
|
|
"monotonic": started,
|
|
"hcas": read_hcas(),
|
|
"gpus": read_gpus(),
|
|
}
|
|
handle.write(json.dumps(record, separators=(",", ":")) + "\n")
|
|
handle.flush()
|
|
time.sleep(max(0.0, args.interval - (time.monotonic() - started)))
|
|
|
|
|
|
def percentile(values: list[float], q: float) -> float | None:
|
|
if not values:
|
|
return None
|
|
ordered = sorted(values)
|
|
index = (len(ordered) - 1) * q
|
|
lower = int(index)
|
|
upper = min(lower + 1, len(ordered) - 1)
|
|
fraction = index - lower
|
|
return ordered[lower] * (1 - fraction) + ordered[upper] * fraction
|
|
|
|
|
|
def monitor_summary(path: Path) -> dict:
|
|
rows = [json.loads(line) for line in path.read_text().splitlines() if line]
|
|
hca_rates: dict[str, dict[str, list[float]]] = {}
|
|
gpu_util, gpu_power = [], []
|
|
for previous, current in zip(rows, rows[1:]):
|
|
elapsed = current["monotonic"] - previous["monotonic"]
|
|
if elapsed <= 0:
|
|
continue
|
|
for hca, counters in current["hcas"].items():
|
|
target = hca_rates.setdefault(hca, {"tx_gbps": [], "rx_gbps": []})
|
|
old = previous["hcas"].get(hca, {})
|
|
for counter, output in (
|
|
("port_xmit_data", "tx_gbps"),
|
|
("port_rcv_data", "rx_gbps"),
|
|
):
|
|
before, after = old.get(counter), counters.get(counter)
|
|
if before is not None and after is not None and after >= before:
|
|
# IB port_*_data counters are in 32-bit words.
|
|
target[output].append((after - before) * 4 * 8 / elapsed / 1e9)
|
|
for gpu in current.get("gpus", []):
|
|
if gpu.get("gpu_util_pct") is not None:
|
|
gpu_util.append(float(gpu["gpu_util_pct"]))
|
|
if gpu.get("power_w") is not None:
|
|
gpu_power.append(float(gpu["power_w"]))
|
|
|
|
hcas = {}
|
|
for hca, rates in hca_rates.items():
|
|
hcas[hca] = {
|
|
name: {
|
|
"mean": statistics.fmean(values) if values else None,
|
|
"p95": percentile(values, 0.95),
|
|
"max": max(values) if values else None,
|
|
}
|
|
for name, values in rates.items()
|
|
}
|
|
return {
|
|
"samples": len(rows),
|
|
"hcas": hcas,
|
|
"gpu_util_pct_mean": statistics.fmean(gpu_util) if gpu_util else None,
|
|
"gpu_util_pct_p95": percentile(gpu_util, 0.95),
|
|
"gpu_power_w_mean": statistics.fmean(gpu_power) if gpu_power else None,
|
|
}
|
|
|
|
|
|
REQ_RE = re.compile(r"ReqTimeStats\(.*?type=(prefill|decode)\): (.*)$")
|
|
VALUE_RE = re.compile(r"([a-z_]+)=([0-9.]+)(ms| GB/s| MB)?")
|
|
|
|
|
|
def request_stats(paths: list[Path]) -> dict:
|
|
by_role: dict[str, dict[str, list[float]]] = {"prefill": {}, "decode": {}}
|
|
for path in paths:
|
|
for line in path.read_text(errors="replace").splitlines():
|
|
match = REQ_RE.search(line)
|
|
if not match:
|
|
continue
|
|
role, payload = match.groups()
|
|
for key, value, unit in VALUE_RE.findall(payload):
|
|
number = float(value)
|
|
if unit == "ms":
|
|
number /= 1000
|
|
by_role[role].setdefault(key, []).append(number)
|
|
result = {}
|
|
for role, fields in by_role.items():
|
|
result[role] = {
|
|
key: {
|
|
"count": len(values),
|
|
"mean_s": statistics.fmean(values),
|
|
"p50_s": percentile(values, 0.50),
|
|
"p95_s": percentile(values, 0.95),
|
|
}
|
|
for key, values in fields.items()
|
|
}
|
|
return result
|
|
|
|
|
|
def benchmark_summary(path: Path) -> dict:
|
|
payload = json.loads(path.read_text().splitlines()[0])
|
|
keys = (
|
|
"duration",
|
|
"completed",
|
|
"total_input_tokens",
|
|
"total_output_tokens",
|
|
"input_throughput",
|
|
"output_throughput",
|
|
"mean_ttft_ms",
|
|
"median_ttft_ms",
|
|
"p95_ttft_ms",
|
|
"mean_tpot_ms",
|
|
"median_tpot_ms",
|
|
"p95_tpot_ms",
|
|
"mean_e2e_latency_ms",
|
|
"p95_e2e_latency_ms",
|
|
"accept_length",
|
|
)
|
|
return {key: payload.get(key) for key in keys}
|
|
|
|
|
|
def analyze(args: argparse.Namespace) -> None:
|
|
root = Path(args.result_root)
|
|
bench_files = sorted((root / "bench").glob("*.jsonl"))
|
|
service_logs = sorted((root / "logs").glob("[pd]_*.log"))
|
|
monitor_files = sorted((root / "telemetry").glob("*.jsonl"))
|
|
summary = {
|
|
"run_id": root.name,
|
|
"benchmarks": {path.name: benchmark_summary(path) for path in bench_files},
|
|
"request_stages": request_stats(service_logs),
|
|
"nodes": {path.stem: monitor_summary(path) for path in monitor_files},
|
|
}
|
|
(root / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
|
|
|
|
rows = []
|
|
for name, item in summary["benchmarks"].items():
|
|
rows.append(
|
|
{
|
|
"file": name,
|
|
"input_tps": item.get("input_throughput"),
|
|
"output_tps": item.get("output_throughput"),
|
|
"ttft_mean_ms": item.get("mean_ttft_ms"),
|
|
"ttft_p95_ms": item.get("p95_ttft_ms"),
|
|
"tpot_mean_ms": item.get("mean_tpot_ms"),
|
|
"e2e_mean_ms": item.get("mean_e2e_latency_ms"),
|
|
"accept_length": item.get("accept_length"),
|
|
}
|
|
)
|
|
with (root / "summary.csv").open("w", newline="") as handle:
|
|
writer = csv.DictWriter(handle, fieldnames=rows[0].keys() if rows else ["file"])
|
|
writer.writeheader()
|
|
writer.writerows(rows)
|
|
print(json.dumps(summary, indent=2))
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
sub = parser.add_subparsers(dest="command", required=True)
|
|
collect = sub.add_parser("sample")
|
|
collect.add_argument("--out", required=True)
|
|
collect.add_argument("--interval", type=float, default=1.0)
|
|
collect.add_argument("--duration", type=float, default=900.0)
|
|
collect.set_defaults(func=sample)
|
|
report = sub.add_parser("analyze")
|
|
report.add_argument("--result-root", required=True)
|
|
report.set_defaults(func=analyze)
|
|
args = parser.parse_args()
|
|
args.func(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|