diff --git a/areal/infra/controller/train_controller.py b/areal/infra/controller/train_controller.py index fe37d1ed8e..c13804d0b7 100644 --- a/areal/infra/controller/train_controller.py +++ b/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, ) ) diff --git a/areal/infra/launcher/sglang_server.py b/areal/infra/launcher/sglang_server.py index 07b76344ca..ffb0022e04 100644 --- a/areal/infra/launcher/sglang_server.py +++ b/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) diff --git a/areal/infra/scheduler/slurm.py b/areal/infra/scheduler/slurm.py index 54696ccd07..b2a582500b 100644 --- a/areal/infra/scheduler/slurm.py +++ b/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 diff --git a/areal/trainer/rl_trainer.py b/areal/trainer/rl_trainer.py index bdb6b22a72..e66a60d28c 100644 --- a/areal/trainer/rl_trainer.py +++ b/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): diff --git a/tests/test_worker_teardown_robustness.py b/tests/test_worker_teardown_robustness.py new file mode 100644 index 0000000000..34985b923f --- /dev/null +++ b/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"])