[Inference API] Propagate product use case http header to EIS #124025

timgrein · 2025-03-04T17:05:55Z

This PR reads the X-elastic-product-use-case header containing the product use case EIS is called with (Assistants etc.). I had to use a workaround to propagate the information through the transport layer: I set the header explicitly in the ThreadContext as we would lose it otherwise when the InferenceActionProxy makes an internal call to InferenceAction or UnifiedCompletionAction (thread context gets stashed and then reconstructed losing most headers; as this is specific to the inference API/EIS we shouldn't add it to Task.HEADERS_TO_COPY - had this discussion with some ES devs).

I was hesitant to pass the InferenceContext through the base methods (doInfer, doUnifiedCompletionInfer etc.) as this would imply changes in all integrations making this PR even larger as it already is, especially considering what it does (passing one value). If we feel like that the product use case information is useful for all integrations (which it probably is) we can still follow-up on this initial change. For now I want to keep it isolated for EIS.

…text)

…ase-header-to-eis' into read-and-propagate-product-use-case-header-to-eis

…ionTests

…o-eis

elasticsearchmachine · 2025-03-07T16:54:27Z

Hi @timgrein, I've created a changelog YAML for you.

jonathan-buttner

Nice work! Left a couple comments.

jonathan-buttner · 2025-03-07T17:04:49Z

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceContext.java

