[ML] throw when definition is requested for pytorch models (#80310) (#80338)

Since pytorch models are not built in Elasticsearch, we don't need to provide supplying the definition when retrieving the trained model.

In fact, some of these definitions are so large, that returning them is prohibitive.

related to #80254

Author: benwtrent

x-pack/plugin/ml/qa/native-multi-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/PyTorchModelIT.java

@@ -417,6 +417,18 @@ public void testInferWithMissingModel() {
         assertThat(ex.getMessage(), containsString("Could not find trained model [missing_model]"));
     }
 
+    public void testGetPytorchModelWithDefinition() throws IOException {
+        String model = "should-fail-get";
+        createTrainedModel(model);
+        putVocabulary(List.of("once", "twice"), model);
+        putModelDefinition(model);
+        Exception ex = expectThrows(
+            Exception.class,
+            () -> client().performRequest(new Request("GET", "_ml/trained_models/" + model + "?include=definition"))
+        );
+        assertThat(ex.getMessage(), containsString("[should-fail-get] is type [pytorch] and does not support retrieving the definition"));
+    }
+
     public void testInferencePipelineAgainstUnallocatedModel() throws IOException {
         String model = "not-deployed";
         createTrainedModel(model);

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/persistence/TrainedModelProvider.java

@@ -24,9 +24,7 @@
 import org.elasticsearch.action.bulk.BulkResponse;
 import org.elasticsearch.action.index.IndexAction;
 import org.elasticsearch.action.index.IndexRequest;
-import org.elasticsearch.action.search.MultiSearchAction;
 import org.elasticsearch.action.search.MultiSearchRequest;
-import org.elasticsearch.action.search.MultiSearchRequestBuilder;
 import org.elasticsearch.action.search.MultiSearchResponse;
 import org.elasticsearch.action.search.SearchAction;
 import org.elasticsearch.action.search.SearchRequest;
@@ -602,26 +600,17 @@ public void getTrainedModel(
         }, finalListener::onFailure);
 
         QueryBuilder queryBuilder = QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds(modelId));
-        MultiSearchRequestBuilder multiSearchRequestBuilder = client.prepareMultiSearch()
-            .add(
-                client.prepareSearch(InferenceIndexConstants.INDEX_PATTERN)
-                    .setQuery(queryBuilder)
-                    // use sort to get the last
-                    .addSort("_index", SortOrder.DESC)
-                    .setSize(1)
-                    .request()
-            );
-
-        if (includes.isIncludeModelDefinition()) {
-            multiSearchRequestBuilder.add(
-                ChunkedTrainedModelRestorer.buildSearch(client, modelId, InferenceIndexConstants.INDEX_PATTERN, MAX_NUM_DEFINITION_DOCS)
-            );
-        }
+        SearchRequest trainedModelConfigSearch = client.prepareSearch(InferenceIndexConstants.INDEX_PATTERN)
+            .setQuery(queryBuilder)
+            // use sort to get the last
+            .addSort("_index", SortOrder.DESC)
+            .setSize(1)
+            .request();
 
