Write the emergency checkpoint when dataloader workers are in use

kempnerforge/data/dataloader.py

@@ -16,6 +16,7 @@
 
 from kempnerforge.config.schema import DataConfig
 from kempnerforge.data.sampler import DistributedSampler, MixtureSampler
+from kempnerforge.resilience.signal_handler import ignore_shutdown_signals_in_worker
 
 logger = logging.getLogger(__name__)
 
@@ -58,6 +59,10 @@ def __init__(
             "prefetch_factor": config.prefetch_factor if config.num_workers > 0 else None,
             "persistent_workers": config.num_workers > 0,
             "drop_last": True,
+            # Workers must survive a group-delivered SIGTERM so the in-flight
+            # step can finish and the emergency checkpoint can be written.
+            # Never called when num_workers == 0.
+            "worker_init_fn": ignore_shutdown_signals_in_worker,
         }
         if collate_fn is not None:
             loader_kwargs["collate_fn"] = collate_fn

kempnerforge/resilience/elastic.py

@@ -17,6 +17,27 @@
 
 logger = logging.getLogger(__name__)
 
+# DCP writes this file last, once every shard is durable, so its presence is
+# the authoritative signal that a checkpoint directory is loadable.
+_DCP_METADATA_FILE = ".metadata"
+
+
+def _dcp_durable(ckpt_dir: Path) -> bool:
+    """Whether ``ckpt_dir`` holds a complete set of DCP shards.
+
+    Accepts both layouts: a flat directory, and the per-stage ``pp{k}/``
+    subdirectories written under pipeline parallelism.
+    """
+    if (ckpt_dir / _DCP_METADATA_FILE).exists():
+        return True
+    if not ckpt_dir.is_dir():
+        return False
+    return any(
+        (d / _DCP_METADATA_FILE).exists()
+        for d in ckpt_dir.iterdir()
+        if d.is_dir() and d.name.startswith("pp")
+    )
+
 
 @dataclass
 class SLURMInfo:
@@ -77,13 +98,20 @@ def resolve_resume_path(checkpoint_dir: str) -> Path | None:
 
     Checks:
       1. ``{checkpoint_dir}/latest`` symlink
-      2. Most recent ``step_N`` directory by step number
+      2. Most recent ``step_N`` directory *whose DCP shards are durable*
+
+    ``CheckpointManager`` only ever points ``latest`` at a durable checkpoint,
+    so (1) needs no further test. The ``step_N`` fallback has no such guarantee:
+    an interrupted save leaves a directory that exists but holds no ``.metadata``,
+    and resuming into it fails in ``dcp.load`` with "metadata is None". Skipping
+    incomplete directories turns an unrecoverable run into one that resumes from
+    the last durable checkpoint, whatever left the partial directory behind.
 
     Args:
         checkpoint_dir: Base checkpoint directory.
 
     Returns:
