wehub-resource-sync
639 lines
23 KiB
Python
639 lines
23 KiB
Python
from unittest import mock
|
|
|
|
import pytest
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.exc import IntegrityError
|
|
|
|
from mlflow.entities.workspace import Workspace, WorkspaceDeletionMode
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.store.artifact.artifact_repo import ArtifactRepository
|
|
from mlflow.store.workspace.dbmodels.models import SqlWorkspace
|
|
from mlflow.store.workspace.sqlalchemy_store import SqlAlchemyStore
|
|
from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME
|
|
|
|
|
|
@pytest.fixture
|
|
def workspace_store(db_uri, monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_ENABLE_WORKSPACES", "true")
|
|
|
|
store = SqlAlchemyStore(db_uri)
|
|
|
|
with store.ManagedSessionMaker(read_only=False) as session:
|
|
try:
|
|
session.add(
|
|
SqlWorkspace(
|
|
name=DEFAULT_WORKSPACE_NAME,
|
|
description="Default workspace",
|
|
)
|
|
)
|
|
session.commit()
|
|
except IntegrityError:
|
|
session.rollback()
|
|
|
|
try:
|
|
yield store
|
|
finally:
|
|
store._engine.dispose()
|
|
|
|
|
|
def _workspace_rows(store):
|
|
with store.ManagedSessionMaker() as session:
|
|
return {
|
|
(row.name, row.description)
|
|
for row in session.query(SqlWorkspace).order_by(SqlWorkspace.name).all()
|
|
}
|
|
|
|
|
|
def test_list_workspaces_returns_all(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description="Team A"))
|
|
workspace_store.create_workspace(Workspace(name="team-b", description=None))
|
|
|
|
workspaces = workspace_store.list_workspaces()
|
|
rows = {(ws.name, ws.description) for ws in workspaces}
|
|
default_description = next(desc for name, desc in rows if name == DEFAULT_WORKSPACE_NAME)
|
|
assert rows == {
|
|
(DEFAULT_WORKSPACE_NAME, default_description),
|
|
("team-a", "Team A"),
|
|
("team-b", None),
|
|
}
|
|
|
|
|
|
def test_get_workspace_success(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description="Team A"))
|
|
|
|
workspace = workspace_store.get_workspace("team-a")
|
|
assert workspace.name == "team-a"
|
|
assert workspace.description == "Team A"
|
|
|
|
|
|
def test_get_workspace_not_found(workspace_store):
|
|
with pytest.raises(MlflowException, match="Workspace 'unknown' not found") as exc:
|
|
workspace_store.get_workspace("unknown")
|
|
assert exc.value.error_code == "RESOURCE_DOES_NOT_EXIST"
|
|
|
|
|
|
def test_create_workspace_persists_record(workspace_store):
|
|
created = workspace_store.create_workspace(
|
|
Workspace(
|
|
name="team-a",
|
|
description="Team A",
|
|
default_artifact_root="s3://root/team-a",
|
|
trace_archival_location="s3://archive/team-a",
|
|
trace_archival_retention="30d",
|
|
),
|
|
)
|
|
assert created.name == "team-a"
|
|
assert created.description == "Team A"
|
|
assert created.default_artifact_root == "s3://root/team-a"
|
|
assert created.trace_archival_location == "s3://archive/team-a"
|
|
assert created.trace_archival_retention == "30d"
|
|
assert ("team-a", "Team A") in _workspace_rows(workspace_store)
|
|
|
|
|
|
def test_create_workspace_duplicate_raises(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description=None))
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="Workspace 'team-a' already exists\\.",
|
|
) as exc:
|
|
workspace_store.create_workspace(Workspace(name="team-a", description=None))
|
|
assert exc.value.error_code == "RESOURCE_ALREADY_EXISTS"
|
|
|
|
|
|
def test_create_workspace_invalid_name_raises(workspace_store):
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="Workspace name 'Team-A' must match the pattern",
|
|
) as exc:
|
|
workspace_store.create_workspace(Workspace(name="Team-A", description=None))
|
|
assert exc.value.error_code == "INVALID_PARAMETER_VALUE"
|
|
|
|
|
|
def test_create_workspace_invalid_trace_archival_location_raises(workspace_store):
|
|
with pytest.raises(MlflowException, match="proxy-only `mlflow-artifacts:` scheme") as exc:
|
|
workspace_store.create_workspace(
|
|
Workspace(
|
|
name="team-a",
|
|
description=None,
|
|
trace_archival_location="mlflow-artifacts:/archive/team-a",
|
|
)
|
|
)
|
|
assert exc.value.error_code == "INVALID_PARAMETER_VALUE"
|
|
|
|
|
|
def test_create_workspace_unsupported_trace_archival_location_raises(workspace_store):
|
|
class UnsupportedArchiveRepo(ArtifactRepository):
|
|
def log_artifact(self, local_file, artifact_path=None):
|
|
raise NotImplementedError
|
|
|
|
def log_artifacts(self, local_dir, artifact_path=None):
|
|
raise NotImplementedError
|
|
|
|
def list_artifacts(self, path=None):
|
|
raise NotImplementedError
|
|
|
|
with mock.patch(
|
|
"mlflow.store.artifact.artifact_repository_registry.get_artifact_repository",
|
|
return_value=UnsupportedArchiveRepo("dbfs:/archive/team-a"),
|
|
):
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="does not support deleting archived payloads",
|
|
) as exc:
|
|
workspace_store.create_workspace(
|
|
Workspace(
|
|
name="team-a",
|
|
description=None,
|
|
trace_archival_location="dbfs:/archive/team-a",
|
|
)
|
|
)
|
|
assert exc.value.error_code == "INVALID_PARAMETER_VALUE"
|
|
|
|
|
|
def test_create_workspace_invalid_trace_archival_retention_raises(workspace_store):
|
|
with pytest.raises(MlflowException, match="Trace archival retention must") as exc:
|
|
workspace_store.create_workspace(
|
|
Workspace(
|
|
name="team-a",
|
|
description=None,
|
|
trace_archival_retention="thirty-days",
|
|
)
|
|
)
|
|
assert exc.value.error_code == "INVALID_PARAMETER_VALUE"
|
|
|
|
|
|
def test_update_workspace_changes_description(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description="old"))
|
|
|
|
updated = workspace_store.update_workspace(
|
|
Workspace(name="team-a", description="new description"),
|
|
)
|
|
assert updated.description == "new description"
|
|
assert ("team-a", "new description") in _workspace_rows(workspace_store)
|
|
|
|
|
|
def test_update_workspace_sets_default_artifact_root(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description="old"))
|
|
|
|
updated = workspace_store.update_workspace(
|
|
Workspace(name="team-a", default_artifact_root="s3://bucket/team-a"),
|
|
)
|
|
assert updated.default_artifact_root == "s3://bucket/team-a"
|
|
fetched = workspace_store.get_workspace("team-a")
|
|
assert fetched.default_artifact_root == "s3://bucket/team-a"
|
|
|
|
|
|
def test_update_workspace_can_clear_default_artifact_root(workspace_store):
|
|
workspace_store.create_workspace(
|
|
Workspace(name="team-a", description="old", default_artifact_root="s3://bucket/team-a")
|
|
)
|
|
|
|
# Empty string signals "clear this field"
|
|
cleared = workspace_store.update_workspace(
|
|
Workspace(name="team-a", default_artifact_root=""),
|
|
)
|
|
assert cleared.default_artifact_root is None
|
|
fetched = workspace_store.get_workspace("team-a")
|
|
assert fetched.default_artifact_root is None
|
|
|
|
|
|
def test_update_workspace_sets_trace_archival_location(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description="old"))
|
|
|
|
updated = workspace_store.update_workspace(
|
|
Workspace(name="team-a", trace_archival_location="s3://archive/team-a")
|
|
)
|
|
assert updated.trace_archival_location == "s3://archive/team-a"
|
|
fetched = workspace_store.get_workspace("team-a")
|
|
assert fetched.trace_archival_location == "s3://archive/team-a"
|
|
|
|
|
|
def test_update_workspace_invalid_trace_archival_location_raises(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description="old"))
|
|
|
|
with pytest.raises(MlflowException, match="proxy-only `mlflow-artifacts:` scheme") as exc:
|
|
workspace_store.update_workspace(
|
|
Workspace(name="team-a", trace_archival_location="mlflow-artifacts:/archive/team-a")
|
|
)
|
|
assert exc.value.error_code == "INVALID_PARAMETER_VALUE"
|
|
|
|
|
|
def test_update_workspace_unsupported_trace_archival_location_raises(workspace_store):
|
|
class UnsupportedArchiveRepo(ArtifactRepository):
|
|
def log_artifact(self, local_file, artifact_path=None):
|
|
raise NotImplementedError
|
|
|
|
def log_artifacts(self, local_dir, artifact_path=None):
|
|
raise NotImplementedError
|
|
|
|
def list_artifacts(self, path=None):
|
|
raise NotImplementedError
|
|
|
|
workspace_store.create_workspace(Workspace(name="team-a", description="old"))
|
|
|
|
with mock.patch(
|
|
"mlflow.store.artifact.artifact_repository_registry.get_artifact_repository",
|
|
return_value=UnsupportedArchiveRepo("dbfs:/archive/team-a"),
|
|
):
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="does not support deleting archived payloads",
|
|
) as exc:
|
|
workspace_store.update_workspace(
|
|
Workspace(name="team-a", trace_archival_location="dbfs:/archive/team-a")
|
|
)
|
|
assert exc.value.error_code == "INVALID_PARAMETER_VALUE"
|
|
|
|
|
|
def test_update_workspace_can_clear_trace_archival_location(workspace_store):
|
|
workspace_store.create_workspace(
|
|
Workspace(
|
|
name="team-a",
|
|
description="old",
|
|
trace_archival_location="s3://archive/team-a",
|
|
)
|
|
)
|
|
|
|
cleared = workspace_store.update_workspace(Workspace(name="team-a", trace_archival_location=""))
|
|
assert cleared.trace_archival_location is None
|
|
fetched = workspace_store.get_workspace("team-a")
|
|
assert fetched.trace_archival_location is None
|
|
|
|
|
|
def test_update_workspace_sets_trace_archival_retention(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description="old"))
|
|
|
|
updated = workspace_store.update_workspace(
|
|
Workspace(name="team-a", trace_archival_retention="14d")
|
|
)
|
|
assert updated.trace_archival_retention == "14d"
|
|
fetched = workspace_store.get_workspace("team-a")
|
|
assert fetched.trace_archival_retention == "14d"
|
|
|
|
|
|
def test_update_workspace_invalid_trace_archival_retention_raises(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description="old"))
|
|
|
|
with pytest.raises(MlflowException, match="Trace archival retention must") as exc:
|
|
workspace_store.update_workspace(
|
|
Workspace(name="team-a", trace_archival_retention="thirty-days")
|
|
)
|
|
assert exc.value.error_code == "INVALID_PARAMETER_VALUE"
|
|
|
|
|
|
def test_update_workspace_can_clear_trace_archival_retention(workspace_store):
|
|
workspace_store.create_workspace(
|
|
Workspace(
|
|
name="team-a",
|
|
description="old",
|
|
trace_archival_retention="30d",
|
|
)
|
|
)
|
|
|
|
cleared = workspace_store.update_workspace(
|
|
Workspace(name="team-a", trace_archival_retention="")
|
|
)
|
|
assert cleared.trace_archival_retention is None
|
|
fetched = workspace_store.get_workspace("team-a")
|
|
assert fetched.trace_archival_retention is None
|
|
|
|
|
|
def test_delete_workspace_removes_empty_workspace(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description=None))
|
|
|
|
workspace_store.delete_workspace("team-a")
|
|
rows = _workspace_rows(workspace_store)
|
|
assert ("team-a", None) not in rows
|
|
default_ws = workspace_store.get_default_workspace()
|
|
assert (DEFAULT_WORKSPACE_NAME, default_ws.description) in rows
|
|
|
|
|
|
def test_delete_default_workspace_rejected(workspace_store):
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=f"Cannot delete the reserved '{DEFAULT_WORKSPACE_NAME}' workspace",
|
|
) as exc:
|
|
workspace_store.delete_workspace(DEFAULT_WORKSPACE_NAME)
|
|
assert exc.value.error_code == "INVALID_STATE"
|
|
|
|
|
|
def test_update_workspace_not_found(workspace_store):
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="Workspace 'unknown' not found",
|
|
) as exc:
|
|
workspace_store.update_workspace(Workspace(name="unknown", description="new description"))
|
|
assert exc.value.error_code == "RESOURCE_DOES_NOT_EXIST"
|
|
|
|
|
|
def test_delete_workspace_not_found(workspace_store):
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="Workspace 'unknown' not found",
|
|
) as exc:
|
|
workspace_store.delete_workspace("unknown")
|
|
assert exc.value.error_code == "RESOURCE_DOES_NOT_EXIST"
|
|
|
|
|
|
def test_resolve_artifact_root_returns_default(workspace_store):
|
|
default_root = "/default/path"
|
|
assert workspace_store.resolve_artifact_root(default_root, DEFAULT_WORKSPACE_NAME) == (
|
|
default_root,
|
|
True,
|
|
)
|
|
workspace_store.create_workspace(Workspace(name="team-a", description=None))
|
|
assert workspace_store.resolve_artifact_root(default_root, workspace_name="team-a") == (
|
|
default_root,
|
|
True,
|
|
)
|
|
|
|
|
|
def test_resolve_artifact_root_prefers_workspace_override(workspace_store):
|
|
workspace_store.create_workspace(
|
|
Workspace(
|
|
name="team-a",
|
|
description=None,
|
|
default_artifact_root="s3://team-a-artifacts",
|
|
)
|
|
)
|
|
|
|
resolved_root, should_append = workspace_store.resolve_artifact_root(
|
|
"/default/path", workspace_name="team-a"
|
|
)
|
|
assert resolved_root == "s3://team-a-artifacts"
|
|
assert not should_append
|
|
|
|
|
|
def test_resolve_artifact_root_cache_updates_on_override_change(workspace_store):
|
|
default_root = "/default/path"
|
|
workspace_store.create_workspace(Workspace(name="team-cache", description=None))
|
|
|
|
assert workspace_store.resolve_artifact_root(default_root, "team-cache") == (
|
|
default_root,
|
|
True,
|
|
)
|
|
|
|
workspace_store.update_workspace(
|
|
Workspace(name="team-cache", default_artifact_root="s3://cache/team")
|
|
)
|
|
|
|
assert workspace_store.resolve_artifact_root(default_root, "team-cache") == (
|
|
"s3://cache/team",
|
|
False,
|
|
)
|
|
|
|
|
|
def test_resolve_artifact_root_cache_handles_delete_and_recreate(workspace_store):
|
|
default_root = "/default/path"
|
|
workspace_store.create_workspace(
|
|
Workspace(name="team-cache", description=None, default_artifact_root="s3://cache/a")
|
|
)
|
|
|
|
assert workspace_store.resolve_artifact_root(default_root, "team-cache") == (
|
|
"s3://cache/a",
|
|
False,
|
|
)
|
|
|
|
workspace_store.delete_workspace("team-cache")
|
|
workspace_store.create_workspace(
|
|
Workspace(name="team-cache", description=None, default_artifact_root="s3://cache/b")
|
|
)
|
|
|
|
assert workspace_store.resolve_artifact_root(default_root, "team-cache") == (
|
|
"s3://cache/b",
|
|
False,
|
|
)
|
|
|
|
|
|
def test_resolve_artifact_root_cache_clears_when_override_removed(workspace_store):
|
|
default_root = "/default/path"
|
|
workspace_store.create_workspace(
|
|
Workspace(name="team-cache", description=None, default_artifact_root="s3://cache/a")
|
|
)
|
|
|
|
assert workspace_store.resolve_artifact_root(default_root, "team-cache") == (
|
|
"s3://cache/a",
|
|
False,
|
|
)
|
|
|
|
workspace_store.update_workspace(Workspace(name="team-cache", default_artifact_root=""))
|
|
|
|
assert workspace_store.resolve_artifact_root(default_root, "team-cache") == (
|
|
default_root,
|
|
True,
|
|
)
|
|
|
|
|
|
def test_resolve_trace_archival_config_returns_defaults(workspace_store):
|
|
config = workspace_store.resolve_trace_archival_config(
|
|
default_trace_archival_root="s3://archive/default",
|
|
default_retention="30d",
|
|
workspace_name=DEFAULT_WORKSPACE_NAME,
|
|
)
|
|
assert config.config.location == "s3://archive/default"
|
|
assert config.append_workspace_prefix
|
|
assert config.config.retention == "30d"
|
|
|
|
|
|
def test_resolve_trace_archival_config_prefers_workspace_overrides(workspace_store):
|
|
workspace_store.create_workspace(
|
|
Workspace(
|
|
name="team-a",
|
|
description=None,
|
|
trace_archival_location="s3://archive/team-a",
|
|
trace_archival_retention="14d",
|
|
)
|
|
)
|
|
|
|
config = workspace_store.resolve_trace_archival_config(
|
|
default_trace_archival_root="s3://archive/default",
|
|
default_retention="30d",
|
|
workspace_name="team-a",
|
|
)
|
|
assert config.config.location == "s3://archive/team-a"
|
|
assert not config.append_workspace_prefix
|
|
assert config.config.retention == "14d"
|
|
|
|
|
|
def test_resolve_trace_archival_config_cache_updates_on_override_change(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-cache", description=None))
|
|
|
|
initial = workspace_store.resolve_trace_archival_config(
|
|
default_trace_archival_root="s3://archive/default",
|
|
default_retention="30d",
|
|
workspace_name="team-cache",
|
|
)
|
|
assert initial.config.location == "s3://archive/default"
|
|
assert initial.append_workspace_prefix
|
|
assert initial.config.retention == "30d"
|
|
|
|
workspace_store.update_workspace(
|
|
Workspace(
|
|
name="team-cache",
|
|
trace_archival_location="s3://archive/team-cache",
|
|
trace_archival_retention="7d",
|
|
)
|
|
)
|
|
|
|
updated = workspace_store.resolve_trace_archival_config(
|
|
default_trace_archival_root="s3://archive/default",
|
|
default_retention="30d",
|
|
workspace_name="team-cache",
|
|
)
|
|
assert updated.config.location == "s3://archive/team-cache"
|
|
assert not updated.append_workspace_prefix
|
|
assert updated.config.retention == "7d"
|
|
|
|
|
|
def test_get_default_workspace_returns_default(workspace_store):
|
|
default_ws = workspace_store.get_default_workspace()
|
|
assert default_ws.name == DEFAULT_WORKSPACE_NAME
|
|
assert default_ws.description is not None
|
|
|
|
|
|
def test_delete_workspace_reassigns_resources_to_default(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description=None))
|
|
|
|
with workspace_store.ManagedSessionMaker(read_only=False) as session:
|
|
session.execute(
|
|
sa.text(
|
|
"INSERT INTO experiments (name, workspace, lifecycle_stage) "
|
|
"VALUES (:name, :ws, 'active')"
|
|
),
|
|
{"name": "exp-in-team-a", "ws": "team-a"},
|
|
)
|
|
|
|
workspace_store.delete_workspace("team-a", mode=WorkspaceDeletionMode.SET_DEFAULT)
|
|
|
|
with workspace_store.ManagedSessionMaker() as session:
|
|
row = session.execute(
|
|
sa.text("SELECT workspace FROM experiments WHERE name = :name"),
|
|
{"name": "exp-in-team-a"},
|
|
).fetchone()
|
|
assert row[0] == DEFAULT_WORKSPACE_NAME
|
|
|
|
|
|
def test_delete_workspace_fails_on_naming_conflict(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description=None))
|
|
|
|
with workspace_store.ManagedSessionMaker(read_only=False) as session:
|
|
session.execute(
|
|
sa.text(
|
|
"INSERT INTO experiments (name, workspace, lifecycle_stage) "
|
|
"VALUES (:name, :ws, 'active')"
|
|
),
|
|
{"name": "shared-exp", "ws": "team-a"},
|
|
)
|
|
session.execute(
|
|
sa.text(
|
|
"INSERT INTO experiments (name, workspace, lifecycle_stage) "
|
|
"VALUES (:name, :ws, 'active')"
|
|
),
|
|
{"name": "shared-exp", "ws": DEFAULT_WORKSPACE_NAME},
|
|
)
|
|
|
|
with pytest.raises(MlflowException, match="already exist in the default workspace") as exc:
|
|
workspace_store.delete_workspace("team-a", mode=WorkspaceDeletionMode.SET_DEFAULT)
|
|
assert exc.value.error_code == "INVALID_STATE"
|
|
|
|
# Workspace should still exist (transaction rolled back)
|
|
ws = workspace_store.get_workspace("team-a")
|
|
assert ws.name == "team-a"
|
|
|
|
|
|
def test_delete_workspace_cascade_removes_resources(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description=None))
|
|
|
|
with workspace_store.ManagedSessionMaker(read_only=False) as session:
|
|
session.execute(
|
|
sa.text(
|
|
"INSERT INTO experiments (name, workspace, lifecycle_stage) "
|
|
"VALUES (:name, :ws, 'active')"
|
|
),
|
|
{"name": "exp-in-team-a", "ws": "team-a"},
|
|
)
|
|
|
|
workspace_store.delete_workspace("team-a", mode=WorkspaceDeletionMode.CASCADE)
|
|
|
|
with workspace_store.ManagedSessionMaker() as session:
|
|
row = session.execute(
|
|
sa.text("SELECT count(*) FROM experiments WHERE name = :name"),
|
|
{"name": "exp-in-team-a"},
|
|
).scalar()
|
|
assert row == 0
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
workspace_store.get_workspace("team-a")
|
|
|
|
|
|
def test_delete_workspace_cascade_removes_experiment_with_runs(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description=None))
|
|
|
|
with workspace_store.ManagedSessionMaker(read_only=False) as session:
|
|
session.execute(
|
|
sa.text(
|
|
"INSERT INTO experiments (experiment_id, name, workspace, lifecycle_stage) "
|
|
"VALUES (:id, :name, :ws, 'active')"
|
|
),
|
|
{"id": 999, "name": "exp-with-runs", "ws": "team-a"},
|
|
)
|
|
session.execute(
|
|
sa.text(
|
|
"INSERT INTO runs (run_uuid, name, experiment_id, lifecycle_stage, status, "
|
|
"source_type, start_time, end_time) "
|
|
"VALUES (:run_id, :name, :exp_id, 'active', 'FINISHED', 'LOCAL', 0, 0)"
|
|
),
|
|
{"run_id": "run-in-team-a", "name": "test-run", "exp_id": 999},
|
|
)
|
|
|
|
workspace_store.delete_workspace("team-a", mode=WorkspaceDeletionMode.CASCADE)
|
|
|
|
with workspace_store.ManagedSessionMaker() as session:
|
|
exp_count = session.execute(
|
|
sa.text("SELECT count(*) FROM experiments WHERE name = :name"),
|
|
{"name": "exp-with-runs"},
|
|
).scalar()
|
|
assert exp_count == 0
|
|
run_count = session.execute(
|
|
sa.text("SELECT count(*) FROM runs WHERE run_uuid = :run_id"),
|
|
{"run_id": "run-in-team-a"},
|
|
).scalar()
|
|
assert run_count == 0
|
|
|
|
|
|
def test_delete_workspace_restrict_blocks_when_resources_exist(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description=None))
|
|
|
|
with workspace_store.ManagedSessionMaker(read_only=False) as session:
|
|
session.execute(
|
|
sa.text(
|
|
"INSERT INTO experiments (name, workspace, lifecycle_stage) "
|
|
"VALUES (:name, :ws, 'active')"
|
|
),
|
|
{"name": "exp-in-team-a", "ws": "team-a"},
|
|
)
|
|
|
|
with pytest.raises(MlflowException, match="still contains") as exc:
|
|
workspace_store.delete_workspace("team-a", mode=WorkspaceDeletionMode.RESTRICT)
|
|
assert exc.value.error_code == "INVALID_STATE"
|
|
|
|
# Workspace and resources should still exist
|
|
ws = workspace_store.get_workspace("team-a")
|
|
assert ws.name == "team-a"
|
|
with workspace_store.ManagedSessionMaker() as session:
|
|
row = session.execute(
|
|
sa.text("SELECT workspace FROM experiments WHERE name = :name"),
|
|
{"name": "exp-in-team-a"},
|
|
).fetchone()
|
|
assert row[0] == "team-a"
|
|
|
|
|
|
def test_delete_workspace_restrict_allows_empty_workspace(workspace_store):
|
|
workspace_store.create_workspace(Workspace(name="team-a", description=None))
|
|
|
|
workspace_store.delete_workspace("team-a", mode=WorkspaceDeletionMode.RESTRICT)
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
workspace_store.get_workspace("team-a")
|