def get_latest_versions(self, name, stages=None):
"""
Latest version models for each requested stage. If no ``stages`` argument is provided,
returns the latest version for each stage.
:param name: Registered model name.
:param stages: List of desired stages. If input list is None, return latest versions for
for 'Staging' and 'Production' stages.
:return: List of :py:class:`mlflow.entities.model_registry.ModelVersion` objects.
"""
with self.ManagedSessionMaker() as session:
sql_registered_model = self._get_registered_model(session, name)
# Convert to RegisteredModel entity first and then extract latest_versions
latest_versions = sql_registered_model.to_mlflow_entity(
).latest_versions
if stages is None or len(stages) == 0:
expected_stages = set([
get_canonical_stage(stage)
for stage in DEFAULT_STAGES_FOR_GET_LATEST_VERSIONS
])
else:
expected_stages = set(
[get_canonical_stage(stage) for stage in stages])
return [
mv for mv in latest_versions
if mv.current_stage in expected_stages
]
def transition_model_version_stage(self, name, version, stage,
archive_existing_versions):
"""
Update model version stage.
:param name: Registered model name.
:param version: Registered model version.
:param new_stage: New desired stage for this model version.
:param archive_existing_versions: If this flag is set to ``True``, all existing model
versions in the stage will be automically moved to the "archived" stage. Only valid
when ``stage`` is ``"staging"`` or ``"production"`` otherwise an error will be raised.
:return: A single :py:class:`mlflow.entities.model_registry.ModelVersion` object.
"""
is_active_stage = get_canonical_stage(
stage) in DEFAULT_STAGES_FOR_GET_LATEST_VERSIONS
if archive_existing_versions and not is_active_stage:
msg_tpl = (
"Model version transition cannot archive existing model versions "
"because '{}' is not an Active stage. Valid stages are {}")
raise MlflowException(
msg_tpl.format(stage, DEFAULT_STAGES_FOR_GET_LATEST_VERSIONS))
with self.ManagedSessionMaker() as session:
last_updated_time = now()
model_versions = []
if archive_existing_versions:
conditions = [
SqlModelVersion.name == name,
SqlModelVersion.version != version,
SqlModelVersion.current_stage == get_canonical_stage(
stage),
]
model_versions = session.query(SqlModelVersion).filter(
*conditions).all()
for mv in model_versions:
mv.current_stage = STAGE_ARCHIVED
mv.last_updated_time = last_updated_time
sql_model_version = self._get_sql_model_version(session=session,
name=name,
version=version)
sql_model_version.current_stage = get_canonical_stage(stage)
sql_model_version.last_updated_time = last_updated_time
sql_registered_model = sql_model_version.registered_model
sql_registered_model.last_updated_time = last_updated_time
self._save_to_db(
session,
[*model_versions, sql_model_version, sql_registered_model])
return sql_model_version.to_mlflow_entity()
def transition_model_version_stage(self, name, version, stage,
archive_existing_versions):
"""
Update model version stage.
:param name: :py:string: Registered model name.
:param version: :py:string: Registered model version.
:param new_stage: New desired stage for this model version.
:param archive_existing_versions: :py:boolean: If this flag is set, all existing model
versions in the stage will be atomically moved to the "archived" stage.
:return: A single :py:class:`mlflow.entities.model_registry.ModelVersion` object.
"""
if archive_existing_versions:
raise MlflowException("'archive_existing_versions' flag is not supported in "
"SqlAlchemyStore. Set it to 'False'")
with self.ManagedSessionMaker() as session:
sql_model_version = self._get_sql_model_version(session=session,
name=name,
version=version)
last_updated_time = now()
sql_model_version.current_stage = get_canonical_stage(stage)
sql_model_version.last_updated_time = last_updated_time
sql_registered_model = sql_model_version.registered_model
sql_registered_model.last_updated_time = last_updated_time
self._save_to_db(session, [sql_model_version, sql_registered_model])
return sql_model_version.to_mlflow_entity()
def update_model_version(self,
model_version,
stage=None,
description=None):
"""
Update metadata associated with a model version in backend.
:param model_version: :py:class:`mlflow.entities.model_registry.ModelVersion` object.
:param stage: New desired stage for this model version.
:param description: New description.
:return: None.
"""
with self.ManagedSessionMaker() as session:
updated_time = now()
sql_model_version = self._get_sql_model_version(
session, model_version)
updated_objs = []
if stage is not None:
sql_model_version.current_stage = get_canonical_stage(stage)
sql_registered_model = sql_model_version.registered_model
sql_registered_model.last_updated_time = updated_time
updated_objs.append(sql_registered_model)
if description is not None:
sql_model_version.description = description
sql_model_version.last_updated_time = updated_time
self._save_to_db(session, updated_objs + [sql_model_version])
def update_model_version(self,
model_version,
stage=None,
description=None):
"""
Update metadata associated with a model version in backend.
:param model_version: :py:class:`mlflow.entities.model_registry.ModelVersion` object.
:param stage: New desired stage for this model version.
:param description: New description.
:return: None.
"""
with self.ManagedSessionMaker() as session:
sql_model_version = self._get_sql_model_version(
session, model_version)
if stage is not None:
sql_model_version.current_stage = get_canonical_stage(stage)
if description is not None:
sql_model_version.description = description
self._save_to_db(session, sql_model_version)
