wehub-resource-sync
126 lines
4.0 KiB
Python
126 lines
4.0 KiB
Python
from __future__ import annotations
|
|
|
|
from contextvars import ContextVar, Token
|
|
|
|
from mlflow.environment_variables import MLFLOW_WORKSPACE
|
|
from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME
|
|
|
|
_WORKSPACE: ContextVar[str | None] = ContextVar("mlflow_active_workspace", default=None)
|
|
_IS_WORKSPACE_RESOLVED: ContextVar[bool] = ContextVar("mlflow_workspace_resolved", default=False)
|
|
|
|
|
|
def get_request_workspace() -> str | None:
|
|
"""
|
|
Return the active workspace for the current execution context.
|
|
|
|
Resolution order:
|
|
1) Request-scoped ContextVar (set by server middleware or explicit setters).
|
|
2) ``MLFLOW_WORKSPACE`` environment variable (client-side fallback, including threads).
|
|
"""
|
|
|
|
if workspace := (_WORKSPACE.get() or "").strip():
|
|
return workspace
|
|
|
|
if env_workspace := (MLFLOW_WORKSPACE.get() or "").strip():
|
|
return env_workspace
|
|
|
|
return None
|
|
|
|
|
|
def is_request_workspace_resolved() -> bool:
|
|
"""Return whether the server resolved the request workspace."""
|
|
return _IS_WORKSPACE_RESOLVED.get()
|
|
|
|
|
|
def _validate_workspace(workspace: str | None) -> None:
|
|
if workspace is not None and workspace != DEFAULT_WORKSPACE_NAME:
|
|
from mlflow.store.workspace.abstract_store import WorkspaceNameValidator
|
|
|
|
WorkspaceNameValidator.validate(workspace)
|
|
|
|
|
|
def set_server_request_workspace(workspace: str | None) -> Token[str | None]:
|
|
"""
|
|
Server-only setter: bind the workspace to the request ContextVar without touching env.
|
|
"""
|
|
|
|
_validate_workspace(workspace)
|
|
_IS_WORKSPACE_RESOLVED.set(True)
|
|
return _WORKSPACE.set(workspace)
|
|
|
|
|
|
def set_workspace(workspace: str | None) -> Token[str | None]:
|
|
"""
|
|
Client setter: binds the workspace to the current thread and persists to env so child
|
|
threads inherit it.
|
|
"""
|
|
|
|
_validate_workspace(workspace)
|
|
token = _WORKSPACE.set(workspace)
|
|
if workspace is None:
|
|
MLFLOW_WORKSPACE.unset()
|
|
else:
|
|
MLFLOW_WORKSPACE.set(workspace)
|
|
return token
|
|
|
|
|
|
def clear_server_request_workspace() -> None:
|
|
"""Clear the request-scoped ContextVar (does not touch the client env)."""
|
|
_IS_WORKSPACE_RESOLVED.set(False)
|
|
_WORKSPACE.set(None)
|
|
|
|
|
|
class WorkspaceContext:
|
|
"""
|
|
Context manager that sets the client workspace (ContextVar + env) for the duration
|
|
of the block. Restores the previous env value on exit.
|
|
"""
|
|
|
|
def __init__(self, workspace: str | None):
|
|
self._workspace = workspace
|
|
self._token: Token[str | None] | None = None
|
|
self._prev_env_raw: str | None = None
|
|
|
|
def __enter__(self) -> str | None:
|
|
self._prev_env_raw = MLFLOW_WORKSPACE.get_raw()
|
|
self._token = set_workspace(self._workspace)
|
|
return self._workspace
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> None:
|
|
if self._token is not None:
|
|
_WORKSPACE.reset(self._token)
|
|
self._token = None
|
|
|
|
if self._prev_env_raw is None:
|
|
MLFLOW_WORKSPACE.unset()
|
|
else:
|
|
MLFLOW_WORKSPACE.set(self._prev_env_raw)
|
|
self._prev_env_raw = None
|
|
|
|
|
|
class ServerWorkspaceContext:
|
|
"""
|
|
Context manager that sets only the server request workspace ContextVars for the duration
|
|
of the block. Unlike ``WorkspaceContext``, it does not touch the process environment.
|
|
"""
|
|
|
|
def __init__(self, workspace: str | None):
|
|
self._workspace = workspace
|
|
self._workspace_token: Token[str | None] | None = None
|
|
self._resolved_token: Token[bool] | None = None
|
|
|
|
def __enter__(self) -> str | None:
|
|
_validate_workspace(self._workspace)
|
|
self._resolved_token = _IS_WORKSPACE_RESOLVED.set(True)
|
|
self._workspace_token = _WORKSPACE.set(self._workspace)
|
|
return self._workspace
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> None:
|
|
if self._workspace_token is not None:
|
|
_WORKSPACE.reset(self._workspace_token)
|
|
self._workspace_token = None
|
|
|
|
if self._resolved_token is not None:
|
|
_IS_WORKSPACE_RESOLVED.reset(self._resolved_token)
|
|
self._resolved_token = None
|