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)
