diff --git a/.github/workflows/nightly-test-nvidia.yml b/.github/workflows/nightly-test-nvidia.yml index bef286886..3a4472c42 100644 --- a/.github/workflows/nightly-test-nvidia.yml +++ b/.github/workflows/nightly-test-nvidia.yml @@ -389,6 +389,7 @@ jobs: env: SGLANG_DIFFUSION_SLACK_TOKEN: ${{ secrets.SGLANG_DIFFUSION_SLACK_TOKEN }} GITHUB_RUN_ID: ${{ github.run_id }} + GPU_CONFIG: "1-gpu-runner" timeout-minutes: 60 run: | @@ -398,6 +399,24 @@ jobs: --partition-id ${{ matrix.part }} \ --total-partitions 2 + - name: Collect diffusion performance metrics + if: always() + run: | + python3 scripts/ci/save_diffusion_metrics.py \ + --gpu-config 1-gpu-runner \ + --run-id ${{ github.run_id }} \ + --output python/diffusion-metrics-1gpu-partition-${{ matrix.part }}.json \ + --results-json python/diffusion-results.json + + - name: Upload diffusion metrics + if: always() + uses: actions/upload-artifact@v4 + with: + name: diffusion-metrics-1gpu-partition-${{ matrix.part }} + path: python/diffusion-metrics-1gpu-partition-${{ matrix.part }}.json + retention-days: 90 + if-no-files-found: ignore + nightly-test-multimodal-server-2-gpu: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-multimodal-server-2-gpu') @@ -422,6 +441,7 @@ jobs: env: SGLANG_DIFFUSION_SLACK_TOKEN: ${{ secrets.SGLANG_DIFFUSION_SLACK_TOKEN }} GITHUB_RUN_ID: ${{ github.run_id }} + GPU_CONFIG: "2-gpu-runner" timeout-minutes: 60 run: | @@ -431,6 +451,24 @@ jobs: --partition-id ${{ matrix.part }} \ --total-partitions 2 + - name: Collect diffusion performance metrics + if: always() + run: | + python3 scripts/ci/save_diffusion_metrics.py \ + --gpu-config 2-gpu-runner \ + --run-id ${{ github.run_id }} \ + --output python/diffusion-metrics-2gpu-partition-${{ matrix.part }}.json \ + --results-json python/diffusion-results.json + + - name: Upload diffusion metrics + if: always() + uses: actions/upload-artifact@v4 + with: + name: diffusion-metrics-2gpu-partition-${{ matrix.part }} + path: python/diffusion-metrics-2gpu-partition-${{ matrix.part }}.json + retention-days: 90 + if-no-files-found: ignore + # B200 Performance tests - 4 GPU nightly-test-perf-4-gpu-b200: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-perf-4-gpu-b200') @@ -475,12 +513,14 @@ jobs: cd test python3 run_suite.py --hw cuda --suite nightly-8-gpu-b200 --nightly --continue-on-error --timeout-per-file 2400 - # Consolidate performance metrics from all 8-GPU jobs + # Consolidate performance metrics from all jobs consolidate-metrics: if: github.repository == 'sgl-project/sglang' && always() needs: - nightly-test-general-8-gpu-h200 - nightly-test-general-8-gpu-b200 + - nightly-test-multimodal-server-1-gpu + - nightly-test-multimodal-server-2-gpu runs-on: ubuntu-latest steps: - name: Checkout code @@ -491,7 +531,7 @@ jobs: - name: Download all partition metrics uses: actions/download-artifact@v4 with: - pattern: metrics-* + pattern: "*metrics-*" path: metrics/ merge-multiple: true diff --git a/docs/performance_dashboard/app.js b/docs/performance_dashboard/app.js index 0d4cb13d2..8bfb12b2e 100644 --- a/docs/performance_dashboard/app.js +++ b/docs/performance_dashboard/app.js @@ -14,12 +14,17 @@ let currentMetricType = 'throughput'; // throughput, latency, ttft, inputThrough // Metric type definitions const metricTypes = { - throughput: { label: 'Overall Throughput', unit: 'tokens/sec', field: 'throughput' }, - outputThroughput: { label: 'Output Throughput', unit: 'tokens/sec', field: 'outputThroughput' }, - inputThroughput: { label: 'Input Throughput', unit: 'tokens/sec', field: 'inputThroughput' }, - latency: { label: 'Latency', unit: 'ms', field: 'latency' }, - ttft: { label: 'Time to First Token', unit: 'ms', field: 'ttft' }, - accLength: { label: 'Accept Length', unit: 'tokens', field: 'accLength', filterInvalid: true } + // Text/VLM metrics + throughput: { label: 'Overall Throughput', unit: 'tokens/sec', field: 'throughput', type: 'text' }, + outputThroughput: { label: 'Output Throughput', unit: 'tokens/sec', field: 'outputThroughput', type: 'text' }, + inputThroughput: { label: 'Input Throughput', unit: 'tokens/sec', field: 'inputThroughput', type: 'text' }, + latency: { label: 'Latency', unit: 'ms', field: 'latency', type: 'text' }, + ttft: { label: 'Time to First Token', unit: 'ms', field: 'ttft', type: 'text' }, + accLength: { label: 'Accept Length', unit: 'tokens', field: 'accLength', filterInvalid: true, type: 'text' }, + // Diffusion metrics + e2eMs: { label: 'End-to-End Time', unit: 'ms', field: 'e2e_ms', type: 'diffusion' }, + avgDenoiseMs: { label: 'Avg Denoise Time', unit: 'ms', field: 'avg_denoise_ms', type: 'diffusion' }, + medianDenoiseMs: { label: 'Median Denoise Time', unit: 'ms', field: 'median_denoise_ms', type: 'diffusion' } }; // Chart.js default configuration for dark theme @@ -142,32 +147,51 @@ async function fetchMetricsForRun(run) { } } +// Helper function to detect if result is diffusion type +function isDiffusionResult(result) { + return result.test_type === 'diffusion' || (result.tests && !result.benchmarks); +} + // Populate filter dropdowns function populateFilters() { const gpuConfigs = new Set(); const models = new Set(); + const testNames = new Set(); // For diffusion tests const batchSizes = new Set(); const ioLengths = new Set(); allMetricsData.forEach(run => { run.results.forEach(result => { gpuConfigs.add(result.gpu_config); - models.add(result.model); - // Try new structure first (benchmarks_by_io_len), fall back to flat benchmarks - if (result.benchmarks_by_io_len) { - Object.entries(result.benchmarks_by_io_len).forEach(([ioKey, ioData]) => { - ioLengths.add(ioKey); - ioData.benchmarks.forEach(bench => { - batchSizes.add(bench.batch_size); + + // Handle diffusion results + if (isDiffusionResult(result)) { + models.add(result.test_suite || 'diffusion'); + if (result.tests) { + result.tests.forEach(test => { + testNames.add(test.test_name); }); - }); - } else if (result.benchmarks) { - result.benchmarks.forEach(bench => { - batchSizes.add(bench.batch_size); - if (bench.input_len && bench.output_len) { - ioLengths.add(`${bench.input_len}_${bench.output_len}`); - } - }); + } + } + // Handle text/VLM results + else { + models.add(result.model); + // Try new structure first (benchmarks_by_io_len), fall back to flat benchmarks + if (result.benchmarks_by_io_len) { + Object.entries(result.benchmarks_by_io_len).forEach(([ioKey, ioData]) => { + ioLengths.add(ioKey); + ioData.benchmarks.forEach(bench => { + batchSizes.add(bench.batch_size); + }); + }); + } else if (result.benchmarks) { + result.benchmarks.forEach(bench => { + batchSizes.add(bench.batch_size); + if (bench.input_len && bench.output_len) { + ioLengths.add(`${bench.input_len}_${bench.output_len}`); + } + }); + } } }); }); @@ -345,7 +369,16 @@ function createMetricTabs() { const tabsContainer = document.getElementById('metric-tabs'); tabsContainer.innerHTML = ''; - Object.entries(metricTypes).forEach(([key, metric], index) => { + // Detect if current data is diffusion or text + const isDiffusion = detectCurrentDataType() === 'diffusion'; + const dataType = isDiffusion ? 'diffusion' : 'text'; + + // Filter metrics based on data type + const relevantMetrics = Object.entries(metricTypes).filter(([key, metric]) => + metric.type === dataType + ); + + relevantMetrics.forEach(([key, metric], index) => { const tab = document.createElement('div'); tab.className = index === 0 ? 'tab active' : 'tab'; tab.textContent = metric.label; @@ -353,6 +386,31 @@ function createMetricTabs() { tab.onclick = () => selectMetricTab(key, tab); tabsContainer.appendChild(tab); }); + + // Set initial metric type + if (relevantMetrics.length > 0) { + currentMetricType = relevantMetrics[0][0]; + } +} + +function detectCurrentDataType() { + // Check if currently selected model/GPU config has diffusion data + const gpuFilter = document.getElementById('gpu-filter')?.value; + const modelFilter = currentModel; + + if (!gpuFilter || !modelFilter) return 'text'; + + for (const run of allMetricsData) { + for (const result of run.results) { + if (result.gpu_config === gpuFilter) { + const resultModel = result.test_suite || result.model; + if (resultModel === modelFilter && isDiffusionResult(result)) { + return 'diffusion'; + } + } + } + } + return 'text'; } function selectMetricTab(metricKey, tabElement) { @@ -374,6 +432,8 @@ function handleModelFilterChange(model) { updateVariantFilter(); // Update IO length filter based on new model selection updateIoLenFilter(); + // Recreate metric tabs in case data type changed (text vs diffusion) + createMetricTabs(); updateCharts(); } @@ -383,6 +443,8 @@ function handleGpuFilterChange() { updateVariantFilter(); // Update IO length filter based on new GPU selection updateIoLenFilter(); + // Recreate metric tabs in case data type changed (text vs diffusion) + createMetricTabs(); updateCharts(); } @@ -518,6 +580,7 @@ function prepareChartData(gpuFilter, modelFilter, variantFilter, ioLenFilter, ba // Prepare chart data grouped by batch size - each batch size is a separate series function prepareChartDataByBatch(gpuFilter, modelFilter, variantFilter, ioLenFilter, batchFilter) { const batchDataMap = new Map(); // batch_size -> Map of variant -> data + const testDataMap = new Map(); // For diffusion: test_name -> data allMetricsData.forEach(run => { const runDate = new Date(run.run_date); @@ -525,6 +588,37 @@ function prepareChartDataByBatch(gpuFilter, modelFilter, variantFilter, ioLenFil run.results.forEach(result => { // Apply filters - GPU and Model are required (no "all" option) if (result.gpu_config !== gpuFilter) return; + + // Handle diffusion results + if (isDiffusionResult(result)) { + const resultModel = result.test_suite || 'diffusion'; + if (resultModel !== modelFilter) return; + + if (result.tests) { + result.tests.forEach(test => { + const testName = test.test_name; + if (!testDataMap.has(testName)) { + testDataMap.set(testName, { + label: testName, + data: [], + model: resultModel, + testName: testName + }); + } + + testDataMap.get(testName).data.push({ + x: runDate, + e2e_ms: test.e2e_ms, + avg_denoise_ms: test.avg_denoise_ms, + median_denoise_ms: test.median_denoise_ms, + runId: run.run_id + }); + }); + } + return; + } + + // Handle text/VLM results if (result.model !== modelFilter) return; if (variantFilter !== 'all' && result.variant !== variantFilter) return; @@ -622,6 +716,17 @@ function prepareChartDataByBatch(gpuFilter, modelFilter, variantFilter, ioLenFil // Sort data points by date and convert to array format const result = {}; + + // For diffusion data, use test names as "batch sizes" + if (testDataMap.size > 0) { + testDataMap.forEach((series, testName) => { + series.data.sort((a, b) => a.x - b.x); + result[testName] = [series]; // Each test is its own series + }); + return result; + } + + // For text/VLM data, use batch sizes batchDataMap.forEach((variantMap, batchSize) => { variantMap.forEach(series => { series.data.sort((a, b) => a.x - b.x); @@ -642,7 +747,16 @@ function updateMetricChart(chartDataByBatch, metricType) { activeCharts = []; const metric = metricTypes[metricType]; - const batchSizes = Object.keys(chartDataByBatch).sort((a, b) => parseInt(a) - parseInt(b)); + const isDiffusion = metric.type === 'diffusion'; + + // For diffusion, keys are test names; for text, keys are batch sizes + const keys = Object.keys(chartDataByBatch); + if (!isDiffusion) { + keys.sort((a, b) => parseInt(a) - parseInt(b)); + } else { + keys.sort(); // Alphabetical sort for test names + } + const batchSizes = keys; // Keep variable name for compatibility if (batchSizes.length === 0) { container.innerHTML = '
No data available for the selected filters
'; @@ -682,7 +796,8 @@ function updateMetricChart(chartDataByBatch, metricType) { const title = document.createElement('div'); title.className = 'batch-chart-title'; - title.textContent = `Batch Size: ${batchSize}`; + // For diffusion, show test name; for text, show batch size + title.textContent = isDiffusion ? `Test: ${batchSize}` : `Batch Size: ${batchSize}`; chartWrapper.appendChild(title); const chartContainer = document.createElement('div'); diff --git a/python/sglang/multimodal_gen/test/server/conftest.py b/python/sglang/multimodal_gen/test/server/conftest.py index 96b49591b..9b5ddde58 100644 --- a/python/sglang/multimodal_gen/test/server/conftest.py +++ b/python/sglang/multimodal_gen/test/server/conftest.py @@ -1,4 +1,97 @@ -_GLOBAL_PERF_RESULTS = [] +import os + +import pytest + +print("[CONFTEST] Loading conftest.py at import time") + + +def pytest_configure(config): + """ + Create the perf results StashKey once and store it in config. + This hook runs once per test session, before module double-import issues. + """ + if not hasattr(config, "_diffusion_perf_key"): + config._diffusion_perf_key = pytest.StashKey[list]() + print(f"[CONFTEST] Created perf_results_key: {config._diffusion_perf_key}") + + +def add_perf_results(config, results: list): + """Add performance results to the shared stash.""" + # Get the shared key from config (created once in pytest_configure) + key = config._diffusion_perf_key + existing = config.stash.get(key, []) + existing.extend(results) + config.stash[key] = existing + print(f"[CONFTEST] Added {len(results)} results, total now: {len(existing)}") + + +@pytest.fixture(scope="session") +def perf_config(request): + """Provide access to pytest config for storing perf results.""" + return request.config + + +def _write_github_step_summary(content: str): + """Write content to GitHub Step Summary if available.""" + summary_file = os.environ.get("GITHUB_STEP_SUMMARY") + if summary_file: + with open(summary_file, "a") as f: + f.write(content) + + +def _write_results_json(results: list, output_path: str = "diffusion-results.json"): + """Write performance results to JSON file for CI artifact collection.""" + import json + + try: + with open(output_path, "w") as f: + json.dump(results, f, indent=2) + print(f"[CONFTEST] Wrote results to {output_path}") + except Exception as e: + print(f"[CONFTEST] Failed to write results JSON: {e}") + + +def _generate_diffusion_markdown_report(results: list) -> str: + """Generate a markdown report for diffusion performance results.""" + if not results: + return "" + + gpu_config = os.environ.get("GPU_CONFIG", "") + header = "## Diffusion Performance Summary" + if gpu_config: + header += f" [{gpu_config}]" + header += "\n\n" + + # Main performance table + markdown = header + markdown += "| Test Suite | Test Name | Modality | E2E (ms) | Avg Denoise (ms) | Median Denoise (ms) |\n" + markdown += "| ---------- | --------- | -------- | -------- | ---------------- | ------------------- |\n" + + for entry in sorted(results, key=lambda x: (x["class_name"], x["test_name"])): + modality = entry.get("modality", "image") + markdown += ( + f"| {entry['class_name']} | {entry['test_name']} | {modality} | " + f"{entry['e2e_ms']:.2f} | {entry['avg_denoise_ms']:.2f} | " + f"{entry['median_denoise_ms']:.2f} |\n" + ) + + # Video-specific metrics table (if any video tests) + video_results = [r for r in results if r.get("modality") == "video"] + if video_results: + markdown += "\n### Video Generation Metrics\n\n" + markdown += "| Test Name | FPS | Total Frames | Avg Frame Time (ms) |\n" + markdown += "| --------- | --- | ------------ | ------------------- |\n" + for entry in video_results: + fps = entry.get("frames_per_second", "N/A") + frames = entry.get("total_frames", "N/A") + avg_frame = entry.get("avg_frame_time_ms", "N/A") + if isinstance(fps, float): + fps = f"{fps:.2f}" + if isinstance(avg_frame, float): + avg_frame = f"{avg_frame:.2f}" + markdown += f"| {entry['test_name']} | {fps} | {frames} | {avg_frame} |\n" + + return markdown def pytest_sessionfinish(session): @@ -6,9 +99,15 @@ def pytest_sessionfinish(session): This hook is called by pytest at the end of the entire test session. It prints a consolidated summary of all performance results. """ - if not _GLOBAL_PERF_RESULTS: + # Get results from stash using the shared key from config + key = session.config._diffusion_perf_key + results = session.config.stash.get(key, []) + print(f"\n[DEBUG] pytest_sessionfinish called, has {len(results)} entries") + if not results: + print("[DEBUG] No results collected, skipping summary output") return + # Print to stdout (existing behavior) print("\n\n" + "=" * 35 + " Performance Summary " + "=" * 35) print( f"{'Test Suite':<30} | {'Test Name':<20} | {'E2E (ms)':>12} | {'Avg Denoise (ms)':>18} | {'Median Denoise (ms)':>20}" @@ -25,7 +124,7 @@ def pytest_sessionfinish(session): + "-" * 20 ) - for entry in sorted(_GLOBAL_PERF_RESULTS, key=lambda x: x["class_name"]): + for entry in sorted(results, key=lambda x: x["class_name"]): print( f"{entry['class_name']:<30} | {entry['test_name']:<20} | {entry['e2e_ms']:>12.2f} | " f"{entry['avg_denoise_ms']:>18.2f} | {entry['median_denoise_ms']:>20.2f}" @@ -34,7 +133,7 @@ def pytest_sessionfinish(session): print("=" * 91) print("\n\n" + "=" * 36 + " Detailed Reports " + "=" * 37) - for entry in sorted(_GLOBAL_PERF_RESULTS, key=lambda x: x["class_name"]): + for entry in sorted(results, key=lambda x: x["class_name"]): print(f"\n--- Details for {entry['class_name']} / {entry['test_name']} ---") stage_report = ", ".join( f"{name}:{duration:.2f}ms" @@ -51,3 +150,11 @@ def pytest_sessionfinish(session): ) print(f" Sampled Steps: {step_report}") print("=" * 91) + + # Write to GitHub Step Summary (new behavior for CI monitoring) + markdown_report = _generate_diffusion_markdown_report(results) + if markdown_report: + _write_github_step_summary(markdown_report) + + # Write results to JSON file for CI artifact collection + _write_results_json(results) diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 9b7b6f4eb..046495a66 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -18,7 +18,7 @@ from openai import OpenAI from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord -from sglang.multimodal_gen.test.server.conftest import _GLOBAL_PERF_RESULTS +from sglang.multimodal_gen.test.server import conftest from sglang.multimodal_gen.test.server.test_server_utils import ( VALIDATOR_REGISTRY, PerformanceValidator, @@ -134,6 +134,7 @@ class DiffusionServerBase: _perf_results: list[dict[str, Any]] = [] _improved_baselines: list[dict[str, Any]] = [] + _pytest_config = None # Store pytest config for stash access @classmethod def setup_class(cls): @@ -142,9 +143,21 @@ class DiffusionServerBase: @classmethod def teardown_class(cls): - for result in cls._perf_results: - result["class_name"] = cls.__name__ - _GLOBAL_PERF_RESULTS.append(result) + print( + f"\n[DEBUG teardown_class] Called for {cls.__name__}, _perf_results has {len(cls._perf_results)} entries" + ) + if cls._pytest_config: + # Add results to pytest stash (shared across all import contexts) + for result in cls._perf_results: + result["class_name"] = cls.__name__ + conftest.add_perf_results(cls._pytest_config, cls._perf_results) + print( + f"[DEBUG teardown_class] Added {len(cls._perf_results)} results to stash" + ) + else: + print( + "[DEBUG teardown_class] No pytest_config available, skipping stash update" + ) if cls._improved_baselines: import json @@ -160,6 +173,11 @@ Consider updating perf_baselines.json with the snippets below: ) print(output) + @pytest.fixture(autouse=True) + def _capture_pytest_config(self, request): + """Capture pytest config for use in teardown_class.""" + self.__class__._pytest_config = request.config + def _client(self, ctx: ServerContext) -> OpenAI: """Get OpenAI client for the server.""" return OpenAI( @@ -261,6 +279,9 @@ Consider updating perf_baselines.json with the snippets below: ) self.__class__._perf_results.append(result) + print( + f"[DEBUG _validate_and_record] Appended result for {case.id}, class {self.__class__.__name__} now has {len(self.__class__._perf_results)} results" + ) def _check_for_improvement( self, diff --git a/scripts/ci/merge_metrics.py b/scripts/ci/merge_metrics.py index 0eb7cad8b..309ae8079 100755 --- a/scripts/ci/merge_metrics.py +++ b/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: diff --git a/scripts/ci/save_diffusion_metrics.py b/scripts/ci/save_diffusion_metrics.py new file mode 100755 index 000000000..c851666c3 --- /dev/null +++ b/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()