Moving back InferenceRunnerConfig to BulkInferenceRunnerConfig

Author: afoucret

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -12,6 +12,7 @@
 import org.apache.lucene.sandbox.document.HalfFloatPoint;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.RemoteException;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.project.ProjectResolver;
@@ -75,8 +76,7 @@
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.NotEquals;
 import org.elasticsearch.xpack.esql.index.EsIndex;
 import org.elasticsearch.xpack.esql.inference.InferenceResolution;
-import org.elasticsearch.xpack.esql.inference.InferenceResolver;
-import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunner;
+import org.elasticsearch.xpack.esql.inference.InferenceService;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.parser.QueryParam;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
@@ -420,8 +420,7 @@ public static LogicalOptimizerContext unboundLogicalOptimizerContext() {
         mock(ProjectResolver.class),
         mock(IndexNameExpressionResolver.class),
         null,
-        mock(InferenceResolver.Factory.class),
-        mock(BulkInferenceRunner.Factory.class)
+        new InferenceService(mock(Client.class))
     );
 
     private EsqlTestUtils() {}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceService.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.inference;
+
+import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunner;
+
+public class InferenceService {
+    private final InferenceResolver.Factory inferenceResolverFactory;
+
+    private final BulkInferenceRunner.Factory bulkInferenceRunnerFactory;
+
+    /**
+     * Creates a new inference service with the given client.
+     *
+     * @param client the Elasticsearch client for inference operations
+     */
+    public InferenceService(Client client) {
+        this(InferenceResolver.factory(client), BulkInferenceRunner.factory(client));
+    }
+
+    private InferenceService(InferenceResolver.Factory inferenceResolverFactory, BulkInferenceRunner.Factory bulkInferenceRunnerFactory) {
+        this.inferenceResolverFactory = inferenceResolverFactory;
+        this.bulkInferenceRunnerFactory = bulkInferenceRunnerFactory;
+    }
+
+    /**
+     * Creates an inference resolver for resolving inference IDs in logical plans.
+     *
+     * @return a new inference resolver instance
+     */
+    public InferenceResolver inferenceResolver() {
+        return inferenceResolverFactory.create();
+    }
+
+    /**
+     * Returns the bulk inference runner factory for executing batched inference operations.
+     *
+     * @return the bulk inference runner factory
+     */
+    public BulkInferenceRunner.Factory bulkInferenceRunnerFactory() {
+        return bulkInferenceRunnerFactory;
+    }
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java

@@ -86,8 +86,8 @@
 import org.elasticsearch.xpack.esql.evaluator.EvalMapper;
 import org.elasticsearch.xpack.esql.evaluator.command.GrokEvaluatorExtracter;
 import org.elasticsearch.xpack.esql.expression.Order;
+import org.elasticsearch.xpack.esql.inference.InferenceService;
 import org.elasticsearch.xpack.esql.inference.XContentRowEncoder;
-import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunner;
 import org.elasticsearch.xpack.esql.inference.completion.CompletionOperator;
 import org.elasticsearch.xpack.esql.inference.rerank.RerankOperator;
 import org.elasticsearch.xpack.esql.plan.logical.Fork;
@@ -159,7 +159,7 @@ public class LocalExecutionPlanner {
     private final Supplier<ExchangeSink> exchangeSinkSupplier;
     private final EnrichLookupService enrichLookupService;
     private final LookupFromIndexService lookupFromIndexService;
-    private final BulkInferenceRunner.Factory bulkInferenceRunnerFactory;
+    private final InferenceService inferenceService;
     private final PhysicalOperationProviders physicalOperationProviders;
     private final List<ShardContext> shardContexts;
 
@@ -175,7 +175,7 @@ public LocalExecutionPlanner(
         Supplier<ExchangeSink> exchangeSinkSupplier,
         EnrichLookupService enrichLookupService,
         LookupFromIndexService lookupFromIndexService,
-        BulkInferenceRunner.Factory bulkInferenceRunnerFactory,
+        InferenceService inferenceService,
         PhysicalOperationProviders physicalOperationProviders,
         List<ShardContext> shardContexts
     ) {
@@ -191,7 +191,7 @@ public LocalExecutionPlanner(
         this.exchangeSinkSupplier = exchangeSinkSupplier;
         this.enrichLookupService = enrichLookupService;
         this.lookupFromIndexService = lookupFromIndexService;
-        this.bulkInferenceRunnerFactory = bulkInferenceRunnerFactory;
+        this.inferenceService = inferenceService;
         this.physicalOperationProviders = physicalOperationProviders;
         this.shardContexts = shardContexts;
     }
@@ -318,7 +318,10 @@ private PhysicalOperation planCompletion(CompletionExec completion, LocalExecuti
             source.layout
         );
 
-        return source.with(new CompletionOperator.Factory(bulkInferenceRunnerFactory, inferenceId, promptEvaluatorFactory), outputLayout);
+        return source.with(
+            new CompletionOperator.Factory(inferenceService.bulkInferenceRunnerFactory(), inferenceId, promptEvaluatorFactory),
+            outputLayout
+        );
     }
 
     private PhysicalOperation planRrfScoreEvalExec(RrfScoreEvalExec rrf, LocalExecutionPlannerContext context) {
@@ -654,7 +657,13 @@ private PhysicalOperation planRerank(RerankExec rerank, LocalExecutionPlannerCon
         int scoreChannel = outputLayout.get(rerank.scoreAttribute().id()).channel();
 
         return source.with(
-            new RerankOperator.Factory(bulkInferenceRunnerFactory, inferenceId, queryText, rowEncoderFactory, scoreChannel),
+            new RerankOperator.Factory(
+                inferenceService.bulkInferenceRunnerFactory(),
+                inferenceId,
+                queryText,
+                rowEncoderFactory,
+                scoreChannel
+            ),
             outputLayout
         );
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java

@@ -52,7 +52,7 @@
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.enrich.EnrichLookupService;
 import org.elasticsearch.xpack.esql.enrich.LookupFromIndexService;
-import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunner;
+import org.elasticsearch.xpack.esql.inference.InferenceService;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeSinkExec;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeSourceExec;
 import org.elasticsearch.xpack.esql.plan.physical.OutputExec;
@@ -131,7 +131,7 @@ public class ComputeService {
     private final DriverTaskRunner driverRunner;
     private final EnrichLookupService enrichLookupService;
     private final LookupFromIndexService lookupFromIndexService;
-    private final BulkInferenceRunner.Factory bulkInferenceRunnerFactory;
+    private final InferenceService inferenceService;
     private final ClusterService clusterService;
     private final ProjectResolver projectResolver;
     private final AtomicLong childSessionIdGenerator = new AtomicLong();
@@ -158,7 +158,7 @@ public ComputeService(
         this.driverRunner = new DriverTaskRunner(transportService, esqlExecutor);
         this.enrichLookupService = enrichLookupService;
         this.lookupFromIndexService = lookupFromIndexService;
-        this.bulkInferenceRunnerFactory = transportActionServices.bulkInferenceRunnerFactory();
+        this.inferenceService = transportActionServices.inferenceService();
         this.clusterService = transportActionServices.clusterService();
         this.projectResolver = transportActionServices.projectResolver();
         this.dataNodeComputeHandler = new DataNodeComputeHandler(
@@ -622,7 +622,7 @@ public SourceProvider createSourceProvider() {
                 context.exchangeSinkSupplier(),
                 enrichLookupService,
                 lookupFromIndexService,
-                bulkInferenceRunnerFactory,
+                inferenceService,
                 physicalOperationProviders,
                 contexts
             );

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportActionServices.java

@@ -14,105 +14,15 @@
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.usage.UsageService;
-import org.elasticsearch.xpack.esql.inference.InferenceResolver;
-import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunner;
-import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunnerConfig;
-
-public class TransportActionServices {
-    private final TransportService transportService;
-    private final SearchService searchService;
-    private final ExchangeService exchangeService;
-    private final ClusterService clusterService;
-    private final ProjectResolver projectResolver;
-    private final IndexNameExpressionResolver indexNameExpressionResolver;
-    private final UsageService usageService;
-    private final InferenceResolver.Factory inferenceResolverFactory;
-    private final BulkInferenceRunner.Factory inferenceRunnerFactory;
-
-    public TransportActionServices(
-        TransportService transportService,
-        SearchService searchService,
-        ExchangeService exchangeService,
-        ClusterService clusterService,
-        ProjectResolver projectResolver,
-        IndexNameExpressionResolver indexNameExpressionResolver,
-        UsageService usageService,
-        InferenceResolver.Factory inferenceResolverFactory,
-        BulkInferenceRunner.Factory bulkInferenceRunnerFactory
-    ) {
-        this.transportService = transportService;
-        this.searchService = searchService;
-        this.exchangeService = exchangeService;
-        this.clusterService = clusterService;
-        this.projectResolver = projectResolver;
-        this.indexNameExpressionResolver = indexNameExpressionResolver;
-        this.usageService = usageService;
-        this.inferenceResolverFactory = inferenceResolverFactory;
-        this.inferenceRunnerFactory = bulkInferenceRunnerFactory;
-    }
-
-    public TransportService transportService() {
-        return transportService;
-    }
-
-    public SearchService searchService() {
-        return searchService;
-    }
-
-    public ExchangeService exchangeService() {
-        return exchangeService;
-    }
-
-    public ClusterService clusterService() {
-        return clusterService;
-    }
-
-    public ProjectResolver projectResolver() {
-        return projectResolver;
-    }
-
-    public IndexNameExpressionResolver indexNameExpressionResolver() {
-        return indexNameExpressionResolver;
-    }
-
-    public UsageService usageService() {
-        return usageService;
-    }
-
-    /**
-     * Creates an inference resolver for collecting used inference IDs from logical plans.
-     *
-     * @return A resolver capable of extracting inference IDs from plans and expressions
-     */
-    public InferenceResolver inferenceResolver() {
-        return inferenceResolverFactory.create();
-    }
-
-    /**
-     * Returns the inference runner factory for creating multiple configured runners.
-     *
-     * @return The factory for creating inference runners
-     */
-    public BulkInferenceRunner.Factory bulkInferenceRunnerFactory() {
-        return inferenceRunnerFactory;
-    }
-
-    /**
-     * Creates an inference runner with the specified execution configuration using the default configuration.
-     *
-     * @return A configured inference runner capable of executing inference requests
-     */
-    public BulkInferenceRunner bulkInferenceRunner() {
-        return bulkInferenceRunner(BulkInferenceRunnerConfig.DEFAULT);
-    }
-
-    /**
-     * Creates an inference runner with the specified execution configuration.
-     *
-     * @param bulkInferenceRunnerConfig Configuration specifying concurrency limits and execution parameters
-     * @return A configured inference runner capable of executing inference requests
-     */
-    public BulkInferenceRunner bulkInferenceRunner(BulkInferenceRunnerConfig bulkInferenceRunnerConfig) {
-        return inferenceRunnerFactory.create(bulkInferenceRunnerConfig);
-    }
-}
+import org.elasticsearch.xpack.esql.inference.InferenceService;
+
+public record TransportActionServices(
+    TransportService transportService,
+    SearchService searchService,
+    ExchangeService exchangeService,
+    ClusterService clusterService,
+    ProjectResolver projectResolver,
+    IndexNameExpressionResolver indexNameExpressionResolver,
+    UsageService usageService,
+    InferenceService inferenceService
+) {}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java

@@ -51,8 +51,7 @@
 import org.elasticsearch.xpack.esql.enrich.LookupFromIndexService;
 import org.elasticsearch.xpack.esql.execution.PlanExecutor;
 import org.elasticsearch.xpack.esql.expression.function.UnsupportedAttribute;
-import org.elasticsearch.xpack.esql.inference.InferenceResolver;
-import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunner;
+import org.elasticsearch.xpack.esql.inference.InferenceService;
 import org.elasticsearch.xpack.esql.session.Configuration;
 import org.elasticsearch.xpack.esql.session.EsqlSession.PlanRunner;
 import org.elasticsearch.xpack.esql.session.Result;
@@ -167,8 +166,7 @@ public TransportEsqlQueryAction(
             projectResolver,
             indexNameExpressionResolver,
             usageService,
-            InferenceResolver.factory(client),
-            BulkInferenceRunner.factory(client)
+            new InferenceService(client)
         );
 
         this.computeService = new ComputeService(

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java

@@ -54,8 +54,7 @@
 import org.elasticsearch.xpack.esql.index.IndexResolution;
 import org.elasticsearch.xpack.esql.index.MappingException;
 import org.elasticsearch.xpack.esql.inference.InferenceResolution;
-import org.elasticsearch.xpack.esql.inference.InferenceResolver;
-import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunner;
+import org.elasticsearch.xpack.esql.inference.InferenceService;
 import org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizer;
 import org.elasticsearch.xpack.esql.optimizer.LogicalPlanPreOptimizer;
 import org.elasticsearch.xpack.esql.optimizer.PhysicalOptimizerContext;
@@ -119,8 +118,7 @@ public interface PlanRunner {
     private final PhysicalPlanOptimizer physicalPlanOptimizer;
     private final PlanTelemetry planTelemetry;
     private final IndicesExpressionGrouper indicesExpressionGrouper;
-    private final BulkInferenceRunner.Factory inferenceRunnerFactory;
-    private final InferenceResolver inferenceResolver;
+    private final InferenceService inferenceService;
     private final RemoteClusterService remoteClusterService;
 
     private boolean explainMode;
@@ -155,8 +153,7 @@ public EsqlSession(
         this.physicalPlanOptimizer = new PhysicalPlanOptimizer(new PhysicalOptimizerContext(configuration));
         this.planTelemetry = planTelemetry;
         this.indicesExpressionGrouper = indicesExpressionGrouper;
-        this.inferenceRunnerFactory = services.bulkInferenceRunnerFactory();
-        this.inferenceResolver = services.inferenceResolver();
+        this.inferenceService = services.inferenceService();
         this.preMapper = new PreMapper(services);
         this.remoteClusterService = services.transportService().getRemoteClusterService();
     }
@@ -749,7 +746,7 @@ private static void resolveFieldNames(LogicalPlan parsed, EnrichResolution enric
     }
 
     private void resolveInferences(LogicalPlan plan, PreAnalysisResult preAnalysisResult, ActionListener<PreAnalysisResult> l) {
-        inferenceResolver.resolveInferenceIds(plan, l.map(preAnalysisResult::withInferenceResolution));
+        inferenceService.inferenceResolver().resolveInferenceIds(plan, l.map(preAnalysisResult::withInferenceResolution));
     }
 
     private PhysicalPlan logicalPlanToPhysicalPlan(LogicalPlan optimizedPlan, EsqlQueryRequest request) {

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/CsvTests.java

@@ -67,7 +67,7 @@
 import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
 import org.elasticsearch.xpack.esql.index.EsIndex;
 import org.elasticsearch.xpack.esql.index.IndexResolution;
-import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunner;
+import org.elasticsearch.xpack.esql.inference.InferenceService;
 import org.elasticsearch.xpack.esql.optimizer.LocalLogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.optimizer.LocalLogicalPlanOptimizer;
 import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalOptimizerContext;
@@ -707,7 +707,7 @@ void executeSubPlan(
             () -> exchangeSink.createExchangeSink(() -> {}),
             mock(EnrichLookupService.class),
             mock(LookupFromIndexService.class),
-            mock(BulkInferenceRunner.Factory.class),
+            mock(InferenceService.class),
             physicalOperationProviders,
             List.of()
         );