util.py

"""
Common utilities for SGLang benchmark scripts.

This module contains shared code for benchmarking different SGLang APIs
including scoring, embeddings, and other endpoints.
"""

import asyncio
import concurrent.futures
import json
import os
import random
from statistics import mean
from typing import Any, Callable, Dict, List, Optional, Tuple

import aiohttp
import numpy as np
from tqdm import tqdm
from transformers import AutoTokenizer


class BenchmarkConfig:
    """Configuration for benchmark parameters."""

    def __init__(self):
        # Common benchmark settings
        self.server_type = "HTTP"
        self.rps_values = [70]
        self.duration_secs_values = [60]
        self.num_unique_requests = 100
        self.distribution = "POISSON"  # Options: "CONSTANT", "POISSON"
        self.profile = False

        # Garbage Collection Control
        self.freeze_gc = True  # Enable/disable garbage collection freezing

        # Profiler configuration
        self.profiler_dir = (
            os.getcwd()
        )  # Default profiler output directory (current working directory)

        # Special token for text generation
        self.special_replicated_token = "<|im_start|>"


def generate_text_with_token_count(
    model_path: str,
    num_tokens: int,
    special_token: str = "<|im_start|>",
    tokenizer: Optional[Any] = None,
) -> str:
    """
    Generate text with precise token count using a replicated token.

    Args:
        model_path: Path to the model for tokenizer
        num_tokens: Target number of tokens
        special_token: Token to replicate
        tokenizer: Optional pre-loaded tokenizer to avoid repeated loading

    Returns:
        Generated text with approximately the target token count
    """
    if tokenizer is None:
        tokenizer = AutoTokenizer.from_pretrained(model_path)

    # Verify token count
    special_token_count = len(tokenizer.encode(special_token, add_special_tokens=False))

    if special_token_count == 1:
        # Simple case: token maps to exactly 1 token
        return special_token * num_tokens
    else:
        print(f"Special token '{special_token}' produces {special_token_count} tokens")
        # Handle case where special token produces multiple tokens
        repetitions = (num_tokens + special_token_count - 1) // special_token_count
        text = special_token * repetitions

        # Verify we got the expected token count
        actual_tokens = len(tokenizer.encode(text, add_special_tokens=False))
        if actual_tokens < num_tokens:
            print(f"Warning: Generated {actual_tokens} tokens, expected {num_tokens}")

        return text


def setup_profiler(config: BenchmarkConfig, benchmark_name: str) -> None:
    """
    Set up profiler environment if profiling is enabled.

    Args:
        config: Benchmark configuration
        benchmark_name: Name of the benchmark (used in directory path)
    """
    if config.profile:
        # Create benchmark-specific subdirectory
        profiler_path = os.path.join(
            config.profiler_dir, benchmark_name.lower().replace("_", "-")
        )
        os.environ["SGLANG_TORCH_PROFILER_DIR"] = profiler_path
        print(f"Profiler enabled. Output directory: {profiler_path}")
    else:
        print("Profiler disabled")


def prepare_all_requests_parallel(
    num_requests: int,
    item_count: int,
    build_request_func: Callable[[int, int], Tuple[int, Any]],
    config: BenchmarkConfig,
    description: str = "requests",
) -> List[Any]:
    """
    Generic function to generate unique requests in parallel, then reuse them.

    Args:
        num_requests: Total number of requests needed
        item_count: Number of items per request (batch size)
        build_request_func: Function that takes (index, item_count) and returns (index, request_data)
        config: Benchmark configuration
        description: Description for progress bars

    Returns:
        List of request data objects
    """

    def build_request_wrapper(index):
        """Wrapper to call the provided build_request_func."""
        try:
            return build_request_func(index, item_count)
        except Exception as e:
            print(f"Error building request {index}: {e}")
            return (index, None)

    # Generate only the unique requests
    unique_requests = [None] * config.num_unique_requests
    max_workers = min(8, os.cpu_count() or 1)  # Limit to 8 threads max

    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = []
        for i in tqdm(
            range(config.num_unique_requests),
            desc=f"Submitting {description} generation tasks",
        ):
            future = executor.submit(build_request_wrapper, i)
            futures.append(future)

        # Collect results as they complete
        for f in tqdm(
            concurrent.futures.as_completed(futures),
            desc=f"Building unique {description}",
            total=config.num_unique_requests,
        ):
            try:
                index, req_data = f.result()
                if req_data is not None:
                    unique_requests[index] = req_data
                else:
                    print(f"Failed to build request {index}")
            except Exception as e:
                print(f"Error processing request result: {e}")

    # Check if we have any valid requests
    valid_requests = [req for req in unique_requests if req is not None]
    if not valid_requests:
        raise RuntimeError("Failed to generate any valid requests")

    print(
        f"Successfully generated {len(valid_requests)} out of "
        f"{config.num_unique_requests} unique {description}"
    )

    # Create the full request list by cycling through unique requests
    print(
        f"Reusing {len(valid_requests)} unique {description} to create "
        f"{num_requests} total requests..."
    )
    all_requests = []
    for i in tqdm(range(num_requests), desc=f"Reusing {description}"):
        unique_index = i % len(valid_requests)
        all_requests.append(valid_requests[unique_index])

    print(f"All {description} prepared.\n")
    return all_requests


