Second attempt

droberts195 · Mar 19, 2024 · d093fb3 · d093fb3
1 parent 475e98e
commit d093fb3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/...src/main/java/org/elasticsearch/xpack/ml/action/TransportGetTrainedModelsStatsAction.java b/...src/main/java/org/elasticsearch/xpack/ml/action/TransportGetTrainedModelsStatsAction.java
@@ -39,7 +39,6 @@
 import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
-import org.elasticsearch.transport.Transports;
 import org.elasticsearch.xpack.core.action.util.ExpandedIdsMatcher;
 import org.elasticsearch.xpack.core.ml.action.GetDeploymentStatsAction;
 import org.elasticsearch.xpack.core.ml.action.GetTrainedModelsAction;
@@ -123,10 +122,12 @@ protected void doExecute(
 
         GetTrainedModelsStatsAction.Response.Builder responseBuilder = new GetTrainedModelsStatsAction.Response.Builder();
 
-        SubscribableListener.<Tuple<Long, Map<String, Set<String>>>>newForked(l -> {
+        SubscribableListener.<String>newForked(l -> {
             // When the request resource is a deployment find the
             // model used in that deployment for the model stats
             String idExpression = addModelsUsedInMatchingDeployments(request.getResourceId(), assignmentMetadata);
+            l.onResponse(idExpression);
+        }).<Tuple<Long, Map<String, Set<String>>>>andThen(executor, null, (l, idExpression) -> {
             logger.debug("Expanded models/deployment Ids request [{}]", idExpression);
 
             // the request id may contain deployment ids
@@ -208,7 +209,6 @@ protected void doExecute(
     }
 
     static String addModelsUsedInMatchingDeployments(String idExpression, TrainedModelAssignmentMetadata assignmentMetadata) {
-        assert Transports.assertNotTransportThread("non-trivial nested loops over cluster state structures");
         if (Strings.isAllOrWildcard(idExpression)) {
             return idExpression;
         } else {

diff --git a/...in/ml/src/main/java/org/elasticsearch/xpack/ml/utils/InferenceProcessorInfoExtractor.java b/...in/ml/src/main/java/org/elasticsearch/xpack/ml/utils/InferenceProcessorInfoExtractor.java
@@ -13,6 +13,7 @@
 import org.elasticsearch.ingest.ConfigurationUtils;
 import org.elasticsearch.ingest.IngestMetadata;
 import org.elasticsearch.ingest.Pipeline;
+import org.elasticsearch.transport.Transports;
 
 import java.util.HashMap;
 import java.util.LinkedHashSet;
@@ -75,6 +76,7 @@ public static int countInferenceProcessors(ClusterState state) {
      */
     @SuppressWarnings("unchecked")
     public static Map<String, Set<String>> pipelineIdsByResource(ClusterState state, Set<String> ids) {
+        assert Transports.assertNotTransportThread("non-trivial nested loops over cluster state structures");
         Map<String, Set<String>> pipelineIdsByModelIds = new HashMap<>();
         Metadata metadata = state.metadata();
         if (metadata == null) {