fix: honor mlflow server artifact_location (#1536) (#1538)

Signed-off-by: Alexander Zhipa <azzhipa@amazon.com>
Co-authored-by: Alexander Zhipa <azzhipa@amazon.com>

Author: clumsy

nemo_rl/utils/logger.py

@@ -65,6 +65,7 @@ class MLflowConfig(TypedDict):
     experiment_name: str
     run_name: str
     tracking_uri: NotRequired[str]
+    artifact_location: NotRequired[str | None]
 
 
 class GPUMonitoringConfig(TypedDict):
@@ -724,31 +725,30 @@ def __init__(self, cfg: MLflowConfig, log_dir: Optional[str] = None):
 
         Args:
             cfg: MLflow configuration
-            log_dir: Optional log directory
+            log_dir: Optional log directory (used as fallback if artifact_location not in cfg)
         """
-        if cfg["tracking_uri"]:
-            mlflow.set_tracking_uri(cfg["tracking_uri"])
+        tracking_uri = cfg.get("tracking_uri")
+        if tracking_uri:
+            mlflow.set_tracking_uri(tracking_uri)
 
-        experiment = mlflow.get_experiment_by_name(cfg["experiment_name"])
+        experiment_name = cfg["experiment_name"]
+        experiment = mlflow.get_experiment_by_name(experiment_name)
         if experiment is None:
-            if log_dir:
-                mlflow.create_experiment(
-                    name=cfg["experiment_name"],
-                    artifact_location=log_dir,
-                )
-            else:
-                mlflow.create_experiment(cfg["experiment_name"])
+            mlflow.create_experiment(
+                name=experiment_name,
+                **{"artifact_location": cfg.get("artifact_location", log_dir)}
+                if "artifact_location" in cfg or log_dir
+                else {},
+            )
         else:
-            mlflow.set_experiment(cfg["experiment_name"])
+            mlflow.set_experiment(experiment_name)
 
         # Start run
-        run_kwargs: dict[str, str] = {}
-        run_kwargs["run_name"] = cfg["run_name"]
-
+        run_name = cfg["run_name"]
+        run_kwargs = {"run_name": run_name}
         self.run = mlflow.start_run(**run_kwargs)
         print(
-            f"Initialized MLflowLogger for experiment {cfg['experiment_name']}, "
-            f"run {cfg['run_name']}"
+            f"Initialized MLflowLogger for experiment {experiment_name}, run {run_name}"
         )
 
     def log_metrics(
@@ -847,8 +847,10 @@ def __init__(self, cfg: LoggerConfig):
             self.loggers.append(tensorboard_logger)
 
         if cfg["mlflow_enabled"]:
-            mlflow_log_dir = os.path.join(self.base_log_dir, "mlflow")
-            os.makedirs(mlflow_log_dir, exist_ok=True)
+            mlflow_log_dir = self.base_log_dir
+            if mlflow_log_dir:
+                mlflow_log_dir = os.path.join(mlflow_log_dir, "mlflow")
+                os.makedirs(mlflow_log_dir, exist_ok=True)
             mlflow_logger = MLflowLogger(cfg["mlflow"], log_dir=mlflow_log_dir)
             self.loggers.append(mlflow_logger)
 

tests/unit/utils/test_logger.py

@@ -548,6 +548,101 @@ def test_cleanup(self, mock_mlflow, temp_dir):
         # Check that end_run was called
         mock_mlflow.end_run.assert_called_once()
 
+    @patch("nemo_rl.utils.logger.mlflow")
+    def test_init_with_none_log_dir(self, mock_mlflow):
+        """Test initialization with None log_dir uses server default artifact location."""
+        cfg = {
+            "experiment_name": "test-experiment",
+            "run_name": "test-run",
+            "tracking_uri": "http://localhost:5000",
+        }
+        mock_mlflow.get_experiment_by_name.return_value = None
+
+        MLflowLogger(cfg, log_dir=None)
+
+        # Verify create_experiment was called without artifact_location
+        mock_mlflow.create_experiment.assert_called_once_with(name="test-experiment")
+        mock_mlflow.start_run.assert_called_once_with(run_name="test-run")
+
+    @patch("nemo_rl.utils.logger.mlflow")
+    def test_init_with_custom_log_dir(self, mock_mlflow):
+        """Test initialization with custom log_dir sets artifact_location."""
+        cfg = {
+            "experiment_name": "test-experiment",
+            "run_name": "test-run",
+            "tracking_uri": "http://localhost:5000",
+        }
+        mock_mlflow.get_experiment_by_name.return_value = None
+
+        MLflowLogger(cfg, log_dir="/custom/path")
+
+        # Verify create_experiment was called with artifact_location
+        mock_mlflow.create_experiment.assert_called_once_with(
+            name="test-experiment", artifact_location="/custom/path"
+        )
+        mock_mlflow.start_run.assert_called_once_with(run_name="test-run")
+
+    @patch("nemo_rl.utils.logger.mlflow")
+    def test_init_with_artifact_location_in_config(self, mock_mlflow):
+        """Test initialization with artifact_location in config takes precedence over log_dir."""
+        cfg = {
+            "experiment_name": "test-experiment",
+            "run_name": "test-run",
+            "tracking_uri": "http://localhost:5000",
+            "artifact_location": "/config/artifact/path",
+        }
+        mock_mlflow.get_experiment_by_name.return_value = None
+
+        MLflowLogger(cfg, log_dir="/fallback/path")
+
+        # Verify create_experiment was called with artifact_location from config
+        mock_mlflow.create_experiment.assert_called_once_with(
+            name=cfg["experiment_name"], artifact_location=cfg["artifact_location"]
+        )
+        mock_mlflow.set_tracking_uri.assert_called_once_with(cfg["tracking_uri"])
+        mock_mlflow.start_run.assert_called_once_with(run_name=cfg["run_name"])
+
+    @patch("nemo_rl.utils.logger.mlflow")
+    def test_init_with_artifact_location_none_in_config(self, mock_mlflow):
+        """Test initialization with artifact_location=None in config uses server default."""
+        cfg = {
+            "experiment_name": "test-experiment",
+            "run_name": "test-run",
+            "tracking_uri": "http://localhost:5000",
+            "artifact_location": None,
+        }
+        mock_mlflow.get_experiment_by_name.return_value = None
+
+        MLflowLogger(cfg, log_dir="/fallback/path")
+
+        # Verify create_experiment was called without artifact_location
+        # (None is explicitly set, so we don't pass it to MLflow)
+        mock_mlflow.create_experiment.assert_called_once_with(
+            name=cfg["experiment_name"], artifact_location=cfg["artifact_location"]
+        )
+        mock_mlflow.set_tracking_uri.assert_called_once_with(cfg["tracking_uri"])
+        mock_mlflow.start_run.assert_called_once_with(run_name=cfg["run_name"])
+
+    @patch("nemo_rl.utils.logger.mlflow")
+    def test_init_without_artifact_location_uses_log_dir(self, mock_mlflow):
+        """Test initialization without artifact_location in config uses log_dir."""
+        cfg = {
+            "experiment_name": "test-experiment",
+            "run_name": "test-run",
+            "tracking_uri": "http://localhost:5000",
+        }
+        mock_mlflow.get_experiment_by_name.return_value = None
+
+        log_dir = "/fallback/path"
+        MLflowLogger(cfg, log_dir=log_dir)
+
+        # Verify create_experiment was called with log_dir as artifact_location
+        mock_mlflow.create_experiment.assert_called_once_with(
+            name=cfg["experiment_name"], artifact_location=log_dir
+        )
+        mock_mlflow.set_tracking_uri.assert_called_once_with(cfg["tracking_uri"])
+        mock_mlflow.start_run.assert_called_once_with(run_name=cfg["run_name"])
+
 
 class TestRayGpuMonitorLogger:
     """Test the RayGpuMonitorLogger class."""