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())