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