genesis-flow 1.0.0__py3-none-any.whl

mlflow/store/db/utils.py

@@ -0,0 +1,314 @@
+import logging
+import os
+import time
+from contextlib import contextmanager
+
+import sqlalchemy
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from sqlalchemy import sql
+
+# We need to import sqlalchemy.pool to convert poolclass string to class object
+from sqlalchemy.pool import (
+    AssertionPool,
+    AsyncAdaptedQueuePool,
+    FallbackAsyncAdaptedQueuePool,
+    NullPool,
+    QueuePool,
+    SingletonThreadPool,
+    StaticPool,
+)
+
+from mlflow.environment_variables import (
+    MLFLOW_MYSQL_SSL_CA,
+    MLFLOW_MYSQL_SSL_CERT,
+    MLFLOW_MYSQL_SSL_KEY,
+    MLFLOW_SQLALCHEMYSTORE_ECHO,
+    MLFLOW_SQLALCHEMYSTORE_MAX_OVERFLOW,
+    MLFLOW_SQLALCHEMYSTORE_POOL_RECYCLE,
+    MLFLOW_SQLALCHEMYSTORE_POOL_SIZE,
+    MLFLOW_SQLALCHEMYSTORE_POOLCLASS,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import (
+    BAD_REQUEST,
+    INTERNAL_ERROR,
+    TEMPORARILY_UNAVAILABLE,
+)
+from mlflow.store.db.db_types import SQLITE
+from mlflow.store.model_registry.dbmodels.models import (
+    SqlModelVersion,
+    SqlModelVersionTag,
+    SqlRegisteredModel,
+    SqlRegisteredModelAlias,
+    SqlRegisteredModelTag,
+)
+from mlflow.store.tracking.dbmodels.initial_models import Base as InitialBase
+from mlflow.store.tracking.dbmodels.models import (
+    SqlDataset,
+    SqlExperiment,
+    SqlExperimentTag,
+    SqlInput,
+    SqlInputTag,
+    SqlLatestMetric,
+    SqlMetric,
+    SqlParam,
+    SqlRun,
+    SqlTag,
+    SqlTraceInfo,
+    SqlTraceMetadata,
+    SqlTraceTag,
+)
+
+_logger = logging.getLogger(__name__)
+
+MAX_RETRY_COUNT = 10
+
+
+def _get_package_dir():
+    """Returns directory containing MLflow python package."""
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    return os.path.normpath(os.path.join(current_dir, os.pardir, os.pardir))
+
+
+def _all_tables_exist(engine):
+    return {
+        t
+        for t in sqlalchemy.inspect(engine).get_table_names()
+        # Filter out alembic tables
+        if not t.startswith("alembic_")
+    } == {
+        SqlExperiment.__tablename__,
+        SqlRun.__tablename__,
+        SqlMetric.__tablename__,
+        SqlParam.__tablename__,
+        SqlTag.__tablename__,
+        SqlExperimentTag.__tablename__,
+        SqlLatestMetric.__tablename__,
+        SqlRegisteredModel.__tablename__,
+        SqlModelVersion.__tablename__,
+        SqlRegisteredModelTag.__tablename__,
+        SqlModelVersionTag.__tablename__,
+        SqlRegisteredModelAlias.__tablename__,
+        SqlDataset.__tablename__,
+        SqlInput.__tablename__,
+        SqlInputTag.__tablename__,
+        SqlTraceInfo.__tablename__,
+        SqlTraceTag.__tablename__,
+        SqlTraceMetadata.__tablename__,
+    }
+
+
+def _initialize_tables(engine):
+    _logger.info("Creating initial MLflow database tables...")
+    InitialBase.metadata.create_all(engine)
+    _upgrade_db(engine)
+
+
+def _get_latest_schema_revision():
+    """Get latest schema revision as a string."""
+    # We aren't executing any commands against a DB, so we leave the DB URL unspecified
+    config = _get_alembic_config(db_url="")
+    script = ScriptDirectory.from_config(config)
+    heads = script.get_heads()
+    if len(heads) != 1:
+        raise MlflowException(
+            f"Migration script directory was in unexpected state. Got {len(heads)} head "
+            f"database versions but expected only 1. Found versions: {heads}"
+        )
+    return heads[0]
+
+
+def _verify_schema(engine):
+    head_revision = _get_latest_schema_revision()
+    current_rev = _get_schema_version(engine)
+    if current_rev != head_revision:
+        raise MlflowException(
+            f"Detected out-of-date database schema (found version {current_rev}, "
+            f"but expected {head_revision}). Take a backup of your database, then run "
+            "'mlflow db upgrade <database_uri>' "
+            "to migrate your database to the latest schema. NOTE: schema migration may "
+            "result in database downtime - please consult your database's documentation for "
+            "more detail."
+        )
+
+
+def _get_managed_session_maker(SessionMaker, db_type):
+    """
+    Creates a factory for producing exception-safe SQLAlchemy sessions that are made available
+    using a context manager. Any session produced by this factory is automatically committed
+    if no exceptions are encountered within its associated context. If an exception is
+    encountered, the session is rolled back. Finally, any session produced by this factory is
+    automatically closed when the session's associated context is exited.
+    """
+
+    @contextmanager
+    def make_managed_session():
+        """Provide a transactional scope around a series of operations."""
+        with SessionMaker() as session:
+            try:
+                if db_type == SQLITE:
+                    session.execute(sql.text("PRAGMA foreign_keys = ON;"))
+                    session.execute(sql.text("PRAGMA busy_timeout = 20000;"))
+                    session.execute(sql.text("PRAGMA case_sensitive_like = true;"))
+                yield session
+                session.commit()
+            except MlflowException:
+                session.rollback()
+                raise
+            except sqlalchemy.exc.OperationalError as e:
+                session.rollback()
+                _logger.exception(
+                    "SQLAlchemy database error. The following exception is caught.\n%s",
+                    e,
+                )
+                raise MlflowException(message=e, error_code=TEMPORARILY_UNAVAILABLE) from e
+            except sqlalchemy.exc.SQLAlchemyError as e:
+                session.rollback()
+                raise MlflowException(message=e, error_code=BAD_REQUEST) from e
+            except Exception as e:
+                session.rollback()
+                raise MlflowException(message=e, error_code=INTERNAL_ERROR) from e
+
+    return make_managed_session
+
+
+def _get_alembic_config(db_url, alembic_dir=None):
+    """
+    Constructs an alembic Config object referencing the specified database and migration script
+    directory.
+
+    Args:
+        db_url: Database URL, like sqlite:///<absolute-path-to-local-db-file>. See
+            https://docs.sqlalchemy.org/en/13/core/engines.html#database-urls for a full list of
+            valid database URLs.
+        alembic_dir: Path to migration script directory. Uses canonical migration script
+            directory under mlflow/alembic if unspecified. TODO: remove this argument in MLflow 1.1,
+            as it's only used to run special migrations for pre-1.0 users to remove duplicate
+            constraint names.
+    """
+    from alembic.config import Config
+
+    final_alembic_dir = (
+        os.path.join(_get_package_dir(), "store", "db_migrations")
+        if alembic_dir is None
+        else alembic_dir
+    )
+    # Escape any '%' that appears in a db_url. This could be in a password,
+    # url, or anything that is part of a potentially complex database url
+    db_url = db_url.replace("%", "%%")
+    config = Config(os.path.join(final_alembic_dir, "alembic.ini"))
+    config.set_main_option("script_location", final_alembic_dir)
+    config.set_main_option("sqlalchemy.url", db_url)
+    return config
+
+
+def _upgrade_db(engine):  # noqa: D417
+    """
+    Upgrade the schema of an MLflow tracking database to the latest supported version.
+    Note that schema migrations can be slow and are not guaranteed to be transactional -
+    we recommend taking a backup of your database before running migrations.
+
+    Args:
+        url: Database URL, like sqlite:///<absolute-path-to-local-db-file>. See
+            https://docs.sqlalchemy.org/en/13/core/engines.html#database-urls for a full list of
+            valid database URLs.
+    """
+    # alembic adds significant import time, so we import it lazily
+    from alembic import command
+
+    db_url = str(engine.url)
+    _logger.info("Updating database tables")
+    config = _get_alembic_config(db_url)
+    # Initialize a shared connection to be used for the database upgrade, ensuring that
+    # any connection-dependent state (e.g., the state of an in-memory database) is preserved
+    # for reference by the upgrade routine. For more information, see
+    # https://alembic.sqlalchemy.org/en/latest/cookbook.html#sharing-a-
+    # connection-with-a-series-of-migration-commands-and-environments
+    with engine.begin() as connection:
+        config.attributes["connection"] = connection
+        command.upgrade(config, "heads")
+
+
+def _get_schema_version(engine):
+    with engine.connect() as connection:
+        mc = MigrationContext.configure(connection)
+        return mc.get_current_revision()
+
+
+def create_sqlalchemy_engine_with_retry(db_uri):
+    attempts = 0
+    while True:
+        attempts += 1
+        engine = create_sqlalchemy_engine(db_uri)
+        try:
+            sqlalchemy.inspect(engine)
+            return engine
+        except Exception as e:
+            if attempts < MAX_RETRY_COUNT:
+                sleep_duration = 0.1 * ((2**attempts) - 1)
+                _logger.warning(
+                    "SQLAlchemy engine could not be created. The following exception is caught.\n"
+                    "%s\nOperation will be retried in %.1f seconds",
+                    e,
+                    sleep_duration,
+                )
+                time.sleep(sleep_duration)
+                continue
+            raise
+
+
+def create_sqlalchemy_engine(db_uri):
+    pool_size = MLFLOW_SQLALCHEMYSTORE_POOL_SIZE.get()
+    pool_max_overflow = MLFLOW_SQLALCHEMYSTORE_MAX_OVERFLOW.get()
+    pool_recycle = MLFLOW_SQLALCHEMYSTORE_POOL_RECYCLE.get()
+    echo = MLFLOW_SQLALCHEMYSTORE_ECHO.get()
+    poolclass = MLFLOW_SQLALCHEMYSTORE_POOLCLASS.get()
+    kwargs = {}
+    # Send argument only if they have been injected.
+    # Some engine does not support them (for example sqllite)
+    if pool_size:
+        kwargs["pool_size"] = pool_size
+    if pool_max_overflow:
+        kwargs["max_overflow"] = pool_max_overflow
+    if pool_recycle:
+        kwargs["pool_recycle"] = pool_recycle
+    if echo:
+        kwargs["echo"] = echo
+    if poolclass:
+        pool_class_map = {
+            "AssertionPool": AssertionPool,
+            "AsyncAdaptedQueuePool": AsyncAdaptedQueuePool,
+            "FallbackAsyncAdaptedQueuePool": FallbackAsyncAdaptedQueuePool,
+            "NullPool": NullPool,
+            "QueuePool": QueuePool,
+            "SingletonThreadPool": SingletonThreadPool,
+            "StaticPool": StaticPool,
+        }
+        if poolclass not in pool_class_map:
+            list_str = " ".join(pool_class_map.keys())
+            err_str = (
+                f"Invalid poolclass parameter: {poolclass}. Set environment variable "
+                f"poolclass to empty or one of the following values: {list_str}"
+            )
+            _logger.warning(err_str)
+            raise ValueError(err_str)
+        kwargs["poolclass"] = pool_class_map[poolclass]
+    if kwargs:
+        _logger.info("Create SQLAlchemy engine with pool options %s", kwargs)
+
+    # Handle MySQL SSL certificates via connect_args
+    if db_uri.startswith("mysql"):
+        connect_args = {
+            k: v
+            for k, v in {
+                "ssl_ca": MLFLOW_MYSQL_SSL_CA.get(),
+                "ssl_cert": MLFLOW_MYSQL_SSL_CERT.get(),
+                "ssl_key": MLFLOW_MYSQL_SSL_KEY.get(),
+            }.items()
+            if v
+        }
+        if connect_args:
+            kwargs["connect_args"] = connect_args
+
+    return sqlalchemy.create_engine(db_uri, pool_pre_ping=True, **kwargs)

mlflow/store/db_migrations/alembic.ini

@@ -0,0 +1,74 @@
+# A generic, single database configuration.
+
+[alembic]
+# path to migration scripts
+script_location = mlflow/store/db_migrations
+
+# template used to generate migration files
+# file_template = %%(rev)s_%%(slug)s
+
+# timezone to use when rendering the date
+# within the migration file as well as the filename.
+# string value is passed to dateutil.tz.gettz()
+# leave blank for localtime
+# timezone =
+
+# max length of characters to apply to the
+# "slug" field
+# truncate_slug_length = 40
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+# set to 'true' to allow .pyc and .pyo files without
+# a source .py file to be detected as revisions in the
+# versions/ directory
+# sourceless = false
+
+# version location specification; this defaults
+# to alembic/versions.  When using multiple version
+# directories, initial revisions must be specified with --version-path
+# version_locations = %(here)s/bar %(here)s/bat alembic/versions
+
+# the output encoding used when revision files
+# are written from script.py.mako
+# output_encoding = utf-8
+
+sqlalchemy.url = ""
+
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S

mlflow/store/db_migrations/env.py

@@ -0,0 +1,84 @@
+from logging.config import fileConfig
+
+from alembic import context
+from sqlalchemy import engine_from_config, pool
+
+# this is the Alembic Config object, which provides
+# access to the values within the .ini file in use.
+config = context.config
+
+# Interpret the config file for Python logging.
+# This line sets up loggers basically.
+fileConfig(config.config_file_name, disable_existing_loggers=False)
+
+# add your model's MetaData object here
+# for 'autogenerate' support
+# from myapp import mymodel
+# target_metadata = mymodel.Base.metadata
+from mlflow.store.db.base_sql_model import Base
+
+target_metadata = Base.metadata
+
+# other values from the config, defined by the needs of env.py,
+# can be acquired:
+# my_important_option = config.get_main_option("my_important_option")
+# ... etc.
+
+
+def run_migrations_offline():
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+    url = config.get_main_option("sqlalchemy.url")
+    # Try https://stackoverflow.com/questions/30378233/sqlite-lack-of-alter-support-alembic-migration-failing-because-of-this-solutio
+    context.configure(
+        url=url, target_metadata=target_metadata, literal_binds=True, render_as_batch=True
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online():
+    """Run migrations in 'online' mode.
+
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+
+    """
+    # If available, use a shared connection for the database upgrade, ensuring that any
+    # connection-dependent state (e.g., the state of an in-memory database) is preserved
+    # for reference by the upgrade routine. For more information, see
+    # https://alembic.sqlalchemy.org/en/latest/cookbook.html#sharing-a-
+    # connection-with-a-series-of-migration-commands-and-environments
+    connection = config.attributes.get("connection")
+    if connection is None:
+        engine = engine_from_config(
+            config.get_section(config.config_ini_section),
+            prefix="sqlalchemy.",
+            poolclass=pool.NullPool,
+        )
+    else:
+        engine = connection.engine
+
+    with engine.connect() as connection:
+        context.configure(
+            connection=connection, target_metadata=target_metadata, render_as_batch=True
+        )
+
+        with context.begin_transaction():
+            context.run_migrations()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()

mlflow/store/db_migrations/versions/0584bdc529eb_add_cascading_deletion_to_datasets_from_experiments.py

@@ -0,0 +1,88 @@
+"""add cascading deletion to datasets from experiments
+
+Revision ID: 0584bdc529eb
+Revises: f5a4f2784254
+Create Date: 2024-11-11 15:27:53.189685
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+from mlflow.exceptions import MlflowException
+from mlflow.store.tracking.dbmodels.models import SqlDataset, SqlExperiment
+
+# revision identifiers, used by Alembic.
+revision = "0584bdc529eb"
+down_revision = "f5a4f2784254"
+branch_labels = None
+depends_on = None
+
+
+def get_datasets_experiment_fk_name():
+    conn = op.get_bind()
+    metadata = sa.MetaData()
+    metadata.bind = conn
+    datasets_table = sa.Table(
+        SqlDataset.__tablename__,
+        metadata,
+        autoload_with=conn,
+    )
+
+    for constraint in datasets_table.foreign_key_constraints:
+        if (
+            constraint.referred_table.name == SqlExperiment.__tablename__
+            and constraint.column_keys[0] == "experiment_id"
+        ):
+            return constraint.name
+
+    raise MlflowException(
+        "Unable to find the foreign key constraint name from datasets to experiments. "
+        "All foreign key constraints in datasets table: \n"
+        f"{datasets_table.foreign_key_constraints}"
+    )
+
+
+def upgrade():
+    dialect_name = op.get_context().dialect.name
+
+    # standardize the constraint to sqlite naming convention
+    new_fk_constraint_name = (
+        f"fk_{SqlDataset.__tablename__}_experiment_id_{SqlExperiment.__tablename__}"
+    )
+
+    if dialect_name == "sqlite":
+        # Only way to drop unnamed fk constraint in sqllite
+        # See https://alembic.sqlalchemy.org/en/latest/batch.html#dropping-unnamed-or-named-foreign-key-constraints
+        with op.batch_alter_table(
+            SqlDataset.__tablename__,
+            schema=None,
+            naming_convention={
+                "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
+            },
+        ) as batch_op:
+            # in SQLite, constraint.name is None, so we have to hardcode it
+            batch_op.drop_constraint(new_fk_constraint_name, type_="foreignkey")
+            # Need to explicitly name the fk constraint with batch alter table
+            batch_op.create_foreign_key(
+                new_fk_constraint_name,
+                SqlExperiment.__tablename__,
+                ["experiment_id"],
+                ["experiment_id"],
+                ondelete="CASCADE",
+            )
+    else:
+        old_fk_constraint_name = get_datasets_experiment_fk_name()
+        op.drop_constraint(old_fk_constraint_name, SqlDataset.__tablename__, type_="foreignkey")
+        op.create_foreign_key(
+            new_fk_constraint_name,
+            SqlDataset.__tablename__,
+            SqlExperiment.__tablename__,
+            ["experiment_id"],
+            ["experiment_id"],
+            ondelete="CASCADE",
+        )
+
+
+def downgrade():
+    pass

mlflow/store/db_migrations/versions/0a8213491aaa_drop_duplicate_killed_constraint.py

@@ -0,0 +1,49 @@
+"""drop_duplicate_killed_constraint
+
+Revision ID: 0a8213491aaa
+Revises: cfd24bdc0731
+Create Date: 2020-01-28 15:26:14.757445
+
+This migration drops a duplicate constraint on the `runs.status` column that was left as a byproduct
+of an erroneous implementation of the `cfd24bdc0731_update_run_status_constraint_with_killed`
+migration in MLflow 1.5. The implementation of this migration has since been fixed.
+"""
+
+import logging
+
+from alembic import op
+
+_logger = logging.getLogger(__name__)
+
+# revision identifiers, used by Alembic.
+revision = "0a8213491aaa"
+down_revision = "cfd24bdc0731"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # Attempt to drop any existing `status` constraints on the `runs` table. This operation
+    # may fail against certain backends with different classes of Exception. For example,
+    # in MySQL <= 8.0.15, dropping constraints produces an invalid `ALTER TABLE` expression.
+    # Further, in certain versions of sqlite, `ALTER` (which is invoked by `drop_constraint`)
+    # is unsupported on `CHECK` constraints. Accordingly, we catch the generic `Exception`
+    # object because the failure modes are not well-enumerated or consistent across database
+    # backends. Because failures automatically stop batch operations and the `drop_constraint()`
+    # operation is expected to fail under certain circumstances, we execute `drop_constraint()`
+    # outside of the batch operation context.
+    try:
+        # For other database backends, the status check constraint is dropped by
+        # cfd24bdc0731_update_run_status_constraint_with_killed.py
+        if op.get_bind().engine.name == "mysql":
+            op.drop_constraint(constraint_name="status", table_name="runs", type_="check")
+    except Exception as e:
+        _logger.warning(
+            "Failed to drop check constraint. Dropping check constraints may not be supported"
+            " by your SQL database. Exception content: %s",
+            e,
+        )
+
+
+def downgrade():
+    pass

mlflow/store/db_migrations/versions/0c779009ac13_add_deleted_time_field_to_runs_table.py

@@ -0,0 +1,24 @@
+"""add deleted_time field to runs table
+
+Revision ID: 0c779009ac13
+Revises: bd07f7e963c5
+Create Date: 2022-07-27 14:13:36.162861
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "0c779009ac13"
+down_revision = "bd07f7e963c5"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.add_column("runs", sa.Column("deleted_time", sa.BigInteger, nullable=True, default=None))
+
+
+def downgrade():
+    pass

mlflow/store/db_migrations/versions/181f10493468_allow_nulls_for_metric_values.py

@@ -0,0 +1,35 @@
+"""allow nulls for metric values
+
+Revision ID: 181f10493468
+Revises: 90e64c465722
+Create Date: 2019-07-10 22:40:18.787993
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "181f10493468"
+down_revision = "90e64c465722"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    with op.batch_alter_table("metrics") as batch_op:
+        batch_op.alter_column("value", type_=sa.types.Float(precision=53), nullable=False)
+        batch_op.add_column(
+            sa.Column(
+                "is_nan", sa.Boolean(create_constraint=False), nullable=False, server_default="0"
+            )
+        )
+        batch_op.drop_constraint(constraint_name="metric_pk", type_="primary")
+        batch_op.create_primary_key(
+            constraint_name="metric_pk",
+            columns=["key", "timestamp", "step", "run_uuid", "value", "is_nan"],
+        )
+
+
+def downgrade():
+    pass