wehub/mlflow--mlflow
1
0
Fork 0
Code Issues Pull Requests Actions 16 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
mlflow--mlflow/tests/store/model_registry/test_sqlalchemy_store.py
T
wehub-resource-sync 4b22cfda96 chore: import upstream snapshot with attribution
2026-07-13 13:22:34 +08:00

2519 lines
94 KiB
Python
Raw Permalink Blame History

import concurrent.futures
import shutil
import time
import uuid
from pathlib import Path
from unittest import mock
import pytest
from sqlalchemy import create_engine, text
from mlflow.entities.model_registry import (
ModelVersion,
ModelVersionTag,
RegisteredModelTag,
)
from mlflow.entities.model_registry.prompt_version import IS_PROMPT_TAG_KEY
from mlflow.entities.webhook import WebhookAction, WebhookEntity, WebhookEvent, WebhookStatus
from mlflow.environment_variables import (
_MLFLOW_GO_STORE_TESTING,
MLFLOW_ENABLE_WORKSPACES,
MLFLOW_TRACKING_URI,
)
from mlflow.exceptions import MlflowException
from mlflow.prompt.constants import PROMPT_TEXT_TAG_KEY
from mlflow.protos.databricks_pb2 import (
INVALID_PARAMETER_VALUE,
RESOURCE_ALREADY_EXISTS,
RESOURCE_DOES_NOT_EXIST,
ErrorCode,
)
from mlflow.store.model_registry.dbmodels.models import (
SqlModelVersion,
SqlModelVersionTag,
SqlRegisteredModel,
SqlRegisteredModelTag,
SqlWebhook,
)
from mlflow.store.model_registry.sqlalchemy_store import SqlAlchemyStore
from mlflow.store.model_registry.sqlalchemy_workspace_store import (
WorkspaceAwareSqlAlchemyStore,
)
from mlflow.utils.workspace_context import WorkspaceContext
from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME
from tests.helper_functions import random_str
pytestmark = pytest.mark.notrackingurimock
GO_MOCK_TIME_TAG = "mock.time.go.testing.tag"
@pytest.fixture(autouse=True, params=[False, True], ids=["workspace-disabled", "workspace-enabled"])
def workspaces_enabled(request, monkeypatch, disable_workspace_mode_by_default):
"""
Run every test in this module with workspaces disabled and enabled to cover both code paths.
"""
enabled = request.param
monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true" if enabled else "false")
if enabled:
with WorkspaceContext(DEFAULT_WORKSPACE_NAME):
yield enabled
else:
yield enabled
@pytest.fixture
def store(tmp_path: Path, cached_db: Path, workspaces_enabled):
store_cls = WorkspaceAwareSqlAlchemyStore if workspaces_enabled else SqlAlchemyStore
if db_uri_env := MLFLOW_TRACKING_URI.get():
s = store_cls(db_uri_env)
yield s
_cleanup_database(s)
else:
db_path = tmp_path / "mlflow.db"
shutil.copy(cached_db, db_path)
db_uri = f"sqlite:///{db_path}"
s = store_cls(db_uri)
yield s
# Dispose the engine to close all pooled connections
s.engine.dispose()
def _cleanup_database(store: SqlAlchemyStore):
with store.ManagedSessionMaker() as session:
# Delete all rows in all tables
for model in (
SqlModelVersionTag,
SqlRegisteredModelTag,
SqlModelVersion,
SqlRegisteredModel,
SqlWebhook,
):
session.query(model).delete()
def _rm_maker(store, name, tags=None, description=None):
return store.create_registered_model(name, tags, description)
def _add_go_test_tags(tags, val):
if _MLFLOW_GO_STORE_TESTING.get():
return tags + [RegisteredModelTag(GO_MOCK_TIME_TAG, val)]
return tags
def _mv_maker(
store,
name,
source="path/to/source",
run_id=uuid.uuid4().hex,
tags=None,
run_link=None,
description=None,
):
return store.create_model_version(
name, source, run_id, tags, run_link=run_link, description=description
)
def _extract_latest_by_stage(latest_versions):
return {mvd.current_stage: str(mvd.version) for mvd in latest_versions}
def test_create_registered_model(store):
name = random_str() + "abCD"
rm1 = _rm_maker(store, name)
assert rm1.name == name
assert rm1.description is None
# error on duplicate
with pytest.raises(
MlflowException, match=rf"Registered Model \(name={name}\) already exists"
) as exception_context:
_rm_maker(store, name)
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS)
# slightly different name is ok
for name2 in [name + "extra", name + name]:
rm2 = _rm_maker(store, name2)
assert rm2.name == name2
# test create model with tags
name2 = random_str() + "tags"
tags = [
RegisteredModelTag("key", "value"),
RegisteredModelTag("anotherKey", "some other value"),
]
rm2 = _rm_maker(store, name2, tags)
rmd2 = store.get_registered_model(name2)
assert rm2.name == name2
assert rm2.tags == {tag.key: tag.value for tag in tags}
assert rmd2.name == name2
assert rmd2.tags == {tag.key: tag.value for tag in tags}
# create with description
name3 = random_str() + "-description"
description = "the best model ever"
rm3 = _rm_maker(store, name3, description=description)
rmd3 = store.get_registered_model(name3)
assert rm3.name == name3
assert rm3.description == description
assert rmd3.name == name3
assert rmd3.description == description
# invalid model name will fail
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'name'"
) as exception_context:
_rm_maker(store, None)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'name'"
) as exception_context:
_rm_maker(store, "")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
def test_get_registered_model(store):
name = "model_1"
tags = [
RegisteredModelTag("key", "value"),
RegisteredModelTag("anotherKey", "some other value"),
]
# use fake clock
with mock.patch("time.time", return_value=1234):
rm = _rm_maker(store, name, _add_go_test_tags(tags, "1234000"))
assert rm.name == name
rmd = store.get_registered_model(name=name)
assert rmd.name == name
assert rmd.creation_timestamp == 1234000
assert rmd.last_updated_timestamp == 1234000
assert rmd.description is None
assert rmd.latest_versions == []
assert rmd.tags == {tag.key: tag.value for tag in tags}
def test_update_registered_model(store):
name = "model_for_update_RM"
rm1 = _rm_maker(store, name)
rmd1 = store.get_registered_model(name=name)
assert rm1.name == name
assert rmd1.description is None
# update description
rm2 = store.update_registered_model(name=name, description="test model")
rmd2 = store.get_registered_model(name=name)
assert rm2.name == "model_for_update_RM"
assert rmd2.name == "model_for_update_RM"
assert rmd2.description == "test model"
def test_rename_registered_model(store):
original_name = "original name"
new_name = "new name"
_rm_maker(store, original_name)
_mv_maker(store, original_name)
_mv_maker(store, original_name)
rm = store.get_registered_model(original_name)
mv1 = store.get_model_version(original_name, 1)
mv2 = store.get_model_version(original_name, 2)
assert rm.name == original_name
assert mv1.name == original_name
assert mv2.name == original_name
# test renaming registered model also updates its model versions
store.rename_registered_model(original_name, new_name)
rm = store.get_registered_model(new_name)
mv1 = store.get_model_version(new_name, 1)
mv2 = store.get_model_version(new_name, 2)
assert rm.name == new_name
assert mv1.name == new_name
assert mv2.name == new_name
# test accessing the model with the old name will fail
with pytest.raises(
MlflowException, match=rf"Registered Model with name={original_name} not found"
) as exception_context:
store.get_registered_model(original_name)
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
# test name another model with the replaced name is ok
_rm_maker(store, original_name)
# cannot rename model to conflict with an existing model
with pytest.raises(
MlflowException,
match=rf"Registered Model \(name={original_name}\) already exists",
) as exception_context:
store.rename_registered_model(new_name, original_name)
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS)
# invalid model name will fail
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'new_name'"
) as exception_context:
store.rename_registered_model(original_name, None)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'new_name'"
) as exception_context:
store.rename_registered_model(original_name, "")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
def test_delete_registered_model(store):
name = "model_for_delete_RM"
_rm_maker(store, name)
_mv_maker(store, name)
rm1 = store.get_registered_model(name=name)
mv1 = store.get_model_version(name, 1)
assert rm1.name == name
assert mv1.name == name
# delete model
store.delete_registered_model(name=name)
# cannot get model
with pytest.raises(
MlflowException, match=rf"Registered Model with name={name} not found"
) as exception_context:
store.get_registered_model(name=name)
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
# cannot update a delete model
with pytest.raises(
MlflowException, match=rf"Registered Model with name={name} not found"
) as exception_context:
store.update_registered_model(name=name, description="deleted")
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
# cannot delete it again
with pytest.raises(
MlflowException, match=rf"Registered Model with name={name} not found"
) as exception_context:
store.delete_registered_model(name=name)
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
# model versions are cascade deleted with the registered model
with pytest.raises(
MlflowException, match=rf"Model Version \(name={name}, version=1\) not found"
) as exception_context:
store.get_model_version(name, 1)
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
def test_get_latest_versions(store):
name = "test_for_latest_versions"
_rm_maker(store, name)
rmd1 = store.get_registered_model(name=name)
assert rmd1.latest_versions == []
mv1 = _mv_maker(store, name)
assert mv1.version == 1
rmd2 = store.get_registered_model(name=name)
assert _extract_latest_by_stage(rmd2.latest_versions) == {"None": "1"}
# add a bunch more
mv2 = _mv_maker(store, name)
assert mv2.version == 2
store.transition_model_version_stage(
name=mv2.name,
version=mv2.version,
stage="Production",
archive_existing_versions=False,
)
mv3 = _mv_maker(store, name)
assert mv3.version == 3
store.transition_model_version_stage(
name=mv3.name,
version=mv3.version,
stage="Production",
archive_existing_versions=False,
)
mv4 = _mv_maker(store, name)
assert mv4.version == 4
store.transition_model_version_stage(
name=mv4.name,
version=mv4.version,
stage="Staging",
archive_existing_versions=False,
)
# test that correct latest versions are returned for each stage
rmd4 = store.get_registered_model(name=name)
assert _extract_latest_by_stage(rmd4.latest_versions) == {
"None": "1",
"Production": "3",
"Staging": "4",
}
assert _extract_latest_by_stage(store.get_latest_versions(name=name, stages=None)) == {
"None": "1",
"Production": "3",
"Staging": "4",
}
assert _extract_latest_by_stage(store.get_latest_versions(name=name, stages=[])) == {
"None": "1",
"Production": "3",
"Staging": "4",
}
assert _extract_latest_by_stage(
store.get_latest_versions(name=name, stages=["Production"])
) == {"Production": "3"}
assert _extract_latest_by_stage(
store.get_latest_versions(name=name, stages=["production"])
) == {"Production": "3"} # The stages are case insensitive.
assert _extract_latest_by_stage(
store.get_latest_versions(name=name, stages=["pROduction"])
) == {"Production": "3"} # The stages are case insensitive.
assert _extract_latest_by_stage(
store.get_latest_versions(name=name, stages=["None", "Production"])
) == {"None": "1", "Production": "3"}
# delete latest Production, and should point to previous one
store.delete_model_version(name=mv3.name, version=mv3.version)
rmd5 = store.get_registered_model(name=name)
assert _extract_latest_by_stage(rmd5.latest_versions) == {
"None": "1",
"Production": "2",
"Staging": "4",
}
assert _extract_latest_by_stage(store.get_latest_versions(name=name, stages=None)) == {
"None": "1",
"Production": "2",
"Staging": "4",
}
assert _extract_latest_by_stage(
store.get_latest_versions(name=name, stages=["Production"])
) == {"Production": "2"}
def test_set_registered_model_tag(store):
name1 = "SetRegisteredModelTag_TestMod"
name2 = "SetRegisteredModelTag_TestMod 2"
initial_tags = [
RegisteredModelTag("key", "value"),
RegisteredModelTag("anotherKey", "some other value"),
]
_rm_maker(store, name1, initial_tags)
_rm_maker(store, name2, initial_tags)
new_tag = RegisteredModelTag("randomTag", "not a random value")
store.set_registered_model_tag(name1, new_tag)
rm1 = store.get_registered_model(name=name1)
all_tags = initial_tags + [new_tag]
assert rm1.tags == {tag.key: tag.value for tag in all_tags}
# test overriding a tag with the same key
overriding_tag = RegisteredModelTag("key", "overriding")
store.set_registered_model_tag(name1, overriding_tag)
all_tags = [tag for tag in all_tags if tag.key != "key"] + [overriding_tag]
rm1 = store.get_registered_model(name=name1)
assert rm1.tags == {tag.key: tag.value for tag in all_tags}
# does not affect other models with the same key
rm2 = store.get_registered_model(name=name2)
assert rm2.tags == {tag.key: tag.value for tag in initial_tags}
# can not set tag on deleted (non-existed) registered model
store.delete_registered_model(name1)
with pytest.raises(
MlflowException, match=rf"Registered Model with name={name1} not found"
) as exception_context:
store.set_registered_model_tag(name1, overriding_tag)
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
# test cannot set tags that are too long
long_tag = RegisteredModelTag("longTagKey", "a" * 100_001)
with pytest.raises(
MlflowException,
match=r"'value' exceeds the maximum length of \d+ characters",
) as exception_context:
store.set_registered_model_tag(name2, long_tag)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# test can set tags that are somewhat long
long_tag = RegisteredModelTag("longTagKey", "a" * 4999)
store.set_registered_model_tag(name2, long_tag)
# can not set invalid tag
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'key'"
) as exception_context:
store.set_registered_model_tag(name2, RegisteredModelTag(key=None, value=""))
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# can not use invalid model name
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'name'"
) as exception_context:
store.set_registered_model_tag(None, RegisteredModelTag(key="key", value="value"))
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
def test_delete_registered_model_tag(store):
name1 = "DeleteRegisteredModelTag_TestMod"
name2 = "DeleteRegisteredModelTag_TestMod 2"
initial_tags = [
RegisteredModelTag("key", "value"),
RegisteredModelTag("anotherKey", "some other value"),
]
_rm_maker(store, name1, initial_tags)
_rm_maker(store, name2, initial_tags)
new_tag = RegisteredModelTag("randomTag", "not a random value")
store.set_registered_model_tag(name1, new_tag)
store.delete_registered_model_tag(name1, "randomTag")
rm1 = store.get_registered_model(name=name1)
assert rm1.tags == {tag.key: tag.value for tag in initial_tags}
# testing deleting a key does not affect other models with the same key
store.delete_registered_model_tag(name1, "key")
rm1 = store.get_registered_model(name=name1)
rm2 = store.get_registered_model(name=name2)
assert rm1.tags == {"anotherKey": "some other value"}
assert rm2.tags == {tag.key: tag.value for tag in initial_tags}
# delete tag that is already deleted does nothing
store.delete_registered_model_tag(name1, "key")
rm1 = store.get_registered_model(name=name1)
assert rm1.tags == {"anotherKey": "some other value"}
# can not delete tag on deleted (non-existed) registered model
store.delete_registered_model(name1)
with pytest.raises(
MlflowException, match=rf"Registered Model with name={name1} not found"
) as exception_context:
store.delete_registered_model_tag(name1, "anotherKey")
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
# can not delete tag with invalid key
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'key'"
) as exception_context:
store.delete_registered_model_tag(name2, None)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# can not use invalid model name
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'name'"
) as exception_context:
store.delete_registered_model_tag(None, "key")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
def test_create_model_version(store):
name = "test_for_update_MV"
_rm_maker(store, name)
run_id = uuid.uuid4().hex
with mock.patch("time.time", return_value=456778):
mv1 = _mv_maker(store, name, "a/b/CD", run_id)
assert mv1.name == name
assert mv1.version == 1
mvd1 = store.get_model_version(mv1.name, mv1.version)
assert mvd1.name == name
assert int(mvd1.version) == 1
assert mvd1.current_stage == "None"
assert mvd1.creation_timestamp == 456778000
assert mvd1.last_updated_timestamp == 456778000
assert mvd1.description is None
assert mvd1.source == "a/b/CD"
assert mvd1.run_id == run_id
assert mvd1.status == "READY"
assert mvd1.status_message is None
assert mvd1.tags == {}
# new model versions for same name autoincrement versions
mv2 = _mv_maker(store, name)
mvd2 = store.get_model_version(name=mv2.name, version=mv2.version)
assert mv2.version == 2
assert int(mvd2.version) == 2
# create model version with tags return model version entity with tags
tags = [
ModelVersionTag("key", "value"),
ModelVersionTag("anotherKey", "some other value"),
]
mv3 = _mv_maker(store, name, tags=tags)
mvd3 = store.get_model_version(name=mv3.name, version=mv3.version)
assert mv3.version == 3
assert mv3.tags == {tag.key: tag.value for tag in tags}
assert int(mvd3.version) == 3
assert mvd3.tags == {tag.key: tag.value for tag in tags}
# create model versions with runLink
run_link = "http://localhost:3000/path/to/run/"
mv4 = _mv_maker(store, name, run_link=run_link)
mvd4 = store.get_model_version(name, mv4.version)
assert mv4.version == 4
assert mv4.run_link == run_link
assert int(mvd4.version) == 4
assert mvd4.run_link == run_link
# create model version with description
description = "the best model ever"
mv5 = _mv_maker(store, name, description=description)
mvd5 = store.get_model_version(name, mv5.version)
assert mv5.version == 5
assert mv5.description == description
assert int(mvd5.version) == 5
assert mvd5.description == description
# create model version without runId
mv6 = _mv_maker(store, name, run_id=None)
mvd6 = store.get_model_version(name, mv6.version)
assert mv6.version == 6
assert mv6.run_id is None
assert int(mvd6.version) == 6
assert mvd6.run_id is None
def test_update_model_version(store):
name = "test_for_update_MV"
_rm_maker(store, name)
mv1 = _mv_maker(store, name)
mvd1 = store.get_model_version(name=mv1.name, version=mv1.version)
assert mvd1.name == name
assert int(mvd1.version) == 1
assert mvd1.current_stage == "None"
# update stage
store.transition_model_version_stage(
name=mv1.name,
version=mv1.version,
stage="Production",
archive_existing_versions=False,
)
mvd2 = store.get_model_version(name=mv1.name, version=mv1.version)
assert mvd2.name == name
assert int(mvd2.version) == 1
assert mvd2.current_stage == "Production"
assert mvd2.description is None
# update description
store.update_model_version(name=mv1.name, version=mv1.version, description="test model version")
mvd3 = store.get_model_version(name=mv1.name, version=mv1.version)
assert mvd3.name == name
assert int(mvd3.version) == 1
assert mvd3.current_stage == "Production"
assert mvd3.description == "test model version"
# only valid stages can be set
with pytest.raises(
MlflowException,
match=(
"Invalid Model Version stage: unknown. "
"Value must be one of None, Staging, Production, Archived."
),
) as exception_context:
store.transition_model_version_stage(
mv1.name, mv1.version, stage="unknown", archive_existing_versions=False
)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# stages are case-insensitive and auto-corrected to system stage names
for stage_name in ["STAGING", "staging", "StAgInG"]:
store.transition_model_version_stage(
name=mv1.name,
version=mv1.version,
stage=stage_name,
archive_existing_versions=False,
)
mvd5 = store.get_model_version(name=mv1.name, version=mv1.version)
assert mvd5.current_stage == "Staging"
def test_transition_model_version_stage_when_archive_existing_versions_is_false(store):
name = "model"
_rm_maker(store, name)
mv1 = _mv_maker(store, name)
mv2 = _mv_maker(store, name)
mv3 = _mv_maker(store, name)
# test that when `archive_existing_versions` is False, transitioning a model version
# to the inactive stages ("Archived" and "None") does not throw.
for stage in ["Archived", "None"]:
store.transition_model_version_stage(name, mv1.version, stage, False)
store.transition_model_version_stage(name, mv1.version, "Staging", False)
store.transition_model_version_stage(name, mv2.version, "Production", False)
store.transition_model_version_stage(name, mv3.version, "Staging", False)
mvd1 = store.get_model_version(name=name, version=mv1.version)
mvd2 = store.get_model_version(name=name, version=mv2.version)
mvd3 = store.get_model_version(name=name, version=mv3.version)
assert mvd1.current_stage == "Staging"
assert mvd2.current_stage == "Production"
assert mvd3.current_stage == "Staging"
store.transition_model_version_stage(name, mv3.version, "Production", False)
mvd1 = store.get_model_version(name=name, version=mv1.version)
mvd2 = store.get_model_version(name=name, version=mv2.version)
mvd3 = store.get_model_version(name=name, version=mv3.version)
assert mvd1.current_stage == "Staging"
assert mvd2.current_stage == "Production"
assert mvd3.current_stage == "Production"
def test_transition_model_version_stage_when_archive_existing_versions_is_true(store):
name = "model"
_rm_maker(store, name)
mv1 = _mv_maker(store, name)
mv2 = _mv_maker(store, name)
mv3 = _mv_maker(store, name)
msg = (
r"Model version transition cannot archive existing model versions "
r"because .+ is not an Active stage"
)
# test that when `archive_existing_versions` is True, transitioning a model version
# to the inactive stages ("Archived" and "None") throws.
for stage in ["Archived", "None"]:
with pytest.raises(MlflowException, match=msg):
store.transition_model_version_stage(name, mv1.version, stage, True)
store.transition_model_version_stage(name, mv1.version, "Staging", False)
store.transition_model_version_stage(name, mv2.version, "Production", False)
store.transition_model_version_stage(name, mv3.version, "Staging", True)
mvd1 = store.get_model_version(name=name, version=mv1.version)
mvd2 = store.get_model_version(name=name, version=mv2.version)
mvd3 = store.get_model_version(name=name, version=mv3.version)
assert mvd1.current_stage == "Archived"
assert mvd2.current_stage == "Production"
assert mvd3.current_stage == "Staging"
assert mvd1.last_updated_timestamp == mvd3.last_updated_timestamp
store.transition_model_version_stage(name, mv3.version, "Production", True)
mvd1 = store.get_model_version(name=name, version=mv1.version)
mvd2 = store.get_model_version(name=name, version=mv2.version)
mvd3 = store.get_model_version(name=name, version=mv3.version)
assert mvd1.current_stage == "Archived"
assert mvd2.current_stage == "Archived"
assert mvd3.current_stage == "Production"
assert mvd2.last_updated_timestamp == mvd3.last_updated_timestamp
for uncanonical_stage_name in ["STAGING", "staging", "StAgInG"]:
store.transition_model_version_stage(mv1.name, mv1.version, "Staging", False)
store.transition_model_version_stage(mv2.name, mv2.version, "None", False)
# stage names are case-insensitive and auto-corrected to system stage names
store.transition_model_version_stage(mv2.name, mv2.version, uncanonical_stage_name, True)
mvd1 = store.get_model_version(name=mv1.name, version=mv1.version)
mvd2 = store.get_model_version(name=mv2.name, version=mv2.version)
assert mvd1.current_stage == "Archived"
assert mvd2.current_stage == "Staging"
def test_delete_model_version(store):
name = "test_for_delete_MV"
initial_tags = [
ModelVersionTag("key", "value"),
ModelVersionTag("anotherKey", "some other value"),
]
_rm_maker(store, name)
mv = _mv_maker(store, name, tags=initial_tags)
mvd = store.get_model_version(name=mv.name, version=mv.version)
assert mvd.name == name
store.delete_model_version(name=mv.name, version=mv.version)
# cannot get a deleted model version
with pytest.raises(
MlflowException,
match=rf"Model Version \(name={mv.name}, version={mv.version}\) not found",
) as exception_context:
store.get_model_version(name=mv.name, version=mv.version)
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
# cannot update a delete
with pytest.raises(
MlflowException,
match=rf"Model Version \(name={mv.name}, version={mv.version}\) not found",
) as exception_context:
store.update_model_version(mv.name, mv.version, description="deleted!")
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
# cannot delete it again
with pytest.raises(
MlflowException,
match=rf"Model Version \(name={mv.name}, version={mv.version}\) not found",
) as exception_context:
store.delete_model_version(name=mv.name, version=mv.version)
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
def test_delete_model_version_redaction(store):
name = "test_for_delete_MV_redaction"
run_link = "http://localhost:5000/path/to/run"
run_id = "12345"
source = "path/to/source"
_rm_maker(store, name)
mv = _mv_maker(store, name, source=source, run_id=run_id, run_link=run_link)
mvd = store.get_model_version(name=name, version=mv.version)
assert mvd.run_link == run_link
assert mvd.run_id == run_id
assert mvd.source == source
# delete the MV now
store.delete_model_version(name, mv.version)
# verify that the relevant fields are redacted
mvd_deleted = store._get_sql_model_version_including_deleted(name=name, version=mv.version)
assert "REDACTED" in mvd_deleted.run_link
assert "REDACTED" in mvd_deleted.source
assert "REDACTED" in mvd_deleted.run_id
def test_get_model_version_download_uri(store):
name = "test_for_update_MV"
_rm_maker(store, name)
source_path = "path/to/source"
mv = _mv_maker(store, name, source=source_path, run_id=uuid.uuid4().hex)
mvd1 = store.get_model_version(name=mv.name, version=mv.version)
assert mvd1.name == name
assert mvd1.source == source_path
# download location points to source
assert store.get_model_version_download_uri(name=mv.name, version=mv.version) == source_path
# download URI does not change even if model version is updated
store.transition_model_version_stage(
name=mv.name,
version=mv.version,
stage="Production",
archive_existing_versions=False,
)
store.update_model_version(name=mv.name, version=mv.version, description="Test for Path")
mvd2 = store.get_model_version(name=mv.name, version=mv.version)
assert mvd2.source == source_path
assert store.get_model_version_download_uri(name=mv.name, version=mv.version) == source_path
# cannot retrieve download URI for deleted model versions
store.delete_model_version(name=mv.name, version=mv.version)
with pytest.raises(
MlflowException,
match=rf"Model Version \(name={mv.name}, version={mv.version}\) not found",
) as exception_context:
store.get_model_version_download_uri(name=mv.name, version=mv.version)
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
def test_search_model_versions(store):
# create some model versions
name = "test_for_search_MV"
_rm_maker(store, name)
run_id_1 = uuid.uuid4().hex
run_id_2 = uuid.uuid4().hex
run_id_3 = uuid.uuid4().hex
mv1 = _mv_maker(store, name=name, source="A/B", run_id=run_id_1)
assert mv1.version == 1
mv2 = _mv_maker(store, name=name, source="A/C", run_id=run_id_2)
assert mv2.version == 2
mv3 = _mv_maker(store, name=name, source="A/D", run_id=run_id_2)
assert mv3.version == 3
mv4 = _mv_maker(store, name=name, source="A/D", run_id=run_id_3)
assert mv4.version == 4
def search_versions(filter_string, max_results=10, order_by=None, page_token=None):
return [
mvd.version
for mvd in store.search_model_versions(filter_string, max_results, order_by, page_token)
]
# search using name should return all 4 versions
assert set(search_versions(f"name='{name}'")) == {1, 2, 3, 4}
# search using version
assert set(search_versions("version_number=2")) == {2}
assert set(search_versions("version_number<=3")) == {1, 2, 3}
# search using run_id_1 should return version 1
assert set(search_versions(f"run_id='{run_id_1}'")) == {1}
# search using run_id_2 should return versions 2 and 3
assert set(search_versions(f"run_id='{run_id_2}'")) == {2, 3}
# search using the IN operator should return all versions
assert set(search_versions(f"run_id IN ('{run_id_1}','{run_id_2}')")) == {1, 2, 3}
# search IN operator is case sensitive
assert set(search_versions(f"run_id IN ('{run_id_1.upper()}','{run_id_2}')")) == {
2,
3,
}
# search IN operator with other conditions
assert set(
search_versions(f"version_number=2 AND run_id IN ('{run_id_1.upper()}','{run_id_2}')")
) == {2}
# search IN operator with right-hand side value containing whitespaces
assert set(search_versions(f"run_id IN ('{run_id_1}', '{run_id_2}')")) == {1, 2, 3}
# search using the IN operator with bad lists should return exceptions
with pytest.raises(
MlflowException,
match=(
r"While parsing a list in the query, "
r"expected string value, punctuation, or whitespace, "
r"but got different type in list"
),
) as exception_context:
search_versions("run_id IN (1,2,3)")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
assert set(search_versions(f"run_id LIKE '{run_id_2[:30]}%'")) == {2, 3}
assert set(search_versions(f"run_id ILIKE '{run_id_2[:30].upper()}%'")) == {2, 3}
# search using the IN operator with empty lists should return exceptions
with pytest.raises(
MlflowException,
match=(
r"While parsing a list in the query, "
r"expected a non-empty list of string values, "
r"but got empty list"
),
) as exception_context:
search_versions("run_id IN ()")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# search using an ill-formed IN operator correctly throws exception
with pytest.raises(
MlflowException, match=r"Invalid clause\(s\) in filter string"
) as exception_context:
search_versions("run_id IN (")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
with pytest.raises(
MlflowException, match=r"Invalid clause\(s\) in filter string"
) as exception_context:
search_versions("run_id IN")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
with pytest.raises(
MlflowException, match=r"Invalid clause\(s\) in filter string"
) as exception_context:
search_versions("name LIKE")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
with pytest.raises(
MlflowException,
match=(
r"While parsing a list in the query, "
r"expected a non-empty list of string values, "
r"but got ill-formed list"
),
) as exception_context:
search_versions("run_id IN (,)")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
with pytest.raises(
MlflowException,
match=(
r"While parsing a list in the query, "
r"expected a non-empty list of string values, "
r"but got ill-formed list"
),
) as exception_context:
search_versions("run_id IN ('runid1',,'runid2')")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# search using source_path "A/D" should return version 3 and 4
assert set(search_versions("source_path = 'A/D'")) == {3, 4}
# search using source_path "A" should not return anything
assert len(search_versions("source_path = 'A'")) == 0
assert len(search_versions("source_path = 'A/'")) == 0
assert len(search_versions("source_path = ''")) == 0
# delete mv4. search should not return version 4
store.delete_model_version(name=mv4.name, version=mv4.version)
assert set(search_versions("")) == {1, 2, 3}
assert set(search_versions(None)) == {1, 2, 3}
assert set(search_versions(f"name='{name}'")) == {1, 2, 3}
assert set(search_versions("source_path = 'A/D'")) == {3}
store.transition_model_version_stage(
name=mv1.name,
version=mv1.version,
stage="production",
archive_existing_versions=False,
)
store.update_model_version(
name=mv1.name, version=mv1.version, description="Online prediction model!"
)
mvds = store.search_model_versions(f"run_id = '{run_id_1}'", max_results=10)
assert len(mvds) == 1
assert isinstance(mvds[0], ModelVersion)
assert mvds[0].current_stage == "Production"
assert mvds[0].run_id == run_id_1
assert mvds[0].source == "A/B"
assert mvds[0].description == "Online prediction model!"
def test_search_model_versions_order_by_simple(store):
# create some model versions
names = ["RM1", "RM2", "RM3", "RM4", "RM1", "RM4"]
sources = ["A"] * 3 + ["B"] * 3
run_ids = [uuid.uuid4().hex for _ in range(6)]
for name in set(names):
_rm_maker(store, name)
for i in range(6):
time.sleep(0.001) # sleep to ensure each model version has a different creation_time
_mv_maker(store, name=names[i], source=sources[i], run_id=run_ids[i])
# by default order by last_updated_timestamp DESC
mvs = store.search_model_versions(filter_string=None)
assert [mv.name for mv in mvs] == names[::-1]
assert [mv.version for mv in mvs] == [2, 2, 1, 1, 1, 1]
# order by name DESC
mvs = store.search_model_versions(filter_string=None, order_by=["name DESC"])
assert [mv.name for mv in mvs] == sorted(names)[::-1]
assert [mv.version for mv in mvs] == [2, 1, 1, 1, 2, 1]
# order by version DESC
mvs = store.search_model_versions(filter_string=None, order_by=["version_number DESC"])
assert [mv.name for mv in mvs] == ["RM1", "RM4", "RM1", "RM2", "RM3", "RM4"]
assert [mv.version for mv in mvs] == [2, 2, 1, 1, 1, 1]
# order by creation_timestamp DESC
mvs = store.search_model_versions(filter_string=None, order_by=["creation_timestamp DESC"])
assert [mv.name for mv in mvs] == names[::-1]
assert [mv.version for mv in mvs] == [2, 2, 1, 1, 1, 1]
# order by last_updated_timestamp ASC
store.update_model_version(names[0], 1, "latest updated")
mvs = store.search_model_versions(filter_string=None, order_by=["last_updated_timestamp ASC"])
assert mvs[-1].name == names[0]
assert mvs[-1].version == 1
def test_search_model_versions_order_by_errors(store):
# create some model versions
name = "RM1"
_rm_maker(store, name)
for _ in range(6):
_mv_maker(store, name=name)
query = "name LIKE 'RM%'"
# test that invalid columns throw even if they come after valid columns
with pytest.raises(
MlflowException, match=r"Invalid attribute key '.+' specified"
) as exception_context:
store.search_model_versions(
query,
page_token=None,
order_by=["name ASC", "run_id DESC"],
max_results=5,
)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# test that invalid columns with random text throw even if they come after valid columns
with pytest.raises(MlflowException, match=r"Invalid order_by clause '.+'") as exception_context:
store.search_model_versions(
query,
page_token=None,
order_by=["name ASC", "last_updated_timestamp DESC blah"],
max_results=5,
)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
def test_search_model_versions_pagination(store):
def search_versions(filter_string, page_token=None, max_results=10):
result = store.search_model_versions(
filter_string=filter_string, page_token=page_token, max_results=max_results
)
return result.to_list(), result.token
name = "test_for_search_MV_pagination"
_rm_maker(store, name)
mvs = [_mv_maker(store, name) for _ in range(50)][::-1]
# test flow with fixed max_results
returned_mvs = []
query = "name LIKE 'test_for_search_MV_pagination%'"
result, token = search_versions(query, page_token=None, max_results=5)
returned_mvs.extend(result)
while token:
result, token = search_versions(query, page_token=token, max_results=5)
returned_mvs.extend(result)
assert mvs == returned_mvs
# test that pagination will return all valid results in sorted order
# by name ascending
result, token1 = search_versions(query, max_results=5)
assert token1 is not None
assert result == mvs[0:5]
result, token2 = search_versions(query, page_token=token1, max_results=10)
assert token2 is not None
assert result == mvs[5:15]
result, token3 = search_versions(query, page_token=token2, max_results=20)
assert token3 is not None
assert result == mvs[15:35]
result, token4 = search_versions(query, page_token=token3, max_results=100)
# assert that page token is None
assert token4 is None
assert result == mvs[35:]
# test that providing a completely invalid page token throws
with pytest.raises(
MlflowException, match=r"Invalid page token, could not base64-decode"
) as exception_context:
search_versions(query, page_token="evilhax", max_results=20)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# test that providing too large of a max_results throws
with pytest.raises(
MlflowException, match=r"Invalid value for max_results\."
) as exception_context:
search_versions(query, page_token="evilhax", max_results=1e15)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
def test_search_model_versions_by_tag(store):
# create some model versions
name = "test_for_search_MV_by_tag"
_rm_maker(store, name)
run_id_1 = uuid.uuid4().hex
run_id_2 = uuid.uuid4().hex
mv1 = _mv_maker(
store,
name=name,
source="A/B",
run_id=run_id_1,
tags=[ModelVersionTag("t1", "abc"), ModelVersionTag("t2", "xyz")],
)
assert mv1.version == 1
mv2 = _mv_maker(
store,
name=name,
source="A/C",
run_id=run_id_2,
tags=[ModelVersionTag("t1", "abc"), ModelVersionTag("t2", "x123")],
)
assert mv2.version == 2
def search_versions(filter_string):
return [mvd.version for mvd in store.search_model_versions(filter_string)]
assert search_versions(f"name = '{name}' and tag.t2 = 'xyz'") == [1]
assert search_versions("name = 'wrong_name' and tag.t2 = 'xyz'") == []
assert search_versions("tag.`t2` = 'xyz'") == [1]
assert search_versions("tag.t3 = 'xyz'") == []
assert search_versions("tag.t2 != 'xy'") == [2, 1]
assert search_versions("tag.t2 LIKE 'xy%'") == [1]
assert search_versions("tag.t2 LIKE 'xY%'") == []
assert search_versions("tag.t2 ILIKE 'xY%'") == [1]
assert search_versions("tag.t2 LIKE 'x%'") == [2, 1]
assert search_versions("tag.T2 = 'xyz'") == []
assert search_versions("tag.t1 = 'abc' and tag.t2 = 'xyz'") == [1]
assert search_versions("tag.t1 = 'abc' and tag.t2 LIKE 'x%'") == [2, 1]
assert search_versions("tag.t1 = 'abc' and tag.t2 LIKE 'y%'") == []
# test filter with duplicated keys
assert search_versions("tag.t2 like 'x%' and tag.t2 != 'xyz'") == [2]
def _assert_workspace_in_main_queries(captured_sql, table_name, context_label):
"""Assert that every main search query (ORDER BY + LIMIT) includes workspace in WHERE."""
main_queries = [s for s in captured_sql if "ORDER BY" in s and table_name in s and "LIMIT" in s]
assert main_queries, f"No main search query found for {context_label}"
for sql in main_queries:
where_clause = sql.split("WHERE", 1)[-1] if "WHERE" in sql else ""
assert "workspace" in where_clause.lower(), (
f"{context_label} missing workspace predicate in WHERE — "
f"causes full table scan on (workspace, ...) PK:\n{sql}"
)
def test_search_model_versions_includes_workspace_predicate(store):
from sqlalchemy import event
name = "test_ws_predicate_mv"
_rm_maker(store, name)
_mv_maker(store, name=name, source="A/B", tags=[ModelVersionTag("t1", "abc")])
for filter_string in [
f"name = '{name}'",
"tag.t1 = 'abc'",
f"name = '{name}' AND tag.t1 = 'abc'",
]:
captured_sql: list[str] = []
def _capture(conn, cursor, statement, parameters, context, executemany):
captured_sql.append(statement)
event.listen(store.engine, "before_cursor_execute", _capture)
try:
store.search_model_versions(filter_string)
finally:
event.remove(store.engine, "before_cursor_execute", _capture)
_assert_workspace_in_main_queries(
captured_sql, "model_versions", f"search_model_versions('{filter_string}')"
)
def test_search_registered_models_includes_workspace_predicate(store):
from sqlalchemy import event
name = "test_ws_predicate_rm"
_rm_maker(store, name, tags=[RegisteredModelTag("t1", "abc")])
for filter_string in [
f"name = '{name}'",
"tag.t1 = 'abc'",
f"name = '{name}' AND tag.t1 = 'abc'",
]:
captured_sql: list[str] = []
def _capture(conn, cursor, statement, parameters, context, executemany):
captured_sql.append(statement)
event.listen(store.engine, "before_cursor_execute", _capture)
try:
store.search_registered_models(filter_string)
finally:
event.remove(store.engine, "before_cursor_execute", _capture)
_assert_workspace_in_main_queries(
captured_sql,
"registered_models",
f"search_registered_models('{filter_string}')",
)
def _search_registered_models(store, filter_string, max_results=10, order_by=None, page_token=None):
result = store.search_registered_models(
filter_string=filter_string,
max_results=max_results,
order_by=order_by,
page_token=page_token,
)
return [registered_model.name for registered_model in result], result.token
def test_search_registered_models(store):
# create some registered models
prefix = "test_for_search_"
names = [prefix + name for name in ["RM1", "RM2", "RM3", "RM4", "RM4A", "RM4ab"]]
for name in names:
_rm_maker(store, name)
# search with no filter should return all registered models
rms, _ = _search_registered_models(store, None)
assert rms == names
# equality search using name should return exactly the 1 name
rms, _ = _search_registered_models(store, f"name='{names[0]}'")
assert rms == [names[0]]
# equality search using name that is not valid should return nothing
rms, _ = _search_registered_models(store, "name='{}'".format(names[0] + "cats"))
assert rms == []
# case-sensitive prefix search using LIKE should return all the RMs
rms, _ = _search_registered_models(store, f"name LIKE '{prefix}%'")
assert rms == names
# case-sensitive prefix search using LIKE with surrounding % should return all the RMs
rms, _ = _search_registered_models(store, "name LIKE '%RM%'")
assert rms == names
# case-sensitive prefix search using LIKE with surrounding % should return all the RMs
# _e% matches test_for_search_ , so all RMs should match
rms, _ = _search_registered_models(store, "name LIKE '_e%'")
assert rms == names
# case-sensitive prefix search using LIKE should return just rm4
rms, _ = _search_registered_models(store, "name LIKE '{}%'".format(prefix + "RM4A"))
assert rms == [names[4]]
# case-sensitive prefix search using LIKE should return no models if no match
rms, _ = _search_registered_models(store, "name LIKE '{}%'".format(prefix + "cats"))
assert rms == []
# confirm that LIKE is not case-sensitive
rms, _ = _search_registered_models(store, "name lIkE '%blah%'")
assert rms == []
rms, _ = _search_registered_models(store, "name like '{}%'".format(prefix + "RM4A"))
assert rms == [names[4]]
# case-insensitive prefix search using ILIKE should return both rm5 and rm6
rms, _ = _search_registered_models(store, "name ILIKE '{}%'".format(prefix + "RM4A"))
assert rms == names[4:]
# case-insensitive postfix search with ILIKE
rms, _ = _search_registered_models(store, "name ILIKE '%RM4a%'")
assert rms == names[4:]
# case-insensitive prefix search using ILIKE should return both rm5 and rm6
rms, _ = _search_registered_models(store, "name ILIKE '{}%'".format(prefix + "cats"))
assert rms == []
# confirm that ILIKE is not case-sensitive
rms, _ = _search_registered_models(store, "name iLike '%blah%'")
assert rms == []
# confirm that ILIKE works for empty query
rms, _ = _search_registered_models(store, "name iLike '%%'")
assert rms == names
rms, _ = _search_registered_models(store, "name ilike '%RM4a%'")
assert rms == names[4:]
# cannot search by invalid comparator types
with pytest.raises(
MlflowException,
match="Parameter value is either not quoted or unidentified quote types used for "
"string value something",
) as exception_context:
_search_registered_models(store, "name!=something")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# cannot search by run_id
with pytest.raises(
MlflowException, match=r"Invalid attribute key 'run_id' specified."
) as exception_context:
_search_registered_models(store, "run_id='somerunID'")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# cannot search by source_path
with pytest.raises(
MlflowException, match=r"Invalid attribute key 'source_path' specified."
) as exception_context:
_search_registered_models(store, "source_path = 'A/D'")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# cannot search by other params
with pytest.raises(
MlflowException, match=r"Invalid clause\(s\) in filter string"
) as exception_context:
_search_registered_models(store, "evilhax = true")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# delete last registered model. search should not return the first 5
store.delete_registered_model(name=names[-1])
assert _search_registered_models(store, None, max_results=1000) == (
names[:-1],
None,
)
# equality search using name should return no names
assert _search_registered_models(store, f"name='{names[-1]}'") == ([], None)
# case-sensitive prefix search using LIKE should return all the RMs
assert _search_registered_models(store, f"name LIKE '{prefix}%'") == (
names[0:5],
None,
)
# case-insensitive prefix search using ILIKE should return both rm5 and rm6
assert _search_registered_models(store, "name ILIKE '{}%'".format(prefix + "RM4A")) == (
[names[4]],
None,
)
def test_search_registered_models_by_tag(store):
name1 = "test_for_search_RM_by_tag1"
name2 = "test_for_search_RM_by_tag2"
tags1 = [
RegisteredModelTag("t1", "abc"),
RegisteredModelTag("t2", "xyz"),
]
tags2 = [
RegisteredModelTag("t1", "abcd"),
RegisteredModelTag("t2", "xyz123"),
RegisteredModelTag("t3", "XYZ"),
]
_rm_maker(store, name1, tags1)
_rm_maker(store, name2, tags2)
rms, _ = _search_registered_models(store, "tag.t3 = 'XYZ'")
assert rms == [name2]
rms, _ = _search_registered_models(store, f"name = '{name1}' and tag.t1 = 'abc'")
assert rms == [name1]
rms, _ = _search_registered_models(store, "tag.t1 LIKE 'ab%'")
assert rms == [name1, name2]
rms, _ = _search_registered_models(store, "tag.t1 ILIKE 'aB%'")
assert rms == [name1, name2]
rms, _ = _search_registered_models(store, "tag.t1 LIKE 'ab%' AND tag.t2 LIKE 'xy%'")
assert rms == [name1, name2]
rms, _ = _search_registered_models(store, "tag.t3 = 'XYz'")
assert rms == []
rms, _ = _search_registered_models(store, "tag.T3 = 'XYZ'")
assert rms == []
rms, _ = _search_registered_models(store, "tag.t1 != 'abc'")
assert rms == [name2]
# test filter with duplicated keys
rms, _ = _search_registered_models(store, "tag.t1 != 'abcd' and tag.t1 LIKE 'ab%'")
assert rms == [name1]
def test_parse_search_registered_models_order_by():
# test that "registered_models.name ASC" is returned by default
parsed = SqlAlchemyStore._parse_search_registered_models_order_by([])
assert [str(x) for x in parsed] == ["registered_models.name ASC"]
# test that the given 'name' replaces the default one ('registered_models.name ASC')
parsed = SqlAlchemyStore._parse_search_registered_models_order_by(["name DESC"])
assert [str(x) for x in parsed] == ["registered_models.name DESC"]
# test that an exception is raised when order_by contains duplicate fields
msg = "`order_by` contains duplicate fields:"
with pytest.raises(MlflowException, match=msg):
SqlAlchemyStore._parse_search_registered_models_order_by([
"last_updated_timestamp",
"last_updated_timestamp",
])
with pytest.raises(MlflowException, match=msg):
SqlAlchemyStore._parse_search_registered_models_order_by(["timestamp", "timestamp"])
with pytest.raises(MlflowException, match=msg):
SqlAlchemyStore._parse_search_registered_models_order_by(
["timestamp", "last_updated_timestamp"],
)
with pytest.raises(MlflowException, match=msg):
SqlAlchemyStore._parse_search_registered_models_order_by(
["last_updated_timestamp ASC", "last_updated_timestamp DESC"],
)
with pytest.raises(MlflowException, match=msg):
SqlAlchemyStore._parse_search_registered_models_order_by(
["last_updated_timestamp", "last_updated_timestamp DESC"],
)
def test_search_registered_model_pagination(store):
rms = [_rm_maker(store, f"RM{i:03}").name for i in range(50)]
# test flow with fixed max_results
returned_rms = []
query = "name LIKE 'RM%'"
result, token = _search_registered_models(store, query, page_token=None, max_results=5)
returned_rms.extend(result)
while token:
result, token = _search_registered_models(store, query, page_token=token, max_results=5)
returned_rms.extend(result)
assert rms == returned_rms
# test that pagination will return all valid results in sorted order
# by name ascending
result, token1 = _search_registered_models(store, query, max_results=5)
assert token1 is not None
assert result == rms[0:5]
result, token2 = _search_registered_models(store, query, page_token=token1, max_results=10)
assert token2 is not None
assert result == rms[5:15]
result, token3 = _search_registered_models(store, query, page_token=token2, max_results=20)
assert token3 is not None
assert result == rms[15:35]
result, token4 = _search_registered_models(store, query, page_token=token3, max_results=100)
# assert that page token is None
assert token4 is None
assert result == rms[35:]
# test that providing a completely invalid page token throws
with pytest.raises(
MlflowException, match=r"Invalid page token, could not base64-decode"
) as exception_context:
_search_registered_models(store, query, page_token="evilhax", max_results=20)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# test that providing too large of a max_results throws
with pytest.raises(
MlflowException, match=r"Invalid value for request parameter max_results"
) as exception_context:
_search_registered_models(store, query, page_token="evilhax", max_results=1e15)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
def test_search_registered_model_order_by(store):
rms = []
# explicitly mock the creation_timestamps because timestamps seem to be unstable in Windows
for i in range(50):
with mock.patch(
"mlflow.store.model_registry.sqlalchemy_store.get_current_time_millis",
return_value=i,
):
rms.append(_rm_maker(store, f"RM{i:03}", _add_go_test_tags([], f"{i}")).name)
# test flow with fixed max_results and order_by (test stable order across pages)
returned_rms = []
query = "name LIKE 'RM%'"
result, token = _search_registered_models(
store, query, page_token=None, order_by=["name DESC"], max_results=5
)
returned_rms.extend(result)
while token:
result, token = _search_registered_models(
store, query, page_token=token, order_by=["name DESC"], max_results=5
)
returned_rms.extend(result)
# name descending should be the opposite order of the current order
assert rms[::-1] == returned_rms
# last_updated_timestamp descending should have the newest RMs first
result, _ = _search_registered_models(
store,
query,
page_token=None,
order_by=["last_updated_timestamp DESC"],
max_results=100,
)
assert rms[::-1] == result
# timestamp returns same result as last_updated_timestamp
result, _ = _search_registered_models(
store, query, page_token=None, order_by=["timestamp DESC"], max_results=100
)
assert rms[::-1] == result
# last_updated_timestamp ascending should have the oldest RMs first
result, _ = _search_registered_models(
store,
query,
page_token=None,
order_by=["last_updated_timestamp ASC"],
max_results=100,
)
assert rms == result
# timestamp returns same result as last_updated_timestamp
result, _ = _search_registered_models(
store, query, page_token=None, order_by=["timestamp ASC"], max_results=100
)
assert rms == result
# timestamp returns same result as last_updated_timestamp
result, _ = _search_registered_models(
store, query, page_token=None, order_by=["timestamp"], max_results=100
)
assert rms == result
# name ascending should have the original order
result, _ = _search_registered_models(
store, query, page_token=None, order_by=["name ASC"], max_results=100
)
assert rms == result
# test that no ASC/DESC defaults to ASC
result, _ = _search_registered_models(
store,
query,
page_token=None,
order_by=["last_updated_timestamp"],
max_results=100,
)
assert rms == result
with mock.patch(
"mlflow.store.model_registry.sqlalchemy_store.get_current_time_millis",
return_value=1,
):
rm1 = _rm_maker(store, "MR1", _add_go_test_tags([], "1")).name
rm2 = _rm_maker(store, "MR2", _add_go_test_tags([], "1")).name
with mock.patch(
"mlflow.store.model_registry.sqlalchemy_store.get_current_time_millis",
return_value=2,
):
rm3 = _rm_maker(store, "MR3", _add_go_test_tags([], "2")).name
rm4 = _rm_maker(store, "MR4", _add_go_test_tags([], "2")).name
query = "name LIKE 'MR%'"
# test with multiple clauses
result, _ = _search_registered_models(
store,
query,
page_token=None,
order_by=["last_updated_timestamp ASC", "name DESC"],
max_results=100,
)
assert result == [rm2, rm1, rm4, rm3]
result, _ = _search_registered_models(
store,
query,
page_token=None,
order_by=["timestamp ASC", "name DESC"],
max_results=100,
)
assert result == [rm2, rm1, rm4, rm3]
# confirm that name ascending is the default, even if ties exist on other fields
result, _ = _search_registered_models(
store, query, page_token=None, order_by=[], max_results=100
)
assert result == [rm1, rm2, rm3, rm4]
# test default tiebreak with descending timestamps
result, _ = _search_registered_models(
store,
query,
page_token=None,
order_by=["last_updated_timestamp DESC"],
max_results=100,
)
assert result == [rm3, rm4, rm1, rm2]
# test timestamp parsing
result, _ = _search_registered_models(
store, query, page_token=None, order_by=["timestamp\tASC"], max_results=100
)
assert result == [rm1, rm2, rm3, rm4]
result, _ = _search_registered_models(
store, query, page_token=None, order_by=["timestamp\r\rASC"], max_results=100
)
assert result == [rm1, rm2, rm3, rm4]
result, _ = _search_registered_models(
store, query, page_token=None, order_by=["timestamp\nASC"], max_results=100
)
assert result == [rm1, rm2, rm3, rm4]
result, _ = _search_registered_models(
store, query, page_token=None, order_by=["timestamp ASC"], max_results=100
)
assert result == [rm1, rm2, rm3, rm4]
# validate order by key is case-insensitive
result, _ = _search_registered_models(
store, query, page_token=None, order_by=["timestamp asc"], max_results=100
)
assert result == [rm1, rm2, rm3, rm4]
result, _ = _search_registered_models(
store, query, page_token=None, order_by=["timestamp aSC"], max_results=100
)
assert result == [rm1, rm2, rm3, rm4]
result, _ = _search_registered_models(
store,
query,
page_token=None,
order_by=["timestamp desc", "name desc"],
max_results=100,
)
assert result == [rm4, rm3, rm2, rm1]
result, _ = _search_registered_models(
store,
query,
page_token=None,
order_by=["timestamp deSc", "name deSc"],
max_results=100,
)
assert result == [rm4, rm3, rm2, rm1]
def test_search_registered_model_order_by_errors(store):
query = "name LIKE 'RM%'"
# test that invalid columns throw even if they come after valid columns
with pytest.raises(
MlflowException, match=r"Invalid order by key '.+' specified"
) as exception_context:
_search_registered_models(
store,
query,
page_token=None,
order_by=["name ASC", "creation_timestamp DESC"],
max_results=5,
)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# test that invalid columns with random text throw even if they come after valid columns
with pytest.raises(MlflowException, match=r"Invalid order_by clause '.+'") as exception_context:
_search_registered_models(
store,
query,
page_token=None,
order_by=["name ASC", "last_updated_timestamp DESC blah"],
max_results=5,
)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
def test_set_model_version_tag(store):
name1 = "SetModelVersionTag_TestMod"
name2 = "SetModelVersionTag_TestMod 2"
initial_tags = [
ModelVersionTag("key", "value"),
ModelVersionTag("anotherKey", "some other value"),
]
_rm_maker(store, name1)
_rm_maker(store, name2)
run_id_1 = uuid.uuid4().hex
run_id_2 = uuid.uuid4().hex
run_id_3 = uuid.uuid4().hex
_mv_maker(store, name1, "A/B", run_id_1, initial_tags)
_mv_maker(store, name1, "A/C", run_id_2, initial_tags)
_mv_maker(store, name2, "A/D", run_id_3, initial_tags)
new_tag = ModelVersionTag("randomTag", "not a random value")
store.set_model_version_tag(name1, 1, new_tag)
all_tags = initial_tags + [new_tag]
rm1mv1 = store.get_model_version(name1, 1)
assert rm1mv1.tags == {tag.key: tag.value for tag in all_tags}
# test overriding a tag with the same key
overriding_tag = ModelVersionTag("key", "overriding")
store.set_model_version_tag(name1, 1, overriding_tag)
all_tags = [tag for tag in all_tags if tag.key != "key"] + [overriding_tag]
rm1mv1 = store.get_model_version(name1, 1)
assert rm1mv1.tags == {tag.key: tag.value for tag in all_tags}
# does not affect other model versions with the same key
rm1mv2 = store.get_model_version(name1, 2)
rm2mv1 = store.get_model_version(name2, 1)
assert rm1mv2.tags == {tag.key: tag.value for tag in initial_tags}
assert rm2mv1.tags == {tag.key: tag.value for tag in initial_tags}
# can not set tag on deleted (non-existed) model version
store.delete_model_version(name1, 2)
with pytest.raises(
MlflowException, match=rf"Model Version \(name={name1}, version=2\) not found"
) as exception_context:
store.set_model_version_tag(name1, 2, overriding_tag)
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
# test cannot set tags that are too long
long_tag = ModelVersionTag("longTagKey", "a" * 100_001)
with pytest.raises(
MlflowException,
match=r"'value' exceeds the maximum length of \d+ characters",
) as exception_context:
store.set_model_version_tag(name1, 1, long_tag)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# test can set tags that are somewhat long
long_tag = ModelVersionTag("longTagKey", "a" * 4999)
store.set_model_version_tag(name1, 1, long_tag)
# can not set invalid tag
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'key'"
) as exception_context:
store.set_model_version_tag(name2, 1, ModelVersionTag(key=None, value=""))
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# can not use invalid model name or version
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'name'"
) as exception_context:
store.set_model_version_tag(None, 1, ModelVersionTag(key="key", value="value"))
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
with pytest.raises(
MlflowException, match=r"Parameter 'version' must be an integer, got 'I am not a version'"
) as exception_context:
store.set_model_version_tag(
name2, "I am not a version", ModelVersionTag(key="key", value="value")
)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
def test_delete_model_version_tag(store):
name1 = "DeleteModelVersionTag_TestMod"
name2 = "DeleteModelVersionTag_TestMod 2"
initial_tags = [
ModelVersionTag("key", "value"),
ModelVersionTag("anotherKey", "some other value"),
]
_rm_maker(store, name1)
_rm_maker(store, name2)
run_id_1 = uuid.uuid4().hex
run_id_2 = uuid.uuid4().hex
run_id_3 = uuid.uuid4().hex
_mv_maker(store, name1, "A/B", run_id_1, initial_tags)
_mv_maker(store, name1, "A/C", run_id_2, initial_tags)
_mv_maker(store, name2, "A/D", run_id_3, initial_tags)
new_tag = ModelVersionTag("randomTag", "not a random value")
store.set_model_version_tag(name1, 1, new_tag)
store.delete_model_version_tag(name1, 1, "randomTag")
rm1mv1 = store.get_model_version(name1, 1)
assert rm1mv1.tags == {tag.key: tag.value for tag in initial_tags}
# testing deleting a key does not affect other model versions with the same key
store.delete_model_version_tag(name1, 1, "key")
rm1mv1 = store.get_model_version(name1, 1)
rm1mv2 = store.get_model_version(name1, 2)
rm2mv1 = store.get_model_version(name2, 1)
assert rm1mv1.tags == {"anotherKey": "some other value"}
assert rm1mv2.tags == {tag.key: tag.value for tag in initial_tags}
assert rm2mv1.tags == {tag.key: tag.value for tag in initial_tags}
# delete tag that is already deleted does nothing
store.delete_model_version_tag(name1, 1, "key")
rm1mv1 = store.get_model_version(name1, 1)
assert rm1mv1.tags == {"anotherKey": "some other value"}
# can not delete tag on deleted (non-existed) model version
store.delete_model_version(name2, 1)
with pytest.raises(
MlflowException, match=rf"Model Version \(name={name2}, version=1\) not found"
) as exception_context:
store.delete_model_version_tag(name2, 1, "key")
assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
# can not delete tag with invalid key
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'key'"
) as exception_context:
store.delete_model_version_tag(name1, 2, None)
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
# can not use invalid model name or version
with pytest.raises(
MlflowException, match=r"Missing value for required parameter 'name'."
) as exception_context:
store.delete_model_version_tag(None, 2, "key")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
with pytest.raises(
MlflowException, match=r"Parameter 'version' must be an integer, got 'I am not a version'"
) as exception_context:
store.delete_model_version_tag(name1, "I am not a version", "key")
assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
def _setup_and_test_aliases(store, model_name):
store.create_registered_model(model_name)
run_id_1 = uuid.uuid4().hex
run_id_2 = uuid.uuid4().hex
store.create_model_version(model_name, "v1", run_id_1)
store.create_model_version(model_name, "v2", run_id_2)
store.set_registered_model_alias(model_name, "test_alias", "2")
model = store.get_registered_model(model_name)
assert model.aliases == {"test_alias": 2}
mv1 = store.get_model_version(model_name, 1)
mv2 = store.get_model_version(model_name, 2)
assert mv1.aliases == []
assert mv2.aliases == ["test_alias"]
def test_set_registered_model_alias(store):
_setup_and_test_aliases(store, "SetRegisteredModelAlias_TestMod")
def test_delete_registered_model_alias(store):
model_name = "DeleteRegisteredModelAlias_TestMod"
_setup_and_test_aliases(store, model_name)
store.delete_registered_model_alias(model_name, "test_alias")
model = store.get_registered_model(model_name)
assert model.aliases == {}
mv2 = store.get_model_version(model_name, 2)
assert mv2.aliases == []
def test_get_model_version_by_alias(store):
model_name = "GetModelVersionByAlias_TestMod"
_setup_and_test_aliases(store, model_name)
mv = store.get_model_version_by_alias(model_name, "test_alias")
assert mv.aliases == ["test_alias"]
def test_delete_model_version_deletes_alias(store):
model_name = "DeleteModelVersionDeletesAlias_TestMod"
_setup_and_test_aliases(store, model_name)
store.delete_model_version(model_name, 2)
model = store.get_registered_model(model_name)
assert model.aliases == {}
with pytest.raises(
MlflowException,
match=r"Registered model alias test_alias not found.",
):
store.get_model_version_by_alias(model_name, "test_alias")
def test_delete_model_deletes_alias(store):
model_name = "DeleteModelDeletesAlias_TestMod"
_setup_and_test_aliases(store, model_name)
store.delete_registered_model(model_name)
with pytest.raises(
MlflowException,
match=rf"Registered Model with name={model_name} not found",
):
store.get_model_version_by_alias(model_name, "test_alias")
@pytest.mark.parametrize("copy_to_same_model", [False, True])
def test_copy_model_version(store, copy_to_same_model):
name1 = "test_for_copy_MV1"
store.create_registered_model(name1)
src_tags = [
ModelVersionTag("key", "value"),
ModelVersionTag("anotherKey", "some other value"),
]
src_mv = _mv_maker(
store,
name1,
tags=src_tags,
run_link="dummylink",
description="test description",
)
# Make some changes to the src MV that won't be copied over
store.transition_model_version_stage(
name1, src_mv.version, "Production", archive_existing_versions=False
)
copy_rm_name = name1 if copy_to_same_model else "test_for_copy_MV2"
copy_mv_version = 2 if copy_to_same_model else 1
timestamp = time.time()
dst_mv = store.copy_model_version(src_mv, copy_rm_name)
assert dst_mv.name == copy_rm_name
assert dst_mv.version == copy_mv_version
copied_mv = store.get_model_version(dst_mv.name, dst_mv.version)
assert copied_mv.name == copy_rm_name
assert int(copied_mv.version) == copy_mv_version
assert copied_mv.current_stage == "None"
assert copied_mv.creation_timestamp >= timestamp
assert copied_mv.last_updated_timestamp >= timestamp
assert copied_mv.description == "test description"
assert copied_mv.source == f"models:/{src_mv.name}/{src_mv.version}"
assert store.get_model_version_download_uri(dst_mv.name, dst_mv.version) == src_mv.source
assert copied_mv.run_link == "dummylink"
assert copied_mv.run_id == src_mv.run_id
assert copied_mv.status == "READY"
assert copied_mv.status_message is None
assert copied_mv.tags == {"key": "value", "anotherKey": "some other value"}
# Copy a model version copy
double_copy_mv = store.copy_model_version(copied_mv, "test_for_copy_MV3")
assert double_copy_mv.source == f"models:/{copied_mv.name}/{copied_mv.version}"
assert store.get_model_version_download_uri(dst_mv.name, dst_mv.version) == src_mv.source
def test_search_prompts(store):
store.create_registered_model("model", tags=[RegisteredModelTag(key="fruit", value="apple")])
store.create_registered_model(
"prompt_1", tags=[RegisteredModelTag(key=IS_PROMPT_TAG_KEY, value="true")]
)
store.create_registered_model(
"prompt_2",
tags=[
RegisteredModelTag(key=IS_PROMPT_TAG_KEY, value="true"),
RegisteredModelTag(key="fruit", value="apple"),
],
)
# By default, should not return prompts
rms = store.search_registered_models(max_results=10)
assert len(rms) == 1
assert rms[0].name == "model"
rms = store.search_registered_models(filter_string="tags.fruit = 'apple'", max_results=10)
assert len(rms) == 1
assert rms[0].name == "model"
rms = store.search_registered_models(filter_string="name = 'prompt_1'", max_results=10)
assert len(rms) == 0
rms = store.search_registered_models(
filter_string="tags.`mlflow.prompt.is_prompt` = 'false'", max_results=10
)
assert len(rms) == 1
assert rms[0].name == "model"
rms = store.search_registered_models(
filter_string="tags.`mlflow.prompt.is_prompt` != 'true'", max_results=10
)
assert len(rms) == 1
assert rms[0].name == "model"
# Search for prompts
rms = store.search_registered_models(
filter_string="tags.`mlflow.prompt.is_prompt` = 'true'", max_results=10
)
assert len(rms) == 2
assert {rm.name for rm in rms} == {"prompt_1", "prompt_2"}
rms = store.search_registered_models(
filter_string="name = 'prompt_1' and tags.`mlflow.prompt.is_prompt` = 'true'",
max_results=10,
)
assert len(rms) == 1
assert rms[0].name == "prompt_1"
rms = store.search_registered_models(
filter_string="tags.`mlflow.prompt.is_prompt` = 'true' and tags.fruit = 'apple'",
max_results=10,
)
assert len(rms) == 1
assert rms[0].name == "prompt_2"
def test_search_prompts_versions(store):
# A Model
store.create_registered_model("model")
store.create_model_version(
"model", "1", "dummy_source", tags=[ModelVersionTag(key="fruit", value="apple")]
)
# A Prompt with 1 version
store.create_registered_model(
"prompt_1", tags=[RegisteredModelTag(key=IS_PROMPT_TAG_KEY, value="true")]
)
store.create_model_version(
"prompt_1", "1", "dummy_source", tags=[ModelVersionTag(key=IS_PROMPT_TAG_KEY, value="true")]
)
# A Prompt with 2 versions
store.create_registered_model(
"prompt_2",
tags=[RegisteredModelTag(key=IS_PROMPT_TAG_KEY, value="true")],
)
store.create_model_version(
"prompt_2",
"1",
"dummy_source",
tags=[
ModelVersionTag(key=IS_PROMPT_TAG_KEY, value="true"),
ModelVersionTag(key="fruit", value="apple"),
],
)
store.create_model_version(
"prompt_2",
"2",
"dummy_source",
tags=[
ModelVersionTag(key=IS_PROMPT_TAG_KEY, value="true"),
ModelVersionTag(key="fruit", value="orange"),
],
)
# Searching model versions should not return prompts by default either
mvs = store.search_model_versions(max_results=10)
assert len(mvs) == 1
assert mvs[0].name == "model"
mvs = store.search_model_versions(filter_string="tags.fruit = 'apple'", max_results=10)
assert len(mvs) == 1
assert mvs[0].name == "model"
mvs = store.search_model_versions(
filter_string="tags.`mlflow.prompt.is_prompt` = 'false'", max_results=10
)
assert len(mvs) == 1
assert mvs[0].name == "model"
mvs = store.search_model_versions(
filter_string="tags.`mlflow.prompt.is_prompt` != 'true'", max_results=10
)
assert len(mvs) == 1
assert mvs[0].name == "model"
# Search for prompts via search_model_versions
mvs = store.search_model_versions(
filter_string="tags.`mlflow.prompt.is_prompt` = 'true'", max_results=10
)
assert len(mvs) == 3
mvs = store.search_model_versions(
filter_string="tags.`mlflow.prompt.is_prompt` = 'true' and name = 'prompt_2'",
max_results=10,
)
assert len(mvs) == 2
mvs = store.search_model_versions(
filter_string="tags.`mlflow.prompt.is_prompt` = 'true' and tags.fruit = 'apple'",
max_results=10,
)
assert len(mvs) == 1
assert mvs[0].name == "prompt_2"
def test_search_prompt_versions(store):
# Create a prompt with 3 versions
store.create_registered_model(
"my_prompt", tags=[RegisteredModelTag(key=IS_PROMPT_TAG_KEY, value="true")]
)
for i in range(1, 4):
store.create_model_version(
"my_prompt",
str(i),
"dummy_source",
tags=[
ModelVersionTag(key=IS_PROMPT_TAG_KEY, value="true"),
ModelVersionTag(key=PROMPT_TEXT_TAG_KEY, value=f"Hello {{{{name}}}} v{i}"),
],
)
# Create a different prompt to verify filtering by name
store.create_registered_model(
"other_prompt", tags=[RegisteredModelTag(key=IS_PROMPT_TAG_KEY, value="true")]
)
store.create_model_version(
"other_prompt",
"1",
"dummy_source",
tags=[
ModelVersionTag(key=IS_PROMPT_TAG_KEY, value="true"),
ModelVersionTag(key=PROMPT_TEXT_TAG_KEY, value="Other prompt text"),
],
)
# Search all versions of my_prompt
results = store.search_prompt_versions("my_prompt")
assert len(results) == 3
# Should be ordered by version descending
assert [pv.version for pv in results] == [3, 2, 1]
assert all(pv.name == "my_prompt" for pv in results)
# Pagination with max_results
page1 = store.search_prompt_versions("my_prompt", max_results=2)
assert len(page1) == 2
assert [pv.version for pv in page1] == [3, 2]
assert page1.token is not None
# Fetch next page
page2 = store.search_prompt_versions("my_prompt", max_results=2, page_token=page1.token)
assert len(page2) == 1
assert page2[0].version == 1
assert page2.token is None
# Search other_prompt returns only its versions
results = store.search_prompt_versions("other_prompt")
assert len(results) == 1
assert results[0].name == "other_prompt"
# Searching a non-existent prompt raises
with pytest.raises(MlflowException, match="not found"):
store.search_prompt_versions("nonexistent_prompt")
# Searching a model (not a prompt) raises
store.create_registered_model("a_model")
with pytest.raises(MlflowException, match="registered as a model, not a prompt"):
store.search_prompt_versions("a_model")
def test_create_registered_model_handle_prompt_properly(store):
prompt_tags = [RegisteredModelTag(key=IS_PROMPT_TAG_KEY, value="true")]
store.create_registered_model("model")
store.create_registered_model("prompt", tags=prompt_tags)
with pytest.raises(MlflowException, match=r"Registered Model \(name=model\) already exists"):
store.create_registered_model("model")
with pytest.raises(MlflowException, match=r"Prompt \(name=prompt\) already exists"):
store.create_registered_model("prompt", tags=prompt_tags)
with pytest.raises(
MlflowException,
match=r"Tried to create a prompt with name 'model', "
r"but the name is already taken by a registered model.",
):
store.create_registered_model("model", tags=prompt_tags)
with pytest.raises(
MlflowException,
match=r"Tried to create a registered model with name 'prompt', "
r"but the name is already taken by a prompt.",
):
store.create_registered_model("prompt")
def test_create_webhook(store):
events = [
WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED),
WebhookEvent(WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED),
]
webhook = store.create_webhook(
name="test_webhook",
url="https://example.com/webhook",
events=events,
description="Test webhook",
secret="secret123",
status=WebhookStatus.ACTIVE,
)
assert webhook.name == "test_webhook"
assert webhook.url == "https://example.com/webhook"
assert webhook.events == events
assert webhook.description == "Test webhook"
assert webhook.status == WebhookStatus.ACTIVE
assert webhook.webhook_id is not None
assert webhook.creation_timestamp is not None
assert webhook.last_updated_timestamp is not None
assert webhook.secret == "secret123"
# Shared test data for invalid webhook names
INVALID_WEBHOOK_NAMES = [
("", r"is invalid"),
(" ", r"is invalid"),
("webhook<script>", r"is invalid"),
("webhook@test", r"is invalid"),
("webhook#hash", r"is invalid"),
("webhook/slash", r"is invalid"),
("webhook\\backslash", r"is invalid"),
("-webhook", r"is invalid"), # Must start with letter or digit
("webhook-", r"is invalid"), # Must end with letter or digit
("_webhook", r"is invalid"), # Must start with letter or digit
("webhook_", r"is invalid"), # Must end with letter or digit
(".webhook", r"is invalid"), # Must start with letter or digit
("webhook.", r"is invalid"), # Must end with letter or digit
("a" * 64, r"is invalid"), # Too long (max 63 chars)
]
# Shared test data for valid webhook names
VALID_WEBHOOK_NAMES = [
"a", # Single character letter
"1", # Single character digit
"a1", # Two characters
"1a", # Start with digit, end with letter
"webhook123", # Alphanumeric
"web_hook", # With underscore
"web-hook", # With hyphen
"web.hook", # With dot
"web_hook-123.test", # Mixed special chars
"A" * 63, # Maximum length
"1" + "a" * 61 + "1", # Maximum length with digit start/end
"WebHook123", # Mixed case
]
@pytest.mark.parametrize(("invalid_name", "expected_match"), INVALID_WEBHOOK_NAMES)
def test_create_webhook_invalid_names(store, invalid_name, expected_match):
with pytest.raises(MlflowException, match=expected_match):
store.create_webhook(
name=invalid_name,
url="https://example.com",
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
)
@pytest.mark.parametrize("valid_name", VALID_WEBHOOK_NAMES)
def test_create_webhook_valid_names(store, valid_name):
webhook = store.create_webhook(
name=valid_name,
url="https://example.com",
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
)
assert webhook.name == valid_name
@pytest.mark.parametrize(
("invalid_url", "expected_match"),
[
("", r"Webhook URL cannot be empty or just whitespace"),
(" ", r"Webhook URL cannot be empty or just whitespace"),
("example.com/webhook", r"Invalid webhook URL"),
("ftp://example.com/webhook", r"Invalid webhook URL scheme"),
("http://[invalid-url", r"Invalid webhook URL"),
("invalid_url", r"Invalid webhook URL"),
],
)
def test_create_webhook_invalid_urls(store, invalid_url, expected_match):
with pytest.raises(MlflowException, match=expected_match):
store.create_webhook(
name="test",
url=invalid_url,
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
)
def test_create_webhook_invalid_events(store):
# Test empty events
with pytest.raises(MlflowException, match="Webhook events must be a non-empty list"):
store.create_webhook(name="test", url="https://example.com", events=[])
# Test non-list events
with pytest.raises(MlflowException, match="Webhook events must be a non-empty list"):
store.create_webhook(name="test", url="https://example.com", events=())
# Test list with non-WebhookEvent items
with pytest.raises(MlflowException, match="Webhook events must be a non-empty list"):
store.create_webhook(name="test", url="https://example.com", events=[1, 2, 3])
def test_get_webhook(store):
events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)]
created_webhook = store.create_webhook(
name="test_webhook", url="https://example.com/webhook", events=events
)
retrieved_webhook = store.get_webhook(created_webhook.webhook_id)
assert retrieved_webhook.webhook_id == created_webhook.webhook_id
assert retrieved_webhook.name == "test_webhook"
assert retrieved_webhook.url == "https://example.com/webhook"
assert retrieved_webhook.events == events
def test_get_webhook_not_found(store):
with pytest.raises(MlflowException, match="Webhook with ID nonexistent not found"):
store.get_webhook("nonexistent")
def test_list_webhooks(store):
# Create multiple webhooks
webhook1 = store.create_webhook(
name="webhook1",
url="https://example.com/1",
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
)
webhook2 = store.create_webhook(
name="webhook2",
url="https://example.com/2",
events=[WebhookEvent(WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED)],
)
webhooks_page = store.list_webhooks()
assert len(webhooks_page) == 2
assert webhooks_page.token is None
webhook_ids = {w.webhook_id for w in webhooks_page}
assert webhook1.webhook_id in webhook_ids
assert webhook2.webhook_id in webhook_ids
def test_list_webhooks_pagination(store):
# Create more webhooks than max_results
for i in range(5):
store.create_webhook(
name=f"webhook{i}",
url=f"https://example.com/{i}",
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
)
# Test pagination with max_results=2
webhooks_page = store.list_webhooks(max_results=2)
assert len(webhooks_page) == 2
assert webhooks_page.token is not None
# Get next page
next_webhooks_page = store.list_webhooks(max_results=2, page_token=webhooks_page.token)
assert len(next_webhooks_page) == 2
assert next_webhooks_page.token is not None
# Verify we don't get duplicates
first_page_ids = {w.webhook_id for w in webhooks_page}
second_page_ids = {w.webhook_id for w in next_webhooks_page}
assert first_page_ids.isdisjoint(second_page_ids)
def test_list_webhooks_invalid_max_results(store):
with pytest.raises(MlflowException, match="max_results must be between 1 and 1000"):
store.list_webhooks(max_results=1001)
def test_update_webhook(store):
events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)]
webhook = store.create_webhook(
name="original_name", url="https://example.com/original", events=events
)
# Update webhook
new_events = [
WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED),
WebhookEvent(WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED),
]
updated_webhook = store.update_webhook(
webhook_id=webhook.webhook_id,
name="updated_name",
url="https://example.com/updated",
events=new_events,
description="Updated description",
secret="new_secret",
status=WebhookStatus.DISABLED,
)
assert updated_webhook.webhook_id == webhook.webhook_id
assert updated_webhook.name == "updated_name"
assert updated_webhook.url == "https://example.com/updated"
assert updated_webhook.events == new_events
assert updated_webhook.description == "Updated description"
assert updated_webhook.status == WebhookStatus.DISABLED
assert updated_webhook.last_updated_timestamp > webhook.last_updated_timestamp
def test_update_webhook_partial(store):
events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)]
webhook = store.create_webhook(
name="original_name", url="https://example.com/original", events=events
)
# Update only name
updated_webhook = store.update_webhook(webhook_id=webhook.webhook_id, name="new_name")
assert updated_webhook.name == "new_name"
assert updated_webhook.url == "https://example.com/original" # Should remain unchanged
assert updated_webhook.events == events # Should remain unchanged
def test_update_webhook_not_found(store):
with pytest.raises(MlflowException, match="Webhook with ID nonexistent not found"):
store.update_webhook(
webhook_id="nonexistent", name="new_name", url="https://example.com/new"
)
def test_update_webhook_invalid_events(store):
# Create a valid webhook first
webhook = store.create_webhook(
name="test_webhook",
url="https://example.com/webhook",
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
)
with pytest.raises(MlflowException, match="Webhook events must be a non-empty list"):
store.update_webhook(webhook_id=webhook.webhook_id, events=[])
# Test non-list events
with pytest.raises(MlflowException, match="Webhook events must be a non-empty list"):
store.update_webhook(webhook_id=webhook.webhook_id, events=())
# Test list with non-WebhookEvent items
with pytest.raises(MlflowException, match="Webhook events must be a non-empty list"):
store.update_webhook(webhook_id=webhook.webhook_id, events=[1, 2, 3])
@pytest.mark.parametrize(("invalid_name", "expected_match"), INVALID_WEBHOOK_NAMES)
def test_update_webhook_invalid_names(store, invalid_name, expected_match):
# Create a valid webhook first
webhook = store.create_webhook(
name="test_webhook",
url="https://example.com/webhook",
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
)
with pytest.raises(MlflowException, match=expected_match):
store.update_webhook(webhook_id=webhook.webhook_id, name=invalid_name)
@pytest.mark.parametrize(
("invalid_url", "expected_match"),
[
(" ", r"Webhook URL cannot be empty or just whitespace"),
("ftp://example.com", r"Invalid webhook URL scheme"),
("http://[invalid", r"Invalid webhook URL"),
],
)
def test_update_webhook_invalid_urls(store, invalid_url, expected_match):
# Create a valid webhook first
webhook = store.create_webhook(
name="test_webhook",
url="https://example.com/webhook",
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
)
with pytest.raises(MlflowException, match=expected_match):
store.update_webhook(webhook_id=webhook.webhook_id, url=invalid_url)
def test_delete_webhook(store):
events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)]
webhook = store.create_webhook(
name="test_webhook",
url="https://example.com/webhook",
events=events,
)
store.delete_webhook(webhook.webhook_id)
with pytest.raises(MlflowException, match=r"Webhook with ID .* not found"):
store.get_webhook(webhook.webhook_id)
webhooks_page = store.list_webhooks()
webhook_ids = {w.webhook_id for w in webhooks_page}
assert webhook.webhook_id not in webhook_ids
def test_delete_webhook_not_found(store):
with pytest.raises(MlflowException, match="Webhook with ID nonexistent not found"):
store.delete_webhook("nonexistent")
def test_webhook_status_transitions(store):
events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)]
webhook = store.create_webhook(
name="test_webhook",
url="https://example.com/webhook",
events=events,
status=WebhookStatus.ACTIVE,
)
assert webhook.status == WebhookStatus.ACTIVE
# Update to inactive
updated_webhook = store.update_webhook(
webhook_id=webhook.webhook_id, status=WebhookStatus.DISABLED
)
assert updated_webhook.status == WebhookStatus.DISABLED
# Update back to active
updated_webhook = store.update_webhook(
webhook_id=webhook.webhook_id, status=WebhookStatus.ACTIVE
)
assert updated_webhook.status == WebhookStatus.ACTIVE
def test_webhook_secret_encryption(store):
store.create_webhook(
name="test_webhook",
url="https://example.com/webhook",
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
secret="my_secret",
)
engine = create_engine(store.db_uri)
with engine.connect() as conn:
(raw_secret,) = conn.execute(text("SELECT secret FROM webhooks")).fetchone()
assert raw_secret is not None
assert raw_secret != "my_secret" # Should be encrypted
def test_create_model_version_with_model_id_and_no_run_id(store):
name = "test_model_with_model_id"
_rm_maker(store, name)
mock_run_id = "mock-run-id-123"
mock_logged_model = mock.MagicMock()
mock_logged_model.source_run_id = mock_run_id
with mock.patch(
"mlflow.store.model_registry.sqlalchemy_store.MlflowClient"
) as mock_client_class:
mock_client = mock.MagicMock()
mock_client_class.return_value = mock_client
mock_client.get_logged_model.return_value = mock_logged_model
mv = store.create_model_version(
name=name,
source="path/to/source",
run_id=None,
model_id="test-model-id-456",
)
mock_client.get_logged_model.assert_called_once_with("test-model-id-456")
assert mv.run_id == mock_run_id
mvd = store.get_model_version(name=mv.name, version=mv.version)
assert mvd.run_id == mock_run_id
def test_create_model_version_concurrent(store):
name = "test_concurrent_mv"
_rm_maker(store, name)
num_threads = 4
versions_per_thread = 5
results = []
def create_versions():
return [
store.create_model_version(name, "path/to/source", uuid.uuid4().hex).version
for _ in range(versions_per_thread)
]
with concurrent.futures.ThreadPoolExecutor(
max_workers=num_threads, thread_name_prefix="create_model_version"
) as executor:
futures = [executor.submit(create_versions) for _ in range(num_threads)]
for f in concurrent.futures.as_completed(futures):
results.extend(f.result())
# All versions should be unique
assert len(results) == len(set(results))
assert len(results) == num_threads * versions_per_thread
Reference in New Issue View Git Blame Copy Permalink