async def sleep_with_distribution(distribution: str, rps: float) -> None:
    """
    Sleep according to the specified distribution pattern.

    Args:
        distribution: "CONSTANT" or "POISSON"
        rps: Requests per second rate
    """
    if distribution == "CONSTANT":
        interval = 1 / rps
        await asyncio.sleep(interval)
    elif distribution == "POISSON":
        # For Poisson process, inter-arrival times follow exponential distribution
        interval = random.expovariate(rps)
        await asyncio.sleep(interval)
    else:
        raise ValueError(
            f"Unknown distribution: {distribution}. Use 'CONSTANT' or 'POISSON'."
        )


def build_http_request_json(request_data: Any) -> str:
    """
    Generic function to build HTTP request JSON.

    Args:
        request_data: The data to serialize to JSON

    Returns:
        JSON string representation of the request data
    """
    return json.dumps(request_data)


async def make_http_call(
    session: aiohttp.ClientSession,
    request_data: Any,
    request_id: int,
    results_queue: asyncio.Queue,
    http_url: str,
    response_validator: Callable[[Dict[str, Any]], bool],
    api_name: str = "API",
) -> None:
    """
    Generic HTTP call function for API requests.

    Args:
        session: aiohttp client session
        request_data: Data to send in the request
        request_id: Unique identifier for this request
        results_queue: Queue to put results
        http_url: URL to send the request to
        response_validator: Function to validate the response JSON
        api_name: Name of the API for error messages
    """
    try:
        start_time = asyncio.get_event_loop().time()

        request_json = build_http_request_json(request_data)
        headers = {"Content-Type": "application/json"}

        async with session.post(http_url, data=request_json, headers=headers) as resp:
            resp_text = await resp.text()

            if resp.status != 200:
                print(
                    f"[HTTP] {api_name} Request {request_id} failed with status "
                    f"{resp.status}: {resp_text}"
                )
                completion_time = asyncio.get_event_loop().time()
                await results_queue.put((request_id, 0, False, completion_time))
                return

            # Parse and validate response
            try:
                response_data = json.loads(resp_text)
                success = response_validator(response_data)
                if not success:
                    print(
                        f"[HTTP] {api_name} Request {request_id} failed response validation"
                    )
            except json.JSONDecodeError:
                print(
                    f"[HTTP] {api_name} Request {request_id} failed to parse JSON response"
                )
                success = False

        completion_time = asyncio.get_event_loop().time()
        elapsed_time = (completion_time - start_time) * 1000
        await results_queue.put((request_id, elapsed_time, success, completion_time))

    except Exception as e:
        print(f"[HTTP] {api_name} Error for request {request_id}: {e}")
        completion_time = asyncio.get_event_loop().time()
        await results_queue.put((request_id, 0, False, completion_time))


