From 5b52728f1243be1a7ab8f3fb95fa24d2d2ce0c10 Mon Sep 17 00:00:00 2001
From: Jan Kuipers <148754765+jan-elastic@users.noreply.github.com>
Date: Fri, 4 Apr 2025 16:50:06 +0200
Subject: [PATCH] Fix deploying custom models with adaptive allocations
 (#126276)

---
 ...andInternalTextEmbeddingServiceSettings.java |  2 +-
 ...ternalTextEmbeddingServiceSettingsTests.java | 17 +++++++++--------
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/CustomElandInternalTextEmbeddingServiceSettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/CustomElandInternalTextEmbeddingServiceSettings.java
index d5f4143b65d36..6395cd36457c2 100644
--- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/CustomElandInternalTextEmbeddingServiceSettings.java
+++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/CustomElandInternalTextEmbeddingServiceSettings.java
@@ -107,7 +107,7 @@ private static CommonFields commonFieldsFromMap(Map<String, Object> map, Validat
     private final DenseVectorFieldMapper.ElementType elementType;
 
     CustomElandInternalTextEmbeddingServiceSettings(
-        int numAllocations,
+        Integer numAllocations,
         int numThreads,
         String modelId,
         AdaptiveAllocationsSettings adaptiveAllocationsSettings,
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/CustomElandInternalTextEmbeddingServiceSettingsTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/CustomElandInternalTextEmbeddingServiceSettingsTests.java
index 3b8ce3a7cc64c..c9eccd02b771b 100644
--- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/CustomElandInternalTextEmbeddingServiceSettingsTests.java
+++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/CustomElandInternalTextEmbeddingServiceSettingsTests.java
@@ -15,6 +15,7 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
 import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
 import org.elasticsearch.xpack.inference.services.ServiceFields;
 
@@ -31,23 +32,23 @@ public class CustomElandInternalTextEmbeddingServiceSettingsTests extends Abstra
     CustomElandInternalTextEmbeddingServiceSettings> {
 
     public static CustomElandInternalTextEmbeddingServiceSettings createRandom() {
-        var numAllocations = randomIntBetween(1, 10);
+        var withAdaptiveAllocations = randomBoolean();
+        var numAllocations = withAdaptiveAllocations ? null : randomIntBetween(1, 10);
+        var adaptiveAllocationsSettings = withAdaptiveAllocations
+            ? new AdaptiveAllocationsSettings(true, randomIntBetween(0, 2), randomIntBetween(2, 5))
+            : null;
         var numThreads = randomIntBetween(1, 10);
         var modelId = randomAlphaOfLength(8);
-        SimilarityMeasure similarityMeasure = SimilarityMeasure.COSINE;
-        Integer dims = null;
+        var similarityMeasure = SimilarityMeasure.COSINE;
         var setDimensions = randomBoolean();
-        if (setDimensions) {
-            dims = 123;
-        }
-
+        var dims = setDimensions ? 123 : null;
         var elementType = randomFrom(DenseVectorFieldMapper.ElementType.values());
 
         return new CustomElandInternalTextEmbeddingServiceSettings(
             numAllocations,
             numThreads,
             modelId,
-            null,
+            adaptiveAllocationsSettings,
             null,
             dims,
             similarityMeasure,