chore: import upstream snapshot with attribution

python/ray/serve/_private/benchmarks/microbenchmark.py

@@ -0,0 +1,182 @@
+# Runs several scenarios with varying max batch size, max concurrent queries,
+# number of replicas, and with intermediate serve handles (to simulate ensemble
+# models) either on or off.
+
+import asyncio
+import logging
+from pprint import pprint
+from typing import Dict, Union
+
+import aiohttp
+from starlette.requests import Request
+
+import ray
+from ray import serve
+from ray.serve._private.benchmarks.common import run_throughput_benchmark
+from ray.serve.handle import DeploymentHandle
+
+NUM_CLIENTS = 8
+CALLS_PER_BATCH = 100
+
+
+async def fetch(session, data):
+    async with session.get("http://localhost:8000/", data=data) as response:
+        response = await response.text()
+        assert response == "ok", response
+
+
+@ray.remote
+class Client:
+    def ready(self):
+        return "ok"
+
+    async def do_queries(self, num, data):
+        async with aiohttp.ClientSession() as session:
+            for _ in range(num):
+                await fetch(session, data)
+
+
+def build_app(
+    intermediate_handles: bool,
+    num_replicas: int,
+    max_batch_size: int,
+    max_ongoing_requests: int,
+):
+    @serve.deployment(max_ongoing_requests=1000)
+    class Upstream:
+        def __init__(self, handle: DeploymentHandle):
+            self._handle = handle
+
+            # Turn off access log.
+            logging.getLogger("ray.serve").setLevel(logging.WARNING)
+
+        async def __call__(self, req: Request):
+            return await self._handle.remote(await req.body())
+
+    @serve.deployment(
+        num_replicas=num_replicas,
+        max_ongoing_requests=max_ongoing_requests,
+    )
+    class Downstream:
+        def __init__(self):
+            # Turn off access log.
+            logging.getLogger("ray.serve").setLevel(logging.WARNING)
+
+        @serve.batch(max_batch_size=max_batch_size)
+        async def batch(self, reqs):
+            return [b"ok"] * len(reqs)
+
+        async def __call__(self, req: Union[bytes, Request]):
+            if max_batch_size > 1:
+                return await self.batch(req)
+            else:
+                return b"ok"
+
+    if intermediate_handles:
+        return Upstream.bind(Downstream.bind())
+    else:
+        return Downstream.bind()
+
+
+async def trial(
+    intermediate_handles: bool,
+    num_replicas: int,
+    max_batch_size: int,
+    max_ongoing_requests: int,
+    data_size: str,
+) -> Dict[str, float]:
+    results = {}
+
+    trial_key_base = (
+        f"replica:{num_replicas}/batch_size:{max_batch_size}/"
+        f"concurrent_queries:{max_ongoing_requests}/"
+        f"data_size:{data_size}/intermediate_handle:{intermediate_handles}"
+    )
+
+    print(
+        f"intermediate_handles={intermediate_handles},"
+        f"num_replicas={num_replicas},"
+        f"max_batch_size={max_batch_size},"
+        f"max_ongoing_requests={max_ongoing_requests},"
+        f"data_size={data_size}"
+    )
+
+    app = build_app(
+        intermediate_handles, num_replicas, max_batch_size, max_ongoing_requests
+    )
+    serve.run(app)
+
+    if data_size == "small":
+        data = None
+    elif data_size == "large":
+        data = b"a" * 1024 * 1024
+    else:
+        raise ValueError("data_size should be 'small' or 'large'.")
+
+    async with aiohttp.ClientSession() as session:
+
+        async def single_client():
+            for _ in range(CALLS_PER_BATCH):
+                await fetch(session, data)
+
+        single_client_avg_tps, single_client_std_tps = await run_throughput_benchmark(
+            single_client,
+            multiplier=CALLS_PER_BATCH,
+        )
+        print(
+            "\t{} {} +- {} requests/s".format(
+                "single client {} data".format(data_size),
+                single_client_avg_tps,
+                single_client_std_tps,
+            )
+        )
+        key = f"num_client:1/{trial_key_base}"
+        results[key] = single_client_avg_tps
+
+    clients = [Client.remote() for _ in range(NUM_CLIENTS)]
+    ray.get([client.ready.remote() for client in clients])
+
+    async def many_clients():
+        ray.get([a.do_queries.remote(CALLS_PER_BATCH, data) for a in clients])
+
+    multi_client_avg_tps, _ = await run_throughput_benchmark(
+        many_clients,
+        multiplier=CALLS_PER_BATCH * len(clients),
+    )
+
+    results[f"num_client:{len(clients)}/{trial_key_base}"] = multi_client_avg_tps
+    return results
+
+
+async def main():
+    results = {}
+    for intermediate_handles in [False, True]:
+        for num_replicas in [1, 8]:
+            for max_batch_size, max_ongoing_requests in [
+                (1, 1),
+                (1, 10000),
+                (10000, 10000),
+            ]:
+                # TODO(edoakes): large data causes broken pipe errors.
+                for data_size in ["small"]:
+                    results.update(
+                        await trial(
+                            intermediate_handles,
+                            num_replicas,
+                            max_batch_size,
+                            max_ongoing_requests,
+                            data_size,
+                        )
+                    )
+
+    print("Results from all conditions:")
+    pprint(results)
+    return results
+
+
+if __name__ == "__main__":
+    ray.init()
+    serve.start()
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(main())