def get_model_version(self, name, version):
"""
Get the model version instance by name and version.
:param name: Registered model name.
:param version: Registered model version.
:return: A single :py:class:`mlflow.entities.model_registry.ModelVersion` object.
"""
req_body = message_to_json(
GetModelVersion(name=name, version=str(version)))
response_proto = self._call_endpoint(GetModelVersion, req_body)
return ModelVersion.from_proto(response_proto.model_version)
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 get_latest_versions(self, name, stages=None):
"""
Latest version models for each requested stage. If no ``stages`` argument is provided,
returns the latest version for each stage.
:param name: Registered model name.
:param stages: List of desired stages. If input list is None, return latest versions for
for 'Staging' and 'Production' stages.
:return: List of :py:class:`mlflow.entities.model_registry.ModelVersion` objects.
"""
req_body = message_to_json(GetLatestVersions(name=name, stages=stages))
response_proto = self._call_endpoint(GetLatestVersions, req_body)
return [
ModelVersion.from_proto(model_version)
for model_version in response_proto.model_versions
]
def search_model_versions(self, filter_string):
"""
Search for model versions in backend that satisfy the filter criteria.
:param filter_string: A filter string expression. Currently supports a single filter
condition either name of model like ``name = 'model_name'`` or
``run_id = '...'``.
:return: PagedList of :py:class:`mlflow.entities.model_registry.ModelVersion`
objects.
"""
req_body = message_to_json(SearchModelVersions(filter=filter_string))
response_proto = self._call_endpoint(SearchModelVersions, req_body)
model_versions = [
ModelVersion.from_proto(mvd)
for mvd in response_proto.model_versions
]
return PagedList(model_versions, response_proto.next_page_token)
def create_model_version(self, name, source, run_id, tags=None):
"""
Create a new model version from given source and run ID.
:param name: Registered model name.
:param source: Source path where the MLflow model is stored.
:param run_id: Run ID from MLflow tracking server that generated the model.
:param tags: A list of :py:class:`mlflow.entities.model_registry.ModelVersionTag`
instances associated with this model version.
:return: A single object of :py:class:`mlflow.entities.model_registry.ModelVersion`
created in the backend.
"""
proto_tags = [tag.to_proto() for tag in tags or []]
req_body = message_to_json(
CreateModelVersion(name=name,
source=source,
run_id=run_id,
tags=proto_tags))
response_proto = self._call_endpoint(CreateModelVersion, req_body)
return ModelVersion.from_proto(response_proto.model_version)
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)