From a1ecae92001f16fe8add1a4054e5da091bbb84fb Mon Sep 17 00:00:00 2001 From: naeemkh Date: Mon, 14 Sep 2026 14:22:07 -0400 Subject: [PATCH] Write the emergency checkpoint when num_workers > 0 --- kempnerforge/data/dataloader.py | 5 + kempnerforge/resilience/elastic.py | 47 ++++++-- kempnerforge/resilience/signal_handler.py | 27 +++++ kempnerforge/training/data_pipeline.py | 3 + kempnerforge/training/loop.py | 28 +++++ tests/unit/test_resilience.py | 130 ++++++++++++++++++++-- tests/unit/test_train_entry.py | 48 ++++++++ 7 files changed, 272 insertions(+), 16 deletions(-) diff --git a/kempnerforge/data/dataloader.py b/kempnerforge/data/dataloader.py index b3663df..fbf7d08 100644 --- a/kempnerforge/data/dataloader.py +++ b/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 diff --git a/kempnerforge/resilience/elastic.py b/kempnerforge/resilience/elastic.py index 8f34c3d..9cb1eb2 100644 --- a/kempnerforge/resilience/elastic.py +++ b/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,7 +125,7 @@ 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 @@ -105,12 +133,17 @@ def resolve_resume_path(checkpoint_dir: str) -> Path | None: 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 diff --git a/kempnerforge/resilience/signal_handler.py b/kempnerforge/resilience/signal_handler.py index 07d37d7..7d84c33 100644 --- a/kempnerforge/resilience/signal_handler.py +++ b/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. diff --git a/kempnerforge/training/data_pipeline.py b/kempnerforge/training/data_pipeline.py index ca2e48b..b50c36f 100644 --- a/kempnerforge/training/data_pipeline.py +++ b/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} " diff --git a/kempnerforge/training/loop.py b/kempnerforge/training/loop.py index c52d859..04a9d9d 100644 --- a/kempnerforge/training/loop.py +++ b/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( diff --git a/tests/unit/test_resilience.py b/tests/unit/test_resilience.py index 7b82326..fe6e43e 100644 --- a/tests/unit/test_resilience.py +++ b/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 diff --git a/tests/unit/test_train_entry.py b/tests/unit/test_train_entry.py index 3510d7a..354cee6 100644 --- a/tests/unit/test_train_entry.py +++ b/tests/unit/test_train_entry.py @@ -234,6 +234,54 @@ def test_reset_forces_a_fresh_iterator(self): stream.reset() assert int(stream.next_batch()["input_ids"]) == 0 + # -- worker teardown (#177) ------------------------------------------- + + def test_close_is_safe_without_a_loader(self): + BatchStream(DataPipeline()).close() # must not raise + + def test_close_is_safe_on_a_worker_free_loader(self): + stream = BatchStream(DataPipeline(dataloader=_fake_batches(3))) + stream.next_batch() + stream.close() + assert stream._iter is None + + def test_close_drains_the_workers_explicitly(self): + """Shielded workers ignore SIGTERM, so terminate() cannot reap them. + + close() must go through the loader's own sentinel path instead, or the + interpreter blocks at exit joining daemonic children that refuse to die. + """ + + class _FakeIterator: + def __init__(self): + self.shutdown_calls = 0 + + def _shutdown_workers(self): + self.shutdown_calls += 1 + + class _FakeInner: + def __init__(self): + self._iterator = _FakeIterator() + + class _FakeLoader: + def __init__(self): + self._dataloader = _FakeInner() + + loader = _FakeLoader() + iterator = loader._dataloader._iterator + stream = BatchStream(DataPipeline(dataloader=loader)) + + stream.close() + + assert iterator.shutdown_calls == 1 + assert loader._dataloader._iterator is None + + def test_close_is_idempotent(self): + stream = BatchStream(DataPipeline(dataloader=_fake_batches(3))) + stream.next_batch() + stream.close() + stream.close() # a second pass must not raise + def test_ensure_started_is_idempotent(self): stream = BatchStream(DataPipeline(dataloader=_fake_batches(3))) stream.ensure_started()