Esempio n. 1
0
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)
Esempio n. 2
0
    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)
Esempio n. 3
0
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)}
Esempio n. 4
0
 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))