[Endpoints] [4/x] Abstract store interface by BenWilson2 · Pull Request #19005 · mlflow/mlflow

BenWilson2 · 2025-11-24T17:57:52Z

🥞 Stacked PR

Use this link to review incremental changes.

stack/endpoints/abstract [Files changed]
- stack/endpoints/sql-store [Files changed]
  - stack/endpoints/rest [Files changed]
    - stack/endpoints/rest-2 [Files changed]
      - stack/endpoints/cache [Files changed]
        
        stack/endpoints/litellm [Files changed]
        
        stack/endpoints/ui-apis [Files changed]
        
        stack/endpoints/ui-create-endpoint [Files changed]
        
        stack/endpoints/ui-tabs [Files changed]
        
        stack/endpoints/key-management [Files changed]
        stack/endpoints/passphrase [Files changed]

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Add the abstract store interface for the endpoints and secrets management APIs

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

No. You can skip the rest of this section.
Yes. Give a description of this change to be included in the release notes for MLflow users.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

How should the PR be classified in the release notes? Choose one:

rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
rn/feature - A new user-facing feature worth mentioning in the release notes
rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?

Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
Bug fixes, doc updates and new features usually go into minor releases.
Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
Bug fixes and doc updates usually go into patch releases.

Yes (this PR will be cherry-picked and included in the next patch release)
No (this PR will be included in the next minor release)

github-actions · 2025-11-26T03:49:07Z

Documentation preview for ad08499 is available at:

https://pr-19005--mlflow-docs-preview.netlify.app/docs/latest/

More info

Ignore this comment if this PR does not change the documentation.
The preview is updated when a new commit is pushed to this PR.
This comment was created by this workflow run.
The documentation was built by this workflow run.

B-Step62 · 2025-12-03T08:24:44Z

mlflow/store/tracking/abstract_store.py

+        raise NotImplementedError(self.__class__.__name__)
+
+    # ===================================================================================
+    # Model Definitions Management


The term "model" is over-loaded. Is there any way to make it clear these APIs are for gateway models? Maybe we can define a mixin?

Yep! I think all of the stores should have mixins for them in order to isolate these and make it even more clear for us that these are UI-only endpoints / methods. I'll refactor to use mixins for the ABC, sql store, and rest store.

Copilot

Pull request overview

This PR introduces abstract store interface methods for secrets management and endpoint configuration as part of the MLflow Gateway/LLM endpoints feature. It provides comprehensive cryptographic infrastructure using envelope encryption (KEK/DEK pattern) following OWASP security best practices.

Key changes:

Adds cryptography utilities with AES-256-GCM encryption and PBKDF2 key derivation
Defines entity models for secrets, endpoints, model definitions, and bindings
Extends abstract store interface with 15 new methods for secrets and endpoints management

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
mlflow/utils/cryptography.py	Core cryptographic utilities implementing envelope encryption with KEK/DEK pattern, following OWASP guidelines (600k PBKDF2 iterations)
mlflow/entities/gateway_secrets.py	GatewaySecret entity for encrypted API key storage
mlflow/entities/gateway_endpoint.py	Entity models for endpoints, model definitions, bindings, and runtime configurations
mlflow/entities/init.py	Exports new gateway-related entities
mlflow/store/tracking/abstract_store.py	Abstract interface methods for CRUD operations on secrets, model definitions, endpoints, and bindings
mlflow/cli/cryptography.py	CLI command for KEK rotation with safety confirmations and progress reporting
mlflow/cli/init.py	Registers crypto CLI group with optional import handling
tests/utils/test_cryptography.py	Comprehensive unit tests for cryptographic functions (606 lines)
tests/entities/test_gateway_secrets.py	Unit tests for GatewaySecret entity
tests/entities/test_gateway_endpoint.py	Unit tests for endpoint-related entities (423 lines)
tests/cli/test_cryptography.py	Extensive CLI tests covering success/failure scenarios (443 lines)
docs/api_reference/api_inventory.txt	Adds new entities to API documentation inventory

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

mlflow/cli/cryptography.py

TomeHirata · 2025-12-05T02:21:27Z

mlflow/store/tracking/gateway/abstract_mixin.py

+from mlflow.store.tracking.gateway.entities import GatewayEndpointConfig
+
+
+class GatewayStoreMixin:


Did we split these methods as a Mixin because tracking store is getting heavy?

Mostly to separate out the ambiguous terms "model" (in MLflow that means something different in our public APIs), but also, yeah, sqlalchemy_store.py is the Titanic of Python modules

Make sense, we should probably do some refactoring of tracking stores using mixin

TomeHirata · 2025-12-05T02:22:53Z

mlflow/store/tracking/gateway/abstract_mixin.py

