Lightning-AI/pytorch-lightning

Allow obtaining num_nodes from ClusterEnvironment

Open

#7,361 opened on May 4, 2021

View on GitHub
 (7 comments) (0 reactions) (0 assignees)Python (3,233 forks)batch import
environmentfeaturegood first issuehelp wanted

Repository metrics

Stars
 (26,687 stars)
PR merge metrics
 (Avg merge 9d 15h) (3 merged PRs in 30d)

Description

🚀 Feature

num_nodes must currently be specified manually by the user. However, the number of nodes is generally known in a cluster environment [1] and could be provided by and initialized from ClusterEnvironment

[1] For example $AWS_BATCH_JOB_NUM_NODES in https://docs.aws.amazon.com/batch/latest/userguide/multi-node-parallel-jobs.html

Pitch

https://github.com/PyTorchLightning/pytorch-lightning/blob/763a9a9495977b23cbd6a57f10253b662fd592a5/pytorch_lightning/plugins/training_type/ddp.py#L62-L75

could be updated to initialize num_nodes from ClusterEnvironment if ClusterEnvironment is provided and implements a num_nodes method.

cc @borda @awaelchli @ananthsub

Contributor guide