fix: RandomizedWeightedBalancer chooses the first upstream with highe…

…r probability than others #666
epam · Jan 31, 2025 · f7a330b · f7a330b
1 parent 0b32f62
commit f7a330b
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 7 deletions.
diff --git a/server/src/main/java/com/epam/aidial/core/server/upstream/RandomizedWeightedBalancer.java b/server/src/main/java/com/epam/aidial/core/server/upstream/RandomizedWeightedBalancer.java
@@ -50,9 +50,16 @@ public Upstream next() {
         if (availableUpstreams.isEmpty()) {
             return null;
         }
+        if (availableUpstreams.size() == 1) {
+            return availableUpstreams.get(0);
+        }
         int total = availableUpstreams.stream().map(Upstream::getWeight).reduce(0, Integer::sum);
         // make sure the upper bound `total` is inclusive
-        int random = generator.nextInt(total + 1);
+        // the lowest bound should be 1 otherwise the 1st upstream with weight 1 has more possibility
+        // to be selected because zero is included in its range
+        // e.g. there are 3 upstreams with equal weight = 1 and the ranges are [0,1], [2,2] and [3,3]
+        // definitely the 1st upstream has higher possibility
+        int random = generator.nextInt(1, total + 1);
         int current = 0;
 
         Upstream result = null;

diff --git a/...er/src/test/java/com/epam/aidial/core/server/upstream/RandomizedWeightedBalancerTest.java b/...er/src/test/java/com/epam/aidial/core/server/upstream/RandomizedWeightedBalancerTest.java
@@ -19,7 +19,7 @@ public class RandomizedWeightedBalancerTest {
 
     @Mock
     private Random generator;
-    
+
     @Test
     void testWeightedLoadBalancer() {
         List<Upstream> upstreams = List.of(
@@ -31,25 +31,25 @@ void testWeightedLoadBalancer() {
 
         RandomizedWeightedBalancer balancer = new RandomizedWeightedBalancer("model1", upstreams, generator);
 
-        when(generator.nextInt(11)).thenReturn(0);
+        when(generator.nextInt(1, 11)).thenReturn(0);
 
         Upstream upstream = balancer.next();
         assertNotNull(upstream);
         assertEquals(upstreams.get(0), upstream);
 
-        when(generator.nextInt(11)).thenReturn(2);
+        when(generator.nextInt(1, 11)).thenReturn(2);
 
         upstream = balancer.next();
         assertNotNull(upstream);
         assertEquals(upstreams.get(1), upstream);
 
-        when(generator.nextInt(11)).thenReturn(6);
+        when(generator.nextInt(1, 11)).thenReturn(6);
 
         upstream = balancer.next();
         assertNotNull(upstream);
         assertEquals(upstreams.get(2), upstream);
 
-        when(generator.nextInt(11)).thenReturn(10);
+        when(generator.nextInt(1, 11)).thenReturn(10);
 
         upstream = balancer.next();
         assertNotNull(upstream);

diff --git a/server/src/test/java/com/epam/aidial/core/server/upstream/TieredBalancerTest.java b/server/src/test/java/com/epam/aidial/core/server/upstream/TieredBalancerTest.java
@@ -111,7 +111,7 @@ void testUpstreamFallback() {
                 .toList();
         model.setUpstreams(upstreams);
         AtomicInteger counter = new AtomicInteger();
-        when(generator.nextInt(5)).thenAnswer(cb -> counter.incrementAndGet());
+        when(generator.nextInt(1, 5)).thenAnswer(cb -> counter.incrementAndGet());
         Supplier<Random> factory = () -> generator;
 
         UpstreamRouteProvider upstreamRouteProvider = new UpstreamRouteProvider(vertx, factory);