使用 Docker Compose 在 MLflow 仓库中实测 PostgreSQL、MySQL、MSSQL 与 SQLite 跟踪存储

使用 Docker Compose 在 MLflow 仓库中实测 PostgreSQL、MySQL、MSSQL 与 SQLite 跟踪存储

【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow

本指南以 MLflow 仓库内 tests/db/README.md 为骨架，完整讲解如何用 Docker Compose 一键拉起 PostgreSQL、MySQL、Microsoft SQL Server、SQLite 四种后端数据库，对 MLflow Tracking 与 Model Registry 的存储层进行真实数据库层面的验证。你将掌握compose.sh的构建、运行、清理全流程，理解compose.yml中每个数据库服务与MLFLOW_TRACKING_URI的对应关系，并借助test_schema.py、check_migration.sh等源码看清 Schema 一致性校验与迁移测试的底层机制，最终具备在本仓库中自行跑通多数据库集成测试、排查数据库兼容性问题的完整实战能力。

这套数据库测试体系要解决什么问题

MLflow 的 Tracking Store 与 Model Registry Store 支持通过 SQLAlchemy 对接多种关系型数据库。数据库方言（Dialect）之间的行为差异——比如唯一约束的反射方式、驱动连接串、事务与错误语义——只有在真实数据库上才能暴露出来，因此仓库在 tests/db 目录下建立了一套完整的容器化测试体系，用 Docker Compose 同时管理四种后端：

• PostgreSQL
• MySQL
• Microsoft SQL Server（MSSQL）
• SQLite

该目录不仅承载测试用例，还包含构建镜像所需的 Dockerfile、Dockerfile.mssql、编排文件 compose.yml、入口脚本 entrypoint.sh、Schema 快照 schemas 以及迁移校验脚本 check_migration.sh，是一套自洽的"数据库即服务"测试环境。

环境前置条件

依据 tests/db/README.md，运行这套测试只需要两样东西：

• Docker：用于构建和运行数据库与测试容器；
• Docker Compose V2：compose.sh内部调用的是docker compose（插件式命令，非旧版docker-compose）。

此外，由于测试容器会以可编辑方式安装当前仓库源码（详见后文 entrypoint），通常需要在仓库根目录下执行命令，并保证本机 Python 环境可用（用于生成依赖列表）。操作系统与 CPU 架构方面，MSSQL 镜像在 compose.yml 中显式声明了platform: linux/amd64，在 Apple Silicon 等 ARM 主机上会通过 Rosetta 模拟，这一点在运行时需要留意。

构建服务镜像

构建单个服务

# Build a service service=mlflow-sqlite ./tests/db/compose.sh build --build-arg DEPENDENCIES="$(python dev/extract_deps.py)" $service

构建全部服务

# Build all services ./tests/db/compose.sh build --build-arg DEPENDENCIES="$(python dev/extract_deps.py)"

这里的DEPENDENCIES构建参数非常关键，它来自 dev/extract_deps.py：该脚本读取仓库根目录pyproject.toml，用正则提取dependencies = [...]列表并通过ast.literal_eval解析，把 MLflow 自身声明的依赖项逐行打印出来。这些依赖会在 Dockerfile 中被写入/tmp/requirements.txt并安装：

FROM python:3.11 COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ ARG DEPENDENCIES RUN uv pip install --system psycopg2 "psycopg[binary]" pymysql mysqlclient psutil pytest pytest-cov pytest-asyncio RUN echo "${DEPENDENCIES}" > /tmp/requirements.txt && \ uv pip install --system -r /tmp/requirements.txt RUN uv pip list

从 Dockerfile 可以看到镜像底座是python:3.11，包管理使用uv，并预装了各数据库驱动与测试工具链：

• PostgreSQL 驱动：psycopg2与psycopg[binary]；
• MySQL 驱动：pymysql、mysqlclient；
• 测试框架：pytest、pytest-cov、pytest-asyncio；
• 以及运行测试所需的psutil。

uv pip list会把最终安装的包清单打印到构建日志中，便于排查依赖缺失。MSSQL 的构建则使用 Dockerfile.mssql，额外安装 SQL Server 的 ODBC 驱动 17 与sqlcmd工具，用于执行初始化脚本。

