Files
ray-project--ray/python/ray/serve/_private/benchmarks/common.py
T
2026-07-13 13:17:40 +08:00

689 lines
22 KiB
Python

import asyncio
import inspect
import logging
import random
import string
import time
from functools import partial
from typing import Any, Callable, Coroutine, Dict, List, Optional, Tuple
import aiohttp
import aiohttp.client_exceptions
import grpc
import numpy as np
import pandas as pd
from starlette.responses import StreamingResponse
from tqdm import tqdm
import ray
from ray import serve
from ray._common.test_utils import SignalActor as _SignalActor
from ray.serve._private.common import DeploymentStatus
from ray.serve.generated import serve_pb2, serve_pb2_grpc
from ray.serve.handle import DeploymentHandle
async def run_latency_benchmark(
f: Callable, num_requests: int, *, num_warmup_requests: int = 100
) -> pd.Series:
if inspect.iscoroutinefunction(f):
to_call = f
else:
async def to_call():
f()
latencies = []
for i in tqdm(range(num_requests + num_warmup_requests)):
start = time.perf_counter()
await to_call()
end = time.perf_counter()
# Don't include warm-up requests.
if i >= num_warmup_requests:
latencies.append(1000 * (end - start))
return pd.Series(latencies)
async def run_throughput_benchmark(
fn: Callable[[], List[float]],
multiplier: int = 1,
num_trials: int = 10,
trial_runtime: float = 1,
) -> Tuple[float, float, pd.Series]:
"""Benchmarks throughput of a function.
Args:
fn: The function to benchmark. If this returns anything, it must
return a list of latencies.
multiplier: The number of requests or tokens (or whatever unit
is appropriate for this throughput benchmark) that is
completed in one call to `fn`.
num_trials: The number of trials to run.
trial_runtime: How long each trial should run for. During the
duration of one trial, `fn` will be repeatedly called.
Returns:
A tuple ``(mean, stddev, latencies)`` summarizing per-trial throughput
across ``num_trials`` runs.
"""
# Warmup
start = time.time()
while time.time() - start < 0.1:
await fn()
# Benchmark
stats = []
latencies = []
for _ in tqdm(range(num_trials)):
start = time.perf_counter()
count = 0
while time.perf_counter() - start < trial_runtime:
res = await fn()
if res:
latencies.extend(res)
count += 1
end = time.perf_counter()
stats.append(multiplier * count / (end - start))
return round(np.mean(stats), 2), round(np.std(stats), 2), pd.Series(latencies)
async def do_single_http_batch(
*,
batch_size: int = 100,
url: str = "http://localhost:8000",
stream: bool = False,
) -> List[float]:
"""Sends a batch of http requests and returns e2e latencies."""
# By default, aiohttp limits the number of client connections to 100.
# We need to use TCPConnector to configure the limit if batch size
# is greater than 100.
connector = aiohttp.TCPConnector(limit=batch_size)
async with aiohttp.ClientSession(
connector=connector, raise_for_status=True
) as session:
async def do_query():
start = time.perf_counter()
try:
async with session.get(url) as r:
if stream:
async for chunk, _ in r.content.iter_chunks():
pass
else:
# Read the response to ensure it's consumed
await r.read()
except aiohttp.client_exceptions.ClientConnectionError:
pass
end = time.perf_counter()
return 1000 * (end - start)
return await asyncio.gather(*[do_query() for _ in range(batch_size)])
async def do_single_grpc_batch(
*, batch_size: int = 100, target: str = "localhost:9000"
):
channel = grpc.aio.insecure_channel(target)
stub = serve_pb2_grpc.RayServeBenchmarkServiceStub(channel)
payload = serve_pb2.StringData(data="")
async def do_query():
start = time.perf_counter()
await stub.grpc_call(payload)
end = time.perf_counter()
return 1000 * (end - start)
return await asyncio.gather(*[do_query() for _ in range(batch_size)])
async def collect_profile_events(coro: Coroutine):
"""Collects profiling events using Viztracer"""
from viztracer import VizTracer
tracer = VizTracer()
tracer.start()
await coro
tracer.stop()
tracer.save()
def generate_payload(size: int = 100, chars=string.ascii_uppercase + string.digits):
return "".join(random.choice(chars) for _ in range(size))
class Blackhole:
def sink(self, o):
pass
@serve.deployment
class Noop:
def __init__(self):
logging.getLogger("ray.serve").setLevel(logging.WARNING)
def __call__(self, *args, **kwargs):
return b""
@serve.deployment
class ModelComp:
def __init__(self, child):
logging.getLogger("ray.serve").setLevel(logging.WARNING)
self._child = child
async def __call__(self, *args, **kwargs):
return await self._child.remote()
@serve.deployment
class GrpcDeployment:
def __init__(self):
logging.getLogger("ray.serve").setLevel(logging.WARNING)
async def grpc_call(self, user_message):
return serve_pb2.ModelOutput(output=9)
async def call_with_string(self, user_message):
return serve_pb2.ModelOutput(output=9)
@serve.deployment
class GrpcModelComp:
def __init__(self, child):
logging.getLogger("ray.serve").setLevel(logging.WARNING)
self._child = child
async def grpc_call(self, user_message):
await self._child.remote()
return serve_pb2.ModelOutput(output=9)
async def call_with_string(self, user_message):
await self._child.remote()
return serve_pb2.ModelOutput(output=9)
@serve.deployment
class Streamer:
def __init__(self, tokens_per_request: int, inter_token_delay_ms: int = 10):
logging.getLogger("ray.serve").setLevel(logging.WARNING)
self._tokens_per_request = tokens_per_request
self._inter_token_delay_s = inter_token_delay_ms / 1000
async def stream(self):
for _ in range(self._tokens_per_request):
await asyncio.sleep(self._inter_token_delay_s)
yield b"hi"
async def __call__(self):
return StreamingResponse(self.stream())
@serve.deployment
class IntermediateRouter:
def __init__(self, handle: DeploymentHandle):
logging.getLogger("ray.serve").setLevel(logging.WARNING)
self._handle = handle.options(stream=True)
async def stream(self):
async for token in self._handle.stream.remote():
yield token
def __call__(self):
return StreamingResponse(self.stream())
@serve.deployment
class Benchmarker:
def __init__(
self,
handle: DeploymentHandle,
stream: bool = False,
):
logging.getLogger("ray.serve").setLevel(logging.WARNING)
self._handle = handle.options(stream=stream)
self._stream = stream
async def do_single_request(self, payload: Any = None) -> float:
"""Completes a single unary request. Returns e2e latency in ms."""
start = time.perf_counter()
if payload is None:
await self._handle.remote()
else:
await self._handle.remote(payload)
end = time.perf_counter()
return 1000 * (end - start)
async def do_single_choose_dispatch(self, payload: Any = None) -> float:
"""Completes a single unary request via choose_replica + dispatch.
Returns e2e latency in ms. With SingletonThreadRouter this involves
two run_coroutine_threadsafe round-trips (one for __aenter__, one
for _dispatch_to_marked_selection) vs. one for ``remote``.
"""
start = time.perf_counter()
if payload is None:
async with self._handle.choose_replica() as sel:
await self._handle.dispatch(sel)
else:
async with self._handle.choose_replica(payload) as sel:
await self._handle.dispatch(sel, payload)
end = time.perf_counter()
return 1000 * (end - start)
async def _do_single_stream(self) -> float:
"""Consumes a single streaming request. Returns e2e latency in ms."""
start = time.perf_counter()
async for r in self._handle.stream.remote():
pass
end = time.perf_counter()
return 1000 * (end - start)
async def _do_single_batch(self, batch_size: int) -> List[float]:
if self._stream:
return await asyncio.gather(
*[self._do_single_stream() for _ in range(batch_size)]
)
else:
return await asyncio.gather(
*[self.do_single_request() for _ in range(batch_size)]
)
async def run_latency_benchmark(
self,
*,
num_requests: int,
payload: Any = None,
mode: str = "remote",
) -> pd.Series:
if mode == "remote":
async def f():
await self.do_single_request(payload)
elif mode == "choose_dispatch":
async def f():
await self.do_single_choose_dispatch(payload)
else:
raise ValueError(f"Unknown mode {mode!r}")
return await run_latency_benchmark(f, num_requests=num_requests)
async def run_throughput_benchmark(
self,
*,
batch_size: int,
num_trials: int,
trial_runtime: float,
tokens_per_request: Optional[float] = None,
) -> Tuple[float, float]:
if self._stream:
assert tokens_per_request
multiplier = tokens_per_request * batch_size
else:
multiplier = batch_size
return await run_throughput_benchmark(
fn=partial(
self._do_single_batch,
batch_size=batch_size,
),
multiplier=multiplier,
num_trials=num_trials,
trial_runtime=trial_runtime,
)
# =============================================================================
# Controller Benchmark
# =============================================================================
# See https://github.com/ray-project/ray/issues/60680 for more details.
CONTROLLER_BENCH_CONFIG = {
"checkpoints": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 3072, 4096],
"marination_period_s": 180,
"sample_interval_s": 5,
}
_CONTROLLER_AUTOSCALING_CONFIG = {
"min_replicas": 1,
"max_replicas": 4096,
"target_ongoing_requests": 1,
"upscale_delay_s": 1,
}
_CONTROLLER_WAITER_TIMEOUT_S = 1200
# SignalActor from ray._common.test_utils; use high max_concurrency for many
# concurrent waiters (up to 4096 in controller benchmark).
_SignalActorForController = _SignalActor.options(max_concurrency=100000)
@serve.deployment(
graceful_shutdown_timeout_s=1,
ray_actor_options={"num_cpus": 0.2},
max_ongoing_requests=100000,
autoscaling_config={
"min_replicas": 5,
"max_replicas": 10,
"target_ongoing_requests": 100000,
"upscale_delay_s": 1,
},
)
class ControllerBenchHelloWorld:
def __init__(self, signal_actor):
self.signal = signal_actor
async def __call__(self):
await self.signal.wait.remote()
return "hello"
@serve.deployment(
autoscaling_config=_CONTROLLER_AUTOSCALING_CONFIG,
max_ongoing_requests=2,
graceful_shutdown_timeout_s=1,
ray_actor_options={"num_cpus": 0.4},
)
class ControllerBenchMetricsGenerator:
"""Autoscaling deployment that generates handle metrics to stress the controller."""
def __init__(self, hello_world: DeploymentHandle):
self.hello_world = hello_world
async def __call__(self):
return await self.hello_world.remote()
def _controller_get_active_nodes() -> int:
"""Get number of active nodes in the cluster."""
return len([n for n in ray.nodes() if n.get("Alive", False)])
async def _controller_get_replica_count(
deployment_name: str = "ControllerBenchMetricsGenerator",
) -> int:
"""Get current number of running replicas for the specified deployment."""
status = serve.status()
for app in status.applications.values():
for name, deployment in app.deployments.items():
if name == deployment_name:
return deployment.replica_states.get("RUNNING", 0)
return 0
async def _controller_get_health_metrics() -> Dict[str, Any]:
"""Get controller health metrics. Fails the run if unavailable."""
client = serve.context._global_client
if client is None:
raise RuntimeError(
"Serve is not connected. get_health_metrics requires an active Serve "
"controller. Ensure Serve is started before running the controller benchmark."
)
controller = client._controller
if not hasattr(controller, "get_health_metrics"):
raise RuntimeError(
"Controller does not have get_health_metrics. This API is required for "
"the controller benchmark. Please use a Ray version that supports "
"controller health metrics."
)
return await controller.get_health_metrics.remote()
def _controller_extract_metrics_row(
health_metrics: Dict[str, Any],
checkpoint: int,
sample: int,
target_replicas: int,
actual_replicas: int,
num_nodes: int,
autoscale_duration_s: float,
) -> Dict[str, Any]:
"""Extract a flat row from health metrics with all available fields."""
def get_stat(d: dict, key: str, stat: str, default=0):
return d.get(key, {}).get(stat, default)
return {
"checkpoint": checkpoint,
"sample": sample,
"target_replicas": target_replicas,
"actual_replicas": actual_replicas,
"num_nodes": num_nodes,
"autoscale_duration_s": round(autoscale_duration_s, 3),
"loop_duration_mean_s": get_stat(health_metrics, "loop_duration_s", "mean"),
"loop_duration_std_s": get_stat(health_metrics, "loop_duration_s", "std"),
"loops_per_second": health_metrics.get("loops_per_second", 0),
"event_loop_delay_s": health_metrics.get("event_loop_delay_s", 0),
"num_asyncio_tasks": health_metrics.get("num_asyncio_tasks", 0),
"deployment_state_update_mean_s": get_stat(
health_metrics, "deployment_state_update_duration_s", "mean"
),
"application_state_update_mean_s": get_stat(
health_metrics, "application_state_update_duration_s", "mean"
),
"proxy_state_update_mean_s": get_stat(
health_metrics, "proxy_state_update_duration_s", "mean"
),
"proxy_state_update_std_s": get_stat(
health_metrics, "proxy_state_update_duration_s", "std"
),
"node_update_mean_s": get_stat(
health_metrics, "node_update_duration_s", "mean"
),
"node_update_std_s": get_stat(health_metrics, "node_update_duration_s", "std"),
"node_update_min_s": get_stat(health_metrics, "node_update_duration_s", "min"),
"handle_metrics_delay_mean_ms": get_stat(
health_metrics, "handle_metrics_delay_ms", "mean"
),
"replica_metrics_delay_mean_ms": get_stat(
health_metrics, "replica_metrics_delay_ms", "mean"
),
"process_memory_mb": health_metrics.get("process_memory_mb", 0),
}
async def _controller_wait_for_replicas_up(target: int, timeout: float = 300) -> float:
start = time.time()
while time.time() - start < timeout:
actual = await _controller_get_replica_count()
if actual >= target:
return time.time() - start
if int(time.time() - start) % 10 == 0:
logging.info(f"Waiting for {target} replicas... {actual}/{target}")
await asyncio.sleep(0.5)
actual = await _controller_get_replica_count()
raise RuntimeError(
f"Timeout: Only {actual}/{target} replicas after {timeout}s. Ending experiment."
)
async def _controller_wait_for_waiters(
signal_actor, expected: int, timeout: float = 300
) -> float:
start = time.time()
while time.time() - start < timeout:
num_waiters = await signal_actor.cur_num_waiters.remote()
if num_waiters >= expected:
return time.time() - start
await asyncio.sleep(0.5)
if int(time.time() - start) % 10 == 0:
logging.info(f"Waiting for {expected} waiters... {num_waiters}/{expected}")
num_waiters = await signal_actor.cur_num_waiters.remote()
raise RuntimeError(
f"Timeout: Only {num_waiters}/{expected} requests reached replicas after "
f"{timeout}s. Ending experiment."
)
async def _controller_wait_for_deployment_healthy(
deployment_name: str = "ControllerBenchMetricsGenerator",
app_name: str = "default",
timeout: float = 60,
) -> None:
"""Wait for the deployment to enter HEALTHY status via serve.status()."""
start = time.time()
while time.time() - start < timeout:
status = serve.status()
app = status.applications.get(app_name)
dep_status = None
if app and deployment_name in app.deployments:
dep = app.deployments[deployment_name]
dep_status = dep.status
if dep_status == DeploymentStatus.HEALTHY:
return
if dep_status == DeploymentStatus.UNHEALTHY:
raise RuntimeError(
f"Deployment {deployment_name} is UNHEALTHY: {getattr(dep, 'message', '')}"
)
if int(time.time() - start) % 10 == 0:
logging.info(
f"Waiting for {deployment_name} to be healthy, current: {dep_status}."
)
await asyncio.sleep(0.5)
raise RuntimeError(
f"Deployment {deployment_name} did not become HEALTHY after {timeout}s."
)
_BATCH_SIZE = 64
async def _controller_run_checkpoint(
handle: DeploymentHandle,
signal_actor,
checkpoint: int,
target_replicas: int,
marination_period_s: int,
sample_interval_s: int,
) -> List[Dict[str, Any]]:
"""Run a single checkpoint and collect metrics."""
start_time = time.time()
num_requests = int(target_replicas)
pending_requests: List[Any] = []
pending_requests.extend([handle.remote() for _ in range(num_requests)])
logging.info(f"Waiting for {num_requests} requests to be up...")
await _controller_wait_for_waiters(
signal_actor, len(pending_requests), timeout=_CONTROLLER_WAITER_TIMEOUT_S
)
logging.info(f"Waiting for {target_replicas} replicas to be up...")
# TODO: This is a hack to allow for some tolerance in the number of replicas.
# This is because the controller may not scale exactly to the target number of replicas.
# This is a bug in the controller autoscaling metrics aggregation logic, needs
# to be investigated further.
# This has the potential to introduce noise in the results from this benchmark.
replica_tolerance = 0.8
await _controller_wait_for_replicas_up(
int(target_replicas * replica_tolerance), timeout=_CONTROLLER_WAITER_TIMEOUT_S
)
logging.info(f"All {target_replicas} replicas are up.")
logging.info("Waiting for deployment to be healthy...")
await _controller_wait_for_deployment_healthy(timeout=_CONTROLLER_WAITER_TIMEOUT_S)
logging.info("Deployment is healthy.")
logging.info(f"Waiting for {marination_period_s} seconds to collect metrics...")
autoscale_duration_s = time.time() - start_time
samples = []
num_samples = marination_period_s // sample_interval_s
for sample_idx in range(num_samples):
health_metrics = await _controller_get_health_metrics()
actual_replicas = await _controller_get_replica_count()
num_nodes = _controller_get_active_nodes()
row = _controller_extract_metrics_row(
health_metrics=health_metrics,
checkpoint=checkpoint,
sample=sample_idx,
target_replicas=target_replicas,
actual_replicas=actual_replicas,
num_nodes=num_nodes,
autoscale_duration_s=autoscale_duration_s,
)
samples.append(row)
if sample_idx < num_samples - 1:
await asyncio.sleep(sample_interval_s)
await signal_actor.send.remote(clear=True)
try:
await asyncio.wait_for(
asyncio.gather(*pending_requests, return_exceptions=True),
timeout=30.0,
)
except asyncio.TimeoutError:
pass
return samples
async def run_controller_benchmark(
config: Optional[Dict[str, Any]] = None,
) -> List[Dict[str, Any]]:
"""
Run the controller health metrics benchmark and return raw samples.
Uses MetricsGenerator (autoscaling) -> HelloWorld (fixed) -> SignalActor
to stress the controller as replicas scale. Fails if get_health_metrics
is unavailable.
Args:
config: Optional benchmark config (checkpoints, marination_period_s,
sample_interval_s). Uses CONTROLLER_BENCH_CONFIG if None.
Returns:
List of sample dicts (one per marination sample). Each sample has
target_replicas, autoscale_duration_s, loop_duration_mean_s,
loops_per_second, event_loop_delay_s, num_asyncio_tasks, etc.
Caller converts to perf_metrics via convert_controller_samples_to_perf_metrics.
"""
cfg = config or CONTROLLER_BENCH_CONFIG
checkpoints = cfg["checkpoints"]
marination_period_s = cfg["marination_period_s"]
sample_interval_s = cfg["sample_interval_s"]
if not ray.is_initialized():
ray.init()
signal_actor = _SignalActorForController.remote()
all_samples: List[Dict[str, Any]] = []
try:
for checkpoint_idx, target_replicas in enumerate(checkpoints):
hello_world = ControllerBenchHelloWorld.bind(signal_actor)
app = ControllerBenchMetricsGenerator.bind(hello_world)
handle = serve.run(app, name="default", route_prefix=None)
samples = await _controller_run_checkpoint(
handle=handle,
signal_actor=signal_actor,
checkpoint=checkpoint_idx,
target_replicas=target_replicas,
marination_period_s=marination_period_s,
sample_interval_s=sample_interval_s,
)
all_samples.extend(samples)
serve.shutdown()
finally:
serve.shutdown()
return all_samples