Conversation
There was a problem hiding this comment.
Pull request overview
This pull request adds tag support for Gateway Endpoints, enabling key-value pair metadata to be associated with endpoints for categorization and filtering purposes. The PR also includes a comprehensive refactoring to rename endpoint-related methods from *_endpoint to *_gateway_endpoint for consistency.
Key changes:
- Introduces a new
GatewayEndpointTagentity and database table for storing endpoint tags - Adds
set_gateway_endpoint_taganddelete_gateway_endpoint_tagmethods across all store implementations - Renames all endpoint-related methods to include "gateway_" prefix for API consistency
- Implements comprehensive test coverage for tag operations including cascading deletes
Reviewed changes
Copilot reviewed 13 out of 15 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
mlflow/store/db_migrations/versions/1bd49d398cd23_add_secrets_tables.py |
Adds endpoint_tags table with foreign key constraints and cascading deletes |
mlflow/store/tracking/dbmodels/models.py |
Implements SqlGatewayEndpointTag model with relationship to SqlGatewayEndpoint |
mlflow/entities/gateway_endpoint.py |
Adds GatewayEndpointTag entity class and integrates tags into GatewayEndpoint |
mlflow/entities/__init__.py |
Exports the new GatewayEndpointTag class |
mlflow/protos/service.proto |
Defines protobuf messages for tag operations and updates RPC method names |
mlflow/protos/service_pb2.pyi |
Updates type stubs for new tag-related messages |
mlflow/store/tracking/gateway/abstract_mixin.py |
Adds abstract methods for tag operations and renames endpoint methods |
mlflow/store/tracking/gateway/sqlalchemy_mixin.py |
Implements tag storage logic using SQLAlchemy with merge/delete operations |
mlflow/store/tracking/gateway/rest_mixin.py |
Implements REST API client for tag operations |
mlflow/server/handlers.py |
Adds HTTP handlers for tag set/delete and updates endpoint method names |
tests/tracking/test_rest_tracking.py |
Adds integration tests for tag operations via REST API |
tests/store/tracking/test_gateway_sql_store.py |
Adds comprehensive unit tests for tag storage including edge cases |
docs/api_reference/api_inventory.txt |
Updates API documentation inventory with new tag-related methods |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
Documentation preview for 8236350 is available at: More info
|
Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>
Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>
41baf4f to
fb2dda8
Compare
Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>
Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>
| from alembic import op | ||
|
|
||
| # revision identifiers, used by Alembic. | ||
| revision = "1b49d398cd23" |
There was a problem hiding this comment.
Oh, so was the revision id incorrect?
There was a problem hiding this comment.
Some weird typo when I was editing... I probably didn't realize and hit delete while the cursor was on that line.
| ), | ||
| ) | ||
|
|
||
| op.create_table( |
There was a problem hiding this comment.
What happens for people who already have revision '1bd49d398cd23' in the database? Do they need to rollback and migrate again? I'm wondering if we should separate a migration file or not in this case.
There was a problem hiding this comment.
Still safe to edit since we haven't released this yet :)
Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>
1196801 to
c930b4d
Compare
Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>
🛠 DevTools 🛠
Install mlflow from this PR
For Databricks, use the following command:
Related Issues/PRs
#xxxWhat changes are proposed in this pull request?
Add Gateway Endpoint tags support (guard against potential disambiguation for now, needed eventually for AI Gateway).
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?
Components
area/tracking: Tracking Service, tracking client APIs, autologgingarea/models: MLmodel format, model serialization/deserialization, flavorsarea/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registryarea/scoring: MLflow Model server, model deployment tools, Spark UDFsarea/evaluation: MLflow model evaluation features, evaluation metrics, and evaluation workflowsarea/gateway: MLflow AI Gateway client APIs, server, and third-party integrationsarea/prompts: MLflow prompt engineering features, prompt templates, and prompt managementarea/tracing: MLflow Tracing features, tracing APIs, and LLM tracing functionalityarea/projects: MLproject format, project running backendsarea/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev serverarea/build: Build and test infrastructure for MLflowarea/docs: MLflow documentation pagesHow 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" sectionrn/breaking-change- The PR will be mentioned in the "Breaking Changes" sectionrn/feature- A new user-facing feature worth mentioning in the release notesrn/bug-fix- A user-facing bug fix worth mentioning in the release notesrn/documentation- A user-facing documentation change worth mentioning in the release notesShould this PR be included in the next patch release?
Yesshould be selected for bug fixes, documentation updates, and other small changes.Noshould 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?
Bug fixes, doc updates and new features usually go into minor releases.
Bug fixes and doc updates usually go into patch releases.