Use a clearer hierarchy of run id -> experiment name -> run name

Signed-off-by: Nathan Azrak <nathan.azrak@gmail.com>

Author: nathan-az

nemo_rl/utils/logger.py

@@ -783,8 +783,21 @@ def __init__(self, cfg: MLflowConfig, log_dir: Optional[str] = None):
         run_name = cfg.get("run_name") or os.getenv("MLFLOW_RUN_NAME")
 
         run = mlflow.active_run()
-        # Start a new run if there is no active run, or if the active run doesn't match the requested run_id
-        if run is None or (run_id and run.info.run_id != run_id):
+        
+        # If run_id is provided, try to use it directly
+        if run_id:
+            # If there is an active run but it's not the one we want, end it
+            if run and run.info.run_id != run_id:
+                mlflow.end_run()
+                run = None
+            
+            # Start/resume the specified run
+            if run is None:
+                run = mlflow.start_run(run_id=run_id)
+        
+        # If no run_id provided, fall back to experiment name logic
+        else:
+            # End any existing active run to start fresh or ensure correct context
             if run:
                 mlflow.end_run()
 
@@ -794,14 +807,13 @@ def __init__(self, cfg: MLflowConfig, log_dir: Optional[str] = None):
                 if experiment is None:
                     mlflow.create_experiment(
                         name=experiment_name,
-                        **{"artifact_location": cfg.get("artifact_location", log_dir)}
-                        if "artifact_location" in cfg or log_dir
-                        else {},
+                        artifact_location=cfg.get("artifact_location") or log_dir,
                     )
                 # set the experiment context manager
                 mlflow.set_experiment(experiment_name)
-            # if run_id is set explicitly, will use. Otherwise, from env var. Otherwise, new run with run name
-            run = mlflow.start_run(run_name=run_name, run_id=run_id)
+            # Start a new run
+            run = mlflow.start_run(run_name=run_name)
+
         self.run = run
         self.run_id = run.info.run_id
         print(

tests/unit/utils/test_logger.py

@@ -503,7 +503,7 @@ def test_init_basic_config(self, mock_mlflow, temp_dir):
         MLflowLogger(cfg, log_dir=temp_dir)
 
         mock_mlflow.set_experiment.assert_called_once_with("test-experiment")
-        mock_mlflow.start_run.assert_called_once_with(run_name="test-run", run_id=None)
+        mock_mlflow.start_run.assert_called_once_with(run_name="test-run")
 
     @patch("nemo_rl.utils.logger.mlflow")
     def test_init_full_config(self, mock_mlflow, temp_dir):
@@ -522,7 +522,7 @@ def test_init_full_config(self, mock_mlflow, temp_dir):
 
         mock_mlflow.set_tracking_uri.assert_called_once_with("http://localhost:5000")
         mock_mlflow.set_experiment.assert_called_once_with("test-experiment")
-        mock_mlflow.start_run.assert_called_once_with(run_name="test-run", run_id=None)
+        mock_mlflow.start_run.assert_called_once_with(run_name="test-run")
 
     @patch("nemo_rl.utils.logger.mlflow")
     def test_log_metrics(self, mock_mlflow, temp_dir):
@@ -619,6 +619,9 @@ def test_cleanup(self, mock_mlflow, temp_dir):
         }
         logger = MLflowLogger(cfg, log_dir=temp_dir)
 
+        # Reset mocks to avoid counting calls from init
+        mock_mlflow.end_run.reset_mock()
+
         # Trigger cleanup
         logger.__del__()
 
@@ -641,8 +644,10 @@ def test_init_with_none_log_dir(self, mock_mlflow):
         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", run_id=None)
+        mock_mlflow.create_experiment.assert_called_once_with(
+            name="test-experiment", artifact_location=None
+        )
+        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):
@@ -663,7 +668,7 @@ def test_init_with_custom_log_dir(self, mock_mlflow):
         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", run_id=None)
+        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):
@@ -689,12 +694,12 @@ def test_init_with_artifact_location_in_config(self, mock_mlflow):
         )
         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"], run_id=None
+            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."""
+        """Test initialization with artifact_location=None in config uses log_dir."""
         # Ensure active_run returns None so initialization logic runs
         mock_mlflow.active_run.return_value = None
         # Mock is_tracking_uri_set to return False so set_tracking_uri is called
@@ -710,14 +715,13 @@ def test_init_with_artifact_location_none_in_config(self, mock_mlflow):
 
         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)
+        # Verify create_experiment was called with log_dir since config is None
         mock_mlflow.create_experiment.assert_called_once_with(
-            name=cfg["experiment_name"], artifact_location=cfg["artifact_location"]
+            name=cfg["experiment_name"], artifact_location="/fallback/path"
         )
         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"], run_id=None
+            run_name=cfg["run_name"]
         )
 
     @patch("nemo_rl.utils.logger.mlflow")
@@ -744,7 +748,7 @@ def test_init_without_artifact_location_uses_log_dir(self, mock_mlflow):
         )
         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"], run_id=None
+            run_name=cfg["run_name"]
         )