Change rate limiter token capacity setting (opensearch-project#1635)

Signed-off-by: Sicheng Song <[email protected]>
b4sjoo · Nov 17, 2023 · a8629cf · a8629cf
1 parent 79335ec
commit a8629cf
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java
@@ -1112,12 +1112,13 @@ private TokenBucket rateLimiterConstructor(Integer eligibleNodeCount, MLModel ml
             TimeUnit rateLimitUnit = mlModel.getRateLimitUnit();
             log
                 .debug(
-                    "Initializing the rate limiter for Model {}, with TPS limit {}, evenly distributed on {} nodes",
+                    "Initializing the rate limiter for Model {}, with TPS limit {} and burst capacity {}, evenly distributed on {} nodes",
                     mlModel.getModelId(),
                     rateLimitNumber / rateLimitUnit.toSeconds(1),
+                    rateLimitNumber,
                     eligibleNodeCount
                 );
-            return new TokenBucket(System::nanoTime, rateLimitNumber / rateLimitUnit.toNanos(1) / eligibleNodeCount, 2);
+            return new TokenBucket(System::nanoTime, rateLimitNumber / rateLimitUnit.toNanos(1) / eligibleNodeCount, rateLimitNumber);
         }
         return null;
     }