Add MLflow SDK support for Model Registry as a Tracking Store, fixes #1225 #1337

dhirajsb · 2025-07-19T16:22:24Z

Description

MLflow SDK supports integrations with non-mlflow Tracking stores by implementing a tracking.AbstractStore API.
This PR implements that API to allow MLflow SDK to directly call Model Registry APIs for storing Experiment and Run metadata.
See the feature request issue #1225 for motivation, design and implementation details.

This PR is based on the Experiment Tracking PR #1318 and must be merged after that PR.
Fixes #1225.

How Has This Been Tested?

Extensive unit and e2e tests have been provided. A local e2e test is provide to test against the go REST server, as well as an e2e test to run against a deployed model registry to verify RBAC support.

Merge criteria:

All the commits have been signed-off (To pass the DCO check)

The commits have meaningful messages
Automated tests are provided as part of the PR for major new functionalities; testing instructions have been added in the PR body (for PRs involving changes that are not immediately obvious).
The developer has manually tested the changes and verified that the changes work.
Code changes follow the kubeflow contribution guidelines.
For first time contributors: Please reach out to the Reviewers to ensure all tests are being run, ensuring the label ok-to-test has been added to the PR.

If you have UI changes

The developer has added tests or explained why testing cannot be added.
Included any necessary screenshots or gifs if it was a UI change.
Verify that UI/UX changes conform the UX guidelines for Kubeflow.

google-oss-prow · 2025-07-19T16:22:33Z

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign andreyvelich for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

dhirajsb · 2025-07-19T16:23:14Z

The long commit history should become much shorter once #1318 is merged in main.

dhirajsb · 2025-07-19T16:26:30Z

To review this PR, focus on the clients/modelregistry_plugin plugin project.

dhirajsb · 2025-09-04T06:14:22Z

Re-based branch on main on top of mlmd removal. Ready for review.

clients/modelregistry_plugin/.gitignore

clients/modelregistry_plugin/README.md

clients/modelregistry_plugin/modelregistry_plugin/api_client.py

clients/modelregistry_plugin/modelregistry_plugin/auth.py

jonburdo

Most of my comments are just suggestions and could be done later. Overall lgtm

clients/modelregistry_plugin/.python-version

clients/modelregistry_plugin/pyproject.toml

clients/modelregistry_plugin/.gitignore

clients/modelregistry_plugin/CONTRIBUTING.md

clients/modelregistry_plugin/tests/test_store.py

clients/modelregistry_plugin/REFACTORING_SUMMARY.md

jonburdo

/lgtm

jonburdo

/lgtm

There are some changes I would like to see us make in the future, but I think this looks good for now. If you were going to make one more change, I would suggest replacing builtins.open with model_registry_mlflow.auth.open and do same for the other patches. Even this should be okay for now though.

clients/model-registry-mlflow/src/model_registry_mlflow/api_client.py

clients/model-registry-mlflow/src/model_registry_mlflow/auth.py

clients/model-registry-mlflow/tests/test_api_client.py

clients/model-registry-mlflow/tests/test_auth.py

dhirajsb · 2025-09-11T19:34:10Z

@jonburdo thanks for the detailed review. I have addressed all the comments. On using mtime for token file caching, I wrote a local benchmark script (using cursor) and it found that with a 95%+ hit rate, using mtime to avoid reading the file would be 4-5X faster.

jonburdo

/lgtm

Thanks for all the hard work on this! Looks great :)

Signed-off-by: Dhiraj Bokde <[email protected]>

…tests for mlflow plugin Signed-off-by: Dhiraj Bokde <[email protected]>

Signed-off-by: Dhiraj Bokde <[email protected]>

… and runs search Signed-off-by: Dhiraj Bokde <[email protected]>

…Store class, add tests for new methods, replace print statements in tests Signed-off-by: Dhiraj Bokde <[email protected]>

Signed-off-by: Dhiraj Bokde <[email protected]>

