wehub-resource-sync
109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
import shutil
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import sqlalchemy
|
|
|
|
from mlflow.environment_variables import MLFLOW_TRACKING_URI
|
|
from mlflow.store.db.db_types import MSSQL, MYSQL, POSTGRES, SQLITE
|
|
from mlflow.store.tracking.dbmodels.models import (
|
|
SqlDataset,
|
|
SqlEntityAssociation,
|
|
SqlEvaluationDataset,
|
|
SqlEvaluationDatasetRecord,
|
|
SqlExperiment,
|
|
SqlExperimentTag,
|
|
SqlInput,
|
|
SqlInputTag,
|
|
SqlLatestMetric,
|
|
SqlLoggedModel,
|
|
SqlLoggedModelMetric,
|
|
SqlLoggedModelParam,
|
|
SqlLoggedModelTag,
|
|
SqlMetric,
|
|
SqlOnlineScoringConfig,
|
|
SqlParam,
|
|
SqlRun,
|
|
SqlTag,
|
|
SqlTraceInfo,
|
|
SqlTraceMetadata,
|
|
SqlTraceTag,
|
|
)
|
|
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
|
|
|
ARTIFACT_URI = "artifact_folder"
|
|
|
|
|
|
def _get_query_to_reset_experiment_id(store: SqlAlchemyStore):
|
|
dialect = store._get_dialect()
|
|
if dialect == POSTGRES:
|
|
return "ALTER SEQUENCE experiments_experiment_id_seq RESTART WITH 1"
|
|
elif dialect == MYSQL:
|
|
return "ALTER TABLE experiments AUTO_INCREMENT = 1"
|
|
elif dialect == MSSQL:
|
|
return "DBCC CHECKIDENT (experiments, RESEED, 0)"
|
|
elif dialect == SQLITE:
|
|
# In SQLite, deleting all experiments resets experiment_id
|
|
return None
|
|
raise ValueError(f"Invalid dialect: {dialect}")
|
|
|
|
|
|
def _cleanup_database(store: SqlAlchemyStore):
|
|
with store.ManagedSessionMaker() as session:
|
|
# Delete all rows in all tables
|
|
for model in (
|
|
SqlLoggedModel,
|
|
SqlLoggedModelMetric,
|
|
SqlLoggedModelParam,
|
|
SqlLoggedModelTag,
|
|
SqlParam,
|
|
SqlMetric,
|
|
SqlLatestMetric,
|
|
SqlTag,
|
|
SqlInputTag,
|
|
SqlInput,
|
|
SqlDataset,
|
|
SqlRun,
|
|
SqlTraceTag,
|
|
SqlTraceMetadata,
|
|
SqlTraceInfo,
|
|
SqlEvaluationDatasetRecord,
|
|
SqlEntityAssociation,
|
|
SqlEvaluationDataset,
|
|
SqlExperimentTag,
|
|
SqlOnlineScoringConfig,
|
|
SqlExperiment,
|
|
):
|
|
session.query(model).delete()
|
|
|
|
# Reset experiment_id to start at 1
|
|
if reset_experiment_id := _get_query_to_reset_experiment_id(store):
|
|
session.execute(sqlalchemy.sql.text(reset_experiment_id))
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def cached_db(tmp_path_factory) -> Path:
|
|
"""Creates and caches a SQLite database to avoid repeated migrations for each test run."""
|
|
tmp_path = tmp_path_factory.mktemp("sqlite_db")
|
|
db_path = tmp_path / "mlflow.db"
|
|
db_uri = f"sqlite:///{db_path}"
|
|
store = SqlAlchemyStore(db_uri, ARTIFACT_URI)
|
|
store.engine.dispose()
|
|
return db_path
|
|
|
|
|
|
@pytest.fixture
|
|
def store(tmp_path: Path, cached_db: Path) -> SqlAlchemyStore:
|
|
artifact_uri = tmp_path / "artifacts"
|
|
artifact_uri.mkdir(exist_ok=True)
|
|
if db_uri_env := MLFLOW_TRACKING_URI.get():
|
|
s = SqlAlchemyStore(db_uri_env, artifact_uri.as_uri())
|
|
yield s
|
|
_cleanup_database(s)
|
|
else:
|
|
db_path = tmp_path / "mlflow.db"
|
|
shutil.copy(cached_db, db_path)
|
|
db_uri = f"sqlite:///{db_path}"
|
|
s = SqlAlchemyStore(db_uri, artifact_uri.as_uri())
|
|
yield s
|