Migrate core MLflow integration from file store to SQLite #8894
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -26,7 +26,7 @@ | |
| from monai.fl.client.monai_algo import MonaiAlgo | ||
| from monai.fl.utils.constants import ExtraItems | ||
| from monai.fl.utils.exchange_object import ExchangeObject | ||
| from monai.utils import path_to_sqlite_uri | ||
| from tests.test_utils import SkipIfNoModule | ||
|
|
||
| _root_dir = Path(__file__).resolve().parents[2] | ||
|
|
@@ -79,7 +79,7 @@ | |
| "save_execute_config": f"{_data_dir}/config_executed.json", | ||
| "trainer": { | ||
| "_target_": "MLFlowHandler", | ||
| "tracking_uri": path_to_sqlite_uri(os.path.join(_data_dir, "mlflow_override.db")), | ||
| "output_transform": "$monai.handlers.from_engine(['loss'], first=True)", | ||
| "close_on_complete": True, | ||
| }, | ||
|
|
@@ -103,7 +103,7 @@ | |
| workflow_type="train", | ||
| logging_file=_logging_file, | ||
| tracking="mlflow", | ||
| tracking_uri=path_to_sqlite_uri(os.path.join(_data_dir, "mlflow_1.db")), | ||
| experiment_name="monai_eval1", | ||
| ), | ||
| "config_filters_filename": os.path.join(_data_dir, "config_fl_filters.json"), | ||
|
|
@@ -119,7 +119,7 @@ | |
| ], | ||
| "eval_kwargs": { | ||
| "tracking": "mlflow", | ||
| "tracking_uri": path_to_sqlite_uri(os.path.join(_data_dir, "mlflow_2.db")), | ||
| "experiment_name": "monai_eval2", | ||
| }, | ||
| "eval_workflow_name": "training", | ||
|
|
@@ -179,6 +179,36 @@ | |
| ] | ||
|
|
||
|
|
||
| def _dispose_sqlite_engines(): | ||
| """Dispose MLflow's open SQLAlchemy SQLite engines so the test ``.db`` files can be removed. | ||
|
|
||
| MLflow keeps a SQLite connection open for the lifetime of its client; on Windows that | ||
| locks the database file and breaks cleanup. ``MLFlowHandler.close()`` releases it, but a | ||
| workflow may finish without closing every handler, so dispose defensively here before | ||
| deleting the files. Scoped to the test's ``mlflow*.db`` backends so unrelated (e.g. | ||
| in-memory) sqlite engines elsewhere in the process are left untouched. | ||
| """ | ||
| import gc | ||
|
|
||
| try: | ||
| from sqlalchemy.engine import Engine | ||
| except ImportError: | ||
| return | ||
| gc.collect() | ||
| for obj in gc.get_objects(): | ||
| if not isinstance(obj, Engine): | ||
| continue | ||
| try: | ||
| url = obj.url | ||
| db = url.database if url.get_backend_name() == "sqlite" else None | ||
| # the test backends are all files named ``mlflow*.db``; match those only so | ||
| # unrelated (e.g. in-memory) sqlite engines in the process are left untouched. | ||
| if db and os.path.basename(db).startswith("mlflow"): | ||
| obj.dispose() | ||
| except Exception: | ||
| pass | ||
|
Comment on lines
+182
to
+211
Contributor
There was a problem hiding this comment. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
# Confirm SQLAlchemy is pinned/available and inspect the URL API.
fd -t f 'requirements.*\.txt|setup\.cfg|pyproject\.toml' | xargs rg -n -i 'sqlalchemy'Repository: Project-MONAI/MONAI Length of output: 47 🌐 Web query:
💡 Result: Yes, the SQLAlchemy URL object exposes both the get_backend_name method and the database attribute in both SQLAlchemy 1.4 and 2.0 [1][2]. Both were formalized as part of the URL object's overhaul in version 1.4, where the URL object was redesigned to be an immutable namedtuple-like structure [3][4]. The get_backend_name method returns the database backend name (the portion of the drivername to the left of the plus sign) [1][2], and the database attribute provides access to the database name specified in the URL [1][2]. Citations:
🏁 Script executed: #!/bin/bash
# Find where SQLAlchemy is declared/pinned in this repo.
rg -n --hidden --no-ignore-vcs -S "sqlalchemy" pyproject.toml setup.cfg setup.py requirements*.txt */requirements*.txt 2>/dev/null || trueRepository: Project-MONAI/MONAI Length of output: 45 🌐 Web query:
💡 Result: MONAI does not appear to pin SQLAlchemy versions. What I found in MONAI packaging/dependency metadata: - MONAI’s declared install requirements in setup.cfg list torch and numpy, plus many optional dependencies under extras; SQLAlchemy is not present in the shown install_requires or extras list [1][2]. - The MONAI PyPI page describes dependency versioning for major releases and points to requirements*.txt files, but the dependency table excerpt shown there does not include SQLAlchemy [3]. So there is no explicit “sqlalchemy version pin” I can cite from MONAI’s published dependency metadata for the monai/fl client algorithm component (monai_algo.py) either—monai_algo.py itself is code for FL client behavior and does not specify SQLAlchemy dependencies [4]. If you meant “MONAI FL (monai_algo) indirectly depends on SQLAlchemy via another library”, tell me which MONAI version you’re using (e.g., 1.3.0/1.4.0/1.5.x) and your install extras (e.g., monai[all]) and I can trace the transitive dependency chain from that exact lockfile/requirements output. Citations:
🏁 Script executed: #!/bin/bash
set -euo pipefail
# List likely dependency declaration files
fd -t f -i 'requirements*.txt|setup.cfg|pyproject.toml|setup.py' . || true
# Search for sqlalchemy/sqlalchemy in those files
for f in $(fd -t f -i 'requirements*.txt|setup.cfg|pyproject.toml|setup.py' .); do
rg -n -i 'sqlalchemy' "$f" || true
doneRepository: Project-MONAI/MONAI Length of output: 146 🏁 Script executed: #!/bin/bash
set -euo pipefail
# Confirm there aren't other dependency declaration files referencing sqlalchemy.
rg -n -i --hidden --no-ignore-vcs "sqlalchemy" \
-g 'requirements*.txt' -g 'setup.cfg' -g 'pyproject.toml' -g 'setup.py' -g 'Pipfile*' -g '*.yml' -g '*.yaml' \
. | head -n 50Repository: Project-MONAI/MONAI Length of output: 45 Narrow the failure mode of
🧰 Tools🪛 Ruff (0.15.15)[error] 208-209: (S110) [warning] 208-208: Do not catch blind exception: (BLE001) 🤖 Prompt for AI Agents |
||
|
|
||
|
|
||
| @SkipIfNoModule("ignite") | ||
| @SkipIfNoModule("mlflow") | ||
| class TestFLMonaiAlgo(unittest.TestCase): | ||
|
|
@@ -202,8 +232,11 @@ def test_train(self, input_params): | |
|
|
||
| # test experiment management | ||
| if "save_execute_config" in algo.train_workflow.parser: | ||
| _dispose_sqlite_engines() # release SQLite handles so the db file can be removed on Windows | ||
| self.assertTrue(os.path.exists(f"{_data_dir}/mlflow_override.db")) | ||
| os.remove(f"{_data_dir}/mlflow_override.db") | ||
| if os.path.isdir(f"{_data_dir}/mlruns"): | ||
| shutil.rmtree(f"{_data_dir}/mlruns") | ||
| self.assertTrue(os.path.exists(f"{_data_dir}/config_executed.json")) | ||
| os.remove(f"{_data_dir}/config_executed.json") | ||
|
|
||
|
|
@@ -225,9 +258,12 @@ def test_evaluate(self, input_params): | |
|
|
||
| # test experiment management | ||
| if "save_execute_config" in algo.eval_workflow.parser: | ||
| _dispose_sqlite_engines() # release SQLite handles so the db files can be removed on Windows | ||
| self.assertGreater(len(list(glob.glob(f"{_data_dir}/mlflow_*"))), 0) | ||
| for f in list(glob.glob(f"{_data_dir}/mlflow_*")): | ||
| shutil.rmtree(f) if os.path.isdir(f) else os.remove(f) | ||
| if os.path.isdir(f"{_data_dir}/mlruns"): | ||
| shutil.rmtree(f"{_data_dir}/mlruns") | ||
| self.assertGreater(len(list(glob.glob(f"{_data_dir}/eval/config_*"))), 0) | ||
| for f in list(glob.glob(f"{_data_dir}/eval/config_*")): | ||
| os.remove(f) | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.