inference runner

henningandersen · May 27, 2024 · 7205c23 · 7205c23
1 parent ce01fd9
commit 7205c23
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/...ugin/ml/src/main/java/org/elasticsearch/xpack/ml/dataframe/inference/InferenceRunner.java b/...ugin/ml/src/main/java/org/elasticsearch/xpack/ml/dataframe/inference/InferenceRunner.java
@@ -17,6 +17,7 @@
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.action.support.UnsafePlainActionFuture;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.client.internal.OriginSettingClient;
 import org.elasticsearch.common.settings.Settings;
@@ -31,6 +32,7 @@
 import org.elasticsearch.xpack.core.ClientHelper;
 import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsConfig;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
+import org.elasticsearch.xpack.ml.MachineLearning;
 import org.elasticsearch.xpack.ml.dataframe.DestinationIndex;
 import org.elasticsearch.xpack.ml.dataframe.stats.DataCountsTracker;
 import org.elasticsearch.xpack.ml.dataframe.stats.ProgressTracker;
@@ -100,7 +102,9 @@ public void run(String modelId) {
 
         LOGGER.info("[{}] Started inference on test data against model [{}]", config.getId(), modelId);
         try {
-            PlainActionFuture<LocalModel> localModelPlainActionFuture = new PlainActionFuture<>();
+            PlainActionFuture<LocalModel> localModelPlainActionFuture = new UnsafePlainActionFuture<>(
+                MachineLearning.UTILITY_THREAD_POOL_NAME
+            );
             modelLoadingService.getModelForInternalInference(modelId, localModelPlainActionFuture);
             InferenceState inferenceState = restoreInferenceState();
             dataCountsTracker.setTestDocsCount(inferenceState.processedTestDocsCount);