+        this(in.readString());
+    }
+
+    public static InferenceContext empty() {


How about we create a static instance that way we don't create multiple empty ones? Something like this:

public static final InferenceContext EMPTY_INSTANCE = new InferenceContext("");

Good point! Adjusted with Replace InferenceContext.empty() with InferenceContext.EMPTY_INSTANCE

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceContext.java

jonathan-buttner · 2025-03-07T17:12:11Z

server/src/main/java/org/elasticsearch/TransportVersions.java

    public static final TransportVersion INCLUDE_INDEX_MODE_IN_GET_DATA_STREAM = def(9_023_0_00);
    public static final TransportVersion MAX_OPERATION_SIZE_REJECTIONS_ADDED = def(9_024_0_00);
    public static final TransportVersion RETRY_ILM_ASYNC_ACTION_REQUIRE_ERROR = def(9_025_0_00);
+    public static final TransportVersion INFERENCE_CONTEXT = def(9_026_0_00);


Just a reminder, if we do want to backport to 8.19 we'll need a TransportVersion for 8.x

for example: COHERE_BIT_EMBEDDING_TYPE_SUPPORT_ADDED_BACKPORT_8_X

We'll also need to change the onAfter() check. Here's an example:
https://github.com/elastic/elasticsearch/blob/main/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/embeddings/CohereEmbeddingType.java#L131-L132

The code in 8.x will look different too (since the 9.x transport version won't exist): https://github.com/elastic/elasticsearch/blob/8.x/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/embeddings/CohereEmbeddingType.java#L131

Thanks for the explanation and the code examples.

Adjusted with Add TransportVersion for 8_X.

In the backport I would then need to replace TransportVersions.INFERENCE_CONTEXT with TransportVersions.INFERENCE_CONTEXT_8_X, right?

jonathan-buttner · 2025-03-07T17:23:46Z

...nce/src/main/java/org/elasticsearch/xpack/inference/action/BaseTransportInferenceAction.java


+            var context = request.getContext();
+            if (Objects.nonNull(context)) {
+                threadPool.getThreadContext().putHeader(InferencePlugin.X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER, context.productUseCase());


Hmm, another option would be to pass this through the infer calls. That'll probably change a ton of files though 🤔 . It does seem a bit strange to parse it out of the header and then put it back in a header when we already have it hmm.

Drilling through a bunch of function calls isn't great either though.

@prwhelan @dan-rubinstein @davidkyle what do you think?

I wonder if we should create a context/components object that is like a catch all for these types of changes. That way in the future we just add it to that class's definition and we don't have to drill it through a ton of places.

I'd prefer that, though I vote we do that in a separate change since this one is already quite large

Hmm, another option would be to pass this through the infer calls. That'll probably change a ton of files though 🤔 .

Yeah, that was basically my reasoning I've put in the PR comment "I was hesitant to pass the InferenceContext through the base methods (doInfer, doUnifiedCompletionInfer etc.) as this would imply changes in all integrations making this PR even larger as it already is, especially considering what it does (passing one value)". But nevertheless I also think it's cleaner to pass it through the methods, as it's then obvious from the signature that there's a context object.

I'd prefer that, though I vote we do that in a separate change since this one is already quite large

I would also prefer this and keep it as is for now 👍

I would also prefer this and keep it as is for now 👍

Sounds good Tim!

Actually could you add a TODO above the line to as a reminder for us to move it to being passed through the various method calls?

Yes of course! Added with Add TODO to remove temporary product use case propagation

jonathan-buttner · 2025-03-07T17:28:18Z

...ugin/inference/src/main/java/org/elasticsearch/xpack/inference/rest/BaseInferenceAction.java

+        }
+
+        // We always get the first value as the header doesn't allow multiple values
+        return productUseCaseHeaders.getFirst();


If we do backport this, it's going to complain about getFirst not being a call in 8.19. Might be worth just leave it as get(0) to avoid all that lol (I've run into it many times from experience).

Good catch 🎣 I've also ran into this in the past, adjusted with Use .get(0) instead of getFirst to avoid compilation errors in backport.

…ase-header-to-eis' into read-and-propagate-product-use-case-header-to-eis

…inference/InferenceContext.java Co-authored-by: Jonathan Buttner <[email protected]>

jonathan-buttner

Thanks for the changes! If you could add a TODO for passing the inference context around that'd be great!

jonathan-buttner · 2025-03-10T12:21:03Z

...nce/src/main/java/org/elasticsearch/xpack/inference/action/BaseTransportInferenceAction.java


+            var context = request.getContext();
+            if (Objects.nonNull(context)) {
+                threadPool.getThreadContext().putHeader(InferencePlugin.X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER, context.productUseCase());


Actually could you add a TODO above the line to as a reminder for us to move it to being passed through the various method calls?

…o-eis

… the same as for others.

…ase-header-to-eis' into read-and-propagate-product-use-case-header-to-eis

…o-eis

elasticsearchmachine · 2025-03-12T11:49:50Z

💔 Backport failed

Status	Branch	Result
❌	8.x	Commit could not be cherrypicked due to conflicts

You can use sqren/backport to manually backport by running backport --upstream elastic/elasticsearch --pr 124025

timgrein · 2025-03-12T16:25:49Z

💚 All backports created successfully

Status	Branch	Result
✅	8.x

Questions ?

Please refer to the Backport tool documentation

…c#124025)

…124025) (#124666)

…c#124025)

WIP

ab712dd

elasticsearchmachine added the v9.1.0 label Mar 4, 2025

timgrein changed the title ~~[Draft] [Inference API] Read and propagate product use case http header~~ [Draft] [Inference API] Read and propagate product use case http header to EIS Mar 4, 2025

elasticsearchmachine and others added 2 commits March 4, 2025 17:14

[CI] Auto commit changes from spotless

b617e0a

Iterate (propagate InferenceContext & put HTTP header into thread con…

2e8231a

…text)

timgrein changed the title ~~[Draft] [Inference API] Read and propagate product use case http header to EIS~~ [Draft] [Inference API] Propagate product use case http header to EIS Mar 6, 2025

timgrein and others added 17 commits March 6, 2025 16:28

Iter (product use case propagation works)

ca79955

Iter (move request metadata extraction to parent request manager class)

2c12778

[CI] Auto commit changes from spotless

cc52207

Add docs to InferenceContext

a1da2f6

Merge remote-tracking branch 'origin/read-and-propagate-product-use-c…

a6e9297

…ase-header-to-eis' into read-and-propagate-product-use-case-header-to-eis

Update InferenceContext.java

82ce713

Remove duplicate context from InferenceAction.Request

e5c5933

Merge remote-tracking branch 'origin/read-and-propagate-product-use-c…

634aaef

…ase-header-to-eis' into read-and-propagate-product-use-case-header-to-eis

Add InferenceContextTests

0302468

Add additional test case for context in InferenceActionRequestTests

ef35491

Add new test cases to InferenceActionRequestTests

3abe303

Add new test cases to UnifiedCompletionActionRequestTests

517c943

Add test to verify that the header is set in the thread context

0bc06df

Remove TODO

7a515bc

Add product use case header extraction test cases to BaseInferenceAct…

cf55cb5

…ionTests

Remove addressed TODO and spotlessApply

8d44488

Add product use case propagation tests in ElasticInferenceServiceTests

139f7a5

timgrein changed the title ~~[Draft] [Inference API] Propagate product use case http header to EIS~~ [Inference API] Propagate product use case http header to EIS Mar 7, 2025

timgrein and others added 2 commits March 7, 2025 15:35

Merge branch 'main' into read-and-propagate-product-use-case-header-t…

280f0cf

…o-eis

Fix compilation error

0a3b0a0

timgrein marked this pull request as ready for review March 7, 2025 14:40

elasticsearchmachine added the needs:triage Requires assignment of a team area label label Mar 7, 2025

prwhelan added :ml Machine learning Team:ML Meta label for the ML team labels Mar 7, 2025

Update docs/changelog/124025.yaml

cd1beee

jonathan-buttner reviewed Mar 7, 2025

View reviewed changes

timgrein and others added 5 commits March 10, 2025 11:04

Replace InferenceContext.empty() with InferenceContext.EMPTY_INSTANCE

1c0a3c4

Merge remote-tracking branch 'origin/read-and-propagate-product-use-c…

d6a7bb9

…ase-header-to-eis' into read-and-propagate-product-use-case-header-to-eis

Update x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/…

89a973a

…inference/InferenceContext.java Co-authored-by: Jonathan Buttner <[email protected]>

Use .get(0) instead of getFirst to avoid compilation errors in backport.

2f881e0

Add TransportVersion for 8_X

e5153a1

jonathan-buttner approved these changes Mar 10, 2025

View reviewed changes

timgrein and others added 13 commits March 10, 2025 14:02

Add TODO to remove temporary product use case propagation

149a094

Add comment with rationale explaining difference in header extraction

69687d5

Ensure that productUseCase field in InferenceContext is non-null

bc1a74f

Merge branch 'main' into read-and-propagate-product-use-case-header-t…

8ab08f9

…o-eis

Merge branch 'main' into read-and-propagate-product-use-case-header-t…

e34359c

…o-eis

spotlessApply

a7d0f27

fix checkstyle errors in xpack core plugin

e494a39

Fix checkstyle errors in inference plugin

5b90efb

[CI] Auto commit changes from spotless

418241f

Fix test in InferenceActionRequestTests and adapt the structure to be…

9163207

… the same as for others.

Merge remote-tracking branch 'origin/read-and-propagate-product-use-c…

5d4b92f

…ase-header-to-eis' into read-and-propagate-product-use-case-header-to-eis

Add equals/hashCode to InferenceContext

fa47c8c

Merge branch 'main' into read-and-propagate-product-use-case-header-t…

0515e1a

…o-eis

timgrein merged commit 0b83425 into elastic:main Mar 12, 2025
16 checks passed

elasticsearchmachine added the backport pending label Mar 12, 2025

timgrein mentioned this pull request Mar 12, 2025

[8.x] [Inference API] Propagate product use case http header to EIS (#124025) #124666

Merged

albertzaharovits pushed a commit to albertzaharovits/elasticsearch that referenced this pull request Mar 13, 2025

[Inference API] Propagate product use case http header to EIS (elasti…

f28ef64

…c#124025)

timgrein added a commit that referenced this pull request Mar 13, 2025

[8.x] [Inference API] Propagate product use case http header to EIS (#…

07552e8

…124025) (#124666)

jfreden pushed a commit to jfreden/elasticsearch that referenced this pull request Mar 13, 2025

[Inference API] Propagate product use case http header to EIS (elasti…

7e58b1c

…c#124025)

[Inference API] Propagate product use case http header to EIS #124025

[Inference API] Propagate product use case http header to EIS #124025

Uh oh!

Conversation

timgrein commented Mar 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

elasticsearchmachine commented Mar 7, 2025

Uh oh!

jonathan-buttner left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jonathan-buttner left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

elasticsearchmachine commented Mar 12, 2025

💔 Backport failed

Uh oh!

timgrein commented Mar 12, 2025

💚 All backports created successfully

Questions ?

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

timgrein commented Mar 4, 2025 •

edited

Loading