sglang/python/sglang/test/ci/ci_stress_utils.py

"""Utilities for running stress tests with bench_serving."""

import subprocess
from typing import List, Optional

from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    is_in_ci,
    popen_launch_server,
    write_github_step_summary,
)


class StressTestRunner:
    """Helper class for running stress tests with bench_serving.

    This class encapsulates common patterns for stress testing, including:
    - Server launch and cleanup
    - Stress test command construction and execution
    - Result collection and reporting
    """

    def __init__(
        self,
        test_name: str,
        base_url: str,
        num_prompts: int = 50000,
        duration_minutes: int = 45,
    ):
        """Initialize the stress test runner.

        Args:
            test_name: Name of the test (used for reporting)
            base_url: Base URL for the server
            num_prompts: Number of prompts to send (default: 50000)
            duration_minutes: Timeout in minutes (default: 45)
        """
        self.test_name = test_name
        self.base_url = base_url
        self.num_prompts = num_prompts
        self.duration_minutes = duration_minutes
        self.full_report = f"## {test_name}\n"

    def build_stress_test_command(
        self,
        model_path: str,
        random_input_len: int,
        random_output_len: int,
        output_file: str,
        random_range_ratio: float = 0.2,
        extra_args: Optional[List[str]] = None,
    ) -> List[str]:
        """Build the bench_serving stress test command.

        Args:
            model_path: Path to the model
            random_input_len: Random input length
            random_output_len: Random output length
            output_file: Output JSONL file path
            random_range_ratio: Random range ratio (default: 0.2)
            extra_args: Optional extra arguments

        Returns:
            List of command arguments ready for subprocess.run()
        """
        command = [
            "python3",
            "-m",
            "sglang.bench_serving",
            "--backend",
            "sglang-oai",
            "--base-url",
            self.base_url,
            "--dataset-name",
            "random",
            "--random-input-len",
            str(random_input_len),
            "--random-output-len",
            str(random_output_len),
            "--random-range-ratio",
            str(random_range_ratio),
            "--num-prompts",
            str(self.num_prompts),
            "--output-file",
            output_file,
        ]

        if extra_args:
            command.extend(extra_args)

        return command

    def run_stress_test_command(
        self, command: List[str], timeout_minutes: Optional[int] = None
    ) -> subprocess.CompletedProcess:
        """Execute the stress test command with timeout.

        Args:
            command: Command to execute
            timeout_minutes: Timeout in minutes (uses class default if None)

        Returns:
            CompletedProcess result
        """
        timeout = (timeout_minutes or self.duration_minutes) * 60
        print(f"Running stress test command (timeout: {timeout}s):")
        print(f"  {' '.join(command)}")

        result = subprocess.run(
            command, capture_output=True, text=True, timeout=timeout
        )

        if result.returncode != 0:
            print(f"Error running stress test:")
            print(result.stderr)
            raise RuntimeError(
                f"Stress test failed with return code {result.returncode}"
            )

        return result

    def run_stress_test_for_model(
        self,
        model_path: str,
        random_input_len: int,
        random_output_len: int,
        output_file: str,
        server_args: Optional[List[str]] = None,
        extra_bench_args: Optional[List[str]] = None,
        timeout_minutes: Optional[int] = None,
    ) -> bool:
        """Run a complete stress test for a single model with server management.

        This method handles:
        - Server launch and cleanup
        - Stress test command construction and execution
        - Error handling and reporting

        Args:
            model_path: Path to the model
            random_input_len: Random input length
            random_output_len: Random output length
            output_file: Output JSONL file path
            server_args: Arguments to pass to server launch
            extra_bench_args: Extra arguments for bench_serving
            timeout_minutes: Timeout in minutes (uses class default if None)

        Returns:
            True if successful, False otherwise
        """
        print(f"\n{'='*60}")
        print(f"Starting stress test for: {model_path}")
        print(f"Input length: {random_input_len}")
        print(f"Output length: {random_output_len}")
        print(f"Num prompts: {self.num_prompts}")
        print(f"{'='*60}\n")

        # Launch server
        process = popen_launch_server(
            model=model_path,
            base_url=self.base_url,
            other_args=server_args or [],
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
        )

        try:
            # Build and run stress test command
            command = self.build_stress_test_command(
                model_path,
                random_input_len,
                random_output_len,
                output_file,
                extra_args=extra_bench_args,
            )

            self.run_stress_test_command(command, timeout_minutes)

            print(f"\nStress test completed successfully for {model_path}")
            self._add_success_to_report(model_path, random_input_len, random_output_len)
            return True

        except Exception as e:
            print(f"\nStress test failed for {model_path}: {e}")
            self._add_failure_to_report(model_path, str(e))
            return False

        finally:
            # Always clean up server process
            print(f"Cleaning up server process...")
            kill_process_tree(process.pid)

    def _add_success_to_report(
        self, model_path: str, input_len: int, output_len: int
    ) -> None:
        """Add success entry to report."""
        model_name = model_path.split("/")[-1]
        self.full_report += f"### {model_name} - Success\n"
        self.full_report += f"- Model: `{model_path}`\n"
        self.full_report += f"- Input Length: {input_len}\n"
        self.full_report += f"- Output Length: {output_len}\n"
        self.full_report += f"- Num Prompts: {self.num_prompts}\n"
        self.full_report += f"- Status: **PASSED**\n\n"

    def _add_failure_to_report(self, model_path: str, error: str) -> None:
        """Add failure entry to report."""
        model_name = model_path.split("/")[-1]
        self.full_report += f"### {model_name} - Failure\n"
        self.full_report += f"- Model: `{model_path}`\n"
        self.full_report += f"- Status: **FAILED**\n"
        self.full_report += f"- Error: {error}\n\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)

    def get_full_report(self) -> str:
        """Get the accumulated full report.

        Returns:
            The full markdown report as a string
        """
        return self.full_report