def _update_model_version(): request_message = _get_request_message(UpdateModelVersion()) new_description = None if request_message.HasField("description"): new_description = request_message.description model_version = _get_model_registry_store().update_model_version( name=request_message.name, version=request_message.version, description=new_description) return _wrap_response(UpdateModelVersion.Response(model_version=model_version.to_proto()))
def _update_model_version(): request_message = _get_request_message(UpdateModelVersion()) new_stage = None new_description = None if request_message.HasField("stage"): new_stage = request_message.stage if request_message.HasField("description"): new_description = request_message.description _get_model_registry_store().update_model_version( ModelVersion.from_proto(request_message.model_version), new_stage, new_description) return _wrap_response(UpdateModelVersion.Response())
def test_update_model_version_stage(self, mock_http): rm = RegisteredModel("model_1") mv = ModelVersion(rm, 5) self.store.update_model_version(model_version=mv, stage="prod") self._verify_requests( mock_http, "model-versions/update", "PATCH", UpdateModelVersion(model_version=mv.to_proto(), stage="prod"))
def _update_model_version(): request_message = _get_request_message(UpdateModelVersion()) new_description = None if request_message.HasField("description"): new_description = request_message.description model_version = _get_model_registry_store().update_model_version( name=request_message.name, version=request_message.version, description=new_description) if request_message.HasField("stage"): model_version = _get_model_registry_store( ).transition_model_version_stage(name=request_message.name, version=request_message.version, stage=request_message.stage, archive_existing_versions=False) return _wrap_response( UpdateModelVersion.Response(model_version=model_version.to_proto()))
def test_update_model_version_decription(self, mock_http): rm = RegisteredModel("model_1") mv = ModelVersion(rm, 5) self.store.update_model_version(model_version=mv, description="test model version") self._verify_requests( mock_http, "model-versions/update", "PATCH", UpdateModelVersion(model_version=mv.to_proto(), description="test model version"))
def test_update_model_version(mock_get_request_message, mock_model_registry_store): rm = RegisteredModel("model1") mv = ModelVersion(registered_model=rm, version=32) mock_get_request_message.return_value = UpdateModelVersion(model_version=mv.to_proto(), stage="Production", description="Great model!") _update_model_version() args, _ = mock_model_registry_store.update_model_version.call_args assert args == (mv, "Production", "Great model!")
def test_update_model_version_decription(self, mock_http): name = "model_1" version = "5" description = "test model version" self.store.update_model_version(name=name, version=version, description=description) self._verify_requests( mock_http, "model-versions/update", "PATCH", UpdateModelVersion(name=name, version=version, description="test model version"), )
def test_update_model_version(mock_get_request_message, mock_model_registry_store): name = "model1" version = "32" description = "Great model!" mock_get_request_message.return_value = UpdateModelVersion(name=name, version=version, description=description) mv = ModelVersion(name=name, version=version, creation_timestamp=123, description=description) mock_model_registry_store.update_model_version.return_value = mv _update_model_version() _, args = mock_model_registry_store.update_model_version.call_args assert args == {"name": name, "version": version, "description": description}
def update_model_version(self, name, version, description): """ Update metadata associated with a model version in backend. :param name: Registered model name. :param version: Registered model version. :param description: New model description. :return: A single :py:class:`mlflow.entities.model_registry.ModelVersion` object. """ req_body = message_to_json(UpdateModelVersion(name=name, version=str(version), description=description)) response_proto = self._call_endpoint(UpdateModelVersion, req_body) return ModelVersion.from_proto(response_proto.model_version)
def update_model_version(self, model_version, stage=None, description=None): """ Update metadata associated with a model version in backend. :param model_version: :py:class:`mlflow.entities.model_registry.ModelVersion` object. :param stage: New desired stage for this model version. :param description: New description. :return: None. """ req_body = message_to_json(UpdateModelVersion(model_version=model_version.to_proto(), stage=stage, description=description)) self._call_endpoint(UpdateModelVersion, req_body)
def update_model_version(self, name, version, description): """ Update metadata associated with a model version in backend. :param name: Registered model name. :param version: Registered model version. :param description: New description. :return: None. """ req_body = message_to_json( UpdateModelVersion(name=name, version=str(version), description=description)) response_proto = self._call_endpoint(UpdateModelVersion, req_body) return ModelVersion.from_proto(response_proto.model_version)