File size: 4,857 Bytes
61ba51e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | #!/usr/bin/env python3
"""Collect and save diffusion performance metrics for artifact collection in CI.
This script reads diffusion test results from the pytest stash and saves them
with metadata for the performance dashboard.
Usage:
python3 scripts/ci/save_diffusion_metrics.py \
--gpu-config 1-gpu-runner \
--run-id 12345678 \
--output test/diffusion-metrics-1gpu.json \
--results-json test/diffusion-results.json
"""
import argparse
import json
import os
import sys
from datetime import datetime, timezone
def load_diffusion_results(results_file: str) -> list[dict]:
"""Load diffusion performance results from JSON file."""
if not os.path.exists(results_file):
print(f"Warning: Results file not found: {results_file}")
return []
try:
with open(results_file, "r", encoding="utf-8") as f:
data = json.load(f)
return data if isinstance(data, list) else [data]
except (json.JSONDecodeError, OSError) as e:
print(f"Warning: Failed to parse {results_file}: {e}")
return []
def transform_diffusion_result(result: dict, gpu_config: str) -> dict:
"""Transform a diffusion result to match dashboard expectations.
Dashboard expects:
- Separate test_name, class_name
- Numeric metrics in consistent units
- Optional modality field
"""
return {
"test_name": result.get("test_name"),
"class_name": result.get("class_name"),
"modality": result.get("modality", "image"),
"e2e_ms": result.get("e2e_ms"),
"avg_denoise_ms": result.get("avg_denoise_ms"),
"median_denoise_ms": result.get("median_denoise_ms"),
"stage_metrics": result.get("stage_metrics", {}),
"sampled_steps": result.get("sampled_steps", {}),
# Video-specific metrics (if present)
"frames_per_second": result.get("frames_per_second"),
"total_frames": result.get("total_frames"),
"avg_frame_time_ms": result.get("avg_frame_time_ms"),
}
def group_results_by_class(results: list[dict], gpu_config: str) -> list[dict]:
"""Group diffusion results by test class (suite).
Returns list with one entry per test class, containing all tests in that class.
"""
groups = {}
for result in results:
class_name = result.get("class_name", "unknown")
if class_name not in groups:
groups[class_name] = {
"gpu_config": gpu_config,
"test_suite": class_name,
"tests": [],
}
transformed = transform_diffusion_result(result, gpu_config)
groups[class_name]["tests"].append(transformed)
return list(groups.values())
def save_metrics(
gpu_config: str,
run_id: str,
output_file: str,
results_file: str,
) -> bool:
"""Collect diffusion metrics and save to output file."""
timestamp = datetime.now(timezone.utc).isoformat()
# Load diffusion results
raw_results = load_diffusion_results(results_file)
print(f"Loaded {len(raw_results)} diffusion test result(s)")
# Group by test class
grouped = group_results_by_class(raw_results, gpu_config)
# Create metrics structure
metrics = {
"run_id": run_id,
"timestamp": timestamp,
"gpu_config": gpu_config,
"test_type": "diffusion",
"results": grouped,
}
# Ensure output directory exists and write output
try:
os.makedirs(os.path.dirname(output_file) or ".", exist_ok=True)
with open(output_file, "w", encoding="utf-8") as f:
json.dump(metrics, f, indent=2)
if not raw_results:
print(f"Created empty metrics file: {output_file}")
else:
print(f"Saved diffusion metrics to: {output_file}")
return True
except OSError as e:
print(f"Error writing metrics file: {e}")
return False
def main():
parser = argparse.ArgumentParser(
description="Collect diffusion performance metrics from test results"
)
parser.add_argument(
"--gpu-config",
required=True,
help="GPU configuration (e.g., 1-gpu-runner, 2-gpu-runner)",
)
parser.add_argument(
"--run-id",
required=True,
help="GitHub Actions run ID",
)
parser.add_argument(
"--output",
required=True,
help="Output file path for metrics JSON",
)
parser.add_argument(
"--results-json",
required=True,
help="Path to diffusion results JSON file",
)
args = parser.parse_args()
success = save_metrics(
gpu_config=args.gpu_config,
run_id=args.run_id,
output_file=args.output,
results_file=args.results_json,
)
sys.exit(0 if success else 1)
if __name__ == "__main__":
main()
|