Conversation
Naeemkh
changed the title
Codecov Report❌ Patch coverage is
🚀 New features to boost your workflow:
|
|
Closes #177 |
Naeemkh
changed the title
| 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, |
There was a problem hiding this comment.
Better to pass the persistent_workers true when the number of workers are more than 0. persistent_workers=config.data.num_workers > 0
persistent_workers (bool, optional) – If True, the data loader will not shut down the worker processes after a dataset has been consumed once. This allows to maintain the workers Dataset instances alive. (default: False)
| return True | ||
| if not ckpt_dir.is_dir(): | ||
| return False | ||
| return any( |
There was a problem hiding this comment.
any calls a PP checkpoint durable when only some stages flushed, which is what
an interrupted PP save looks like (all ranks mkdir their own pp{k}/ up front).
_resolve_dcp_load_dir then re-tests per rank, so stage 0 resumes at step_20 and
stage 1 at step_10, silently.
| f"produced it did not complete" | ||
| ) | ||
|
|
||
| return None |
There was a problem hiding this comment.
Before this PR the function always returned a path when step_N directories existed. Now, if none of them looks durable, it returns nothing, and the entry path reads that as "no checkpoint to resume" and trains from scratch with no error.
| path = step_dirs[-1] | ||
| logger.info(f"Auto-resume: found checkpoint at {path}") | ||
| return path | ||
| for path in step_dirs: |
There was a problem hiding this comment.
Possible hang: Every rank now stats .metadata, which only the DCP coordinator writes, with no broadcast. The old name-based fallback always agreed across ranks. If one rank resolves None it returns early at entry.py:207 and skips the collective dcp.load while the others block in it, hanging the job at startup.
| Args: | ||
| worker_id: Worker index, supplied by ``DataLoader``. Unused. | ||
| """ | ||
| for sig in _SHUTDOWN_SIGNALS: |
There was a problem hiding this comment.
in close(), w.kill() any worker still alive after _shutdown_workers().
2 participants
Summary
When SLURM preempts a job it signals every process in the group, not just the main one. That kills the dataloader's worker processes first, so the main process fails trying to fetch its next batch before the shutdown handler ever ran. No emergency checkpoint was written and nothing appeared in the log. With
data.num_workers = 0the same drill worked, which is what pinned the cause. Workers now ignore the shutdown signals, so the step in progress can finish and the checkpoint is written on the normal path. That has a side effect worth knowing: the usual way to stop a worker is also a signal, so a shielded worker can't be stopped that way and the process would hang on exit waiting for them. They are now shut down explicitly when the training loop ends.Also, on restart, auto-resume picked the newest checkpoint folder without checking whether it had finished being written. A run whose only checkpoint was a failed emergency save resumed into it and died. It now skips unfinished folders and takes the newest complete one. #194 already made the loader recover from this, but the resume path is also used to read the saved step before loading, and that read doesn't notice an unfinished folder, so resolving to a complete one closes the remaining gap.
Testing
uv run ruff check kempnerforge/ tests/passesuv run ruff format --check kempnerforge/ tests/ scripts/passesuv run pyright kempnerforge/passes (0 errors)uv run pytest tests/unit/ -v --timeout=60passesuv run torchrun --nproc_per_node=4 -m pytest tests/distributed/ -vuv run pytest tests/e2e/ --e2e -vCloses #