def _create_model_version():
request_message = _get_request_message(CreateModelVersion())
model_version = _get_model_registry_store().create_model_version(
request_message.name, request_message.source, request_message.run_id)
response_message = CreateModelVersion.Response(
model_version=model_version.to_proto())
return _wrap_response(response_message)
def test_create_model_version(self, mock_http):
self.store.create_model_version("model_1", "path/to/source")
self._verify_requests(
mock_http,
"model-versions/create",
"POST",
CreateModelVersion(name="model_1", source="path/to/source"),
)
# test optional fields
run_id = uuid.uuid4().hex
tags = [
ModelVersionTag(key="key", value="value"),
ModelVersionTag(key="anotherKey", value="some other value"),
]
run_link = "localhost:5000/path/to/run"
description = "version description"
self.store.create_model_version(
"model_1", "path/to/source", run_id, tags, run_link=run_link, description=description,
)
self._verify_requests(
mock_http,
"model-versions/create",
"POST",
CreateModelVersion(
name="model_1",
source="path/to/source",
run_id=run_id,
run_link=run_link,
tags=[tag.to_proto() for tag in tags],
description=description,
),
)
def _create_model_version():
request_message = _get_request_message(CreateModelVersion())
username = _get_username()
model_version = _get_model_registry_store().create_model_version(
name=request_message.name,
source=request_message.source,
run_id=request_message.run_id,
run_link=request_message.run_link,
tags=request_message.tags,
user_id=username)
response_message = CreateModelVersion.Response(model_version=model_version.to_proto())
return _wrap_response(response_message)
def test_create_model_version(mock_get_request_message,
mock_model_registry_store):
run_id = uuid.uuid4().hex
tags = [
ModelVersionTag(key="key", value="value"),
ModelVersionTag(key="anotherKey", value="some other value")
]
run_link = "localhost:5000/path/to/run"
mock_get_request_message.return_value = CreateModelVersion(
name="model_1",
source="A/B",
run_id=run_id,
run_link=run_link,
tags=[tag.to_proto() for tag in tags])
mv = ModelVersion(name="model_1",
version="12",
creation_timestamp=123,
tags=tags,
run_link=run_link)
mock_model_registry_store.create_model_version.return_value = mv
resp = _create_model_version()
_, args = mock_model_registry_store.create_model_version.call_args
assert args["name"] == "model_1"
assert args["source"] == "A/B"
assert args["run_id"] == run_id
assert {tag.key: tag.value
for tag in args["tags"]} == {tag.key: tag.value
for tag in tags}
assert args["run_link"] == run_link
assert json.loads(resp.get_data()) == {"model_version": jsonify(mv)}
def create_model_version(self,
name,
source,
run_id,
tags=None,
run_link=None,
description=None):
"""
Create a new model version from given source and run ID.
:param name: Registered model name.
:param source: Source path where the MLflow model is stored.
:param run_id: Run ID from MLflow tracking server that generated the model.
:param tags: A list of :py:class:`mlflow.entities.model_registry.ModelVersionTag`
instances associated with this model version.
:param run_link: Link to the run from an MLflow tracking server that generated this model.
:param description: Description of the version.
:return: A single object of :py:class:`mlflow.entities.model_registry.ModelVersion`
created in the backend.
"""
proto_tags = [tag.to_proto() for tag in tags or []]
req_body = message_to_json(
CreateModelVersion(
name=name,
source=source,
run_id=run_id,
run_link=run_link,
tags=proto_tags,
description=description,
))
response_proto = self._call_endpoint(CreateModelVersion, req_body)
return ModelVersion.from_proto(response_proto.model_version)
def test_create_model_version(self, mock_http):
run_id = uuid.uuid4().hex
self.store.create_model_version("model_1", "path/to/source", run_id)
self._verify_requests(
mock_http, "model-versions/create", "POST",
CreateModelVersion(name="model_1",
source="path/to/source",
run_id=run_id))
def test_create_model_version(mock_get_request_message, mock_model_registry_store):
run_id = uuid.uuid4().hex
mock_get_request_message.return_value = CreateModelVersion(name="model_1", source="A/B",
run_id=run_id)
mv = ModelVersion(name="model_1", version="12", creation_timestamp=123)
mock_model_registry_store.create_model_version.return_value = mv
resp = _create_model_version()
_, args = mock_model_registry_store.create_model_version.call_args
assert args == {"name": "model_1", "source": "A/B", "run_id": run_id}
assert json.loads(resp.get_data()) == {"model_version": jsonify(mv)}
def test_create_model_version(mock_get_request_message, mock_model_registry_store):
run_id = uuid.uuid4().hex
mock_get_request_message.return_value = CreateModelVersion(name="model_1", source="A/B",
run_id=run_id)
mv = ModelVersion(registered_model=RegisteredModel(name="model_1"), version=12)
mock_model_registry_store.create_model_version.return_value = mv
resp = _create_model_version()
args, _ = mock_model_registry_store.create_model_version.call_args
assert args == ("model_1", "A/B", run_id)
assert json.loads(resp.get_data()) == {"model_version": jsonify(mv)}
def create_model_version(self, name, source, run_id):
"""
Create a new model version from given source and run ID.
:param name: Name ID for containing registered model.
:param source: Source path where the MLflow model is stored.
:param run_id: Run ID from MLflow tracking server that generated the model
:return: A single object of :py:class:`mlflow.entities.model_registry.ModelVersion`
created in the backend.
"""
req_body = message_to_json(CreateModelVersion(name=name, source=source, run_id=run_id))
response_proto = self._call_endpoint(CreateModelVersion, req_body)
return ModelVersion.from_proto(response_proto.model_version)
def test_create_model_version(self, mock_http):
tags = [
ModelVersionTag(key="key", value="value"),
ModelVersionTag(key="anotherKey", value="some other value")
]
run_id = uuid.uuid4().hex
self.store.create_model_version("model_1", "path/to/source", run_id,
tags)
self._verify_requests(
mock_http, "model-versions/create", "POST",
CreateModelVersion(name="model_1",
source="path/to/source",
run_id=run_id,
tags=[tag.to_proto() for tag in tags]))
