Lightning-AI / Lightning-AI/pytorch-lightning
pytorch lightning causes slurm nodes to drain
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
Hello! When I train with DDP strategy, any type of crashes like `Out Of Memory (OOM)` error or `scancel` slurm job results in slurm nodes to drain due to `Kill task failed` which means that the pytorch lightning process running on these nodes failed to clean up after termination. I was wondering how I could fix this?
### How to reproduce the bug
```python
callbacks_list = [
lr_monitor,
checkpoint_callback,
swa_ensemble,
PrintCallback(),
callbacks.ModelSummary(),
callbacks.DeviceStatsMonitor(cpu_stats=True),
]
logger = [
loggers.tensorboard.TensorBoardLogger(save_dir="./logs", version=slurm_job_id),
]
trainer = pl.Trainer(
gpus=_cfg.slurm_job.gpus_per_node,
num_nodes=_cfg.slurm_job.number_of_nodes,
accelerator="gpu",
strategy=DDPStrategy(find_unused_parameters=False),
plugins=pl.plugins.environments.SLURMEnvironment(auto_requeue=False),
max_epochs=_cfg.training.fit.epochs,
callbacks=callbacks_list,
logger=logger,
accumulate_grad_batches=_cfg.training.fit.accumulation_steps,
profiler="simple",
)
data_module = DataModule(_cfg)
model_module = LitModelModule(_cfg)
trainer.fit(
model=model_module,
datamodule=data_module,
ckpt_path=None
)
```
### Error messages and logs
```
# Error messages and logs here please
```
### Environment
```
- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): Trainer, LightningModule
- PyTorch Lightning Version (e.g., 1.5.0): 1.7.7
- PyTorch Version (e.g., 1.10): 1.12.0
- Python version (e.g., 3.9): 3.9.12
- OS (e.g., Linux): Linux --RHEL7.4
- CUDA/cuDNN version: 11.7
- GPU models and configuration: RTX 5000's
- How you installed Lightning(`conda`, `pip`, source): pip
-
```
### More info
_No response_
cc @awaelchli
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
The report names no repository files or tests. Start by reproducing the DDP run with Trainer and SLURMEnvironment on the stated versions, then inspect termination handling around SLURMEnvironment and the Trainer process lifecycle. Done means OOM or scancel exits without triggering “Kill task failed” or draining the SLURM nodes.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- linux, python, pytorch
- Domain
- devops, distributed-systems, infrastructure
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100