def _get_registered_model():
request_message = _get_request_message(GetRegisteredModel())
registered_model = _get_model_registry_store().get_registered_model(
name=request_message.name)
response_message = GetRegisteredModel.Response(
registered_model=registered_model.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 test_get_registered_model(mock_get_request_message, mock_model_registry_store):
name = "model1"
mock_get_request_message.return_value = GetRegisteredModel(name=name)
rmd = RegisteredModel(name=name, creation_timestamp=111,
last_updated_timestamp=222, description="Test model",
latest_versions=[])
mock_model_registry_store.get_registered_model.return_value = rmd
resp = _get_registered_model()
_, args = mock_model_registry_store.get_registered_model.call_args
assert args == {"name": name}
assert json.loads(resp.get_data()) == {"registered_model": jsonify(rmd)}
def test_get_registered_model(self, mock_http):
name = "model_1"
self.store.get_registered_model(name=name)
self._verify_requests(mock_http, "registered-models/get", "GET",
GetRegisteredModel(name=name))