lmcache-cli 0.4.5.dev0__py3-none-any.whl

lmcache/cli/commands/bench/engine_bench/stats.py

@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stats collection, aggregation, and export for ``lmcache bench engine``."""
+
+# Standard
+from dataclasses import asdict, dataclass
+from pathlib import Path
+import csv
+import json
+import threading
+import time
+
+# First Party
+from lmcache.cli.commands.bench.engine_bench.config import EngineBenchConfig
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class RequestResult:
+    """Raw per-request result collected by the request sender."""
+
+    request_id: str
+    successful: bool
+    ttft: float  # time to first token (seconds)
+    request_latency: float  # total request time (seconds)
+    num_input_tokens: int  # from server usage report
+    num_output_tokens: int  # tokens generated
+    decode_speed: float  # output tokens / decode time (tok/s)
+    submit_time: float  # absolute timestamp
+    first_token_time: float  # absolute timestamp
+    finish_time: float  # absolute timestamp
+    error: str  # empty string if successful
+
+
+@dataclass
+class AggregatedStats:
+    """Snapshot of aggregated statistics (running totals)."""
+
+    total_requests: int
+    successful_requests: int
+    failed_requests: int
+    elapsed_time: float  # seconds since benchmark start
+
+    mean_ttft_ms: float
+    mean_decode_speed: float  # tok/s
+    mean_request_latency_ms: float
+
+    input_throughput: float  # total input tokens / elapsed time
+    output_throughput: float  # total output tokens / elapsed time
+
+    total_input_tokens: int
+    total_output_tokens: int
+
+
+@dataclass
+class FinalStats(AggregatedStats):
+    """Final statistics with percentiles. Extends AggregatedStats."""
+
+    p50_ttft_ms: float = 0.0
+    p90_ttft_ms: float = 0.0
+    p99_ttft_ms: float = 0.0
+    p50_decode_speed: float = 0.0
+    p90_decode_speed: float = 0.0
+    p99_decode_speed: float = 0.0
+    p50_request_latency_ms: float = 0.0
+    p90_request_latency_ms: float = 0.0
+    p99_request_latency_ms: float = 0.0
+
+
+class StatsCollector:
+    """Thread-safe stats aggregation for benchmark results.
+
+    Receives ``RequestResult`` objects from the request sender,
+    maintains running totals, and produces final summaries.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._results: list[RequestResult] = []
+        self._start_time: float = time.monotonic()
+
+        # Running accumulators (updated under lock)
+        self._successful: int = 0
+        self._failed: int = 0
+        self._sum_ttft: float = 0.0
+        self._sum_decode_speed: float = 0.0
+        self._sum_request_latency: float = 0.0
+        self._total_input_tokens: int = 0
+        self._total_output_tokens: int = 0
+
+    def on_request_finished(self, result: RequestResult) -> None:
+        """Record a completed request. Thread-safe."""
+        with self._lock:
+            self._results.append(result)
+            if result.successful:
+                self._successful += 1
+                self._sum_ttft += result.ttft
+                self._sum_decode_speed += result.decode_speed
+                self._sum_request_latency += result.request_latency
+            else:
+                self._failed += 1
+            self._total_input_tokens += result.num_input_tokens
+            self._total_output_tokens += result.num_output_tokens
+        logger.debug(
+            "Recorded result for %s (successful=%s)",
+            result.request_id,
+            result.successful,
+        )
+
+    def reset(self) -> None:
+        """Clear all accumulated results and restart the timer.
+
+        Used between warmup and benchmark phases so warmup stats
+        don't pollute benchmark results. Thread-safe.
+        """
+        with self._lock:
+            self._results.clear()
+            self._start_time = time.monotonic()
+            self._successful = 0
+            self._failed = 0
+            self._sum_ttft = 0.0
+            self._sum_decode_speed = 0.0
+            self._sum_request_latency = 0.0
+            self._total_input_tokens = 0
+            self._total_output_tokens = 0
+        logger.debug("Stats collector reset")
+
+    def get_current_stats(self) -> AggregatedStats:
+        """Return current aggregated stats snapshot. Thread-safe."""
+        with self._lock:
+            successful = self._successful
+            failed = self._failed
+            elapsed = time.monotonic() - self._start_time
+            sum_ttft = self._sum_ttft
+            sum_decode = self._sum_decode_speed
+            sum_latency = self._sum_request_latency
+            total_in = self._total_input_tokens
+            total_out = self._total_output_tokens
+
+        safe_successful = max(successful, 1)
+
+        return AggregatedStats(
+            total_requests=successful + failed,
+            successful_requests=successful,
+            failed_requests=failed,
+            elapsed_time=elapsed,
+            mean_ttft_ms=(sum_ttft / safe_successful) * 1000.0,
+            mean_decode_speed=sum_decode / safe_successful,
+            mean_request_latency_ms=(sum_latency / safe_successful) * 1000.0,
+            input_throughput=total_in / max(elapsed, 1e-9),
+            output_throughput=total_out / max(elapsed, 1e-9),
+            total_input_tokens=total_in,
+            total_output_tokens=total_out,
+        )
+
+    def get_final_stats(self) -> FinalStats:
+        """Compute and return final stats with percentiles.
+
+        Should be called once after the benchmark completes.
+        """
+        with self._lock:
+            results = list(self._results)
+
+        successful_results = [r for r in results if r.successful]
+        current = self.get_current_stats()
+
+        if not successful_results:
+            return FinalStats(
+                total_requests=current.total_requests,
+                successful_requests=current.successful_requests,
+                failed_requests=current.failed_requests,
+                elapsed_time=current.elapsed_time,
+                mean_ttft_ms=current.mean_ttft_ms,
+                mean_decode_speed=current.mean_decode_speed,
+                mean_request_latency_ms=current.mean_request_latency_ms,
+                input_throughput=current.input_throughput,
+                output_throughput=current.output_throughput,
+                total_input_tokens=current.total_input_tokens,
+                total_output_tokens=current.total_output_tokens,
+            )
+
+        ttfts = sorted(r.ttft * 1000.0 for r in successful_results)
+        decode_speeds = sorted(r.decode_speed for r in successful_results)
+        latencies = sorted(r.request_latency * 1000.0 for r in successful_results)
+
+        return FinalStats(
+            total_requests=current.total_requests,
+            successful_requests=current.successful_requests,
+            failed_requests=current.failed_requests,
+            elapsed_time=current.elapsed_time,
+            mean_ttft_ms=current.mean_ttft_ms,
+            mean_decode_speed=current.mean_decode_speed,
+            mean_request_latency_ms=current.mean_request_latency_ms,
+            input_throughput=current.input_throughput,
+            output_throughput=current.output_throughput,
+            total_input_tokens=current.total_input_tokens,
+            total_output_tokens=current.total_output_tokens,
+            p50_ttft_ms=_percentile(ttfts, 50),
+            p90_ttft_ms=_percentile(ttfts, 90),
+            p99_ttft_ms=_percentile(ttfts, 99),
+            p50_decode_speed=_percentile(decode_speeds, 50),
+            p90_decode_speed=_percentile(decode_speeds, 90),
+            p99_decode_speed=_percentile(decode_speeds, 99),
+            p50_request_latency_ms=_percentile(latencies, 50),
+            p90_request_latency_ms=_percentile(latencies, 90),
+            p99_request_latency_ms=_percentile(latencies, 99),
+        )
+
+    def get_all_results(self) -> list[RequestResult]:
+        """Return all raw results for CSV export."""
+        with self._lock:
+            return list(self._results)
+
+    def export_csv(self, path: str) -> None:
+        """Write per-request results to a CSV file."""
+        results = self.get_all_results()
+        fieldnames = [
+            "request_id",
+            "successful",
+            "ttft",
+            "request_latency",
+            "num_input_tokens",
+            "num_output_tokens",
+            "decode_speed",
+            "submit_time",
+            "first_token_time",
+            "finish_time",
+            "error",
+        ]
+        Path(path).parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=fieldnames)
+            writer.writeheader()
+            for result in results:
+                writer.writerow(asdict(result))
+        logger.debug("Exported %d results to CSV: %s", len(results), path)
+
+    def export_json(self, path: str, config: EngineBenchConfig) -> None:
+        """Write summary JSON with config and aggregated metrics."""
+        final = self.get_final_stats()
+        output = {
+            "config": asdict(config),
+            "results": asdict(final),
+        }
+        with open(path, "w") as f:
+            json.dump(output, f, indent=2)
+        logger.debug("Exported JSON summary to: %s", path)
+
+
+def _percentile(sorted_data: list[float], p: float) -> float:
+    """Compute the p-th percentile using linear interpolation.
+
+    Uses the Tensormesh-Benchmark V1 method:
+    ``k = (len(sorted_data) - 1) * p / 100``
+
+    Args:
+        sorted_data: Pre-sorted list of values.
+        p: Percentile value (0-100).
+
+    Returns:
+        Interpolated percentile value. Returns 0.0 for empty data.
+    """
+    if not sorted_data:
+        return 0.0
+    n = len(sorted_data)
+    if n == 1:
+        return sorted_data[0]
+    k = (n - 1) * p / 100.0
+    floor_k = int(k)
+    ceil_k = min(floor_k + 1, n - 1)
+    fraction = k - floor_k
+    return sorted_data[floor_k] + fraction * (
+        sorted_data[ceil_k] - sorted_data[floor_k]
+    )

