def _get_latest_versions():
request_message = _get_request_message(GetLatestVersions())
latest_versions = _get_model_registry_store().get_latest_versions(
RegisteredModel.from_proto(request_message.registered_model), request_message.stages)
response_message = GetLatestVersions.Response()
response_message.model_versions_detailed.extend([e.to_proto() for e in latest_versions])
return _wrap_response(response_message)
def search_registered_models(self,
filter_string=None,
max_results=None,
order_by=None,
page_token=None):
"""
Search for registered models in backend that satisfy the filter criteria.
:param filter_string: Filter query string, defaults to searching all registered models.
:param max_results: Maximum number of registered models desired.
:param order_by: List of column names with ASC|DESC annotation, to be used for ordering
matching search results.
:param page_token: Token specifying the next page of results. It should be obtained from
a ``search_registered_models`` call.
:return: A PagedList of :py:class:`mlflow.entities.model_registry.RegisteredModel` objects
that satisfy the search expressions. The pagination token for the next page can be
obtained via the ``token`` attribute of the object.
"""
req_body = message_to_json(
SearchRegisteredModels(filter=filter_string,
max_results=max_results,
order_by=order_by,
page_token=page_token))
response_proto = self._call_endpoint(SearchRegisteredModels, req_body)
registered_models = [
RegisteredModel.from_proto(registered_model)
for registered_model in response_proto.registered_models
]
return PagedList(registered_models, response_proto.next_page_token)
def _get_registered_model_details():
request_message = _get_request_message(GetRegisteredModelDetails())
registered_model_detailed = _get_model_registry_store().get_registered_model_details(
RegisteredModel.from_proto(request_message.registered_model))
response_message = GetRegisteredModelDetails.Response(
registered_model_detailed=registered_model_detailed.to_proto())
return _wrap_response(response_message)
def get_registered_model(self, name):
"""
:param name: Registered model name.
:return: A single :py:class:`mlflow.entities.model_registry.RegisteredModel` object.
"""
req_body = message_to_json(GetRegisteredModel(name=name))
response_proto = self._call_endpoint(GetRegisteredModel, req_body)
return RegisteredModel.from_proto(response_proto.registered_model)
def from_proto(cls, proto):
# input: mlflow.protos.model_registry_pb2.ModelVersionDetailed
# returns: ModelVersionDetailed entity
return cls(
RegisteredModel.from_proto(proto.model_version.registered_model),
proto.model_version.version, proto.creation_timestamp,
proto.last_updated_timestamp, proto.description, proto.user_id,
proto.current_stage, proto.source, proto.run_id,
ModelVersionStatus.to_string(proto.status), proto.status_message)
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)
def update_registered_model(self, name, description):
"""
Update description of the registered model.
:param name: Registered model name.
:param description: New description.
:return: A single updated :py:class:`mlflow.entities.model_registry.RegisteredModel` object.
"""
req_body = message_to_json(
UpdateRegisteredModel(name=name, description=description))
response_proto = self._call_endpoint(UpdateRegisteredModel, req_body)
return RegisteredModel.from_proto(response_proto.registered_model)
def list_registered_models(self):
"""
List of all registered models.
:return: List of :py:class:`mlflow.entities.model_registry.RegisteredModel` objects.
"""
req_body = message_to_json(ListRegisteredModels())
response_proto = self._call_endpoint(ListRegisteredModels, req_body)
return [
RegisteredModel.from_proto(registered_model)
for registered_model in response_proto.registered_models
]
def create_registered_model(self, name):
"""
Create a new registered model in backend store.
:param name: Name of the new model. This is expected to be unique in the backend store.
:return: A single object of :py:class:`mlflow.entities.model_registry.RegisteredModel`
created in the backend.
"""
req_body = message_to_json(CreateRegisteredModel(name=name))
response_proto = self._call_endpoint(CreateRegisteredModel, req_body)
return RegisteredModel.from_proto(response_proto.registered_model)
def _update_registered_model():
request_message = _get_request_message(UpdateRegisteredModel())
new_name = None
new_description = None
if request_message.HasField("name"):
new_name = request_message.name
if request_message.HasField("description"):
new_description = request_message.description
registered_model = _get_model_registry_store().update_registered_model(
RegisteredModel.from_proto(request_message.registered_model), new_name,
new_description)
response_message = UpdateRegisteredModel.Response(
registered_model=registered_model.to_proto())
return _wrap_response(response_message)
def create_registered_model(self, name, tags=None):
"""
Create a new registered model in backend store.
:param name: Name of the new model. This is expected to be unique in the backend store.
:param tags: A list of :py:class:`mlflow.entities.model_registry.RegisteredModelTag`
instances associated with this registered model.
:return: A single object of :py:class:`mlflow.entities.model_registry.RegisteredModel`
created in the backend.
"""
proto_tags = [tag.to_proto() for tag in tags or []]
req_body = message_to_json(CreateRegisteredModel(name=name, tags=proto_tags))
response_proto = self._call_endpoint(CreateRegisteredModel, req_body)
return RegisteredModel.from_proto(response_proto.registered_model)
def update_registered_model(self, registered_model, new_name=None, description=None):
"""
Updates metadata for RegisteredModel entity. Either ``new_name`` or ``description`` should
be non-None. Backend raises exception if a registered model with given name does not exist.
:param registered_model: :py:class:`mlflow.entities.model_registry.RegisteredModel` object.
:param new_name: (Optional) New proposed name for the registered model.
:param description: (Optional) New description.
:return: A single updated :py:class:`mlflow.entities.model_registry.RegisteredModel` object.
"""
req_body = message_to_json(UpdateRegisteredModel(
registered_model=registered_model.to_proto(), name=new_name, description=description))
response_proto = self._call_endpoint(UpdateRegisteredModel, req_body)
return RegisteredModel.from_proto(response_proto.registered_model)
def list_registered_models(self, max_results, page_token):
"""
List of all registered models.
:param max_results: Maximum number of registered models desired.
:param page_token: Token specifying the next page of results. It should be obtained from
a ``list_registered_models`` call.
:return: PagedList of :py:class:`mlflow.entities.model_registry.RegisteredModel` objects.
"""
req_body = message_to_json(
ListRegisteredModels(page_token=page_token,
max_results=max_results))
response_proto = self._call_endpoint(ListRegisteredModels, req_body)
return PagedList([
RegisteredModel.from_proto(registered_model)
for registered_model in response_proto.registered_models
], response_proto.next_page_token)
def from_proto(cls, proto):
# input: mlflow.protos.model_registry_pb2.ModelVersion
# returns: ModelVersion entity
return cls(RegisteredModel.from_proto(proto.registered_model),
proto.version)
def _delete_registered_model():
request_message = _get_request_message(DeleteRegisteredModel())
_get_model_registry_store().delete_registered_model(
RegisteredModel.from_proto(request_message.registered_model))
return _wrap_response(DeleteRegisteredModel.Response())
