test_nightly_text_models_perf.py

import os
import subprocess
import time
import unittest

from sglang.bench_one_batch_server import BenchmarkResult
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
    _parse_int_list_env,
    is_in_ci,
    parse_models,
    popen_launch_server,
    write_github_step_summary,
)

PROFILE_DIR = "performance_profiles_text_models"


class TestNightlyTextModelsPerformance(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.model_groups = [
            (parse_models("meta-llama/Llama-3.1-8B-Instruct"), False, False),
            (parse_models("Qwen/Qwen2-57B-A14B-Instruct"), False, True),
            # (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1), False, False),
            # (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2), False, True),
            # (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP1), True, False),
            # (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP2), True, True),
        ]
        cls.base_url = DEFAULT_URL_FOR_TEST
        cls.batch_sizes = [1, 1, 8, 16, 64]
        cls.input_lens = tuple(_parse_int_list_env("NIGHTLY_INPUT_LENS", "4096"))
        cls.output_lens = tuple(_parse_int_list_env("NIGHTLY_OUTPUT_LENS", "512"))
        os.makedirs(PROFILE_DIR, exist_ok=True)
        cls.full_report = f"## {cls.__name__}\n" + BenchmarkResult.help_str()

    def test_bench_one_batch(self):
        all_benchmark_results = []

        for model_group, is_fp8, is_tp2 in self.model_groups:
            for model in model_group:
                benchmark_results = []
                with self.subTest(model=model):
                    process = popen_launch_server(
                        model=model,
                        base_url=self.base_url,
                        other_args=["--tp", "2"] if is_tp2 else [],
                        timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
                    )
                    try:

                        profile_filename = (
                            f"{model.replace('/', '_')}_{int(time.time())}"
                        )
                        profile_path_prefix = os.path.join(
                            PROFILE_DIR, profile_filename
                        )
                        json_output_file = (
                            f"results_{model.replace('/', '_')}_{int(time.time())}.json"
                        )

                        command = [
                            "python3",
                            "-m",
                            "sglang.bench_one_batch_server",
                            "--model",
                            model,
                            "--base-url",
                            self.base_url,
                            "--batch-size",
                            *[str(x) for x in self.batch_sizes],
                            "--input-len",
                            *[str(x) for x in self.input_lens],
                            "--output-len",
                            *[str(x) for x in self.output_lens],
                            "--show-report",
                            "--profile",
                            "--profile-by-stage",
                            "--profile-filename-prefix",
                            profile_path_prefix,
                            f"--output-path={json_output_file}",
                            "--no-append-to-github-summary",
                        ]

                        print(f"Running command: {' '.join(command)}")
                        result = subprocess.run(command, capture_output=True, text=True)

                        if result.returncode != 0:
                            print(
                                f"Error running benchmark for {model} with batch size:"
                            )
                            print(result.stderr)
                            # Continue to next batch size even if one fails
                            continue

                        # Load and deserialize JSON results
                        if os.path.exists(json_output_file):
                            import json

                            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)
                                all_benchmark_results.append(benchmark_result)
                                benchmark_results.append(benchmark_result)

                            print(
                                f"Loaded {len(benchmark_results)} benchmark results from {json_output_file}"
                            )

                            # Clean up JSON file
                            os.remove(json_output_file)
                        else:
                            print(
                                f"Warning: JSON output file {json_output_file} not found"
                            )

                    finally:
                        kill_process_tree(process.pid)

                    report_part = BenchmarkResult.generate_markdown_report(
                        PROFILE_DIR, benchmark_results
                    )
                    self.full_report += report_part + "\n"

        if is_in_ci():
            write_github_step_summary(self.full_report)


if __name__ == "__main__":
    unittest.main()