-        ActionListener<MultiSearchResponse> multiSearchResponseActionListener = ActionListener.wrap(multiSearchResponse -> {
+        ActionListener<SearchResponse> trainedModelSearchHandler = ActionListener.wrap(modelSearchResponse -> {
             TrainedModelConfig.Builder builder;
             try {
-                builder = handleSearchItem(multiSearchResponse.getResponses()[0], modelId, this::parseModelConfigLenientlyFromSource);
+                builder = handleHits(modelSearchResponse.getHits().getHits(), modelId, this::parseModelConfigLenientlyFromSource).get(0);
             } catch (ResourceNotFoundException ex) {
                 getTrainedModelListener.onFailure(
                     new ResourceNotFoundException(Messages.getMessage(Messages.INFERENCE_NOT_FOUND, modelId))
@@ -631,46 +620,58 @@ public void getTrainedModel(
                 getTrainedModelListener.onFailure(ex);
                 return;
             }
-
-            if (includes.isIncludeModelDefinition()) {
-                try {
-                    List<TrainedModelDefinitionDoc> docs = handleSearchItems(
-                        multiSearchResponse.getResponses()[1],
+            if (includes.isIncludeModelDefinition() == false) {
+                getTrainedModelListener.onResponse(builder);
+                return;
+            }
+            if (builder.getModelType() == TrainedModelType.PYTORCH && includes.isIncludeModelDefinition()) {
+                finalListener.onFailure(
+                    ExceptionsHelper.badRequestException(
+                        "[{}] is type [{}] and does not support retrieving the definition",
                         modelId,
-                        (bytes, resourceId) -> ChunkedTrainedModelRestorer.parseModelDefinitionDocLenientlyFromSource(
-                            bytes,
-                            resourceId,
-                            xContentRegistry
-                        )
-                    );
+                        builder.getModelType()
+                    )
+                );
+                return;
+            }
+            executeAsyncWithOrigin(
+                client,
+                ML_ORIGIN,
+                SearchAction.INSTANCE,
+                ChunkedTrainedModelRestorer.buildSearch(client, modelId, InferenceIndexConstants.INDEX_PATTERN, MAX_NUM_DEFINITION_DOCS),
+                ActionListener.wrap(definitionSearchResponse -> {
                     try {
-                        BytesReference compressedData = getDefinitionFromDocs(docs, modelId);
-                        builder.setDefinitionFromBytes(compressedData);
-                    } catch (ElasticsearchException elasticsearchException) {
-                        getTrainedModelListener.onFailure(elasticsearchException);
+                        List<TrainedModelDefinitionDoc> docs = handleHits(
+                            definitionSearchResponse.getHits().getHits(),
+                            modelId,
+                            (bytes, resourceId) -> ChunkedTrainedModelRestorer.parseModelDefinitionDocLenientlyFromSource(
+                                bytes,
+                                resourceId,
+                                xContentRegistry
+                            )
+                        );
+                        try {
+                            BytesReference compressedData = getDefinitionFromDocs(docs, modelId);
+                            builder.setDefinitionFromBytes(compressedData);
+                        } catch (ElasticsearchException elasticsearchException) {
+                            getTrainedModelListener.onFailure(elasticsearchException);
+                            return;
+                        }
+
+                    } catch (ResourceNotFoundException ex) {
+                        getTrainedModelListener.onFailure(
+                            new ResourceNotFoundException(Messages.getMessage(Messages.MODEL_DEFINITION_NOT_FOUND, modelId))
+                        );
+                        return;
+                    } catch (Exception ex) {
+                        getTrainedModelListener.onFailure(ex);
                         return;
                     }
-
-                } catch (ResourceNotFoundException ex) {
-                    getTrainedModelListener.onFailure(
-                        new ResourceNotFoundException(Messages.getMessage(Messages.MODEL_DEFINITION_NOT_FOUND, modelId))
-                    );
-                    return;
-                } catch (Exception ex) {
-                    getTrainedModelListener.onFailure(ex);
-                    return;
-                }
-            }
-            getTrainedModelListener.onResponse(builder);
+                    getTrainedModelListener.onResponse(builder);
+                }, getTrainedModelListener::onFailure)
+            );
         }, getTrainedModelListener::onFailure);
-
-        executeAsyncWithOrigin(
-            client,
-            ML_ORIGIN,
-            MultiSearchAction.INSTANCE,
-            multiSearchRequestBuilder.request(),
-            multiSearchResponseActionListener
-        );
+        executeAsyncWithOrigin(client, ML_ORIGIN, SearchAction.INSTANCE, trainedModelConfigSearch, trainedModelSearchHandler);
     }
 
     public void getTrainedModels(
@@ -1204,6 +1205,9 @@ private static <T> List<T> handleHits(
         String resourceId,
         CheckedBiFunction<BytesReference, String, T, Exception> parseLeniently
     ) throws Exception {
+        if (hits.length == 0) {
+            throw new ResourceNotFoundException(resourceId);
+        }
         List<T> results = new ArrayList<>(hits.length);
         String initialIndex = hits[0].getIndex();
         for (SearchHit hit : hits) {