Browse Source

[Inference API] Fix upgrade tests by accounting for breaking change of renaming models (#110547)

* Fix upgrade tests by accounting for breaking change of renaming "models" to "endpoints"

* improvements from review

---------

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>
Max Hniebergall 1 năm trước cách đây
mục cha
commit
688870ca8b

+ 10 - 6
x-pack/plugin/inference/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/xpack/application/AzureOpenAiServiceUpgradeIT.java

@@ -49,35 +49,39 @@ public class AzureOpenAiServiceUpgradeIT extends InferenceUpgradeTestCase {
     @AwaitsFix(bugUrl = "Cannot set the URL in the tests")
     public void testOpenAiEmbeddings() throws IOException {
         var openAiEmbeddingsSupported = getOldClusterTestVersion().onOrAfter(OPEN_AI_AZURE_EMBEDDINGS_ADDED);
+        // `gte_v` indicates that the cluster version is Greater Than or Equal to MODELS_RENAMED_TO_ENDPOINTS
+        String oldClusterEndpointIdentifier = oldClusterHasFeature("gte_v" + MODELS_RENAMED_TO_ENDPOINTS) ? "endpoints" : "models";
         assumeTrue("Azure OpenAI embedding service added in " + OPEN_AI_AZURE_EMBEDDINGS_ADDED, openAiEmbeddingsSupported);
 
         final String oldClusterId = "old-cluster-embeddings";
         final String upgradedClusterId = "upgraded-cluster-embeddings";
 
+        var testTaskType = TaskType.TEXT_EMBEDDING;
+
         if (isOldCluster()) {
             // queue a response as PUT will call the service
             openAiEmbeddingsServer.enqueue(new MockResponse().setResponseCode(200).setBody(OpenAiServiceUpgradeIT.embeddingResponse()));
-            put(oldClusterId, embeddingConfig(getUrl(openAiEmbeddingsServer)), TaskType.TEXT_EMBEDDING);
+            put(oldClusterId, embeddingConfig(getUrl(openAiEmbeddingsServer)), testTaskType);
 
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterId).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get(oldClusterEndpointIdentifier);
             assertThat(configs, hasSize(1));
         } else if (isMixedCluster()) {
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterId).get("endpoints");
+            var configs = getConfigsWithBreakingChangeHandling(testTaskType, oldClusterId);
             assertEquals("azureopenai", configs.get(0).get("service"));
 
             assertEmbeddingInference(oldClusterId);
         } else if (isUpgradedCluster()) {
             // check old cluster model
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterId).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get("endpoints");
             var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
 
             // Inference on old cluster model
             assertEmbeddingInference(oldClusterId);
 
             openAiEmbeddingsServer.enqueue(new MockResponse().setResponseCode(200).setBody(OpenAiServiceUpgradeIT.embeddingResponse()));
-            put(upgradedClusterId, embeddingConfig(getUrl(openAiEmbeddingsServer)), TaskType.TEXT_EMBEDDING);
+            put(upgradedClusterId, embeddingConfig(getUrl(openAiEmbeddingsServer)), testTaskType);
 
-            configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, upgradedClusterId).get("endpoints");
+            configs = (List<Map<String, Object>>) get(testTaskType, upgradedClusterId).get("endpoints");
             assertThat(configs, hasSize(1));
 
             // Inference on the new config

+ 64 - 43
x-pack/plugin/inference/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/xpack/application/CohereServiceUpgradeIT.java

@@ -15,6 +15,8 @@ import org.elasticsearch.test.http.MockResponse;
 import org.elasticsearch.test.http.MockWebServer;
 import org.elasticsearch.xpack.inference.services.cohere.embeddings.CohereEmbeddingType;
 import org.hamcrest.Matchers;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
 
 import java.io.IOException;
 import java.util.List;
@@ -39,7 +41,7 @@ public class CohereServiceUpgradeIT extends InferenceUpgradeTestCase {
         super(upgradedNodes);
     }
 