-        Path to the latest checkpoint, or None if none found.
+        Path to the latest usable checkpoint, or None if none found.
     """
     base = Path(checkpoint_dir)
     if not base.exists():
@@ -97,20 +125,25 @@ def resolve_resume_path(checkpoint_dir: str) -> Path | None:
             logger.info(f"Auto-resume: found latest checkpoint at {resolved}")
             return resolved
 
-    # Fall back to most recent step_N directory
+    # Fall back to the newest durable step_N directory, newest first.
     step_dirs = sorted(
         (
             d
             for d in base.iterdir()
             if d.is_dir() and d.name.startswith("step_") and d.name.split("_")[1].isdigit()
         ),
         key=lambda d: int(d.name.split("_")[1]),
+        reverse=True,
     )
 
-    if step_dirs:
-        path = step_dirs[-1]
-        logger.info(f"Auto-resume: found checkpoint at {path}")
-        return path
+    for path in step_dirs:
+        if _dcp_durable(path):
+            logger.info(f"Auto-resume: found checkpoint at {path}")
+            return path
+        logger.warning(
+            f"Auto-resume: skipping {path} — no DCP .metadata, so the save that "
+            f"produced it did not complete"
+        )
 
     return None
 

kempnerforge/resilience/signal_handler.py

@@ -21,6 +21,33 @@
 _SHUTDOWN_SIGNALS = (signal.SIGTERM, signal.SIGUSR1)
 
 
+def ignore_shutdown_signals_in_worker(worker_id: int) -> None:
+    """``worker_init_fn`` that makes a DataLoader worker ignore shutdown signals.
+
+    SLURM preemption and ``torchrun`` teardown signal the whole process group,
+    not a single PID, so SIGTERM reaches every DataLoader worker as well as the
+    rank's main process. Workers install no handler, take the default action and
+    die at once. The main process — which handles the signal cooperatively, by
+    setting a flag and finishing its step — then fails fetching its next
+    micro-batch, and that failure escapes before the loop reaches its
+    ``should_shutdown()`` check, so no emergency checkpoint is written.
+
+    Ignoring the shutdown signals in the worker keeps it serving long enough for
+    the step to finish, which puts the emergency save back on the loop's normal
+    path rather than on an exception path.
+
+    The cost: ``multiprocessing.Process.terminate()`` is itself SIGTERM, so a
+    shielded worker cannot be reaped that way and the interpreter's exit hook
+    would block joining it. Workers must therefore be shut down explicitly
+    through the loader's sentinel path — see ``BatchStream.close``.
+
+    Args:
+        worker_id: Worker index, supplied by ``DataLoader``. Unused.
+    """
+    for sig in _SHUTDOWN_SIGNALS:
+        signal.signal(sig, signal.SIG_IGN)
+
+
 class ShutdownHandler:
     """Cooperative shutdown handler for long-running training jobs.
 

kempnerforge/training/data_pipeline.py

@@ -27,6 +27,7 @@
 from kempnerforge.data.sampler import DistributedSampler, MixtureSampler
 from kempnerforge.distributed.utils import get_dp_info
 from kempnerforge.metrics.logger import get_logger
+from kempnerforge.resilience.signal_handler import ignore_shutdown_signals_in_worker
 from kempnerforge.training.eval import should_build_eval_dataloader
 from kempnerforge.training.runtime import RuntimeContext
 
@@ -286,6 +287,8 @@ def _build_hf_pipeline(config: JobConfig, dp_rank: int, dp_size: int) -> DataPip
             num_workers=config.data.num_workers,
             pin_memory=config.data.pin_memory,
             prefetch_factor=(config.data.prefetch_factor if config.data.num_workers > 0 else None),
