Lightning-AI / Lightning-AI/pytorch-lightning

pytorch lightning causes slurm nodes to drain

Open
#15,008 14 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

environment: slurm question
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.