Add workspace CRUD support

Author: mprahl

mlflow/environment_variables.py

@@ -106,6 +106,10 @@ def get(self):
 #: (default: ``None``)
 MLFLOW_REGISTRY_URI = _EnvironmentVariable("MLFLOW_REGISTRY_URI", str, None)
 
+#: Specifies the workspace provider backend URI.
+#: Defaults to the tracking URI when unset.
+MLFLOW_WORKSPACE_URI = _EnvironmentVariable("MLFLOW_WORKSPACE_URI", str, None)
+
 #: Specifies the ``dfs_tmpdir`` parameter to use for ``mlflow.spark.save_model``,
 #: ``mlflow.spark.log_model`` and ``mlflow.spark.load_model``. See
 #: https://www.mlflow.org/docs/latest/python_api/mlflow.spark.html#mlflow.spark.save_model
@@ -435,6 +439,10 @@ def get(self):
     "MLFLOW_ENABLE_DBFS_FUSE_ARTIFACT_REPO", True
 )
 
+#: Enables workspace-aware multi-tenancy features on the MLflow server.
+#: (default: ``False``)
+MLFLOW_ENABLE_WORKSPACES = _BooleanEnvironmentVariable("MLFLOW_ENABLE_WORKSPACES", False)
+
 #: Specifies whether or not to use UC Volume FUSE mount to store artifacts on Databricks
 #: (default: ``True``)
 MLFLOW_ENABLE_UC_VOLUME_FUSE_ARTIFACT_REPO = _BooleanEnvironmentVariable(

mlflow/server/handlers.py

@@ -28,6 +28,7 @@
     Param,
     RunTag,
     ViewType,
+    Workspace,
 )
 from mlflow.entities.logged_model import LoggedModel
 from mlflow.entities.logged_model_input import LoggedModelInput
@@ -57,7 +58,9 @@
 from mlflow.protos import databricks_pb2
 from mlflow.protos.databricks_pb2 import (
     BAD_REQUEST,
+    FEATURE_DISABLED,
     INVALID_PARAMETER_VALUE,
+    INVALID_STATE,
     RESOURCE_DOES_NOT_EXIST,
 )
 from mlflow.protos.mlflow_artifacts_pb2 import (
@@ -183,6 +186,7 @@
     WebhookService,
 )
 from mlflow.server.validation import _validate_content_type
+from mlflow.server.workspace_helpers import _get_workspace_store, _workspaces_enabled_flag
 from mlflow.store.artifact.artifact_repo import MultipartUploadMixin
 from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
 from mlflow.store.db.db_types import DATABASE_ENGINES
@@ -191,6 +195,7 @@
 from mlflow.store.model_registry.rest_store import RestStore as ModelRegistryRestStore
 from mlflow.store.tracking.abstract_store import AbstractStore as AbstractTrackingStore
 from mlflow.store.tracking.databricks_rest_store import DatabricksTracingRestStore
