Derive artifact location from MLFLOW_TRACKING_URI env var too

When MLFLOW_TRACKING_URI was set as an environment variable (not the
tracking_uri argument) to a local SQLite database, the handler did not
co-locate run artifacts next to the db file, so they went to MLflow's
default mlruns directory relative to the working directory. An
argument-passed SQLite URI did get a sibling mlruns. Resolve an
effective tracking URI from the argument or the env var and key the
artifact-location logic off it, while still passing only tracking_uri to
the client so MLflow keeps resolving the env var. Add a regression test
for the env-var SQLite case.

Signed-off-by: Lawson Darrow <lawson.darrow@gmail.com>

Author: Lawson-Darrow

monai/handlers/mlflow_handler.py

@@ -130,9 +130,10 @@ class MLFlowHandler:
             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 ``tracking_uri`` is used, 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.
+            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.
 
@@ -184,16 +185,23 @@ 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
-        # Only fall back to the SQLite default when the caller gave no tracking_uri and the
-        # `MLFLOW_TRACKING_URI` environment variable is unset, so that env-var configuration keeps
-        # working. When it is set, `tracking_uri` stays None and MLflow resolves the env var.
-        if not tracking_uri and not os.environ.get("MLFLOW_TRACKING_URI"):
-            tracking_uri = path_to_sqlite_uri(os.path.join(os.getcwd(), "mlruns.db"))
+        # 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.
-        if self.artifact_location is None and tracking_uri and tracking_uri.startswith("sqlite:///"):
-            db_path = Path(tracking_uri[len("sqlite:///") :])
+        # 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)

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
@@ -289,6 +289,25 @@ def test_explicit_sqlite_tracking_uri_colocates_artifacts(self):
             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.