async def send_profile_request(
    profile_text: str, http_url: str, session: Optional[aiohttp.ClientSession] = None
) -> None:
    """
    Send a profile request (START_PROFILE or STOP_PROFILE) and wait for completion.

    Args:
        profile_text: "START_PROFILE" or "STOP_PROFILE"
        http_url: Base HTTP URL (will derive profile endpoints from this)
        session: Optional aiohttp session to use
    """
    try:
        if session:
            print(f"Sending {profile_text} request via HTTP...")

            # Determine the correct endpoint
            if "/v1/" in http_url:
                base_url = http_url.rsplit("/v1/", 1)[0]  # Remove /v1/xxx
            else:
                base_url = http_url.rsplit("/", 1)[0]  # Remove last path component

            if profile_text == "START_PROFILE":
                endpoint_url = f"{base_url}/start_profile"
            elif profile_text == "STOP_PROFILE":
                endpoint_url = f"{base_url}/stop_profile"
            else:
                print(f"Unknown profile request: {profile_text}")
                return

            headers = {"Content-Type": "application/json"}

            async with session.post(endpoint_url, headers=headers) as resp:
                resp_text = await resp.text()
                if resp.status == 200:
                    print(f"{profile_text} request completed")
                else:
                    print(
                        f"{profile_text} request failed with status "
                        f"{resp.status}: {resp_text}"
                    )
        else:
            print(f"Cannot send {profile_text} request - missing session")

    except Exception as e:
        print(f"Error sending {profile_text} request: {e}")


async def call_freeze_gc_http(session: aiohttp.ClientSession, http_url: str) -> None:
    """
    Call the /freeze_gc HTTP endpoint.

    Args:
        session: aiohttp client session
        http_url: Base HTTP URL to derive the freeze_gc endpoint from
    """
    try:
        # Derive freeze_gc endpoint from the API URL
        if "/v1/" in http_url:
            freeze_gc_url = http_url.rsplit("/v1/", 1)[0] + "/freeze_gc"
        else:
            freeze_gc_url = http_url.rsplit("/", 1)[0] + "/freeze_gc"

        print(f"Calling freeze_gc endpoint: {freeze_gc_url}")

        async with session.post(freeze_gc_url) as resp:
            if resp.status == 200:
                print("freeze_gc called successfully")
            else:
                resp_text = await resp.text()
                print(f"freeze_gc failed with status {resp.status}: {resp_text}")

    except Exception as e:
        print(f"Failed to call freeze_gc: {e}")


async def send_warmup_requests(
    session: aiohttp.ClientSession,
    http_url: str,
    build_warmup_request_func: Callable[[], Any],
    num_warmup: int = 3,
) -> None:
    """
    Send warmup requests to HTTP server.

    Args:
        session: aiohttp client session
        http_url: URL to send warmup requests to
        build_warmup_request_func: Function that returns a warmup request object
        num_warmup: Number of warmup requests to send
    """
    print(f"Sending {num_warmup} HTTP warmup requests...")

    for i in range(num_warmup):
        try:
            warmup_data = build_warmup_request_func()
            request_json = build_http_request_json(warmup_data)
            headers = {"Content-Type": "application/json"}

            async with session.post(
                http_url, data=request_json, headers=headers
            ) as resp:
                if resp.status == 200:
                    print(f"Warmup request {i+1}/{num_warmup} completed successfully")
                else:
                    print(
                        f"Warmup request {i+1}/{num_warmup} failed with status {resp.status}"
                    )

        except Exception as e:
            print(f"Warmup request {i+1}/{num_warmup} failed with error: {e}")

    print("HTTP warmup requests completed")


async def perform_global_warmup_and_freeze(
    config: BenchmarkConfig,
    http_url: str,
    build_warmup_request_func: Callable[[], Any],
) -> None:
    """
    Perform warmup and optionally GC freeze operations once before all benchmark runs.

    Args:
        config: Benchmark configuration
        http_url: URL for API requests
        build_warmup_request_func: Function that returns a warmup request object
    """
    print("=" * 80)
    print(f"PERFORMING GLOBAL WARMUP{' AND GC FREEZE' if config.freeze_gc else ''}")
    print("=" * 80)

    print(f"Performing HTTP warmup{' and GC freeze' if config.freeze_gc else ''}...")
    async with aiohttp.ClientSession() as session:
        await send_warmup_requests(session, http_url, build_warmup_request_func)
        if config.freeze_gc:
            await call_freeze_gc_http(session, http_url)
        print(
            f"HTTP warmup{' and GC freeze' if config.freeze_gc else ''} completed successfully."
        )

    print(
        f"Global warmup{' and GC freeze' if config.freeze_gc else ''} operations completed."
    )
    print("=" * 80)


