Skip to content

Commit 4cd31fe

Browse files
committed
Fix MLFlowLogger.save_dir Windows file URI handling (#20972)
- Replace simple string slicing with urllib.parse.urlparse and url2pathname - Properly handle Windows absolute file URIs (e.g., file:///C:/path) - Add comprehensive tests for various file URI formats - Fix malformed paths like ///C:/path becoming C:/path on Windows Fixes #20972
1 parent a0ce930 commit 4cd31fe

File tree

2 files changed

+37
-1
lines changed

2 files changed

+37
-1
lines changed

src/lightning/pytorch/loggers/mlflow.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -300,7 +300,11 @@ def save_dir(self) -> Optional[str]:
300300
301301
"""
302302
if self._tracking_uri.startswith(LOCAL_FILE_URI_PREFIX):
303-
return self._tracking_uri[len(LOCAL_FILE_URI_PREFIX) :]
303+
from urllib.parse import urlparse
304+
from urllib.request import url2pathname
305+
306+
parsed_uri = urlparse(self._tracking_uri)
307+
return url2pathname(parsed_uri.path)
304308
return None
305309

306310
@property

tests/tests_pytorch/loggers/test_mlflow.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -427,3 +427,35 @@ def test_set_tracking_uri(mlflow_mock):
427427
mlflow_mock.set_tracking_uri.assert_not_called()
428428
_ = logger.experiment
429429
mlflow_mock.set_tracking_uri.assert_called_with("the_tracking_uri")
430+
431+
432+
@mock.patch("lightning.pytorch.loggers.mlflow._get_resolve_tags", Mock())
433+
def test_mlflow_logger_save_dir_file_uri_handling(mlflow_mock):
434+
"""Test that save_dir correctly handles file URIs, especially on Windows."""
435+
# Test Unix-style absolute file URI
436+
logger = MLFlowLogger(tracking_uri="file:///home/user/mlruns")
437+
expected_unix = "/home/user/mlruns"
438+
assert logger.save_dir == expected_unix
439+
440+
# Test Windows-style absolute file URI
441+
logger_win = MLFlowLogger(tracking_uri="file:///C:/Dev/example/mlruns")
442+
# On Windows, url2pathname converts file:///C:/path to C:\path
443+
# On Unix, it converts to /C:/path, but we test the actual behavior
444+
import platform
445+
446+
expected_win = "C:\\Dev\\example\\mlruns" if platform.system() == "Windows" else "/C:/Dev/example/mlruns"
447+
assert logger_win.save_dir == expected_win
448+
449+
# Test relative file URI
450+
logger_rel = MLFlowLogger(tracking_uri="file:./mlruns")
451+
expected_rel = "./mlruns"
452+
assert logger_rel.save_dir == expected_rel
453+
454+
# Test non-file URI (should return None)
455+
logger_http = MLFlowLogger(tracking_uri="http://localhost:8080")
456+
assert logger_http.save_dir is None
457+
458+
# Test file URI with special characters and spaces
459+
logger_special = MLFlowLogger(tracking_uri="file:///path/with%20spaces/mlruns")
460+
expected_special = "/path/with spaces/mlruns"
461+
assert logger_special.save_dir == expected_special

0 commit comments

Comments
 (0)