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

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

@@ -0,0 +1,281 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Long-document Q&A workload for ``lmcache bench engine``."""
+
+# Standard
+from dataclasses import dataclass
+import asyncio
+import random
+
+# 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.cli.commands.bench.engine_bench.workloads.base import BaseWorkload
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class LongDocQAConfig:
+    """Workload-specific config for the long-doc-qa workload."""
+
+    document_length: int = 10000
+    query_per_document: int = 2
+    num_documents: int = 1
+    shuffle_policy: str = "random"
+    num_inflight_requests: int = 3
+
+    def __post_init__(self) -> None:
+        if self.document_length <= 0:
+            raise ValueError(
+                f"document_length must be positive, got {self.document_length}"
+            )
+        if self.query_per_document < 1:
+            raise ValueError(
+                f"query_per_document must be >= 1, got {self.query_per_document}"
+            )
+        if self.num_inflight_requests < 1:
+            raise ValueError(
+                f"num_inflight_requests must be >= 1, got {self.num_inflight_requests}"
+            )
+        if self.shuffle_policy not in ("random", "tile"):
+            raise ValueError(
+                f"shuffle_policy must be 'random' or 'tile', "
+                f"got {self.shuffle_policy!r}"
+            )
+
+    @classmethod
+    def resolve(
+        cls,
+        kv_cache_volume_gb: float,
+        tokens_per_gb_kvcache: int,
+        document_length: int = 10000,
+        query_per_document: int = 2,
+        shuffle_policy: str = "random",
+        num_inflight_requests: int = 3,
+    ) -> "LongDocQAConfig":
+        """Create a config with ``num_documents`` computed from KV cache budget.
+
+        Args:
+            kv_cache_volume_gb: Target active KV cache volume in GB.
+            tokens_per_gb_kvcache: Tokens fitting in 1 GB of KV cache.
+            document_length: Token length of each document.
+            query_per_document: Number of questions per document.
+            shuffle_policy: Request ordering — ``"random"`` or ``"tile"``.
+            num_inflight_requests: Max concurrent in-flight requests.
+
+        Returns:
+            A fully-resolved LongDocQAConfig with computed num_documents.
+        """
+        num_documents = max(
+            int(kv_cache_volume_gb * tokens_per_gb_kvcache / document_length),
+            1,
+        )
+        logger.debug(
+            "Computed num_documents=%d from kv_cache_volume_gb=%.1f, "
+            "tokens_per_gb_kvcache=%d, document_length=%d",
+            num_documents,
+            kv_cache_volume_gb,
+            tokens_per_gb_kvcache,
+            document_length,
+        )
+        return cls(
+            document_length=document_length,
+            query_per_document=query_per_document,
+            num_documents=num_documents,
+            shuffle_policy=shuffle_policy,
+            num_inflight_requests=num_inflight_requests,
+        )
+
+
+_QUESTIONS = [
+    "What is this document about?",
+    "Summarize the key points.",
+    "What is the main topic discussed?",
+    "Provide a brief overview.",
+]
+
+
+class LongDocQAWorkload(BaseWorkload):
+    """Workload that simulates repeated Q&A over long documents.
+
+    Generates synthetic documents, builds a request schedule, warms up
+    the KV cache by sending each document once, then dispatches benchmark
+    requests with semaphore-controlled concurrency.
+    """
+
+    def __init__(
+        self,
+        config: LongDocQAConfig,
+        request_sender: RequestSender,
+        stats_collector: StatsCollector,
+        progress_monitor: ProgressMonitor,
+        seed: int = 42,
+    ) -> None:
+        super().__init__(request_sender, stats_collector, progress_monitor)
+        self._config = config
+        self._seed = seed
+
+        self._documents = self._generate_documents()
+        self._schedule = self._build_schedule()
+        self._schedule_index = 0
+
+        self._semaphore = asyncio.Semaphore(config.num_inflight_requests)
+        self._pending_tasks: set[asyncio.Task] = set()
+
+    def log_config(self) -> None:
+        """Log key workload config before the benchmark starts."""
+        c = self._config
+        B = "\033[1m"  # bold
+        C = "\033[96m"  # cyan
+        Y = "\033[93m"  # yellow
+        R = "\033[0m"  # reset
+        total = c.num_documents * c.query_per_document
+        print(
+            f"{B}{'═' * 50}{R}\n"
+            f"{B} Workload: {C}long-doc-qa{R}\n"
+            f"{B}{'─' * 50}{R}\n"
+            f"  Documents:        {Y}{c.num_documents}{R}\n"
+            f"  Queries/doc:      {Y}{c.query_per_document}{R}\n"
+            f"  Total requests:   {Y}{total}{R}\n"
+            f"  Document length:  {Y}{c.document_length}{R} tokens\n"
+            f"  Max inflight:     {Y}{c.num_inflight_requests}{R}\n"
+            f"  Shuffle policy:   {Y}{c.shuffle_policy}{R}\n"
+            f"{B}{'═' * 50}{R}"
+        )
+
+    # ------------------------------------------------------------------
+    # Data generation
+    # ------------------------------------------------------------------
+
+    def _generate_documents(self) -> list[str]:
+        """Generate synthetic documents of approximately ``document_length`` tokens."""
+        documents = []
+        for doc_id in range(self._config.num_documents):
+            prefix = f"Document {doc_id}: "
+            body = " ".join(["hi"] * max(self._config.document_length - 10, 1))
+            documents.append(prefix + body)
+        logger.debug(
+            "Generated %d documents of ~%d tokens each",
+            len(documents),
+            self._config.document_length,
+        )
+        return documents
+
+    def _build_schedule(self) -> list[tuple[int, int]]:
+        """Build the request schedule as ``(doc_index, query_index)`` pairs.
+
+        Tile policy uses query-major order: all documents for query 0,
+        then all documents for query 1, etc.  Random policy shuffles
+        the same pairs with a seeded RNG.
+        """
+        schedule: list[tuple[int, int]] = []
+        for q_idx in range(self._config.query_per_document):
+            for doc_idx in range(self._config.num_documents):
+                schedule.append((doc_idx, q_idx))
+
+        if self._config.shuffle_policy == "random":
+            rng = random.Random(self._seed)
+            rng.shuffle(schedule)
+
+        logger.debug(
+            "Built schedule with %d requests (policy=%s)",
+            len(schedule),
+            self._config.shuffle_policy,
+        )
+        return schedule
+
+    def _build_messages(
+        self,
+        doc_index: int,
+        query_index: int,
+    ) -> list[dict[str, str]]:
+        """Build chat messages for a benchmark request."""
+        document = self._documents[doc_index]
+        question = _QUESTIONS[query_index % len(_QUESTIONS)]
+        content = f"{document}\n\nQuestion {query_index}: {question}"
+        return [{"role": "user", "content": content}]
+
+    # ------------------------------------------------------------------
+    # Warmup
+    # ------------------------------------------------------------------
+
+    async def warmup(self) -> None:
+        """Send each document once with ``max_tokens=1`` to populate KV cache."""
+        num_docs = self._config.num_documents
+        for doc_idx in range(num_docs):
+            request_id = f"warmup_doc{doc_idx}"
+            messages = [{"role": "user", "content": self._documents[doc_idx]}]
+            self._progress_monitor.log_message(f"Warmup {doc_idx + 1}/{num_docs}")
+            self._progress_monitor.on_request_sent(request_id)
+            result = await self._request_sender.send_warmup_request(
+                request_id,
+                messages,
+            )
+            if not result.successful:
+                self._progress_monitor.log_message(
+                    f"Warmup {request_id} failed: {result.error}"
+                )
+        self._progress_monitor.log_message(
+            f"Warmup complete: {num_docs} documents sent",
+        )
+
+    # ------------------------------------------------------------------
+    # Benchmark dispatch
+    # ------------------------------------------------------------------
+
+    async def step(self, time_offset: float) -> float:
+        """Dispatch the next request if semaphore allows."""
+        if self._schedule_index < len(self._schedule):
+            await self._semaphore.acquire()
+            doc_idx, q_idx = self._schedule[self._schedule_index]
+            self._schedule_index += 1
+
+            task = asyncio.create_task(self._dispatch(doc_idx, q_idx))
+            self._pending_tasks.add(task)
+            task.add_done_callback(self._on_task_done)
+            return 0.0  # immediate re-call
+
+        # All dispatched — wait for pending tasks.
+        if self._pending_tasks:
+            await asyncio.wait(
+                self._pending_tasks,
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+            return 0.0
+
+        return -1.0  # all done
+
+    async def _dispatch(
+        self,
+        doc_index: int,
+        query_index: int,
+    ) -> None:
+        """Send a single benchmark request, then release the semaphore."""
+        request_id = f"doc{doc_index}_q{query_index}"
+        messages = self._build_messages(doc_index, query_index)
+        self._progress_monitor.on_request_sent(request_id)
+        self._progress_monitor.log_message(
+            f"Dispatched request {request_id} (doc {doc_index}, query {query_index})"
+        )
+        try:
+            await self._request_sender.send_request(
+                request_id,
+                messages,
+            )
+        finally:
+            self._semaphore.release()
+
+    def _on_task_done(self, task: asyncio.Task) -> None:
+        """Clean up completed tasks and log unexpected errors."""
+        self._pending_tasks.discard(task)
+        if not task.cancelled():
+            exc = task.exception()
+            if exc is not None:
+                self._progress_monitor.log_message(f"Dispatch task failed: {exc}")
+
+    def on_request_finished(self, request_id: str, output: str) -> None:
+        """No-op — this workload is stateless."""

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

@@ -0,0 +1,337 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-round chat workload for ``lmcache bench engine``."""
+
+# Standard
+from dataclasses import dataclass, field
+import asyncio
+
+# 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.cli.commands.bench.engine_bench.workloads.base import BaseWorkload
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class MultiRoundChatConfig:
+    """Workload-specific config for the multi-round-chat workload."""
+
+    shared_prompt_length: int = 2000
+    chat_history_length: int = 10000
+    user_input_length: int = 50
+    output_length: int = 200
+    qps: float = 1.0
+    duration: float = 60.0
+    num_concurrent_users: int = 1
+
+    def __post_init__(self) -> None:
+        if self.shared_prompt_length <= 0:
+            raise ValueError(
+                f"shared_prompt_length must be positive, "
+                f"got {self.shared_prompt_length}"
+            )
+        if self.chat_history_length <= 0:
+            raise ValueError(
+                f"chat_history_length must be positive, got {self.chat_history_length}"
+            )
+        if self.user_input_length < 1:
+            raise ValueError(
+                f"user_input_length must be >= 1, got {self.user_input_length}"
+            )
+        if self.output_length < 1:
+            raise ValueError(f"output_length must be >= 1, got {self.output_length}")
+        if self.qps <= 0:
+            raise ValueError(f"qps must be positive, got {self.qps}")
+        if self.duration <= 0:
+            raise ValueError(f"duration must be positive, got {self.duration}")
+        if self.num_concurrent_users < 1:
+            raise ValueError(
+                f"num_concurrent_users must be >= 1, got {self.num_concurrent_users}"
+            )
+
+    @classmethod
+    def resolve(
+        cls,
+        kv_cache_volume_gb: float,
+        tokens_per_gb_kvcache: int,
+        shared_prompt_length: int = 2000,
+        chat_history_length: int = 10000,
+        user_input_length: int = 50,
+        output_length: int = 200,
+        qps: float = 1.0,
+        duration: float = 60.0,
+    ) -> "MultiRoundChatConfig":
+        """Create a config with ``num_concurrent_users`` computed from KV cache budget.
+
+        Args:
+            kv_cache_volume_gb: Target active KV cache volume in GB.
+            tokens_per_gb_kvcache: Tokens fitting in 1 GB of KV cache.
+            shared_prompt_length: Token length of the system prompt.
+            chat_history_length: Token length of pre-filled history.
+            user_input_length: Token length per user query.
+            output_length: Max tokens to generate per response.
+            qps: Queries per second.
+            duration: Benchmark duration in seconds.
+
+        Returns:
+            A fully-resolved MultiRoundChatConfig.
+        """
+        tokens_per_session = shared_prompt_length + chat_history_length
+        total_tokens = kv_cache_volume_gb * tokens_per_gb_kvcache
+        num_users = max(1, int(total_tokens / tokens_per_session))
+        logger.debug(
+            "Computed num_concurrent_users=%d from kv_cache_volume_gb=%.1f, "
+            "tokens_per_gb_kvcache=%d, tokens_per_session=%d",
+            num_users,
+            kv_cache_volume_gb,
+            tokens_per_gb_kvcache,
+            tokens_per_session,
+        )
+        return cls(
+            shared_prompt_length=shared_prompt_length,
+            chat_history_length=chat_history_length,
+            user_input_length=user_input_length,
+            output_length=output_length,
+            qps=qps,
+            duration=duration,
+            num_concurrent_users=num_users,
+        )
+
+
+@dataclass
+class Session:
+    """A single stateful chat session.
+
+    Maintains a system prompt, pre-filled history, and accumulated
+    Q&A exchanges.  The ``in_flight`` flag tracks whether a request
+    is currently pending for this session.
+    """
+
+    session_id: int
+    system_prompt: str
+    history_text: str
+    exchanges: list[tuple[str, str]] = field(default_factory=list)
+    in_flight: bool = False
+
+    def build_messages(self, query: str) -> list[dict[str, str]]:
+        """Construct OpenAI-format messages for a request."""
+        messages: list[dict[str, str]] = [
+            {"role": "system", "content": self.system_prompt},
+        ]
+        if self.history_text:
+            messages.append(
+                {"role": "user", "content": self.history_text},
+            )
+            messages.append(
+                {
+                    "role": "assistant",
+                    "content": "Understood, I have read the context above.",
+                },
+            )
+        for q, a in self.exchanges:
+            messages.append({"role": "user", "content": q})
+            messages.append({"role": "assistant", "content": a})
+        messages.append({"role": "user", "content": query})
+        return messages
+
+    def record_answer(self, query: str, answer: str) -> None:
+        """Record a completed exchange and mark session as ready."""
+        self.exchanges.append((query, answer))
+        self.in_flight = False
+
+
+class MultiRoundChatWorkload(BaseWorkload):
+    """Workload that simulates multi-round chat with stateful sessions.
+
+    Creates multiple concurrent user sessions, dispatches requests at a
+    fixed QPS rate using round-robin scheduling, and records responses
+    in session history so subsequent queries include prior context.
+    """
+
+    def __init__(
+        self,
+        config: MultiRoundChatConfig,
+        request_sender: RequestSender,
+        stats_collector: StatsCollector,
+        progress_monitor: ProgressMonitor,
+        seed: int = 42,
+    ) -> None:
+        super().__init__(request_sender, stats_collector, progress_monitor)
+        self._config = config
+        self._seed = seed
+
+        self._sessions = self._create_sessions()
+        self._global_index = 0
+        self._interval = 1.0 / config.qps
+        self._pending_info: dict[str, tuple[int, str]] = {}
+        self._pending_tasks: set[asyncio.Task] = set()
+
+    def log_config(self) -> None:
+        """Log key workload config before the benchmark starts."""
+        c = self._config
+        B = "\033[1m"  # bold
+        C = "\033[96m"  # cyan
+        Y = "\033[93m"  # yellow
+        R = "\033[0m"  # reset
+        print(
+            f"{B}{'═' * 50}{R}\n"
+            f"{B} Workload: {C}multi-round-chat{R}\n"
+            f"{B}{'─' * 50}{R}\n"
+            f"  Concurrent users: {Y}{c.num_concurrent_users}{R}\n"
+            f"  Prompt length:    {Y}{c.shared_prompt_length}{R} tokens\n"
+            f"  History length:   {Y}{c.chat_history_length}{R} tokens\n"
+            f"  Query length:     {Y}{c.user_input_length}{R} tokens\n"
+            f"  Output length:    {Y}{c.output_length}{R} tokens\n"
+            f"  QPS:              {Y}{c.qps}{R}\n"
+            f"  Duration:         {Y}{c.duration}s{R}\n"
+            f"{B}{'═' * 50}{R}"
+        )
+
+    # ------------------------------------------------------------------
+    # Session creation
+    # ------------------------------------------------------------------
+
+    def _create_sessions(self) -> list[Session]:
+        """Create sessions with synthetic prompts and history."""
+        sessions: list[Session] = []
+        for i in range(self._config.num_concurrent_users):
+            system_prompt = self._generate_system_prompt(i)
+            history = self._generate_history(i)
+            sessions.append(
+                Session(
+                    session_id=i,
+                    system_prompt=system_prompt,
+                    history_text=history,
+                )
+            )
+        logger.debug(
+            "Created %d sessions (prompt=%d, history=%d tokens each)",
+            len(sessions),
+            self._config.shared_prompt_length,
+            self._config.chat_history_length,
+        )
+        return sessions
+
+    def _generate_system_prompt(self, session_id: int) -> str:
+        """Generate a system prompt of approximately ``shared_prompt_length`` tokens."""
+        prefix = f"Session {session_id}. You are a helpful assistant. "
+        remaining = max(0, self._config.shared_prompt_length - len(prefix.split()))
+        return prefix + " ".join(["help"] * remaining)
+
+    def _generate_history(self, session_id: int) -> str:
+        """
+        Generate pre-filled history of approximately ``chat_history_length`` tokens.
+        """
+        prefix = f"[Session {session_id} history] "
+        remaining = max(0, self._config.chat_history_length - len(prefix.split()))
+        return prefix + " ".join(["hi"] * remaining)
+
+    def _generate_query(self) -> str:
+        """Generate a user query of approximately ``user_input_length`` tokens."""
+        return " ".join(["tell"] * self._config.user_input_length)
+
+    # ------------------------------------------------------------------
+    # Warmup
+    # ------------------------------------------------------------------
+
+    async def warmup(self) -> None:
+        """Send one warmup request per session with ``max_tokens=1``."""
+        num_sessions = len(self._sessions)
+        for session in self._sessions:
+            request_id = f"warmup_s{session.session_id}"
+            messages = session.build_messages("Hello")
+            self._progress_monitor.log_message(
+                f"Warmup {session.session_id + 1}/{num_sessions}"
+            )
+            self._progress_monitor.on_request_sent(request_id)
+            result = await self._request_sender.send_warmup_request(
+                request_id,
+                messages,
+            )
+            if not result.successful:
+                self._progress_monitor.log_message(
+                    f"Warmup session {session.session_id} failed: {result.error}"
+                )
+        self._progress_monitor.log_message(f"Warmup complete: {num_sessions} sessions")
+
+    # ------------------------------------------------------------------
+    # Benchmark dispatch
+    # ------------------------------------------------------------------
+
+    async def step(self, time_offset: float) -> float:
+        """Dispatch the next request at the QPS-controlled rate.
+
+        Returns:
+            Next wakeup time offset, or negative when done.
+        """
+        # Check duration — stop dispatching new requests
+        if time_offset >= self._config.duration:
+            if self._pending_tasks:
+                await asyncio.wait(
+                    self._pending_tasks,
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                return 0.0
+            return -1.0
+
+        # Find ready session via round-robin
+        target_idx = self._global_index % len(self._sessions)
+        session = self._sessions[target_idx]
+
+        if session.in_flight:
+            # Session busy — sleep briefly, let loop drain queue
+            return time_offset + 0.01
+
+        # Dispatch request
+        query = self._generate_query()
+        request_id = f"s{session.session_id}_r{self._global_index}"
+        messages = session.build_messages(query)
+
+        session.in_flight = True
+        self._pending_info[request_id] = (session.session_id, query)
+        self._progress_monitor.on_request_sent(request_id)
+        self._progress_monitor.log_message(
+            f"Session {session.session_id} dispatched request {self._global_index}"
+        )
+
+        task = asyncio.create_task(
+            self._dispatch(request_id, messages),
+        )
+        self._pending_tasks.add(task)
+        task.add_done_callback(self._on_task_done)
+
+        self._global_index += 1
+        return self._global_index * self._interval
+
+    async def _dispatch(
+        self,
+        request_id: str,
+        messages: list[dict[str, str]],
+    ) -> None:
+        """Send a single benchmark request."""
+        await self._request_sender.send_request(
+            request_id,
+            messages,
+            max_tokens=self._config.output_length,
+        )
+
+    def _on_task_done(self, task: asyncio.Task) -> None:
+        """Clean up completed tasks and log unexpected errors."""
+        self._pending_tasks.discard(task)
+        if not task.cancelled():
+            exc = task.exception()
+            if exc is not None:
+                self._progress_monitor.log_message(f"Dispatch task failed: {exc}")
+
+    def on_request_finished(self, request_id: str, output: str) -> None:
+        """Record the response in the session's conversation history."""
+        info = self._pending_info.pop(request_id, None)
+        if info is None:
+            return  # warmup request or already processed
+        session_id, query = info
+        self._sessions[session_id].record_answer(query, output)