-    // @BeforeClass
+    @BeforeClass
     public static void startWebServer() throws IOException {
         cohereEmbeddingsServer = new MockWebServer();
         cohereEmbeddingsServer.start();
@@ -48,58 +50,74 @@ public class CohereServiceUpgradeIT extends InferenceUpgradeTestCase {
         cohereRerankServer.start();
     }
 
-    // @AfterClass // for the awaitsfix
+    @AfterClass
     public static void shutdown() {
         cohereEmbeddingsServer.close();
         cohereRerankServer.close();
     }
 
     @SuppressWarnings("unchecked")
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/107887")
     public void testCohereEmbeddings() throws IOException {
         var embeddingsSupported = getOldClusterTestVersion().onOrAfter(COHERE_EMBEDDINGS_ADDED);
+        // `gte_v` indicates that the cluster version is Greater Than or Equal to MODELS_RENAMED_TO_ENDPOINTS
+        String oldClusterEndpointIdentifier = oldClusterHasFeature("gte_v" + MODELS_RENAMED_TO_ENDPOINTS) ? "endpoints" : "models";
         assumeTrue("Cohere embedding service added in " + COHERE_EMBEDDINGS_ADDED, embeddingsSupported);
 
         final String oldClusterIdInt8 = "old-cluster-embeddings-int8";
         final String oldClusterIdFloat = "old-cluster-embeddings-float";
 
+        var testTaskType = TaskType.TEXT_EMBEDDING;
+
         if (isOldCluster()) {
             // queue a response as PUT will call the service
             cohereEmbeddingsServer.enqueue(new MockResponse().setResponseCode(200).setBody(embeddingResponseByte()));
-            put(oldClusterIdInt8, embeddingConfigInt8(getUrl(cohereEmbeddingsServer)), TaskType.TEXT_EMBEDDING);
+            put(oldClusterIdInt8, embeddingConfigInt8(getUrl(cohereEmbeddingsServer)), testTaskType);
             // float model
             cohereEmbeddingsServer.enqueue(new MockResponse().setResponseCode(200).setBody(embeddingResponseFloat()));
-            put(oldClusterIdFloat, embeddingConfigFloat(getUrl(cohereEmbeddingsServer)), TaskType.TEXT_EMBEDDING);
+            put(oldClusterIdFloat, embeddingConfigFloat(getUrl(cohereEmbeddingsServer)), testTaskType);
 
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterIdInt8).get("endpoints");
-            assertThat(configs, hasSize(1));
-            assertEquals("cohere", configs.get(0).get("service"));
-            var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
-            assertThat(serviceSettings, hasEntry("model_id", "embed-english-light-v3.0"));
-            var embeddingType = serviceSettings.get("embedding_type");
-            // An upgraded node will report the embedding type as byte, the old node int8
-            assertThat(embeddingType, Matchers.is(oneOf("int8", "byte")));
-
-            assertEmbeddingInference(oldClusterIdInt8, CohereEmbeddingType.BYTE);
-            assertEmbeddingInference(oldClusterIdFloat, CohereEmbeddingType.FLOAT);
+            {
+                var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterIdInt8).get(oldClusterEndpointIdentifier);
+                assertThat(configs, hasSize(1));
+                assertEquals("cohere", configs.get(0).get("service"));
+                var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
+                assertThat(serviceSettings, hasEntry("model_id", "embed-english-light-v3.0"));
+                var embeddingType = serviceSettings.get("embedding_type");
+                // An upgraded node will report the embedding type as byte, the old node int8
+                assertThat(embeddingType, Matchers.is(oneOf("int8", "byte")));
+                assertEmbeddingInference(oldClusterIdInt8, CohereEmbeddingType.BYTE);
+            }
+            {
+                var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterIdFloat).get(oldClusterEndpointIdentifier);
+                assertThat(configs, hasSize(1));
+                assertEquals("cohere", configs.get(0).get("service"));
+                var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
+                assertThat(serviceSettings, hasEntry("model_id", "embed-english-light-v3.0"));
+                assertThat(serviceSettings, hasEntry("embedding_type", "float"));
+                assertEmbeddingInference(oldClusterIdFloat, CohereEmbeddingType.FLOAT);
+            }
         } else if (isMixedCluster()) {
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterIdInt8).get("endpoints");
-            assertEquals("cohere", configs.get(0).get("service"));
-            var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
-            assertThat(serviceSettings, hasEntry("model_id", "embed-english-light-v3.0"));
-            var embeddingType = serviceSettings.get("embedding_type");
-            // An upgraded node will report the embedding type as byte, an old node int8
-            assertThat(embeddingType, Matchers.is(oneOf("int8", "byte")));
-
-            configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterIdFloat).get("endpoints");
-            serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
-            assertThat(serviceSettings, hasEntry("embedding_type", "float"));
-
-            assertEmbeddingInference(oldClusterIdInt8, CohereEmbeddingType.BYTE);
-            assertEmbeddingInference(oldClusterIdFloat, CohereEmbeddingType.FLOAT);
+            {
+                var configs = getConfigsWithBreakingChangeHandling(testTaskType, oldClusterIdInt8);
+                assertEquals("cohere", configs.get(0).get("service"));
+                var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
+                assertThat(serviceSettings, hasEntry("model_id", "embed-english-light-v3.0"));
+                var embeddingType = serviceSettings.get("embedding_type");
+                // An upgraded node will report the embedding type as byte, an old node int8
+                assertThat(embeddingType, Matchers.is(oneOf("int8", "byte")));
+                assertEmbeddingInference(oldClusterIdInt8, CohereEmbeddingType.BYTE);
+            }
+            {
+                var configs = getConfigsWithBreakingChangeHandling(testTaskType, oldClusterIdFloat);
+                assertEquals("cohere", configs.get(0).get("service"));
+                var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
+                assertThat(serviceSettings, hasEntry("model_id", "embed-english-light-v3.0"));
+                assertThat(serviceSettings, hasEntry("embedding_type", "float"));
+                assertEmbeddingInference(oldClusterIdFloat, CohereEmbeddingType.FLOAT);
+            }
         } else if (isUpgradedCluster()) {
             // check old cluster model
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterIdInt8).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterIdInt8).get("endpoints");
             var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
             assertThat(serviceSettings, hasEntry("model_id", "embed-english-light-v3.0"));
             assertThat(serviceSettings, hasEntry("embedding_type", "byte"));
