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)