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 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_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 __init__(
self,
name,
version,
creation_timestamp,
last_updated_timestamp=None,
description=None,
user_id=None,
current_stage=None,
source=None,
run_id=None,
status=ModelVersionStatus.to_string(ModelVersionStatus.READY),
status_message=None,
tags=None,
run_link=None,
):
super().__init__()
self._name = name
self._version = version
self._creation_time = creation_timestamp
self._last_updated_timestamp = last_updated_timestamp
self._description = description
self._user_id = user_id
self._current_stage = current_stage
self._source = source
self._run_id = run_id
self._run_link = run_link
self._status = status
self._status_message = status_message
self._tags = {tag.key: tag.value for tag in (tags or [])}
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 from_proto(cls, proto):
# input: mlflow.protos.model_registry_pb2.ModelVersion
# returns: ModelVersion entity
return cls(proto.name, proto.version, proto.creation_timestamp,
proto.last_updated_timestamp, proto.description,
proto.user_id, proto.current_stage, proto.source,
proto.run_id, ModelVersionStatus.to_string(proto.status),
proto.status_message)
class SqlModelVersion(Base):
__tablename__ = "model_versions"
name = Column(String(256),
ForeignKey("registered_models.name", onupdate="cascade"))
version = Column(Integer, nullable=False)
creation_time = Column(BigInteger, default=lambda: int(time.time() * 1000))
last_updated_time = Column(BigInteger, nullable=True, default=None)
description = Column(String(5000), nullable=True)
user_id = Column(String(256), nullable=True, default=None)
current_stage = Column(String(20), default=STAGE_NONE)
source = Column(String(500), nullable=True, default=None)
run_id = Column(String(32), nullable=False)
run_link = Column(String(500), nullable=True, default=None)
status = Column(String(20),
default=ModelVersionStatus.to_string(
ModelVersionStatus.READY))
status_message = Column(String(500), nullable=True, default=None)
# linked entities
registered_model = relationship("SqlRegisteredModel",
backref=backref("model_versions",
cascade="all"))
__table_args__ = (PrimaryKeyConstraint("name",
"version",
name="model_version_pk"), )
# entity mappers
def to_mlflow_entity(self):
return ModelVersion(
self.name,
self.version,
self.creation_time,
self.last_updated_time,
self.description,
self.user_id,
self.current_stage,
self.source,
self.run_id,
self.status,
self.status_message,
[tag.to_mlflow_entity() for tag in self.model_version_tags],
self.run_link,
)
def create_model_version(
self,
name,
source,
run_id,
tags=None,
run_link=None,
description=None,
await_creation_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
):
"""
Create a new model version from given source.
:param name: Name of the 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.
:param tags: A dictionary of key-value pairs that are converted into
:py:class:`mlflow.entities.model_registry.ModelVersionTag` objects.
:param run_link: Link to the run from an MLflow tracking server that generated this model.
:param description: Description of the version.
:param await_creation_for: Number of seconds to wait for the model version to finish being
created and is in ``READY`` status. By default, the function
waits for five minutes. Specify 0 or None to skip waiting.
Wait until the model version is finished being created and is in ``READY`` status.
:return: Single :py:class:`mlflow.entities.model_registry.ModelVersion` object created by
backend.
"""
tags = tags if tags else {}
tags = [
ModelVersionTag(key, str(value)) for key, value in tags.items()
]
mv = self.store.create_model_version(name, source, run_id, tags,
run_link, description)
if await_creation_for and await_creation_for > 0:
_logger.info(
"Waiting up to %d seconds for model version to finish creation. \
Model name: %s, version %s",
await_creation_for,
name,
mv.version,
)
max_datetime = datetime.utcnow() + timedelta(
seconds=await_creation_for)
pending_status = ModelVersionStatus.to_string(
ModelVersionStatus.PENDING_REGISTRATION)
while mv.status == pending_status:
if datetime.utcnow() > max_datetime:
raise MlflowException(
"Exceeded max wait time for model name: {} version: {} to become READY. \
Status: {} Wait Time: {}".format(
mv.name, mv.version, mv.status,
await_creation_for))
mv = self.get_model_version(mv.name, mv.version)
sleep(AWAIT_MODEL_VERSION_CREATE_SLEEP_DURATION_SECONDS)
return mv
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_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_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_create_model_version_non_ready_model(mock_registry_store):
run_id = "runid"
client = MlflowClient(tracking_uri="http://10.123.1231.11")
mock_registry_store.create_model_version.return_value = ModelVersion(
"name",
1,
0,
1,
source="source",
run_id=run_id,
status=ModelVersionStatus.to_string(ModelVersionStatus.FAILED_REGISTRATION),
)
with pytest.raises(MlflowException, match="Model version creation failed for model name"):
client.create_model_version("name", "source")
class SqlModelVersion(Base):
__tablename__ = 'model_versions'
name = Column(String(256),
ForeignKey('registered_models.name', onupdate='cascade'))
version = Column(Integer, nullable=False)
creation_time = Column(BigInteger, default=lambda: int(time.time() * 1000))
last_updated_time = Column(BigInteger, nullable=True, default=None)
description = Column(String(5000), nullable=True)
user_id = Column(String(256), nullable=True, default=None)
current_stage = Column(String(20), default=STAGE_NONE)
source = Column(String(500), nullable=True, default=None)
run_id = Column(String(32), nullable=False)
status = Column(String(20),
default=ModelVersionStatus.to_string(
ModelVersionStatus.READY))
status_message = Column(String(500), nullable=True, default=None)
# linked entities
registered_model = relationship('SqlRegisteredModel',
backref=backref('model_versions',
cascade='all'))
__table_args__ = (PrimaryKeyConstraint('name',
'version',
name='model_version_pk'), )
# entity mappers
def to_mlflow_entity(self):
return ModelVersion(self.registered_model.to_mlflow_entity(),
self.version)
def to_mlflow_detailed_entity(self):
return ModelVersionDetailed(self.registered_model.to_mlflow_entity(),
self.version, self.creation_time,
self.last_updated_time, self.description,
self.user_id, self.current_stage,
self.source, self.run_id, self.status,
self.status_message)
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 from_proto(cls, proto):
# input: mlflow.protos.model_registry_pb2.ModelVersion
# returns: ModelVersion entity
model_version = cls(proto.name,
proto.version,
proto.creation_timestamp,
proto.last_updated_timestamp,
proto.description,
proto.user_id,
proto.current_stage,
proto.source,
proto.run_id,
ModelVersionStatus.to_string(proto.status),
proto.status_message,
run_link=proto.run_link)
for tag in proto.tags:
model_version._add_tag(ModelVersionTag.from_proto(tag))
return model_version
def upgrade():
bind = op.get_bind()
session = orm.Session(bind=bind)
_logger.info(
"Adding registered_models and model_versions tables to database.")
op.create_table(
SqlRegisteredModel.__tablename__,
Column("name", String(256), unique=True, nullable=False),
Column("creation_time",
BigInteger,
default=lambda: int(time.time() * 1000)),
Column("last_updated_time", BigInteger, nullable=True, default=None),
Column("description", String(5000), nullable=True),
PrimaryKeyConstraint("name", name="registered_model_pk"),
)
op.create_table(
SqlModelVersion.__tablename__,
Column("name", String(256),
ForeignKey("registered_models.name", onupdate="cascade")),
Column("version", Integer, nullable=False),
Column("creation_time",
BigInteger,
default=lambda: int(time.time() * 1000)),
Column("last_updated_time", BigInteger, nullable=True, default=None),
Column("description", String(5000), nullable=True),
Column("user_id", String(256), nullable=True, default=None),
Column("current_stage", String(20), default=STAGE_NONE),
Column("source", String(500), nullable=True, default=None),
Column("run_id", String(32), nullable=False),
Column("status",
String(20),
default=ModelVersionStatus.to_string(ModelVersionStatus.READY)),
Column("status_message", String(500), nullable=True, default=None),
PrimaryKeyConstraint("name", "version", name="model_version_pk"),
)
session.commit()
_logger.info("Migration complete!")
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 upgrade():
bind = op.get_bind()
session = orm.Session(bind=bind)
_logger.info(
"Adding registered_models and model_versions tables to database.")
op.create_table(
SqlRegisteredModel.__tablename__,
Column('name', String(256), unique=True, nullable=False),
Column('creation_time',
BigInteger,
default=lambda: int(time.time() * 1000)),
Column('last_updated_time', BigInteger, nullable=True, default=None),
Column('description', String(5000), nullable=True),
PrimaryKeyConstraint('name', name='registered_model_pk'))
op.create_table(
SqlModelVersion.__tablename__,
Column('name', String(256),
ForeignKey('registered_models.name', onupdate='cascade')),
Column('version', Integer, nullable=False),
Column('creation_time',
BigInteger,
default=lambda: int(time.time() * 1000)),
Column('last_updated_time', BigInteger, nullable=True, default=None),
Column('description', String(5000), nullable=True),
Column('user_id', String(256), nullable=True, default=None),
Column('current_stage', String(20), default=STAGE_NONE),
Column('source', String(500), nullable=True, default=None),
Column('run_id', String(32), nullable=False),
Column('status',
String(20),
default=ModelVersionStatus.to_string(ModelVersionStatus.READY)),
Column('status_message', String(500), nullable=True, default=None),
PrimaryKeyConstraint('name', 'version', name='model_version_pk'))
session.commit()
_logger.info("Migration complete!")
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