@@ -114,9 +132,9 @@ public class CohereServiceUpgradeIT extends InferenceUpgradeTestCase {
                 final String upgradedClusterIdByte = "upgraded-cluster-embeddings-byte";
 
                 cohereEmbeddingsServer.enqueue(new MockResponse().setResponseCode(200).setBody(embeddingResponseByte()));
-                put(upgradedClusterIdByte, embeddingConfigByte(getUrl(cohereEmbeddingsServer)), TaskType.TEXT_EMBEDDING);
+                put(upgradedClusterIdByte, embeddingConfigByte(getUrl(cohereEmbeddingsServer)), testTaskType);
 
-                configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, upgradedClusterIdByte).get("endpoints");
+                configs = (List<Map<String, Object>>) get(testTaskType, upgradedClusterIdByte).get("endpoints");
                 serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
                 assertThat(serviceSettings, hasEntry("embedding_type", "byte"));
 
@@ -127,9 +145,9 @@ public class CohereServiceUpgradeIT extends InferenceUpgradeTestCase {
                 final String upgradedClusterIdInt8 = "upgraded-cluster-embeddings-int8";
 
                 cohereEmbeddingsServer.enqueue(new MockResponse().setResponseCode(200).setBody(embeddingResponseByte()));
-                put(upgradedClusterIdInt8, embeddingConfigInt8(getUrl(cohereEmbeddingsServer)), TaskType.TEXT_EMBEDDING);
+                put(upgradedClusterIdInt8, embeddingConfigInt8(getUrl(cohereEmbeddingsServer)), testTaskType);
 
-                configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, upgradedClusterIdInt8).get("endpoints");
+                configs = (List<Map<String, Object>>) get(testTaskType, upgradedClusterIdInt8).get("endpoints");
                 serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
                 assertThat(serviceSettings, hasEntry("embedding_type", "byte")); // int8 rewritten to byte
 
