Repository: pytorch
Issue: Exception in distributed context doesn't propagate to child processes launched with multiprocessing
Labels: oncall: distributed

### 🐛 Describe the bug

I'm using `torch.distributed` to launch multi GPU trainings and `torch.multiprocessing` in data preparation but I have an issue if an exception is raised in a child process. 

Here is an example:
- I use 2 devices (launching my script using `torchrun --nproc_per_node=2 myscript.py`)
- on each device, the main process (called main:0 and main:1) starts a child process (called child:0 and child:1)
- if child raises an exception, I catch this exception to raise it in main aswell
- then, main:0 is killed thanks to SIGTERM
- but, child:0 isn't killed and thus it becomes a zombie process that keeps using resources and blocks the end of the whole process

Here a snippet that you can run using `torchrun --nproc_per_node=2 tmp.py` on a machine with at least 2 GPUs to reproduce this issue:

```python
# tmp.py
import time

import torch
from torch import multiprocessing


def main():
    torch.distributed.init_process_group("nccl")
    device = torch.device(torch.distributed.get_rank())
    torch.cuda.set_device(device)
    ctx = multiprocessing.get_context("spawn")
    exc_queue = ctx.Queue(1)
    torch.Tensor([1]).to(device)  # useful to display active processes in nvidia-smi
    process = ctx.Process(target=subprocess, args=(device, exc_queue), daemon=True)
    process.start()
    print(f"main:{device.index} started")
    process.join()
    print(f"child:{device.index} has joined")
    if not exc_queue.empty():
        exc = exc_queue.get()
        raise RuntimeError("subprocess failed") from exc
    print(f"main:{device.index} finished")


def subprocess(device: torch.device, exc_queue: multiprocessing.Queue):
    try:
        torch.Tensor([1]).to(device)  # useful to display active processes in nvidia-smi
        print(f"child:{device.index} started")
        if device.index:
            # emulating an exception that occurs only on a single device
            time.sleep(3)
            raise RuntimeError()
        while True:
            # emulating the "normal" behavior of the subprocess such as reading data
            print(f"child:{device.index} in loop")
            time.sleep(4)
    except Exception as exc:
        exc_queue.put(exc)
        raise exc


if __name__ == "__main__":
    main()
```

Here is the traceback I got:

```pyhon
$ torchrun --nproc_per_node=2 tmp.py
WARNING:torch.distributed.run:
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
*****************************************
main:0 started
main:1 started
child:0 started
child:0 in loop
child:1 started
Process SpawnProcess-1:
Traceback (most recent call last):
  File "/opt/rh/rh-python38/root/usr/lib64/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/opt/rh/rh-python38/root/usr/lib64/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "./tmp.py", line 40, in subprocess
    raise exc
  File "./tmp.py", line 33, in subprocess
    raise RuntimeError()
RuntimeError
child:1 has joined
RuntimeError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "tmp.py", line 44, in <module>
    main()
  File "tmp.py", line 22, in main
    raise RuntimeError("subprocess failed") from exc
RuntimeError: subprocess failed
child:0 in loop
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 178624 closing signal SIGTERM
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 1 (pid: 178625) of binary: ./venv/bin/python
Traceback (most recent call last):
  File "./venv/bin/torchrun", line 8, in <module>
    sys.exit(main())
  File "./venv/lib64/python3.8/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 345, in wrapper
    return f(*args, **kwargs)
  File "./venv/lib64/python3.8/site-packages/torch/distributed/run.py", line 761, in main
    run(args)
  File "./venv/lib64/python3.8/site-packages/torch/distributed/run.py", line 752, in run
    elastic_launch(
  File "./venv/lib64/python3.8/site-packages/torch/distributed/launcher/api.py", line 131, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "./venv/lib64/python3.8/site-packages/torch/distributed/launcher/api.py", line 245, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
============================================================
tmp.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2023-01-13_07:23:42
  host      : bv4sxk2.pnp.melodis.com
  rank      : 1 (local_rank: 1)
  exitcode  : 1 (pid: 178625)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
$ child:0 in loop
child:0 in loop
child:0 in loop
... # the loop keeps running endlessly
```

When looking at nvidia-smi after all processes have been killed except child:0 I see child:0 process still occupying memory.

Is there a way to make sure all children processes are killed when one of the main processes is killed in a distributed context ?

Thanks

### Versions

```
Collecting environment information...
PyTorch version: 1.12.1+cu116
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 9.3.1 20200408 (Red Hat 9.3.1-2)
Clang version: 14.0.6 ( 14.0.6-4sh.el7)
CMake version: version 3.24.3
Libc version: glibc-2.17

Python version: 3.8.13 (default, Aug 16 2022, 12:16:29)  [GCC 9.3.1 20200408 (Red Hat 9.3.1-2)] (64-bit runtime)
Python platform: Linux-3.10.0-693.2.2.el7.x86_64-x86_64-with-glibc2.2.5
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: 
GPU models and configuration: 
GPU 0: NVIDIA Tesla P100-PCIE-16GB
GPU 1: NVIDIA Tesla P100-PCIE-16GB

Nvidia driver version: 465.19.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.991
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.4
[pip3] torch==1.12.1+cu116
[pip3] torch-summary==1.4.5
[pip3] torch-tb-profiler==0.4.0
[pip3] torchaudio==0.12.1+cu116
[conda] Could not collect
```

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu