Lightning-AI / Lightning-AI/pytorch-lightning
Cannot run in SLURM Interactive Session
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
I'm trying to run pytorch lightning on the SLURM cluster.
It runs into MPI initialization errors first and after specifying SLURM job name to 'bash', as suggested in this issue https://github.com/Lightning-AI/pytorch-lightning/issues/16730 , I can successfully run my scripts using `sbatch`. However, I still can't run the script in interactive sessions (both 'bash' or 'interactive' job_names failed).
### What version are you seeing the problem on?
v2.2
### How to reproduce the bug
```python
import pytorch_lightning as pl
import numpy as np
import torch
from torch.nn import MSELoss
from torch.optim import Adam
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
class SimpleDataset(Dataset):
def __init__(self):
X = np.arange(10000)
y = X * 2
X = [[_] for _ in X]
y = [[_] for _ in y]
self.X = torch.Tensor(X)
self.y = torch.Tensor(y)
def __len__(self):
return len(self.y)
def __getitem__(self, idx):
return {"X": self.X[idx], "y": self.y[idx]}
class MyModel(pl.LightningModule):
def __init__(self):
super().__init__()
self.fc = nn.Linear(1, 1)
self.criterion = MSELoss()
def forward(self, inputs_id, labels=None):
outputs = self.fc(inputs_id)
loss = 0
if labels is not None:
loss = self.criterion(outputs, labels)
return loss, outputs
def train_dataloader(self):
dataset = SimpleDataset()
return DataLoader(dataset, batch_size=1000)
def training_step(self, batch, batch_idx):
input_ids = batch["X"]
labels = batch["y"]
loss, outputs = self(input_ids, labels)
return {"loss": loss}
def configure_optimizers(self):
optimizer = Adam(self.parameters())
return optimizer
if __name__ == '__main__':
model = MyModel()
print('Starts trainer initialization')
trainer = pl.Trainer(max_epochs=2000, accelerator='gpu')
trainer.fit(model)
X = torch.Tensor([[1.0], [51.0], [89.0]])
_, y = model(X)
print(y)
# python src/scripts/pl_test.py
```
### Error messages and logs
Bash commands and errors
```bash
# Please note that I've alreday loaded openmpi/4.0.4.
(pl_dbg) jianan.zhao@cn-g009:~/scratch/INC$ python src/scripts/pl_test.py
Starts trainer initialization
/home/mila/j/jianan.zhao/scratch/miniconda3/envs/pl_dbg/lib/python3.10/site-packages/lightning_fabric/plugins/environments/slurm.py:204: The `srun` command is available on your system but is not used. HINT: If your intention is to run Lightning on SLURM, prepend your python command with `srun` like so: srun python src/scripts/pl_test.py ...
[cn-g009.server.mila.quebec:910946] OPAL ERROR: Unreachable in file pmix3x_client.c at line 111
--------------------------------------------------------------------------
The application appears to have been direct launched using "srun",
but OMPI was not built with SLURM's PMI support and therefore cannot
execute. There are several options for building PMI support under
SLURM, depending upon the SLURM version you are using:
version 16.05 or later: you can use SLURM's PMIx support. This
requires that you configure and build SLURM --with-pmix.
Versions earlier than 16.05: you must use either SLURM's PMI-1 or
PMI-2 support. SLURM builds PMI-1 by default, or you can manually
install PMI-2. You must then build Open MPI using --with-pmi pointing
to the SLURM PMI library location.
Please configure as appropriate and try again.
--------------------------------------------------------------------------
*** An error occurred in MPI_Init_thread
*** on a NULL communicator
*** MPI_ERRORS_ARE_FATAL (processes in this communicator will now abort,
*** and potentially your MPI job)
[cn-g009.server.mila.quebec:910946] Local abort before MPI_INIT completed completed successfully, but am not able to aggregate error messages, and not able to guarantee that all other processes were killed!
(pl_dbg) jianan.zhao@cn-g009:~/scratch/INC$ echo $SLURM_JOB_NAME
bash
```
### Environment
Current environment
* CUDA:
- GPU:
- NVIDIA A100-SXM4-80GB MIG 2g.20gb
- available: True
- version: 12.1
* Lightning:
- lightning: 2.2.1
- lightning-utilities: 0.11.2
- pytorch-lightning: 2.2.1
- torch: 2.2.1
- torchaudio: 2.2.1
- torchdata: 0.7.1
- torchmetrics: 1.3.2
- torchvision: 0.17.1
* Packages:
- aiohttp: 3.9.3
- aiosignal: 1.3.1
- annotated-types: 0.6.0
- antlr4-python3-runtime: 4.9.3
- anyio: 4.3.0
- appdirs: 1.4.4
- argon2-cffi: 23.1.0
- argon2-cffi-bindings: 21.2.0
- arrow: 1.3.0
- asttokens: 2.4.1
- async-lru: 2.0.4
- async-timeout: 4.0.3
- attrs: 23.2.0
- babel: 2.14.0
- beautifulsoup4: 4.12.3
- bleach: 6.1.0
- brotli: 1.1.0
- cached-property: 1.5.2
- certifi: 2024.2.2
- cffi: 1.16.0
- charset-normalizer: 3.3.2
- click: 8.1.7
- colorama: 0.4.6
- colorlog: 6.8.2
- comm: 0.2.2
- datasets: 2.18.0
- debugpy: 1.8.1
- decorator: 5.1.1
- defusedxml: 0.7.1
- dgl: 2.1.0+cu118
- dill: 0.3.8
- docker-pycreds: 0.4.0
- easydict: 1.13
- einops: 0.7.0
- entrypoints: 0.4
- exceptiongroup: 1.2.0
- executing: 2.0.1
- fastjsonschema: 2.19.1
- filelock: 3.13.3
- fqdn: 1.5.1
- frozenlist: 1.4.1
- fsspec: 2024.2.0
- gitdb: 4.0.11
- gitpython: 3.1.43
- gmpy2: 2.1.2
- h11: 0.14.0
- h2: 4.1.0
- hpack: 4.0.0
- httpcore: 1.0.5
- httpx: 0.27.0
- huggingface-hub: 0.22.2
- hydra-colorlog: 1.2.0
- hydra-core: 1.3.2
- hyperframe: 6.0.1
- idna: 3.6
- importlib-metadata: 7.1.0
- importlib-resources: 6.4.0
- ipykernel: 6.29.3
- ipython: 8.22.2
- ipywidgets: 8.1.2
- isoduration: 20.11.0
- jedi: 0.19.1
- jinja2: 3.1.3
- joblib: 1.3.2
- json5: 0.9.24
- jsonpointer: 2.4
- jsonschema: 4.21.1
- jsonschema-specifications: 2023.12.1
- jupyter: 1.0.0
- jupyter-client: 8.6.1
- jupyter-console: 6.6.3
- jupyter-core: 5.7.2
- jupyter-events: 0.10.0
- jupyter-lsp: 2.2.4
- jupyter-server: 2.13.0
- jupyter-server-terminals: 0.5.3
- jupyterlab: 4.1.5
- jupyterlab-pygments: 0.3.0
- jupyterlab-server: 2.25.4
- jupyterlab-widgets: 3.0.10
- lightning: 2.2.1
- lightning-utilities: 0.11.2
- littleutils: 0.2.2
- markdown-it-py: 3.0.0
- markupsafe: 2.1.5
- matplotlib-inline: 0.1.6
- mdurl: 0.1.2
- mistune: 3.0.2
- mpi4py: 3.1.5
- mpmath: 1.3.0
- multidict: 6.0.5
- multiprocess: 0.70.16
- nbclient: 0.10.0
- nbconvert: 7.16.3
- nbformat: 5.10.4
- nest-asyncio: 1.6.0
- networkx: 3.3
- notebook: 7.1.2
- notebook-shim: 0.2.4
- numpy: 1.26.4
- nvidia-htop: 1.2.0
- ogb: 1.3.6
- omegaconf: 2.3.0
- outdated: 0.2.2
- overrides: 7.7.0
- packaging: 24.0
- pandas: 2.2.1
- pandocfilters: 1.5.0
- parso: 0.8.4
- pathtools: 0.1.2
- pexpect: 4.9.0
- pickleshare: 0.7.5
- pillow: 9.4.0
- pip: 24.0
- pkgutil-resolve-name: 1.3.10
- platformdirs: 4.2.0
- prometheus-client: 0.20.0
- prompt-toolkit: 3.0.42
- protobuf: 4.25.3
- psutil: 5.9.8
- ptyprocess: 0.7.0
- pure-eval: 0.2.2
- pyarrow: 15.0.2
- pyarrow-hotfix: 0.6
- pycparser: 2.22
- pydantic: 2.6.4
- pydantic-core: 2.16.3
- pygments: 2.17.2
- pysocks: 1.7.1
- python-dateutil: 2.9.0
- python-dotenv: 1.0.1
- python-json-logger: 2.0.7
- pytorch-lightning: 2.2.1
- pytz: 2024.1
- pyyaml: 6.0.1
- pyzmq: 25.1.2
- qtconsole: 5.5.1
- qtpy: 2.4.1
- referencing: 0.34.0
- regex: 2023.12.25
- requests: 2.31.0
- rfc3339-validator: 0.1.4
- rfc3986-validator: 0.1.1
- rich: 13.7.1
- rootutils: 1.0.7
- rpds-py: 0.18.0
- safetensors: 0.4.2
- scikit-learn: 1.4.1.post1
- scipy: 1.13.0
- send2trash: 1.8.2
- sentry-sdk: 1.44.1
- setproctitle: 1.3.3
- setuptools: 69.2.0
- six: 1.16.0
- smmap: 5.0.0
- sniffio: 1.3.1
- soupsieve: 2.5
- stack-data: 0.6.2
- sympy: 1.12
- termcolor: 2.4.0
- terminado: 0.18.1
- threadpoolctl: 3.4.0
- tinycss2: 1.2.1
- tokenizers: 0.15.2
- tomli: 2.0.1
- torch: 2.2.1
- torchaudio: 2.2.1
- torchdata: 0.7.1
- torchmetrics: 1.3.2
- torchvision: 0.17.1
- tornado: 6.4
- tqdm: 4.66.2
- traitlets: 5.14.2
- transformers: 4.39.3
- triton: 2.2.0
- types-python-dateutil: 2.9.0.20240316
- typing-extensions: 4.11.0
- typing-utils: 0.1.0
- tzdata: 2024.1
- uri-template: 1.3.0
- urllib3: 2.2.1
- wandb: 0.16.5
- wcwidth: 0.2.13
- webcolors: 1.13
- webencodings: 0.5.1
- websocket-client: 1.7.0
- wheel: 0.43.0
- widgetsnbextension: 4.0.10
- xxhash: 3.4.1
- yarl: 1.9.4
- zipp: 3.17.0
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.10.14
- release: 5.15.0-101-generic
- version: #111-Ubuntu SMP Tue Mar 5 20:16:58 UTC 2024
### More info
_No response_
cc @lantiga
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
Start with lightning_fabric/plugins/environments/slurm.py, especially the warning shown at line 204, and reproduce the command from the issue inside an interactive SLURM session. Compare the `python` and `srun python` entry points and the reported MPI initialization failure; done means the supplied script runs successfully in the interactive session without that failure.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100