Semantic Text Chunking Indexing Pressure

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

@@ -8,11 +8,24 @@
 package org.elasticsearch.xpack.inference;
 
 import org.elasticsearch.ElasticsearchException;
-import org.elasticsearch.ElasticsearchWrapperException;
+import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.rest.RestStatus;
 
-public class InferenceException extends ElasticsearchException implements ElasticsearchWrapperException {
+public class InferenceException extends ElasticsearchException {
     public InferenceException(String message, Throwable cause, Object... args) {
         super(message, cause, args);
     }
 
+    @Override
+    public RestStatus status() {
+        // Override status so that we get the status of the cause while retaining the message of the inference exception when emitting to
+        // XContent
+        RestStatus status = RestStatus.INTERNAL_SERVER_ERROR;
+        Throwable cause = getCause();
+        if (cause != null) {
+            status = ExceptionsHelper.status(cause);
+        }
+
+        return status;
+    }
 }

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

@@ -17,25 +17,29 @@
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
+import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.IndexScopedSettings;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.settings.SettingsFilter;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.index.mapper.Mapper;
 import org.elasticsearch.index.mapper.MetadataFieldMapper;
 import org.elasticsearch.indices.SystemIndexDescriptor;
+import org.elasticsearch.indices.breaker.BreakerSettings;
 import org.elasticsearch.inference.InferenceServiceExtension;
 import org.elasticsearch.inference.InferenceServiceRegistry;
 import org.elasticsearch.license.License;
 import org.elasticsearch.license.LicensedFeature;
 import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.node.PluginComponentBinding;
 import org.elasticsearch.plugins.ActionPlugin;
+import org.elasticsearch.plugins.CircuitBreakerPlugin;
 import org.elasticsearch.plugins.ClusterPlugin;
 import org.elasticsearch.plugins.ExtensiblePlugin;
 import org.elasticsearch.plugins.MapperPlugin;
@@ -157,7 +161,8 @@ public class InferencePlugin extends Plugin
         MapperPlugin,
         SearchPlugin,
         InternalSearchPlugin,
-        ClusterPlugin {
+        ClusterPlugin,
+        CircuitBreakerPlugin {
 
     /**
      * When this setting is true the verification check that
@@ -175,6 +180,15 @@ public class InferencePlugin extends Plugin
         Setting.Property.Dynamic
     );
 
+    public static final String INFERENCE_BYTES_CIRCUIT_BREAKER_NAME = "inference.bytes";
+    public static final double INFERENCE_BYTES_CIRCUIT_BREAKER_OVERHEAD = 1.0;
+    public static final Setting<ByteSizeValue> INFERENCE_BYTES_CIRCUIT_BREAKER_LIMIT_SETTING = Setting.memorySizeSetting(
+        "indices.breaker.inference.bytes.limit",
+        "25%",
+        Setting.Property.Dynamic,
+        Setting.Property.NodeScope
+    );
+
     public static final LicensedFeature.Momentary INFERENCE_API_FEATURE = LicensedFeature.momentary(
         "inference",
         "api",
@@ -199,6 +213,7 @@ public class InferencePlugin extends Plugin
     private final SetOnce<ElasticInferenceServiceComponents> elasticInferenceServiceComponents = new SetOnce<>();
     private final SetOnce<InferenceServiceRegistry> inferenceServiceRegistry = new SetOnce<>();
     private final SetOnce<ShardBulkInferenceActionFilter> shardBulkInferenceActionFilter = new SetOnce<>();
+    private final SetOnce<CircuitBreaker> inferenceBytesBreaker = new SetOnce<>();
     private List<InferenceServiceExtension> inferenceServiceExtensions;
 
     public InferencePlugin(Settings settings) {
@@ -324,8 +339,14 @@ public Collection<?> createComponents(PluginServices services) {
         inferenceServiceRegistry.set(serviceRegistry);
 
         var actionFilter = new ShardBulkInferenceActionFilter(services.clusterService(), serviceRegistry, modelRegistry, getLicenseState());
+        actionFilter.setInferenceBytesCircuitBreaker(inferenceBytesBreaker.get());
         shardBulkInferenceActionFilter.set(actionFilter);
 
+        services.clusterService().getClusterSettings().addSettingsUpdateConsumer(
+            INFERENCE_BYTES_CIRCUIT_BREAKER_LIMIT_SETTING,
+            this::updateCircuitBreaker
+        );
+
         var meterRegistry = services.telemetryProvider().getMeterRegistry();
         var inferenceStats = new PluginComponentBinding<>(InferenceStats.class, InferenceStats.create(meterRegistry));
 
@@ -467,6 +488,7 @@ public List<Setting<?>> getSettings() {
         settings.add(SKIP_VALIDATE_AND_START);
         settings.add(INDICES_INFERENCE_BATCH_SIZE);
         settings.addAll(ElasticInferenceServiceSettings.getSettingsDefinitions());
+        settings.add(INFERENCE_BYTES_CIRCUIT_BREAKER_LIMIT_SETTING);
 
         return settings;
     }
@@ -538,6 +560,32 @@ public Map<String, Highlighter> getHighlighters() {
         return Map.of(SemanticTextHighlighter.NAME, new SemanticTextHighlighter());
     }
 
+    @Override
+    public BreakerSettings getCircuitBreaker(Settings settings) {
+        return new BreakerSettings(
+            INFERENCE_BYTES_CIRCUIT_BREAKER_NAME,
+            INFERENCE_BYTES_CIRCUIT_BREAKER_LIMIT_SETTING.get(settings).getBytes(),
+            INFERENCE_BYTES_CIRCUIT_BREAKER_OVERHEAD,
+            CircuitBreaker.Type.MEMORY,
+            CircuitBreaker.Durability.TRANSIENT
+        );
+    }
+
+    @Override
+    public void setCircuitBreaker(CircuitBreaker circuitBreaker) {
+        assert circuitBreaker.getName().equals(INFERENCE_BYTES_CIRCUIT_BREAKER_NAME);
+        this.inferenceBytesBreaker.set(circuitBreaker);
+    }
+
+    private void updateCircuitBreaker(ByteSizeValue inferenceBytesLimit) {
+        CircuitBreaker circuitBreaker = this.inferenceBytesBreaker.get();
+        if (circuitBreaker == null) {
+            throw new IllegalStateException("[" + INFERENCE_BYTES_CIRCUIT_BREAKER_NAME + "] circuit breaker not set");
+        }
+
+        circuitBreaker.setLimitAndOverhead(inferenceBytesLimit.getBytes(), INFERENCE_BYTES_CIRCUIT_BREAKER_OVERHEAD);
+    }
+
     @Override
     public void onNodeStarted() {
         var registry = inferenceServiceRegistry.get();