运行服务与测试

运行单个服务

# Run a service (`pytest tests/db` is executed by default) ./tests/db/compose.sh run --rm $service

容器默认命令由 compose.yml 中base服务的command: pytest tests/db定义，也就是说直接run一个mlflow-*服务就会执行整目录的数据库测试。若要指定测试范围，可以覆盖命令：

# Run tests ./tests/db/compose.sh run --rm $service pytest /path/to/directory/or/script # Run a python script ./tests/db/compose.sh run --rm $service python /path/to/script

运行全部服务

# Run all services for service in $(./tests/db/compose.sh config --services | grep '^mlflow-') do ./tests/db/compose.sh run --rm "$service" done

compose.sh config --services会输出 compose.yml 中定义的全部服务名，再用grep '^mlflow-'过滤出 MLflow 测试容器（跳过postgresql、mysql、mssql等纯数据库容器），逐个以run --rm方式执行，确保退出后即清理容器。仓库还提供了 update_schemas.sh 作为一键化流程：它先通过tests/store/dump_schema.py生成基础 Schema，再构建并依次对每个mlflow-*服务执行python tests/db/test_schema.py，用于统一刷新 Schema 快照。

compose.sh 做了什么

tests/db/compose.sh 本身是一个极简包装器：

#!/bin/bash set -ex docker compose --project-directory tests/db down --volumes --remove-orphans > /dev/null 2>&1 docker compose --project-directory tests/db "$@"

它做了两件事：

1. 每次执行前先强制down --volumes --remove-orphans，确保旧的数据库数据卷、容器与孤儿容器被清空，让测试始终从干净状态开始（这也是为什么 README 建议构建/运行/清理都用这个脚本而不是直接docker compose）；
2. 把后续参数透传给docker compose --project-directory tests/db，使 compose 在tests/db目录下解析compose.yml，并且会把仓库根目录挂载进容器。

容器内发生了什么：entrypoint.sh

entrypoint.sh 是每个 MLflow 测试容器的入口，它按顺序执行三件事：

#!/bin/bash set -ex # Install mlflow (assuming the repository root is mounted to the working directory) if [ "$INSTALL_MLFLOW_FROM_REPO" = "true" ]; then uv pip install --system --no-deps -e . fi # For Microsoft SQL server, wait until the database is up and running if [[ $MLFLOW_TRACKING_URI == mssql* ]]; then ./tests/db/init-mssql-db.sh fi # Execute the command exec "$@"

• 当环境变量INSTALL_MLFLOW_FROM_REPO=true（mlflow-*服务均设置了该值）时，用uv pip install --system --no-deps -e .以可编辑模式安装仓库根目录的 MLflow，保证被测的是当前源码而非 PyPI 包；
• 对mssql+pyodbc://开头的 Tracking URI，执行 init-mssql-db.sh。该脚本用 0、1、2、4、8、16 秒的指数退避轮询sqlcmd，等待 SQL Server 就绪后执行 init-mssql-db.sql 完成建库建用户；
• 最后exec "$@"执行外部传入的命令（默认pytest tests/db）。

四种数据库的 compose 服务拓扑

tests/db/compose.yml 中每个数据库都包含一个纯数据库容器和一个（或两个）MLflow 测试容器，并通过extends: service: base继承公共配置：

base服务还设置了DISABLE_RESET_MLFLOW_URI_FIXTURE: "true"，表示不使用测试框架中重置 Tracking URI 的 fixture，让容器通过环境变量MLFLOW_TRACKING_URI直接决定后端。值得注意的细节有：

• MySQL 唯一约束：compose.yml 为 MySQL 容器追加了--log-bin-trust-function-creators=1，以允许在二进制日志开启时创建函数，规避迁移脚本执行中的权限限制；
• MSSQL 驱动选择：Tracking URI 使用mssql+pyodbc://并显式指定driver=ODBC+Driver+17+for+SQL+Server。compose.yml 中的注释说明：若改用 ODBC Driver 18，可在 URI 追加LongAsMax=Yes以规避 SQLAlchemy 2.0 以下版本将长字符串参数发送为TEXT/NTEXT导致的 "varchar and ntext are incompatible in the equal to operator" 报错（详见 SQLAlchemy 官方 MSSQL 方言文档）；
• SQLite 的绝对路径：sqlite:////tmp/mlflowdb是四斜杠写法，前两个表示空主机名（本地文件），后两个是/tmp/mlflowdb的绝对路径。

若未设置MLFLOW_TRACKING_URI，tests/db/conftest.py 中的 autouse fixture 会兜底：将 Tracking URI 指向tmp_path下的mlruns.sqlite文件，保证测试在本地开发环境下（无数据库容器时）也能直接运行。

数据库层级的跟踪操作测试

tests/db/test_tracking_operations.py 是四个数据库共同执行的"冒烟测试"（该文件标记了pytest.mark.notrackingurimock，避免全局 URI mock 干扰真实后端连接）：

1.test_search_runs：mlflow.start_run()内写入参数、指标、标签，并log_model注册模型，随后mlflow.search_runs按param.start_time DESC排序查询并读取 run，验证四种数据库的读写与排序行为。

2.test_set_run_status_to_killed：将 run 状态置为KILLED，专门回归两条迁移脚本cfd24bdc0731_update_run_status_constraint_with_killed.py与0a8213491aaa_drop_duplicate_killed_constraint.py引入的状态约束，确认所有方言下都能正确写入并终止 run。

3.test_database_operational_error（仅 SQLite 执行，其余后端pytest.skip）：通过 monkeypatch 包装 SQLAlchemy SQLite 方言的dbapi/import_dbapi与connect，在写入特定参数值时人为抛出sqlite3.OperationalError，验证 MLflow 将数据库操作错误（PEP 249 定义的连接断开、事务失败等瞬时错误）与通用SQLAlchemyError区分处理——REST 客户端场景下，TEMPORARILY_UNAVAILABLE会触发重试而BAD_REQUEST不会。测试同时断言mlflow.store.db.utils._logger.exception被调用且日志包含 "SQLAlchemy database error" 与 "sqlite3.OperationalError"。

这三个用例覆盖了真实数据库上的写入、查询、状态流转与错误处理四条关键路径，是理解 MLflow 存储层行为的最直接样本。

Schema 一致性校验：test_schema.py

tests/db/test_schema.py 负责保证 MLflow 通过 SQLAlchemy ORM 创建的表结构与 schemas 目录下按方言存放的 Schema 快照完全一致。目录中的四个快照文件 postgresql.sql、mysql.sql、mssql.sql、sqlite.sql（后者包含alembic_version、budget_policies、entity_associations、evaluation_datasets等约 787 行建表语句）就是各数据库的"黄金基准"。

核心流程在test_schema_is_up_to_date中：

1. initialize_database()以mlflow.start_run()触发建库；
2. dump_schema(tracking_uri)用 SQLAlchemy 的MetaData.reflect反射出全部表，再经_reattach_missing_unique_constraints补充反射遗漏的唯一约束，最后用CreateTable序列化为规范文本；
3. 与schemas/<dialect>.sql比较，不一致则输出包含 EXPECTED / ACTUAL / DIFF / HOW TO FIX 四段内容的详细报告。

其中_reattach_missing_unique_constraints专门处理跨方言差异：MySQL 把唯一约束反射成唯一索引，MSSQL 根本不实现get_unique_constraints，因此代码先尝试get_unique_constraints，失败或为空时回退到get_indexes中unique=True的索引，并以_DIALECT_REFLECTED_UNIQUE_CONSTRAINTS中记录的约束名（如uq_experiments_workspace_name）为准补回UniqueConstraint对象。schema_equal则用正则把CREATE TABLE拆成"表名 + 列集合"后再比较，忽略方言间的空白与格式差异。

当 Schema 漂移时，该脚本会给出修复命令：

docker compose -f tests/db/compose.yml run --rm mlflow-<dialect> python tests/db/test_schema.py

即重新初始化数据库后用main()分支把最新 Schema 写回快照文件（执行前会断言MLFLOW_TRACKING_URI已设置且确实指向数据库 URI）。这也是贡献者修改 ORM 模型后刷新快照的标准姿势。

