Repository: Lightning-AI/pytorch-lightning
Pull request: support launching Lightning ddp with traditional command
## What does this PR do?

Fixes #7003

The Lightning environment provides a convenient way to launch DDP multi-gpu experiments. It launches the required number of processes automatically under the hood as explained in the docs. However, there is currently only a hacky way for the user to prevent this if they wish to launch all processes manually through the command line or with utilities like `torch.distributed.launch`. This PR adds detection of `LOCAL_RANK` variable in the environment and determines accordingly if processes need to be launched or not.


## Before submitting
- [x] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [x] Did you make sure to update the documentation with your changes? (if necessary)
- [x] Did you write any new necessary tests? (not for typos and docs)
- [x] Did you verify new and existing tests pass locally with your changes?
- [x] Did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

 - [x] Is this pull request ready for review? (if not, please submit in draft mode)
 - [x] Check that all items from **Before submitting** are resolved
 - [x] Make sure the title is self-explanatory and the description concisely explains the PR
 - [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?
I made sure I had fun coding 🙃


Problem statement:
PyTorch Lightning ignores traditional WORLD_SIZE/RANK specifications in environment and doesn't document replacement
## 🐛 Bug

Standard torch.distributed environment variables seem to be handled differently.

### To Reproduce

```
$ MASTER_ADDR=192.168.1.3 MASTER_PORT=1234 WORLD_SIZE=2 RANK=0 python3 boring_model.py
```

### Expected behavior

Should wait for second job to connect to the MASTER_PORT. Instead, just starts training. No combination of arguments and environment variables seems to change this behavior.

It appears that PL handles startup for distributed processing differently somehow and internally overrides environment variables. This works pretty well for ddp on single nodes and (presumably) automates deployment under Slurm.

But I can't figure out from the documentation or from the samples how to arrange for the traditional PyTorch behavior in which all jobs are started up manually. 

For comparison, this code behaves as expected:

```
$ cat > simple.py
import torch

print("init")
torch.distributed.init_process_group("gloo")
print("done", torch.distributed.get_rank(), torch.distributed.get_world_size())
$ MASTER_ADDR=192.168.1.3 MASTER_PORT=1234 WORLD_SIZE=2 RANK=0 python3 simple.py & sleep 3; MASTER_ADDR=192.168.1.3 MASTER_PORT=1234 WORLD_SIZE=2 RANK=1 python3 simple.py
```

That is, the two jobs rendezvous as expected and then exit.

### Suggested Behavior

- document how the traditional PyTorch behavior can be reproduced (i.e., no calculations of nodes/size/... in DDP)
- perhaps provide some kind of command line flag that restores the traditional behavior (e.g., "--accelerator ddp_plain")
- maybe something else

### Environment

 - PyTorch Version (e.g., 1.0): 1.8.1
 - OS (e.g., Linux): Ubuntu 20.04
 - How you installed PyTorch (`conda`, `pip`, source): pip3
 - Build command you used (if compiling from source):
 - Python version: 3.8
 - CUDA/cuDNN version: 11.1
 - GPU models and configuration: 1080
 - Any other relevant information:

### Additional context

<!-- Add any other context about the problem here. -->

Hints from issue discussion:
Hi, did you follow these steps here?
https://pytorch-lightning.readthedocs.io/en/latest/advanced/cluster.html

I think you can also force the behavior you want like so:


```python
from pytorch_lightning.plugins.environments import LightningEnvironment


class MyCluster(LightningEnvironment):

    def creates_children(self):
        return True

trainer = Trainer(plugins=[MyCluster()])
```


Thanks. That page says to pass num_nodes explicitly to Trainer, which I can do. But that still doesn't work. I.e., if I say --num_nodes=2 on both nodes, each node believes its GLOBAL_RANK=0, even if I pass GLOBAL_RANK=1, RANK=1, etc. in the environment. They seem to connect to the master node, but nothing starts. So, I simply don't know how to get two nodes to rendezvous or even change the GLOBAL_RANK of any node.

As for your sugggested solution,  LightningEnvironment doesn't exist in 1.2.7, but apparently has been added in 1.2.8. I have added my own implementation in 1.2.7:

```Python
def I(value, name):
    print(name, value)
    return value


class MyCluster(environments.ClusterEnvironment):
    def __init__(self):
        super().__init__()
        self._master_addr = I(os.environ["MASTER_ADDR"], "MASTER_ADDR")
        self._master_port = I(int(os.environ["MASTER_PORT"]), "MASTER_PORT")
        self._world_size = I(int(os.environ["WORLD_SIZE"]), "WORLD_SIZE")
        self._global_rank = I(int(os.environ["RANK"]), "RANK")

    def creates_childern(self):
        return False

    def master_address(self):
        return self._master_addr

    def master_port(self):
        return self._master_port

    def world_size(self):
        return I(self._world_size, "world_size")

    def global_rank(self):
        return self._global_rank

    def set_world_size(self, *args):
        raise Exception()

    def set_global_rank(self, *args):
        raise Exception()

    def local_rank(self):
        return I(0, "LOCAL_RANK")

    def node_rank(self):
        return self._global_rank
```

Unfortunately, this doesn't quite work either. It appears that I also need to set `num_nodes` to `world_size` as well. But even doing that, I can't get the jobs to rendezvous and train.

So, in summary, DDP training works great single node, letting LIghtning spawn the subprocesses. I don't have access to Slurm, so I can't test that.

And I don't have a working solution based on Lightning for manuall starting up multiple jobs and having them all connect to a master node, like I used to be able to with regular torch.distributed.

Any help would be appreciated.
I believe you need to return True from creates_children, because otherwise Lightning will spawn the processes. But as you say, you want to start them manually on both nodes.
Sorry that there is not much documentation. These environments are new, we are working on better documentation.
@awaelchli the `creates_children` naming is confusing. many readers would assume its the environment within lightning that creates children processes vs processes being created before trainer.fit()
Thanks; I have tried returning both values, and neither works. Even if I stop ddp.py from creating subprocesses, I still don't know how to match the actual configuration of nodes and GPUs that I have. And even if I could somehow communicate what I want to it, it wouldn't be using the standard PyTorch distributed environment variables, which Lightning seems to use differently. 

It looks to me that in order to fit Lightning into an existing distributed training environment, I actually need both a new Environment class and a new `training_type`. The `training_type` should just wrap the model into DDP, but otherwise leave things alone: no subprocesses, no modification of the environment, etc.

I suspect I'm not the only user who has existing non-Lightning distributed code for which I need a solution that doesn't break/override the existing distributed startup code.

I'll investigate a bit further. 
@tmbdev the torchelastic environment in Lightning does not create subprocesses nor does it modify the environment; ddp+torchelastic does just wrap the model in DDP. Does this launcher work for you? https://pytorch.org/elastic/0.2.2/distributed.html
@tmbdev In Lightning you don't need to set the global, local and world size. It gets computed automatically. You only have to launch the script on every node and set the NODE_RANK like so (and set the `num_nodes` trainer argument):

`NODE_RANK=1 MASTER_ADDR=192.168.1.3 MASTER_PORT=1234  python your_script.py`

> @awaelchli the creates_children naming is confusing. many readers would assume its the environment within lightning that creates children processes vs processes being created before trainer.fit()

Apologies, this was introduced recently by me. We can change the name when we do the property/setter refactor. #6303 
Thanks. I got something to work like this, but only on a local machine

```
# first job
$ CUDA_VISIBLE_DEVICES=0 MASTER_ADDR=127.0.0.1 MASTER_PORT=9888 NODE_RANK=0 python3 train.py --num_nodes=2 --gpus 1 --accelerator=ddp -b 32
# different window same host
$ CUDA_VISIBLE_DEVICES=1 MASTER_ADDR=127.0.0.1 MASTER_PORT=9888 NODE_RANK=1 python3 train.py --num_nodes=2 --gpus 1 --accelerator=ddp -b 32
```

When I use it on different hosts, things still fail:

```
192.168.1.3$ PL_TORCH_DISTRIBUTED_BACKEND=gloo MASTER_ADDR=192.168.1.3 MASTER_PORT=9888 NODE_RANK=0 python3 train.py --num_nodes=2 --gpus 1 --accelerator=ddp -b 32
training_urls =  ./shards/imagenet-train-{000000..000146}.tar
val_urls =  ./shards/imagenet-val-{000000..000006}.tar
batch_size 32 num_workers 6 world_size 0
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
initializing ddp: GLOBAL_RANK: 0, MEMBER: 1/2
Traceback (most recent call last):
  File "train.py", line 245, in <module>
    main(args)
  File "train.py", line 234, in main
    trainer.fit(model, data)
  File "/home/tmb/proj/webdataset-lightning/venv/src/pytorch-lightning/pytorch_lightning/trainer/trainer.py", line 458, in fit
    self.accelerator.setup_environment()
  File "/home/tmb/proj/webdataset-lightning/venv/src/pytorch-lightning/pytorch_lightning/accelerators/accelerator.py", line 80, in setup_environment
    self.training_type_plugin.setup_environment()
  File "/home/tmb/proj/webdataset-lightning/venv/src/pytorch-lightning/pytorch_lightning/plugins/training_type/ddp.py", line 108, in setup_environment
    self.setup_distributed()
  File "/home/tmb/proj/webdataset-lightning/venv/src/pytorch-lightning/pytorch_lightning/plugins/training_type/ddp.py", line 196, in setup_distributed
    self.init_ddp_connection()
  File "/home/tmb/proj/webdataset-lightning/venv/src/pytorch-lightning/pytorch_lightning/plugins/training_type/ddp.py", line 270, in init_ddp_connection
    torch_distrib.init_process_group(self.torch_distributed_backend, rank=global_rank, world_size=world_size)
  File "/home/tmb/proj/webdataset-lightning/venv/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 503, in init_process_group
    _update_default_pg(_new_process_group_helper(
  File "/home/tmb/proj/webdataset-lightning/venv/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 588, in _new_process_group_helper
    pg = ProcessGroupGloo(
RuntimeError: [/pytorch/third_party/gloo/gloo/transport/tcp/pair.cc:769] connect [127.0.1.1]:6042: Connection refused
192.168.1.3$ 
```

```
192.168.1.9$ PL_TORCH_DISTRIBUTED_BACKEND=gloo MASTER_ADDR=192.168.1.3 MASTER_PORT=9888 NODE_RANK=1 python3 train.py --num_nodes=2 --gpus 1 --accelerator=ddp -b 32
training_urls =  ./shards/imagenet-train-{000000..000146}.tar
val_urls =  ./shards/imagenet-val-{000000..000006}.tar
batch_size 32 num_workers 6 world_size 0
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
initializing ddp: GLOBAL_RANK: 1, MEMBER: 2/2
... hangs ...
```

It looks like the jobs ignore the MASTER_ADDR and just pick their own on localhost anyway.

Source is here: https://github.com/tmbdev/webdataset-lightning/blob/f4d837dda1216ca0be8049c1a4297c49486b8ba2/train.py#L221


You need to select more than one GPU, otherwise it it will not run with the DDP plugin. This is why it ignores env variables that you supply. Can you try that? Then we can think if we can  support the multi node single gpu case.

Commit message: add tab to report print

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f344f490de6c1..7a29e0661fc2a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -140,6 +140,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `FastForwardSampler` and `CaptureIterableDataset` ([#8307](https://github.com/PyTorchLightning/pytorch-lightning/pull/8307))
 
 
+- Enabled traditional/manual launching of DDP processes through `LOCAL_RANK` and `NODE_RANK` environment variable assignments ([#7480](https://github.com/PyTorchLightning/pytorch-lightning/pull/7480))
+
+
 ### Changed
 
 
diff --git a/pytorch_lightning/plugins/environments/lightning_environment.py b/pytorch_lightning/plugins/environments/lightning_environment.py
index 25da0cfb691e8..077ebf995eebf 100644
--- a/pytorch_lightning/plugins/environments/lightning_environment.py
+++ b/pytorch_lightning/plugins/environments/lightning_environment.py
@@ -23,12 +23,16 @@ class LightningEnvironment(ClusterEnvironment):
     """
     The default environment used by Lightning for a single node or free cluster (not managed).
 
-    The master process must be launched by the user and Lightning will spawn new
-    worker processes for distributed training, either in a single node or across multiple nodes.
+    There are two modes the Lightning environment can operate with:
+
+    1.  The user only launches the main process by :code:`python train.py ...` with no additional environment variables
+        set. Lightning will spawn new worker processes for distributed training in the current node.
+    2.  The user launches all processes manually or with utilities like :code:`torch.distributed.launch`.
+        The appropriate environment variables need to be set, and at minimum :code:`LOCAL_RANK`.
 
     If the master address and port are not provided, the default environment will choose them
     automatically. It is recommended to use this default environment for single-node distributed
-    training as it provides the most convenient way to launch the training script.
+    training as it provides a convenient way to launch the training script.
     """
 
     def __init__(self):
@@ -38,7 +42,12 @@ def __init__(self):
         self._world_size: int = 1
 
     def creates_children(self) -> bool:
-        return False
+        """
+        Returns whether the cluster creates the processes or not.
+        If at least :code:`LOCAL_RANK` is available as environment variable, Lightning assumes the user acts as the
+        process launcher/job scheduler and Lightning will not launch new processes.
+        """
+        return "LOCAL_RANK" in os.environ
 
     def master_address(self) -> str:
         return os.environ.get("MASTER_ADDR", "127.0.0.1")
diff --git a/pytorch_lightning/plugins/training_type/ddp.py b/pytorch_lightning/plugins/training_type/ddp.py
index bbea7c9d4d514..55401d006be81 100644
--- a/pytorch_lightning/plugins/training_type/ddp.py
+++ b/pytorch_lightning/plugins/training_type/ddp.py
@@ -157,7 +157,7 @@ def _is_single_process_single_device(self) -> bool:
 
     def setup_environment(self) -> None:
         # start the other scripts
-        if not self.cluster_environment.creates_children() and os.environ.get("PL_IN_DDP_SUBPROCESS", "0") != "1":
+        if not self.cluster_environment.creates_children():
             self._call_children_scripts()
 
         # set the task idx
@@ -209,8 +209,6 @@ def _call_children_scripts(self):
         if self.parallel_devices is None:
             raise MisconfigurationException("you selected (distribute_backend = ddp) but did not set Trainer(gpus=?)")
 
-        os.environ["PL_IN_DDP_SUBPROCESS"] = "1"
-
         os.environ["WORLD_SIZE"] = f"{self.num_processes * self.num_nodes}"
 
         self.interactive_ddp_procs = []
diff --git a/pytorch_lightning/trainer/connectors/accelerator_connector.py b/pytorch_lightning/trainer/connectors/accelerator_connector.py
index 2e9a59b6181b9..ef9d089e48b4f 100644
--- a/pytorch_lightning/trainer/connectors/accelerator_connector.py
+++ b/pytorch_lightning/trainer/connectors/accelerator_connector.py
@@ -452,11 +452,6 @@ def select_training_type_plugin(self) -> TrainingTypePlugin:
             use_ddp_sharded_spawn = self._distrib_type == DistributedType.DDP_SHARDED_SPAWN
             use_ddp_fully_sharded = self._distrib_type == DistributedType.DDP_FULLY_SHARDED
 
-            # TODO: decouple from TE
-            # ddp script mode uses the same flags as TE
-            if os.environ.get("PL_IN_DDP_SUBPROCESS", False):
-                use_torchelastic_ddp = False
-
             if use_tpu_spawn:
                 ddp_plugin_cls = TPUSpawnPlugin
             elif use_ddp_sharded:
diff --git a/tests/plugins/environments/test_lightning_environment.py b/tests/plugins/environments/test_lightning_environment.py
index 8ebcec953fcc8..90c7403f1fe60 100644
--- a/tests/plugins/environments/test_lightning_environment.py
+++ b/tests/plugins/environments/test_lightning_environment.py
@@ -1,6 +1,8 @@
 import os
 from unittest import mock
 
+import pytest
+
 from pytorch_lightning.plugins.environments import LightningEnvironment
 
 
@@ -37,6 +39,24 @@ def test_attributes_from_environment_variables():
     assert env.world_size() == 100
 
 
+@pytest.mark.parametrize(
+    "environ, creates_children", [
+        ({}, False),
+        ({
+            "LOCAL_RANK": "2"
+        }, True),
+        ({
+            "NODE_RANK": "1"
+        }, False),
+    ]
+)
+def test_manual_user_launch(environ, creates_children):
+    """ Test that the environment switches to manual user mode when LOCAL_RANK env variable detected. """
+    with mock.patch.dict(os.environ, environ):
+        env = LightningEnvironment()
+        assert env.creates_children() == creates_children
+
+
 @mock.patch.dict(os.environ, {
     "GROUP_RANK": "1",
 })
diff --git a/tests/special_tests.sh b/tests/special_tests.sh
index 95311fb2df515..2786f30bf4ab9 100755
--- a/tests/special_tests.sh
+++ b/tests/special_tests.sh
@@ -78,6 +78,12 @@ if [ $? -eq 0 ]; then
     report+="Ran\ttests/utilities/test_warnings.py\n"
 fi
 
+# test that a user can manually launch individual processes
+args="--trainer.gpus 2 --trainer.accelerator ddp --trainer.max_epochs=1 --data.batch_size=32 --trainer.limit_train_batches=2 --trainer.limit_val_batches=2"
+MASTER_ADDR="localhost" MASTER_PORT=1234 LOCAL_RANK=1 python pl_examples/basic_examples/simple_image_classifier.py ${args} &
+MASTER_ADDR="localhost" MASTER_PORT=1234 LOCAL_RANK=0 python pl_examples/basic_examples/simple_image_classifier.py ${args}
+report+="Ran\tmanual ddp launch test\n"
+
 # echo test report
 printf '=%.s' {1..80}
 printf "\n$report"


Code review comments:
diff_hunk: @@ -78,6 +78,12 @@ if [ $? -eq 0 ]; then
     report+="Ran\ttests/utilities/test_warnings.py\n"
 fi
 
+# test that a user can manually launch individual processes
+args="--trainer.gpus 2 --trainer.accelerator ddp --trainer.max_epochs=1 --data.batch_size=32 --trainer.limit_train_batches=2 --trainer.limit_val_batches=2"
path: tests/special_tests.sh
text: @user1:
For the purpose of this test, this should be enough

```suggestion
args="--trainer.gpus 2 --trainer.accelerator ddp --trainer.fast_dev_run=True
```

@author:
I took this from the existing tests