def _transition_stage():
request_message = _get_request_message(TransitionModelVersionStage())
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=request_message.archive_existing_versions)
return _wrap_response(TransitionModelVersionStage.Response(
model_version=model_version.to_proto()))
def test_transition_model_version_stage(mock_get_request_message, mock_model_registry_store):
name = "model1"
version = "32"
stage = "Production"
mock_get_request_message.return_value = TransitionModelVersionStage(name=name, version=version,
stage=stage)
mv = ModelVersion(name=name, version=version, creation_timestamp=123, current_stage=stage)
mock_model_registry_store.transition_model_version_stage.return_value = mv
_transition_stage()
_, args = mock_model_registry_store.transition_model_version_stage.call_args
assert args == {"name": name, "version": version, "stage": stage,
"archive_existing_versions": False}
def test_transition_model_version_stage(self, mock_http):
name = "model_1"
version = "5"
self.store.transition_model_version_stage(
name=name,
version=version,
stage="prod",
archive_existing_versions=True)
self._verify_requests(
mock_http, "model-versions/transition-stage", "POST",
TransitionModelVersionStage(name=name,
version=version,
stage="prod",
archive_existing_versions=True))
def transition_model_version_stage(self, name, version, stage,
archive_existing_versions):
"""
Update model version stage.
:param name: Registered model name.
:param version: Registered model version.
:param new_stage: New desired stage for this model version.
:param archive_existing_versions: If this flag is set to ``True``, all existing model
versions in the stage will be automically moved to the "archived" stage. Only valid
when ``stage`` is ``"staging"`` or ``"production"`` otherwise an error will be raised.
:return: A single :py:class:`mlflow.entities.model_registry.ModelVersion` object.
"""
req_body = message_to_json(TransitionModelVersionStage(
name=name, version=str(version),
stage=stage,
archive_existing_versions=archive_existing_versions))
response_proto = self._call_endpoint(TransitionModelVersionStage, req_body)
return ModelVersion.from_proto(response_proto.model_version)