Stabilize first-run cudagraph timing and colorize terminal output
- _cudagraph_plugin: warm up autotune/JIT explicitly before graph capture (internal 5-iter warmup is too short), tag fallback markers with the failing phase; first-run latency no longer jitters run-to-run. - _device_guard_plugin (new): set GEMS_VENDOR via torch probe before importing flag_gems, avoiding its timeout-less nvidia-smi subprocess probe that can hang import in fork-broken environments. - _pretty_report_plugin (new) + _term_style (new): fold inputs identical across all result rows into a legend line, color status/plugin tags/markers on the live terminal; run.log is ANSI-stripped and keeps upstream SUCCESS/column wording for grep compatibility. - run_pytest.sh: make USE_FLAGTUNE overridable, group all knobs into a config section with Chinese comments, add start/end banners. - README: document the warmup semantics, new plugins/env vars, and the A/B rule that both sides must use the same USE_FLAGTUNE.
This commit is contained in:
@@ -0,0 +1,148 @@
|
||||
"""pytest plugin: readable, colorized benchmark result table.
|
||||
|
||||
Replaces FlagGems' BenchmarkResult.__str__ (benchmark/consts.py), whose rows
|
||||
print every input's torch.Size(...) on every line — for the MoE ops that is 11
|
||||
tensors and ~400 chars per row, ~90% of it identical across rows. This plugin:
|
||||
|
||||
- folds inputs identical across all rows into one legend line above the table,
|
||||
so each row's Size Detail keeps only what varies between rows;
|
||||
- shapes keep their original torch.Size([...]) spelling (no compaction);
|
||||
- colors the table title and Status (SUCCESS green / FAILED red);
|
||||
- keeps the literal "SUCCESS"/"FAILED" words and upstream column titles, so
|
||||
run.log greps and __str__-wrapping hooks (vllm column renames) still work.
|
||||
|
||||
Color only reaches the live terminal: run_pytest.sh strips ANSI from run.log,
|
||||
and _term_style disables color for non-tty stdout unless forced.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
from _term_style import BOLD, CYAN, DIM, GREEN, RED, paint, tag
|
||||
|
||||
|
||||
def _shape_rows(metrics_list) -> list:
|
||||
"""One list of per-input cell strings per metrics row (upstream spelling)."""
|
||||
rows = []
|
||||
for m in metrics_list:
|
||||
sd = m.shape_detail
|
||||
if isinstance(sd, (list, tuple)):
|
||||
rows.append([str(e) for e in sd])
|
||||
else:
|
||||
rows.append([str(sd) if sd is not None else "N/A"])
|
||||
return rows
|
||||
|
||||
|
||||
def _factor_static(rows):
|
||||
"""Fold columns identical across all rows into a legend.
|
||||
|
||||
Only kicks in when it actually helps: >=2 rows of equal arity >=3, with
|
||||
>=2 static positions and at least one varying position left. Returns
|
||||
(legend_cells_or_None, per_row_varying_strings_or_None).
|
||||
"""
|
||||
if len(rows) >= 2:
|
||||
arities = {len(r) for r in rows}
|
||||
if len(arities) == 1:
|
||||
n = next(iter(arities))
|
||||
if n >= 3:
|
||||
static = {i for i in range(n) if len({r[i] for r in rows}) == 1}
|
||||
if 2 <= len(static) < n:
|
||||
legend = [rows[0][i] for i in sorted(static)]
|
||||
kept = [
|
||||
", ".join(r[i] for i in range(n) if i not in static)
|
||||
for r in rows
|
||||
]
|
||||
return legend, kept
|
||||
return None, None
|
||||
|
||||
|
||||
def _num_cell(value, fmt: str, width: int) -> str:
|
||||
s = f"{value:{fmt}}" if value is not None else "N/A"
|
||||
return f"{s:>{width}}"
|
||||
|
||||
|
||||
def _pretty_str(self) -> str:
|
||||
metrics = self.result or []
|
||||
title = (
|
||||
"\n"
|
||||
+ paint(f"Operator: {self.op_name}", BOLD, CYAN)
|
||||
+ paint(f" (dtype={self.dtype}, mode={self.mode}, level={self.level})", DIM)
|
||||
+ "\n"
|
||||
)
|
||||
if not metrics:
|
||||
return title + "(no results)\n"
|
||||
|
||||
# Same optional-column conditions as upstream.
|
||||
with_tflops = bool(metrics[0].tflops)
|
||||
with_gbps = metrics[0].gbps is not None
|
||||
|
||||
legend, varying = _factor_static(_shape_rows(metrics))
|
||||
legend_line = ""
|
||||
if legend:
|
||||
legend_line = (
|
||||
paint(
|
||||
f"{len(legend)} inputs identical across all rows: "
|
||||
+ ", ".join(legend),
|
||||
DIM,
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
shape_cells = varying
|
||||
else:
|
||||
# No folding possible: print the full shape_detail, upstream style.
|
||||
shape_cells = [
|
||||
str(m.shape_detail) if m.shape_detail is not None else "N/A"
|
||||
for m in metrics
|
||||
]
|
||||
|
||||
cols = [
|
||||
("Status", 8, "<"),
|
||||
("Torch Latency (ms)", 19, ">"),
|
||||
("Gems Latency (ms)", 18, ">"),
|
||||
("Gems Speedup", 13, ">"),
|
||||
]
|
||||
if with_tflops:
|
||||
cols.append(("TFLOPS", 13, ">"))
|
||||
if with_gbps:
|
||||
cols.append(("Torch GBPS", 12, ">"))
|
||||
cols.append(("Gems GBPS", 12, ">"))
|
||||
head_plain = (
|
||||
" ".join(f"{name:{align}{width}}" for name, width, align in cols)
|
||||
+ " Size Detail"
|
||||
)
|
||||
header = paint(head_plain, BOLD) + "\n" + paint("-" * len(head_plain), DIM) + "\n"
|
||||
|
||||
lines = []
|
||||
for m, shape_cell in zip(metrics, shape_cells):
|
||||
ok = m.error_msg is None
|
||||
cells = [paint(f"{'SUCCESS' if ok else 'FAILED':<8}", GREEN if ok else RED)]
|
||||
cells.append(_num_cell(m.latency_base, ".6f", 19))
|
||||
cells.append(_num_cell(m.latency, ".6f", 18))
|
||||
cells.append(_num_cell(m.speedup, ".3f", 13))
|
||||
if with_tflops:
|
||||
cells.append(_num_cell(m.tflops, ".3f", 13))
|
||||
if with_gbps:
|
||||
cells.append(_num_cell(m.gbps_base, ".3f", 12))
|
||||
cells.append(_num_cell(m.gbps, ".3f", 12))
|
||||
lines.append(" ".join(cells) + " " + shape_cell)
|
||||
|
||||
return title + legend_line + header + "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
try:
|
||||
from benchmark import consts as fg_consts
|
||||
except Exception as exc:
|
||||
print(
|
||||
f"{tag('[pretty-report-plugin]')} disabled "
|
||||
f"(cannot import benchmark.consts: {exc})",
|
||||
file=sys.stderr, flush=True,
|
||||
)
|
||||
return
|
||||
fg_consts.BenchmarkResult.__str__ = _pretty_str
|
||||
print(
|
||||
f"{tag('[pretty-report-plugin]')} BenchmarkResult table -> folded + color "
|
||||
"(SUCCESS/FAILED words and column titles unchanged)",
|
||||
file=sys.stderr, flush=True,
|
||||
)
|
||||
Reference in New Issue
Block a user