WIP: Add workspace CRUD support#11
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the WalkthroughAdds workspace multi-tenancy: new environment variables and CLI flags, server endpoints and helpers, workspace store registry and providers (SQLAlchemy/REST), MlflowClient workspace methods, Workspace.to_dict(), name validation, utilities for workspace URI resolution, DB schema integration, and comprehensive tests. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant User
participant MlflowClient
participant WorkspaceProviderClient
participant Registry as WorkspaceStoreRegistry
participant Store as WorkspaceStore
participant Backend
User->>MlflowClient: list_workspaces()
MlflowClient->>MlflowClient: resolve workspace_uri
MlflowClient->>WorkspaceProviderClient: ensure provider (cached)
WorkspaceProviderClient->>Registry: get_workspace_store(workspace_uri)
Registry->>Registry: resolve scheme & cache
alt SQL scheme
Registry->>Store: instantiate SqlAlchemyStore
else HTTP(s) scheme
Registry->>Store: instantiate RestWorkspaceStore
end
WorkspaceProviderClient->>Store: list_workspaces()
Store->>Backend: query or HTTP request
Backend-->>Store: workspaces
Store-->>WorkspaceProviderClient: Workspace objects
WorkspaceProviderClient-->>MlflowClient: Workspace objects
MlflowClient-->>User: return list
sequenceDiagram
autonumber
participant Client as MLflow Server
participant Handler as handlers.py
participant Helper as workspace_helpers.py
participant Registry as WorkspaceStoreRegistry
participant Store as WorkspaceStore
Client->>Handler: GET /api/2.0/mlflow/workspaces
Handler->>Handler: check feature flag
alt disabled
Handler-->>Client: FEATURE_DISABLED error
else enabled
Handler->>Helper: _get_workspace_store()
Helper->>Registry: get_workspace_store(resolved_uri)
Registry-->>Helper: WorkspaceStore
Handler->>Store: list_workspaces()
Store-->>Handler: workspaces
Handler-->>Client: JSON response
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Areas requiring extra attention:
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
mprahl
force-pushed
the
workspace-crud-api
branch
from
4796b92 to
ad8dca3
Compare
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
mprahl
force-pushed
the
workspace-crud-api
branch
4 times, most recently
from
17c0a79 to
b9fabb9
Compare
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
mlflow/server/handlers.py (1)
31-65: Use an HTTP status code rather thandatabricks_pb2.FEATURE_DISABLEDin the guardThe review comment is accurate.
_disable_if_workspaces_disabledpassesdatabricks_pb2.FEATURE_DISABLED(which equals1005) directly as the HTTP status code toResponse(). HTTP status codes must be in the range 100–599; 1005 is invalid and will produce a malformed HTTP response.In contrast, the existing
_disable_unless_serve_artifactsguard (lines 751–761) correctly uses503. Additionally, the helper_workspace_not_supported(line 805) demonstrates the correct pattern by raisingMlflowException(message, FEATURE_DISABLED), letting the error handler convert the Protobuf error code to a proper HTTP status.The decorator should either:
- Return
Response(..., 503)directly, or- Raise
MlflowException(..., error_code=FEATURE_DISABLED)to leverage the existing exception-to-HTTP-response mapping.
🧹 Nitpick comments (13)
mlflow/tracking/_workspace/client.py (1)
35-36: Potentially unnecessary list conversion.The return type annotation
list[Workspace]suggests the provider already returns a list, makinglist(...)conversion potentially redundant. Iflist_workspaces()returns an iterator or generator, this is correct; otherwise, it's unnecessary overhead.Verify whether the provider returns a list or an iterator. If it's already a list, simplify:
- return list(self.provider.list_workspaces()) + return self.provider.list_workspaces()tests/tracking/test_client_workspace.py (1)
1-52: Good focused coverage of workspace URI resolution inMlflowClientThis test cleanly validates the three key paths (env var, explicit
workspace_uri, and fallback to tracking URI) by patching the internal clients and resolvers and recording constructor arguments. It’s isolated from real backends and should be robust to internal changes as long asWorkspaceProviderClientis still constructed with aworkspace_uri. If you want to tighten it further, you could also assertrecorded["tracking_uri"]in each scenario to guard against regressions in tracking URI resolution, but that’s optional.mlflow/server/workspace_helpers.py (2)
15-40: Clarify singleton caching semantics for_get_workspace_store
_get_workspace_storememoizes the first resolved workspace store in_workspace_storeand returns it for all subsequent calls, ignoring laterworkspace_uri/tracking_uriarguments. This is perfectly fine if the server is designed to use a single workspace store per process (which seems to be the intent), but it makes the parameters effectively “first-call only”. Consider documenting this in the function docstring or renaming the helper to make the singleton nature explicit, so future callers don’t assume that passing a different URI will have an effect.
15-15: Confirm Python version compatibility for PEP 604 union type hintsThe signature
def _get_workspace_store(workspace_uri: str | None = None, tracking_uri: str | None = None)uses theX | Yunion syntax. Withoutfrom __future__ import annotationsin this module, that requires Python ≥ 3.10 to import successfully. If MLflow still supports older Python versions (e.g., 3.8/3.9) in this branch, you’ll need to either addfrom __future__ import annotationsat the top of the file or switch toOptional[str]fromtyping. Please confirm against the project’s supported Python versions.mlflow/tracking/client.py (1)
45-46: Workspace client wiring and URI resolution look coherent
- Importing
Workspaceand wiringWorkspaceProviderClientviaworkspace_utils.resolve_workspace_uri()gives a clear, consistent precedence: explicit arg → global / env → tracking URI.- Lazy initialization in
_get_workspace_client()with a specializedUnsupportedWorkspaceStoreURIExceptiontranslated to aMlflowExceptionwithFEATURE_DISABLEDmatches the existing registry client pattern and should produce good UX.- The new
get_workspace_uri()helper and the CRUD methods (list_workspace*,create_*,get_*,update_*,delete_*) are thin, type-annotated pass-throughs, which is exactly what you want at this layer.Unless you plan to treat workspaces as experimental, these additions look ready as-is. If workspaces are still considered preview, consider decorating the new public methods with
@experimentalfor consistency, but that’s optional and policy-driven.Please confirm whether workspace APIs are meant to be GA or should be annotated as experimental to match other new surfaces (e.g., datasets, prompts).
Also applies to: 110-113, 203-208, 211-237, 242-244, 250-257, 295-339
mlflow/tracking/_workspace/registry.py (1)
18-29: Workspace store registry is solid; consider aligning error code with other Unsupported exceptions*
- The registry design (entry-point group, default DB/REST registrations, LRU-cached
get_storewith a build lock) is consistent with existing MLflow store registries and looks correct.UnsupportedWorkspaceStoreURIExceptionprovides a clear message and exposessupported_uri_schemes, which is useful for callers likeMlflowClient._get_workspace_client().One small consistency suggestion: other “unsupported store URI” paths (e.g., model registry) typically use
FEATURE_DISABLEDas the error code. Here the custom exception usesINVALID_PARAMETER_VALUE, while callers re-wrap toFEATURE_DISABLED. If you expect external callers to catchUnsupportedWorkspaceStoreURIExceptiondirectly, you might want to switch itserror_codetoFEATURE_DISABLEDto match the rest of the API surface; otherwise, current behavior is acceptable.Also applies to: 31-63, 66-79, 82-92, 94-112
mlflow/cli/__init__.py (1)
23-28: Workspace server flags and envvar wiring are consistent; minor UX edge case
- Importing
MLFLOW_ENABLE_WORKSPACES/MLFLOW_WORKSPACE_URIand exposing--enable-workspaces/--workspace-store-urivia their.nameattributes matches the existing experiment CLI patterns.- The
server()logic correctly:
- Forces
MLFLOW_ENABLE_WORKSPACES=truewhen workspaces are enabled, and- Propagates
workspace_store_uriintoMLFLOW_WORKSPACE_URIonly when workspaces are enabled, otherwise warning and ignoring it.One UX edge case: when
MLFLOW_WORKSPACE_URIis set only via environment variable (andMLFLOW_ENABLE_WORKSPACESis left false),workspace_store_uriwill still be populated by click and you’ll emit the “Ignoring --workspace-store-uri” message even though the user didn’t pass the flag explicitly. If that’s undesirable, you could checkctx.get_parameter_source("workspace_store_uri")to distinguish CLI vs envvar and adjust the message accordingly, but behavior is otherwise sound.Also applies to: 466-483, 506-508, 553-563
mlflow/tracking/_workspace/utils.py (1)
5-18: Tighten resolve_workspace_uri truthiness checks to match type hintsThe override/envvar handling looks good and the precedence is clear. One small refinement:
resolve_workspace_uri()currently uses truthiness checks:
if workspace_uri:if configured_uri:Given the annotated types are
str | None, it’s safer and less surprising to treat an empty string as an explicit (albeit likely invalid) value rather than falling through to the tracking URI. You can switch to explicitis not Nonechecks:def resolve_workspace_uri( workspace_uri: str | None = None, tracking_uri: str | None = None ) -> str | None: @@ - if workspace_uri: + if workspace_uri is not None: return workspace_uri @@ - configured_uri = get_workspace_uri() - if configured_uri: - return configured_uri + configured_uri = get_workspace_uri() + if configured_uri is not None: + return configured_uriThis better matches the type hints and avoids subtle fallback behavior if an empty string is ever passed in.
Also applies to: 21-29, 31-52
mlflow/store/workspace/rest_store.py (1)
12-98: REST workspace store behavior looks consistent; note 201/empty-body fallback semanticsThe overall implementation (list/get/create/update/delete/get_default) is consistent with other REST stores and correctly uses
http_request+verify_rest_response, URL-escapes workspace names, and tolerates empty JSON bodies by fabricating aWorkspacefrom the request data on 201 responses.Only behavioral nuance to be aware of: for
create_workspaceandupdate_workspace, if the server returns a 2xx with an empty body, the store will fall back to the client-sideWorkspacevalues instead of failing. That seems reasonable, but if the REST provider later adds extra server-populated fields, they won’t be surfaced in that edge case.tests/store/workspace/test_sqlalchemy_store.py (1)
13-160: Good coverage and realistic expectations for SQLAlchemy workspace storeThe fixture and tests exercise the full CRUD surface plus default handling and
resolve_artifact_root:
- Fixture correctly initializes full MLflow schema, seeds the default workspace, and disposes both engines.
- Tests validate both the error messages and
error_codevalues (RESOURCE_DOES_NOT_EXIST,RESOURCE_ALREADY_EXISTS,INVALID_STATE), which tightly couples toSqlAlchemyStorebehavior.test_resolve_artifact_root_returns_defaultexplicitly codifies thatresolve_artifact_rootignoresworkspace_nameand just returns(default_root, True); this is an intentional contract you’ll want to remember if you later add per-workspace artifact roots.Looks solid overall.
mlflow/store/workspace/abstract_store.py (1)
12-129: Abstract interface and name validator are well-structured; consider DRY’ing validationThe
AbstractStoreAPI surface is clear, and the defaultresolve_artifact_rootplusget_default_workspacecontract is straightforward and aligns with the SQL/REST implementations.
WorkspaceNameValidatorenforces length, pattern, and reserved names cleanly and exposes bothis_validandvalidate. There is some duplication between the two (the same checks are reimplemented), so if this grows further it may be worth centralizing the logic (e.g., haveis_validcall a shared helper andvalidatewrap that with error messages). Not urgent as-is.mlflow/server/handlers.py (1)
808-877: Workspace CRUD handlers are correct; consider minor polish on validation and responsesThe workspace handlers follow existing patterns:
- All are wrapped in
@catch_mlflow_exceptionand gated by_disable_if_workspaces_disabled, so store-levelMlflowExceptions get serialized consistently._create_workspace_handler:
- Validates presence of
"name"and runsWorkspaceNameValidator.validate(name).- Correctly maps
NotImplementedErrorfrom the store to aFEATURE_DISABLED-codedMlflowExceptionvia_workspace_not_supported.- Returns
201with a JSON body ofworkspace.to_dict()._update_workspace_handler:
- Rejects empty payloads and unknown keys.
- Enforces that only
"description"is allowed and safely accessespayload["description"]after that check._delete_workspace_handlerand_get_workspace_handlersimply delegate to the store, withNotImplementedErrorfrom delete mapped to feature-disabled errors.Two optional tweaks you might consider:
- Apply
WorkspaceNameValidator.validate(workspace_name)in the path-parameter handlers (_get_workspace_handler,_update_workspace_handler,_delete_workspace_handler) for symmetry with create, unless you expect to support legacy names that don’t pass the validator.- For create, optionally add a
Locationheader pointing to/api/2.0/mlflow/workspaces/<workspace_name>(and its AJAX variant), which can simplify client code.mlflow/store/workspace/sqlalchemy_store.py (1)
24-100: SQL workspace store implementation is solid; import path difference is the only thing to double-checkThe
SqlAlchemyStoreimplementation looks correct overall:
- Engine/session management:
- Uses
create_sqlalchemy_engine_with_retryand a managed session maker fromdb_utils, consistent with other MLflow SQL stores.- CRUD semantics:
list_workspacesorders bynameand converts ORM rows viato_mlflow_entity().get_workspaceand_get_workspacecentralize not-found behavior and raiseMlflowException("Workspace '<name>' not found", RESOURCE_DOES_NOT_EXIST), matching the tests.create_workspacecatchesIntegrityErrorand maps it toRESOURCE_ALREADY_EXISTSwith a clear message.update_workspaceonly mutatesdescription, flushes, logs, and returns the updated entity.delete_workspacecorrectly rejects deletion ofDEFAULT_WORKSPACE_NAMEwithINVALID_STATEbefore touching the DB.get_default_workspaceandresolve_artifact_rootline up with the tests’ expectations.The only thing worth double-checking is the import:
from mlflow.store.workspace.dbmodels import SqlWorkspacewhereas tests import
SqlWorkspacefrommlflow.store.workspace.dbmodels.models. Ifdbmodels/__init__.pydoes not re-exportSqlWorkspace, this will raise anImportError. If it does, then everything is fine, but it’d be good to keep the import style consistent across modules.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (22)
docs/api_reference/api_inventory.txt(6 hunks)mlflow/cli/__init__.py(4 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/environment_variables.py(2 hunks)mlflow/server/handlers.py(7 hunks)mlflow/server/workspace_helpers.py(1 hunks)mlflow/store/db/utils.py(2 hunks)mlflow/store/workspace/__init__.py(1 hunks)mlflow/store/workspace/abstract_store.py(1 hunks)mlflow/store/workspace/rest_store.py(1 hunks)mlflow/store/workspace/sqlalchemy_store.py(1 hunks)mlflow/tracking/_workspace/__init__.py(1 hunks)mlflow/tracking/_workspace/client.py(1 hunks)mlflow/tracking/_workspace/registry.py(1 hunks)mlflow/tracking/_workspace/utils.py(1 hunks)mlflow/tracking/client.py(5 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/store/workspace/test_sqlalchemy_store.py(1 hunks)tests/store/workspace/test_workspace_validator.py(1 hunks)tests/tracking/test_client_workspace.py(1 hunks)tests/tracking/test_workspace_registry.py(1 hunks)tests/tracking/test_workspace_utils.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (22)
- GitHub Check: windows (1)
- GitHub Check: windows (4)
- GitHub Check: windows (3)
- GitHub Check: pyfunc (3)
- GitHub Check: pyfunc (1)
- GitHub Check: pyfunc (2)
- GitHub Check: windows (2)
- GitHub Check: pyfunc (4)
- GitHub Check: evaluate (2)
- GitHub Check: evaluate (1)
- GitHub Check: build (dev)
- GitHub Check: genai
- GitHub Check: build
- GitHub Check: python (2)
- GitHub Check: python (1)
- GitHub Check: python (3)
- GitHub Check: database
- GitHub Check: models (2)
- GitHub Check: models (1)
- GitHub Check: python-skinny
- GitHub Check: flavors
- GitHub Check: java
🔇 Additional comments (12)
docs/api_reference/api_inventory.txt (1)
43-43: LGTM!The API inventory correctly documents the new workspace management methods on
MlflowClientand theWorkspace.to_dictserialization method. The additions are consistent with the broader workspace CRUD feature.Also applies to: 65-65, 89-90, 99-99, 154-154, 538-538
mlflow/utils/workspace_utils.py (1)
1-1: LGTM!The updated docstring "Workspace-related utility constants" is clearer and more accurate than the previous description.
mlflow/store/db/utils.py (1)
68-68: LGTM!The workspace table integration follows the established pattern for schema validation. The import and table name addition properly extend the canonical schema to include workspace support.
Also applies to: 109-109
mlflow/entities/workspace.py (1)
15-16: LGTM!The
to_dictmethod provides straightforward serialization of the Workspace entity. The return type correctly reflects that description can be None.tests/store/workspace/test_workspace_validator.py (1)
7-41: LGTM!The test coverage for
WorkspaceNameValidatoris comprehensive, covering:
- Valid cases: length boundaries, allowed patterns, numeric names
- Invalid cases: length violations, reserved words, case sensitivity, invalid patterns
- Both
is_validandvalidatecode pathsThe parameterized approach makes the tests clear and maintainable.
tests/tracking/test_workspace_registry.py (1)
14-40: LGTM!The workspace store registry tests provide good coverage:
- Proper test isolation via cache-clearing fixture
- SQLAlchemy, REST, and error cases are validated
- Resource cleanup with
store._engine.dispose()prevents test pollutiontests/tracking/test_workspace_utils.py (1)
7-38: LGTM!The workspace URI resolution tests are well-structured:
- Proper isolation with autouse fixture clearing both state and environment
- Complete coverage of the resolution precedence chain
- Clean use of monkeypatch for environment variable testing
mlflow/environment_variables.py (2)
109-112: Workspace URI environment variable is correctly introduced
MLFLOW_WORKSPACE_URIfollows existing env-var patterns (public name,strtype,Nonedefault) and the docstring matches the expected “fallback to tracking URI” behavior implemented in the resolver utilities. No changes needed.
442-445: Feature flag for workspaces is consistent with usage
MLFLOW_ENABLE_WORKSPACESis defined as a boolean env var with a sensible default (False) and is used to gate server-side workspace APIs (seemlflow.server.workspace_helpers._get_workspace_store). This aligns with the PR’s multi‑tenancy feature-flagging goals.mlflow/store/workspace/__init__.py (1)
1-11: Workspace store facade re-exports look correctThe re-exports cleanly surface
Workspace,AbstractStore, andRestWorkspaceStore, and__all__matches the imported symbols. This is a good, minimal public API aggregation point.mlflow/tracking/_workspace/__init__.py (1)
1-19: Workspace tracking API aggregation is consistent and minimalThe module re-exports the key workspace tracking primitives (
WorkspaceProviderClient, registry, store getter, and URI helpers) and__all__matches them exactly. This keeps callers from depending on deeper module paths and looks like the right public surface.mlflow/server/handlers.py (1)
880-892: Endpoint registration for workspace routes is consistent with existing routing
_workspace_endpointscorrectly:
- Uses
_get_paths("/mlflow/workspaces")and_get_paths("/mlflow/workspaces/<workspace_name>")so both REST (/api/2.0/...) and AJAX (/ajax-api/...) paths are registered.- Associates:
GET/POSTwith the collection path, andGET/PATCH/DELETEwith the item path.And
get_endpointssimply appends_workspace_endpoints()to the existing service endpoint lists, matching the pattern used elsewhere in this module.This wiring looks correct and should make the new workspace APIs available through both REST and AJAX routes without interfering with protobuf-derived endpoints.
Also applies to: 3939-3946
mprahl
force-pushed
the
workspace-crud-api
branch
5 times, most recently
from
65bd3d6 to
93b4ac7
Compare
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
mprahl
force-pushed
the
workspace-crud-api
branch
3 times, most recently
from
f38013a to
70edd1d
Compare
- Introduce the workspace store abstraction (SQLAlchemy + REST), name validation, and registry/URI utilities with CLI flags and env vars. - Expose workspace CRUD endpoints in the Flask server with proper guards/feature flag wiring plus MlflowClient surface methods. - Add the Workspace entity, tests for the stores/registry/client, and docs entries so the new multi-tenant provider can be exercised end-to-end. Signed-off-by: mprahl <mprahl@users.noreply.github.com>
mprahl
force-pushed
the
workspace-crud-api
branch
from
70edd1d to
3d850ff
Compare
1 participant
🛠 DevTools 🛠
Install mlflow from this PR
For Databricks, use the following command:
Related Issues/PRs
#xxxWhat 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?
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.
Summary by CodeRabbit
New Features
Tests
✏️ Tip: You can customize this high-level summary in your review settings.