def test_with_latest_model_versions(self):
name = random_str()
mvd_1 = ModelVersion(
name,
"1",
1000,
2000,
"version 1",
"user 1",
"Production",
"source 1",
"run ID 1",
"PENDING_REGISTRATION",
"Model version is in production!",
)
mvd_2 = ModelVersion(
name,
"4",
1300,
2002,
"version 4",
"user 2",
"Staging",
"source 4",
"run ID 12",
"READY",
"Model copied over!",
)
as_dict = {
"name": name,
"creation_timestamp": 1,
"last_updated_timestamp": 4000,
"description": random_str(),
"latest_versions": [mvd_1, mvd_2],
"tags": [],
}
rmd_1 = RegisteredModel.from_dictionary(as_dict)
as_dict["tags"] = {}
self.assertEqual(dict(rmd_1), as_dict)
proto = rmd_1.to_proto()
self.assertEqual(proto.creation_timestamp, 1)
self.assertEqual(proto.last_updated_timestamp, 4000)
self.assertEqual(set([mvd.version for mvd in proto.latest_versions]), set(["1", "4"]))
self.assertEqual(set([mvd.name for mvd in proto.latest_versions]), set([name]))
self.assertEqual(
set([mvd.current_stage for mvd in proto.latest_versions]),
set(["Production", "Staging"]),
)
self.assertEqual(
set([mvd.last_updated_timestamp for mvd in proto.latest_versions]), set([2000, 2002])
)
self.assertEqual(
set([mvd.creation_timestamp for mvd in proto.latest_versions]), set([1300, 1000])
)
def test_creation_and_hydration(self):
name = random_str()
t1, t2 = 100, 150
source = "path/to/source"
run_id = uuid.uuid4().hex
run_link = "http://localhost:5000/path/to/run"
tags = [ModelVersionTag("key", "value"), ModelVersionTag("randomKey", "not a random value")]
mvd = ModelVersion(name, "5", t1, t2, "version five", "user 1", "Production",
source, run_id, "READY", "Model version #5 is ready to use.", tags,
run_link)
self._check(mvd, name, "5", t1, t2, "version five", "user 1",
"Production", source, run_id, "READY",
"Model version #5 is ready to use.",
{tag.key: tag.value for tag in (tags or [])})
expected_dict = {
"name": name,
"version": "5",
"creation_timestamp": t1,
"last_updated_timestamp": t2,
"description": "version five",
"user_id": "user 1",
"current_stage": "Production",
"source": source,
"run_id": run_id,
"run_link": run_link,
"status": "READY",
"status_message": "Model version #5 is ready to use.",
"tags": {tag.key: tag.value for tag in (tags or [])}}
model_version_as_dict = dict(mvd)
self.assertEqual(model_version_as_dict, expected_dict)
proto = mvd.to_proto()
self.assertEqual(proto.name, name)
self.assertEqual(proto.version, "5")
self.assertEqual(proto.status, ModelVersionStatus.from_string("READY"))
self.assertEqual(proto.status_message, "Model version #5 is ready to use.")
self.assertEqual(set([tag.key for tag in proto.tags]),
set(["key", "randomKey"]))
self.assertEqual(set([tag.value for tag in proto.tags]),
set(["value", "not a random value"]))
mvd_2 = ModelVersion.from_proto(proto)
self._check(mvd_2, name, "5", t1, t2, "version five", "user 1",
"Production", source, run_id, "READY",
"Model version #5 is ready to use.",
{tag.key: tag.value for tag in (tags or [])})
expected_dict.update({"registered_model": RegisteredModel(name)})
expected_dict["tags"] = tags
mvd_3 = ModelVersion.from_dictionary(expected_dict)
self._check(mvd_3, name, "5", t1, t2, "version five", "user 1",
"Production", source, run_id, "READY",
"Model version #5 is ready to use.",
{tag.key: tag.value for tag in (tags or [])})
def test_string_repr(self):
model_version = ModelVersion(
name="myname",
version="43",
creation_timestamp=12,
last_updated_timestamp=100,
description="This is a test model.",
user_id="user one",
current_stage="Archived",
source="path/to/a/notebook",
run_id="some run",
run_link="http://localhost:5000/path/to/run",
status="PENDING_REGISTRATION",
status_message="Copying!",
tags=[],
)
assert (
str(model_version) == "<ModelVersion: creation_timestamp=12, "
"current_stage='Archived', description='This is a test "
"model.', last_updated_timestamp=100, "
"name='myname', "
"run_id='some run', run_link='http://localhost:5000/path/"
"to/run', source='path/to/a/notebook', "
"status='PENDING_REGISTRATION', status_message='Copying!', "
"tags={}, user_id='user one', version='43'>"
)
def test_creation_and_hydration(self):
name = random_str()
t1, t2 = 100, 150
source = "path/to/source"
run_id = uuid.uuid4().hex
mvd = ModelVersion(name, "5", t1, t2, "version five", "user 1",
"Production", source, run_id, "READY",
"Model version #5 is ready to use.")
self._check(mvd, name, "5", t1, t2, "version five", "user 1",
"Production", source, run_id, "READY",
"Model version #5 is ready to use.")
expected_dict = {
"name": name,
"version": "5",
"creation_timestamp": t1,
"last_updated_timestamp": t2,
"description": "version five",
"user_id": "user 1",
"current_stage": "Production",
"source": source,
"run_id": run_id,
"status": "READY",
"status_message": "Model version #5 is ready to use."
}
model_version_as_dict = dict(mvd)
self.assertEqual(model_version_as_dict, expected_dict)
proto = mvd.to_proto()
self.assertEqual(proto.name, name)
self.assertEqual(proto.version, "5")
self.assertEqual(proto.status, ModelVersionStatus.from_string("READY"))
self.assertEqual(proto.status_message,
"Model version #5 is ready to use.")
mvd_2 = ModelVersion.from_proto(proto)
self._check(mvd_2, name, "5", t1, t2, "version five", "user 1",
"Production", source, run_id, "READY",
"Model version #5 is ready to use.")
expected_dict.update({"registered_model": RegisteredModel(name)})
mvd_3 = ModelVersion.from_dictionary(expected_dict)
self._check(mvd_3, name, "5", t1, t2, "version five", "user 1",
"Production", source, run_id, "READY",
"Model version #5 is ready to use.")
def from_proto(cls, proto):
# input: mlflow.protos.model_registry_pb2.RegisteredModel
# returns RegisteredModel entity
registered_model = cls(
proto.name, proto.creation_timestamp, proto.last_updated_timestamp,
proto.description,
[ModelVersion.from_proto(mvd) for mvd in proto.latest_versions])
for tag in proto.tags:
registered_model._add_tag(RegisteredModelTag.from_proto(tag))
return registered_model
def test_creation_and_hydration(self):
name = random_str()
rm = RegisteredModel(name)
model_version = ModelVersion(rm, 100)
self._check(model_version, name, 100)
expected_dict = {"version": 100}
model_version_as_dict = dict(model_version)
self.assertEqual(model_version_as_dict["version"], 100)
self.assertIsInstance(model_version_as_dict["registered_model"],
RegisteredModel)
self.assertEqual(model_version_as_dict["registered_model"].name, name)
model_version_as_dict.pop("registered_model")
self.assertEqual(model_version_as_dict, expected_dict)
proto = model_version.to_proto()
self.assertEqual(proto.registered_model.name, name)
self.assertEqual(proto.version, 100)
model_version_2 = ModelVersion.from_proto(proto)
self._check(model_version_2, name, 100)
expected_dict.update({"registered_model": RegisteredModel(name)})
model_version_3 = ModelVersion.from_dictionary(expected_dict)
self._check(model_version_3, name, 100)
def test_string_repr(self):
model_version = ModelVersion(name="myname",
version="43",
creation_timestamp=12,
last_updated_timestamp=100,
description="This is a test model.",
user_id="user one",
current_stage="Archived",
source="path/to/a/notebook",
run_id="some run",
status="PENDING_REGISTRATION",
status_message="Copying!")
assert str(model_version) == "<ModelVersion: creation_timestamp=12, " \
"current_stage='Archived', description='This is a test " \
"model.', last_updated_timestamp=100, " \
"name='myname', " \
"run_id='some run', source='path/to/a/notebook', " \
"status='PENDING_REGISTRATION', status_message='Copying!', " \
"user_id='user one', version='43'>"
def _properties(cls):
# aggregate with base class properties since cls.__dict__ does not do it automatically
return sorted(cls._get_properties_helper() +
ModelVersion._properties())
def test_string_repr(self):
model_version = ModelVersion(RegisteredModel(name="myname"), version=4)
assert str(model_version) == "<ModelVersion: registered_model=<RegisteredModel: " \
"name='myname'>, version=4>"