迁移兼容性校验：check_migration.sh 与 check_migration.py

迁移测试用于验证"旧版本数据升级到新 Schema 后数据不丢失、结构正确"。入口是 check_migration.sh：

#!/bin/bash set -ex cd tests/db # Install the lastest version of mlflow from PyPI uv pip install --system mlflow python check_migration.py pre-migration # Install mlflow from the repository uv pip install --system -e ../.. mlflow db upgrade $MLFLOW_TRACKING_URI python check_migration.py post-migration

执行分四步：

1. 安装 PyPI 最新版 MLflow（uv pip install --system mlflow），即"旧版本"；
2. python check_migration.py pre-migration用旧版本写入一批数据并打快照；
3. 切回仓库源码（uv pip install --system -e ../..）并执行mlflow db upgrade $MLFLOW_TRACKING_URI，让 Alembic 把旧 Schema 升级到当前版本；
4. python check_migration.py post-migration校验升级后的数据与快照一致。

check_migration.py（用 Click 实现 CLI）中：

• pre_migration循环 5 次调用log_everything()——写入 experiment、run、params、metrics、tags，注册模型与模型版本并打prod别名，创建 webhook，并直接向evaluation_datasets、jobs表插入工作区数据——然后把 11 张核心表（SqlExperiment、SqlRun、SqlMetric、SqlParam、SqlTag、SqlExperimentTag、SqlLatestMetric、SqlRegisteredModel、SqlModelVersion、SqlRegisteredModelTag、SqlModelVersionTag，均来自 mlflow/store/tracking/dbmodels/models.py 与 mlflow/store/model_registry/dbmodels/models.py）导出为 pickle 快照存入tests/db/snapshots/；
• post_migration用pd.testing.assert_frame_equal逐表比对迁移后的数据与快照，并对WORKSPACE_TABLES（experiments、registered_models、model_versions、registered_model_tags、model_version_tags、registered_model_aliases、evaluation_datasets、webhooks、jobs）断言两件事：不存在 NULL workspace，且 workspace 全部等于default——这正是工作区（workspace）回填迁移的核心验收标准。

清理环境

测试结束后按 README 建议用compose.sh统一清理，而不是直接执行docker compose down：

# Clean up containers, networks, and volumes ./tests/db/compose.sh down --volumes --remove-orphans # Clean up containers, networks, volumes, and images ./tests/db/compose.sh down --volumes --remove-orphans --rmi all

第一条删除容器、网络与数据卷（含数据库持久化数据，保证下次测试从干净库开始），第二条追加--rmi all连同相关镜像一并删除。由于compose.sh本身在每次调用前就会执行一次down --volumes --remove-orphans，重复执行清理命令是幂等且安全的。

其他常用命令

查看某个数据库服务的实时日志，便于定位连接失败或 SQL 报错：

# View database logs ./tests/db/compose.sh logs --follow <database service>

<database service>可以是postgresql、mysql、mssql等数据库容器名，日志中可见服务启动、账号认证与客户端连接记录；排查问题时也可与测试容器的 pytest 输出对照分析。

小结：从 README 到源码的完整链路

回看 tests/db/README.md 的六段命令，其背后是一条清晰的验证链：

1. 构建（compose.sh build）→ 用 dev/extract_deps.py 拉取仓库真实依赖，与数据库驱动一起固化进镜像；
2. 运行（compose.sh run --rm）→ entrypoint.sh 以可编辑模式安装当前源码，MSSQL 额外等待并初始化数据库，随后执行 test_tracking_operations.py 等功能测试；
3. Schema 校验→ test_schema.py 对比 schemas 快照，捕捉 ORM 模型变更与方言反射差异；
4. 迁移校验→ check_migration.sh 用旧版写数据、新版跑mlflow db upgrade，再由 check_migration.py 验证数据完整性与工作区回填；
5. 清理（compose.sh down）→ 干净回收容器、卷与镜像。

这套体系既服务于 MLflow 开发者日常的数据库兼容性回归，也为后续新增数据库后端（如新的方言）提供了可直接套用的容器化测试模板——复制服务定义、补充 Schema 快照与迁移校验，即可把新后端纳入 CI 覆盖范围。

【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow