chore: import upstream snapshot with attribution

python/sglang/test/nightly_utils.py

@@ -0,0 +1,364 @@
+"""Utilities for running nightly performance benchmarks with profiling."""
+
+import json
+import os
+import subprocess
+import time
+from typing import List, Optional, Tuple
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.nightly_bench_utils import BenchmarkResult, generate_markdown_report
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    is_in_ci,
+    popen_launch_server,
+    write_github_step_summary,
+)
+
+
+class NightlyBenchmarkRunner:
+    """Helper class for running nightly performance benchmarks with profiling.
+
+    This class encapsulates common patterns used across nightly performance tests,
+    including profile directory management, benchmark command construction,
+    result parsing, and report generation.
+    """
+
+    def __init__(
+        self,
+        profile_dir: str,
+        test_name: str,
+        base_url: str,
+        gpu_config: str = None,
+    ):
+        """Initialize the benchmark runner.
+
+        Args:
+            profile_dir: Directory to store performance profiles
+            test_name: Name of the test (used for reporting)
+            base_url: Base URL for the server
+            gpu_config: Optional GPU configuration string (e.g., "2-gpu-h100", "8-gpu-b200")
+        """
+        self.profile_dir = profile_dir
+        self.test_name = test_name
+        self.base_url = base_url
+        self.gpu_config = gpu_config or os.environ.get("GPU_CONFIG", "")
+
+        # Include GPU config in report header if available
+        header = f"## {test_name}"
+        if self.gpu_config:
+            header += f" ({self.gpu_config})"
+        header += "\n"
+        self.full_report = header + BenchmarkResult.help_str()
+
+    def setup_profile_directory(self) -> None:
+        """Create the profile directory if it doesn't exist."""
+        os.makedirs(self.profile_dir, exist_ok=True)
+
+    def generate_profile_filename(
+        self, model_path: str, variant: str = ""
+    ) -> Tuple[str, str]:
+        """Generate unique profile filename and path for the model.
+
+        Args:
+            model_path: Path to the model (e.g., "deepseek-ai/DeepSeek-V3.1")
+            variant: Optional variant suffix (e.g., "basic", "mtp", "dsa")
+
+        Returns:
+            Tuple of (profile_path_prefix, json_output_file)
+        """
+        timestamp = int(time.time())
+        model_safe_name = model_path.replace("/", "_")
+
+        # Build filename with optional variant
+        if variant:
+            profile_filename = f"{model_safe_name}_{variant}_{timestamp}"
+            json_filename = f"results_{model_safe_name}_{variant}_{timestamp}.json"
+        else:
+            profile_filename = f"{model_safe_name}_{timestamp}"
+            json_filename = f"results_{model_safe_name}_{timestamp}.json"
+
+        profile_path_prefix = os.path.join(self.profile_dir, profile_filename)
+
+        return profile_path_prefix, json_filename
+
+    def build_benchmark_command(
+        self,
+        model_path: str,
+        batch_sizes: List[int],
+        input_lens: Tuple[int, ...],
+        output_lens: Tuple[int, ...],
+        profile_path_prefix: str,
+        json_output_file: str,
+        extra_args: Optional[List[str]] = None,
+        server_args: Optional[List[str]] = None,
+        enable_profile: bool = True,
+    ) -> List[str]:
+        """Build the benchmark command with all required arguments.
+
+        Args:
+            model_path: Path to the model
+            batch_sizes: List of batch sizes to test
+            input_lens: Tuple of input lengths to test
+            output_lens: Tuple of output lengths to test
+            profile_path_prefix: Prefix for profile output files
+            json_output_file: Path to JSON output file
+            extra_args: Optional extra arguments to append to command
+            server_args: Optional server launch arguments to record in metrics
+            enable_profile: Whether to enable profiling (default True for NVIDIA)
+
+        Returns:
+            List of command arguments ready for subprocess.run()
+        """
+        command = [
+            "python3",
+            "-m",
+            "sglang.bench_one_batch_server",
+            "--model",
+            model_path,
+            "--base-url",
+            self.base_url,
+            "--batch-size",
+            *[str(x) for x in batch_sizes],
+            "--input-len",
+            *[str(x) for x in input_lens],
+            "--output-len",
+            *[str(x) for x in output_lens],
+            "--show-report",
+            f"--pydantic-result-filename={json_output_file}",
+            "--no-append-to-github-summary",
+            "--trust-remote-code",
+        ]
+
+        # Add profiling flags only if enabled (disabled for AMD tests)
+        if enable_profile and profile_path_prefix:
+            command.extend(
+                [
+                    "--profile",
+                    "--profile-by-stage",
+                    "--profile-output-dir",
+                    profile_path_prefix,
+                ]
+            )
+
+        if extra_args:
+            command.extend(extra_args)
+
+        # Record server launch arguments in metrics for tracking configurations
+        if server_args:
+            command.append("--server-args-for-metrics")
+            command.extend(server_args)
+
+        return command
+
+    def run_benchmark_command(
+        self, command: List[str], model_description: str = ""
+    ) -> Tuple[subprocess.CompletedProcess, bool]:
+        """Execute the benchmark command and return the result.
+
+        Args:
+            command: Command to execute
+            model_description: Description for logging (e.g., "model_name (variant)")
+
+        Returns:
+            Tuple of (CompletedProcess, success_bool)
+        """
+        print(f"Running command: {' '.join(command)}")
+        result = subprocess.run(command, capture_output=True, text=True)
+
+        if result.returncode != 0:
+            desc = model_description or "benchmark"
+            print(f"Error running benchmark for {desc}:")
+            print(result.stderr)
+            return result, False
+
+        return result, True
+
+    def load_benchmark_results(
+        self, json_output_file: str, model_description: str = ""
+    ) -> Tuple[List[BenchmarkResult], bool]:
+        """Load and parse benchmark results from JSON file.
+
+        Args:
+            json_output_file: Path to JSON output file
+            model_description: Description for logging
+
+        Returns:
+            Tuple of (list of BenchmarkResult objects, success_bool)
+        """
+        benchmark_results = []
+
+        if not os.path.exists(json_output_file):
+            desc = model_description or "model"
+            print(f"Warning: JSON output file {json_output_file} not found for {desc}")
+            return benchmark_results, False
+
+        try:
+            with open(json_output_file, "r") as f:
+                json_data = json.load(f)
+
+            # Convert JSON data to BenchmarkResult objects
+            for data in json_data:
+                benchmark_result = BenchmarkResult(**data)
+                benchmark_results.append(benchmark_result)
+
+            print(
+                f"Loaded {len(benchmark_results)} benchmark results from {json_output_file}"
+            )
+
+            # Note: JSON files are preserved for metrics collection by CI scripts
+            # They will be collected by scripts/ci/utils/save_metrics.py
+
+            return benchmark_results, True
+
+        except Exception as e:
+            desc = model_description or "model"
+            print(f"Error loading benchmark results for {desc}: {e}")
+            return benchmark_results, False
+
+    def run_benchmark_for_model(
+        self,
+        model_path: str,
+        batch_sizes: List[int],
+        input_lens: Tuple[int, ...],
+        output_lens: Tuple[int, ...],
+        other_args: Optional[List[str]] = None,
+        variant: str = "",
+        extra_bench_args: Optional[List[str]] = None,
+        enable_profile: bool = True,
+        timeout: Optional[int] = None,
+        env: Optional[dict] = None,
+    ) -> Tuple[List[BenchmarkResult], bool, Optional[float]]:
+        """Run a complete benchmark for a single model with server management.
+
+        This method handles:
+        - Server launch and cleanup
+        - Profile filename generation
+        - Benchmark command construction and execution
+        - Result loading and parsing
+        - Fetching speculative decoding accept length (for MTP/EAGLE)
+
+        Args:
+            model_path: Path to the model
+            batch_sizes: List of batch sizes to test
+            input_lens: Tuple of input lengths
+            output_lens: Tuple of output lengths
+            other_args: Arguments to pass to server launch
+            variant: Optional variant suffix (e.g., "basic", "mtp")
+            extra_bench_args: Extra arguments for the benchmark command
+            enable_profile: Whether to enable profiling (default True for NVIDIA)
+            timeout: Optional timeout for server launch (defaults to DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH)
+            env: Environment dict for subprocess
+
+        Returns:
+            Tuple of (list of BenchmarkResult objects, success_bool, avg_spec_accept_length or None)
+        """
+        benchmark_results = []
+        avg_spec_accept_length = None
+        model_description = f"{model_path}" + (f" ({variant})" if variant else "")
+
+        process = None
+        try:
+            # Launch server
+            process = popen_launch_server(
+                model=model_path,
+                base_url=self.base_url,
+                other_args=other_args or [],
+                timeout=(
+                    timeout
+                    if timeout is not None
+                    else DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+                ),
+                env=env,
+            )
+
+            # Generate filenames
+            profile_path_prefix, json_output_file = self.generate_profile_filename(
+                model_path, variant
+            )
+
+            # Build and run benchmark command
+            # Prepare extra args with run_name if variant is specified
+            bench_args = list(extra_bench_args) if extra_bench_args else []
+            if variant:
+                bench_args.extend(["--run-name", variant])
+
+            command = self.build_benchmark_command(
+                model_path,
+                batch_sizes,
+                input_lens,
+                output_lens,
+                profile_path_prefix,
+                json_output_file,
+                extra_args=bench_args,
+                server_args=other_args,
+                enable_profile=enable_profile,
+            )
+
+            result, cmd_success = self.run_benchmark_command(command, model_description)
+
+            if not cmd_success:
+                return benchmark_results, False, None
+
+            # Load results
+            benchmark_results, load_success = self.load_benchmark_results(
+                json_output_file, model_description
+            )
+
+            # Fetch speculative decoding accept length before killing server
+            avg_spec_accept_length = self._get_spec_accept_length()
+
+            return benchmark_results, load_success, avg_spec_accept_length
+
+        finally:
+            # Always clean up server process
+            if process is not None:
+                kill_process_tree(process.pid)
+
+    def _get_spec_accept_length(self) -> Optional[float]:
+        """Query the server for avg_spec_accept_length metric.
+
+        Returns:
+            The average speculative decoding accept length, or None if not available.
+        """
+        try:
+            response = requests.get(f"{self.base_url}/server_info", timeout=10)
+            if response.status_code == 200:
+                server_info = response.json()
+                internal_states = server_info.get("internal_states", [])
+                if internal_states and len(internal_states) > 0:
+                    accept_length = internal_states[0].get("avg_spec_accept_length")
+                    if accept_length is not None:
+                        print(f"  avg_spec_accept_length={accept_length:.2f}")
+                        return accept_length
+        except Exception as e:
+            print(f"  Warning: Could not fetch spec accept length: {e}")
+        return None
+
+    def add_report(
+        self, results: List[BenchmarkResult], variant: Optional[str] = None
+    ) -> None:
+        """Add benchmark results to the full report.
+
+        Args:
+            results: List of BenchmarkResult objects to add to report
+        """
+        if results:
+            report_part = generate_markdown_report(self.profile_dir, results, variant)
+            self.full_report += report_part + "\n"
+
+    def write_final_report(self) -> None:
+        """Write the final report to GitHub summary if in CI."""
+        if is_in_ci():
+            write_github_step_summary(self.full_report)
+        print(self.full_report)
+
+    def get_full_report(self) -> str:
+        """Get the accumulated full report.
+
+        Returns:
+            The full markdown report as a string
+        """
+        return self.full_report