wehub-resource-sync
119 lines
3.4 KiB
Python
119 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
from mlflow.environment_variables import MLFLOW_WORKSPACE, MLFLOW_WORKSPACE_STORE_URI
|
|
|
|
_workspace_store_uri: str | None = None
|
|
|
|
DEFAULT_WORKSPACE_NAME = "default"
|
|
WORKSPACES_DIR_NAME = "workspaces"
|
|
WORKSPACE_HEADER_NAME = "X-MLFLOW-WORKSPACE"
|
|
|
|
|
|
def _normalize_workspace(workspace: str | None) -> str | None:
|
|
"""Normalize a workspace identifier.
|
|
|
|
Args:
|
|
workspace: Raw workspace value, possibly ``None`` or whitespace padded.
|
|
|
|
Returns:
|
|
The trimmed workspace string or ``None`` when the input is empty.
|
|
"""
|
|
if workspace is None:
|
|
return None
|
|
value = workspace.strip()
|
|
return value or None
|
|
|
|
|
|
def resolve_entity_workspace_name(workspace: str | None) -> str:
|
|
"""Determine the workspace to associate with client-facing entities.
|
|
|
|
Preference order:
|
|
1. Explicit ``workspace`` argument provided by the backend store.
|
|
2. Active workspace bound via ``mlflow.set_workspace``.
|
|
3. ``MLFLOW_WORKSPACE`` environment variable.
|
|
4. ``DEFAULT_WORKSPACE_NAME``.
|
|
|
|
Args:
|
|
workspace: Optional workspace name provided by the store layer.
|
|
|
|
Returns:
|
|
A normalized workspace name honoring the preference order.
|
|
"""
|
|
|
|
if candidate := _normalize_workspace(workspace):
|
|
return candidate
|
|
|
|
from mlflow.utils.workspace_context import get_request_workspace
|
|
|
|
if candidate := _normalize_workspace(get_request_workspace()):
|
|
return candidate
|
|
|
|
if candidate := _normalize_workspace(MLFLOW_WORKSPACE.get()):
|
|
return candidate
|
|
|
|
return DEFAULT_WORKSPACE_NAME
|
|
|
|
|
|
def set_workspace_store_uri(uri: str | None) -> None:
|
|
"""Set the global workspace provider URI override.
|
|
|
|
Args:
|
|
uri: URI of the workspace provider or ``None`` to clear the override.
|
|
"""
|
|
|
|
global _workspace_store_uri
|
|
_workspace_store_uri = uri
|
|
if uri is None:
|
|
MLFLOW_WORKSPACE_STORE_URI.unset()
|
|
else:
|
|
MLFLOW_WORKSPACE_STORE_URI.set(uri)
|
|
|
|
|
|
def resolve_workspace_store_uri(
|
|
workspace_store_uri: str | None = None, tracking_uri: str | None = None
|
|
) -> str | None:
|
|
"""Resolve the workspace provider URI according to precedence rules.
|
|
|
|
Args:
|
|
workspace_store_uri: URI provided explicitly by the caller.
|
|
tracking_uri: Tracking URI used as the final fallback.
|
|
|
|
Returns:
|
|
The workspace provider URI chosen in this order:
|
|
|
|
1. Explicit ``workspace_store_uri`` argument.
|
|
2. Value configured via :func:`set_workspace_store_uri`
|
|
or ``MLFLOW_WORKSPACE_STORE_URI``.
|
|
3. The resolved tracking URI.
|
|
"""
|
|
|
|
if workspace_store_uri is not None:
|
|
return workspace_store_uri
|
|
|
|
if configured_uri := get_workspace_store_uri():
|
|
return configured_uri
|
|
|
|
# Lazy import to avoid circular dependency during module import.
|
|
from mlflow.tracking._tracking_service import utils as tracking_utils
|
|
|
|
return tracking_utils._resolve_tracking_uri(tracking_uri)
|
|
|
|
|
|
def get_workspace_store_uri() -> str | None:
|
|
"""Get the current workspace provider URI override, if any.
|
|
|
|
Returns:
|
|
The globally configured workspace provider URI, or ``None``.
|
|
"""
|
|
return _workspace_store_uri or MLFLOW_WORKSPACE_STORE_URI.get()
|
|
|
|
|
|
__all__ = [
|
|
"DEFAULT_WORKSPACE_NAME",
|
|
"WORKSPACES_DIR_NAME",
|
|
"WORKSPACE_HEADER_NAME",
|
|
"resolve_entity_workspace_name",
|
|
"set_workspace_store_uri",
|
|
"get_workspace_store_uri",
|
|
]
|