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()
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_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 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)