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)