+            # See StatefulDataLoader: workers survive a group-delivered SIGTERM.
+            worker_init_fn=ignore_shutdown_signals_in_worker,
         )
         logger.info(
             f"Dataset: streaming from {config.data.hf_dataset_name} "

kempnerforge/training/loop.py

@@ -80,6 +80,32 @@ def reset(self) -> None:
         """
         self._iter = None
 
+    def close(self) -> None:
+        """Shut the loader's worker processes down explicitly.
+
+        Workers ignore SIGTERM (see ``ignore_shutdown_signals_in_worker``) so a
+        preemption signal cannot kill them out from under an in-flight step.
+        The same shield means ``Process.terminate()``, which is SIGTERM, cannot
+        reap them either: left to the interpreter's exit hook, the process would
+        block forever joining daemonic children that refuse to die. Draining
+        them through the loader's own sentinel path avoids that.
+
+        Rank-local and idempotent — it issues no collectives, so it is safe on
+        the exception path in ``train_loop``'s ``finally``.
+        """
+        self._iter = None
+        self._source = None
+        loader = self.pipeline.dataloader
+        if loader is None:
+            return
+        # StatefulDataLoader wraps a torch DataLoader; a plain loader is itself.
+        inner = getattr(loader, "_dataloader", loader)
+        iterator = getattr(inner, "_iterator", None)
+        if iterator is not None and hasattr(iterator, "_shutdown_workers"):
+            iterator._shutdown_workers()
+        if hasattr(inner, "_iterator"):
+            inner._iterator = None
+
     def next_batch(self) -> dict[str, torch.Tensor]:
         if self.dataloader is None:
             raise RuntimeError("BatchStream has no dataloader; check has_data first")
@@ -651,6 +677,8 @@ def run_training_loop(
             prof.stop()
             if runtime.rank == 0:
                 print_profiler_summary(prof, trace_dir=config.profiling.trace_dir)
+        # Dataloader workers ignore SIGTERM, so nothing else will reap them.
+        session.batches.close()
 
     if completed_normally and not config.checkpoint.should_save(step):
         ckpt_mgr.save(

tests/unit/test_resilience.py

@@ -19,7 +19,10 @@
     NaNDetector,
     check_gpu_health,
 )
-from kempnerforge.resilience.signal_handler import ShutdownHandler
+from kempnerforge.resilience.signal_handler import (
+    ShutdownHandler,
+    ignore_shutdown_signals_in_worker,
+)
 
 # ---------------------------------------------------------------------------
 # ShutdownHandler
@@ -493,6 +496,24 @@ def test_log_job_info_requeued(self):
                     os.environ.pop(k, None)
 
 
+def _durable_step_dir(base, step: int, pp_stages: int = 0):
+    """Create a ``step_N`` directory that looks like a completed DCP save.
+
+    ``resolve_resume_path`` treats DCP's ``.metadata`` as the durability marker,
+    so a directory without one stands for an interrupted save.
+    """
+    d = base / f"step_{step}"
+    d.mkdir()
+    if pp_stages:
+        for k in range(pp_stages):
+            stage = d / f"pp{k}"
+            stage.mkdir()
+            (stage / ".metadata").write_text("")
+    else:
+        (d / ".metadata").write_text("")
+    return d
+
+
 class TestResumePathResolution:
     def test_no_checkpoint_dir(self, tmp_path):
         assert resolve_resume_path(str(tmp_path / "nonexistent")) is None
@@ -517,15 +538,14 @@ def test_latest_symlink(self, tmp_path):
     def test_fallback_to_highest_step(self, tmp_path):
         # Create step directories without "latest" symlink
         for step in [10, 50, 30]:
-            d = tmp_path / f"step_{step}"
-            d.mkdir()
+            _durable_step_dir(tmp_path, step)
 
         result = resolve_resume_path(str(tmp_path))
         assert result is not None
         assert result.name == "step_50"
 
     def test_ignores_non_step_dirs(self, tmp_path):
-        (tmp_path / "step_100").mkdir()
+        _durable_step_dir(tmp_path, 100)
         (tmp_path / "other_dir").mkdir()
         (tmp_path / "step_abc").mkdir()  # Not a valid step dir (will cause error)
 
@@ -538,8 +558,7 @@ def test_ignores_non_step_dirs(self, tmp_path):
 
     def test_latest_broken_symlink_falls_back(self, tmp_path):
         # Create a checkpoint directory
-        step_dir = tmp_path / "step_50"
-        step_dir.mkdir()
+        _durable_step_dir(tmp_path, 50)
 
         # Create broken "latest" symlink pointing to nonexistent
         latest = tmp_path / "latest"
@@ -558,7 +577,7 @@ def test_latest_broken_symlink_falls_back(self, tmp_path):
     def test_latest_symlink_takes_priority_over_higher_step(self, tmp_path):
         """latest symlink should be used even if higher step dirs exist."""
         for step in [10, 50, 100]:
-            (tmp_path / f"step_{step}").mkdir()
+            _durable_step_dir(tmp_path, step)
 
         # Point latest at step_50 (not the highest)
         latest = tmp_path / "latest"
@@ -570,7 +589,7 @@ def test_latest_symlink_takes_priority_over_higher_step(self, tmp_path):
 
     def test_single_step_dir(self, tmp_path):
         """Works with exactly one step directory."""
-        (tmp_path / "step_1").mkdir()
+        _durable_step_dir(tmp_path, 1)
 
         result = resolve_resume_path(str(tmp_path))
         assert result is not None
@@ -579,9 +598,102 @@ def test_single_step_dir(self, tmp_path):
     def test_step_dirs_with_large_numbers(self, tmp_path):
         """Handles large step numbers correctly (sorts numerically, not lexically)."""
         for step in [9, 100, 1000, 20]:
-            (tmp_path / f"step_{step}").mkdir()
+            _durable_step_dir(tmp_path, step)
 
         result = resolve_resume_path(str(tmp_path))
         assert result is not None
         # Numerically highest is 1000, not lexically highest "step_9"
         assert result.name == "step_1000"
+
+    # -- durability of the step_N fallback (#177) --------------------------
+
+    def test_skips_incomplete_newest_and_falls_back(self, tmp_path):
+        """An interrupted save must not be selected over an older durable one.
+
+        This is the #177 failure: the emergency save dies part-way, leaving a
+        step dir with no DCP `.metadata`, and auto-resume picks it and fails in
+        dcp.load with "metadata is None".
+        """
+        _durable_step_dir(tmp_path, 10)
+        (tmp_path / "step_20").mkdir()  # interrupted save: no .metadata
+
+        result = resolve_resume_path(str(tmp_path))
+        assert result is not None
+        assert result.name == "step_10"
+
+    def test_returns_none_when_no_step_dir_is_durable(self, tmp_path):
+        """Better to start fresh than to resume into a directory that cannot load."""
+        for step in [10, 20]:
+            (tmp_path / f"step_{step}").mkdir()
+
+        assert resolve_resume_path(str(tmp_path)) is None
+
+    def test_accepts_pipeline_parallel_layout(self, tmp_path):
+        """Under PP each stage writes its own pp{k}/.metadata, not a flat one."""
+        _durable_step_dir(tmp_path, 30, pp_stages=2)
+
+        result = resolve_resume_path(str(tmp_path))
+        assert result is not None
+        assert result.name == "step_30"
+
+    def test_skips_incomplete_pipeline_parallel_dir(self, tmp_path):
+        """A pp{k}/ subdir without .metadata is still an interrupted save."""
+        _durable_step_dir(tmp_path, 10, pp_stages=2)
+        partial = tmp_path / "step_20"
+        (partial / "pp0").mkdir(parents=True)  # stage dir exists, never finished
+
+        result = resolve_resume_path(str(tmp_path))
+        assert result is not None
+        assert result.name == "step_10"
+
+
+# ---------------------------------------------------------------------------
+# DataLoader worker signal shielding (#177)
+# ---------------------------------------------------------------------------
+
+
+class TestWorkerSignalShielding:
+    """SIGTERM is delivered to the process group, so it reaches DataLoader
+    workers too. Workers that die take the emergency checkpoint with them:
+    the main process fails fetching its next batch and raises before the loop
+    reaches its should_shutdown() check.
+    """
+
+    def _restoring(self, fn):
+        """Run fn with the process's shutdown-signal handlers restored after."""
+        saved = {s: signal.getsignal(s) for s in (signal.SIGTERM, signal.SIGUSR1)}
+        try:
+            fn()
+        finally:
+            for s, h in saved.items():
+                signal.signal(s, h)
+
+    def test_ignores_both_shutdown_signals(self):
+        def check():
+            ignore_shutdown_signals_in_worker(0)
+            assert signal.getsignal(signal.SIGTERM) is signal.SIG_IGN
+            assert signal.getsignal(signal.SIGUSR1) is signal.SIG_IGN
+
+        self._restoring(check)
+
+    def test_sigterm_does_not_kill_the_caller(self):
+        """The point of the shield: the signal arrives and is survived."""
+
+        def check():
+            ignore_shutdown_signals_in_worker(0)
+            os.kill(os.getpid(), signal.SIGTERM)  # would terminate by default
+
+        self._restoring(check)
+
+    def test_dataloader_installs_the_shield(self):
+        """StatefulDataLoader must wire the shield, not just export it."""
+        from torch.utils.data import TensorDataset
+
+        from kempnerforge.config.schema import DataConfig
+        from kempnerforge.data.dataloader import StatefulDataLoader
+
+        dataset = TensorDataset(torch.arange(8).float().unsqueeze(1))
+        loader = StatefulDataLoader(
+            dataset, batch_size=2, config=DataConfig(num_workers=2, pin_memory=False)
+        )
+        assert loader._dataloader.worker_init_fn is ignore_shutdown_signals_in_worker