Repository: Lightning-AI/pytorch-lightning
Pull request: [bugfix] Add mechanism to prevent deadlock for DDP on Exception Trigger
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

This PR introduces a mechanism to prevent deadlock.

Mechanism:
* On ddp instantiation, make all processes aware of pids
* Share a temporary directory used to save some synchronisation state

On failure for a given process:
* Save a dummy object to the temporary
* Wait a little bit
* Check if the other processes have actually wrote to the folder. 
* If the num files == world_size, all good ! 
* else: kill all process and raised DeadLock Exception.

Other approaches tried:
* Use timetout for TCPStore -> result in cuda failure
* Use signal -> Very flicky and un-stable 

Fixes #8130

## Before submitting
- [ ] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [ ] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [ ] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [ ] Did you make sure to update the documentation with your changes? (if necessary)
- [ ] 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?
Make sure you had fun coding 🙃


Problem statement:
DDP script hangs forever if one child process dies.
## 🐛 Bug

When using multiple gpus with `accelerator=ddp`, if one child processes throws an error, the main script hangs indefinitely and never terminates. The error does appear in the stdout but does not cause the main script to terminate as it should.

Note, that this bug does not appear if one uses `accelerator=ddp_spawn`. 

## To Reproduce

Below a simple .py file that demonstrates the error. This file uses the example model provided in the docs but forces the child process on local_rank 1 to throw an error in the validation step. When running this script on 2 gpus, this hangs forever (despite the errror appearing in the log). 

```
import os

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import MNIST
from torchvision import transforms
import pytorch_lightning as pl
from pytorch_lightning.metrics.functional import accuracy
from pl_bolts.datasets import DummyDataset

train = DummyDataset((1, 28, 28), (1,))
train = DataLoader(train, batch_size=32)
val = DummyDataset((1, 28, 28), (1,))
val = DataLoader(val, batch_size=32)
test = DummyDataset((1, 28, 28), (1,))
test = DataLoader(test, batch_size=32)

class LitAutoEncoder(pl.LightningModule):

    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))

    def training_step(self, batch, batch_idx):
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('train_loss', loss)
        return loss
        

    def validation_step(self, batch, batch_idx):
        # --------------------------
        # CODE MODIFICATION FOR BUG REPORT HERE
        if self.local_rank == 1:
            raise ValueError 
        # --------------------------

        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('val_loss', loss)


    def test_step(self, batch, batch_idx):
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('test_loss', loss)


    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

if __name__ == "__main__":
    # init model
    ae = LitAutoEncoder()

    # Initialize a trainer
    trainer = pl.Trainer(gpus=2, max_epochs=3, progress_bar_refresh_rate=20, accelerator="ddp")

    # Train the model ⚡
    trainer.fit(ae, train, val)
```
Note that when using `accelerator = ddp_spawn` on the other hand fails correctly.

### Expected behavior

If one child processes fails, the main script should fail and exit, not run forever.

### Environment

 - PyTorch Version (e.g., 1.0): 1.8.0
 - OS (e.g., Linux): Linux
 - Python version: 3.7.3
 - PL version: 1.3.7
 - GPU models and configuration: 2 GPUS

Hints from issue discussion:
Hi @melanibe, thanks for the issue. 

I'll look into it, but depending on how the child process fails, there may not be much we can do.

Let me quickly explain why and also why it works with spawn:

When using spawn, we actually use python's multiprocessing to spawn the new processes, that means they are connected and there is a communication broadcast. However, this comes with some other limitations. This is why with plain ddp, we circumvent this by relying on an OS-call to the interpreter to run the same script again. That however, has the disadvantage that python does not know, that these processes are connected to each other (i.e. that the one is the other's child).
Dear @melanibe @justusschock,

I am working on a mechanism to prevent this behaviour and it will kill the processes if it detects a deadlock.
Here is the code associated with it: https://github.com/PyTorchLightning/pytorch-lightning/blob/3da64e5336c19bd9cdff51d6e2924c47dbcc54b0/pytorch_lightning/trainer/trainer.py#L1010.

@justusschock I tried using signal to emulate a keyboardInterrupt and enable the code to keep going, but it seems to work only 1/10. Killing seemed the only stable option.

Best,
T.C

Commit message: update on comments

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 980d2a450f786..867ad7cfa3c70 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -220,6 +220,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - `Trainer(resume_from_checkpoint=...)` now restores the model directly after `LightningModule.setup()`, which is before `LightningModule.configure_sharded_model()` ([#7652](https://github.com/PyTorchLightning/pytorch-lightning/pull/7652))
 
 
+- Added a mechanism to detect `deadlock` for `DDP` and `kill the processes` when it happens ([#8167](https://github.com/PyTorchLightning/pytorch-lightning/pull/8167))
+
+
 ### Deprecated
 
 
diff --git a/pl_examples/basic_examples/autoencoder.py b/pl_examples/basic_examples/autoencoder.py
index 70e46b037e2e1..94e4fbfcf7ae2 100644
--- a/pl_examples/basic_examples/autoencoder.py
+++ b/pl_examples/basic_examples/autoencoder.py
@@ -124,7 +124,7 @@ def predict_dataloader(self):
 
 
 def cli_main():
-    cli = LightningCLI(LitAutoEncoder, MyDataModule, seed_everything_default=1234)
+    cli = LightningCLI(LitAutoEncoder, MyDataModule, seed_everything_default=1234, save_config_overwrite=True)
     cli.trainer.test(cli.model, datamodule=cli.datamodule)
     predictions = cli.trainer.predict(cli.model, datamodule=cli.datamodule)
     print(predictions[0])
diff --git a/pl_examples/basic_examples/backbone_image_classifier.py b/pl_examples/basic_examples/backbone_image_classifier.py
index 0d43748cecdfd..381cda088ea9d 100644
--- a/pl_examples/basic_examples/backbone_image_classifier.py
+++ b/pl_examples/basic_examples/backbone_image_classifier.py
@@ -135,7 +135,7 @@ def predict_dataloader(self):
 
 
 def cli_main():
-    cli = LightningCLI(LitClassifier, MyDataModule, seed_everything_default=1234)
+    cli = LightningCLI(LitClassifier, MyDataModule, seed_everything_default=1234, save_config_overwrite=True)
     cli.trainer.test(cli.model, datamodule=cli.datamodule)
     predictions = cli.trainer.predict(cli.model, datamodule=cli.datamodule)
     print(predictions[0])
diff --git a/pl_examples/basic_examples/simple_image_classifier.py b/pl_examples/basic_examples/simple_image_classifier.py
index ffb6434352b2e..70aaa35931f8e 100644
--- a/pl_examples/basic_examples/simple_image_classifier.py
+++ b/pl_examples/basic_examples/simple_image_classifier.py
@@ -76,7 +76,7 @@ def configure_optimizers(self):
 
 
 def cli_main():
-    cli = LightningCLI(LitClassifier, MNISTDataModule, seed_everything_default=1234)
+    cli = LightningCLI(LitClassifier, MNISTDataModule, seed_everything_default=1234, save_config_overwrite=True)
     cli.trainer.test(cli.model, datamodule=cli.datamodule)
 
 
diff --git a/pytorch_lightning/plugins/training_type/ddp.py b/pytorch_lightning/plugins/training_type/ddp.py
index b855d100b1f12..8c4ae6878ad4d 100644
--- a/pytorch_lightning/plugins/training_type/ddp.py
+++ b/pytorch_lightning/plugins/training_type/ddp.py
@@ -13,8 +13,12 @@
 # limitations under the License.
 import logging
 import os
+import shutil
+import signal
 import subprocess
 import sys
+import tempfile
+import time
 from time import sleep
 from typing import Any, Dict, List, Optional, Union
 
@@ -38,7 +42,7 @@
     rank_zero_warn,
 )
 from pytorch_lightning.utilities.distributed import rank_zero_info, rank_zero_only, ReduceOp, sync_ddp_if_available
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.exceptions import DeadlockDetectedException, MisconfigurationException
 from pytorch_lightning.utilities.seed import reset_seed
 
 if _HYDRA_AVAILABLE:
@@ -94,6 +98,8 @@ def __init__(
         self._ddp_comm_state = ddp_comm_state
         self._ddp_comm_hook = ddp_comm_hook
         self._ddp_comm_wrapper = ddp_comm_wrapper
+        self._pids: Optional[List[int]] = None
+        self._sync_dir: Optional[str] = None
         self.set_world_ranks()
 
     @property
@@ -142,7 +148,6 @@ def setup_environment(self) -> None:
         self.setup_distributed()
 
     def _call_children_scripts(self):
-
         # bookkeeping of spawned processes
         assert self.local_rank == 0
         self._check_can_spawn_children()
@@ -156,6 +161,9 @@ def _call_children_scripts(self):
         os.environ["NODE_RANK"] = str(self.cluster_environment.node_rank())
         os.environ["LOCAL_RANK"] = str(self.cluster_environment.local_rank())
 
+        # create a temporary directory used to synchronize processes on deadlock.
+        os.environ["PL_DDP_SYNC_TMPDIR"] = self._sync_dir = tempfile.mkdtemp()
+
         # Check if the current calling command looked like `python a/b/c.py` or `python -m a.b.c`
         # See https://docs.python.org/3/reference/import.html#main-spec
         if __main__.__spec__ is None:  # pragma: no-cover
@@ -318,7 +326,8 @@ def pre_dispatch(self):
 
         self.configure_ddp()
 
-        self.barrier()
+        # share ddp pids to all processes
+        self._share_information_to_prevent_deadlock()
 
     def post_dispatch(self) -> None:
         self.cluster_environment.teardown()
@@ -381,3 +390,45 @@ def register_plugins(cls, plugin_registry: Dict) -> None:
             description="DDP Plugin with `find_unused_parameters` as False",
             find_unused_parameters=False
         )
+
+    def _share_information_to_prevent_deadlock(self):
+        self._share_pids()
+
+        # remove `PL_DDP_SYNC_TMPDIR` from os.environ
+        self._sync_dir = os.environ.pop("PL_DDP_SYNC_TMPDIR", None)
+
+    def _share_pids(self):
+        """
+        Make all DDP processes aware of all processes pids.
+        """
+        self.barrier()
+        pids = self.all_gather(torch.tensor(os.getpid(), device=self.root_device))
+        pids = pids.cpu().numpy().tolist()
+        self._pids = pids if isinstance(pids, list) else [pids]
+
+    def reconciliate_processes(self, trace: str):
+        if self.world_size > 1:
+            sync_dir = self._sync_dir
+
+            if not os.path.exists(sync_dir):
+                # avoid race condition
+                try:
+                    os.makedirs(sync_dir)
+                except FileExistsError:
+                    pass
+
+            # save a file locally.
+            torch.save(trace, os.path.join(sync_dir, f"{self.global_rank}.pl"))
+
+            # sleep for a short time
+            time.sleep(3)
+
+            # return if all processes wrote a file in the `sync_dir`.
+            if len(os.listdir(sync_dir)) == self.world_size:
+                return
+
+            for pid in self._pids:
+                if pid != os.getpid():
+                    os.kill(pid, signal.SIGKILL)
+                shutil.rmtree(sync_dir)
+                raise DeadlockDetectedException(f"DeadLock detected from rank: {self.global_rank} \n {trace}")
diff --git a/pytorch_lightning/plugins/training_type/training_type_plugin.py b/pytorch_lightning/plugins/training_type/training_type_plugin.py
index e7ca73bc9f40d..cc4d1aa115e74 100644
--- a/pytorch_lightning/plugins/training_type/training_type_plugin.py
+++ b/pytorch_lightning/plugins/training_type/training_type_plugin.py
@@ -57,6 +57,12 @@ def setup_environment(self) -> None:
     def setup(self, model: Module) -> None:
         """Called by the accelerator to finish setup."""
 
+    def reconciliate_processes(self, trace: str):
+        """
+        Function to re-conciliate processes on failure
+        """
+        pass
+
     @property
     @abstractmethod
     def on_gpu(self) -> bool:
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
index 4d097a2de2763..dd7cdd3b7311e 100644
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Trainer to automate the training."""
 import logging
+import traceback
 import warnings
 from datetime import timedelta
 from pathlib import Path
@@ -992,6 +993,8 @@ def _run_train(self) -> None:
                 self.state.stage = None
         except BaseException:
             self.state.status = TrainerStatus.INTERRUPTED
+            # try syncing remaing processes, kill otherwise
+            self.training_type_plugin.reconciliate_processes(traceback.format_exc())
             # give accelerators a chance to finish
             self.accelerator.on_train_end()
             # reset bookkeeping
diff --git a/pytorch_lightning/utilities/exceptions.py b/pytorch_lightning/utilities/exceptions.py
index 01b1e8c053950..4f9b553ca2e5e 100644
--- a/pytorch_lightning/utilities/exceptions.py
+++ b/pytorch_lightning/utilities/exceptions.py
@@ -15,3 +15,7 @@
 
 class MisconfigurationException(Exception):
     pass
+
+
+class DeadlockDetectedException(Exception):
+    pass
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index d91f8d0a427e7..14c6f9c1942ce 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -39,7 +39,7 @@
 from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities import DeviceType, DistributedType
 from pytorch_lightning.utilities.cloud_io import load as pl_load
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.exceptions import DeadlockDetectedException, MisconfigurationException
 from pytorch_lightning.utilities.seed import seed_everything
 from tests.base import EvalModelTemplate
 from tests.helpers import BoringModel, RandomDataset
@@ -1937,3 +1937,32 @@ def test_exception_when_lightning_module_is_not_set_on_trainer():
         trainer.test()
     with pytest.raises(MisconfigurationException, match=r"`model` must be provided.*predict"):
         trainer.predict()
+
+
+@RunIf(min_gpus=2, special=True)
+def test_ddp_terminate_when_deadlock_is_detected(tmpdir):
+
+    class CustomException(Exception):
+        pass
+
+    class TestModel(BoringModel):
+
+        def training_step(self, batch, batch_idx):
+            if batch_idx == 1 and self.trainer.is_global_zero:
+                raise CustomException
+            return super().training_step(batch, batch_idx)
+
+    model = TestModel()
+
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=1,
+        limit_train_batches=5,
+        num_sanity_val_steps=0,
+        gpus=2,
+        accelerator="ddp",
+    )
+
+    # simulate random failure in training_step on rank 0
+    with pytest.raises(DeadlockDetectedException, match="CustomException"):
+        trainer.fit(model)


Code review comments:
diff_hunk: @@ -381,3 +390,45 @@ def register_plugins(cls, plugin_registry: Dict) -> None:
             description="DDP Plugin with `find_unused_parameters` as False",
             find_unused_parameters=False
         )
+
+    def _share_information_to_prevent_deadlock(self):
+        self._share_pids()
+
+        # remove `PL_DDP_SYNC_TMPDIR` from os.environ
+        self._sync_dir = os.environ.pop("PL_DDP_SYNC_TMPDIR", None)
+
+    def _share_pids(self):
+        """
+        Make all DDP processes aware of all processes pids.
+        """
+        self.barrier()
+        pids = self.all_gather(torch.tensor(os.getpid(), device=self.root_device))
+        pids = pids.cpu().numpy().tolist()
+        self._pids = pids if isinstance(pids, list) else [pids]
+
+    def reconciliate_processes(self, trace: str):
+        if self.world_size > 1:
+            sync_dir = self._sync_dir
+
+            if not os.path.exists(sync_dir):
+                # avoid race condition
+                try:
+                    os.makedirs(sync_dir)
+                except FileExistsError:
+                    pass
+
+            # save a file locally.
+            torch.save(trace, os.path.join(sync_dir, f"{self.global_rank}.pl"))
path: pytorch_lightning/plugins/training_type/ddp.py
text: @user1:
why are we saving the trace to a torch pickled file?
shouldn't it be a text file?

@author:
Removed the trace as it is not needed.

diff_hunk: @@ -57,6 +57,12 @@ def setup_environment(self) -> None:
     def setup(self, model: Module) -> None:
         """Called by the accelerator to finish setup."""
 
+    def reconciliate_processes(self, trace: str):
path: pytorch_lightning/plugins/training_type/training_type_plugin.py
text: @user2:
we should consider adding this only to the parallel plugin and not to the base class.
what do you think @user1 ?

@user1:
@user2 I thought about that as well, but then it wouldn't make sense to call this in the trainer unconditonally. Maybe we should rename this to `reconcilate_training`?

@author:
Agree. Done !

diff_hunk: @@ -381,3 +390,45 @@ def register_plugins(cls, plugin_registry: Dict) -> None:
             description="DDP Plugin with `find_unused_parameters` as False",
             find_unused_parameters=False
         )
+
+    def _share_information_to_prevent_deadlock(self):
+        self._share_pids()
+
+        # remove `PL_DDP_SYNC_TMPDIR` from os.environ
+        self._sync_dir = os.environ.pop("PL_DDP_SYNC_TMPDIR", None)
+
+    def _share_pids(self):
+        """
+        Make all DDP processes aware of all processes pids.
+        """
+        self.barrier()
+        pids = self.all_gather(torch.tensor(os.getpid(), device=self.root_device))
+        pids = pids.cpu().numpy().tolist()
+        self._pids = pids if isinstance(pids, list) else [pids]
+
+    def reconciliate_processes(self, trace: str):
+        if self.world_size > 1:
+            sync_dir = self._sync_dir
path: pytorch_lightning/plugins/training_type/ddp.py
text: @user1:
I suggest to invert the if condition so we have less indentation levels below

@author:
Done !

diff_hunk: @@ -381,3 +390,45 @@ def register_plugins(cls, plugin_registry: Dict) -> None:
             description="DDP Plugin with `find_unused_parameters` as False",
             find_unused_parameters=False
         )
+
+    def _share_information_to_prevent_deadlock(self):
+        self._share_pids()
+
+        # remove `PL_DDP_SYNC_TMPDIR` from os.environ
+        self._sync_dir = os.environ.pop("PL_DDP_SYNC_TMPDIR", None)
+
+    def _share_pids(self):
+        """
+        Make all DDP processes aware of all processes pids.
+        """
+        self.barrier()
+        pids = self.all_gather(torch.tensor(os.getpid(), device=self.root_device))
+        pids = pids.cpu().numpy().tolist()
+        self._pids = pids if isinstance(pids, list) else [pids]
+
+    def reconciliate_processes(self, trace: str):
+        if self.world_size > 1:
+            sync_dir = self._sync_dir
+
+            if not os.path.exists(sync_dir):
+                # avoid race condition
+                try:
+                    os.makedirs(sync_dir)
+                except FileExistsError:
+                    pass
+
+            # save a file locally.
+            torch.save(trace, os.path.join(sync_dir, f"{self.global_rank}.pl"))
+
+            # sleep for a short time
+            time.sleep(3)
+
path: pytorch_lightning/plugins/training_type/ddp.py
text: @user1:
you are waiting until all processese have written the file, correct?
so it is a sort of timeout. would it make sense to make this a parameter to the function with a default value?
how did you determine the value 3?

@author:
It was already working with `1`. Just added a 3 x rule for security. But there is no special reason there.

@author:
Done !

diff_hunk: @@ -381,3 +390,45 @@ def register_plugins(cls, plugin_registry: Dict) -> None:
             description="DDP Plugin with `find_unused_parameters` as False",
             find_unused_parameters=False
         )
+
+    def _share_information_to_prevent_deadlock(self):
+        self._share_pids()
+
+        # remove `PL_DDP_SYNC_TMPDIR` from os.environ
+        self._sync_dir = os.environ.pop("PL_DDP_SYNC_TMPDIR", None)
+
+    def _share_pids(self):
+        """
+        Make all DDP processes aware of all processes pids.
+        """
+        self.barrier()
+        pids = self.all_gather(torch.tensor(os.getpid(), device=self.root_device))
+        pids = pids.cpu().numpy().tolist()
+        self._pids = pids if isinstance(pids, list) else [pids]
+
+    def reconciliate_processes(self, trace: str):
+        if self.world_size > 1:
+            sync_dir = self._sync_dir
+
+            if not os.path.exists(sync_dir):
path: pytorch_lightning/plugins/training_type/ddp.py
text: @user1:
this check is redundant as you ahve the try-except block below anyway. 

@author:
Removed !