wehub-resource-sync
124 lines
3.5 KiB
Python
124 lines
3.5 KiB
Python
import abc
|
|
import asyncio
|
|
import enum
|
|
import logging
|
|
import time
|
|
from typing import Tuple, Union
|
|
|
|
import numpy as np
|
|
|
|
from ray.actor import ActorHandle
|
|
from ray.runtime_env import RuntimeEnv
|
|
from ray.serve._private.benchmarks.common import Blackhole, run_throughput_benchmark
|
|
from ray.serve._private.benchmarks.serialization.common import PayloadPydantic
|
|
from ray.serve.handle import DeploymentHandle
|
|
|
|
GRPC_DEBUG_RUNTIME_ENV = RuntimeEnv(
|
|
env_vars={"GRPC_TRACE": "http", "GRPC_VERBOSITY": "debug"},
|
|
)
|
|
|
|
|
|
class IOMode(enum.Enum):
|
|
SYNC = "SYNC"
|
|
ASYNC = "ASYNC"
|
|
|
|
|
|
class Endpoint:
|
|
def __init__(self, tokens_per_request: int):
|
|
self._tokens_per_request = tokens_per_request
|
|
# Switch off logging to minimize its impact
|
|
logging.getLogger("ray").setLevel(logging.WARNING)
|
|
logging.getLogger("ray.serve").setLevel(logging.WARNING)
|
|
|
|
def stream(self):
|
|
payload = PayloadPydantic(
|
|
text="Test output",
|
|
floats=[float(f) for f in range(1, 100)],
|
|
ints=list(range(1, 100)),
|
|
ts=time.time(),
|
|
reason="Success!",
|
|
)
|
|
|
|
for i in range(self._tokens_per_request):
|
|
yield payload
|
|
|
|
async def aio_stream(self):
|
|
payload = PayloadPydantic(
|
|
text="Test output",
|
|
floats=[float(f) for f in range(1, 100)],
|
|
ints=list(range(1, 100)),
|
|
ts=time.time(),
|
|
reason="Success!",
|
|
)
|
|
|
|
for i in range(self._tokens_per_request):
|
|
yield payload
|
|
|
|
|
|
class Caller(Blackhole):
|
|
def __init__(
|
|
self,
|
|
downstream: Union[ActorHandle, DeploymentHandle],
|
|
*,
|
|
mode: IOMode,
|
|
tokens_per_request: int,
|
|
batch_size: int,
|
|
num_trials: int,
|
|
trial_runtime: float,
|
|
):
|
|
self._h = downstream
|
|
self._mode = mode
|
|
self._tokens_per_request = tokens_per_request
|
|
self._batch_size = batch_size
|
|
self._num_trials = num_trials
|
|
self._trial_runtime = trial_runtime
|
|
self._durations = []
|
|
|
|
# Switch off logging to minimize its impact
|
|
logging.getLogger("ray").setLevel(logging.WARNING)
|
|
logging.getLogger("ray.serve").setLevel(logging.WARNING)
|
|
|
|
def _get_remote_method(self):
|
|
if self._mode == IOMode.SYNC:
|
|
return self._h.stream
|
|
elif self._mode == IOMode.ASYNC:
|
|
return self._h.aio_stream
|
|
else:
|
|
raise NotImplementedError(f"Streaming mode not supported ({self._mode})")
|
|
|
|
@abc.abstractmethod
|
|
async def _consume_single_stream(self):
|
|
pass
|
|
|
|
async def _do_single_batch(self):
|
|
durations = await asyncio.gather(
|
|
*[
|
|
self._execute(self._consume_single_stream)
|
|
for _ in range(self._batch_size)
|
|
]
|
|
)
|
|
|
|
self._durations.extend(durations)
|
|
|
|
async def _execute(self, fn):
|
|
start = time.monotonic()
|
|
await fn()
|
|
dur_s = time.monotonic() - start
|
|
return dur_s * 1000 # ms
|
|
|
|
async def run_benchmark(self) -> Tuple[float, float]:
|
|
coro = run_throughput_benchmark(
|
|
fn=self._do_single_batch,
|
|
multiplier=self._batch_size * self._tokens_per_request,
|
|
num_trials=self._num_trials,
|
|
trial_runtime=self._trial_runtime,
|
|
)
|
|
# total_runtime = await collect_profile_events(coro)
|
|
total_runtime = await coro
|
|
|
|
p50, p75, p99 = np.percentile(self._durations, [50, 75, 99])
|
|
|
|
print(f"Individual request quantiles:\n\tP50={p50}\n\tP75={p75}\n\tP99={p99}")
|
|
|
|
return total_runtime
|