chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,688 @@
|
||||
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
|
||||
Reference in New Issue
Block a user