+from mlflow.store.workspace.abstract_store import WorkspaceNameValidator
 from mlflow.tracing.utils.artifact_utils import (
     TRACE_DATA_FILE_NAME,
     get_artifact_uri_for_trace,
@@ -779,6 +784,122 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def _disable_if_workspaces_disabled(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        if not _workspaces_enabled_flag():
+            return Response(
+                (
+                    f"Endpoint: {request.url_rule} disabled because the server is running "
+                    "without multi-tenancy support. To enable workspace functionality, run "
+                    "`mlflow server` with `--enable-workspaces`"
+                ),
+                databricks_pb2.FEATURE_DISABLED,
+            )
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def _workspace_to_response_payload(workspace: Workspace) -> dict[str, str | None]:
+    return {"name": workspace.name, "description": workspace.description}
+
+
+def _workspace_response(workspace: Workspace, status: int = 200) -> Response:
+    response = jsonify(_workspace_to_response_payload(workspace))
+    response.status_code = status
+    return response
+
+
+
+def _workspace_not_supported(message: str) -> MlflowException:
+    return MlflowException(message, FEATURE_DISABLED)
+
+
+@catch_mlflow_exception
+@_disable_if_workspaces_disabled
+def _list_workspaces_handler():
+    workspaces = _get_workspace_store().list_workspaces(request)
+    return jsonify({"workspaces": [_workspace_to_response_payload(ws) for ws in workspaces]})
+
+
+@catch_mlflow_exception
+@_disable_if_workspaces_disabled
+def _create_workspace_handler():
+    payload = _get_request_json(request) or {}
+    name = payload.get("name")
+    if not name:
+        raise MlflowException.invalid_parameter_value("Workspace name must be provided")
+
+    WorkspaceNameValidator.validate(name)
+
+    description = payload.get("description")
+    store = _get_workspace_store()
+    try:
+        workspace = store.create_workspace(Workspace(name=name, description=description), request)
+    except NotImplementedError:
+        raise _workspace_not_supported("Workspace creation is not supported by this provider")
+
+    return _workspace_response(workspace, status=201)
+
+
+@catch_mlflow_exception
+@_disable_if_workspaces_disabled
+def _get_workspace_handler(workspace_name: str):
+    workspace = _get_workspace_store().get_workspace(workspace_name, request)
+    return _workspace_response(workspace)
+
+
+@catch_mlflow_exception
+@_disable_if_workspaces_disabled
+def _update_workspace_handler(workspace_name: str):
+    payload = _get_request_json(request) or {}
+
+    if not payload.keys():
+        raise MlflowException.invalid_parameter_value("Workspace update must have at least one key")
+
+    invalid_keys = payload.keys() - {"description"}
+    if invalid_keys:
+        raise MlflowException.invalid_parameter_value(
+            f"Workspace update had the following invalid keys: {', '.join(invalid_keys)}"
+        )
+
+    store = _get_workspace_store()
+    try:
+        workspace = store.update_workspace(
+            Workspace(name=workspace_name, description=payload["description"]), request
+        )
+    except NotImplementedError:
+        raise _workspace_not_supported("Workspace updates are not supported by this provider")
+
+    return _workspace_response(workspace)
+
+
+@catch_mlflow_exception
+@_disable_if_workspaces_disabled
+def _delete_workspace_handler(workspace_name: str):
+    store = _get_workspace_store()
+    try:
+        store.delete_workspace(workspace_name, request)
+    except NotImplementedError:
+        raise _workspace_not_supported("Workspace deletion is not supported by this provider")
+    return Response(status=204)
+
+
+def _workspace_endpoints():
+    endpoints = []
+    for path in _get_paths("/mlflow/workspaces"):
+        endpoints.append((path, _list_workspaces_handler, ["GET"]))
+        endpoints.append((path, _create_workspace_handler, ["POST"]))
+
+    for path in _get_paths("/mlflow/workspaces/<workspace_name>"):
+        endpoints.append((path, _get_workspace_handler, ["GET"]))
+        endpoints.append((path, _update_workspace_handler, ["PATCH"]))
+        endpoints.append((path, _delete_workspace_handler, ["DELETE"]))
+
+    return endpoints
+
+
 @catch_mlflow_exception
 def get_artifact_handler():
     run_id = request.args.get("run_id") or request.args.get("run_uuid")
@@ -3830,6 +3951,7 @@ def get_endpoints(get_handler=get_handler):
         + get_service_endpoints(MlflowArtifactsService, get_handler)
         + get_service_endpoints(WebhookService, get_handler)
         + [(_add_static_prefix("/graphql"), _graphql, ["GET", "POST"])]
+        + _workspace_endpoints()
     )
 
 

mlflow/server/workspace_helpers.py

@@ -0,0 +1,47 @@
+import logging
+import os
+
+from mlflow.environment_variables import MLFLOW_ENABLE_WORKSPACES, MLFLOW_WORKSPACE_URI
+from mlflow.exceptions import MlflowException
+from mlflow.protos import databricks_pb2
+from mlflow.tracking._workspace import utils as workspace_utils
+from mlflow.tracking._workspace.registry import get_workspace_store
+
+_logger = logging.getLogger(__name__)
+
+_workspace_store = None
+
+
+def _workspaces_enabled_flag() -> bool:
+    return bool(MLFLOW_ENABLE_WORKSPACES.get())
+
+
+def _get_workspace_store(workspace_uri: str | None = None, tracking_uri: str | None = None):
+    if not _workspaces_enabled_flag():
+        raise MlflowException(
+            "Workspace APIs are not available: multi-tenancy is not enabled on this server",
+            databricks_pb2.FEATURE_DISABLED,
+        )
+
+    global _workspace_store
+    if _workspace_store is not None:
+        return _workspace_store
+
+    from mlflow.server import BACKEND_STORE_URI_ENV_VAR
+
+    resolved_tracking_uri = tracking_uri or os.environ.get(BACKEND_STORE_URI_ENV_VAR)
+    resolved_workspace_uri = workspace_utils.resolve_workspace_uri(
+        workspace_uri, tracking_uri=resolved_tracking_uri
+    )
+    if resolved_workspace_uri is None:
+        raise MlflowException.invalid_parameter_value(
+            "Workspace URI could not be resolved. Provide --workspace-store-uri or set "
+            f"{MLFLOW_WORKSPACE_URI.name}."
+        )
+
+    _workspace_store = get_workspace_store(workspace_uri=resolved_workspace_uri)
+    return _workspace_store
+
+
+__all__ = ["_get_workspace_store", "_workspaces_enabled_flag"]
+

mlflow/store/db/utils.py

@@ -46,6 +46,7 @@
     SqlRegisteredModelAlias,
     SqlRegisteredModelTag,
 )
+from mlflow.store.workspace.dbmodels.models import SqlWorkspace
 from mlflow.store.tracking.dbmodels.initial_models import Base as InitialBase
 from mlflow.store.tracking.dbmodels.models import (
     SqlDataset,
@@ -105,6 +106,7 @@ def _all_tables_exist(engine):
         SqlScorer.__tablename__,
         SqlScorerVersion.__tablename__,
         SqlJob.__tablename__,
+        SqlWorkspace.__tablename__,
     }
 
 