+        """
+        raise NotImplementedError(self.__class__.__name__)
+
+    def get_secret_info(


Can we make the resource name consistent?

Suggested change

def get_secret_info(

def get_secret(

I selected the 'info' part simply because I wanted to convey that this returns the metadata and not the actual secret. I'm not convinced that it's needed... curious to hear your thoughts on my reasoning (perhaps I'm just being too paranoid ;) )

Got it. I proposed get_secret in order to be consistent about the terminology. Since we return GatewaySecret, the method should be called get_secret or get_gateway_secrete. If the GatewaySecret entity does not include actual secret, GatewaySecret might be a misleading name. So I'd call it GatewaySecretInfo. What do you think?

Also, if we call this method get_secret_info, then can we change list_secrets into list_secret_info to be consistent?

good call! Since that is metadata as well :)

TomeHirata · 2025-12-05T02:36:45Z

mlflow/store/tracking/gateway/abstract_mixin.py

+        self,
+        resource_type: str,
+        resource_id: str,
+    ) -> GatewayEndpointConfig:


Is endpoint uniquely identified by resource_type and resource_id? If the composite key (endpoint_id, resource_type, resource_id) is unique, should the returned value list[GatewayEndpointConfig] or list[GatewayEndpoint]?

This is a really great point. For the initial version I've been focused on the concept of multiple resources -> single endpoint. But there's no reason to restrict this since the DB model supports this without any issues.

Got it, can we return multiple endpoints here to be future proof? We might want to have multiple endpoints attached to a single resource in the future. We can add the single restriction in the online-scorer specific logic in the downstream if needed.

yep making the change is in progress right now ;)

B-Step62

LGTM, #19005 (comment) is not blocking but I imagine it would be cleaner to move the method outside the store if it has special access requirement.

B-Step62 · 2025-12-05T04:45:16Z

mlflow/store/tracking/gateway/abstract_mixin.py

+        """
+        raise NotImplementedError(self.__class__.__name__)
+
+    def get_resource_endpoint_configs(


How to downstream invoke this API? Usually store methods are invoked through MlflowClient, but in this case, it is a bit weird because we don't want anything outside server to call this.

Yeah having it isolated and not even in the store implementation makes sense. I put the entities that back the responses here into the store layer, so this seems like a logical location to consolidate that access

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

BenWilson2 changed the title ~~WIP~~ [Endpoints] [4/x] Abstract store interface Nov 24, 2025

BenWilson2 marked this pull request as ready for review November 24, 2025 18:25

BenWilson2 mentioned this pull request Nov 24, 2025

[Endpoints] [8/x] Add credential cache #19014

Merged

29 tasks

github-actions bot added area/tracking Tracking service, tracking client APIs, autologging rn/feature Mention under Features in Changelogs. labels Nov 24, 2025

BenWilson2 force-pushed the stack/endpoints/abstract branch 3 times, most recently from 75dd05b to 2b4ec92 Compare November 26, 2025 03:39

BenWilson2 mentioned this pull request Nov 26, 2025

[Endpoints] [11/x] Main Gateway page and create endpoint #19042

Closed

29 tasks

This was referenced Nov 27, 2025

[Endpoints] [12/x] Add Left hand nav and management stub #19070

Closed

[Endpoints] [13/x] Add the API Key management pages #19071

Closed

BenWilson2 force-pushed the stack/endpoints/abstract branch from 2b4ec92 to 139d2d1 Compare November 27, 2025 06:31

BenWilson2 mentioned this pull request Nov 27, 2025

[Endpoints] [14/x] Models management UI #19074

Closed

29 tasks

BenWilson2 force-pushed the stack/endpoints/abstract branch 2 times, most recently from 66167b9 to bcecc8f Compare December 2, 2025 02:14

BenWilson2 added the team-review Trigger a team review request label Dec 2, 2025

github-actions bot requested review from B-Step62, daniellok-db, harupy, kevin-lyn and serena-ruan December 2, 2025 16:03

github-actions bot requested review from TomeHirata, WeichenXu123 and xq-yin December 2, 2025 16:03

BenWilson2 mentioned this pull request Dec 2, 2025

[Endpoints] [14/x] Add encryption availability endpoints #19176

Closed

29 tasks

BenWilson2 force-pushed the stack/endpoints/abstract branch 2 times, most recently from 9f711e6 to 62710ce Compare December 3, 2025 01:52

B-Step62 reviewed Dec 3, 2025

View reviewed changes

BenWilson2 force-pushed the stack/endpoints/abstract branch 3 times, most recently from 3a572c2 to 45dcd4f Compare December 4, 2025 01:54

Copilot AI review requested due to automatic review settings December 4, 2025 01:54

Copilot started reviewing on behalf of BenWilson2 December 4, 2025 01:55 View session

Copilot finished reviewing on behalf of BenWilson2 December 4, 2025 01:56

Copilot AI reviewed Dec 4, 2025

View reviewed changes

mlflow/cli/cryptography.py Outdated Show resolved Hide resolved

BenWilson2 force-pushed the stack/endpoints/abstract branch 2 times, most recently from 49a8458 to 903c1e8 Compare December 4, 2025 23:40

TomeHirata reviewed Dec 5, 2025

View reviewed changes

BenWilson2 force-pushed the stack/endpoints/abstract branch from 903c1e8 to 5bdfd7f Compare December 5, 2025 04:29

B-Step62 approved these changes Dec 5, 2025

View reviewed changes

BenWilson2 added 2 commits December 5, 2025 00:30

WIP

b1f9420

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

abstract store

ad08499

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

BenWilson2 force-pushed the stack/endpoints/abstract branch from 5bdfd7f to ad08499 Compare December 5, 2025 06:50

BenWilson2 added this pull request to the merge queue Dec 5, 2025

Merged via the queue into mlflow:master with commit a6167b1 Dec 5, 2025
50 checks passed

BenWilson2 deleted the stack/endpoints/abstract branch December 5, 2025 16:24

		from mlflow.store.tracking.gateway.entities import GatewayEndpointConfig


		class GatewayStoreMixin:

Conversation

BenWilson2 commented Nov 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🥞 Stacked PR

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

How should the PR be classified in the release notes? Choose one:

Should this PR be included in the next patch release?

Uh oh!

github-actions bot commented Nov 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

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!

TomeHirata Dec 5, 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!

TomeHirata Dec 5, 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!

B-Step62 left a comment • 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!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

BenWilson2 commented Nov 24, 2025 •

edited

Loading

github-actions bot commented Nov 26, 2025 •

edited

Loading

TomeHirata Dec 5, 2025 •

edited

Loading

TomeHirata Dec 5, 2025 •

edited

Loading

B-Step62 left a comment •

edited

Loading