[Inference API] add service and task type aware rate limiting

Author: brendan-jugan-elastic

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -208,6 +208,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion PROJECT_ID_IN_SNAPSHOT = def(9_040_0_00);
     public static final TransportVersion INDEX_STATS_AND_METADATA_INCLUDE_PEAK_WRITE_LOAD = def(9_041_0_00);
     public static final TransportVersion REPOSITORIES_METADATA_AS_PROJECT_CUSTOM = def(9_042_0_00);
+    public static final TransportVersion INFERENCE_REQUEST_SERVICE_TASK_TYPE_RATE_LIMITING = def(9_043_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
 import org.elasticsearch.xpack.inference.external.http.sender.AmazonBedrockRequestExecutorService;
@@ -89,8 +90,8 @@ protected AmazonBedrockRequestSender(
     }
 
     @Override
-    public void updateRateLimitDivisor(int rateLimitDivisor) {
-        executorService.updateRateLimitDivisor(rateLimitDivisor);
+    public void updateRateLimitDivisor(String serviceName, TaskType taskType, int rateLimitDivisor) {
+        executorService.updateRateLimitDivisor(serviceName, taskType, rateLimitDivisor);
     }
 
     @Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/amazonbedrock/AmazonBedrockRequestSender.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.RequestManager;
 
@@ -21,7 +22,7 @@ public interface RequestExecutor {
 
     void shutdown();
 
-    void updateRateLimitDivisor(int newDivisor);
+    void updateRateLimitDivisor(String serviceName, TaskType taskType, int newDivisor);
 
     boolean isShutdown();
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/RequestExecutor.java

@@ -15,7 +15,14 @@
 abstract class AlibabaCloudSearchRequestManager extends BaseRequestManager {
 
     protected AlibabaCloudSearchRequestManager(ThreadPool threadPool, AlibabaCloudSearchModel model) {
-        super(threadPool, model.getInferenceEntityId(), RateLimitGrouping.of(model), model.rateLimitServiceSettings().rateLimitSettings());
+        super(
+            threadPool,
+            model.getInferenceEntityId(),
+            AlibabaCloudSearchRequestManager.RateLimitGrouping.of(model),
+            model.rateLimitServiceSettings().rateLimitSettings(),
+            model.getConfigurations().getService(),
+            model.getConfigurations().getTaskType()
+        );
     }
 
     record RateLimitGrouping(int apiKeyHash) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/AlibabaCloudSearchRequestManager.java

@@ -15,13 +15,21 @@
 
 import java.util.Objects;
 
-public abstract class AmazonBedrockRequestManager implements RequestManager {
+public abstract class AmazonBedrockRequestManager extends BaseRequestManager {
 
     protected final ThreadPool threadPool;
     protected final TimeValue timeout;
     private final AmazonBedrockModel baseModel;
 
     protected AmazonBedrockRequestManager(AmazonBedrockModel baseModel, ThreadPool threadPool, @Nullable TimeValue timeout) {
+        super(
+            threadPool,
+            baseModel.getInferenceEntityId(),
+            AmazonBedrockRequestManager.RateLimitGrouping.of(baseModel),
+            baseModel.rateLimitSettings(),
+            baseModel.getConfigurations().getService(),
+            baseModel.getConfigurations().getTaskType()
+        );
         this.baseModel = Objects.requireNonNull(baseModel);
         this.threadPool = Objects.requireNonNull(threadPool);
         this.timeout = timeout;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/AmazonBedrockRequestManager.java

@@ -16,7 +16,14 @@
 abstract class AnthropicRequestManager extends BaseRequestManager {
 
     protected AnthropicRequestManager(ThreadPool threadPool, AnthropicModel model) {
-        super(threadPool, model.getInferenceEntityId(), RateLimitGrouping.of(model), model.rateLimitServiceSettings().rateLimitSettings());
+        super(
+            threadPool,
+            model.getInferenceEntityId(),
+            AnthropicRequestManager.RateLimitGrouping.of(model),
+            model.rateLimitServiceSettings().rateLimitSettings(),
+            model.getConfigurations().getService(),
+            model.getConfigurations().getTaskType()
+        );
     }
 
     record RateLimitGrouping(int accountHash, int modelIdHash) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/AnthropicRequestManager.java

@@ -15,7 +15,14 @@
 public abstract class AzureAiStudioRequestManager extends BaseRequestManager {
 
     protected AzureAiStudioRequestManager(ThreadPool threadPool, AzureAiStudioModel model) {
-        super(threadPool, model.getInferenceEntityId(), AzureAiStudioRequestManager.RateLimitGrouping.of(model), model.rateLimitSettings());
+        super(
+            threadPool,
+            model.getInferenceEntityId(),
+            AzureAiStudioRequestManager.RateLimitGrouping.of(model),
+            model.rateLimitSettings(),
+            model.getConfigurations().getService(),
+            model.getConfigurations().getTaskType()
+        );
     }
 
     record RateLimitGrouping(int targetHashcode) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/AzureAiStudioRequestManager.java

@@ -14,7 +14,14 @@
 
 public abstract class AzureOpenAiRequestManager extends BaseRequestManager {
     protected AzureOpenAiRequestManager(ThreadPool threadPool, AzureOpenAiModel model) {
-        super(threadPool, model.getInferenceEntityId(), RateLimitGrouping.of(model), model.rateLimitServiceSettings().rateLimitSettings());
+        super(
+            threadPool,
+            model.getInferenceEntityId(),
+            RateLimitGrouping.of(model),
+            model.rateLimitServiceSettings().rateLimitSettings(),
+            model.getConfigurations().getService(),
+            model.getConfigurations().getTaskType()
+        );
     }
 
     record RateLimitGrouping(int resourceNameHash, int deploymentIdHash) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/AzureOpenAiRequestManager.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.inference.external.http.sender;
 
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
@@ -23,17 +24,28 @@ abstract class BaseRequestManager implements RequestManager {
     // the rate and the other inference endpoint's rate will be ignored
     private final EndpointGrouping endpointGrouping;
     private final RateLimitSettings rateLimitSettings;
+    private final String service;
+    private final TaskType taskType;
 
-    BaseRequestManager(ThreadPool threadPool, String inferenceEntityId, Object rateLimitGroup, RateLimitSettings rateLimitSettings) {
+    BaseRequestManager(
+        ThreadPool threadPool,
+        String inferenceEntityId,
+        Object rateLimitGroup,
+        RateLimitSettings rateLimitSettings,
+        String service,
+        TaskType taskType
+    ) {
         this.threadPool = Objects.requireNonNull(threadPool);
         this.inferenceEntityId = Objects.requireNonNull(inferenceEntityId);
 
         Objects.requireNonNull(rateLimitSettings);
         this.endpointGrouping = new EndpointGrouping(Objects.requireNonNull(rateLimitGroup).hashCode(), rateLimitSettings);
         this.rateLimitSettings = rateLimitSettings;
+        this.service = service;
+        this.taskType = taskType;
     }
 
-    BaseRequestManager(ThreadPool threadPool, RateLimitGroupingModel rateLimitGroupingModel) {
+    BaseRequestManager(ThreadPool threadPool, RateLimitGroupingModel rateLimitGroupingModel, String service, TaskType taskType) {
         this.threadPool = Objects.requireNonNull(threadPool);
         Objects.requireNonNull(rateLimitGroupingModel);
 
@@ -43,6 +55,8 @@ abstract class BaseRequestManager implements RequestManager {
             rateLimitGroupingModel.rateLimitSettings()
         );
         this.rateLimitSettings = rateLimitGroupingModel.rateLimitSettings();
+        this.service = service;
+        this.taskType = taskType;
     }
 
     protected void execute(Runnable runnable) {
@@ -64,5 +78,15 @@ public RateLimitSettings rateLimitSettings() {
         return rateLimitSettings;
     }
 
+    @Override
+    public String service() {
+        return this.service;
+    }
+
+    @Override
+    public TaskType taskType() {
+        return this.taskType;
+    }
+
     private record EndpointGrouping(int group, RateLimitSettings settings) {}
 }