def to_proto(self):
# input: ModelVersion entity
# returns mlflow.protos.model_registry_pb2.ModelVersion
model_version = ProtoModelVersion()
model_version.name = self.name
model_version.version = str(self.version)
model_version.creation_timestamp = self.creation_timestamp
if self.last_updated_timestamp is not None:
model_version.last_updated_timestamp = self.last_updated_timestamp
if self.description is not None:
model_version.description = self.description
if self.user_id is not None:
model_version.user_id = self.user_id
if self.current_stage is not None:
model_version.current_stage = self.current_stage
if self.source is not None:
model_version.source = str(self.source)
if self.run_id is not None:
model_version.run_id = str(self.run_id)
if self.run_link is not None:
model_version.run_link = str(self.run_link)
if self.status is not None:
model_version.status = ModelVersionStatus.from_string(self.status)
if self.status_message:
model_version.status_message = self.status_message
model_version.tags.extend(
[ProtoModelVersionTag(key=key, value=value) for key, value in self._tags.items()]
)
return model_version
def wait_until_ready(client, model_name, model_version):
for _ in range(10):
model_version_details = client.get_model_version(name=model_name,
version=model_version)
status = ModelVersionStatus.from_string(model_version_details.status)
print("Model status: %s" % ModelVersionStatus.to_string(status))
if status == ModelVersionStatus.READY:
break
time.sleep(1)
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 wait_until_ready(model_name, model_version):
client = MlflowClient(registry_uri=cmr_uri)
for _ in range(20):
model_version_details = client.get_model_version(
name=model_name,
version=model_version,
)
status = ModelVersionStatus.from_string(model_version_details.status)
print("Model status: %s" % ModelVersionStatus.to_string(status))
if status == ModelVersionStatus.READY:
break
time.sleep(5)
def wait_until_version_is_ready(client, model_name, model_version, sleep_time=1, iterations=100):
""" Due to blob eventual consistency, wait until a newly created version is in READY state """
start = time.time()
for _ in range(iterations):
version = client.get_model_version(model_name, model_version.version)
status = ModelVersionStatus.from_string(version.status)
_show_version(version)
if status == ModelVersionStatus.READY:
break
time.sleep(sleep_time)
end = time.time()
print(f"Waited {round(end-start,2)} seconds")
def wait_until_version_ready(model_name, model_version, sleep_time=1, iterations=100):
start = time.time()
for _ in range(iterations):
version = client.get_model_version(model_name, model_version.version)
status = ModelVersionStatus.from_string(version.status)
dt = time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime(round(time.time())))
print(f"{dt}: Version {version.version} status: {ModelVersionStatus.to_string(status)}")
if status == ModelVersionStatus.READY:
break
time.sleep(sleep_time)
end = time.time()
print(f"Waited {round(end-start,2)} seconds")
def wait_until_ready() -> bool:
for _ in range(60):
model_version_details = client.get_model_version(
name=model_name,
version=model_details.version,
)
status = ModelVersionStatus.from_string(model_version_details.status)
print("Model status: %s" % ModelVersionStatus.to_string(status))
if status == ModelVersionStatus.READY:
return True
time.sleep(5)
return False
def wait_until_version_ready(model_name,
model_version,
sleep_time=1,
iterations=100):
start = time.time()
for _ in range(iterations):
version = client.get_model_version(model_name, model_version.version)
status = ModelVersionStatus.from_string(version.status)
show_version(version)
if status == ModelVersionStatus.READY:
break
time.sleep(sleep_time)
end = time.time()
print(f"Waited {round(end-start,2)} seconds")
def test_creation_and_hydration(self):
name = random_str()
rm = RegisteredModel(name)
t1, t2 = 100, 150
source = "path/to/source"
run_id = uuid.uuid4().hex
mvd = ModelVersionDetailed(rm, 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 = {
"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.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 = mvd.to_proto()
self.assertEqual(proto.model_version.registered_model.name, name)
self.assertEqual(proto.model_version.version, 5)
self.assertEqual(proto.status, ModelVersionStatus.from_string("READY"))
self.assertEqual(proto.status_message,
"Model version #5 is ready to use.")
mvd_2 = ModelVersionDetailed.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 = ModelVersionDetailed.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 to_proto(self):
# input: ModelVersionDetailed entity
# returns mlflow.protos.model_registry_pb2.ModelVersionDetailed
model_version_detailed = ProtoModelVersionDetailed()
model_version_detailed.model_version.MergeFrom(
super(ModelVersionDetailed, self).to_proto())
model_version_detailed.creation_timestamp = self.creation_timestamp
if self.last_updated_timestamp:
model_version_detailed.last_updated_timestamp = self.last_updated_timestamp
if self.description:
model_version_detailed.description = self.description
if self.user_id:
model_version_detailed.user_id = self.user_id
model_version_detailed.current_stage = self.current_stage
model_version_detailed.source = self.source
model_version_detailed.run_id = self.run_id
model_version_detailed.status = ModelVersionStatus.from_string(
self.status)
if self.status_message:
model_version_detailed.status_message = self.status_message
return model_version_detailed
def to_proto(self):
# input: ModelVersion entity
# returns mlflow.protos.model_registry_pb2.ModelVersion
model_version = ProtoModelVersion()
model_version.name = self.name
model_version.version = str(self.version)
model_version.creation_timestamp = self.creation_timestamp
if self.last_updated_timestamp is not None:
model_version.last_updated_timestamp = self.last_updated_timestamp
if self.description is not None:
model_version.description = self.description
if self.user_id is not None:
model_version.user_id = self.user_id
if self.current_stage is not None:
model_version.current_stage = self.current_stage
if self.source is not None:
model_version.source = str(self.source)
if self.run_id is not None:
model_version.run_id = str(self.run_id)
if self.status is not None:
model_version.status = ModelVersionStatus.from_string(self.status)
if self.status_message:
model_version.status_message = self.status_message
return model_version