lmcache/cli/commands/bench/engine_bench/workloads/__init__.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Workload definitions and factory for ``lmcache bench engine``.
+
+Each workload module defines its own config dataclass and workload
+class. The ``create_workload`` factory selects the right workload
+based on ``EngineBenchConfig.workload``, resolves the workload-specific
+config from CLI args, and returns the workload instance.
+"""
+
+# Standard
+import argparse
+
+# First Party
+from lmcache.cli.commands.bench.engine_bench.config import EngineBenchConfig
+from lmcache.cli.commands.bench.engine_bench.progress import ProgressMonitor
+from lmcache.cli.commands.bench.engine_bench.request_sender import (
+    RequestSender,
+)
+from lmcache.cli.commands.bench.engine_bench.stats import StatsCollector
+from lmcache.cli.commands.bench.engine_bench.workloads.base import BaseWorkload
+from lmcache.cli.commands.bench.engine_bench.workloads.long_doc_permutator import (
+    LongDocPermutatorConfig,
+    LongDocPermutatorWorkload,
+)
+from lmcache.cli.commands.bench.engine_bench.workloads.long_doc_qa import (
+    LongDocQAConfig,
+    LongDocQAWorkload,
+)
+from lmcache.cli.commands.bench.engine_bench.workloads.multi_round_chat import (
+    MultiRoundChatConfig,
+    MultiRoundChatWorkload,
+)
+from lmcache.cli.commands.bench.engine_bench.workloads.random_prefill import (
+    RandomPrefillConfig,
+    RandomPrefillWorkload,
+)
+
+__all__ = [
+    "BaseWorkload",
+    "LongDocPermutatorConfig",
+    "LongDocPermutatorWorkload",
+    "LongDocQAConfig",
+    "LongDocQAWorkload",
+    "MultiRoundChatConfig",
+    "MultiRoundChatWorkload",
+    "RandomPrefillConfig",
+    "RandomPrefillWorkload",
+    "create_workload",
+]
+
+_WORKLOAD_NAMES = (
+    "long-doc-permutator",
+    "long-doc-qa",
+    "multi-round-chat",
+    "random-prefill",
+)
+
+
+def create_workload(
+    config: EngineBenchConfig,
+    args: argparse.Namespace,
+    request_sender: RequestSender,
+    stats_collector: StatsCollector,
+    progress_monitor: ProgressMonitor,
+) -> BaseWorkload:
+    """Resolve workload-specific config and create the workload instance.
+
+    Dispatches on ``config.workload`` to the appropriate workload module,
+    resolves the workload-specific config from ``args`` and ``config``,
+    and returns the workload instance ready to ``run()``.
+
+    Args:
+        config: Fully-resolved general benchmark config.
+        args: Raw CLI args namespace (contains workload-specific flags).
+        request_sender: Shared request sender instance.
+        stats_collector: Shared stats collector instance.
+        progress_monitor: Shared progress monitor instance.
+
+    Returns:
+        A concrete BaseWorkload instance.
+
+    Raises:
+        ValueError: If the workload name is not recognized.
+    """
+    if config.workload == "long-doc-permutator":
+        ldp_workload_config = LongDocPermutatorConfig.resolve(
+            num_contexts=args.ldp_num_contexts,
+            context_length=args.ldp_context_length,
+            system_prompt_length=args.ldp_system_prompt_length,
+            num_permutations=args.ldp_num_permutations,
+            vocab_size=8000,
+            num_inflight_requests=args.ldp_num_inflight_requests,
+        )
+        return LongDocPermutatorWorkload(
+            config=ldp_workload_config,
+            request_sender=request_sender,
+            stats_collector=stats_collector,
+            progress_monitor=progress_monitor,
+            seed=config.seed,
+        )
+
+    if config.workload == "long-doc-qa":
+        ld_workload_config = LongDocQAConfig.resolve(
+            kv_cache_volume_gb=config.kv_cache_volume_gb,
+            tokens_per_gb_kvcache=config.tokens_per_gb_kvcache,
+            document_length=args.ldqa_document_length,
+            query_per_document=args.ldqa_query_per_document,
+            shuffle_policy=args.ldqa_shuffle_policy,
+            num_inflight_requests=args.ldqa_num_inflight_requests,
+        )
+        return LongDocQAWorkload(
+            config=ld_workload_config,
+            request_sender=request_sender,
+            stats_collector=stats_collector,
+            progress_monitor=progress_monitor,
+            seed=config.seed,
+        )
+
+    if config.workload == "multi-round-chat":
+        mr_workload_config = MultiRoundChatConfig.resolve(
+            kv_cache_volume_gb=config.kv_cache_volume_gb,
+            tokens_per_gb_kvcache=config.tokens_per_gb_kvcache,
+            shared_prompt_length=args.mrc_shared_prompt_length,
+            chat_history_length=args.mrc_chat_history_length,
+            user_input_length=args.mrc_user_input_length,
+            output_length=args.mrc_output_length,
+            qps=args.mrc_qps,
+            duration=args.mrc_duration,
+        )
+        return MultiRoundChatWorkload(
+            config=mr_workload_config,
+            request_sender=request_sender,
+            stats_collector=stats_collector,
+            progress_monitor=progress_monitor,
+            seed=config.seed,
+        )
+
+    if config.workload == "random-prefill":
+        rp_workload_config = RandomPrefillConfig.resolve(
+            request_length=args.rp_request_length,
+            num_requests=args.rp_num_requests,
+        )
+        return RandomPrefillWorkload(
+            config=rp_workload_config,
+            request_sender=request_sender,
+            stats_collector=stats_collector,
+            progress_monitor=progress_monitor,
+            seed=config.seed,
+        )
+
+    raise ValueError(
+        f"Unknown workload {config.workload!r}. Available: {', '.join(_WORKLOAD_NAMES)}"
+    )

lmcache/cli/commands/bench/engine_bench/workloads/base.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Abstract base class for engine benchmark workloads."""
+
+# Standard
+from abc import ABC, abstractmethod
+import asyncio
+import queue
+import time
+
+# First Party
+from lmcache.cli.commands.bench.engine_bench.progress import ProgressMonitor
+from lmcache.cli.commands.bench.engine_bench.request_sender import RequestSender
+from lmcache.cli.commands.bench.engine_bench.stats import StatsCollector
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+class BaseWorkload(ABC):
+    """Abstract base class for all engine benchmark workloads.
+
+    Owns the internal dispatch loop that calls ``step()`` on the
+    concrete workload.  Provides a thread-safe ``request_finished``
+    callback that enqueues completed requests for the loop to drain.
+    """
+
+    def __init__(
+        self,
+        request_sender: RequestSender,
+        stats_collector: StatsCollector,
+        progress_monitor: ProgressMonitor,
+    ) -> None:
+        self._request_sender = request_sender
+        self._stats_collector = stats_collector
+        self._progress_monitor = progress_monitor
+        self._finished_queue: queue.Queue[tuple[str, str]] = queue.Queue()
+
+    # ------------------------------------------------------------------
+    # Abstract methods
+    # ------------------------------------------------------------------
+
+    @abstractmethod
+    async def warmup(self) -> None:
+        """Run warmup requests.  Blocks until all warmup is done."""
+
+    @abstractmethod
+    async def step(self, time_offset: float) -> float:
+        """Execute one step of the workload.
+
+        Args:
+            time_offset: seconds since benchmark start.
+
+        Returns:
+            Next wakeup time offset (absolute, from benchmark start).
+            Return a negative value to signal that the workload is done.
+        """
+
+    @abstractmethod
+    def log_config(self) -> None:
+        """Log key workload config before the benchmark starts."""
+
+    @abstractmethod
+    def on_request_finished(self, request_id: str, output: str) -> None:
+        """Called when a request finishes (from the loop thread).
+
+        Stateless workloads can implement this as a no-op.
+        Stateful workloads use it to record responses in session history.
+        """
+
+    # ------------------------------------------------------------------
+    # Concrete methods
+    # ------------------------------------------------------------------
+
+    def run(self) -> None:
+        """Run the full workload: warmup, then the benchmark loop.
+
+        Blocks until the workload is complete.
+        """
+        asyncio.run(self._run_async())
+
+    async def _run_async(self) -> None:
+        """Internal async implementation of the run loop."""
+        self._progress_monitor.log_message("Starting warmup phase")
+        await self.warmup()
+
+        self._progress_monitor.log_message(
+            "Warmup complete, starting benchmark",
+        )
+        self._stats_collector.reset()
+        self._drain_finished_queue()  # discard warmup completions
+
+        start_time = time.monotonic()
+
+        while True:
+            self._drain_finished_queue()
+            time_offset = time.monotonic() - start_time
+            next_wakeup = await self.step(time_offset)
+            if next_wakeup < 0:
+                break
+            sleep_duration = max(0.0, next_wakeup - (time.monotonic() - start_time))
+            if sleep_duration > 0:
+                await asyncio.sleep(sleep_duration)
+
+        self._drain_finished_queue()  # final drain
+        self._progress_monitor.log_message("Benchmark complete")
+
+    def request_finished(self, result, response_text: str) -> None:
+        """Thread-safe callback matching ``OnFinishedCallback`` signature.
+
+        Enqueues ``(request_id, response_text)`` for the loop to drain.
+        Registered on ``RequestSender.on_finished`` by the orchestrator.
+        """
+        self._finished_queue.put((result.request_id, response_text))
+
+    def _drain_finished_queue(self) -> None:
+        """Drain all completed requests and call ``on_request_finished``."""
+        while True:
+            try:
+                request_id, output = self._finished_queue.get_nowait()
+            except queue.Empty:
+                break
+            self.on_request_finished(request_id, output)