Abstract hf code and fix plot bug

This commit is contained in:
Satyam Srivastava
2026-04-27 23:12:42 -07:00
parent 3ff4a8d2d2
commit d451e61749
3 changed files with 314 additions and 240 deletions
+10 -59
View File
@@ -18,6 +18,8 @@ from huggingface_hub import HfApi, snapshot_download
from datetime import datetime, timezone
from typing import Any
from hf_store import sync_from_hf, upload_record, load_records_for_model, sanitize, safe_float
# Use the env var passed by Modal, fallback to a default if needed
HF_REPO_ID = os.environ.get("HF_REPO_ID", "FastVideo/performance-tracking")
HF_TOKEN = os.environ.get("HF_API_KEY")
@@ -32,7 +34,6 @@ TRACKING_ROOT = os.environ.get(
)
MAX_REGRESSION = float(os.environ.get("PERF_MAX_REGRESSION", "0.15"))
def _should_persist_tracking() -> bool:
# test_scope = os.environ.get("TEST_SCOPE", "")
# branch = os.environ.get("BUILDKITE_BRANCH", "")
@@ -119,58 +120,7 @@ def _sync_baselines_from_hf():
# If the repo is empty or doesn't exist yet, we just start fresh
print(f"Note: Could not sync from HF (may be an empty repo): {e}")
def _upload_to_hf(local_path: str, record: dict[str, Any]):
"""Uploads result to the FastVideo organization repo."""
if not HF_TOKEN:
print("HF_API_KEY not found. Skipping upload.")
return
api = HfApi(token=HF_TOKEN)
# Store in folders by model_id: e.g., "fastvideo-v1/2026-04-23_commitsha.json"
path_in_repo = f"{_sanitize(record['model_id'])}/{os.path.basename(local_path)}"
try:
api.upload_file(
path_or_fileobj=local_path,
path_in_repo=path_in_repo,
repo_id=HF_REPO_ID,
repo_type="dataset",
commit_message=f"Perf: {record['model_id']} at {record['commit_sha'][:7]}"
)
print(f"Successfully synced to FastVideo/performance-tracking: {path_in_repo}")
except Exception as e:
print(f"HF Sync failed: {e}")
def _load_baseline_records(model_id: str) -> list[dict[str, Any]]:
model_dir = os.path.join(TRACKING_ROOT, _sanitize(model_id))
if not os.path.exists(model_dir):
return []
# Get all historical JSONs synced from HF
all_paths = sorted(glob.glob(os.path.join(model_dir, "*.json")))
baseline: list[dict[str, Any]] = []
for path in all_paths:
try:
with open(path, encoding="utf-8") as f:
data = json.load(f)
# IMPORTANT: Only use successful runs for the baseline mean
# This prevents "Performance Drift"
if data.get("success", True):
baseline.append(data)
except (OSError, json.JSONDecodeError):
continue
return baseline
def _filter_by_gpu_type(
records: list[dict[str, Any]],
gpu_type: str,
) -> list[dict[str, Any]]:
return [r for r in records if r.get("gpu_type") == gpu_type]
# Todo: Maybe we can go with the median and not mean.
def _mean_metric(records: list[dict[str, Any]], key: str) -> float | None:
values = [
_safe_float(r.get(key))
@@ -297,6 +247,7 @@ def _build_markdown_summary(
return "\n".join(lines) + "\n"
# Todo: Put this in post run hookup since buildkite agent won't be available in the Modal function environment.
def _emit_markdown_summary(markdown: str) -> None:
print("\n" + markdown)
@@ -321,7 +272,7 @@ def _emit_markdown_summary(markdown: str) -> None:
def main() -> int:
# Pull the current state of the world from HF
_sync_baselines_from_hf()
sync_from_hf(TRACKING_ROOT)
current_results = _load_current_results()
if not current_results:
@@ -340,9 +291,10 @@ def main() -> int:
for raw in current_results:
record = _normalize_record(raw)
baseline_records = _load_baseline_records(record["model_id"])
baseline_records = _filter_by_gpu_type(baseline_records, record["gpu_type"])
baseline_records = baseline_records[-5:]
baseline_records = load_records_for_model(
TRACKING_ROOT, record["model_id"], record["gpu_type"],
last_n=5, successful_only=True
)
failures = _check_regressions(record, baseline_records, MAX_REGRESSION)
@@ -366,7 +318,7 @@ def main() -> int:
# This writes the JSON with the "success" field to /tmp
current_path = _write_tracking_record(record)
# This pushes it to the FastVideo/performance-tracking repo
_upload_to_hf(current_path, record)
upload_record(current_path, record)
summary_row = _build_summary_row(record, baseline_records, bool(failures))
@@ -384,6 +336,5 @@ def main() -> int:
print("Performance baseline comparison passed")
return 0
if __name__ == "__main__":
sys.exit(main())
+50 -181
View File
@@ -1,151 +1,60 @@
# SPDX-License-Identifier: Apache-2.0
import os
import json
import glob
import pandas as pd
from datetime import datetime, timedelta, timezone
import shutil
import subprocess
from datetime import datetime
import plotly.express as px
from huggingface_hub import snapshot_download
HF_REPO_ID = os.environ.get("HF_REPO_ID", "FastVideo/performance-tracking")
HF_TOKEN = os.environ.get("HF_API_KEY")
import pandas as pd
from hf_store import sync_from_hf, load_as_dataframe
# -----------------------------
# 1. Load HF JSON artifacts
# -----------------------------
# def load_hf_dataset(days: int = 30) -> pd.DataFrame:
# """
# Loads only recent benchmark JSONs from HF repo.
# Avoids full history scan when possible.
# """
# api = HfApi(token=HF_TOKEN)
# files = api.list_repo_files(
# repo_id=HF_REPO_ID,
# repo_type="dataset",
# )
# cutoff = datetime.now(timezone.utc) - timedelta(days=days)
# records = []
# for f in files:
# if not f.endswith(".json"):
# continue
# path = hf_hub_download(
# repo_id=HF_REPO_ID,
# repo_type="dataset",
# filename=f,
# token=HF_TOKEN,
# )
# try:
# with open(path, "r") as fp:
# data = json.load(fp)
# # fast skip: timestamp filter BEFORE dataframe creation
# ts = pd.to_datetime(data.get("timestamp"), utc=True)
# if ts < cutoff:
# continue
# records.append(data)
# except Exception:
# continue
# return pd.DataFrame(records)
def load_hf_dataset(days: int = 30) -> pd.DataFrame:
cutoff = datetime.now(timezone.utc) - timedelta(days=days)
local_dir = snapshot_download(
repo_id=HF_REPO_ID,
repo_type="dataset",
token=HF_TOKEN,
allow_patterns="*.json"
)
records = []
for path in glob.glob(f"{local_dir}/**/*.json", recursive=True):
try:
with open(path, "r") as fp:
data = json.load(fp)
ts = pd.to_datetime(data.get("timestamp"), utc=True, errors="coerce")
if pd.isna(ts) or ts < cutoff:
continue
records.append(data)
except Exception:
continue
return pd.DataFrame(records)
# -----------------------------
# 2. Normalize schema (your CI format)
# -----------------------------
def normalize(df: pd.DataFrame) -> pd.DataFrame:
if df.empty:
return df
df["timestamp"] = pd.to_datetime(df["timestamp"], utc=True)
# config = commit + optional future knobs
df["config_id"] = df["commit_sha"].fillna("unknown").str[:7]
# ensure numeric safety
for c in ["latency", "throughput", "memory"]:
df[c] = pd.to_numeric(df[c], errors="coerce")
return df
# -----------------------------
# 3. Grouping (your requested design)
# 1. Grouping
# -----------------------------
def group_data(df: pd.DataFrame):
keys = ["model_id", "gpu_type", "config_id"]
# Group only by model+GPU so each group produces a time-series line.
# config_id (commit SHA) is carried as a column for hover/color use.
keys = ["model_id", "gpu_type"]
return df.groupby(keys, dropna=False)
# -----------------------------
# 4. Plot builder (clean + scalable)
# 2. Plot builder
# -----------------------------
def build_plots(df: pd.DataFrame):
def build_plots(df: pd.DataFrame) -> list:
figs = []
grouped = group_data(df)
for name, g in grouped:
for (model_id, gpu_type), g in group_data(df):
g = g.sort_values("timestamp")
model_id, gpu_type, config_id = name
# One chart per metric so the y-axes aren't on wildly different scales
for metric in ("latency", "throughput", "memory"):
if g[metric].isna().all():
continue
title = f"{model_id} | {gpu_type} | {config_id}"
fig = px.line(
g,
x="timestamp",
y=["latency", "throughput", "memory"],
markers=True,
title=title,
)
figs.append(fig)
fig = px.line(
g,
x="timestamp",
y=metric,
markers=True,
hover_data=["config_id", "commit_sha"],
title=f"{model_id} | {gpu_type} | {metric}",
labels={"timestamp": "Time", metric: metric},
)
figs.append(fig)
return figs
# -----------------------------
# 5. Render HTML dashboard
# 3. Render HTML dashboard
# -----------------------------
def render_html(figs):
def render_html(figs: list, days: int) -> str:
html_parts = [
"<html><head><meta charset='utf-8'></head><body>",
f"<h2>Performance Dashboard (last run)</h2>",
"<html>",
"<head><meta charset='utf-8'>",
"<style>body { font-family: sans-serif; margin: 2rem; }</style>",
"</head><body>",
f"<h2>Performance Dashboard (last {days} days)</h2>",
]
for fig in figs:
@@ -154,100 +63,60 @@ def render_html(figs):
html_parts.append("</body></html>")
return "\n".join(html_parts)
# -----------------------------
# 6. Buildkite annotation (fixed)
# 4. Buildkite helpers
# -----------------------------
def annotate_buildkite(html: str):
def annotate_buildkite(html: str) -> None:
if not os.environ.get("BUILDKITE"):
return
import subprocess
subprocess.run(
["buildkite-agent", "annotate", "--style", "info", "--context", "perf-dashboard"],
input=html.encode(),
check=False,
)
# -----------------------------
# 6. Buildkite upload
# -----------------------------
# def upload_artifact(path: str):
# if not os.environ.get("BUILDKITE"):
# return
# import subprocess
# subprocess.run(
# ["buildkite-agent", "artifact", "upload", path],
# check=False,
# )
def upload_artifact(path: str):
import shutil
import subprocess
def upload_artifact(path: str) -> None:
if shutil.which("buildkite-agent") is None:
print("buildkite-agent not found, skipping artifact upload")
return
subprocess.run(["buildkite-agent", "artifact", "upload", path], check=True)
subprocess.run(
["buildkite-agent", "artifact", "upload", path],
check=True,
)
# -----------------------------
# 7. Main
# 5. Main
# -----------------------------
def main():
def main() -> None:
days = int(os.environ.get("DASHBOARD_DAYS", "30"))
df = load_hf_dataset(days=days)
local_dir = sync_from_hf("/tmp/perf-tracking")
df = load_as_dataframe(local_dir, days=days)
if df.empty:
print("No data found")
return
df = normalize(df)
# Sanity-check: log what we actually loaded
print(f"Loaded {len(df)} records across {df['model_id'].nunique()} model(s), "
f"{df['gpu_type'].nunique()} GPU type(s), "
f"date range: {df['timestamp'].min()} → {df['timestamp'].max()}")
figs = build_plots(df)
html = render_html(figs)
html = render_html(figs, days)
# Get the commit SHA from environment (passed via MODAL_ENV in your shell script)
commit_sha = os.environ.get("BUILDKITE_COMMIT", "unknown")[:7]
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
report_dir = "/root/data/perf_reports"
os.makedirs(report_dir, exist_ok=True)
# Include the SHA in the filename!
# Example: dashboard_a1b2c3d_20240426_032528.html
filename = f"dashboard_{commit_sha}_{timestamp}.html"
output_file = os.path.join(report_dir, filename)
with open(output_file, "w") as f:
with open(output_file, "w", encoding="utf-8") as f:
f.write(html)
print(f"Dashboard generated: {output_file}")
# report_dir = "/root/data/perf_reports"
# # 2. Create the directory (if it doesn't exist)
# os.makedirs(report_dir, exist_ok=True)
# # 3. Define the full file path
# output_file = os.path.join(report_dir, f"dashboard_{datetime.now().strftime('%Y%m%d_%H%M%S')}.html")
# with open(output_file, "w") as f:
# f.write(html)
# annotate_buildkite(html)
# print(f"Dashboard generated: {output_file}")
# upload_artifact(output_file)
# print(f"Dashboard uploaded as Buildkite artifact: {output_file}")
if __name__ == "__main__":
main()
+254
View File
@@ -0,0 +1,254 @@
# SPDX-License-Identifier: Apache-2.0
"""Shared HuggingFace storage utilities for performance tracking.
Provides a single place for:
- Syncing the HF dataset repo to a local directory
- Loading raw JSON records (with optional recency filter)
- Loading records as a normalized pandas DataFrame
- Uploading individual result files back to HF
- Common helpers: sanitize, safe_float
"""
import glob
import json
import os
import re
from datetime import datetime, timedelta, timezone
from typing import Any
import pandas as pd
from huggingface_hub import HfApi, snapshot_download
# ---------------------------------------------------------------------------
# Configuration — read once at import time, shared across both consumers
# ---------------------------------------------------------------------------
HF_REPO_ID: str = os.environ.get("HF_REPO_ID", "FastVideo/performance-tracking")
HF_TOKEN: str | None = os.environ.get("HF_API_KEY")
# ---------------------------------------------------------------------------
# Low-level helpers
# ---------------------------------------------------------------------------
def sanitize(value: str) -> str:
"""Return a filesystem- and HF-path-safe version of *value*."""
return re.sub(r"[^A-Za-z0-9._-]", "_", value)
def safe_float(value: Any) -> float | None:
"""Coerce *value* to float, returning None on failure."""
if value is None:
return None
try:
return float(value)
except (TypeError, ValueError):
return None
# ---------------------------------------------------------------------------
# HF I/O
# ---------------------------------------------------------------------------
def sync_from_hf(local_dir: str) -> str:
"""Download the HF dataset repo snapshot to *local_dir*.
Returns *local_dir* so callers can chain: ``load_records(sync_from_hf(...))``.
On failure (empty repo, no credentials, network error) the function logs a
warning and returns *local_dir* unchanged so the caller can still work with
whatever is already on disk.
"""
if not HF_REPO_ID:
print("hf_store: HF_REPO_ID not set, skipping sync.")
return local_dir
print(f"hf_store: syncing from {HF_REPO_ID} → {local_dir}")
try:
snapshot_download(
repo_id=HF_REPO_ID,
repo_type="dataset",
local_dir=local_dir,
token=HF_TOKEN,
allow_patterns="*.json",
)
except Exception as exc:
print(f"hf_store: sync skipped — {exc}")
return local_dir
def upload_record(local_path: str, record: dict[str, Any]) -> None:
"""Upload *local_path* to the HF repo under ``<model_id>/<filename>``.
Silently skips if HF_TOKEN is absent so local/CI runs without credentials
don't crash.
"""
if not HF_TOKEN:
print("hf_store: HF_API_KEY not set, skipping upload.")
return
model_id = record.get("model_id", "unknown")
path_in_repo = f"{sanitize(model_id)}/{os.path.basename(local_path)}"
commit_sha = (record.get("commit_sha") or "unknown")[:7]
api = HfApi(token=HF_TOKEN)
try:
api.upload_file(
path_or_fileobj=local_path,
path_in_repo=path_in_repo,
repo_id=HF_REPO_ID,
repo_type="dataset",
commit_message=f"Perf: {model_id} at {commit_sha}",
)
print(f"hf_store: uploaded → {HF_REPO_ID}/{path_in_repo}")
except Exception as exc:
print(f"hf_store: upload failed — {exc}")
# ---------------------------------------------------------------------------
# Record loading
# ---------------------------------------------------------------------------
def load_records(
local_dir: str,
*,
days: int | None = None,
successful_only: bool = False,
) -> list[dict[str, Any]]:
"""Return raw JSON dicts from *local_dir*.
Args:
local_dir: Root directory previously populated by :func:`sync_from_hf`.
days: When set, discard records whose ``timestamp`` is older than this
many days. Records with a missing/unparseable timestamp are kept.
successful_only: When True, only records with ``success=True`` are
returned. Useful when building a regression baseline.
Returns:
List of raw dicts sorted by ``timestamp`` ascending (records that could
not be parsed are silently skipped).
"""
cutoff: datetime | None = None
if days is not None:
cutoff = datetime.now(timezone.utc) - timedelta(days=days)
records: list[dict[str, Any]] = []
for path in sorted(glob.glob(os.path.join(local_dir, "**", "*.json"), recursive=True)):
try:
with open(path, encoding="utf-8") as fh:
data: dict[str, Any] = json.load(fh)
except (OSError, json.JSONDecodeError):
continue
if successful_only and not data.get("success", True):
continue
if cutoff is not None:
raw_ts = data.get("timestamp")
if raw_ts:
try:
ts = datetime.fromisoformat(raw_ts)
if ts.tzinfo is None:
ts = ts.replace(tzinfo=timezone.utc)
if ts < cutoff:
continue
except ValueError:
pass # keep records with unparseable timestamps
records.append(data)
return records
def load_records_for_model(
local_dir: str,
model_id: str,
gpu_type: str | None = None,
*,
last_n: int | None = None,
successful_only: bool = True,
) -> list[dict[str, Any]]:
"""Return records for a specific *model_id*, optionally filtered by GPU.
Args:
local_dir: Root directory previously populated by :func:`sync_from_hf`.
model_id: Matches the ``model_id`` field inside each JSON record.
gpu_type: When set, only records whose ``gpu_type`` matches are returned.
last_n: When set, return only the most recent *n* records (after all
other filters). Useful for sliding-window baseline calculations.
successful_only: Passed through to :func:`load_records`.
Returns:
List of matching dicts sorted by timestamp ascending.
"""
model_dir = os.path.join(local_dir, sanitize(model_id))
if not os.path.isdir(model_dir):
return []
records = load_records(model_dir, successful_only=successful_only)
if gpu_type is not None:
records = [r for r in records if r.get("gpu_type") == gpu_type]
if last_n is not None:
records = records[-last_n:]
return records
# ---------------------------------------------------------------------------
# DataFrame helpers (dashboard / analytics consumers)
# ---------------------------------------------------------------------------
_NUMERIC_COLS = ("latency", "throughput", "memory")
def normalize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
"""Apply standard type coercions to a raw records DataFrame.
- Parses ``timestamp`` to UTC-aware datetime.
- Coerces ``latency``, ``throughput``, ``memory`` to float.
- Adds a ``config_id`` column (first 7 chars of ``commit_sha``).
Returns the mutated DataFrame (also modifies in place for efficiency).
"""
if df.empty:
return df
df["timestamp"] = pd.to_datetime(df["timestamp"], utc=True, errors="coerce")
df["config_id"] = df.get("commit_sha", pd.Series(dtype=str)).fillna("unknown").str[:7]
for col in _NUMERIC_COLS:
if col in df.columns:
df[col] = pd.to_numeric(df[col], errors="coerce")
return df
def load_as_dataframe(
local_dir: str,
*,
days: int | None = None,
successful_only: bool = False,
) -> pd.DataFrame:
"""Load and normalize records from *local_dir* into a pandas DataFrame.
Combines :func:`load_records` + :func:`normalize_dataframe` into a single
call for consumers (e.g. the dashboard) that work exclusively with
DataFrames.
Args:
local_dir: Root directory previously populated by :func:`sync_from_hf`.
days: Passed through to :func:`load_records`.
successful_only: Passed through to :func:`load_records`.
Returns:
Normalized DataFrame, or an empty DataFrame if no records were found.
"""
records = load_records(local_dir, days=days, successful_only=successful_only)
if not records:
return pd.DataFrame()
df = pd.DataFrame(records)
return normalize_dataframe(df)