set location for packaged model correctly (#95399)

set the location for packaged models correctly, previously it used the default .ml-inference-... due to the order the fields get set. This change sets the config according to the model type correctly, with the option to use a dedicated index for packaged models in future.

Author: Hendrik Muhs

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ModelPackageConfigTests.java

@@ -8,15 +8,23 @@
 package org.elasticsearch.xpack.core.ml.inference.trainedmodel;
 
 import org.elasticsearch.TransportVersion;
+import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.xcontent.XContentHelper;
+import org.elasticsearch.xcontent.ToXContent;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.core.ml.AbstractBWCSerializationTestCase;
+import org.elasticsearch.xpack.core.ml.inference.TrainedModelType;
 
 import java.io.IOException;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Map;
 
 public class ModelPackageConfigTests extends AbstractBWCSerializationTestCase<ModelPackageConfig> {
 
@@ -31,9 +39,9 @@ public static ModelPackageConfig randomModulePackageConfig() {
                 : null,
             randomLongBetween(0, Long.MAX_VALUE - 100),
             randomBoolean() ? randomAlphaOfLength(10) : null,
+            randomInferenceConfigAsMap(),
             randomBoolean() ? Collections.singletonMap(randomAlphaOfLength(10), randomAlphaOfLength(10)) : null,
-            randomBoolean() ? Collections.singletonMap(randomAlphaOfLength(10), randomAlphaOfLength(10)) : null,
-            randomBoolean() ? randomAlphaOfLength(10) : null,
+            randomFrom(TrainedModelType.values()).toString(),
             randomBoolean() ? Arrays.asList(generateRandomStringArray(randomIntBetween(0, 5), 15, false)) : null,
             randomBoolean() ? randomAlphaOfLength(10) : null
         );
@@ -104,4 +112,30 @@ protected ModelPackageConfig mutateInstance(ModelPackageConfig instance) {
     protected ModelPackageConfig mutateInstanceForVersion(ModelPackageConfig instance, TransportVersion version) {
         return instance;
     }
+
+    private static Map<String, Object> randomInferenceConfigAsMap() {
+        InferenceConfig inferenceConfig = randomFrom(
+            new InferenceConfig[] {
+                ClassificationConfigTests.randomClassificationConfig(),
+                RegressionConfigTests.randomRegressionConfig(),
+                NerConfigTests.createRandom(),
+                PassThroughConfigTests.createRandom(),
+                TextClassificationConfigTests.createRandom(),
+                FillMaskConfigTests.createRandom(),
+                TextEmbeddingConfigTests.createRandom(),
+                QuestionAnsweringConfigTests.createRandom(),
+                TextSimilarityConfigTests.createRandom(),
+                TextExpansionConfigTests.createRandom() }
+        );
+
+        try (XContentBuilder xContentBuilder = XContentFactory.jsonBuilder()) {
+            XContentBuilder content = inferenceConfig.toXContent(xContentBuilder, ToXContent.EMPTY_PARAMS);
+            return Collections.singletonMap(
+                inferenceConfig.getWriteableName(),
+                XContentHelper.convertToMap(BytesReference.bytes(content), true, XContentType.JSON).v2()
+            );
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportPutTrainedModelAction.java

@@ -313,7 +313,7 @@ protected void masterOperation(
                     return;
                 }
                 modelPackageConfigHolder.set(resolvedModelPackageConfig);
-                setTrainedModelConfigFieldsFromPackagedModel(trainedModelConfig, resolvedModelPackageConfig);
+                setTrainedModelConfigFieldsFromPackagedModel(trainedModelConfig, resolvedModelPackageConfig, xContentRegistry);
 
                 checkModelIdAgainstTags(trainedModelConfig.getModelId(), modelIdTagCheckListener);
             }, listener::onFailure));
@@ -322,26 +322,6 @@ protected void masterOperation(
         }
     }
 
-    private void setTrainedModelConfigFieldsFromPackagedModel(
-        TrainedModelConfig.Builder trainedModelConfig,
-        ModelPackageConfig resolvedModelPackageConfig
-    ) throws IOException {
-        trainedModelConfig.setDescription(resolvedModelPackageConfig.getDescription());
-        trainedModelConfig.setModelType(TrainedModelType.fromString(resolvedModelPackageConfig.getModelType()));
-        trainedModelConfig.setMetadata(resolvedModelPackageConfig.getMetadata());
-        trainedModelConfig.setInferenceConfig(
-            parseInferenceConfigFromModelPackage(
-                resolvedModelPackageConfig.getInferenceConfigSource(),
-                xContentRegistry,
-                LoggingDeprecationHandler.INSTANCE
-            )
-        );
-        trainedModelConfig.setTags(resolvedModelPackageConfig.getTags());
-        trainedModelConfig.setModelPackageConfig(
-            new ModelPackageConfig.Builder(resolvedModelPackageConfig).resetPackageOnlyFields().build()
-        );
-    }
-
     private void triggerModelFetchIfNecessary(
         String modelId,
         ModelPackageConfig modelPackageConfig,
@@ -435,6 +415,29 @@ protected void doExecute(Task task, Request request, ActionListener<Response> li
         }
     }
 
+    static void setTrainedModelConfigFieldsFromPackagedModel(
+        TrainedModelConfig.Builder trainedModelConfig,
+        ModelPackageConfig resolvedModelPackageConfig,
+        NamedXContentRegistry xContentRegistry
+    ) throws IOException {
+        trainedModelConfig.setDescription(resolvedModelPackageConfig.getDescription());
+        trainedModelConfig.setModelType(TrainedModelType.fromString(resolvedModelPackageConfig.getModelType()));
+        trainedModelConfig.setMetadata(resolvedModelPackageConfig.getMetadata());
+        trainedModelConfig.setInferenceConfig(
+            parseInferenceConfigFromModelPackage(
+                resolvedModelPackageConfig.getInferenceConfigSource(),
+                xContentRegistry,
+                LoggingDeprecationHandler.INSTANCE
+            )
+        );
+        trainedModelConfig.setTags(resolvedModelPackageConfig.getTags());
+        trainedModelConfig.setModelPackageConfig(
+            new ModelPackageConfig.Builder(resolvedModelPackageConfig).resetPackageOnlyFields().build()
+        );
+
+        trainedModelConfig.setLocation(trainedModelConfig.getModelType().getDefaultLocation(trainedModelConfig.getModelId()));
+    }
+
     static InferenceConfig parseInferenceConfigFromModelPackage(
         Map<String, Object> source,
         NamedXContentRegistry namedXContentRegistry,

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportPutTrainedModelActionTests.java

@@ -17,9 +17,14 @@
 import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.core.ml.inference.MlInferenceNamedXContentProvider;
+import org.elasticsearch.xpack.core.ml.inference.TrainedModelConfig;
+import org.elasticsearch.xpack.core.ml.inference.TrainedModelInputTests;
+import org.elasticsearch.xpack.core.ml.inference.TrainedModelType;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.ClassificationConfigTests;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.FillMaskConfigTests;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.InferenceConfig;
+import org.elasticsearch.xpack.core.ml.inference.trainedmodel.ModelPackageConfig;
+import org.elasticsearch.xpack.core.ml.inference.trainedmodel.ModelPackageConfigTests;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.NerConfigTests;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.PassThroughConfigTests;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.QuestionAnsweringConfigTests;
@@ -67,6 +72,35 @@ public void testParseInferenceConfigFromModelPackage() throws IOException {
         assertEquals(inferenceConfig, parsedInferenceConfig);
     }
 
+    public void testSetTrainedModelConfigFieldsFromPackagedModel() throws IOException {
+        ModelPackageConfig packageConfig = ModelPackageConfigTests.randomModulePackageConfig();
+
+        TrainedModelConfig.Builder trainedModelConfigBuilder = new TrainedModelConfig.Builder().setModelId(
+            "." + packageConfig.getPackagedModelId()
+        ).setInput(TrainedModelInputTests.createRandomInput());
+
+        TransportPutTrainedModelAction.setTrainedModelConfigFieldsFromPackagedModel(
+            trainedModelConfigBuilder,
+            packageConfig,
+            xContentRegistry()
+        );
+
+        TrainedModelConfig trainedModelConfig = trainedModelConfigBuilder.build();
+
+        assertEquals(packageConfig.getModelType(), trainedModelConfig.getModelType().toString());
+        assertEquals(packageConfig.getDescription(), trainedModelConfig.getDescription());
+        assertEquals(packageConfig.getMetadata(), trainedModelConfig.getMetadata());
+        assertEquals(packageConfig.getTags(), trainedModelConfig.getTags());
+
+        // fully tested in {@link #testParseInferenceConfigFromModelPackage}
+        assertNotNull(trainedModelConfig.getInferenceConfig());
+
+        assertEquals(
+            TrainedModelType.fromString(packageConfig.getModelType()).getDefaultLocation(trainedModelConfig.getModelId()),
+            trainedModelConfig.getLocation()
+        );
+    }
+
     @Override
     protected NamedXContentRegistry xContentRegistry() {
         return new NamedXContentRegistry(new MlInferenceNamedXContentProvider().getNamedXContentParsers());