propagate clustersettings to sageMaker

Author: Samiul-TheSoccerFan

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

@@ -329,7 +329,8 @@ public Collection<?> createComponents(PluginServices services) {
                     ),
                     sageMakerSchemas,
                     services.threadPool(),
-                    sageMakerConfigurations::getOrCompute
+                    sageMakerConfigurations::getOrCompute,
+                    serviceComponents.get()
                 )
             )
         );

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/sagemaker/SageMakerService.java

@@ -28,7 +28,9 @@
 import org.elasticsearch.inference.UnifiedCompletionRequest;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.inference.InferencePlugin;
 import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
+import org.elasticsearch.xpack.inference.services.ServiceComponents;
 import org.elasticsearch.xpack.inference.services.sagemaker.model.SageMakerModel;
 import org.elasticsearch.xpack.inference.services.sagemaker.model.SageMakerModelBuilder;
 import org.elasticsearch.xpack.inference.services.sagemaker.schema.SageMakerSchemas;
@@ -55,13 +57,15 @@ public class SageMakerService implements InferenceService {
     private final SageMakerSchemas schemas;
     private final ThreadPool threadPool;
     private final LazyInitializable<InferenceServiceConfiguration, RuntimeException> configuration;
+    private final ServiceComponents serviceComponents;
 
     public SageMakerService(
         SageMakerModelBuilder modelBuilder,
         SageMakerClient client,
         SageMakerSchemas schemas,
         ThreadPool threadPool,
-        CheckedSupplier<Map<String, SettingsConfiguration>, RuntimeException> configurationMap
+        CheckedSupplier<Map<String, SettingsConfiguration>, RuntimeException> configurationMap,
+        ServiceComponents serviceComponents
     ) {
         this.modelBuilder = modelBuilder;
         this.client = client;
@@ -74,6 +78,7 @@ public SageMakerService(
                 .setConfigurations(configurationMap.get())
                 .build()
         );
+        this.serviceComponents = serviceComponents;
     }
 
     @Override
@@ -146,6 +151,10 @@ public void infer(
 
         var inferenceRequest = new SageMakerInferenceRequest(query, returnDocuments, topN, input, stream, inputType);
 
+        if (timeout == null) {
+            timeout = serviceComponents.clusterService().getClusterSettings().get(InferencePlugin.SEMANTIC_TEXT_INFERENCE_TIMEOUT);
+        }
+
         try {
             var sageMakerModel = ((SageMakerModel) model).override(taskSettings);
             var regionAndSecrets = regionAndSecrets(sageMakerModel);
@@ -156,7 +165,7 @@ public void infer(
                 client.invokeStream(
                     regionAndSecrets,
                     request,
-                    timeout != null ? timeout : DEFAULT_TIMEOUT,
+                    timeout,
                     ActionListener.wrap(
                         response -> listener.onResponse(schema.streamResponse(sageMakerModel, response)),
                         e -> listener.onFailure(schema.error(sageMakerModel, e))
@@ -168,7 +177,7 @@ public void infer(
                 client.invoke(
                     regionAndSecrets,
                     request,
-                    timeout != null ? timeout : DEFAULT_TIMEOUT,
+                    timeout,
                     ActionListener.wrap(
                         response -> listener.onResponse(schema.response(sageMakerModel, response, threadPool.getThreadContext())),
                         e -> listener.onFailure(schema.error(sageMakerModel, e))

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/sagemaker/SageMakerServiceTests.java

@@ -7,6 +7,8 @@
 
 package org.elasticsearch.xpack.inference.services.sagemaker;
 
+import org.elasticsearch.xpack.inference.services.ServiceComponents;
+
 import software.amazon.awssdk.services.sagemakerruntime.model.InvokeEndpointResponse;
 import software.amazon.awssdk.services.sagemakerruntime.model.InvokeEndpointWithResponseStreamResponse;
 
@@ -75,16 +77,18 @@ public class SageMakerServiceTests extends ESTestCase {
     private SageMakerClient client;
     private SageMakerSchemas schemas;
     private SageMakerService sageMakerService;
+    private ServiceComponents serviceComponents;
 
     @Before
     public void init() {
         modelBuilder = mock();
         client = mock();
         schemas = mock();
+        serviceComponents = mock();
         ThreadPool threadPool = mock();
         when(threadPool.executor(anyString())).thenReturn(EsExecutors.DIRECT_EXECUTOR_SERVICE);
         when(threadPool.getThreadContext()).thenReturn(new ThreadContext(Settings.EMPTY));
-        sageMakerService = new SageMakerService(modelBuilder, client, schemas, threadPool, Map::of);
+        sageMakerService = new SageMakerService(modelBuilder, client, schemas, threadPool, Map::of, serviceComponents);
     }
 
     public void testSupportedTaskTypes() {