[ML] Improve EIS authorization to perform requests on a periodic basis instead of only once #123639

jonathan-buttner · 2025-02-27T20:00:05Z

This PR refactors the authorization logic to accomplish a few things:

I removed the revocation logic since we don't need it right now and it complicates things
Authorization happens after the node boots up and on a interval of every 10 minutes + between 0 - 5 minutes of jitter

Testing

I manually tested by spinning up an EIS gateway and ensuring that the authorization requests were received to simulate success. I also tested by modifying the gateway to return an empty response to simulate revoking authorization.

jonathan-buttner · 2025-02-27T20:03:34Z

...st/java/org/elasticsearch/xpack/inference/integration/InferenceRevokeDefaultEndpointsIT.java

@@ -1,288 +0,0 @@
-/*


We're not revoking anymore so we don't need these tests.

jonathan-buttner · 2025-02-27T20:04:55Z

...rc/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java

        );
    }

-    private record DefaultModelConfig(Model model, MinimalServiceSettings settings) {}


Basically moved all this logic to the new ElasticInferenceServiceAuthorizationHandler class.

jonathan-buttner · 2025-02-27T20:05:15Z

...rc/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java

-        return authorizedModels;
-    }
-
-    private void handleRevokedDefaultConfigs(Set<String> authorizedDefaultModelIds) {


This got removed since we're no longer revoking.

@jonathan-buttner I think I was imprecise here and misunderstood how the endpoint configurations work. We still need to be able to revoke default endpoints if we disable EIS for a project. I thought a restart would still give us this behavior, but with periodic tries in place restarts do not do anything special.
I think we need to keep the revokations.

Ah ok! Thanks for clarifying. I'll add the functionality back 👍

jonathan-buttner · 2025-02-27T20:05:51Z

...rc/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java

                responseString = responseString + " " + useChatCompletionUrlMessage(model);
            }
            listener.onFailure(new ElasticsearchStatusException(responseString, RestStatus.BAD_REQUEST));
+            return;


Noticed we had a fall through bug that luckily is getting picked up by the if-block below but figured I'd fix it anyway.

jonathan-buttner · 2025-02-27T20:07:07Z

...ck/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandler.java


 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.elasticsearch.ElasticsearchWrapperException;


I would just consider this class completely new and don't even look at the old code when reviewing.

jonathan-buttner · 2025-02-27T20:11:18Z

...ck/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandler.java

+    private synchronized void setAuthorizedContent(ElasticInferenceServiceAuthorizationModel auth) {
+        logger.debug("Received authorization response");
+        var authorizedTaskTypesAndModels = authorizedContent.get().taskTypesAndModels.merge(auth)
+            .newLimitedToTaskTypes(EnumSet.copyOf(implementedTaskTypes));


As a performance improvement we could check if the new authorization response is the same as the previous on and skip the rest of the logic here.

jonathan-buttner · 2025-02-27T20:12:53Z

...rence/services/elastic/authorization/ElasticInferenceServiceAuthorizationRequestHandler.java

@@ -0,0 +1,137 @@
+/*


This is the previous ElasticInferenceServiceAuthorizationHandler class.

jonathan-buttner · 2025-02-27T20:14:19Z

...inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationModelTests.java

+        assertThat(limitedAuth, is(ElasticInferenceServiceAuthorizationModel.newDisabledService()));
+    }
+
+    public void testMerge_CombinesCorrectly() {


New merge method tests.

jonathan-buttner · 2025-02-27T20:15:00Z

.../services/elastic/authorization/ElasticInferenceServiceAuthorizationRequestHandlerTests.java

@@ -0,0 +1,272 @@
+/*


This was renamed from ElasticInferenceServiceAuthorizationHandlerTests

jonathan-buttner · 2025-02-27T20:15:33Z

...ference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandlerTests.java


 package org.elasticsearch.xpack.inference.services.elastic.authorization;

-import org.apache.logging.log4j.Logger;


Just review this as a totally new file. Ignore the previous code.

…-cache

elasticsearchmachine · 2025-02-27T21:27:17Z

Pinging @elastic/ml-core (Team:ML)

prwhelan · 2025-02-28T13:22:20Z

...rc/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java

     * @param waitTime the max time to wait
     * @throws IllegalStateException if the wait time is exceeded or the call receives an {@link InterruptedException}
     */
    public void waitForAuthorizationToComplete(TimeValue waitTime) {


This isn't called anywhere other than tests, right?

Yeah correct, I'll make it package private

prwhelan · 2025-02-28T13:32:40Z

...ck/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandler.java

-        this.threadPool = Objects.requireNonNull(threadPool);
-        logger = LogManager.getLogger(ElasticInferenceServiceAuthorizationHandler.class);
+        configuration = new AtomicReference<>(
+            new ElasticInferenceService.Configuration(authorizedContent.get().taskTypesAndModels.getAuthorizedTaskTypes())


Won't this always return and set an empty set? If so, I see it's already in the previous iteration, maybe it's something we fix after this PR?

At this point it'll always be an empty set, but when we get an authorization response back it updates it:

https://github.com/elastic/elasticsearch/pull/123639/files#diff-c37e92ab207cd0222cb473c4e45b5af53a479972c4bf7b885c0303fa99d0dba7R264

…-cache

elasticsearchmachine · 2025-03-03T20:02:12Z

💔 Backport failed

Status	Branch	Result
❌	9.0	Commit could not be cherrypicked due to conflicts
❌	8.18	Commit could not be cherrypicked due to conflicts
❌	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 123639

…s instead of only once (elastic#123639) * Refactoring * Add internal cluster setting to aid testing * [CI] Auto commit changes from spotless * Allowing the auth interval to be configurable via a setting * Removing unused code * Adding revocation functionality back --------- Co-authored-by: elasticsearchmachine <[email protected]> (cherry picked from commit a88d645) # Conflicts: # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandler.java # x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandlerTests.java

…s instead of only once (elastic#123639) * Refactoring * Add internal cluster setting to aid testing * [CI] Auto commit changes from spotless * Allowing the auth interval to be configurable via a setting * Removing unused code * Adding revocation functionality back --------- Co-authored-by: elasticsearchmachine <[email protected]> (cherry picked from commit a88d645) # Conflicts: # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandler.java

jonathan-buttner · 2025-03-03T20:33:34Z

💚 All backports created successfully

Status	Branch	Result
✅	8.x
✅	9.0
✅	8.18

Questions ?

Please refer to the Backport tool documentation

…s instead of only once (elastic#123639) * Refactoring * Add internal cluster setting to aid testing * [CI] Auto commit changes from spotless * Allowing the auth interval to be configurable via a setting * Removing unused code * Adding revocation functionality back --------- Co-authored-by: elasticsearchmachine <[email protected]> (cherry picked from commit a88d645) # Conflicts: # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandler.java # x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandlerTests.java

…ic basis instead of only once (#123639) (#123920) * [ML] Improve EIS authorization to perform requests on a periodic basis instead of only once (#123639) * Refactoring * Add internal cluster setting to aid testing * [CI] Auto commit changes from spotless * Allowing the auth interval to be configurable via a setting * Removing unused code * Adding revocation functionality back --------- Co-authored-by: elasticsearchmachine <[email protected]> (cherry picked from commit a88d645) # Conflicts: # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandler.java # x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandlerTests.java * Fixing tests

…c basis instead of only once (#123639) (#123918) * [ML] Improve EIS authorization to perform requests on a periodic basis instead of only once (#123639) * Refactoring * Add internal cluster setting to aid testing * [CI] Auto commit changes from spotless * Allowing the auth interval to be configurable via a setting * Removing unused code * Adding revocation functionality back --------- Co-authored-by: elasticsearchmachine <[email protected]> (cherry picked from commit a88d645) # Conflicts: # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandler.java * Fixing tests * Fixing task type any test

…c basis instead of only once (#123639) (#123916) * [ML] Improve EIS authorization to perform requests on a periodic basis instead of only once (#123639) * Refactoring * Add internal cluster setting to aid testing * [CI] Auto commit changes from spotless * Allowing the auth interval to be configurable via a setting * Removing unused code * Adding revocation functionality back --------- Co-authored-by: elasticsearchmachine <[email protected]> (cherry picked from commit a88d645) # Conflicts: # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandler.java # x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandlerTests.java * Removing getFirst calls

…s instead of only once (elastic#123639) * Refactoring * Add internal cluster setting to aid testing * [CI] Auto commit changes from spotless * Allowing the auth interval to be configurable via a setting * Removing unused code * Adding revocation functionality back --------- Co-authored-by: elasticsearchmachine <[email protected]>

jonathan-buttner added 2 commits February 25, 2025 16:51

Refactoring

94612b5

Add internal cluster setting to aid testing

b8bb59a

jonathan-buttner added >non-issue :ml Machine learning Team:ML Meta label for the ML team auto-backport Automatically create backport pull requests when merged Feature:GenAI Features around GenAI v9.0.0 v8.18.0 v8.18.1 v8.19.0 v9.0.1 v9.1.0 labels Feb 27, 2025

[CI] Auto commit changes from spotless

7d6f525

jonathan-buttner commented Feb 27, 2025

View reviewed changes

Merge branch 'main' of github.com:elastic/elasticsearch into eis-auth…

13af714

…-cache

jonathan-buttner requested review from maxjakob and prwhelan February 27, 2025 21:26

jonathan-buttner marked this pull request as ready for review February 27, 2025 21:26

prwhelan approved these changes Feb 28, 2025

View reviewed changes

jonathan-buttner added 3 commits February 28, 2025 13:43

Allowing the auth interval to be configurable via a setting

5361c2b

Merge branch 'main' of github.com:elastic/elasticsearch into eis-auth…

340f0c6

…-cache

Removing unused code

14ece26

jonathan-buttner mentioned this pull request Feb 28, 2025

[ML] Adding DNS look up logic for triaging EIS timeouts #123767

Closed

jonathan-buttner added 2 commits March 3, 2025 11:34

Adding revocation functionality back

6b110dc

Merge branch 'main' of github.com:elastic/elasticsearch into eis-auth…

690c58e

…-cache

jonathan-buttner requested a review from prwhelan March 3, 2025 18:37

prwhelan approved these changes Mar 3, 2025

View reviewed changes

jonathan-buttner merged commit a88d645 into elastic:main Mar 3, 2025
17 checks passed

elasticsearchmachine added the backport pending label Mar 3, 2025

This was referenced Mar 3, 2025

[8.x] [ML] Improve EIS authorization to perform requests on a periodic basis instead of only once (#123639) #123916

Merged

[9.0] [ML] Improve EIS authorization to perform requests on a periodic basis instead of only once (#123639) #123918

Merged

jonathan-buttner mentioned this pull request Mar 3, 2025

[8.18] [ML] Improve EIS authorization to perform requests on a periodic basis instead of only once (#123639) #123920

Merged

jonathan-buttner mentioned this pull request Aug 7, 2025

[ML] Improve EIS auth call logs and fix revocation bug #132546

Merged


		package org.elasticsearch.xpack.inference.services.elastic.authorization;

		import org.apache.logging.log4j.Logger;

[ML] Improve EIS authorization to perform requests on a periodic basis instead of only once #123639

[ML] Improve EIS authorization to perform requests on a periodic basis instead of only once #123639

Uh oh!

Conversation

jonathan-buttner commented Feb 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Testing

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 Feb 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

elasticsearchmachine commented Feb 27, 2025

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!

Uh oh!

elasticsearchmachine commented Mar 3, 2025

💔 Backport failed

Uh oh!

jonathan-buttner commented Mar 3, 2025

💚 All backports created successfully

Questions ?

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

jonathan-buttner commented Feb 27, 2025 •

edited

Loading

jonathan-buttner Feb 27, 2025 •

edited

Loading