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_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, tag):
"""
Set a tag for the registered model.
:param name: Registered model name.
:param tag: :py:class:`mlflow.entities.model_registry.RegisteredModelTag` instance to log.
:return: None
"""
req_body = message_to_json(SetRegisteredModelTag(name=name, key=tag.key, value=tag.value))
self._call_endpoint(SetRegisteredModelTag, req_body)