From d1018350a52e7b57576754a75f67e3d8ddd48d34 Mon Sep 17 00:00:00 2001 From: "R. Garcia-Dias" Date: Thu, 30 Jul 2026 17:13:15 +0100 Subject: [PATCH 1/5] fix: address Dependabot alerts for mlflow, transformers, setuptools - mlflow: bump floor to >=3.14.0, closing CVEs across the 2.x/early-3.x line. mlflow>=3.13 hard-errors on the local file-store backend (#8891); rather than opting back into the file store, MLFlowHandler now defaults to a local SQLite backend: when no tracking_uri is given and MLFLOW_TRACKING_URI is unset it uses sqlite:////mlruns.db with artifacts under /mlruns. Adds monai.utils.path_to_sqlite_uri and an artifact_location argument, updates the bundle defaults, and releases the SQLite engine on close. - transformers: bump floor to >=5.5.0, closing two HIGH severity CVEs. MultiModal now builds a real transformers.BertConfig with _attn_implementation="eager" for transformers>=4.48's attention dispatch. - setuptools: bump the build-system floor to >=78.1.1, closing one HIGH severity CVE. The old <=79.0.1 cap (#8439) is no longer needed since the legacy fetch_build_eggs CLI usage is gone from setup.py. The torch.float8_e8m0fnu NGC-image concern behind the old transformers<5.0 cap should be re-verified against the current NGC base image before merging. Signed-off-by: R. Garcia-Dias --- docs/requirements.txt | 4 +- monai/bundle/scripts.py | 3 +- monai/bundle/utils.py | 6 +- monai/handlers/mlflow_handler.py | 95 ++++++++++++- monai/networks/nets/transchex.py | 7 +- monai/utils/__init__.py | 1 + monai/utils/misc.py | 16 +++ pyproject.toml | 10 +- tests/fl/monai_algo/test_fl_monai_algo.py | 52 ++++++- tests/handlers/test_handler_mlflow.py | 129 ++++++++++++++++-- .../test_integration_bundle_run.py | 11 +- 11 files changed, 295 insertions(+), 39 deletions(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index 9e023cec5eb..d49f659c680 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -20,8 +20,8 @@ sphinxcontrib-serializinghtml sphinx-autodoc-typehints==1.11.1 pandas einops -transformers>=4.53.0 -mlflow>=2.12.2,<3.13 +transformers>=5.5.0 +mlflow>=3.14.0 clearml>=1.10.0rc0 tensorboardX imagecodecs; platform_system == "Linux" or platform_system == "Darwin" diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 1007cedc1ec..c285c8b3ab7 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -997,7 +997,8 @@ def run( common parameters shown below will be added and can be passed through the `override` parameter of this method. - ``"output_dir"``: the path to save mlflow tracking outputs locally, default to "/eval". - - ``"tracking_uri"``: uri to save mlflow tracking outputs, default to "/output_dir/mlruns". + - ``"tracking_uri"``: uri to save mlflow tracking outputs, default to a local SQLite database + at "/mlruns.db" with run artifacts kept under "/mlruns". - ``"experiment_name"``: experiment name for this run, default to "monai_experiment". - ``"run_name"``: the name of current run. - ``"save_execute_config"``: whether to save the executed config files. It can be `False`, `/path/to/artifacts` diff --git a/monai/bundle/utils.py b/monai/bundle/utils.py index 81f76d0435b..ebd521a93f1 100644 --- a/monai/bundle/utils.py +++ b/monai/bundle/utils.py @@ -116,8 +116,10 @@ "configs": { # if no "output_dir" in the bundle config, default to "/eval" "output_dir": "$@bundle_root + '/eval'", - # use URI to support linux, mac and windows os - "tracking_uri": "$monai.utils.path_to_uri(@output_dir) + '/mlruns'", + # MLflow 3.13+ rejects the filesystem (file store) tracking backend, so default tracking + # to a local SQLite database. The handler keeps run artifacts under "/mlruns" + # (next to the db). A URI is used so the path is valid on linux, mac and windows os. + "tracking_uri": "$monai.utils.path_to_sqlite_uri(@output_dir + '/mlruns.db')", "experiment_name": "monai_experiment", "run_name": None, # may fill it at runtime diff --git a/monai/handlers/mlflow_handler.py b/monai/handlers/mlflow_handler.py index 2ea54cc06ab..7560823d76d 100644 --- a/monai/handlers/mlflow_handler.py +++ b/monai/handlers/mlflow_handler.py @@ -22,7 +22,16 @@ from torch.utils.data import Dataset from monai.apps.utils import get_logger -from monai.utils import CommonKeys, IgniteInfo, ensure_tuple, flatten_dict, min_version, optional_import +from monai.utils import ( + CommonKeys, + IgniteInfo, + ensure_tuple, + flatten_dict, + min_version, + optional_import, + path_to_sqlite_uri, + path_to_uri, +) Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events") mlflow, _ = optional_import("mlflow", descriptor="Please install mlflow before using MLFlowHandler.") @@ -68,7 +77,14 @@ class MLFlowHandler: tracking_uri: connects to a tracking URI. can also set the `MLFLOW_TRACKING_URI` environment variable to have MLflow find a URI from there. in both cases, the URI can either be an HTTP/HTTPS URI for a remote server, a database connection string, or a local path - to log data to a directory. The URI defaults to path `mlruns`. + to log data to a directory. When no ``tracking_uri`` is provided and the + ``MLFLOW_TRACKING_URI`` environment variable is unset, the handler now + defaults to a local SQLite database backend at ``sqlite:////mlruns.db`` with + artifacts stored under ``/mlruns``. The default was changed from the filesystem + (file store) backend because MLflow 3.13+ raises an exception for the file store unless + ``MLFLOW_ALLOW_FILE_STORE=true`` is set; SQLite is the backend MLflow recommends and it + does not raise. Any explicitly provided ``tracking_uri`` (including a local file path or + ``file://`` URI) is passed through unchanged. for more details: https://mlflow.org/docs/latest/python_api/mlflow.html#mlflow.set_tracking_uri. iteration_log: whether to log data to MLFlow when iteration completed, default to `True`. ``iteration_log`` can be also a function and it will be interpreted as an event filter @@ -113,6 +129,11 @@ class MLFlowHandler: optimizer_param_names: parameter names in the optimizer that need to be recorded during running the workflow, default to `'lr'`. close_on_complete: whether to close the mlflow run in `complete` phase in workflow, default to False. + artifact_location: the location to store run artifacts in, passed to MLflow when the experiment is + created. When ``None`` and a local SQLite backend is used (from the ``tracking_uri`` argument + or the ``MLFLOW_TRACKING_URI`` environment variable), it defaults to an ``mlruns`` directory + next to the database file; for other backends ``None`` lets MLflow decide based on the + ``tracking_uri``. Has no effect if the experiment already exists. For more details of MLFlow usage, please refer to: https://mlflow.org/docs/latest/index.html. @@ -141,6 +162,7 @@ def __init__( artifacts: str | Sequence[Path] | None = None, optimizer_param_names: str | Sequence[str] = "lr", close_on_complete: bool = False, + artifact_location: str | None = None, ) -> None: self.iteration_log = iteration_log self.epoch_log = epoch_log @@ -156,6 +178,31 @@ def __init__( self.experiment_param = experiment_param self.artifacts = ensure_tuple(artifacts) self.optimizer_param_names = ensure_tuple(optimizer_param_names) + # When no tracking_uri is provided, default to a local SQLite backend instead of the + # filesystem (file store) backend. MLflow 3.13+ raises for the file store unless + # `MLFLOW_ALLOW_FILE_STORE=true` is set, while SQLite is the recommended backend and does + # not raise. Artifacts cannot live inside a database, so by default they are stored under + # the `./mlruns` directory (where the previous file store default kept them) via the + # experiment `artifact_location`. Any explicitly provided tracking_uri is left unchanged. + self.artifact_location = artifact_location + # Resolve the effective tracking URI from the argument or the `MLFLOW_TRACKING_URI` + # environment variable, so both configure the artifact location the same way. + effective_tracking_uri = tracking_uri or os.environ.get("MLFLOW_TRACKING_URI") + # When neither is set, fall back to the local SQLite default described above. + if not effective_tracking_uri: + tracking_uri = effective_tracking_uri = path_to_sqlite_uri(os.path.join(os.getcwd(), "mlruns.db")) + # For a local SQLite backend, keep run artifacts in an `mlruns` directory next to the + # database file (mirroring the previous file-store layout) unless the caller set + # `artifact_location`. Other backends (e.g. a remote server) are left to MLflow to decide. + # Only `tracking_uri` is passed to the client, so an `MLFLOW_TRACKING_URI` env var is + # still resolved by MLflow itself. + if ( + self.artifact_location is None + and effective_tracking_uri + and effective_tracking_uri.startswith("sqlite:///") + ): + db_path = Path(effective_tracking_uri[len("sqlite:///") :]) + self.artifact_location = path_to_uri(db_path.parent / "mlruns") self.client = mlflow.MlflowClient(tracking_uri=tracking_uri if tracking_uri else None) self.run_finish_status = mlflow.entities.RunStatus.to_string(mlflow.entities.RunStatus.FINISHED) self.close_on_complete = close_on_complete @@ -246,7 +293,12 @@ def _set_experiment(self): try: experiment = self.client.get_experiment_by_name(self.experiment_name) if not experiment: - experiment_id = self.client.create_experiment(self.experiment_name) + # pass an explicit artifact_location (set for the default SQLite backend, or + # by the caller) so artifacts land in the intended directory; when it is + # None MLflow decides based on the tracking_uri. + experiment_id = self.client.create_experiment( + self.experiment_name, artifact_location=self.artifact_location + ) experiment = self.client.get_experiment(experiment_id) break except MlflowException as e: @@ -338,14 +390,43 @@ def complete(self) -> None: for artifact in artifact_list: self.client.log_artifact(self.cur_run.info.run_id, artifact) + def _dispose_sqlite_store(self) -> None: + """ + Release MLflow's SQLAlchemy engine when a local SQLite tracking backend is used. + + MLflow keeps the SQLite connection open for the lifetime of the client, which on + Windows prevents the database file from being deleted. MLflow exposes no public + client close/dispose API, so this reaches into its internals defensively to release + the engine. It is a no-op for non-SQLite backends. + """ + tracking_uri = getattr(self.client, "tracking_uri", "") + if not isinstance(tracking_uri, str) or not tracking_uri.startswith("sqlite:"): + return + store = getattr(getattr(self.client, "_tracking_client", None), "store", None) + if store is None: + return + dispose = getattr(store, "_dispose_engine", None) + if callable(dispose): + dispose() + else: + engine = getattr(store, "engine", None) + if engine is not None: + engine.dispose() + read_engine = getattr(store, "read_engine", None) + if read_engine is not None: + read_engine.dispose() + def close(self) -> None: """ - Stop current running logger of MLFlow. + Stop current running logger of MLFlow and release local SQLite resources. """ - if self.cur_run: - self.client.set_terminated(self.cur_run.info.run_id, self.run_finish_status) - self.cur_run = None + try: + if self.cur_run: + self.client.set_terminated(self.cur_run.info.run_id, self.run_finish_status) + self.cur_run = None + finally: + self._dispose_sqlite_store() def epoch_completed(self, engine: Engine) -> None: """ diff --git a/monai/networks/nets/transchex.py b/monai/networks/nets/transchex.py index dfe71dd4bad..cd71b307422 100644 --- a/monai/networks/nets/transchex.py +++ b/monai/networks/nets/transchex.py @@ -23,6 +23,7 @@ transformers = optional_import("transformers") load_tf_weights_in_bert = optional_import("transformers", name="load_tf_weights_in_bert")[0] cached_file = optional_import("transformers.utils", name="cached_file")[0] +BertConfig = optional_import("transformers", name="BertConfig")[0] BertEmbeddings = optional_import("transformers.models.bert.modeling_bert", name="BertEmbeddings")[0] BertLayer = optional_import("transformers.models.bert.modeling_bert", name="BertLayer")[0] @@ -222,7 +223,11 @@ def __init__( """ super().__init__() - self.config = type("obj", (object,), bert_config) + self.config = BertConfig(**bert_config) + # explicitly select the eager attention path: transformers>=4.48 dispatches attention + # implementations via `config._attn_implementation`, which is otherwise left unset since + # `bert_config` above does not come from a `from_pretrained` call. + self.config._attn_implementation = "eager" self.embeddings = BertEmbeddings(self.config) self.language_encoder = nn.ModuleList([BertLayer(self.config) for _ in range(num_language_layers)]) self.vision_encoder = nn.ModuleList([BertLayer(self.config) for _ in range(num_vision_layers)]) diff --git a/monai/utils/__init__.py b/monai/utils/__init__.py index d1a705205ce..944f34aab77 100644 --- a/monai/utils/__init__.py +++ b/monai/utils/__init__.py @@ -89,6 +89,7 @@ is_sqrt, issequenceiterable, list_to_dict, + path_to_sqlite_uri, path_to_uri, pprint_edges, progress_bar, diff --git a/monai/utils/misc.py b/monai/utils/misc.py index 10f9c774439..ea9cf753e89 100644 --- a/monai/utils/misc.py +++ b/monai/utils/misc.py @@ -69,6 +69,7 @@ "save_obj", "label_union", "path_to_uri", + "path_to_sqlite_uri", "pprint_edges", "check_key_duplicates", "CheckKeyDuplicatesYamlLoader", @@ -727,6 +728,21 @@ def path_to_uri(path: PathLike) -> str: return Path(path).absolute().as_uri() +def path_to_sqlite_uri(path: PathLike) -> str: + """ + Convert a database file path to a SQLite connection URI, e.g. for use as an MLflow + ``tracking_uri``. If not an absolute path, it is converted to an absolute path first. + + A forward-slash (POSIX) path is used so the URI is valid on Windows as well as POSIX: + on Windows this yields ``sqlite:///C:/path/db.sqlite`` and on POSIX ``sqlite:////path/db.sqlite``. + + Args: + path: input database file path, can be a string or `Path` object. + + """ + return f"sqlite:///{Path(path).absolute().as_posix()}" + + def pprint_edges(val: Any, n_lines: int = 20) -> str: """ Pretty print the head and tail ``n_lines`` of ``val``, and omit the middle part if the part has more than 3 lines. diff --git a/pyproject.toml b/pyproject.toml index 9c5f8922835..830ced4d241 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [build-system] requires = [ - "setuptools", + "setuptools>=78.1.1", "wheel", "versioneer[toml]", "more-itertools>=8.0", @@ -71,7 +71,7 @@ all = [ "lpips==0.1.4", "matplotlib>=3.6.3", "MetricsReloaded @ git+https://github.com/Project-MONAI/MetricsReloaded@monai-support", - "mlflow>=2.12.2,<3.13", + "mlflow>=3.14.0", "nibabel", "ninja", "nni; platform_system == 'Linux' and 'arm' not in platform_machine and 'aarch' not in platform_machine", @@ -103,7 +103,7 @@ all = [ "torchio", "torchvision", "tqdm>=4.47.0", - "transformers>=4.53.0, <5.0", + "transformers>=5.5.0", "zarr" ] clearml = ["clearml>=1.10.0rc0"] @@ -126,7 +126,7 @@ lmdb = ["lmdb"] lpips = ["lpips==0.1.4"] matplotlib = ["matplotlib>=3.6.3"] metrics_reloaded = ["MetricsReloaded @ git+https://github.com/Project-MONAI/MetricsReloaded@monai-support"] -mlflow = ["mlflow>=2.12.2,<3.13"] +mlflow = ["mlflow>=3.14.0"] nibabel = ["nibabel"] nni = [ "nni; platform_system == 'Linux' and 'arm' not in platform_machine and 'aarch' not in platform_machine", @@ -155,7 +155,7 @@ tifffile = ["tifffile; platform_system == 'Linux' or platform_system == 'Darwin' torchio = ["torchio"] torchvision = ["torchvision"] tqdm = ["tqdm>=4.47.0"] -transformers = ["transformers>=4.53.0, <5.0"] # 5.x references torch.float8_e8m0fnu absent in older PyTorch builds +transformers = ["transformers>=5.5.0"] # 5.x needs the transchex BertLayer/BertConfig updates; re-verify the NGC image float8 concern zarr = ["zarr"] # these dependencies are for testing/building only, they aren't needed for regular use so don't appear in "all" testing = [ diff --git a/tests/fl/monai_algo/test_fl_monai_algo.py b/tests/fl/monai_algo/test_fl_monai_algo.py index 2c1a8488cc0..a55dcc45605 100644 --- a/tests/fl/monai_algo/test_fl_monai_algo.py +++ b/tests/fl/monai_algo/test_fl_monai_algo.py @@ -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_uri +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_uri(_data_dir) + "/mlflow_override", + "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_uri(_data_dir) + "/mlflow_1", + 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_uri(_data_dir) + "/mlflow_2", + "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,38 @@ ] +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(): + # gc.get_objects() can include dead weakref proxies, whose isinstance() raises + # ReferenceError, so guard the whole inspection (ReferenceError is an Exception). + try: + if not isinstance(obj, Engine): + continue + 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 + + @SkipIfNoModule("ignite") @SkipIfNoModule("mlflow") class TestFLMonaiAlgo(unittest.TestCase): @@ -202,8 +234,11 @@ def test_train(self, input_params): # test experiment management if "save_execute_config" in algo.train_workflow.parser: - self.assertTrue(os.path.exists(f"{_data_dir}/mlflow_override")) - shutil.rmtree(f"{_data_dir}/mlflow_override") + _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 +260,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) + 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) diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py index 80630e6f5a2..ab94d908ae4 100644 --- a/tests/handlers/test_handler_mlflow.py +++ b/tests/handlers/test_handler_mlflow.py @@ -17,7 +17,7 @@ import tempfile import unittest from concurrent.futures import ThreadPoolExecutor -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import numpy as np from ignite.engine import Engine, Events @@ -26,7 +26,7 @@ from monai.apps import download_and_extract from monai.bundle import ConfigWorkflow, download from monai.handlers import MLFlowHandler -from monai.utils import optional_import, path_to_uri +from monai.utils import optional_import, path_to_sqlite_uri, path_to_uri from tests.test_utils import skip_if_downloading_fails, skip_if_quick _, has_dataset_tracking = optional_import("mlflow", "2.4.0") @@ -55,7 +55,7 @@ def _train_func(engine, batch): handler = MLFlowHandler( iteration_log=False, epoch_log=True, - tracking_uri=path_to_uri(test_path), + tracking_uri=path_to_sqlite_uri(test_path), state_attributes=["test"], close_on_complete=True, ) @@ -95,7 +95,7 @@ def _update_metric(engine): handler = MLFlowHandler( iteration_log=False, epoch_log=True, - tracking_uri=path_to_uri(test_path), + tracking_uri=path_to_sqlite_uri(test_path), state_attributes=["test"], close_on_complete=True, ) @@ -106,6 +106,117 @@ def _update_metric(engine): # the run count should equal to the times of creating engine self.assertEqual(create_engine_times, run_cnt) + def test_default_tracking_uri_is_sqlite(self): + # when no tracking_uri is provided, the handler should default to a local SQLite backend + # rather than the filesystem (file store) backend, which raises on mlflow 3.13+. + with tempfile.TemporaryDirectory() as tempdir: + cwd = os.getcwd() + os.chdir(tempdir) + handler = None + try: + handler = MLFlowHandler(iteration_log=False, epoch_log=False) + self.assertTrue(handler.client.tracking_uri.startswith("sqlite:///")) + self.assertTrue(handler.client.tracking_uri.endswith("mlruns.db")) + # artifacts should still default to a `./mlruns`-style directory + self.assertIsNotNone(handler.artifact_location) + self.assertTrue(handler.artifact_location.endswith("mlruns")) + finally: + if handler is not None: + handler.close() # release the SQLite handle so Windows can delete the db + os.chdir(cwd) + + def test_explicit_tracking_uri_is_preserved(self): + # an explicitly provided tracking_uri must be passed through unchanged, including file paths. + with tempfile.TemporaryDirectory() as tempdir: + explicit_uri = path_to_uri(os.path.join(tempdir, "mlflow_explicit")) + handler = MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri=explicit_uri) + self.assertEqual(handler.client.tracking_uri, explicit_uri) + self.assertIsNone(handler.artifact_location) + + def test_remote_tracking_uri_leaves_artifact_location_unset(self): + # a non-local (e.g. remote) tracking_uri must not get a local artifact_location injected, + # so the remote backend keeps deciding where artifacts go. + handler = MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri="http://localhost:5000") + self.assertEqual(handler.client.tracking_uri, "http://localhost:5000") + self.assertIsNone(handler.artifact_location) + + def test_explicit_sqlite_tracking_uri_colocates_artifacts(self): + # an explicit local SQLite tracking_uri should still co-locate artifacts next to the db. + with tempfile.TemporaryDirectory() as tempdir: + uri = path_to_sqlite_uri(os.path.join(tempdir, "sub", "mlruns.db")) + handler = MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri=uri) + try: + self.assertEqual(handler.client.tracking_uri, uri) + self.assertIsNotNone(handler.artifact_location) + self.assertTrue(handler.artifact_location.endswith("mlruns")) + finally: + handler.close() # release the SQLite handle so Windows can delete the db + + def test_env_var_sqlite_tracking_uri_colocates_artifacts(self): + # a SQLite `MLFLOW_TRACKING_URI` env var should co-locate artifacts next to the db, the + # same as an explicit `tracking_uri` argument. The env var itself is left for MLflow to + # resolve, so the handler does not pass it to the client. + with tempfile.TemporaryDirectory() as tempdir: + uri = path_to_sqlite_uri(os.path.join(tempdir, "sub", "mlruns.db")) + handler = None + with patch.dict(os.environ, {"MLFLOW_TRACKING_URI": uri}): + try: + handler = MLFlowHandler(iteration_log=False, epoch_log=False) + self.assertTrue(handler.client.tracking_uri.endswith("mlruns.db")) + self.assertIsNotNone(handler.artifact_location) + self.assertTrue(handler.artifact_location.endswith("mlruns")) + # co-located with the db file (the `sub` dir), not a cwd-relative `./mlruns` + self.assertIn("sub", handler.artifact_location) + finally: + if handler is not None: + handler.close() # release the SQLite handle so Windows can delete the db + + def test_explicit_artifact_location_is_used(self): + # an explicitly provided artifact_location should be kept even with the default SQLite + # backend, so callers (e.g. the bundle defaults) can co-locate artifacts with the db. + with tempfile.TemporaryDirectory() as tempdir: + cwd = os.getcwd() + os.chdir(tempdir) + handler = None + try: + art = path_to_uri(os.path.join(tempdir, "artifacts")) + handler = MLFlowHandler(iteration_log=False, epoch_log=False, artifact_location=art) + self.assertEqual(handler.artifact_location, art) + finally: + if handler is not None: + handler.close() # release the SQLite handle so Windows can delete the db + os.chdir(cwd) + + def test_default_sqlite_run_flow(self): + # a basic log/run flow should work with the default SQLite backend (no tracking_uri given). + with tempfile.TemporaryDirectory() as tempdir: + cwd = os.getcwd() + os.chdir(tempdir) + try: + + def _train_func(engine, batch): + return [batch + 1.0] + + engine = Engine(_train_func) + + @engine.on(Events.EPOCH_COMPLETED) + def _update_metric(engine): + current_metric = engine.state.metrics.get("acc", 0.1) + engine.state.metrics["acc"] = current_metric + 0.1 + + # close_on_complete=False so cur_run stays available after the run for the metric + # check below; the run is closed explicitly afterwards. + handler = MLFlowHandler(iteration_log=False, epoch_log=True, close_on_complete=False) + handler.attach(engine) + engine.run(range(3), max_epochs=2) + cur_run = handler.client.get_run(handler.cur_run.info.run_id) + self.assertTrue("acc" in cur_run.data.metrics.keys()) + handler.close() + # the default backend should have created a SQLite database file in the cwd + self.assertTrue(os.path.exists(os.path.join(tempdir, "mlruns.db"))) + finally: + os.chdir(cwd) + def test_metrics_track(self): experiment_param = {"backbone": "efficientnet_b0"} with tempfile.TemporaryDirectory() as tempdir: @@ -137,7 +248,7 @@ def _update_metric(engine): handler = MLFlowHandler( iteration_log=False, epoch_log=True, - tracking_uri=path_to_uri(test_path), + tracking_uri=path_to_sqlite_uri(test_path), state_attributes=["test"], experiment_param=experiment_param, artifacts=[artifact_path], @@ -173,7 +284,7 @@ def _update_metric(engine): handler = MLFlowHandler( iteration_log=False, epoch_log=epoch_log, - tracking_uri=path_to_uri(test_path), + tracking_uri=path_to_sqlite_uri(test_path), state_attributes=["test"], experiment_param=experiment_param, close_on_complete=True, @@ -212,7 +323,7 @@ def _update_metric(engine): handler = MLFlowHandler( iteration_log=iteration_log, epoch_log=False, - tracking_uri=path_to_uri(test_path), + tracking_uri=path_to_sqlite_uri(test_path), state_attributes=["test"], experiment_param=experiment_param, close_on_complete=True, @@ -271,7 +382,7 @@ def test_dataset_tracking(self): final_id="finalize", ) - tracking_path = os.path.join(bundle_root, "eval") + tracking_path = os.path.join(tempdir, "mlflow_dataset.db") workflow.bundle_root = bundle_root workflow.dataset_dir = data_dir workflow.initialize() @@ -280,7 +391,7 @@ def test_dataset_tracking(self): iteration_log=False, epoch_log=False, dataset_dict={"test": infer_dataset}, - tracking_uri=path_to_uri(tracking_path), + tracking_uri=path_to_sqlite_uri(tracking_path), ) mlflow_handler.attach(workflow.evaluator) workflow.run() diff --git a/tests/integration/test_integration_bundle_run.py b/tests/integration/test_integration_bundle_run.py index 7f366d47458..67f44562593 100644 --- a/tests/integration/test_integration_bundle_run.py +++ b/tests/integration/test_integration_bundle_run.py @@ -29,7 +29,7 @@ from monai.bundle import ConfigParser from monai.bundle.utils import DEFAULT_HANDLERS_ID from monai.transforms import LoadImage -from monai.utils import path_to_uri +from monai.utils import path_to_sqlite_uri from tests.test_utils import command_line_tests TESTS_PATH = Path(__file__).parents[1] @@ -175,7 +175,7 @@ def test_shape(self, config_file, expected_shape): "no_epoch": True, # test override config in the settings file "evaluator": { "_target_": "MLFlowHandler", - "tracking_uri": "$monai.utils.path_to_uri(@output_dir) + '/mlflow_override1'", + "tracking_uri": "$monai.utils.path_to_sqlite_uri(@output_dir + '/mlflow_override1.db')", "iteration_log": "@no_epoch", }, }, @@ -208,16 +208,17 @@ def test_shape(self, config_file, expected_shape): command_line_tests(la + ["--args_file", def_args_file] + ["--tracking", settings_file]) loader = LoadImage(image_only=True) self.assertTupleEqual(loader(os.path.join(tempdir, "image", "image_seg.nii.gz")).shape, expected_shape) - self.assertTrue(os.path.exists(f"{tempdir}/mlflow_override1")) + self.assertTrue(os.path.exists(f"{tempdir}/mlflow_override1.db")) - tracking_uri = path_to_uri(tempdir) + "/mlflow_override2" # test override experiment management configs + # test override experiment management configs + tracking_uri = path_to_sqlite_uri(os.path.join(tempdir, "mlflow_override2.db")) # here test the script with `google fire` tool as CLI cmd = "-m fire monai.bundle.scripts run --tracking mlflow --evaluator#amp False" cmd += f" --tracking_uri {tracking_uri} {override} --output_dir {tempdir} --device {device}" la = ["coverage", "run"] + cmd.split(" ") + ["--meta_file", meta_file] + ["--config_file", config_file] command_line_tests(la) self.assertTupleEqual(loader(os.path.join(tempdir, "image", "image_trans.nii.gz")).shape, expected_shape) - self.assertTrue(os.path.exists(f"{tempdir}/mlflow_override2")) + self.assertTrue(os.path.exists(f"{tempdir}/mlflow_override2.db")) # test the saved execution configs self.assertTrue(len(glob(f"{tempdir}/config_*.json")), 2) From 901e04660134ce10b4cf38e0b77883923976b5c4 Mon Sep 17 00:00:00 2001 From: "R. Garcia-Dias" Date: Fri, 28 Aug 2026 14:37:03 +0100 Subject: [PATCH 2/5] fix(test): drop file-store tracking URI test and handle db-file cleanup mlflow>=3.14 eagerly resolves the tracking store on MlflowClient construction and rejects the filesystem (file store) backend, so passing an explicit file:// tracking_uri now raises instead of being preserved. Remove that test; explicit-URI passthrough is still covered by the remote and SQLite tests. Also make tearDown remove SQLite db files (not just directories), fixing NotADirectoryError on Windows when the SQLite default backend creates a db file instead of a directory. Signed-off-by: R. Garcia-Dias --- tests/handlers/test_handler_mlflow.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py index ab94d908ae4..6ce89f83da8 100644 --- a/tests/handlers/test_handler_mlflow.py +++ b/tests/handlers/test_handler_mlflow.py @@ -71,7 +71,8 @@ def setUp(self): def tearDown(self): for tmpdir in self.tmpdir_list: if tmpdir and os.path.exists(tmpdir): - shutil.rmtree(tmpdir) + # the SQLite default backend creates a db file rather than a directory + shutil.rmtree(tmpdir) if os.path.isdir(tmpdir) else os.remove(tmpdir) def test_multi_run(self): with tempfile.TemporaryDirectory() as tempdir: @@ -125,14 +126,6 @@ def test_default_tracking_uri_is_sqlite(self): handler.close() # release the SQLite handle so Windows can delete the db os.chdir(cwd) - def test_explicit_tracking_uri_is_preserved(self): - # an explicitly provided tracking_uri must be passed through unchanged, including file paths. - with tempfile.TemporaryDirectory() as tempdir: - explicit_uri = path_to_uri(os.path.join(tempdir, "mlflow_explicit")) - handler = MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri=explicit_uri) - self.assertEqual(handler.client.tracking_uri, explicit_uri) - self.assertIsNone(handler.artifact_location) - def test_remote_tracking_uri_leaves_artifact_location_unset(self): # a non-local (e.g. remote) tracking_uri must not get a local artifact_location injected, # so the remote backend keeps deciding where artifacts go. From 1395e5874ccfe97751e821a65567e3d851b7bb0c Mon Sep 17 00:00:00 2001 From: "R. Garcia-Dias" Date: Fri, 28 Aug 2026 15:09:31 +0100 Subject: [PATCH 3/5] fix: address PR #9032 review feedback - pyproject.toml, docs/requirements.txt: raise mlflow floor to >=3.15.2, closing the unauthenticated webhook SSRF fixed in 3.15.0 - monai/handlers/mlflow_handler.py: reject file-store tracking URIs (local paths and file://) with an actionable error instead of the opaque MlflowException mlflow>=3.14 raises at client construction - monai/utils/misc.py: percent-encode URI-special characters in path_to_sqlite_uri so ?/# in a path are not misparsed by SQLAlchemy - tests: add docstrings to new handler tests, fix the tearDown temp-dir leak, drop the obsolete mlflow>=2.4.0 guard, and cover file-store rejection plus sqlite-URI escaping Signed-off-by: R. Garcia-Dias --- docs/requirements.txt | 2 +- monai/handlers/mlflow_handler.py | 15 +++++++++-- monai/utils/misc.py | 5 +++- pyproject.toml | 4 +-- tests/handlers/test_handler_mlflow.py | 33 +++++++++++------------ tests/utils/misc/test_monai_utils_misc.py | 17 +++++++++++- 6 files changed, 52 insertions(+), 24 deletions(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index d49f659c680..6598722dcab 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -21,7 +21,7 @@ sphinx-autodoc-typehints==1.11.1 pandas einops transformers>=5.5.0 -mlflow>=3.14.0 +mlflow>=3.15.2 clearml>=1.10.0rc0 tensorboardX imagecodecs; platform_system == "Linux" or platform_system == "Darwin" diff --git a/monai/handlers/mlflow_handler.py b/monai/handlers/mlflow_handler.py index 7560823d76d..0e9604d16cd 100644 --- a/monai/handlers/mlflow_handler.py +++ b/monai/handlers/mlflow_handler.py @@ -83,8 +83,9 @@ class MLFlowHandler: artifacts stored under ``/mlruns``. The default was changed from the filesystem (file store) backend because MLflow 3.13+ raises an exception for the file store unless ``MLFLOW_ALLOW_FILE_STORE=true`` is set; SQLite is the backend MLflow recommends and it - does not raise. Any explicitly provided ``tracking_uri`` (including a local file path or - ``file://`` URI) is passed through unchanged. + does not raise. Any explicitly provided ``tracking_uri`` is passed through unchanged, + except local file paths and ``file://`` URIs, which are rejected because MLflow no + longer supports the filesystem (file store) tracking backend. for more details: https://mlflow.org/docs/latest/python_api/mlflow.html#mlflow.set_tracking_uri. iteration_log: whether to log data to MLFlow when iteration completed, default to `True`. ``iteration_log`` can be also a function and it will be interpreted as an event filter @@ -203,6 +204,16 @@ def __init__( ): db_path = Path(effective_tracking_uri[len("sqlite:///") :]) self.artifact_location = path_to_uri(db_path.parent / "mlruns") + # MLflow 3.13+ refuses the filesystem (file store) tracking backend, and 3.14+ resolves + # the store eagerly at client construction, so a local path or ``file://`` URI would raise + # an opaque MlflowException. Reject those here with an actionable message instead. + if effective_tracking_uri.startswith("file://") or "://" not in effective_tracking_uri: + raise ValueError( + "MLflow no longer supports the filesystem (file store) tracking backend; got " + f"tracking_uri={effective_tracking_uri!r}. Use a SQLite URI " + "(sqlite:////mlruns.db) or a remote tracking URI, or set " + "MLFLOW_ALLOW_FILE_STORE=true to opt back in." + ) self.client = mlflow.MlflowClient(tracking_uri=tracking_uri if tracking_uri else None) self.run_finish_status = mlflow.entities.RunStatus.to_string(mlflow.entities.RunStatus.FINISHED) self.close_on_complete = close_on_complete diff --git a/monai/utils/misc.py b/monai/utils/misc.py index ea9cf753e89..f11e070f9df 100644 --- a/monai/utils/misc.py +++ b/monai/utils/misc.py @@ -27,6 +27,7 @@ from math import log10 from pathlib import Path from typing import TYPE_CHECKING, Any, TypeVar, cast, overload +from urllib.parse import quote import numpy as np import torch @@ -735,12 +736,14 @@ def path_to_sqlite_uri(path: PathLike) -> str: A forward-slash (POSIX) path is used so the URI is valid on Windows as well as POSIX: on Windows this yields ``sqlite:///C:/path/db.sqlite`` and on POSIX ``sqlite:////path/db.sqlite``. + URI-special characters in the path (e.g. ``?``, ``#``) are percent-encoded so they are not + misparsed as query/fragment components by SQLAlchemy. Args: path: input database file path, can be a string or `Path` object. """ - return f"sqlite:///{Path(path).absolute().as_posix()}" + return f"sqlite:///{quote(Path(path).absolute().as_posix(), safe='/:')}" def pprint_edges(val: Any, n_lines: int = 20) -> str: diff --git a/pyproject.toml b/pyproject.toml index 830ced4d241..fcd57adac5f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,7 +71,7 @@ all = [ "lpips==0.1.4", "matplotlib>=3.6.3", "MetricsReloaded @ git+https://github.com/Project-MONAI/MetricsReloaded@monai-support", - "mlflow>=3.14.0", + "mlflow>=3.15.2", "nibabel", "ninja", "nni; platform_system == 'Linux' and 'arm' not in platform_machine and 'aarch' not in platform_machine", @@ -126,7 +126,7 @@ lmdb = ["lmdb"] lpips = ["lpips==0.1.4"] matplotlib = ["matplotlib>=3.6.3"] metrics_reloaded = ["MetricsReloaded @ git+https://github.com/Project-MONAI/MetricsReloaded@monai-support"] -mlflow = ["mlflow>=3.14.0"] +mlflow = ["mlflow>=3.15.2"] nibabel = ["nibabel"] nni = [ "nni; platform_system == 'Linux' and 'arm' not in platform_machine and 'aarch' not in platform_machine", diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py index 6ce89f83da8..11913a5a48a 100644 --- a/tests/handlers/test_handler_mlflow.py +++ b/tests/handlers/test_handler_mlflow.py @@ -26,11 +26,9 @@ from monai.apps import download_and_extract from monai.bundle import ConfigWorkflow, download from monai.handlers import MLFlowHandler -from monai.utils import optional_import, path_to_sqlite_uri, path_to_uri +from monai.utils import path_to_sqlite_uri, path_to_uri from tests.test_utils import skip_if_downloading_fails, skip_if_quick -_, has_dataset_tracking = optional_import("mlflow", "2.4.0") - def get_event_filter(e): def event_filter(_, event): @@ -71,8 +69,9 @@ def setUp(self): def tearDown(self): for tmpdir in self.tmpdir_list: if tmpdir and os.path.exists(tmpdir): - # the SQLite default backend creates a db file rather than a directory - shutil.rmtree(tmpdir) if os.path.isdir(tmpdir) else os.remove(tmpdir) + # tmpdir_list may hold a directory or (from dummy_train) a SQLite db file; remove + # the containing directory either way so the mkdtemp parent does not leak. + shutil.rmtree(tmpdir if os.path.isdir(tmpdir) else os.path.dirname(tmpdir)) def test_multi_run(self): with tempfile.TemporaryDirectory() as tempdir: @@ -108,8 +107,7 @@ def _update_metric(engine): self.assertEqual(create_engine_times, run_cnt) def test_default_tracking_uri_is_sqlite(self): - # when no tracking_uri is provided, the handler should default to a local SQLite backend - # rather than the filesystem (file store) backend, which raises on mlflow 3.13+. + """Verify the handler defaults to a local SQLite backend, not the file store, without a tracking URI.""" with tempfile.TemporaryDirectory() as tempdir: cwd = os.getcwd() os.chdir(tempdir) @@ -127,14 +125,19 @@ def test_default_tracking_uri_is_sqlite(self): os.chdir(cwd) def test_remote_tracking_uri_leaves_artifact_location_unset(self): - # a non-local (e.g. remote) tracking_uri must not get a local artifact_location injected, - # so the remote backend keeps deciding where artifacts go. + """Verify a remote tracking URI gets no local artifact location injected.""" handler = MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri="http://localhost:5000") self.assertEqual(handler.client.tracking_uri, "http://localhost:5000") self.assertIsNone(handler.artifact_location) + def test_file_store_tracking_uri_is_rejected(self): + """Verify local paths and file:// URIs are rejected with an actionable error.""" + for uri in ("/tmp/mlruns", path_to_uri(os.path.join("some", "dir"))): + with self.assertRaises(ValueError): + MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri=uri) + def test_explicit_sqlite_tracking_uri_colocates_artifacts(self): - # an explicit local SQLite tracking_uri should still co-locate artifacts next to the db. + """Verify an explicit SQLite tracking URI co-locates artifacts next to the database.""" with tempfile.TemporaryDirectory() as tempdir: uri = path_to_sqlite_uri(os.path.join(tempdir, "sub", "mlruns.db")) handler = MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri=uri) @@ -146,9 +149,7 @@ def test_explicit_sqlite_tracking_uri_colocates_artifacts(self): handler.close() # release the SQLite handle so Windows can delete the db def test_env_var_sqlite_tracking_uri_colocates_artifacts(self): - # a SQLite `MLFLOW_TRACKING_URI` env var should co-locate artifacts next to the db, the - # same as an explicit `tracking_uri` argument. The env var itself is left for MLflow to - # resolve, so the handler does not pass it to the client. + """Verify a SQLite ``MLFLOW_TRACKING_URI`` env var co-locates artifacts next to the db.""" with tempfile.TemporaryDirectory() as tempdir: uri = path_to_sqlite_uri(os.path.join(tempdir, "sub", "mlruns.db")) handler = None @@ -165,8 +166,7 @@ def test_env_var_sqlite_tracking_uri_colocates_artifacts(self): handler.close() # release the SQLite handle so Windows can delete the db def test_explicit_artifact_location_is_used(self): - # an explicitly provided artifact_location should be kept even with the default SQLite - # backend, so callers (e.g. the bundle defaults) can co-locate artifacts with the db. + """Verify an explicit artifact location is preserved with the default SQLite backend.""" with tempfile.TemporaryDirectory() as tempdir: cwd = os.getcwd() os.chdir(tempdir) @@ -181,7 +181,7 @@ def test_explicit_artifact_location_is_used(self): os.chdir(cwd) def test_default_sqlite_run_flow(self): - # a basic log/run flow should work with the default SQLite backend (no tracking_uri given). + """Verify a basic run flow works end-to-end with the default SQLite backend.""" with tempfile.TemporaryDirectory() as tempdir: cwd = os.getcwd() os.chdir(tempdir) @@ -347,7 +347,6 @@ def test_multi_thread(self): self.assertTrue(len(glob.glob(res)) > 0) @skip_if_quick - @unittest.skipUnless(has_dataset_tracking, reason="Requires mlflow version >= 2.4.0.") def test_dataset_tracking(self): test_bundle_name = "endoscopic_tool_segmentation" with tempfile.TemporaryDirectory() as tempdir: diff --git a/tests/utils/misc/test_monai_utils_misc.py b/tests/utils/misc/test_monai_utils_misc.py index f4eb5d3956d..2c9e2f89a71 100644 --- a/tests/utils/misc/test_monai_utils_misc.py +++ b/tests/utils/misc/test_monai_utils_misc.py @@ -16,7 +16,13 @@ from parameterized import parameterized -from monai.utils.misc import MONAIEnvVars, check_kwargs_exist_in_class_init, run_cmd, to_tuple_of_dictionaries +from monai.utils.misc import ( + MONAIEnvVars, + check_kwargs_exist_in_class_init, + path_to_sqlite_uri, + run_cmd, + to_tuple_of_dictionaries, +) TO_TUPLE_OF_DICTIONARIES_TEST_CASES = [ ({}, tuple(), tuple()), @@ -99,5 +105,14 @@ def test_run_cmd(self): self.assertNotIn("\\t", str(cm.exception)) +class TestPathToSqliteUri(unittest.TestCase): + def test_path_to_sqlite_uri(self): + """Verify a sqlite:/// URI is built from the absolute path with special chars escaped.""" + self.assertTrue(path_to_sqlite_uri("/tmp/mlruns.db").startswith("sqlite:///")) + self.assertTrue(path_to_sqlite_uri("/tmp/mlruns.db").endswith("mlruns.db")) + self.assertIn("%3F", path_to_sqlite_uri("a?b.db")) + self.assertIn("%23", path_to_sqlite_uri("a#b.db")) + + if __name__ == "__main__": unittest.main() From 06e0c3116fde4ed22874dd3dd80f0c8b674134b3 Mon Sep 17 00:00:00 2001 From: "R. Garcia-Dias" Date: Fri, 28 Aug 2026 16:50:19 +0100 Subject: [PATCH 4/5] fix(test): root MLflow multi-thread temp dirs in a top-level TemporaryDirectory Follow the repo pattern instead of mkdtemp + manual rmtree: test_multi_thread now creates one top-level TemporaryDirectory, passes it to dummy_train, and lets the context manager clean up. Removes the now-unused setUp/tearDown and shutil import. Signed-off-by: R. Garcia-Dias --- tests/handlers/test_handler_mlflow.py | 33 ++++++++------------------- 1 file changed, 10 insertions(+), 23 deletions(-) diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py index 11913a5a48a..1e222d5e75c 100644 --- a/tests/handlers/test_handler_mlflow.py +++ b/tests/handlers/test_handler_mlflow.py @@ -13,7 +13,6 @@ import glob import os -import shutil import tempfile import unittest from concurrent.futures import ThreadPoolExecutor @@ -39,9 +38,7 @@ def event_filter(_, event): return event_filter -def dummy_train(tracking_folder): - tempdir = tempfile.mkdtemp() - +def dummy_train(tracking_folder, tempdir): # set up engine def _train_func(engine, batch): return [batch + 1.0] @@ -63,16 +60,6 @@ def _train_func(engine, batch): class TestHandlerMLFlow(unittest.TestCase): - def setUp(self): - self.tmpdir_list = [] - - def tearDown(self): - for tmpdir in self.tmpdir_list: - if tmpdir and os.path.exists(tmpdir): - # tmpdir_list may hold a directory or (from dummy_train) a SQLite db file; remove - # the containing directory either way so the mkdtemp parent does not leak. - shutil.rmtree(tmpdir if os.path.isdir(tmpdir) else os.path.dirname(tmpdir)) - def test_multi_run(self): with tempfile.TemporaryDirectory() as tempdir: # set up the train function for engine @@ -336,15 +323,15 @@ def _update_metric(engine): def test_multi_thread(self): test_uri_list = ["monai_mlflow_test1", "monai_mlflow_test2"] - with ThreadPoolExecutor(2, "Training") as executor: - futures = {} - for t in test_uri_list: - futures[t] = executor.submit(dummy_train, t) - - for _, future in futures.items(): - res = future.result() - self.tmpdir_list.append(res) - self.assertTrue(len(glob.glob(res)) > 0) + with tempfile.TemporaryDirectory() as tempdir: + with ThreadPoolExecutor(2, "Training") as executor: + futures = {} + for t in test_uri_list: + futures[t] = executor.submit(dummy_train, t, tempdir) + + for _, future in futures.items(): + res = future.result() + self.assertTrue(len(glob.glob(res)) > 0) @skip_if_quick def test_dataset_tracking(self): From 711c29201e46b5b6327f00ca6e081f0dd953c2cd Mon Sep 17 00:00:00 2001 From: "R. Garcia-Dias" Date: Fri, 28 Aug 2026 17:52:26 +0100 Subject: [PATCH 5/5] fix: honor MLFLOW_TRACKING_URI priority and clarify file-store rejection - Resolve the effective tracking URI with MLFLOW_TRACKING_URI taking precedence over the tracking_uri argument, so an env var can override a hard-coded value (None is passed to the client so MLflow resolves the env var itself). - Drop the redundant truthiness check on effective_tracking_uri in the SQLite artifact-location branch, since it is guaranteed set after the default fallback. - Remove the misleading MLFLOW_ALLOW_FILE_STORE opt-in mention from the file-store rejection message: the handler disallows the file store regardless, so that env var can no longer opt back in. - Add a test covering env-var-over-argument priority. Signed-off-by: R. Garcia-Dias --- monai/handlers/mlflow_handler.py | 30 +++++++++++++-------------- tests/handlers/test_handler_mlflow.py | 14 +++++++++++++ 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/monai/handlers/mlflow_handler.py b/monai/handlers/mlflow_handler.py index 0e9604d16cd..1cd26d5287a 100644 --- a/monai/handlers/mlflow_handler.py +++ b/monai/handlers/mlflow_handler.py @@ -83,9 +83,10 @@ class MLFlowHandler: artifacts stored under ``/mlruns``. The default was changed from the filesystem (file store) backend because MLflow 3.13+ raises an exception for the file store unless ``MLFLOW_ALLOW_FILE_STORE=true`` is set; SQLite is the backend MLflow recommends and it - does not raise. Any explicitly provided ``tracking_uri`` is passed through unchanged, - except local file paths and ``file://`` URIs, which are rejected because MLflow no - longer supports the filesystem (file store) tracking backend. + does not raise. Any explicitly provided ``tracking_uri`` is passed through unchanged + unless ``MLFLOW_TRACKING_URI`` is set (which takes precedence); local file paths and + ``file://`` URIs are rejected because MLflow no longer supports the filesystem (file + store) tracking backend. for more details: https://mlflow.org/docs/latest/python_api/mlflow.html#mlflow.set_tracking_uri. iteration_log: whether to log data to MLFlow when iteration completed, default to `True`. ``iteration_log`` can be also a function and it will be interpreted as an event filter @@ -186,22 +187,18 @@ def __init__( # the `./mlruns` directory (where the previous file store default kept them) via the # experiment `artifact_location`. Any explicitly provided tracking_uri is left unchanged. self.artifact_location = artifact_location - # Resolve the effective tracking URI from the argument or the `MLFLOW_TRACKING_URI` - # environment variable, so both configure the artifact location the same way. - effective_tracking_uri = tracking_uri or os.environ.get("MLFLOW_TRACKING_URI") + # Resolve the effective tracking URI. The `MLFLOW_TRACKING_URI` environment variable takes + # priority so it can override a hard-coded `tracking_uri` argument; both configure the + # artifact location the same way. + env_tracking_uri = os.environ.get("MLFLOW_TRACKING_URI") + effective_tracking_uri = env_tracking_uri or tracking_uri # When neither is set, fall back to the local SQLite default described above. if not effective_tracking_uri: tracking_uri = effective_tracking_uri = path_to_sqlite_uri(os.path.join(os.getcwd(), "mlruns.db")) # For a local SQLite backend, keep run artifacts in an `mlruns` directory next to the # database file (mirroring the previous file-store layout) unless the caller set # `artifact_location`. Other backends (e.g. a remote server) are left to MLflow to decide. - # Only `tracking_uri` is passed to the client, so an `MLFLOW_TRACKING_URI` env var is - # still resolved by MLflow itself. - if ( - self.artifact_location is None - and effective_tracking_uri - and effective_tracking_uri.startswith("sqlite:///") - ): + if self.artifact_location is None and effective_tracking_uri.startswith("sqlite:///"): db_path = Path(effective_tracking_uri[len("sqlite:///") :]) self.artifact_location = path_to_uri(db_path.parent / "mlruns") # MLflow 3.13+ refuses the filesystem (file store) tracking backend, and 3.14+ resolves @@ -211,10 +208,11 @@ def __init__( raise ValueError( "MLflow no longer supports the filesystem (file store) tracking backend; got " f"tracking_uri={effective_tracking_uri!r}. Use a SQLite URI " - "(sqlite:////mlruns.db) or a remote tracking URI, or set " - "MLFLOW_ALLOW_FILE_STORE=true to opt back in." + "(sqlite:////mlruns.db) or a remote tracking URI instead." ) - self.client = mlflow.MlflowClient(tracking_uri=tracking_uri if tracking_uri else None) + # Only the argument is passed to the client; when `MLFLOW_TRACKING_URI` took priority it + # is left None so MLflow resolves the environment variable itself. + self.client = mlflow.MlflowClient(tracking_uri=None if env_tracking_uri else tracking_uri) self.run_finish_status = mlflow.entities.RunStatus.to_string(mlflow.entities.RunStatus.FINISHED) self.close_on_complete = close_on_complete self.experiment = None diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py index 1e222d5e75c..a396227eb9a 100644 --- a/tests/handlers/test_handler_mlflow.py +++ b/tests/handlers/test_handler_mlflow.py @@ -152,6 +152,20 @@ def test_env_var_sqlite_tracking_uri_colocates_artifacts(self): if handler is not None: handler.close() # release the SQLite handle so Windows can delete the db + def test_env_var_tracking_uri_takes_priority_over_argument(self): + """Verify ``MLFLOW_TRACKING_URI`` overrides an explicit ``tracking_uri`` argument.""" + with tempfile.TemporaryDirectory() as tempdir: + env_uri = path_to_sqlite_uri(os.path.join(tempdir, "env.db")) + arg_uri = path_to_sqlite_uri(os.path.join(tempdir, "arg.db")) + handler = None + with patch.dict(os.environ, {"MLFLOW_TRACKING_URI": env_uri}): + try: + handler = MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri=arg_uri) + self.assertTrue(handler.client.tracking_uri.endswith("env.db")) + finally: + if handler is not None: + handler.close() # release the SQLite handle so Windows can delete the db + def test_explicit_artifact_location_is_used(self): """Verify an explicit artifact location is preserved with the default SQLite backend.""" with tempfile.TemporaryDirectory() as tempdir: