wehub-resource-sync
145 lines
3.9 KiB
Python
145 lines
3.9 KiB
Python
import time
|
|
from enum import Enum
|
|
from typing import (
|
|
AsyncGenerator,
|
|
Callable,
|
|
List,
|
|
Set,
|
|
TypeVar,
|
|
)
|
|
|
|
from ray.util import metrics
|
|
|
|
# Histogram buckets for short-range latencies measurements:
|
|
# Min=1ms, Max=30s
|
|
#
|
|
# NOTE: Number of buckets have to be bounded (and not exceed 30)
|
|
# to avoid overloading metrics sub-system
|
|
SHORT_RANGE_LATENCY_HISTOGRAM_BUCKETS_MS: List[float] = [
|
|
1,
|
|
5,
|
|
10,
|
|
25,
|
|
50,
|
|
100,
|
|
150,
|
|
250,
|
|
500,
|
|
1000,
|
|
1500,
|
|
2500,
|
|
5000,
|
|
7500,
|
|
10000,
|
|
20000,
|
|
30000,
|
|
]
|
|
|
|
# Histogram buckets for long-range latencies measurements:
|
|
# Min=10ms, Max=300s
|
|
LONG_RANGE_LATENCY_HISTOGRAM_BUCKETS_MS = [
|
|
x * 10 for x in SHORT_RANGE_LATENCY_HISTOGRAM_BUCKETS_MS
|
|
]
|
|
|
|
|
|
class ClockUnit(int, Enum):
|
|
ms = 1000
|
|
s = 1
|
|
|
|
|
|
class MsClock:
|
|
"""A clock that tracks intervals in milliseconds"""
|
|
|
|
def __init__(self, unit: ClockUnit = ClockUnit.ms):
|
|
self.reset()
|
|
self.unit = unit.value
|
|
self.start_time = time.perf_counter()
|
|
|
|
def reset(self):
|
|
self.start_time = time.perf_counter()
|
|
|
|
def interval(self):
|
|
return (time.perf_counter() - self.start_time) * self.unit
|
|
|
|
def reset_interval(self):
|
|
interval = self.interval()
|
|
self.reset()
|
|
return interval
|
|
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class InstrumentTokenAsyncGenerator:
|
|
"""This class instruments an asynchronous generator.
|
|
|
|
It gathers 3 metrics:
|
|
1. Time to first time
|
|
2. Time between tokens
|
|
3. Total completion time
|
|
|
|
Usage:
|
|
|
|
@InstrumentTokenAsyncGenerator("my_special_fn")
|
|
async def to_instrument():
|
|
yield ...
|
|
"""
|
|
|
|
all_instrument_names: Set[str] = set()
|
|
|
|
def __init__(
|
|
self, generator_name: str, latency_histogram_buckets: List[float] = None
|
|
):
|
|
self.generator_name = f"rayllm_{generator_name}"
|
|
|
|
target_latency_histogram_buckets = (
|
|
latency_histogram_buckets or SHORT_RANGE_LATENCY_HISTOGRAM_BUCKETS_MS
|
|
)
|
|
|
|
assert (
|
|
self.generator_name not in self.all_instrument_names
|
|
), "This generator name was already used elsewhere. Please specify another one."
|
|
self.all_instrument_names.add(self.generator_name)
|
|
|
|
self.token_latency_histogram = metrics.Histogram(
|
|
f"{self.generator_name}_per_token_latency_ms",
|
|
f"Generator metrics for {self.generator_name}",
|
|
boundaries=target_latency_histogram_buckets,
|
|
)
|
|
|
|
self.first_token_latency_histogram = metrics.Histogram(
|
|
f"{self.generator_name}_first_token_latency_ms",
|
|
f"Generator metrics for {self.generator_name}",
|
|
boundaries=target_latency_histogram_buckets,
|
|
)
|
|
self.total_latency_histogram = metrics.Histogram(
|
|
f"{self.generator_name}_total_latency_ms",
|
|
f"Generator metrics for {self.generator_name}",
|
|
boundaries=target_latency_histogram_buckets,
|
|
)
|
|
|
|
def __call__(
|
|
self, async_generator_fn: Callable[..., AsyncGenerator[T, None]]
|
|
) -> Callable[..., AsyncGenerator[T, None]]:
|
|
async def new_gen(*args, **kwargs):
|
|
interval_clock = MsClock()
|
|
total_clock = MsClock()
|
|
is_first_token = True
|
|
try:
|
|
async for x in async_generator_fn(*args, **kwargs):
|
|
if is_first_token:
|
|
self.first_token_latency_histogram.observe(
|
|
total_clock.interval()
|
|
)
|
|
interval_clock.reset()
|
|
is_first_token = False
|
|
else:
|
|
self.token_latency_histogram.observe(
|
|
interval_clock.reset_interval()
|
|
)
|
|
yield x
|
|
finally:
|
|
self.total_latency_histogram.observe(total_clock.interval())
|
|
|
|
return new_gen
|