ES|QL Inference runner refactoring

[afoucret]

This PR implements a batch of evolution in the ES|QL inference runner that are required in the context of the TEXT_EMBEDDING function implementation.

Changes:

• Full rewrite of the tasks runner to get rid of a race condition
• Improve the way the inference resolution is working:
  • More flexible implementation, so it will be easier to resolve inference ids from plans and functions
  • Reduced footprint in the PreAnalyzer
• InferenceOperator simplification
  • Remove the extra step in addInput
  • Simplification of the output building
  • Better test coverages (using randomized input/output columns)

Related Issue:

• ES|QL: Add TEXT_EMBEDDING function #131022

[elasticsearchmachine]

Pinging @elastic/es-analytical-engine (Team:Analytics)

[afoucret]

ℹ️ Inference resolution is moved after function resolution, so we can resolve inference ids use in functions like TEXT_EMBEDDING

[afoucret]

ℹ️ Now a separate class as the logic of inference resolution is becoming much more complex.

[tteofili]

nice :)

[afoucret]

ℹ️ This is a full rewrite of the InferenceRunner that was causing race conditions.

[ioanatia]

There is a race condition mentioned in the PR description.
What is the race condition?
Do we need to backport a fix to 9.1 if we found a race condition since COMPLETION is already released for 9.1?
Is it necessary to rewrite the implementation of BulkInferenceExecutor to fix the race condition? Or can we have a more surgical fix that we can review separately?

[ioanatia]

Is there any difference between this implementation and the new one you added?
I don't think it is - they both do the same thing, even if this one overrides the rule method and the new one the apply one.

So this diff here looks like it was entirely unnecessary.
If we needed to make a change to ResolveInference, we could have made it here, not move it entirely later in the file. This forces the reviewers to check each line to figure out what actually changed.

Maybe the one thing that we needed in Analyzer.java was the order in which we apply rules. But even that is debatable, because it's not necessary yet, we could have switched the order when we actually add the text_embedding function. 🤷‍♀️

[afoucret]

There is no major difference but the new method is easier to extends since I can not chain more transform.

return plan.transformDown(InferencePlan.class, p -> resolveInferencePlan(p, context));

will become when implementing the resolution of InferenceFunction

return plan.transformDown(InferencePlan.class, p -> resolveInferencePlan(p, context))
    .transformExpressionsOnly(InferenceFunction.class, f -> resolveInferenceFunction(f, context));

This is the whole point of this PR to anticipate change that are required in the existing framework and to isolate these changes, so we can focus on testing potential regression.

BTW, I am 100% sure that we will have to change the order of the resolution and that I will use this change to finish the implementation of TEXT_EMBEDDING.

[afoucret]

The race condition can happen when submitted several batch of inferences.
If previous batches have already exhausted the number of allocated permits, the newly submitted batch will never be started and. As a consequence the listener will never be called and the request will timeout.

In our execution model we are submitted one batch per page (then each request inside the batch are executed in parallel) up to 10 concurrent batches (this last point is managed by the AsyncOperator). It means that the problem affects mostly big request that are handling multiple pages (hundreds / thousands of row). So not the most likely case.

I have added a test to verify this parallel execution that I will be adding to the 9.1 branch :

    public void testParallelBulkExecution() throws Exception {
        int batches = between(50, 100);
        CountDownLatch latch = new CountDownLatch(batches);

        for (int i = 0; i < batches; i++) {
            List<InferenceAction.Request> requests = randomInferenceRequestList(between(1, 1_000));
            List<InferenceAction.Response> responses = randomInferenceResponseList(requests.size());

            Client client = mockClient(invocation -> {
                runWithRandomDelay(() -> {
                    ActionListener<InferenceAction.Response> l = invocation.getArgument(2);
                    l.onResponse(responses.get(requests.indexOf(invocation.getArgument(1, InferenceAction.Request.class))));
                });
                return null;
            });

            ActionListener<List<InferenceAction.Response>> listener = ActionListener.wrap(r -> {
                assertThat(r, equalTo(responses));
                LogManager.getLogger(BulkInferenceRunnerTests.class).warn("Received [{}] responses", responses.size());
                latch.countDown();
            }, ESTestCase::fail);

            inferenceRunnerFactory(client).create(randomBulkExecutionConfig()).executeBulk(requestIterator(requests), listener);
        }

        latch.await();
    }

I think a surgical patch will be not be to complicated.

[afoucret]

@ioanatia Here is the patch for the race condition in 9.1: https://github.com/elastic/elasticsearch/pull/130991/files

[tteofili]

LGTM, nice work Aurelien 💯

[tteofili]

nice :)

[tteofili]

is it always a BytesRefBlock ? probably it is, but I wonder if we should be safe and do a check here?

[afoucret]

There is a small issue here but I will handle it as part of another PR because it does also applies to branch 9.1 and 8.19, so it will be easier to backport.