Co-authored-by: Paul Boyd <[email protected]> Signed-off-by: Dhiraj Bokde <[email protected]>

…del_registry_mlflow to match pypi and python pkg naming conventions, removed redunant code in tests, general code cleanup Signed-off-by: Dhiraj Bokde <[email protected]>

…ild system Signed-off-by: Dhiraj Bokde <[email protected]>

Signed-off-by: Dhiraj Bokde <[email protected]>

google-oss-prow · 2025-09-11T20:15:09Z

New changes are detected. LGTM label has been removed.

dhirajsb · 2025-09-11T20:46:01Z

@tarilabs @Al-Pragliola please add an approval so we can close this.

Signed-off-by: Dhiraj Bokde <[email protected]>

rareddy · 2025-09-11T21:43:29Z

@dhirajsb please hold on this with merging; we did not have any conversations in the community on the direction; we need to establish what we are signing up here. Supporting MLflow this way is big commitment, especially given experiment tracking KEP was closed without resolution in the KFP recently.

github-actions · 2025-12-12T04:17:58Z

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

github-actions · 2026-01-01T04:27:22Z

This pull request has been automatically closed because it has not had recent activity. Please comment "/reopen" to reopen it.

google-oss-prow bot requested a review from Al-Pragliola July 19, 2025 16:22

google-oss-prow bot requested a review from andreyvelich July 19, 2025 16:22

github-actions bot added Area/MR Python client Area/Go REST server labels Jul 19, 2025

google-oss-prow bot added the size/XXL label Jul 19, 2025

dhirajsb marked this pull request as draft July 19, 2025 16:28

google-oss-prow bot added the do-not-merge/work-in-progress label Jul 19, 2025

dhirajsb force-pushed the feat/mlflow-plugin branch 2 times, most recently from 60c43ad to 60ad248 Compare September 4, 2025 06:12

dhirajsb changed the title ~~[Draft] Add MLflow SDK support for Model Registry as a Tracking Store, fixes #1225~~ Add MLflow SDK support for Model Registry as a Tracking Store, fixes #1225 Sep 4, 2025

dhirajsb marked this pull request as ready for review September 4, 2025 06:13

google-oss-prow bot removed the do-not-merge/work-in-progress label Sep 4, 2025

pboyd reviewed Sep 4, 2025

View reviewed changes

jonburdo reviewed Sep 5, 2025

View reviewed changes

jonburdo approved these changes Sep 10, 2025

View reviewed changes

google-oss-prow bot assigned jonburdo Sep 10, 2025

google-oss-prow bot added lgtm and removed lgtm labels Sep 10, 2025

jonburdo approved these changes Sep 11, 2025

View reviewed changes

google-oss-prow bot added lgtm and removed lgtm labels Sep 11, 2025

jonburdo approved these changes Sep 11, 2025

View reviewed changes

google-oss-prow bot added the lgtm label Sep 11, 2025

dhirajsb added 2 commits September 11, 2025 13:14

feat: mlflow tracking store plugin for Kubeflow model registry, wip

a28f1c9

Signed-off-by: Dhiraj Bokde <[email protected]>

fix: tracking store plugin with model artifact support, updated mock …

225f5e6

…tests for mlflow plugin Signed-off-by: Dhiraj Bokde <[email protected]>

dhirajsb and others added 19 commits September 11, 2025 13:14

chore: ruff format

1b3149f

Signed-off-by: Dhiraj Bokde <[email protected]>

fix: fixed big refactor, updated tests for missing params, etc.

bb385d2

Signed-off-by: Dhiraj Bokde <[email protected]>

fix: fix missing params in search_logged_models

61bb526

Signed-off-by: Dhiraj Bokde <[email protected]>

fix: rename store.py to tracking_store.py

6e909cd

Signed-off-by: Dhiraj Bokde <[email protected]>

fix: fix run status to use str, fix default view type for experiments…

abad6b3