mlflow/store/workspace/__init__.py

@@ -0,0 +1,11 @@
+"""Public workspace store facade and re-exports."""
+
+from mlflow.entities.workspace import Workspace
+from mlflow.store.workspace.abstract_store import AbstractStore
+from mlflow.store.workspace.rest_store import RestWorkspaceStore
+
+__all__ = [
+    "Workspace",
+    "AbstractStore",
+    "RestWorkspaceStore",
+]

mlflow/store/workspace/abstract_store.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import re
+from abc import ABC, abstractmethod
+from typing import Iterable
+
+from mlflow.entities import Workspace
+from mlflow.exceptions import MlflowException
+from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME
+
+
+class AbstractStore(ABC):
+    """Interface for resolving and managing workspaces in the tracking server."""
+
+    @abstractmethod
+    def list_workspaces(self, request) -> Iterable[Workspace]:
+        """
+        Return the workspaces visible to the current request context.
+
+        Implementations may inspect the request (e.g., for authN/Z context) to
+        determine which workspaces to expose.
+        """
+
+    @abstractmethod
+    def get_workspace(self, workspace_name: str, request) -> Workspace:
+        """
+        Validate access to ``workspace_name`` and return its metadata.
+
+        Implementations should raise ``MlflowException`` with
+        ``RESOURCE_DOES_NOT_EXIST`` if the workspace cannot be found and
+        ``PERMISSION_DENIED`` if the caller lacks access.
+        """
+
+    def create_workspace(self, workspace: Workspace, request) -> Workspace:
+        """Provision a new workspace.
+
+        Raises ``NotImplementedError`` when the active provider is read-only.
+        Implementations should raise ``MlflowException`` with
+        ``RESOURCE_ALREADY_EXISTS`` when the workspace already exists or
+        ``INVALID_PARAMETER_VALUE`` when validation fails.
+        """
+
+        raise NotImplementedError
+
+    def update_workspace(self, workspace: Workspace, request) -> Workspace:
+        """Update metadata for an existing workspace."""
+
+        raise NotImplementedError
+
+    def delete_workspace(self, workspace_name: str, request) -> None:
+        """Delete an existing workspace."""
+
+        raise NotImplementedError
+
+    def get_default_workspace(self, request) -> Workspace:
+        """
+        Return the workspace to select when none is explicitly supplied.
+
+        Implementations that require an explicit workspace should raise an
+        ``MlflowException`` with ``INVALID_PARAMETER_VALUE``.
+        """
+
+        raise NotImplementedError
+
+    def resolve_artifact_root(
+        self, default_artifact_root: str, workspace_name: str | None = None
+    ) -> tuple[str, bool]:
+        """
+        Allow a provider to customize artifact storage roots per workspace.
+
+        Returns:
+            A tuple ``(root, append_workspace_prefix)`` where ``root`` is the base artifact
+            location to use for the workspace, and ``append_workspace_prefix`` controls whether
+            MLflow should append the ``/workspaces/<workspace_name>`` suffix automatically.
+        """
+
+        return default_artifact_root, True
+
+
+class WorkspaceNameValidator:
+    _PATTERN = r"^[a-z0-9][-a-z0-9]*[a-z0-9]$"
+    _MIN_LENGTH = 2
+    _MAX_LENGTH = 63
+    _RESERVED = {DEFAULT_WORKSPACE_NAME, "workspaces", "api", "ajax-api", "static-files"}
+
+    @classmethod
+    def pattern(cls) -> str:
+        return cls._PATTERN
+
+    @classmethod
+    def is_valid(cls, name: str) -> bool:
+        if not isinstance(name, str):
+            return False
+
+        if not (cls._MIN_LENGTH <= len(name) <= cls._MAX_LENGTH):
+            return False
+
+        if not re.match(cls._PATTERN, name):
+            return False
+
+        if name in cls._RESERVED:
+            return False
+
+        return True
+
+    @classmethod
+    def validate(cls, name: str) -> None:
+        if not isinstance(name, str):
+            raise MlflowException.invalid_parameter_value(
+                f"Workspace name must be a string, got {type(name).__name__!s}."
+            )
+
+        if not (cls._MIN_LENGTH <= len(name) <= cls._MAX_LENGTH):
+            raise MlflowException.invalid_parameter_value(
+                f"Workspace name '{name}' must be between {cls._MIN_LENGTH} and "
+                f"{cls._MAX_LENGTH} characters."
+            )
+
+        if not re.match(cls._PATTERN, name):
+            raise MlflowException.invalid_parameter_value(
+                f"Workspace name '{name}' must match the pattern {cls.pattern()} "
+                "(lowercase alphanumeric with optional internal hyphens)."
+            )
+
+        if name in cls._RESERVED:
+            raise MlflowException.invalid_parameter_value(
+                f"Workspace name '{name}' is reserved and cannot be used."
+            )
+