[ML] Fix loss of context in the inference API for streaming APIs #118999
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
jonathan-buttner
added
>bug
:ml
|
Hi @jonathan-buttner, I've created a changelog YAML for you. |
|
Hi @jonathan-buttner, I've updated the changelog YAML for you. |
jonathan-buttner
changed the title
| var responseConsumer = new AsyncInferenceResponseConsumer(); | ||
| request.setOptions(RequestOptions.DEFAULT.toBuilder().setHttpAsyncResponseConsumerFactory(() -> responseConsumer).build()); | ||
| var latch = new CountDownLatch(1); | ||
| client().performRequestAsync(request, new ResponseListener() { | ||
| @Override | ||
| public void onSuccess(Response response) { | ||
| if (responseConsumerCallback != null) { |
There was a problem hiding this comment.
Adding a way to get the response so we can check the headers
| @@ -469,7 +477,7 @@ public void testSupportedStream() throws Exception { | |||
|
|
|||
| var input = IntStream.range(1, 2 + randomInt(8)).mapToObj(i -> randomAlphanumericOfLength(5)).toList(); | |||
| try { | |||
| var events = streamInferOnMockService(modelId, TaskType.COMPLETION, input); | |||
| var events = streamInferOnMockService(modelId, TaskType.COMPLETION, input, VALIDATE_ELASTIC_PRODUCT_HEADER_CONSUMER); | |||
There was a problem hiding this comment.
Unfortunately, even if I remove my change this test still passes. I opted for keeping these tests just in case something really goes wrong in the future but they wouldn't have caught the original issue 😞
| @@ -115,6 +121,12 @@ private void initializeStream(InferenceAction.Response response) { | |||
| ) | |||
| ); | |||
| }); | |||
|
|
|||
| nextBodyPartListener = ContextPreservingActionListener.wrapPreservingContext( | |||
There was a problem hiding this comment.
When I manually test this, it fixes the issue. I always get the header back as expected. Without the fix the _stream api usually does not return the header. I believe this is because we are create a new listener that maybe gets executed with a different thread context. I think if we responded in the onResponse implementation of this class it'd probably avoid the issue. But that's not how this was designed so I'm keeping the structure as it was.
…ming-lost-context
jonathan-buttner
requested review from
davidkyle,
maxhniebergall and
prwhelan
|
Pinging @elastic/ml-core (Team:ML) |
maxhniebergall
approved these changes
Dec 19, 2024
jonathan-buttner
added a commit
to jonathan-buttner/elasticsearch
that referenced
this pull request
Dec 23, 2024
…stic#118999) * Adding context preserving fix * Update docs/changelog/118999.yaml * Update docs/changelog/118999.yaml * Using a setonce and adding a test * Updating the changelog (cherry picked from commit 7ba3cb9) # Conflicts: # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java
jonathan-buttner
added a commit
to jonathan-buttner/elasticsearch
that referenced
this pull request
Dec 23, 2024
…stic#118999) * Adding context preserving fix * Update docs/changelog/118999.yaml * Update docs/changelog/118999.yaml * Using a setonce and adding a test * Updating the changelog (cherry picked from commit 7ba3cb9) # Conflicts: # x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceBaseRestTest.java # x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rest/RestUnifiedCompletionInferenceAction.java # x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rest/RestUnifiedCompletionInferenceActionTests.java
💚 All backports created successfully
Questions ?Please refer to the Backport tool documentation |
jonathan-buttner
added a commit
to jonathan-buttner/elasticsearch
that referenced
this pull request
Dec 23, 2024
…stic#118999) * Adding context preserving fix * Update docs/changelog/118999.yaml * Update docs/changelog/118999.yaml * Using a setonce and adding a test * Updating the changelog (cherry picked from commit 7ba3cb9) # Conflicts: # x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceBaseRestTest.java # x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rest/RestUnifiedCompletionInferenceAction.java # x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rest/RestUnifiedCompletionInferenceActionTests.java
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR fixes an issue where the
X-elastic-produceheader was not being returned sometimes. I believe this was a concurrency issue in someway because when testing using the_streamAPI the header is usually not returned. But when testing with the_unifiedAPI the header is usually returned.I believe the issue relates to the fact that the rest handler's listener doesn't immediately return the response. The
onResponsemethod creates a new listener that is called sometime later (maybe on a different thread?).One thing I noticed is that the headers are only returned once which is why I only preserve the context on the first action listener.
TODO
SetOnceinstead of theThreadPoolto avoid a potential issue where thecreateComponents()call gets moved after the route retrieval. ✅Fixes #119000