def _list_registered_models():
_get_request_message(ListRegisteredModels())
registered_models_detailed = _get_model_registry_store().list_registered_models()
response_message = ListRegisteredModels.Response()
response_message.registered_models_detailed.extend([e.to_proto()
for e in registered_models_detailed])
return _wrap_response(response_message)
def _list_registered_models():
request_message = _get_request_message(ListRegisteredModels())
registered_models = _get_model_registry_store().list_registered_models(
request_message.max_results, request_message.page_token)
response_message = ListRegisteredModels.Response()
response_message.registered_models.extend(
[e.to_proto() for e in registered_models])
if registered_models.token:
response_message.next_page_token = registered_models.token
return _wrap_response(response_message)
def test_list_registered_models(mock_get_request_message,
mock_model_registry_store):
mock_get_request_message.return_value = ListRegisteredModels(
max_results=50)
rmds = PagedList(
[
RegisteredModel(
name="model_1",
creation_timestamp=111,
last_updated_timestamp=222,
description="Test model",
latest_versions=[],
),
RegisteredModel(
name="model_2",
creation_timestamp=111,
last_updated_timestamp=333,
description="Another model",
latest_versions=[],
),
],
"next_pt",
)
mock_model_registry_store.list_registered_models.return_value = rmds
resp = _list_registered_models()
args, _ = mock_model_registry_store.list_registered_models.call_args
assert args == (50, "")
assert json.loads(resp.get_data()) == {
"next_page_token": "next_pt",
"registered_models": jsonify(rmds),
}
def test_list_registered_model(self, mock_http):
self.store.list_registered_models(max_results=50, page_token=None)
self._verify_requests(
mock_http,
"registered-models/list",
"GET",
ListRegisteredModels(page_token=None, max_results=50),
)
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 [RegisteredModelDetailed.from_proto(registered_model_detailed)
for registered_model_detailed in response_proto.registered_models_detailed]
def test_list_registered_models(mock_get_request_message, mock_model_registry_store):
mock_get_request_message.return_value = ListRegisteredModels()
rmds = [
RegisteredModel(name="model_1", creation_timestamp=111,
last_updated_timestamp=222, description="Test model",
latest_versions=[]),
RegisteredModel(name="model_2", creation_timestamp=111,
last_updated_timestamp=333, description="Another model",
latest_versions=[]),
]
mock_model_registry_store.list_registered_models.return_value = rmds
resp = _list_registered_models()
args, _ = mock_model_registry_store.list_registered_models.call_args
assert args == ()
assert json.loads(resp.get_data()) == {"registered_models": jsonify(rmds)}
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 test_list_registered_model(self, mock_http):
self.store.list_registered_models()
self._verify_requests(mock_http, "registered-models/list", "GET",
ListRegisteredModels())
