[ML] Inference API Rate limiter #106330

jonathan-buttner · 2024-03-13T19:26:45Z

This PR adds a RateLimiter class. It is currently unused but will be leveraged once the queuing and threading of the external services is refactored.

It implements the token bucket algorithm: https://en.wikipedia.org/wiki/Token_bucket

elasticsearchmachine · 2024-03-13T20:20:56Z

Pinging @elastic/ml-core (Team:ML)

davidkyle · 2024-03-18T12:31:31Z

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/RateLimiter.java

+            throw new IllegalArgumentException("Accumulated tokens limit must be greater than or equal to 0");
+        }
+
+        if (newAccumulatedTokensLimit == Double.POSITIVE_INFINITY) {


Suggested change

if (newAccumulatedTokensLimit == Double.POSITIVE_INFINITY) {

if (Double.isInfinite(newAccumulatedTokensLimit)) {

davidkyle · 2024-03-18T13:47:59Z

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/RateLimiter.java

+                return Double.POSITIVE_INFINITY;
+            }
+
+            return Double.NEGATIVE_INFINITY;


In the accumulateTokens code:

var newTokens = tokensPerNanos * elapsedTimeNanos; accumulatedTokens = Math.min(accumulatedTokensLimit, newTokens);

Math.min(a_positive_number, Double.NEGATIVE_INFINITY) returns Double.NEGATIVE_INFINITY. If accumulatedTokens becomes a -ve number I think that could cause errors.

One option is to return 0 ( not +ve or -ve infinity). Using ChronoUnit.MICRO or ChronoUnit.MILLIS reduces the chance of an arithmetic overflow

Good point. I'll switch it to 0 and use micros instead.

davidkyle · 2024-03-18T13:49:17Z

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/RateLimiter.java

+
+    private static double nanosBetweenExact(Instant start, Instant end) {
+        try {
+            return ChronoUnit.NANOS.between(start, end);


TemporalUnit.between() returns a long not double

davidkyle · 2024-03-18T13:51:18Z

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/RateLimiter.java

+    private void accumulateTokens() {
+        var now = Instant.now(clock);
+        if (now.isAfter(nextTokenAvailability)) {
+            var elapsedTimeNanos = nanosBetweenExact(nextTokenAvailability, now);


This is called in the ctor via setRate() at which point nextTokenAvailability == Instant.MIN. Because the calculated elapsedTimeNanos is high the class will be initialised with accumulatedTokens == accumulatedTokensLimit.

That seems reasonable to me, or at least as good as initialising accumulatedTokens to 0. Just want to check that is the intention

Yeah that was intentional. My thinking was that the first request can move forward without having to wait for tokens to accumulate if the limit was set to a positive number. If we always want it to start as 0 that's fine with me too though.

jonathan-buttner · 2024-03-18T15:14:12Z

@elasticmachine run elasticsearch-ci/part-1

jonathan-buttner · 2024-03-18T18:37:28Z

@elasticmachine merge upstream

davidkyle · 2024-03-18T21:18:37Z

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/RateLimiter.java

+    private void accumulateTokens() {
+        var now = Instant.now(clock);
+        if (now.isAfter(nextTokenAvailability)) {
+            var elapsedTimeNanos = microsBetweenExact(nextTokenAvailability, now);


Suggested change

var elapsedTimeNanos = microsBetweenExact(nextTokenAvailability, now);

var elapsedTimeMicros = microsBetweenExact(nextTokenAvailability, now);

Ugh sorry for all the missed nanos find replace 🤦‍♂️

davidkyle · 2024-03-18T21:18:46Z

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/RateLimiter.java

+        var now = Instant.now(clock);
+        if (now.isAfter(nextTokenAvailability)) {
+            var elapsedTimeNanos = microsBetweenExact(nextTokenAvailability, now);
+            var newTokens = tokensPerMicros * elapsedTimeNanos;


Suggested change

var newTokens = tokensPerMicros * elapsedTimeNanos;

var newTokens = tokensPerMicros * elapsedTimeMicros;

davidkyle · 2024-03-18T21:28:10Z

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/RateLimiter.java

+
+        accumulatedTokensLimit = newAccumulatedTokensLimit;
+
+        var unitsInNanos = newUnit.toMicros(1);


Suggested change

var unitsInNanos = newUnit.toMicros(1);

var unitsInMicros = newUnit.toMicros(1);

davidkyle · 2024-03-18T21:28:19Z

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/RateLimiter.java

+        accumulatedTokensLimit = newAccumulatedTokensLimit;
+
+        var unitsInNanos = newUnit.toMicros(1);
+        tokensPerMicros = newTokensPerTimeUnit / unitsInNanos;


Suggested change

tokensPerMicros = newTokensPerTimeUnit / unitsInNanos;

tokensPerMicros = newTokensPerTimeUnit / unitsInMicros;

davidkyle · 2024-03-18T21:31:17Z

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/RateLimiter.java

+        if (now.isAfter(nextTokenAvailability)) {
+            var elapsedTimeNanos = microsBetweenExact(nextTokenAvailability, now);
+            var newTokens = tokensPerMicros * elapsedTimeNanos;
+            accumulatedTokens = Math.min(accumulatedTokensLimit, newTokens);


Should this include the previously accumulated tokens?

Suggested change

accumulatedTokens = Math.min(accumulatedTokensLimit, newTokens);

accumulatedTokens = Math.min(accumulatedTokensLimit, accumulatedTokens + newTokens);

Yep thanks for that.

…search into ml-token-bucket

jonathan-buttner · 2024-03-19T13:22:05Z

@elasticmachine merge upstream

jonathan-buttner · 2024-03-19T13:53:51Z

@elasticmachine run elasticsearch-ci/part-3

davidkyle

LGTM

jonathan-buttner · 2024-03-19T19:02:20Z

@elasticmachine merge upstream

jonathan-buttner · 2024-03-19T19:33:59Z

@elasticmachine merge upstream

maxhniebergall · 2024-03-28T15:47:09Z

This PR might be relevant to this issue #106877

jonathan-buttner added 3 commits March 12, 2024 17:44

Working tests

9f13884

Adding more tests

eb0b086

Adding comment

f15bd38

jonathan-buttner added >non-issue :ml Machine learning Team:ML Meta label for the ML team v8.14.0 labels Mar 13, 2024

jonathan-buttner marked this pull request as ready for review March 13, 2024 20:20

jonathan-buttner requested review from davidkyle and maxhniebergall March 13, 2024 20:20

davidkyle reviewed Mar 18, 2024

View reviewed changes

Switching to micros and addressing feedback

8cc99c9

Merge branch 'main' into ml-token-bucket

830e586

davidkyle reviewed Mar 18, 2024

View reviewed changes

jonathan-buttner added 2 commits March 19, 2024 08:56

Removing nanos and adding test for bug fix

0496a89

Merge branch 'ml-token-bucket' of github.com:jonathan-buttner/elastic…

dc975a8

…search into ml-token-bucket

Merge branch 'main' into ml-token-bucket

4a65523

davidkyle approved these changes Mar 19, 2024

View reviewed changes

Merge branch 'main' into ml-token-bucket

9ba1a0d

Merge branch 'main' into ml-token-bucket

9bc0e12

jonathan-buttner merged commit edbff94 into elastic:main Mar 19, 2024
14 checks passed

jonathan-buttner deleted the ml-token-bucket branch March 19, 2024 21:44

lkts mentioned this pull request Mar 21, 2024

Add DownsampleMetrics #106637

Closed

This was referenced Mar 25, 2024

Set index mode earlier for new downsample index #106728

Merged

Added initial metrics for synthetic source #106732

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ML] Inference API Rate limiter #106330

[ML] Inference API Rate limiter #106330

jonathan-buttner commented Mar 13, 2024

elasticsearchmachine commented Mar 13, 2024

davidkyle Mar 18, 2024

davidkyle Mar 18, 2024

jonathan-buttner Mar 18, 2024

davidkyle Mar 18, 2024

davidkyle Mar 18, 2024

jonathan-buttner Mar 18, 2024 •

edited

Loading

jonathan-buttner commented Mar 18, 2024

jonathan-buttner commented Mar 18, 2024

davidkyle Mar 18, 2024

jonathan-buttner Mar 19, 2024

davidkyle Mar 18, 2024

davidkyle Mar 18, 2024

davidkyle Mar 18, 2024

davidkyle Mar 18, 2024

jonathan-buttner Mar 19, 2024

jonathan-buttner commented Mar 19, 2024

jonathan-buttner commented Mar 19, 2024

davidkyle left a comment

jonathan-buttner commented Mar 19, 2024

jonathan-buttner commented Mar 19, 2024

maxhniebergall commented Mar 28, 2024

	if (newAccumulatedTokensLimit == Double.POSITIVE_INFINITY) {
	if (Double.isInfinite(newAccumulatedTokensLimit)) {

	var elapsedTimeNanos = microsBetweenExact(nextTokenAvailability, now);
	var elapsedTimeMicros = microsBetweenExact(nextTokenAvailability, now);

	var newTokens = tokensPerMicros * elapsedTimeNanos;
	var newTokens = tokensPerMicros * elapsedTimeMicros;


		accumulatedTokensLimit = newAccumulatedTokensLimit;

		var unitsInNanos = newUnit.toMicros(1);

	var unitsInNanos = newUnit.toMicros(1);
	var unitsInMicros = newUnit.toMicros(1);

	tokensPerMicros = newTokensPerTimeUnit / unitsInNanos;
	tokensPerMicros = newTokensPerTimeUnit / unitsInMicros;

	accumulatedTokens = Math.min(accumulatedTokensLimit, newTokens);
	accumulatedTokens = Math.min(accumulatedTokensLimit, accumulatedTokens + newTokens);

[ML] Inference API Rate limiter #106330

[ML] Inference API Rate limiter #106330

Conversation

jonathan-buttner commented Mar 13, 2024

elasticsearchmachine commented Mar 13, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jonathan-buttner Mar 18, 2024 • edited Loading

Choose a reason for hiding this comment

jonathan-buttner commented Mar 18, 2024

jonathan-buttner commented Mar 18, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jonathan-buttner commented Mar 19, 2024

jonathan-buttner commented Mar 19, 2024

davidkyle left a comment

Choose a reason for hiding this comment

jonathan-buttner commented Mar 19, 2024

jonathan-buttner commented Mar 19, 2024

maxhniebergall commented Mar 28, 2024

jonathan-buttner Mar 18, 2024 •

edited

Loading