fix(tracking): finalize MLflow runs on SIGTERM instead of leaving them RUNNING

skyrl/train/utils/tracking.py

@@ -15,9 +15,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import atexit
 import dataclasses
 import pprint
+import signal
 import traceback
+import weakref
 from enum import Enum
 from functools import partial
 from pathlib import Path
@@ -89,6 +92,49 @@ def __init__(
             self.logger["console"] = self.console_logger
 
         self._exception_logged = False
+        self._finished = False
+
+        # `__del__` is not guaranteed to run on signal-driven termination (e.g.
+        # SLURM `scancel` sending SIGTERM), which leaves MLflow runs stuck in
+        # RUNNING. Register an atexit hook and a SIGTERM handler that finalize the
+        # run on abnormal exit. Both reference self weakly (and are removed in
+        # `finish`) so they don't keep the instance alive or disable `__del__`.
+        self_ref = weakref.ref(self)
+
+        def _atexit_handler():
+            tracker = self_ref()
+            if tracker is not None:
+                tracker.finish()
+
+        def _sigterm_handler(signum, frame):
+            tracker = self_ref()
+            if tracker is not None:
+                tracker._handle_sigterm(signum, frame)
+
+        self._atexit_handler = _atexit_handler
+        self._sigterm_handler = _sigterm_handler
+        atexit.register(_atexit_handler)
+
+        # `signal.signal` only works from the main thread of the main interpreter.
+        # Tracking may be constructed inside a Ray worker thread, so guard against
+        # the ValueError and simply skip the handler there (atexit still applies).
+        self._previous_sigterm_handler = None
+        try:
+            self._previous_sigterm_handler = signal.signal(signal.SIGTERM, _sigterm_handler)
+        except ValueError as e:
+            logger.warning(f"Could not register SIGTERM handler for tracking (not main thread?): {e}")
+
+    def _handle_sigterm(self, signum, frame):
+        self.finish(status="KILLED")
+        # Chain to the previous handler so we don't swallow the signal's effect
+        # (e.g. process teardown). If there was no Python-level handler, restore
+        # the default disposition and re-raise SIGTERM.
+        previous = self._previous_sigterm_handler
+        if callable(previous):
+            previous(signum, frame)
+        elif previous == signal.SIG_DFL:
+            signal.signal(signal.SIGTERM, signal.SIG_DFL)
+            signal.raise_signal(signal.SIGTERM)
 
     def log(self, data, step, commit=False):
         for logger_name, logger_instance in self.logger.items():
@@ -97,14 +143,30 @@ def log(self, data, step, commit=False):
             else:
                 logger_instance.log(data=data, step=step)
 
-    def finish(self):
+    def finish(self, status: str = "FINISHED"):
+        # Idempotent: signal handler, atexit hook, and __del__ can all fire, but
+        # the underlying backends should only be finalized once.
+        if self._finished:
+            return
+        self._finished = True
+
+        # Remove the lifecycle hooks now that the run is finalized.
+        atexit.unregister(self._atexit_handler)
+        if self._previous_sigterm_handler is not None:
+            try:
+                signal.signal(signal.SIGTERM, self._previous_sigterm_handler)
+            except ValueError:
+                pass
+
         for logger_name, logger_instance in self.logger.items():
             # NOTE (sumanthrh): We use a try-except block here while finishing tracking.
             # This is because wandb often errors out with a BrokenPipeError when closing.
             # https://github.com/wandb/wandb/issues/6449
             try:
                 if logger_name == "wandb":
                     logger_instance.finish(exit_code=0)
+                elif logger_name == "mlflow":
+                    logger_instance.finish(status=status)
                 elif logger_name != "console":
                     logger_instance.finish()
             except Exception as e:
@@ -225,9 +287,9 @@ def log(self, data, step):
         results = {k.replace("@", "_at_"): v for k, v in data.items()}
         self.mlflow.log_metrics(metrics=results, step=step)
 
-    def finish(self):
+    def finish(self, status: str = "FINISHED"):
         if self.we_created_mlflow:
-            self.mlflow.end_run()
+            self.mlflow.end_run(status=status)
 
 
 def _compute_mlflow_params_from_objects(params) -> Dict[str, Any]:

tests/train/test_tracking.py

@@ -2,11 +2,33 @@
 uv run --isolated --extra dev --extra skyrl-train pytest -s tests/train/test_tracking.py
 """
 
+import gc
+import signal
+import weakref
 from unittest.mock import MagicMock, patch
 
+import pytest
+
 from skyrl.train.utils.tracking import Tracking
 
 
+@pytest.fixture
+def mlflow_mock():
+    """Provide a mocked mlflow module with no active run so the adapter starts one."""
+    mlflow = MagicMock()
+    mlflow.active_run.return_value = None
+    with patch.dict("sys.modules", {"mlflow": mlflow}):
+        yield mlflow
+
+
+@pytest.fixture(autouse=True)
+def restore_sigterm():
+    """Tracking installs a SIGTERM handler; restore the default after each test."""
+    original = signal.getsignal(signal.SIGTERM)
+    yield
+    signal.signal(signal.SIGTERM, original)
+
+
 def test_wandb_init_receives_tags():
     """Tags passed to Tracking are forwarded to wandb.init."""
     with patch.dict("sys.modules", {"wandb": MagicMock()}) as mocked:
@@ -39,3 +61,62 @@ def test_wandb_init_tags_default_none():
 
         wandb_mock.init.assert_called_once()
         assert wandb_mock.init.call_args.kwargs["tags"] is None
+
+
+def test_mlflow_finish_default_status_finished(mlflow_mock):
+    """Tracking.finish() ends the MLflow run with FINISHED by default."""
+    tracker = Tracking(project_name="proj", experiment_name="exp", backends=["mlflow"], config=None)
+
+    tracker.finish()
+
+    mlflow_mock.end_run.assert_called_once_with(status="FINISHED")
+
+
+def test_mlflow_finish_forwards_status(mlflow_mock):
+    """A status passed to Tracking.finish is forwarded to mlflow.end_run."""
+    tracker = Tracking(project_name="proj", experiment_name="exp", backends=["mlflow"], config=None)
+
+    tracker.finish(status="KILLED")
+
+    mlflow_mock.end_run.assert_called_once_with(status="KILLED")
+
+
+def test_mlflow_finish_is_idempotent(mlflow_mock):
+    """Repeated finish() calls (signal + atexit + __del__) only end the run once."""
+    tracker = Tracking(project_name="proj", experiment_name="exp", backends=["mlflow"], config=None)
+
+    tracker.finish(status="KILLED")
+    tracker.finish()
+
+    mlflow_mock.end_run.assert_called_once_with(status="KILLED")
+
+
+def test_sigterm_handler_finishes_run_as_killed(mlflow_mock):
+    """A SIGTERM after Tracking init ends the MLflow run with status KILLED.
+
+    Regression test for runs left in RUNNING when the process is killed by a
+    signal (e.g. SLURM scancel) rather than exiting cleanly via __del__.
+    """
+    tracker = Tracking(project_name="proj", experiment_name="exp", backends=["mlflow"], config=None)
+
+    # Tracking should have installed its SIGTERM handler.
+    handler = signal.getsignal(signal.SIGTERM)
+    assert handler == tracker._sigterm_handler
+
+    # Previous disposition was the default; chaining would re-raise SIGTERM and
+    # kill the test process, so swap it for a no-op before invoking the handler.
+    tracker._previous_sigterm_handler = lambda *args: None
+    handler(signal.SIGTERM, None)
+
+    mlflow_mock.end_run.assert_called_once_with(status="KILLED")
+
+
+def test_finish_releases_instance_for_gc(mlflow_mock):
+    """After finish(), the atexit/SIGTERM hooks must not keep the instance alive."""
+    tracker = Tracking(project_name="proj", experiment_name="exp", backends=["mlflow"], config=None)
+    ref = weakref.ref(tracker)
+    tracker.finish()
+
+    del tracker
+    gc.collect()
+    assert ref() is None