From 33fa1bbece00ea40aa1348cd18adc14be9e443c4 Mon Sep 17 00:00:00 2001 From: Joongi Kim Date: Thu, 5 Dec 2024 01:13:14 +0900 Subject: [PATCH 1/3] fix: Optimize weighted random choice in route selection --- src/ai/backend/wsproxy/proxy/backend/http.py | 18 ++---------------- src/ai/backend/wsproxy/proxy/backend/tcp.py | 18 ++---------------- 2 files changed, 4 insertions(+), 32 deletions(-) diff --git a/src/ai/backend/wsproxy/proxy/backend/http.py b/src/ai/backend/wsproxy/proxy/backend/http.py index e02710af67..f4a18504cf 100644 --- a/src/ai/backend/wsproxy/proxy/backend/http.py +++ b/src/ai/backend/wsproxy/proxy/backend/http.py @@ -34,22 +34,8 @@ def selected_route(self) -> RouteInfo: if selected_route.traffic_ratio == 0: raise WorkerNotAvailable else: - routes = [ - r for r in sorted(self.routes, key=lambda r: r.traffic_ratio) if r.traffic_ratio > 0 - ] - ranges: list[float] = [] - ratio_sum = 0.0 - for route in routes: - ratio_sum += route.traffic_ratio - ranges.append(ratio_sum) - rand = random.random() * ranges[-1] - for i in range(len(ranges)): - ceiling = ranges[i] - if (i == 0 and rand < ceiling) or (ranges[i - 1] <= rand and rand < ceiling): - selected_route = routes[i] - break - else: - selected_route = routes[-1] + ratios: list[float] = [r.traffic_ratio for r in self.routes] + selected_route = random.choices(self.routes, weights=ratios, k=1)[0] return selected_route def get_x_forwarded_proto(self, request: web.Request) -> str: diff --git a/src/ai/backend/wsproxy/proxy/backend/tcp.py b/src/ai/backend/wsproxy/proxy/backend/tcp.py index d8f8ec95c2..701e9f3de7 100644 --- a/src/ai/backend/wsproxy/proxy/backend/tcp.py +++ b/src/ai/backend/wsproxy/proxy/backend/tcp.py @@ -31,22 +31,8 @@ def selected_route(self) -> RouteInfo: if selected_route.traffic_ratio == 0: raise WorkerNotAvailable else: - routes = [ - r for r in sorted(self.routes, key=lambda r: r.traffic_ratio) if r.traffic_ratio > 0 - ] - ranges: list[float] = [] - ratio_sum = 0.0 - for route in routes: - ratio_sum += route.traffic_ratio - ranges.append(ratio_sum) - rand = random.random() * ranges[-1] - for i in range(len(ranges)): - ceiling = ranges[i] - if (i == 0 and rand < ceiling) or (ranges[i - 1] <= rand and rand < ceiling): - selected_route = routes[i] - break - else: - selected_route = routes[-1] + ratios: list[float] = [r.traffic_ratio for r in self.routes] + selected_route = random.choices(self.routes, weights=ratios, k=1)[0] return selected_route async def bind( From 8d5f585cc88d8d7be7c0d2ff8b0266a0d6c80c5d Mon Sep 17 00:00:00 2001 From: Joongi Kim Date: Thu, 5 Dec 2024 01:20:12 +0900 Subject: [PATCH 2/3] doc: Add news fragment --- changes/.fix.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/.fix.md diff --git a/changes/.fix.md b/changes/.fix.md new file mode 100644 index 0000000000..d7036015ab --- /dev/null +++ b/changes/.fix.md @@ -0,0 +1 @@ +Optimize the route selection in App Proxy using `random.choices()` based on the native C implementation in CPython From 4d22a1dd9c6110b4d935efe9dec18a9a0885f5d5 Mon Sep 17 00:00:00 2001 From: Joongi Kim Date: Thu, 5 Dec 2024 01:22:16 +0900 Subject: [PATCH 3/3] doc: Rename news fragment manually --- changes/{.fix.md => 3199.fix.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename changes/{.fix.md => 3199.fix.md} (100%) diff --git a/changes/.fix.md b/changes/3199.fix.md similarity index 100% rename from changes/.fix.md rename to changes/3199.fix.md