fix: adding performance logging for nightly diffusion (#18023)

scripts/ci/merge_metrics.py

@@ -22,8 +22,14 @@ from datetime import datetime, timezone
 
 def find_partition_files(input_dir: str) -> list[str]:
     """Find all partition metric files in the input directory."""
-    pattern = os.path.join(input_dir, "**/metrics-*.json")
-    return glob.glob(pattern, recursive=True)
+    patterns = [
+        os.path.join(input_dir, "**/metrics-*.json"),
+        os.path.join(input_dir, "**/diffusion-metrics-*.json"),
+    ]
+    files = set()
+    for pattern in patterns:
+        files.update(glob.glob(pattern, recursive=True))
+    return list(files)
 
 
 def load_partition_metrics(filepath: str) -> dict | None:

scripts/ci/save_diffusion_metrics.py

@@ -0,0 +1,163 @@
+#!/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()