def _create_registered_model(): request_message = _get_request_message(CreateRegisteredModel()) registered_model = _get_model_registry_store().create_registered_model( name=request_message.name, tags=request_message.tags) response_message = CreateRegisteredModel.Response( registered_model=registered_model.to_proto()) return _wrap_response(response_message)
def test_create_registered_model(mock_get_request_message, mock_model_registry_store): mock_get_request_message.return_value = CreateRegisteredModel(name="model_1") rm = RegisteredModel("model_1") mock_model_registry_store.create_registered_model.return_value = rm resp = _create_registered_model() _, args = mock_model_registry_store.create_registered_model.call_args assert args == {"name": "model_1"} assert json.loads(resp.get_data()) == {"registered_model": jsonify(rm)}
def test_create_registered_model(self, mock_http): tags = [ RegisteredModelTag(key="key", value="value"), RegisteredModelTag(key="anotherKey", value="some other value") ] self.store.create_registered_model("model_1", tags) self._verify_requests( mock_http, "registered-models/create", "POST", CreateRegisteredModel(name="model_1", tags=[tag.to_proto() for tag in tags]))
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 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 test_create_registered_model(mock_get_request_message, mock_model_registry_store): tags = [ RegisteredModelTag(key="key", value="value"), RegisteredModelTag(key="anotherKey", value="some other value"), ] mock_get_request_message.return_value = CreateRegisteredModel( name="model_1", tags=[tag.to_proto() for tag in tags] ) rm = RegisteredModel("model_1", tags=tags) mock_model_registry_store.create_registered_model.return_value = rm resp = _create_registered_model() _, args = mock_model_registry_store.create_registered_model.call_args assert args["name"] == "model_1" assert {tag.key: tag.value for tag in args["tags"]} == {tag.key: tag.value for tag in tags} assert json.loads(resp.get_data()) == {"registered_model": jsonify(rm)}
def test_create_registered_model(self, mock_http): self.store.create_registered_model("model_1") self._verify_requests(mock_http, "registered-models/create", "POST", CreateRegisteredModel(name="model_1"))