wehub-resource-sync
322 lines
11 KiB
Python
322 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Public API for Batched Stream, a wrapper of Streams that are batched together.
|
|
"""
|
|
|
|
# Future
|
|
from __future__ import annotations
|
|
|
|
# Standard
|
|
from collections.abc import Callable, Sequence
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from typing import Any
|
|
import time
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.cli.metrics import Metrics, StreamHandler, get_formatter
|
|
import lmcache.sdk.stream as lmc_stream
|
|
|
|
|
|
class LMCacheBatchedStreamError(RuntimeError):
|
|
"""Raised when a BatchedStream operation fails."""
|
|
|
|
|
|
class LMCacheBatchedStream:
|
|
"""A collection of LMCacheStreams run together as one batch."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Create a batch.
|
|
|
|
Attributes:
|
|
streams: Member streams keyed by stream id.
|
|
perf_metrics: Latest StreamPerfMetrics per stream id, overwritten
|
|
on each call to run_streams.
|
|
"""
|
|
self.streams: dict[str, lmc_stream.LMCacheStream] = {}
|
|
self.perf_metrics: dict[str, lmc_stream.StreamPerfMetrics] = {}
|
|
|
|
def add(self, stream: lmc_stream.LMCacheStream) -> None:
|
|
"""Add a stream to the batch, keyed by its stream id.
|
|
|
|
Args:
|
|
stream: The stream to add.
|
|
|
|
Raises:
|
|
LMCacheBatchedStreamError: If a stream with the same id is
|
|
in the batch.
|
|
"""
|
|
if stream.stream_id() in self.streams:
|
|
raise LMCacheBatchedStreamError(
|
|
f"Stream {stream.stream_id()} already exists."
|
|
)
|
|
self.streams[stream.stream_id()] = stream
|
|
|
|
def get_stream(self, stream_id: str) -> lmc_stream.LMCacheStream:
|
|
"""Return the stream registered under stream_id.
|
|
|
|
Args:
|
|
stream_id: The id of the stream to fetch.
|
|
|
|
Returns:
|
|
The matching LMCacheStream.
|
|
|
|
Raises:
|
|
LMCacheBatchedStreamError: If no stream has that id.
|
|
"""
|
|
if stream_id not in self.streams:
|
|
raise LMCacheBatchedStreamError(
|
|
f"Stream with id {stream_id} does not exist."
|
|
)
|
|
return self.streams[stream_id]
|
|
|
|
def run_streams(
|
|
self, sampling_params: dict[str, Any], stream_ids: list[str] | None = None
|
|
) -> float:
|
|
"""Decode every selected stream concurrently and store their metrics.
|
|
|
|
Submits stream.generate(sampling_params) on a thread pool (one worker
|
|
per stream) and records each StreamPerfMetrics in perf_metrics.
|
|
|
|
Args:
|
|
sampling_params: Engine sampling params (e.g. max_tokens,
|
|
temperature) passed to every stream's generate().
|
|
stream_ids: Streams to run; None means all streams in the batch.
|
|
|
|
Returns:
|
|
Wall-clock duration of the batch, in seconds (0.0 if no streams).
|
|
|
|
Raises:
|
|
LMCacheBatchedStreamError: If any requested stream id is unknown.
|
|
"""
|
|
if stream_ids is None:
|
|
stream_ids = list(self.streams.keys())
|
|
|
|
for stream_id in stream_ids:
|
|
if stream_id not in self.streams:
|
|
raise LMCacheBatchedStreamError(f"Stream {stream_id} does not exist.")
|
|
|
|
if not stream_ids:
|
|
return 0.0
|
|
|
|
self.perf_metrics.clear()
|
|
start_time = time.perf_counter()
|
|
with ThreadPoolExecutor(max_workers=len(stream_ids)) as executor:
|
|
futures = {}
|
|
for stream_id in stream_ids:
|
|
stream = self.streams[stream_id]
|
|
future = executor.submit(
|
|
stream.generate, sampling_params=sampling_params
|
|
)
|
|
futures[future] = stream_id
|
|
# collect metrics
|
|
for future in as_completed(futures):
|
|
stream_id = futures[future]
|
|
self.perf_metrics[stream_id] = future.result()
|
|
|
|
return time.perf_counter() - start_time # in seconds
|
|
|
|
def modify_stream(
|
|
self,
|
|
fn: Callable[[torch.Tensor, Sequence[int]], tuple[torch.Tensor, Sequence[int]]],
|
|
stream_ids: list[str] | None = None,
|
|
timeout: float = 30.0,
|
|
poll_interval: float = 0.2,
|
|
) -> float:
|
|
"""Apply KV edit to every selected stream concurrently.
|
|
|
|
Submits stream.modify_kv(fn) on a thread pool (one worker per stream).
|
|
|
|
Args:
|
|
fn: KV editor applied to each stream. Given (kv, tokens) and
|
|
returning (new_kv, new_tokens). See LMCacheStream.modify_kv().
|
|
stream_ids: Streams to modify. None means all streams in the batch.
|
|
timeout: Max seconds to wait for the cached KV to appear.
|
|
poll_interval: Seconds between retrieve attempts.
|
|
|
|
Returns:
|
|
Wall-clock duration of the batch, in seconds (0.0 if no streams).
|
|
|
|
Raises:
|
|
LMCacheBatchedStreamError: If any requested stream id is unknown.
|
|
"""
|
|
if stream_ids is None:
|
|
stream_ids = list(self.streams.keys())
|
|
|
|
for stream_id in stream_ids:
|
|
if stream_id not in self.streams:
|
|
raise LMCacheBatchedStreamError(f"Stream {stream_id} does not exist.")
|
|
|
|
if not stream_ids:
|
|
return 0.0
|
|
|
|
start_time = time.perf_counter()
|
|
|
|
with ThreadPoolExecutor(max_workers=len(stream_ids)) as executor:
|
|
futures = {}
|
|
for stream_id in stream_ids:
|
|
stream = self.streams[stream_id]
|
|
future = executor.submit(
|
|
stream.modify_kv, fn, timeout=timeout, poll_interval=poll_interval
|
|
)
|
|
futures[future] = stream_id
|
|
for future in as_completed(futures):
|
|
stream_id = futures[future]
|
|
future.result()
|
|
|
|
return time.perf_counter() - start_time # in seconds
|
|
|
|
def get_perf_metrics(
|
|
self,
|
|
duration: float,
|
|
fmt: str,
|
|
width: int,
|
|
mode: str,
|
|
stream_ids: list[str] | None = None,
|
|
) -> Metrics:
|
|
"""Aggregate the stored per-stream metrics into a Metrics report.
|
|
|
|
Args:
|
|
duration: Batch wall-clock duration in seconds, for throughput.
|
|
fmt: Formatter name for the output handler (e.g. "terminal").
|
|
width: Output width passed to the formatter.
|
|
mode: "prefill" (input tokens + TTFT) or "decode" (output tokens +
|
|
TPOT + decoding speed).
|
|
stream_ids: Streams to include. None means all streams in batch.
|
|
|
|
Returns:
|
|
The aggregated Metrics report.
|
|
|
|
Raises:
|
|
LMCacheBatchedStreamError: If mode is not "prefill" or "decode".
|
|
"""
|
|
if stream_ids is None:
|
|
stream_ids = list(self.streams.keys())
|
|
|
|
if mode not in ["prefill", "decode"]:
|
|
raise LMCacheBatchedStreamError(f"Invalid mode {mode}.")
|
|
|
|
metrics = Metrics(title=f"Batched Stream Metrics ({mode})")
|
|
metrics.add_handler(StreamHandler(get_formatter(fmt, width=width)))
|
|
|
|
cfg_section = metrics.add_section("config", "Configuration")
|
|
cfg_section.add("num_streams", "Number of Streams", len(self.streams))
|
|
|
|
result_section = metrics.add_section("results", "Results")
|
|
result_section.add("duration", "Total Duration (s)", duration)
|
|
|
|
input_tokens = 0
|
|
output_tokens = 0
|
|
for stream_id, perf_metrics in self.perf_metrics.items():
|
|
if stream_id in stream_ids:
|
|
input_tokens += perf_metrics.input_tokens
|
|
output_tokens += perf_metrics.output_tokens
|
|
|
|
if mode == "prefill":
|
|
result_section.add("input_tokens", "Total Input Tokens", input_tokens)
|
|
result_section.add(
|
|
"input_tput",
|
|
"Input Throughput (tokens/s)",
|
|
input_tokens / duration if duration > 0 else 0.0,
|
|
)
|
|
|
|
if mode == "decode":
|
|
result_section.add("output_tokens", "Total Output Tokens", output_tokens)
|
|
result_section.add(
|
|
"output_tput",
|
|
"Decode Throughput (tokens/s)",
|
|
output_tokens / duration if duration > 0 else 0.0,
|
|
)
|
|
|
|
return metrics
|
|
|
|
def prefill(
|
|
self,
|
|
sampling_params: dict[str, Any],
|
|
fmt: str = "terminal",
|
|
width: int = 80,
|
|
stream_ids: list[str] | None = None,
|
|
) -> Metrics:
|
|
"""Prefill every stream once (max_tokens forced to 1) and report.
|
|
|
|
Args:
|
|
sampling_params: Engine sampling params, max_tokens is set 1.
|
|
fmt: Formatter name for the output handler.
|
|
width: Output width passed to the formatter.
|
|
stream_ids: Streams to run; None means all streams in the batch.
|
|
|
|
Returns:
|
|
A prefill-mode Metrics report (input tokens, throughput, TTFT).
|
|
"""
|
|
sampling_params = {**sampling_params, "max_tokens": 1}
|
|
duration = self.run_streams(
|
|
sampling_params=sampling_params, stream_ids=stream_ids
|
|
)
|
|
return self.get_perf_metrics(
|
|
duration=duration,
|
|
fmt=fmt,
|
|
width=width,
|
|
mode="prefill",
|
|
stream_ids=stream_ids,
|
|
)
|
|
|
|
def modify(
|
|
self,
|
|
fn: Callable[[torch.Tensor, Sequence[int]], tuple[torch.Tensor, Sequence[int]]],
|
|
fmt: str = "terminal",
|
|
width: int = 80,
|
|
stream_ids: list[str] | None = None,
|
|
) -> Metrics:
|
|
"""Apply a KV edit to every stream and report the time taken.
|
|
|
|
Args:
|
|
fn: KV editor applied to each stream (see modify_stream).
|
|
fmt: Formatter name for the output handler.
|
|
width: Output width passed to the formatter.
|
|
stream_ids: Streams to modify. None means all streams in the batch.
|
|
|
|
Returns:
|
|
A Metrics report containing the total modify duration.
|
|
"""
|
|
duration = self.modify_stream(fn, stream_ids=stream_ids)
|
|
|
|
metrics = Metrics(title="Batched Stream Modify Metrics")
|
|
metrics.add_handler(StreamHandler(get_formatter(fmt, width=width)))
|
|
result_section = metrics.add_section("results", "Results")
|
|
result_section.add("duration", "Total Duration (s)", duration)
|
|
|
|
return metrics
|
|
|
|
def decode(
|
|
self,
|
|
sampling_params: dict[str, Any],
|
|
fmt: str = "terminal",
|
|
width: int = 80,
|
|
stream_ids: list[str] | None = None,
|
|
) -> Metrics:
|
|
"""Decode every stream and report decode metrics.
|
|
|
|
Args:
|
|
sampling_params: Engine sampling params (e.g. max_tokens,
|
|
temperature) passed to every stream's generate().
|
|
fmt: Formatter name for the output handler.
|
|
width: Output width passed to the formatter.
|
|
stream_ids: Streams to run. None means all streams in the batch.
|
|
|
|
Returns:
|
|
A decode-mode Metrics report (output tokens, throughput, TPOT,
|
|
decoding speed).
|
|
"""
|
|
duration = self.run_streams(
|
|
sampling_params=sampling_params, stream_ids=stream_ids
|
|
)
|
|
return self.get_perf_metrics(
|
|
duration=duration,
|
|
fmt=fmt,
|
|
width=width,
|
|
mode="decode",
|
|
stream_ids=stream_ids,
|
|
)
|