@@ -139,9 +157,9 @@ public class CohereServiceUpgradeIT extends InferenceUpgradeTestCase {
             {
                 final String upgradedClusterIdFloat = "upgraded-cluster-embeddings-float";
                 cohereEmbeddingsServer.enqueue(new MockResponse().setResponseCode(200).setBody(embeddingResponseFloat()));
-                put(upgradedClusterIdFloat, embeddingConfigFloat(getUrl(cohereEmbeddingsServer)), TaskType.TEXT_EMBEDDING);
+                put(upgradedClusterIdFloat, embeddingConfigFloat(getUrl(cohereEmbeddingsServer)), testTaskType);
 
-                configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, upgradedClusterIdFloat).get("endpoints");
+                configs = (List<Map<String, Object>>) get(testTaskType, upgradedClusterIdFloat).get("endpoints");
                 serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
                 assertThat(serviceSettings, hasEntry("embedding_type", "float"));
 
@@ -169,22 +187,25 @@ public class CohereServiceUpgradeIT extends InferenceUpgradeTestCase {
     }
 
     @SuppressWarnings("unchecked")
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/107887")
     public void testRerank() throws IOException {
         var rerankSupported = getOldClusterTestVersion().onOrAfter(COHERE_RERANK_ADDED);
+        String old_cluster_endpoint_identifier = oldClusterHasFeature("gte_v" + MODELS_RENAMED_TO_ENDPOINTS) ? "endpoints" : "models";
         assumeTrue("Cohere rerank service added in " + COHERE_RERANK_ADDED, rerankSupported);
 
         final String oldClusterId = "old-cluster-rerank";
         final String upgradedClusterId = "upgraded-cluster-rerank";
 
+        var testTaskType = TaskType.RERANK;
+
         if (isOldCluster()) {
-            put(oldClusterId, rerankConfig(getUrl(cohereRerankServer)), TaskType.RERANK);
-            var configs = (List<Map<String, Object>>) get(TaskType.RERANK, oldClusterId).get("endpoints");
+            put(oldClusterId, rerankConfig(getUrl(cohereRerankServer)), testTaskType);
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get(old_cluster_endpoint_identifier);
             assertThat(configs, hasSize(1));
 
             assertRerank(oldClusterId);
         } else if (isMixedCluster()) {
-            var configs = (List<Map<String, Object>>) get(TaskType.RERANK, oldClusterId).get("endpoints");
+            var configs = getConfigsWithBreakingChangeHandling(testTaskType, oldClusterId);
+
             assertEquals("cohere", configs.get(0).get("service"));
             var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
             assertThat(serviceSettings, hasEntry("model_id", "rerank-english-v3.0"));
@@ -195,7 +216,7 @@ public class CohereServiceUpgradeIT extends InferenceUpgradeTestCase {
 
         } else if (isUpgradedCluster()) {
             // check old cluster model
-            var configs = (List<Map<String, Object>>) get(TaskType.RERANK, oldClusterId).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get("endpoints");
             assertEquals("cohere", configs.get(0).get("service"));
             var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
             assertThat(serviceSettings, hasEntry("model_id", "rerank-english-v3.0"));
@@ -205,7 +226,7 @@ public class CohereServiceUpgradeIT extends InferenceUpgradeTestCase {
             assertRerank(oldClusterId);
 
             // New endpoint
-            put(upgradedClusterId, rerankConfig(getUrl(cohereRerankServer)), TaskType.RERANK);
+            put(upgradedClusterId, rerankConfig(getUrl(cohereRerankServer)), testTaskType);
             configs = (List<Map<String, Object>>) get(upgradedClusterId).get("endpoints");
             assertThat(configs, hasSize(1));
 

+ 23 - 15
x-pack/plugin/inference/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/xpack/application/HuggingFaceServiceUpgradeIT.java

@@ -13,6 +13,8 @@ import org.elasticsearch.common.Strings;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.test.http.MockResponse;
 import org.elasticsearch.test.http.MockWebServer;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
 
 import java.io.IOException;
 import java.util.List;
@@ -34,7 +36,7 @@ public class HuggingFaceServiceUpgradeIT extends InferenceUpgradeTestCase {
         super(upgradedNodes);
     }
 
-    // @BeforeClass
+    @BeforeClass
     public static void startWebServer() throws IOException {
         embeddingsServer = new MockWebServer();
         embeddingsServer.start();
@@ -43,47 +45,51 @@ public class HuggingFaceServiceUpgradeIT extends InferenceUpgradeTestCase {
         elserServer.start();
     }
 
-    // @AfterClass for the awaits fix
+    @AfterClass
     public static void shutdown() {
         embeddingsServer.close();
         elserServer.close();
     }
 
     @SuppressWarnings("unchecked")
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/107887")
     public void testHFEmbeddings() throws IOException {
         var embeddingsSupported = getOldClusterTestVersion().onOrAfter(HF_EMBEDDINGS_ADDED);
+        // `gte_v` indicates that the cluster version is Greater Than or Equal to MODELS_RENAMED_TO_ENDPOINTS
+        String oldClusterEndpointIdentifier = oldClusterHasFeature("gte_v" + MODELS_RENAMED_TO_ENDPOINTS) ? "endpoints" : "models";
         assumeTrue("Hugging Face embedding service added in " + HF_EMBEDDINGS_ADDED, embeddingsSupported);
 
         final String oldClusterId = "old-cluster-embeddings";
         final String upgradedClusterId = "upgraded-cluster-embeddings";
 
+        var testTaskType = TaskType.TEXT_EMBEDDING;
+
         if (isOldCluster()) {
             // queue a response as PUT will call the service
             embeddingsServer.enqueue(new MockResponse().setResponseCode(200).setBody(embeddingResponse()));
-            put(oldClusterId, embeddingConfig(getUrl(embeddingsServer)), TaskType.TEXT_EMBEDDING);
+            put(oldClusterId, embeddingConfig(getUrl(embeddingsServer)), testTaskType);
 
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterId).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get(oldClusterEndpointIdentifier);
             assertThat(configs, hasSize(1));
 
             assertEmbeddingInference(oldClusterId);
         } else if (isMixedCluster()) {
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterId).get("endpoints");
+            var configs = getConfigsWithBreakingChangeHandling(testTaskType, oldClusterId);
+
             assertEquals("hugging_face", configs.get(0).get("service"));
 
             assertEmbeddingInference(oldClusterId);
         } else if (isUpgradedCluster()) {
             // check old cluster model
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterId).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get("endpoints");
             assertEquals("hugging_face", configs.get(0).get("service"));
 
             // Inference on old cluster model
             assertEmbeddingInference(oldClusterId);
 
             embeddingsServer.enqueue(new MockResponse().setResponseCode(200).setBody(embeddingResponse()));
-            put(upgradedClusterId, embeddingConfig(getUrl(embeddingsServer)), TaskType.TEXT_EMBEDDING);
+            put(upgradedClusterId, embeddingConfig(getUrl(embeddingsServer)), testTaskType);
 
-            configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, upgradedClusterId).get("endpoints");
+            configs = (List<Map<String, Object>>) get(testTaskType, upgradedClusterId).get("endpoints");
             assertThat(configs, hasSize(1));
 
             assertEmbeddingInference(upgradedClusterId);
@@ -100,27 +106,29 @@ public class HuggingFaceServiceUpgradeIT extends InferenceUpgradeTestCase {
     }
 
     @SuppressWarnings("unchecked")
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/107887")
     public void testElser() throws IOException {
         var supported = getOldClusterTestVersion().onOrAfter(HF_ELSER_ADDED);
+        String old_cluster_endpoint_identifier = oldClusterHasFeature("gte_v" + MODELS_RENAMED_TO_ENDPOINTS) ? "endpoints" : "models";
         assumeTrue("HF elser service added in " + HF_ELSER_ADDED, supported);
 
         final String oldClusterId = "old-cluster-elser";
         final String upgradedClusterId = "upgraded-cluster-elser";
 
+        var testTaskType = TaskType.SPARSE_EMBEDDING;
+
         if (isOldCluster()) {
-            put(oldClusterId, elserConfig(getUrl(elserServer)), TaskType.SPARSE_EMBEDDING);
-            var configs = (List<Map<String, Object>>) get(TaskType.SPARSE_EMBEDDING, oldClusterId).get("endpoints");
+            put(oldClusterId, elserConfig(getUrl(elserServer)), testTaskType);
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get(old_cluster_endpoint_identifier);
             assertThat(configs, hasSize(1));
 
             assertElser(oldClusterId);
         } else if (isMixedCluster()) {
-            var configs = (List<Map<String, Object>>) get(TaskType.SPARSE_EMBEDDING, oldClusterId).get("endpoints");
+            var configs = getConfigsWithBreakingChangeHandling(testTaskType, oldClusterId);
             assertEquals("hugging_face", configs.get(0).get("service"));
             assertElser(oldClusterId);
         } else if (isUpgradedCluster()) {
             // check old cluster model
-            var configs = (List<Map<String, Object>>) get(TaskType.SPARSE_EMBEDDING, oldClusterId).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get("endpoints");
             assertEquals("hugging_face", configs.get(0).get("service"));
             var taskSettings = (Map<String, Object>) configs.get(0).get("task_settings");
             assertThat(taskSettings.keySet(), empty());
@@ -128,7 +136,7 @@ public class HuggingFaceServiceUpgradeIT extends InferenceUpgradeTestCase {
             assertElser(oldClusterId);
 
             // New endpoint
-            put(upgradedClusterId, elserConfig(getUrl(elserServer)), TaskType.SPARSE_EMBEDDING);
+            put(upgradedClusterId, elserConfig(getUrl(elserServer)), testTaskType);
             configs = (List<Map<String, Object>>) get(upgradedClusterId).get("endpoints");
             assertThat(configs, hasSize(1));
 

+ 17 - 0
x-pack/plugin/inference/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/xpack/application/InferenceUpgradeTestCase.java

@@ -16,13 +16,17 @@ import org.elasticsearch.test.http.MockWebServer;
 import org.elasticsearch.upgrades.AbstractRollingUpgradeTestCase;
 
 import java.io.IOException;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 import static org.elasticsearch.core.Strings.format;
 
 public class InferenceUpgradeTestCase extends AbstractRollingUpgradeTestCase {
 
+    static final String MODELS_RENAMED_TO_ENDPOINTS = "8.15.0";
+
     public InferenceUpgradeTestCase(@Name("upgradedNodes") int upgradedNodes) {
         super(upgradedNodes);
     }
@@ -104,4 +108,17 @@ public class InferenceUpgradeTestCase extends AbstractRollingUpgradeTestCase {
         var response = client().performRequest(request);
         assertOKAndConsume(response);
     }
+
+    @SuppressWarnings("unchecked")
+    // in version 8.15, there was a breaking change where "models" was renamed to "endpoints"
+    LinkedList<Map<String, Object>> getConfigsWithBreakingChangeHandling(TaskType testTaskType, String oldClusterId) throws IOException {
+
+        LinkedList<Map<String, Object>> configs;
+        configs = new LinkedList<>(
+            (List<Map<String, Object>>) Objects.requireNonNullElse((get(testTaskType, oldClusterId).get("endpoints")), List.of())
+        );
+        configs.addAll(Objects.requireNonNullElse((List<Map<String, Object>>) get(testTaskType, oldClusterId).get("models"), List.of()));
+
+        return configs;
+    }
 }

+ 28 - 16
x-pack/plugin/inference/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/xpack/application/OpenAiServiceUpgradeIT.java

@@ -12,6 +12,8 @@ import org.elasticsearch.common.Strings;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.test.http.MockResponse;
 import org.elasticsearch.test.http.MockWebServer;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
 
 import java.io.IOException;
 import java.util.List;
@@ -35,7 +37,7 @@ public class OpenAiServiceUpgradeIT extends InferenceUpgradeTestCase {
         super(upgradedNodes);
     }
 
-    // @BeforeClass
+    @BeforeClass
     public static void startWebServer() throws IOException {
         openAiEmbeddingsServer = new MockWebServer();
         openAiEmbeddingsServer.start();
@@ -44,33 +46,37 @@ public class OpenAiServiceUpgradeIT extends InferenceUpgradeTestCase {
         openAiChatCompletionsServer.start();
     }
 
-    // @AfterClass for the awaits fix
+    @AfterClass
     public static void shutdown() {
         openAiEmbeddingsServer.close();
         openAiChatCompletionsServer.close();
     }
 
     @SuppressWarnings("unchecked")
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/107887")
     public void testOpenAiEmbeddings() throws IOException {
         var openAiEmbeddingsSupported = getOldClusterTestVersion().onOrAfter(OPEN_AI_EMBEDDINGS_ADDED);
+        // `gte_v` indicates that the cluster version is Greater Than or Equal to MODELS_RENAMED_TO_ENDPOINTS
+        String oldClusterEndpointIdentifier = oldClusterHasFeature("gte_v" + MODELS_RENAMED_TO_ENDPOINTS) ? "endpoints" : "models";
         assumeTrue("OpenAI embedding service added in " + OPEN_AI_EMBEDDINGS_ADDED, openAiEmbeddingsSupported);
 
         final String oldClusterId = "old-cluster-embeddings";
         final String upgradedClusterId = "upgraded-cluster-embeddings";
 
+        var testTaskType = TaskType.TEXT_EMBEDDING;
+
         if (isOldCluster()) {
             String inferenceConfig = oldClusterVersionCompatibleEmbeddingConfig();
             // queue a response as PUT will call the service
             openAiEmbeddingsServer.enqueue(new MockResponse().setResponseCode(200).setBody(embeddingResponse()));
-            put(oldClusterId, inferenceConfig, TaskType.TEXT_EMBEDDING);
+            put(oldClusterId, inferenceConfig, testTaskType);
 
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterId).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get(oldClusterEndpointIdentifier);
             assertThat(configs, hasSize(1));
 
             assertEmbeddingInference(oldClusterId);
         } else if (isMixedCluster()) {
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterId).get("endpoints");
+            var configs = getConfigsWithBreakingChangeHandling(testTaskType, oldClusterId);
+
             assertEquals("openai", configs.get(0).get("service"));
             var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
             var taskSettings = (Map<String, Object>) configs.get(0).get("task_settings");
@@ -80,7 +86,7 @@ public class OpenAiServiceUpgradeIT extends InferenceUpgradeTestCase {
             assertEmbeddingInference(oldClusterId);
         } else if (isUpgradedCluster()) {
             // check old cluster model
-            var configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, oldClusterId).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get("endpoints");
             var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
             // model id is moved to service settings
             assertThat(serviceSettings, hasEntry("model_id", "text-embedding-ada-002"));
@@ -92,9 +98,9 @@ public class OpenAiServiceUpgradeIT extends InferenceUpgradeTestCase {
 
             String inferenceConfig = embeddingConfigWithModelInServiceSettings(getUrl(openAiEmbeddingsServer));
             openAiEmbeddingsServer.enqueue(new MockResponse().setResponseCode(200).setBody(embeddingResponse()));
-            put(upgradedClusterId, inferenceConfig, TaskType.TEXT_EMBEDDING);
+            put(upgradedClusterId, inferenceConfig, testTaskType);
 
-            configs = (List<Map<String, Object>>) get(TaskType.TEXT_EMBEDDING, upgradedClusterId).get("endpoints");
+            configs = (List<Map<String, Object>>) get(testTaskType, upgradedClusterId).get("endpoints");
             assertThat(configs, hasSize(1));
 
             assertEmbeddingInference(upgradedClusterId);
@@ -111,23 +117,29 @@ public class OpenAiServiceUpgradeIT extends InferenceUpgradeTestCase {
     }
 
     @SuppressWarnings("unchecked")
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/107887")
     public void testOpenAiCompletions() throws IOException {
         var openAiEmbeddingsSupported = getOldClusterTestVersion().onOrAfter(OPEN_AI_COMPLETIONS_ADDED);
+        String old_cluster_endpoint_identifier = oldClusterHasFeature("gte_v" + MODELS_RENAMED_TO_ENDPOINTS) ? "endpoints" : "models";
         assumeTrue("OpenAI completions service added in " + OPEN_AI_COMPLETIONS_ADDED, openAiEmbeddingsSupported);
 
         final String oldClusterId = "old-cluster-completions";
         final String upgradedClusterId = "upgraded-cluster-completions";
 
+        var testTaskType = TaskType.COMPLETION;
+
         if (isOldCluster()) {
-            put(oldClusterId, chatCompletionsConfig(getUrl(openAiChatCompletionsServer)), TaskType.COMPLETION);
+            put(oldClusterId, chatCompletionsConfig(getUrl(openAiChatCompletionsServer)), testTaskType);
 
-            var configs = (List<Map<String, Object>>) get(TaskType.COMPLETION, oldClusterId).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get(old_cluster_endpoint_identifier);
             assertThat(configs, hasSize(1));
 
             assertCompletionInference(oldClusterId);
         } else if (isMixedCluster()) {
-            var configs = (List<Map<String, Object>>) get(TaskType.COMPLETION, oldClusterId).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get("endpoints");
+            if (oldClusterHasFeature("gte_v" + MODELS_RENAMED_TO_ENDPOINTS) == false) {
+                configs.addAll((List<Map<String, Object>>) get(testTaskType, oldClusterId).get(old_cluster_endpoint_identifier));
+                // in version 8.15, there was a breaking change where "models" was renamed to "endpoints"
+            }
             assertEquals("openai", configs.get(0).get("service"));
             var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
             assertThat(serviceSettings, hasEntry("model_id", "gpt-4"));
@@ -137,7 +149,7 @@ public class OpenAiServiceUpgradeIT extends InferenceUpgradeTestCase {
             assertCompletionInference(oldClusterId);
         } else if (isUpgradedCluster()) {
             // check old cluster model
-            var configs = (List<Map<String, Object>>) get(TaskType.COMPLETION, oldClusterId).get("endpoints");
+            var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get("endpoints");
             var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
             assertThat(serviceSettings, hasEntry("model_id", "gpt-4"));
             var taskSettings = (Map<String, Object>) configs.get(0).get("task_settings");
@@ -145,8 +157,8 @@ public class OpenAiServiceUpgradeIT extends InferenceUpgradeTestCase {
 
             assertCompletionInference(oldClusterId);
 
-            put(upgradedClusterId, chatCompletionsConfig(getUrl(openAiChatCompletionsServer)), TaskType.COMPLETION);
-            configs = (List<Map<String, Object>>) get(TaskType.COMPLETION, upgradedClusterId).get("endpoints");
+            put(upgradedClusterId, chatCompletionsConfig(getUrl(openAiChatCompletionsServer)), testTaskType);
+            configs = (List<Map<String, Object>>) get(testTaskType, upgradedClusterId).get("endpoints");
             assertThat(configs, hasSize(1));
 
             // Inference on the new config

+ 1 - 1
x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/cohere/CohereRankedResponseEntity.java

@@ -161,5 +161,5 @@ public class CohereRankedResponseEntity {
 
     private CohereRankedResponseEntity() {}
 
-    static String FAILED_TO_FIND_FIELD_TEMPLATE = "Failed to find required field [%s] in Cohere embeddings response";
+    static String FAILED_TO_FIND_FIELD_TEMPLATE = "Failed to find required field [%s] in Cohere rerank response";
 }