… and runs search Signed-off-by: Dhiraj Bokde <[email protected]>

fix: add missing store methods not defined by mlflow in base Abstract…

7e344e8

…Store class, add tests for new methods, replace print statements in tests Signed-off-by: Dhiraj Bokde <[email protected]>

feat: add support for custom CA cert bundles

d0b50ac

Signed-off-by: Dhiraj Bokde <[email protected]>

feat: update mlflow plugin to use embedmd datastore

886704f

Signed-off-by: Dhiraj Bokde <[email protected]>

Replace SSL/TLS with TLS in README

23321ee

Co-authored-by: Paul Boyd <[email protected]> Signed-off-by: Dhiraj Bokde <[email protected]>

Replace SSL with TLS in README

6f29be1

Co-authored-by: Paul Boyd <[email protected]> Signed-off-by: Dhiraj Bokde <[email protected]>

Replace SSL with TLS in exception message

b02ace3

Co-authored-by: Paul Boyd <[email protected]> Signed-off-by: Dhiraj Bokde <[email protected]>

fix: renamed plugin module to model-registry-mlflow and package to mo…

8ac52e9

…del_registry_mlflow to match pypi and python pkg naming conventions, removed redunant code in tests, general code cleanup Signed-off-by: Dhiraj Bokde <[email protected]>

fix: move plugin pkg to src, switch from hatchling to uv_build for bu…

82a078d

…ild system Signed-off-by: Dhiraj Bokde <[email protected]>

fix: cache auth token to reduce file read overhead

d420a50

Signed-off-by: Dhiraj Bokde <[email protected]>

chore: update pyproject metadata

02530dc

Signed-off-by: Dhiraj Bokde <[email protected]>

fix: fix pyproject uv info

2fef6bd

Signed-off-by: Dhiraj Bokde <[email protected]>

fix: make k8s token file refresh detection robust

876622f

Signed-off-by: Dhiraj Bokde <[email protected]>

fix: made k8s token cache class singleton, fixed tests

1348e4a

Signed-off-by: Dhiraj Bokde <[email protected]>

fix: use temp token file instead of mocking file io

465cbce

Signed-off-by: Dhiraj Bokde <[email protected]>

dhirajsb force-pushed the feat/mlflow-plugin branch from f75fd6d to 465cbce Compare September 11, 2025 20:15

google-oss-prow bot removed the lgtm label Sep 11, 2025

dhirajsb added 2 commits September 11, 2025 13:58

fix: fix clean target

7b6bde3

Signed-off-by: Dhiraj Bokde <[email protected]>

fix: typo in README

2401adb

Signed-off-by: Dhiraj Bokde <[email protected]>

github-actions bot added the lifecycle/stale label Dec 12, 2025

github-actions bot closed this Jan 1, 2026

Add MLflow SDK support for Model Registry as a Tracking Store, fixes #1225 #1337

Add MLflow SDK support for Model Registry as a Tracking Store, fixes #1225 #1337

Uh oh!

Conversation

dhirajsb commented Jul 19, 2025

Description

How Has This Been Tested?

Merge criteria:

Uh oh!

google-oss-prow bot commented Jul 19, 2025

Uh oh!

dhirajsb commented Jul 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dhirajsb commented Jul 19, 2025

Uh oh!

dhirajsb commented Sep 4, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jonburdo left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jonburdo left a comment

Choose a reason for hiding this comment

Uh oh!

jonburdo left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dhirajsb commented Sep 11, 2025

Uh oh!

jonburdo left a comment

Choose a reason for hiding this comment

Uh oh!

google-oss-prow bot commented Sep 11, 2025

Uh oh!

dhirajsb commented Sep 11, 2025

Uh oh!

rareddy commented Sep 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Dec 12, 2025

Uh oh!

github-actions bot commented Jan 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

dhirajsb commented Jul 19, 2025 •

edited

Loading

rareddy commented Sep 11, 2025 •

edited

Loading