wehub-resource-sync
471 lines
19 KiB
Python
471 lines
19 KiB
Python
import hashlib
|
|
import logging
|
|
import os
|
|
import re
|
|
import sqlite3
|
|
import tempfile
|
|
import time
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
|
|
import sqlalchemy
|
|
from alembic.migration import MigrationContext
|
|
from alembic.script import ScriptDirectory
|
|
from packaging.version import Version
|
|
from sqlalchemy import event, sql
|
|
|
|
# We need to import sqlalchemy.pool to convert poolclass string to class object
|
|
from sqlalchemy.pool import (
|
|
AssertionPool,
|
|
AsyncAdaptedQueuePool,
|
|
NullPool,
|
|
QueuePool,
|
|
SingletonThreadPool,
|
|
StaticPool,
|
|
)
|
|
|
|
# CRITICAL: Import ORM modules to register all table metadata with Base.metadata.
|
|
# _all_tables_exist() depends on Base.metadata.tables being fully populated.
|
|
# If these imports are removed, Base.metadata.tables will be incomplete and
|
|
# _all_tables_exist() will return True even when tables are missing, silently
|
|
# skipping DB initialization/migration.
|
|
import mlflow.store.model_registry.dbmodels.models # noqa: F401
|
|
import mlflow.store.tracking.dbmodels.models # noqa: F401
|
|
import mlflow.store.workspace.dbmodels.models # noqa: F401
|
|
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.base_sql_model import Base
|
|
from mlflow.store.db.db_types import SQLITE
|
|
from mlflow.store.tracking.dbmodels.initial_models import Base as InitialBase
|
|
|
|
_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):
|
|
# Check if all ORM-defined tables exist in the database.
|
|
# Derived from Base.metadata.tables so that newly added ORM models are
|
|
# automatically included without requiring a manual update here.
|
|
# Using issubset() instead of equality so that additional tables added by migrations
|
|
# don't cause this check to fail. This prevents unnecessary calls to _initialize_tables
|
|
# which can cause migration errors like "Can't locate revision identified by 'xxx'".
|
|
expected_tables = set(Base.metadata.tables)
|
|
actual_tables = {
|
|
t for t in sqlalchemy.inspect(engine).get_table_names() if not t.startswith("alembic_")
|
|
}
|
|
return expected_tables.issubset(actual_tables)
|
|
|
|
|
|
def _is_empty_database(engine):
|
|
# Alembic's bookkeeping table is ignored so a previously failed upgrade attempt
|
|
# doesn't make the DB look populated.
|
|
return all(t.startswith("alembic_") for t in sqlalchemy.inspect(engine).get_table_names())
|
|
|
|
|
|
def _initialize_tables(engine):
|
|
_logger.info("Creating initial MLflow database tables...")
|
|
InitialBase.metadata.create_all(engine)
|
|
_upgrade_db(engine)
|
|
|
|
|
|
def _safe_initialize_tables(engine: sqlalchemy.engine.Engine) -> None:
|
|
from mlflow.utils.file_utils import ExclusiveFileLock
|
|
|
|
if os.name == "nt":
|
|
if not _all_tables_exist(engine):
|
|
_initialize_tables(engine)
|
|
return
|
|
|
|
url_hash = hashlib.md5(
|
|
str(engine.url).encode("utf-8"),
|
|
usedforsecurity=False,
|
|
).hexdigest()
|
|
with ExclusiveFileLock(f"{tempfile.gettempdir()}/db_init_lock-{url_hash}"):
|
|
if not _all_tables_exist(engine):
|
|
_initialize_tables(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.
|
|
"""
|
|
|
|
# Default to read_only=True to avoid accidentally treating read query as write
|
|
# and send to the main workers. The test time validation detects when the
|
|
# actual write query is marked read-only.
|
|
@contextmanager
|
|
def make_managed_session(read_only=True):
|
|
"""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;"))
|
|
|
|
# Validation that only triggers while testing to make sure the
|
|
# read_only flag is properly set.
|
|
if read_only and os.environ.get("PYTEST_CURRENT_TEST"):
|
|
|
|
def _reject_flush(session, flush_context, instances):
|
|
raise MlflowException(
|
|
"Write operation detected on a read-only session. "
|
|
"ManagedSessionMaker creates read-only sessions "
|
|
"by default. Pass read_only=False to the constructor "
|
|
"to mark it as a write operation.",
|
|
error_code=INTERNAL_ERROR,
|
|
)
|
|
|
|
event.listen(session, "before_flush", _reject_flush)
|
|
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_routing_session_maker(write_session_maker, read_session_maker, db_type):
|
|
"""
|
|
Creates a routing-aware managed session factory that directs read operations to a
|
|
read replica and write operations to the primary database.
|
|
|
|
When a read replica URI is configured, read-only operations (marked with
|
|
``read_only=True``) use the read replica's session, while all other operations
|
|
use the primary (write) session. If no read replica is configured, both session
|
|
makers should be identical and all operations go to the primary database.
|
|
|
|
Args:
|
|
write_session_maker: SQLAlchemy ``sessionmaker`` bound to the primary (write) engine.
|
|
read_session_maker: SQLAlchemy ``sessionmaker`` bound to the read replica engine.
|
|
db_type: Database type string (e.g. ``"sqlite"``) forwarded to
|
|
``_get_managed_session_maker`` for dialect-specific session configuration.
|
|
|
|
Returns:
|
|
A context-manager factory accepting an optional ``read_only`` keyword argument.
|
|
"""
|
|
_write = _get_managed_session_maker(write_session_maker, db_type)
|
|
_read = _get_managed_session_maker(read_session_maker, db_type)
|
|
|
|
@contextmanager
|
|
def make_managed_session(read_only=True):
|
|
with (_read if read_only else _write)(read_only=read_only) as session:
|
|
yield session
|
|
|
|
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 _check_sqlite_version(db_url: str) -> None:
|
|
"""
|
|
Check if SQLite version supports required features.
|
|
|
|
MLflow requires SQLite 3.31.0 or higher for computed columns support.
|
|
Raises MlflowException if the version is too old.
|
|
"""
|
|
if not db_url.startswith("sqlite:"):
|
|
return
|
|
|
|
version_str = sqlite3.sqlite_version
|
|
try:
|
|
sqlite_version = Version(version_str)
|
|
except Exception:
|
|
return
|
|
min_version_str = "3.31.0"
|
|
if sqlite_version < Version(min_version_str):
|
|
raise MlflowException(
|
|
f"MLflow requires SQLite >= {min_version_str} for SQL based "
|
|
f"store, but found {version_str}. Please upgrade your SQLite. "
|
|
f"See https://www.sqlite.org/download.html for installation options."
|
|
)
|
|
|
|
|
|
def _upgrade_db(engine):
|
|
"""
|
|
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)
|
|
# Check SQLite version before running migrations
|
|
_check_sqlite_version(db_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 _make_parent_dirs_if_sqlite(db_uri: str) -> None:
|
|
"""Create parent directories for SQLite database file if they don't exist."""
|
|
if not db_uri.startswith("sqlite:///"):
|
|
return
|
|
# SQLite URI format: sqlite:///path/to/file.db (relative) or sqlite:////abs/path (Unix)
|
|
# Remove the 'sqlite:///' prefix to get the path
|
|
# Skip in-memory databases (:memory: or empty path)
|
|
db_path = db_uri.removeprefix("sqlite:///")
|
|
if db_path and db_path != ":memory:":
|
|
Path(db_path).parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def create_sqlalchemy_engine_with_retry(db_uri):
|
|
_make_parent_dirs_if_sqlite(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,
|
|
"NullPool": NullPool,
|
|
"QueuePool": QueuePool,
|
|
"SingletonThreadPool": SingletonThreadPool,
|
|
"StaticPool": StaticPool,
|
|
}
|
|
try:
|
|
# FallbackAsyncAdaptedQueuePool was removed in SQLAlchemy 2.1
|
|
from sqlalchemy.pool import FallbackAsyncAdaptedQueuePool
|
|
|
|
pool_class_map["FallbackAsyncAdaptedQueuePool"] = FallbackAsyncAdaptedQueuePool
|
|
except ImportError:
|
|
pass
|
|
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
|
|
|
|
engine = sqlalchemy.create_engine(db_uri, pool_pre_ping=True, **kwargs)
|
|
|
|
# Register custom functions for SQLite
|
|
if db_uri.startswith("sqlite"):
|
|
|
|
@event.listens_for(engine, "connect")
|
|
def _register_sqlite_functions(dbapi_conn, connection_record):
|
|
# Register REGEXP function to enable RLIKE operator support
|
|
def regexp(pattern, string):
|
|
"""Custom REGEXP function for SQLite that uses Python's re module."""
|
|
if string is None or pattern is None:
|
|
return False
|
|
try:
|
|
return re.search(pattern, string) is not None
|
|
except re.error:
|
|
return False
|
|
|
|
dbapi_conn.create_function("regexp", 2, regexp)
|
|
|
|
# Register PERCENTILE aggregate function (PERCENTILE_CONT with linear interpolation)
|
|
class PercentileAggregate:
|
|
def __init__(self):
|
|
self.values = []
|
|
self.percentile_value = None
|
|
|
|
def step(self, value, percentile):
|
|
if value is not None:
|
|
self.values.append(float(value))
|
|
if self.percentile_value is None and percentile is not None:
|
|
self.percentile_value = float(percentile)
|
|
|
|
def finalize(self):
|
|
if not self.values or self.percentile_value is None:
|
|
return None
|
|
|
|
sorted_values = sorted(self.values)
|
|
n = len(sorted_values)
|
|
percentile_ratio = self.percentile_value / 100.0
|
|
|
|
# Linear interpolation: index = percentile_ratio * (n - 1)
|
|
index = percentile_ratio * (n - 1)
|
|
lower_idx = int(index)
|
|
upper_idx = min(lower_idx + 1, n - 1)
|
|
fraction = index - lower_idx
|
|
|
|
return sorted_values[lower_idx] + fraction * (
|
|
sorted_values[upper_idx] - sorted_values[lower_idx]
|
|
)
|
|
|
|
dbapi_conn.create_aggregate("percentile", 2, PercentileAggregate)
|
|
|
|
# SQLAlchemy's MSSQL dialect calls the base fetch_clause() which generates
|
|
# "FETCH FIRST n ROWS ONLY" but SQL Server requires "FETCH NEXT n ROWS ONLY".
|
|
# Register an event listener to rewrite the generated SQL before execution.
|
|
if db_uri.startswith("mssql"):
|
|
|
|
@event.listens_for(engine, "before_cursor_execute", retval=True)
|
|
def _fix_mssql_fetch_syntax(conn, cursor, statement, parameters, context, executemany):
|
|
if " FETCH FIRST " in statement:
|
|
statement = statement.replace(" FETCH FIRST ", " FETCH NEXT ")
|
|
return statement, parameters
|
|
|
|
return engine
|