async def process_results(
    results_queue: asyncio.Queue,
    num_requests: int,
    send_duration: float,
    total_duration: float,
    rps: int,
    duration_secs: int,
    item_count: int,
    test_start_time: float,
    config: BenchmarkConfig,
    http_mode: str = "UNKNOWN",
) -> List[Dict[str, Any]]:
    """
    Process benchmark results and group them by minute intervals.

    Args:
        results_queue: Queue containing result tuples
        num_requests: Total number of requests sent
        send_duration: Time taken to send all requests
        total_duration: Total time for all requests to complete
        rps: Target requests per second
        duration_secs: Test duration in seconds
        item_count: Number of items per request
        test_start_time: Start time of the test
        config: Benchmark configuration
        http_mode: Description of the HTTP mode/API being tested

    Returns:
        List of dictionaries containing minute-by-minute results
    """
    all_results = []

    # Collect all results
    for _ in range(num_requests):
        result = await results_queue.get()
        request_id, elapsed_time, success, completion_time = result
        all_results.append(
            {
                "request_id": request_id,
                "elapsed_time": elapsed_time,
                "success": success,
                "completion_time": completion_time,
            }
        )

    # Group results by minute intervals
    minute_results = []
    num_minutes = int(duration_secs // 60) + (1 if duration_secs % 60 > 0 else 0)

    for minute in range(num_minutes):
        minute_start = test_start_time + (minute * 60)
        minute_end = test_start_time + ((minute + 1) * 60)

        # Filter results that completed in this minute
        minute_data = [
            r for r in all_results if minute_start <= r["completion_time"] < minute_end
        ]

        response_times = [r["elapsed_time"] for r in minute_data if r["success"]]
        successful_requests = len([r for r in minute_data if r["success"]])
        failed_requests = len([r for r in minute_data if not r["success"]])

        avg_response_time = mean(response_times) if response_times else 0

        # Calculate percentiles using numpy
        if response_times:
            p50 = np.percentile(response_times, 50)
            p90 = np.percentile(response_times, 90)
            p99 = np.percentile(response_times, 99)
        else:
            p50 = p90 = p99 = 0

        minute_result = {
            "test_duration_secs": duration_secs,
            "minute_interval": minute + 1,
            "target_rps": rps,
            "item_count": item_count,
            "server_type": config.server_type,
            "distribution": config.distribution,
            "unique_requests": config.num_unique_requests,
            "total_requests": len(minute_data),
            "successful_requests": successful_requests,
            "failed_requests": failed_requests,
            "send_duration_secs": send_duration,
            "total_duration_secs": total_duration,
            "avg_response_time_ms": avg_response_time,
            "p50_response_time_ms": p50,
            "p90_response_time_ms": p90,
            "p99_response_time_ms": p99,
        }

        minute_results.append(minute_result)

        print(
            f"\nMinute {minute + 1} Summary for RPS {rps}, "
            f"Duration {duration_secs}s, Item Count {item_count}:"
        )
        print(f"  Requests completed in minute: {len(minute_data)}")
        print(f"  Successful requests:   {successful_requests}")
        print(f"  Failed requests:       {failed_requests}")
        print(f"  Average response time: {avg_response_time:.2f} ms")
        print(f"  P50 response time:     {p50:.2f} ms")
        print(f"  P90 response time:     {p90:.2f} ms")
        print(f"  P99 response time:     {p99:.2f} ms")

    # Print overall summary
    all_response_times = [r["elapsed_time"] for r in all_results if r["success"]]
    total_successful = len([r for r in all_results if r["success"]])
    total_failed = len([r for r in all_results if not r["success"]])

    overall_avg = mean(all_response_times) if all_response_times else 0
    if all_response_times:
        overall_p50 = np.percentile(all_response_times, 50)
        overall_p90 = np.percentile(all_response_times, 90)
        overall_p99 = np.percentile(all_response_times, 99)
    else:
        overall_p50 = overall_p90 = overall_p99 = 0

    print(
        f"\nOverall Summary for RPS {rps}, Duration {duration_secs}s, "
        f"Item Count {item_count}:"
    )
    print(f"  Test duration:         {duration_secs} seconds")
    print(f"  Server type:           {config.server_type}")
    print(f"  HTTP mode:             {http_mode}")
    print(f"  Target RPS:            {rps}")
    print(f"  Item count:            {item_count}")
    print(f"  Distribution:          {config.distribution}")
    print(f"  Unique requests generated: {config.num_unique_requests}")
    print(f"  Total requests sent:   {num_requests}")
    print(f"  Successful requests:   {total_successful}")
    print(f"  Failed requests:       {total_failed}")
    print(f"  Time to send all requests: {send_duration:.2f} seconds")
    print(f"  Time for all requests to complete: {total_duration:.2f} seconds")
    print(f"  Average response time: {overall_avg:.2f} ms")
    print(f"  P50 response time:     {overall_p50:.2f} ms")
    print(f"  P90 response time:     {overall_p90:.2f} ms")
    print(f"  P99 response time:     {overall_p99:.2f} ms\n")

    return minute_results


def print_csv_results(all_results: List[Dict[str, Any]]) -> None:
    """
    Print benchmark results in CSV format.

    Args:
        all_results: List of result dictionaries from process_results
    """
    print("\n" + "=" * 80)
    print("FINAL CSV RESULTS:")
    print("=" * 80)

    # CSV Header
    headers = [
        "test_duration_secs",
        "minute_interval",
        "target_rps",
        "item_count",
        "server_type",
        "distribution",
        "unique_requests",
        "total_requests",
        "successful_requests",
        "failed_requests",
        "send_duration_secs",
        "total_duration_secs",
        "avg_response_time_ms",
        "p50_response_time_ms",
        "p90_response_time_ms",
        "p99_response_time_ms",
    ]
    print(",".join(headers))

    # CSV Data
    for result in all_results:
        row = [
            result["test_duration_secs"],
            result["minute_interval"],
            result["target_rps"],
            result["item_count"],
            result["server_type"],
            result["distribution"],
            result["unique_requests"],
            result["total_requests"],
            result["successful_requests"],
            result["failed_requests"],
            f"{result['send_duration_secs']:.2f}",
            f"{result['total_duration_secs']:.2f}",
            f"{result['avg_response_time_ms']:.2f}",
            f"{result['p50_response_time_ms']:.2f}",
            f"{result['p90_response_time_ms']:.2f}",
            f"{result['p99_response_time_ms']:.2f}",
        ]
        print(",".join(map(str, row)))


async def run_benchmark_main(
    config: BenchmarkConfig,
    run_single_benchmark_func,
    benchmark_name: str,
    http_url: str,
    item_count_values: List[int],
    additional_info: Optional[Dict[str, Any]] = None,
    build_warmup_request_func: Optional[Callable[[], Any]] = None,
) -> None:
    """
    Main benchmark orchestration function.

    Args:
        config: Benchmark configuration
        run_single_benchmark_func: Async function to run a single benchmark
        benchmark_name: Name of the benchmark (e.g., "SCORING", "EMBEDDINGS")
        http_url: URL of the API endpoint
        item_count_values: List of item counts to test
        additional_info: Additional information to print in the header
        build_warmup_request_func: Optional function to build warmup requests
    """
    total_combinations = (
        len(config.duration_secs_values)
        * len(config.rps_values)
        * len(item_count_values)
    )

    print(
        f"Running benchmarks for {len(config.duration_secs_values)} duration "
        f"values, {len(config.rps_values)} RPS values, and "
        f"{len(item_count_values)} item count values = "
        f"{total_combinations} total combinations"
    )
    print(f"Server Type: {config.server_type}")
    print(f"HTTP Mode: {benchmark_name}")
    print(f"API URL: {http_url}")

    if additional_info:
        for key, value in additional_info.items():
            print(f"{key}: {value}")

    print(f"Items per request (batch size): {item_count_values}")
    print(f"Profiling Enabled: {config.profile}")
    print(f"Duration values: {config.duration_secs_values}")
    print(f"RPS values: {config.rps_values}")
    print(f"Item count values: {item_count_values}")
    print("=" * 80)

    # Set up profiler environment
    setup_profiler(config, benchmark_name)

    # Perform global warmup and GC freeze operations if warmup function is provided
    if build_warmup_request_func is not None:
        await perform_global_warmup_and_freeze(
            config, http_url, build_warmup_request_func
        )

    all_results = []

    for duration_secs in config.duration_secs_values:
        for rps in config.rps_values:
            for item_count in item_count_values:
                result = await run_single_benchmark_func(rps, duration_secs, item_count)
                all_results.extend(result)  # Extend with minute results

    print_csv_results(all_results)


async def run_generic_benchmark(
    rps: int,
    duration_secs: int,
    item_count: int,
    config: BenchmarkConfig,
    http_url: str,
    build_request_func: Callable[[int, int], Tuple[int, Any]],
    response_validator: Callable[[Dict[str, Any]], bool],
    api_name: str,
    request_description: str = "requests",
) -> List[Dict[str, Any]]:
    """
    Generic benchmark runner that can be used for different APIs.

    Args:
        rps: Requests per second
        duration_secs: Duration of the test in seconds
        item_count: Number of items per request (batch size)
        config: Benchmark configuration
        http_url: URL of the API endpoint
        build_request_func: Function to build individual requests
        response_validator: Function to validate API responses
        api_name: Name of the API for logging
        request_description: Description for progress bars

    Returns:
        List of dictionaries containing minute-by-minute results
    """
    num_requests = int(rps * duration_secs)
    print(
        f"Starting benchmark with RPS={rps}, Duration={duration_secs}s, "
        f"Item Count={item_count}, num_requests={num_requests}"
    )
    print(f"Server Type: {config.server_type}")
    print(f"HTTP Mode: {api_name}")
    print(f"Profiling Enabled: {config.profile}")

    # Build requests in parallel (unmeasured)
    all_requests = prepare_all_requests_parallel(
        num_requests, item_count, build_request_func, config, request_description
    )

    results_queue = asyncio.Queue()
    tasks = []

    # Track timing for sending requests
    send_start_time = asyncio.get_event_loop().time()

    # HTTP implementation
    async with aiohttp.ClientSession(
        timeout=aiohttp.ClientTimeout(total=300)
    ) as session:

        # Send START_PROFILE if profiling is enabled
        if config.profile:
            await send_profile_request("START_PROFILE", http_url, session=session)

        # Add progress bar for sending requests
        with tqdm(
            total=len(all_requests),
            desc=f"Sending HTTP {request_description} at {rps} RPS",
            unit="req",
        ) as pbar:
            for i, request_data in enumerate(all_requests):
                request_id = i + 1
                tasks.append(
                    asyncio.create_task(
                        make_http_call(
                            session,
                            request_data,
                            request_id,
                            results_queue,
                            http_url,
                            response_validator,
                            api_name,
                        )
                    )
                )

                # Update progress bar
                pbar.update(1)

                # Throttle based on distribution
                if i < len(all_requests) - 1:
                    await sleep_with_distribution(config.distribution, rps)

        send_end_time = asyncio.get_event_loop().time()
        send_duration = send_end_time - send_start_time

        # Wait for all requests to complete with progress tracking
        print(f"Waiting for {len(tasks)} HTTP {request_description} to complete...")
        with tqdm(
            total=len(tasks), desc=f"Completing HTTP {request_description}", unit="req"
        ) as completion_pbar:
            completed_tasks = []
            for task in asyncio.as_completed(tasks):
                await task
                completed_tasks.append(task)
                completion_pbar.update(1)

        # Send STOP_PROFILE if profiling is enabled
        if config.profile:
            await send_profile_request("STOP_PROFILE", http_url, session=session)

    completion_end_time = asyncio.get_event_loop().time()
    total_duration = completion_end_time - send_start_time

    return await process_results(
        results_queue,
        num_requests,
        send_duration,
        total_duration,
        rps,
        duration_secs,
        item_count,
        send_start_time,
        config,
        api_name,
    )