fix(infra): reap orphaned workers and survive partial teardown

areal/infra/controller/train_controller.py

@@ -565,6 +565,17 @@ async def _call_workers(
         """Send dispatched inputs to workers. DP heads get slices, others empty."""
         tasks = []
         dp_idx = 0
+        # Save/load are long blocking ops — the HF saver's TP coalesced
+        # all-gather can occupy a worker for tens of seconds, during which its
+        # RPC server may briefly refuse connections (ClientConnectorError). The
+        # default 3-retry/~3s budget would misjudge such a busy worker as dead
+        # and tear down the whole step.
+        # Widen the connection-retry budget for these ops so a busy-but-alive
+        # worker is given time to finish; a genuinely dead worker still fails
+        # after the (larger) retry budget is exhausted, so this does not mask
+        # real crashes.
+        long_op = method in ("save", "load")
+        retry_kw = dict(max_retries=8, retry_delay=2.0) if long_op else {}
         for idx, worker in enumerate(self.workers):
             if self.workers_is_dp_head[idx]:
                 worker_args = [splits[dp_idx] for splits in dp_split_args]
@@ -583,6 +594,7 @@ async def _call_workers(
                     self._engine_name(idx),
                     *worker_args,
                     rpc_meta=rpc_meta,
+                    **retry_kw,
                     **worker_kwargs,
                 )
             )

areal/infra/launcher/sglang_server.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
+import signal
 import subprocess
 import sys
 import time
@@ -256,8 +257,21 @@ def launch_sglang_server(argv):
 
 
 def main(argv):
+    # SIGTERM terminates Python without running finally blocks, orphaning the
+    # sglang child tree (its multiprocessing scheduler and detokenizer workers)
+    # and leaving the node occupied. Convert SIGTERM/SIGINT into SystemExit so
+    # the finally-block kill_process_tree below reaps the whole tree.
+    def _term_handler(signum, _frame):
+        logger.warning(f"sglang launcher received signal {signum}, cleaning up...")
+        raise SystemExit(128 + signum)
+
+    signal.signal(signal.SIGTERM, _term_handler)
+    signal.signal(signal.SIGINT, _term_handler)
+
     try:
         launch_sglang_server(argv)
+    except SystemExit:
+        raise
     except Exception:
         logger.error(traceback.format_exc())
         sys.exit(1)

areal/infra/scheduler/slurm.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import getpass
+import os
 import re
 import shlex
 import subprocess
