diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index 353af6d12..b5657baf7 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -13,6 +13,16 @@ from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset +def _barrier() -> None: + """``torch.distributed.barrier()`` that is a no-op outside a process group. + + Single-process runs (the CPU smoke test, ``devices=1`` without DDP) never + call ``init_process_group``; an unguarded barrier raises there. + """ + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.barrier() + + class ModelWrapper(LightningModule): """ Wrapper class around robomimic models to ensure compatibility with Pytorch Lightning. @@ -224,7 +234,7 @@ def on_validation_end(self): f"Rank {self.global_rank} on validation end, waiting for all ranks to synchronize", flush=True, ) - torch.distributed.barrier() + _barrier() print( f"Rank {self.global_rank} on validation end, all ranks synchronized", flush=True, @@ -283,7 +293,7 @@ def on_fit_start(self): f"Rank {self.global_rank} on fit start, waiting for all ranks to synchronize", flush=True, ) - torch.distributed.barrier() + _barrier() print( f"Rank {self.global_rank} on fit start, all ranks synchronized", flush=True ) diff --git a/egomimic/pl_utils/preemption.py b/egomimic/pl_utils/preemption.py new file mode 100644 index 000000000..83ce174ae --- /dev/null +++ b/egomimic/pl_utils/preemption.py @@ -0,0 +1,164 @@ +"""Checkpoint-on-preemption for Slurm runs. + +Why this exists +--------------- +Slurm delivers a warning signal before it preempts or times out a job. Two +things already listen for it: + +* **submitit** (``hydra -m`` launches) installs a handler on ``USR2`` (or + ``$SUBMITIT_PREEMPT_SIGNAL``) at process start. It re-pickles the Hydra task + and calls ``scontrol requeue`` -- but it never saves a Lightning checkpoint. +* **Lightning** would save a checkpoint and requeue, but its + ``SignalConnector`` refuses to register on a signal that already has a + handler, so under submitit it silently does nothing. + +Result: every preemption restarts from the last *periodic* ``last.ckpt``. + +``PreemptionCheckpoint`` installs its own handler at ``on_fit_start`` that +saves a checkpoint first and then hands the signal to the handler that was +installed before it (submitit's requeue). The saved path defaults to +``/checkpoints/last.ckpt``, which is what ``trainHydra.py`` +resumes from when ``SLURM_RESTART_COUNT`` is set. + +The window available is the partition's preemption ``GraceTime`` (300 s on +``overcap``; ``rl2-lab`` does not preempt), or submitit's warning (120 s by +default) before the time limit. +""" + +from __future__ import annotations + +import logging +import os +import signal +from pathlib import Path +from typing import Optional + +from lightning import Trainer +from lightning.pytorch.callbacks import Callback, ModelCheckpoint + +log = logging.getLogger(__name__) + +_SUBMITIT_SIGNAL_ENV = "SUBMITIT_PREEMPT_SIGNAL" + + +def preempt_signal() -> signal.Signals: + """The signal submitit listens for: ``$SUBMITIT_PREEMPT_SIGNAL`` or ``USR2``.""" + name = os.environ.get(_SUBMITIT_SIGNAL_ENV, "USR2") + return signal.Signals[f"SIG{name}"] + + +class PreemptionCheckpoint(Callback): + """Save a checkpoint on the Slurm preemption signal, then chain to the + previously installed handler (submitit's requeue). + + Args: + ckpt_path: where to save. Defaults to + ``/checkpoints/last.ckpt``. + signum: signal to handle. Defaults to :func:`preempt_signal`. + """ + + def __init__( + self, + ckpt_path: Optional[str] = None, + signum: Optional[signal.Signals] = None, + ) -> None: + self.ckpt_path = ckpt_path + self.signum = signal.Signals(signum) if signum is not None else preempt_signal() + self._trainer: Optional[Trainer] = None + self._previous = None + self._installed = False + self._pid: Optional[int] = None + + # -- Callback hooks ------------------------------------------------------ + + def on_fit_start(self, trainer: Trainer, pl_module) -> None: + self._trainer = trainer + self._pid = os.getpid() + if self.ckpt_path is None: + self.ckpt_path = str( + Path(trainer.default_root_dir) / "checkpoints" / "last.ckpt" + ) + self._previous = signal.getsignal(self.signum) + signal.signal(self.signum, self._handle) + self._installed = True + if callable(self._previous): + log.info( + "PreemptionCheckpoint: on %s save %s, then chain to the existing handler", + self.signum.name, + self.ckpt_path, + ) + else: + log.warning( + "PreemptionCheckpoint: on %s save %s; no other handler is installed, " + "so nothing will requeue this job", + self.signum.name, + self.ckpt_path, + ) + + def teardown(self, trainer: Trainer, pl_module, stage: str) -> None: + if self._installed: + signal.signal(self.signum, self._previous) + self._installed = False + + # -- handler ------------------------------------------------------------- + + def _handle(self, signum, frame) -> None: + if os.getpid() != self._pid: + # Slurm signals every process in the step. Forked DataLoader + # workers inherit this handler; only the trainer process may act. + return + trainer = self._trainer + assert trainer is not None and self.ckpt_path is not None + log.warning( + "Caught %s: saving checkpoint to %s before requeue", + signal.Signals(signum).name, + self.ckpt_path, + ) + try: + self._claim_last_ckpt(trainer) + # Write next to the target and rename, so a SIGKILL mid-write + # (GraceTime expiry) cannot leave a truncated last.ckpt behind. + tmp = self.ckpt_path + ".preempt" + # Collective: every rank must get here. Saving from inside the + # signal handler is the same design as Lightning's own SLURM + # handler, with the same caveat: a rank interrupted mid-NCCL + # collective cannot join this one, and the step then hangs until + # GraceTime's SIGKILL. All shipped launchers run world size 1; a + # multi-GPU follow-up should flag here and save from + # on_train_batch_end instead. + trainer.save_checkpoint(tmp) + if trainer.is_global_zero: + os.replace(tmp, self.ckpt_path) + except Exception: + # A failed save (quota, NFS) must not cancel the requeue: the job + # would otherwise end with nothing saved and nothing rescheduled. + # If only rank 0 fails, the other ranks stay in save_checkpoint's + # barrier until the requeue tears the step down; acceptable. + log.exception("Preemption checkpoint failed; requeuing anyway") + finally: + if callable(self._previous): + # submitit's checkpoint_and_try_requeue (it requeues on rank 0 + # only and may raise SystemExit; that propagates, as it would + # without this callback). + self._previous(signum, frame) + + def _claim_last_ckpt(self, trainer: Trainer) -> None: + """Tell every ``ModelCheckpoint(save_last=True)`` writing into the same + directory that ``last.ckpt`` is its file. + + Otherwise its version counter sees a ``last.ckpt`` it did not write and + every later periodic save goes to ``last-v1.ckpt``, ``last-v2.ckpt``, ... + while ``trainHydra`` keeps resuming from the stale ``last.ckpt``. The + path is part of ``ModelCheckpoint``'s saved state, so this also holds + in the resumed process. Compared through ``realpath`` because Lightning + resolves ``dirpath`` that way but not ``default_root_dir``. + """ + assert self.ckpt_path is not None + target = Path(self.ckpt_path) + target_dir = os.path.realpath(target.parent) + for cb in trainer.checkpoint_callbacks: + if not (isinstance(cb, ModelCheckpoint) and cb.save_last and cb.dirpath): + continue + last_name = cb.CHECKPOINT_NAME_LAST + cb.FILE_EXTENSION + if os.path.realpath(cb.dirpath) == target_dir and target.name == last_name: + cb.last_model_path = os.path.join(cb.dirpath, last_name) diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 53c0f012a..a0fe205ac 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -1,6 +1,5 @@ import copy import os -import signal from typing import Any, Dict, List, Optional, Tuple import hydra @@ -19,6 +18,7 @@ import egomimic.utils.hydra_resolvers # noqa: F401 -- registers OmegaConf resolvers from egomimic.eval.eval import Eval from egomimic.pl_utils.pl_model import ModelWrapper +from egomimic.pl_utils.preemption import PreemptionCheckpoint from egomimic.rldb.zarr.utils import set_global_seed from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset from egomimic.utils.aws.aws_data_utils import load_env @@ -234,10 +234,12 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: if cfg.get("mmap_checkpoint", True): plugins.append(MmapCheckpointIO()) if os.environ.get("SLURM_JOB_ID"): - # requeue_signal is a single signal, not a list -- SignalConnector passes - # it straight to signal.getsignal(), which raises TypeError on a list. - plugins.append(SLURMEnvironment(requeue_signal=signal.SIGUSR1)) - print("SLURM REQUEUE ENABLED") + # Preemption is handled by PreemptionCheckpoint (save last.ckpt, then + # chain to submitit's requeue). Lightning's own auto-requeue is disabled: + # it listened on SIGUSR1, which neither submitit nor Slurm ever sends. + plugins.append(SLURMEnvironment(auto_requeue=False)) + callbacks.append(PreemptionCheckpoint()) + log.info("Slurm preemption checkpointing enabled") trainer: Trainer = hydra.utils.instantiate( cfg.trainer, callbacks=callbacks, logger=logger, plugins=plugins or None ) @@ -262,8 +264,15 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: last_ckpt_path = os.path.join( trainer.default_root_dir, "checkpoints", "last.ckpt" ) - log.info("Detected SLURM requeue — resuming from 'last.ckpt'") - cfg.ckpt_path = last_ckpt_path + if os.path.exists(last_ckpt_path): + log.info(f"Detected SLURM requeue — resuming from {last_ckpt_path}") + cfg.ckpt_path = last_ckpt_path + else: + # Preempted before the first checkpoint (e.g. during dataset setup). + log.warning( + f"Detected SLURM requeue but {last_ckpt_path} does not exist; " + "starting from scratch" + ) os.makedirs(os.path.join(trainer.default_root_dir, "videos"), exist_ok=True) diff --git a/tests/unit/test_pl_model_barriers.py b/tests/unit/test_pl_model_barriers.py new file mode 100644 index 000000000..5f863c0b1 --- /dev/null +++ b/tests/unit/test_pl_model_barriers.py @@ -0,0 +1,33 @@ +"""ModelWrapper hooks must not require a torch.distributed process group. + +The CPU smoke test (``trainer.accelerator=cpu devices=1``) never initialises a +process group, and every hook that unconditionally calls +``torch.distributed.barrier()`` crashes there with +``RuntimeError: Default process group has not been initialized``. +""" + +import torch +import torch.nn as nn + +from egomimic.pl_utils.pl_model import ModelWrapper + + +class _StubAlgo: + """Minimal stand-in for an ``Algo``: just enough for ``ModelWrapper.__init__``.""" + + def __init__(self): + self.nets = nn.ModuleDict({"policy": nn.Linear(1, 1)}) + self.device = None + + +def _wrapper() -> ModelWrapper: + assert not torch.distributed.is_initialized() + return ModelWrapper(robomimic_model=_StubAlgo()) + + +def test_on_fit_start_runs_without_process_group(): + _wrapper().on_fit_start() + + +def test_on_validation_end_runs_without_process_group(): + _wrapper().on_validation_end() diff --git a/tests/unit/test_preemption.py b/tests/unit/test_preemption.py new file mode 100644 index 000000000..3a756f433 --- /dev/null +++ b/tests/unit/test_preemption.py @@ -0,0 +1,204 @@ +"""PreemptionCheckpoint: save a checkpoint on the Slurm preemption signal, then requeue. + +Under ``hydra -m`` submitit owns the preemption signal (``USR2`` by default, or +``$SUBMITIT_PREEMPT_SIGNAL``) and requeues the job *without* saving a Lightning +checkpoint. Lightning refuses to register its own handler on a signal that already +has one, so the callback installs a handler that saves first and then chains to +the handler that was installed before it (submitit's requeue). +""" + +import os +import signal + +import pytest +import torch +from lightning import Trainer +from lightning.pytorch.callbacks import Callback, ModelCheckpoint +from lightning.pytorch.demos.boring_classes import BoringModel + +from egomimic.pl_utils.preemption import PreemptionCheckpoint + +SIG = signal.SIGUSR2 + + +class _RaiseSignal(Callback): + """Deliver ``signum`` to this process from inside the training loop.""" + + def __init__(self, signum: signal.Signals): + self.signum = signum + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + if batch_idx == 1: + signal.raise_signal(self.signum) + + +def _fit(tmp_path, *callbacks, max_steps=4, ckpt_path=None): + trainer = Trainer( + default_root_dir=str(tmp_path), + max_steps=max_steps, + limit_train_batches=max_steps, + accelerator="cpu", + devices=1, + logger=False, + enable_checkpointing=any(isinstance(c, ModelCheckpoint) for c in callbacks), + enable_progress_bar=False, + enable_model_summary=False, + callbacks=[*callbacks, _RaiseSignal(SIG)], + ) + trainer.fit(BoringModel(), ckpt_path=ckpt_path) + return trainer + + +@pytest.fixture +def restore_sig(): + old = signal.getsignal(SIG) + yield + signal.signal(SIG, old) + + +def test_saves_checkpoint_then_chains_to_previous_handler(tmp_path, restore_sig): + seen = [] + prev = lambda signum, frame: seen.append(signum) # noqa: E731 + signal.signal(SIG, prev) + ckpt = tmp_path / "ckpts" / "preempt.ckpt" + + _fit(tmp_path, PreemptionCheckpoint(ckpt_path=str(ckpt), signum=SIG)) + + assert ckpt.exists(), "checkpoint was not written on the preemption signal" + assert seen == [SIG], "previous (submitit) handler was not chained" + assert signal.getsignal(SIG) is prev, "previous handler was not restored after fit" + + +def test_default_ckpt_path_is_last_ckpt_under_trainer_root(tmp_path, restore_sig): + signal.signal(SIG, lambda s, f: None) + + _fit(tmp_path, PreemptionCheckpoint(signum=SIG)) + + assert (tmp_path / "checkpoints" / "last.ckpt").exists() + + +def test_default_signal_follows_submitit_env(monkeypatch): + monkeypatch.delenv("SUBMITIT_PREEMPT_SIGNAL", raising=False) + assert PreemptionCheckpoint().signum == signal.SIGUSR2 + + monkeypatch.setenv("SUBMITIT_PREEMPT_SIGNAL", "USR1") + assert PreemptionCheckpoint().signum == signal.SIGUSR1 + + +def test_save_failure_still_chains_to_requeue(tmp_path, restore_sig, monkeypatch): + seen = [] + signal.signal(SIG, lambda signum, frame: seen.append(signum)) + + def boom(self, *a, **k): + raise OSError("disk quota exceeded") + + monkeypatch.setattr(Trainer, "save_checkpoint", boom) + cb = PreemptionCheckpoint(ckpt_path=str(tmp_path / "last.ckpt"), signum=SIG) + + _fit(tmp_path, cb) # the OSError must not escape the handler either + + assert seen == [SIG], "a failed save must not cancel the (submitit) requeue" + + +def test_no_temp_file_left_next_to_checkpoint(tmp_path, restore_sig): + signal.signal(SIG, lambda s, f: None) + ckpt = tmp_path / "last.ckpt" + + _fit(tmp_path, PreemptionCheckpoint(ckpt_path=str(ckpt), signum=SIG)) + + assert ckpt.exists() + assert not (tmp_path / "last.ckpt.preempt").exists() + + +def test_coexists_with_model_checkpoint_save_last(tmp_path, restore_sig): + """Preempt before ModelCheckpoint's first periodic save (the real config saves + every 100 epochs), then resume. ModelCheckpoint must keep writing last.ckpt + rather than versioning to last-v1.ckpt, which trainHydra never resumes from.""" + signal.signal(SIG, lambda s, f: None) + ckpt_dir = tmp_path / "checkpoints" + + # Same arguments in both fits, as in a real resume: Lightning only restores + # ModelCheckpoint state into a callback with an identical state key. + def mc(): + return ModelCheckpoint( + dirpath=str(ckpt_dir), save_last=True, every_n_train_steps=4, save_top_k=0 + ) + + # Signal fires at global_step 2, before the first periodic save at step 4. + _fit(tmp_path, PreemptionCheckpoint(signum=SIG), mc(), max_steps=3) + assert (ckpt_dir / "last.ckpt").exists() + + _fit(tmp_path, mc(), max_steps=8, ckpt_path=str(ckpt_dir / "last.ckpt")) + + files = sorted(p.name for p in ckpt_dir.glob("*.ckpt")) + assert files == ["last.ckpt"], f"ModelCheckpoint versioned last.ckpt: {files}" + step = torch.load(ckpt_dir / "last.ckpt", map_location="cpu", weights_only=False)[ + "global_step" + ] + assert step == 8, f"last.ckpt is stale (global_step={step})" + + +def test_coexists_with_model_checkpoint_through_symlinked_run_dir( + tmp_path, restore_sig +): + """Lightning realpath()s ModelCheckpoint.dirpath but only normpath()s + default_root_dir; the claim must still match through a symlink.""" + signal.signal(SIG, lambda s, f: None) + real = tmp_path / "real_run" + real.mkdir() + link = tmp_path / "link_run" + link.symlink_to(real, target_is_directory=True) + ckpt_dir = link / "checkpoints" + + def mc(): + return ModelCheckpoint( + dirpath=str(ckpt_dir), save_last=True, every_n_train_steps=4, save_top_k=0 + ) + + _fit(link, PreemptionCheckpoint(signum=SIG), mc(), max_steps=3) + _fit(link, mc(), max_steps=8, ckpt_path=str(ckpt_dir / "last.ckpt")) + + files = sorted(p.name for p in (real / "checkpoints").glob("*.ckpt")) + assert files == ["last.ckpt"], f"ModelCheckpoint versioned last.ckpt: {files}" + + +class _ForkAndSignal(Callback): + """Deliver the signal in a forked child (a DataLoader worker stand-in).""" + + def __init__(self, signum): + self.signum = signum + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + if batch_idx != 1: + return + pid = os.fork() + if pid == 0: # child: inherits the parent's handler table + signal.raise_signal(self.signum) + os._exit(0) + _, status = os.waitpid(pid, 0) + assert os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0 + + +def test_forked_worker_does_not_checkpoint(tmp_path, restore_sig): + """Slurm signals every process in the step, including DataLoader workers, + which inherit the handler. Only the process that installed it may act.""" + signal.signal(SIG, signal.SIG_IGN) + ckpt = tmp_path / "last.ckpt" + trainer = Trainer( + default_root_dir=str(tmp_path), + max_steps=4, + limit_train_batches=4, + accelerator="cpu", + devices=1, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + enable_model_summary=False, + callbacks=[ + PreemptionCheckpoint(ckpt_path=str(ckpt), signum=SIG), + _ForkAndSignal(SIG), + ], + ) + trainer.fit(BoringModel()) + + assert not ckpt.exists(), "a forked child wrote the preemption checkpoint"