def _rename_registered_model():
request_message = _get_request_message(RenameRegisteredModel())
name = request_message.name
new_name = request_message.new_name
registered_model = _get_model_registry_store().rename_registered_model(
name=name, new_name=new_name)
response_message = RenameRegisteredModel.Response(registered_model=registered_model.to_proto())
return _wrap_response(response_message)
def test_update_registered_model_name(self, mock_http):
name = "model_1"
new_name = "model_2"
self.store.rename_registered_model(name=name, new_name=new_name)
self._verify_requests(
mock_http, "registered-models/rename", "POST",
RenameRegisteredModel(name=name, new_name=new_name))
def test_rename_registered_model(mock_get_request_message, mock_model_registry_store):
name = "model_1"
new_name = "model_2"
mock_get_request_message.return_value = RenameRegisteredModel(name=name, new_name=new_name)
rm2 = RegisteredModel(new_name)
mock_model_registry_store.rename_registered_model.return_value = rm2
resp = _rename_registered_model()
_, args = mock_model_registry_store.rename_registered_model.call_args
assert args == {"name": name, "new_name": new_name}
assert json.loads(resp.get_data()) == {"registered_model": jsonify(rm2)}
def rename_registered_model(self, name, new_name):
"""
Rename the registered model.
:param name: Registered model name.
:param new_name: New proposed name.
:return: A single updated :py:class:`mlflow.entities.model_registry.RegisteredModel` object.
"""
req_body = message_to_json(
RenameRegisteredModel(name=name, new_name=new_name))
response_proto = self._call_endpoint(RenameRegisteredModel, req_body)
return RegisteredModel.from_proto(response_proto.registered_model)
