[ML] Integrate SageMaker with OpenAI Embeddings (elastic#126856) (elastic#127610)

Integrating with SageMaker.

Current design:
- SageMaker accepts any byte payload, which can be text, csv, or json. `api` represents the structure of the payload that we will send, for example `openai`, `elastic`, `common`, probably `cohere` or `huggingface` as well.
- `api` implementations are extensions of `SageMakerSchemaPayload`, which supports:
  - "extra" service and task settings specific to the payload structure, so `cohere` would require `embedding_type` and `openai` would require `dimensions` in the `service_settings`
  - conversion logic from model, service settings, task settings, and input to `SdkBytes`
  - conversion logic from responding `SdkBytes` to `InferenceServiceResults`
- Everything else is tunneling, there are a number of base `service_settings` and `task_settings` that are independent of the api format that we will store and set
- We let the SDK do the bulk of the work in terms of connection details, rate limiting, retries, etc.

Author: prwhelan

docs/changelog/126856.yaml

@@ -0,0 +1,5 @@
+pr: 126856
+summary: "[ML] Integrate SageMaker with OpenAI Embeddings"
+area: Machine Learning
+type: enhancement
+issues: []

gradle/verification-metadata.xml

@@ -4802,6 +4802,11 @@
             <sha256 value="c83dd82a9d82ff8c7d2eb1bdb2ae9f9505b312dad9a6bf0b80bc0136653a3a24" origin="Generated by Gradle"/>
          </artifact>
       </component>
+      <component group="software.amazon.awssdk" name="sagemakerruntime" version="2.30.38">
+         <artifact name="sagemakerruntime-2.30.38.jar">
+            <sha256 value="b26ee73fa06d047eab9a174e49627972e646c0bbe909f479c18dbff193b561f5" origin="Generated by Gradle"/>
+         </artifact>
+      </component>
       <component group="software.amazon.awssdk" name="sdk-core" version="2.30.38">
          <artifact name="sdk-core-2.30.38.jar">
             <sha256 value="556463b8c353408d93feab74719d141fcfda7fd3d7b7d1ad3a8a548b7cc2982d" origin="Generated by Gradle"/>

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

@@ -207,6 +207,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion SEMANTIC_TEXT_CHUNKING_CONFIG_8_19 = def(8_841_0_18);
     public static final TransportVersion BATCHED_QUERY_PHASE_VERSION_BACKPORT_8_X = def(8_841_0_19);
     public static final TransportVersion SEARCH_INCREMENTAL_TOP_DOCS_NULL_BACKPORT_8_19 = def(8_841_0_20);
+    public static final TransportVersion ML_INFERENCE_SAGEMAKER_8_19 = def(8_841_0_21);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/common/ValidationException.java

@@ -53,6 +53,12 @@ public final List<String> validationErrors() {
         return validationErrors;
     }
 
+    public final void throwIfValidationErrorsExist() {
+        if (validationErrors().isEmpty() == false) {
+            throw this;
+        }
+    }
+
     @Override
     public final String getMessage() {
         StringBuilder sb = new StringBuilder();

x-pack/plugin/inference/build.gradle

@@ -63,6 +63,7 @@ dependencies {
 
   /* AWS SDK v2 */
   implementation("software.amazon.awssdk:bedrockruntime:${versions.awsv2sdk}")
+  implementation("software.amazon.awssdk:sagemakerruntime:${versions.awsv2sdk}")
   api "software.amazon.awssdk:protocol-core:${versions.awsv2sdk}"
   api "software.amazon.awssdk:aws-json-protocol:${versions.awsv2sdk}"
   api "software.amazon.awssdk:third-party-jackson-core:${versions.awsv2sdk}"
@@ -143,6 +144,7 @@ tasks.named("dependencyLicenses").configure {
   mapping from: /json-utils.*/, to: 'aws-sdk-2'
   mapping from: /endpoints-spi.*/, to: 'aws-sdk-2'
   mapping from: /bedrockruntime.*/, to: 'aws-sdk-2'
+  mapping from: /sagemakerruntime.*/, to: 'aws-sdk-2'
   mapping from: /netty-nio-client/, to: 'aws-sdk-2'
   /* Cannot use REGEX to match netty-* because netty-nio-client is an AWS package */
   mapping from: /netty-buffer/, to: 'netty'

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java

@@ -14,169 +14,165 @@
 import org.elasticsearch.inference.TaskType;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
 import static org.elasticsearch.xpack.inference.InferenceBaseRestTest.assertStatusOkOrCreated;
+import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.equalTo;
 
 public class InferenceGetServicesIT extends BaseMockEISAuthServerTest {
 
-    @SuppressWarnings("unchecked")
     public void testGetServicesWithoutTaskType() throws IOException {
         List<Object> services = getAllServices();
-        assertThat(services.size(), equalTo(21));
-
-        String[] providers = new String[services.size()];
-        for (int i = 0; i < services.size(); i++) {
-            Map<String, Object> serviceConfig = (Map<String, Object>) services.get(i);
-            providers[i] = (String) serviceConfig.get("service");
-        }
-
-        assertArrayEquals(
-            List.of(
-                "alibabacloud-ai-search",
-                "amazonbedrock",
-                "anthropic",
-                "azureaistudio",
-                "azureopenai",
-                "cohere",
-                "deepseek",
-                "elastic",
-                "elasticsearch",
-                "googleaistudio",
-                "googlevertexai",
-                "hugging_face",
-                "jinaai",
-                "mistral",
-                "openai",
-                "streaming_completion_test_service",
-                "test_reranking_service",
-                "test_service",
-                "text_embedding_test_service",
-                "voyageai",
-                "watsonxai"
-            ).toArray(),
-            providers
+        assertThat(services.size(), equalTo(22));
+
+        var providers = providers(services);
+
+        assertThat(
+            providers,
+            containsInAnyOrder(
+                List.of(
+                    "alibabacloud-ai-search",
+                    "amazonbedrock",
+                    "anthropic",
+                    "azureaistudio",
+                    "azureopenai",
+                    "cohere",
+                    "deepseek",
+                    "elastic",
+                    "elasticsearch",
+                    "googleaistudio",
+                    "googlevertexai",
+                    "hugging_face",
+                    "jinaai",
+                    "mistral",
+                    "openai",
+                    "streaming_completion_test_service",
+                    "test_reranking_service",
+                    "test_service",
+                    "text_embedding_test_service",
+                    "voyageai",
+                    "watsonxai",
+                    "sagemaker"
+                ).toArray()
+            )
         );
     }
 
     @SuppressWarnings("unchecked")
+    private Iterable<String> providers(List<Object> services) {
+        return services.stream().map(service -> {
+            var serviceConfig = (Map<String, Object>) service;
+            return (String) serviceConfig.get("service");
+        }).toList();
+    }
+
     public void testGetServicesWithTextEmbeddingTaskType() throws IOException {
         List<Object> services = getServices(TaskType.TEXT_EMBEDDING);
-        assertThat(services.size(), equalTo(15));
-
-        String[] providers = new String[services.size()];
-        for (int i = 0; i < services.size(); i++) {
-            Map<String, Object> serviceConfig = (Map<String, Object>) services.get(i);
-            providers[i] = (String) serviceConfig.get("service");
-        }
-
-        assertArrayEquals(
-            List.of(
-                "alibabacloud-ai-search",
-                "amazonbedrock",
-                "azureaistudio",
-                "azureopenai",
-                "cohere",
-                "elasticsearch",
-                "googleaistudio",
-                "googlevertexai",
-                "hugging_face",
-                "jinaai",
-                "mistral",
-                "openai",
-                "text_embedding_test_service",
-                "voyageai",
-                "watsonxai"
-            ).toArray(),
-            providers
+        assertThat(services.size(), equalTo(16));
+
+        var providers = providers(services);
+
+        assertThat(
+            providers,
+            containsInAnyOrder(
+                List.of(
+                    "alibabacloud-ai-search",
+                    "amazonbedrock",
+                    "azureaistudio",
+                    "azureopenai",
+                    "cohere",
+                    "elasticsearch",
+                    "googleaistudio",
+                    "googlevertexai",
+                    "hugging_face",
+                    "jinaai",
+                    "mistral",
+                    "openai",
+                    "text_embedding_test_service",
+                    "voyageai",
+                    "watsonxai",
+                    "sagemaker"
+                ).toArray()
+            )
         );
     }
 
-    @SuppressWarnings("unchecked")
     public void testGetServicesWithRerankTaskType() throws IOException {
         List<Object> services = getServices(TaskType.RERANK);
         assertThat(services.size(), equalTo(7));
 
-        String[] providers = new String[services.size()];
-        for (int i = 0; i < services.size(); i++) {
-            Map<String, Object> serviceConfig = (Map<String, Object>) services.get(i);
-            providers[i] = (String) serviceConfig.get("service");
-        }
-
-        assertArrayEquals(
-            List.of("alibabacloud-ai-search", "cohere", "elasticsearch", "googlevertexai", "jinaai", "test_reranking_service", "voyageai")
-                .toArray(),
-            providers
+        var providers = providers(services);
+
+        assertThat(
+            providers,
+            containsInAnyOrder(
+                List.of(
+                    "alibabacloud-ai-search",
+                    "cohere",
+                    "elasticsearch",
+                    "googlevertexai",
+                    "jinaai",
+                    "test_reranking_service",
+                    "voyageai"
+                ).toArray()
+            )
         );
     }
 
-    @SuppressWarnings("unchecked")
     public void testGetServicesWithCompletionTaskType() throws IOException {
         List<Object> services = getServices(TaskType.COMPLETION);
         assertThat(services.size(), equalTo(10));
 
-        String[] providers = new String[services.size()];
-        for (int i = 0; i < services.size(); i++) {
-            Map<String, Object> serviceConfig = (Map<String, Object>) services.get(i);
-            providers[i] = (String) serviceConfig.get("service");
-        }
-
-        var providerList = new ArrayList<>(
-            List.of(
-                "alibabacloud-ai-search",
-                "amazonbedrock",
-                "anthropic",
-                "azureaistudio",
-                "azureopenai",
-                "cohere",
-                "deepseek",
-                "googleaistudio",
-                "openai",
-                "streaming_completion_test_service"
+        var providers = providers(services);
+
+        assertThat(
+            providers,
+            containsInAnyOrder(
+                List.of(
+                    "alibabacloud-ai-search",
+                    "amazonbedrock",
+                    "anthropic",
+                    "azureaistudio",
+                    "azureopenai",
+                    "cohere",
+                    "deepseek",
+                    "googleaistudio",
+                    "openai",
+                    "streaming_completion_test_service"
+                ).toArray()
             )
         );
-
-        assertArrayEquals(providers, providerList.toArray());
     }
 
-    @SuppressWarnings("unchecked")
     public void testGetServicesWithChatCompletionTaskType() throws IOException {
         List<Object> services = getServices(TaskType.CHAT_COMPLETION);
         assertThat(services.size(), equalTo(4));
 
-        String[] providers = new String[services.size()];
-        for (int i = 0; i < services.size(); i++) {
-            Map<String, Object> serviceConfig = (Map<String, Object>) services.get(i);
-            providers[i] = (String) serviceConfig.get("service");
-        }
+        var providers = providers(services);
 
-        assertArrayEquals(List.of("deepseek", "elastic", "openai", "streaming_completion_test_service").toArray(), providers);
+        assertThat(providers, containsInAnyOrder(List.of("deepseek", "elastic", "openai", "streaming_completion_test_service").toArray()));
     }
 
-    @SuppressWarnings("unchecked")
     public void testGetServicesWithSparseEmbeddingTaskType() throws IOException {
         List<Object> services = getServices(TaskType.SPARSE_EMBEDDING);
         assertThat(services.size(), equalTo(6));
 
-        String[] providers = new String[services.size()];
-        for (int i = 0; i < services.size(); i++) {
-            Map<String, Object> serviceConfig = (Map<String, Object>) services.get(i);
-            providers[i] = (String) serviceConfig.get("service");
-        }
-
-        assertArrayEquals(
-            List.of(
-                "alibabacloud-ai-search",
-                "elastic",
-                "elasticsearch",
-                "hugging_face",
-                "streaming_completion_test_service",
-                "test_service"
-            ).toArray(),
-            providers
+        var providers = providers(services);
+
+        assertThat(
+            providers,
+            containsInAnyOrder(
+                List.of(
+                    "alibabacloud-ai-search",
+                    "elastic",
+                    "elasticsearch",
+                    "hugging_face",
+                    "streaming_completion_test_service",
+                    "test_service"
+                ).toArray()
+            )
         );
     }
 

x-pack/plugin/inference/src/main/java/module-info.java

@@ -36,6 +36,7 @@
     requires org.elasticsearch.logging;
     requires org.elasticsearch.sslconfig;
     requires org.apache.commons.text;
+    requires software.amazon.awssdk.services.sagemakerruntime;
 
     exports org.elasticsearch.xpack.inference.action;
     exports org.elasticsearch.xpack.inference.registry;

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

@@ -92,6 +92,8 @@
 import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionTaskSettings;
 import org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsTaskSettings;
+import org.elasticsearch.xpack.inference.services.sagemaker.model.SageMakerModel;
+import org.elasticsearch.xpack.inference.services.sagemaker.schema.SageMakerSchemas;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 import org.elasticsearch.xpack.inference.services.voyageai.VoyageAIServiceSettings;
 import org.elasticsearch.xpack.inference.services.voyageai.embeddings.VoyageAIEmbeddingsServiceSettings;
@@ -157,6 +159,8 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
 
         namedWriteables.addAll(StreamingTaskManager.namedWriteables());
         namedWriteables.addAll(DeepSeekChatCompletionModel.namedWriteables());
+        namedWriteables.addAll(SageMakerModel.namedWriteables());
+        namedWriteables.addAll(SageMakerSchemas.namedWriteables());
 
         return namedWriteables;
     }