def test_create_registered_model(self, mock_http):
tags = [
RegisteredModelTag(key="key", value="value"),
RegisteredModelTag(key="anotherKey", value="some other value")
]
self.store.create_registered_model("model_1", tags)
self._verify_requests(
mock_http, "registered-models/create", "POST",
CreateRegisteredModel(name="model_1",
tags=[tag.to_proto() for tag in tags]))
def test_create_registered_model(mock_get_request_message, mock_model_registry_store):
tags = [
RegisteredModelTag(key="key", value="value"),
RegisteredModelTag(key="anotherKey", value="some other value"),
]
mock_get_request_message.return_value = CreateRegisteredModel(
name="model_1", tags=[tag.to_proto() for tag in tags]
)
rm = RegisteredModel("model_1", tags=tags)
mock_model_registry_store.create_registered_model.return_value = rm
resp = _create_registered_model()
_, args = mock_model_registry_store.create_registered_model.call_args
assert args["name"] == "model_1"
assert {tag.key: tag.value for tag in args["tags"]} == {tag.key: tag.value for tag in tags}
assert json.loads(resp.get_data()) == {"registered_model": jsonify(rm)}
def _set_registered_model_tag():
request_message = _get_request_message(SetRegisteredModelTag())
tag = RegisteredModelTag(key=request_message.key, value=request_message.value)
_get_model_registry_store().set_registered_model_tag(
name=request_message.name,
tag=tag)
return _wrap_response(SetRegisteredModelTag.Response())
def test_set_registered_model_tag(self, mock_http):
name = "model_1"
tag = RegisteredModelTag(key="key", value="value")
self.store.set_registered_model_tag(name=name, tag=tag)
self._verify_requests(
mock_http, "registered-models/set-tag", "POST",
SetRegisteredModelTag(name=name, key=tag.key, value=tag.value))
def test_create_registered_model(mock_store):
tags_dict = {"key": "value", "another key": "some other value"}
tags = [RegisteredModelTag(key, value) for key, value in tags_dict.items()]
mock_store.create_registered_model.return_value = RegisteredModel("Model 1", tags=tags)
result = newModelRegistryClient().create_registered_model("Model 1", tags_dict)
mock_store.create_registered_model.assert_called_once_with("Model 1", tags)
assert result.name == "Model 1"
assert result.tags == tags_dict
def test_set_registered_model_tag(mock_get_request_message,
mock_model_registry_store):
name = "model1"
tag = RegisteredModelTag(key="some weird key", value="some value")
mock_get_request_message.return_value = SetRegisteredModelTag(
name=name, key=tag.key, value=tag.value)
_set_registered_model_tag()
_, args = mock_model_registry_store.set_registered_model_tag.call_args
assert args == {"name": name, "tag": tag}
def set_registered_model_tag(self, name, key, value):
"""
Set a tag for the registered model.
:param name: Registered model name.
:param key: Tag key to log.
:param value: Tag value log.
:return: None
"""
self.store.set_registered_model_tag(name, RegisteredModelTag(key, str(value)))
def test_get_registered_model_details(mock_store):
name = "Model 1"
tags = [
RegisteredModelTag("key", "value"),
RegisteredModelTag("another key", "some other value")
]
mock_store.get_registered_model.return_value = RegisteredModel(
name,
"1263283747835",
"1283168374623874",
"I am a model", [
_model_version("Model 1", 3, "None"),
_model_version("Model 1", 2, "Staging"),
_model_version("Model 1", 1, "Production")
],
tags=tags)
result = newModelRegistryClient().get_registered_model(name)
mock_store.get_registered_model.assert_called_once()
assert result.name == name
assert len(result.latest_versions) == 3
assert result.tags == {tag.key: tag.value for tag in tags}
def create_registered_model(self, name, tags=None):
"""
Create a new registered model in backend store.
:param name: Name of the new model. This is expected to be unique in the backend store.
:param tags: A dictionary of key-value pairs that are converted into
:py:class:`mlflow.entities.model_registry.RegisteredModelTag` objects.
:return: A single object of :py:class:`mlflow.entities.model_registry.RegisteredModel`
created by backend.
"""
# TODO: Do we want to validate the name is legit here - non-empty without "/" and ":" ?
# Those are constraints applicable to any backend, given the model URI format.
tags = tags if tags else {}
tags = [RegisteredModelTag(key, str(value)) for key, value in tags.items()]
return self.store.create_registered_model(name, tags)
def to_mlflow_entity(self):
return RegisteredModelTag(self.key, self.value)