chore: import upstream snapshot with attribution

doc/source/serve/doc_code/custom_request_router_app.py

@@ -0,0 +1,164 @@
+# flake8: noqa
+
+# __begin_deploy_app_with_uniform_request_router__
+from ray import serve
+from ray.serve.request_router import ReplicaID
+import time
+from collections import defaultdict
+from ray.serve.context import _get_internal_replica_context
+from typing import Any, Dict
+from ray.serve.config import RequestRouterConfig
+
+
+@serve.deployment(
+    request_router_config=RequestRouterConfig(
+        request_router_class="custom_request_router:UniformRequestRouter",
+    ),
+    num_replicas=10,
+    ray_actor_options={"num_cpus": 0},
+)
+class UniformRequestRouterApp:
+    def __init__(self):
+        context = _get_internal_replica_context()
+        self.replica_id: ReplicaID = context.replica_id
+
+    async def __call__(self):
+        return self.replica_id
+
+
+handle = serve.run(UniformRequestRouterApp.bind())
+response = handle.remote().result()
+print(f"Response from UniformRequestRouterApp: {response}")
+# Example output:
+#   Response from UniformRequestRouterApp:
+#   Replica(id='67vc4ts5', deployment='UniformRequestRouterApp', app='default')
+# __end_deploy_app_with_uniform_request_router__
+
+
+# __begin_deploy_app_with_throughput_aware_request_router__
+def _time_ms() -> int:
+    return int(time.time() * 1000)
+
+
+@serve.deployment(
+    request_router_config=RequestRouterConfig(
+        request_router_class="custom_request_router:ThroughputAwareRequestRouter",
+        request_routing_stats_period_s=1,
+        request_routing_stats_timeout_s=1,
+    ),
+    num_replicas=3,
+    ray_actor_options={"num_cpus": 0},
+)
+class ThroughputAwareRequestRouterApp:
+    def __init__(self):
+        self.throughput_buckets: Dict[int, int] = defaultdict(int)
+        self.last_throughput_buckets = _time_ms()
+        context = _get_internal_replica_context()
+        self.replica_id: ReplicaID = context.replica_id
+
+    def __call__(self):
+        self.update_throughput()
+        return self.replica_id
+
+    def update_throughput(self):
+        current_timestamp_ms = _time_ms()
+
+        # Under high concurrency, requests can come in at different times. This
+        # early return helps to skip if the current_timestamp_ms is more than a
+        # second older than the last throughput bucket.
+        if current_timestamp_ms < self.last_throughput_buckets - 1000:
+            return
+
+        # Record the request to the bucket
+        self.throughput_buckets[current_timestamp_ms] += 1
+        self.last_throughput_buckets = current_timestamp_ms
+
+    def record_routing_stats(self) -> Dict[str, Any]:
+        current_timestamp_ms = _time_ms()
+        throughput = 0
+
+        for t, c in list(self.throughput_buckets.items()):
+            if t < current_timestamp_ms - 1000:
+                # Remove the bucket if it is older than 1 second
+                self.throughput_buckets.pop(t)
+            else:
+                throughput += c
+
+        return {
+            "throughput": throughput,
+        }
+
+
+handle = serve.run(ThroughputAwareRequestRouterApp.bind())
+response = handle.remote().result()
+print(f"Response from ThroughputAwareRequestRouterApp: {response}")
+# Example output:
+#   Response from ThroughputAwareRequestRouterApp:
+#   Replica(id='tkywafya', deployment='ThroughputAwareRequestRouterApp', app='default')
+# __end_deploy_app_with_throughput_aware_request_router__
+
+
+# __begin_deploy_app_with_round_robin_router__
+@serve.deployment(
+    request_router_config=RequestRouterConfig(
+        request_router_class=(
+            "ray.serve.experimental.round_robin_router:RoundRobinRouter"
+        ),
+    ),
+    num_replicas=4,
+    ray_actor_options={"num_cpus": 0},
+)
+class RoundRobinRouterApp:
+    def __init__(self):
+        context = _get_internal_replica_context()
+        self.replica_id: ReplicaID = context.replica_id
+
+    async def __call__(self):
+        return self.replica_id
+
+
+handle = serve.run(RoundRobinRouterApp.bind())
+response = handle.remote().result()
+print(f"Response from RoundRobinRouterApp: {response}")
+# __end_deploy_app_with_round_robin_router__
+
+
+# __begin_deploy_app_with_consistent_hash_router__
+import requests
+from starlette.requests import Request
+
+@serve.deployment(
+    request_router_config=RequestRouterConfig(
+        request_router_class=(
+            "ray.serve.experimental.consistent_hash_router:ConsistentHashRouter"
+        ),
+        request_router_kwargs={
+            "num_virtual_nodes": 100,
+            "num_fallback_replicas": 2,
+        },
+    ),
+    num_replicas=4,
+    ray_actor_options={"num_cpus": 0},
+)
+class ConsistentHashRouterApp:
+    def __init__(self):
+        context = _get_internal_replica_context()
+        self.replica_id: ReplicaID = context.replica_id
+
+    async def __call__(self, request: Request) -> str:
+        return str(self.replica_id)
+
+
+serve.run(ConsistentHashRouterApp.bind())
+
+# Clients pin a session to a replica by sending the same `x-session-id`
+# on every request.
+response = requests.get(
+    "http://localhost:8000/",
+    headers={"x-session-id": "example-session-id"},
+)
+print(f"Response from ConsistentHashRouterApp: {response.text}")
+# Example output:
+#   Response from ConsistentHashRouterApp:
+#   Replica(id='...', deployment='ConsistentHashRouterApp', app='default')
+# __end_deploy_app_with_consistent_hash_router__