chore: import upstream snapshot with attribution

python/ray/serve/_private/benchmarks/streaming/common.py

@@ -0,0 +1,123 @@
+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