@@ -904,6 +905,17 @@ def _generate_sbatch_script(
             final_cmd += f" {env_string}"
             final_cmd += f" {spec.image}"
             final_cmd += f" {cmd}"
+            # Stagger container creation across each node's tasks: concurrent
+            # singularity mounts on one node race for the kernel's loop devices
+            # and intermittently die with "failed to find loop device".
+            # AREAL_APPTAINER_STAGGER_SECONDS had been exported by launch
+            # scripts for a while but nothing consumed it. srun does not go
+            # through a shell, so wrap in bash -c for the sleep.
+            stagger = int(os.environ.get("AREAL_APPTAINER_STAGGER_SECONDS", "0"))
+            if stagger > 0:
+                final_cmd = "bash -c " + shlex.quote(
+                    f"sleep $((SLURM_LOCALID * {stagger})); exec {final_cmd}"
+                )
         else:  # native
             final_cmd = cmd
 

areal/trainer/rl_trainer.py

@@ -880,24 +880,47 @@ def train(
             self._save_perf_tracer(step=global_step)
 
     def close(self):
-        self.saver.finalize()
-        if hasattr(self, "_train_rdataset") and self._train_rdataset is not None:
-            self._train_rdataset.close()
-        if hasattr(self, "_valid_rdataset") and self._valid_rdataset is not None:
-            self._valid_rdataset.close()
-        if hasattr(self, "data_controller") and self.data_controller is not None:
-            self.data_controller.destroy()
-        self.stats_logger.close()
-        if self.eval_rollout is not None:
-            self.eval_rollout.destroy()
-        self.rollout.destroy()
-        if self.teacher is not None:
-            self.teacher.destroy()
-        if self.ref is not None:
-            self.ref.destroy()
-        if self.critic is not None:
-            self.critic.destroy()
-        self.actor.destroy()
+        # Must tolerate a partially-constructed trainer (called from
+        # __init__'s failure path), and one engine's destroy() failure must
+        # not keep the remaining workers alive.
+        saver = getattr(self, "saver", None)
+        if saver is not None:
+            try:
+                saver.finalize()
+            except Exception:
+                logger.warning("saver.finalize() failed during close", exc_info=True)
+        for attr in ("_train_rdataset", "_valid_rdataset"):
+            rdataset = getattr(self, attr, None)
+            if rdataset is not None:
+                try:
+                    rdataset.close()
+                except Exception:
+                    logger.warning(f"{attr}.close() failed during close", exc_info=True)
+        data_controller = getattr(self, "data_controller", None)
+        if data_controller is not None:
+            try:
+                data_controller.destroy()
+            except Exception:
+                logger.warning(
+                    "data_controller.destroy() failed during close", exc_info=True
+                )
+        stats_logger = getattr(self, "stats_logger", None)
+        if stats_logger is not None:
+            try:
+                stats_logger.close()
+            except Exception:
+                logger.warning(
+                    "stats_logger.close() failed during close", exc_info=True
+                )
+        for attr in ("eval_rollout", "rollout", "teacher", "ref", "critic", "actor"):
+            engine = getattr(self, attr, None)
+            if engine is not None:
+                try:
+                    engine.destroy()
+                except Exception:
+                    logger.warning(
+                        f"{attr}.destroy() failed during close", exc_info=True
+                    )
         perf_tracer.save(force=True)
 
     def _config_perf_tracer(self):

tests/test_worker_teardown_robustness.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for worker teardown and long-op robustness."""
+
+import os
+import signal
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+
+from areal.infra.launcher import sglang_server
+from areal.trainer.rl_trainer import PPOTrainer
+
+
+class TestSglangLauncherReapsItsChildren:
+    def test_sigterm_is_converted_into_systemexit(self):
+        with (
+            mock.patch.object(sglang_server, "launch_sglang_server") as launch,
+            mock.patch.object(sglang_server, "kill_process_tree") as reap,
+        ):
+            launch.side_effect = lambda argv: os.kill(os.getpid(), signal.SIGTERM)
+
+            with pytest.raises(SystemExit) as exc:
+                sglang_server.main([])
+
+        assert exc.value.code == 128 + signal.SIGTERM
+        reap.assert_called_once()
+
+    def test_unexpected_errors_still_reap_the_tree(self):
+        with (
+            mock.patch.object(sglang_server, "launch_sglang_server") as launch,
+            mock.patch.object(sglang_server, "kill_process_tree") as reap,
+        ):
+            launch.side_effect = RuntimeError("boom")
+
+            with pytest.raises(SystemExit):
+                sglang_server.main([])
+
+        reap.assert_called_once()
+
+
+class TestTrainerCloseToleratesPartialConstruction:
+    def test_close_on_a_bare_trainer_does_not_raise(self):
+        trainer = object.__new__(PPOTrainer)
+
+        PPOTrainer.close(trainer)
+
+    def test_one_failing_component_does_not_skip_the_others(self):
+        trainer = object.__new__(PPOTrainer)
+        closed = []
+        trainer.saver = SimpleNamespace(
+            finalize=lambda: (_ for _ in ()).throw(RuntimeError("saver down"))
+        )
+        trainer.stats_logger = SimpleNamespace(close=lambda: closed.append("stats"))
+
+        PPOTrainer.close(trainer)
+
+        assert "stats" in closed, (
+            "a failing saver prevented the remaining components from closing"
+        )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])