Repository: pytorch
Issue: race condition of agent's TCP store on exit_barrier()
Labels: triaged,Stale,oncall: r2p

## 🐛 Bug

The agent uses the store to make consensus on exiting the job. Each agent put its rank on the store and gathers all the ranks from the store. Since `get()` does not return until the key-value pair exists, the barrier guarantees that all workers exit. 

However, at the same time, the heartbeat mechanism of the rendezvous handler is still working. The heartbeat mechanism needs to write the timestamp to the store and get the timestamp of other workers from the store. 

This results in a race condition - the same store client are issuing multiple requests. When receving a response from the store server, the client does not know which request the response belones to, leading to a crash. 

Here is the full error log.
```py
INFO:torch.distributed.launcher.api:Starting elastic_operator with launch configs:
  entrypoint       : main.py
  min_nodes        : 2
  max_nodes        : 2
  nproc_per_node   : 1
  run_id           : 1234
  rdzv_backend     : c10d
  rdzv_endpoint    : 10.28.1.16
  rdzv_configs     : {'timeout': 900}
  max_restarts     : 0
  monitor_interval : 5
  log_dir          : None
  metrics_cfg      : {}

INFO:torch.distributed.elastic.agent.server.local_elastic_agent:log directory set to: /tmp/torchelastic_w9bu0zk3/1234_a1f0w34w
INFO:torch.distributed.elastic.agent.server.api:[default] starting workers for entrypoint: python3
INFO:torch.distributed.elastic.agent.server.api:[default] Rendezvous'ing worker group
INFO:torch.distributed.elastic.agent.server.api:[default] Rendezvous complete for workers. Result:
  restart_count=0
  master_addr=net-g1
  master_port=32951
  group_rank=0
  group_world_size=2
  local_ranks=[0]
  role_ranks=[0]
  global_ranks=[0]
  role_world_sizes=[2]
  global_world_sizes=[2]

INFO:torch.distributed.elastic.agent.server.api:[default] Starting worker group
INFO:torch.distributed.elastic.multiprocessing:Setting worker0 reply file to: /tmp/torchelastic_w9bu0zk3/1234_a1f0w34w/attempt_0/0/error.json
INFO:torch.distributed.elastic.agent.server.api:[default] worker group successfully finished. Waiting 300 seconds for other agents to finish.
INFO:torch.distributed.elastic.agent.server.api:Local worker group finished (SUCCEEDED). Waiting 300 seconds for other agents to finish
ERROR:torch.distributed.elastic.agent.server.api:Error waiting on exit barrier. Elapsed: 5.0016467571258545 seconds
Traceback (most recent call last):
  File "/home/yczhong/.conda/envs/yczhong/lib/python3.6/site-packages/torch/distributed/elastic/agent/server/api.py", line 904, in _exit_barrier
    barrier_timeout=self._exit_barrier_timeout,
  File "/home/yczhong/.conda/envs/yczhong/lib/python3.6/site-packages/torch/distributed/elastic/utils/store.py", line 67, in barrier
    synchronize(store, data, rank, world_size, key_prefix, barrier_timeout)
  File "/home/yczhong/.conda/envs/yczhong/lib/python3.6/site-packages/torch/distributed/elastic/utils/store.py", line 53, in synchronize
    agent_data = get_all(store, key_prefix, world_size)
  File "/home/yczhong/.conda/envs/yczhong/lib/python3.6/site-packages/torch/distributed/elastic/utils/store.py", line 31, in get_all
    data = store.get(f"{prefix}{idx}")
RuntimeError: Stop_waiting response is expected
WARNING:torch.distributed.elastic.rendezvous.dynamic_rendezvous:The node 'net-g1_31840_0' has failed to send a keep-alive heartbeat to the rendezvous '1234' due to an error of type RendezvousStateError.
```

As shown in the log, the client should have received the stop_waiting response, but due to the heartbeat, it received the response to the heartbeat request, causing the failure. 

## To Reproduce


Here is a minimal reproducible code.

On server 1:
```py
import os
import time

import torch


def main():
    local_rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(local_rank)
    torch.distributed.init_process_group(
        'nccl'
    )

if __name__ == '__main__':
    main()
```

On server 2:
```py
import os
import time

import torch


def main():
    local_rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(local_rank)
    torch.distributed.init_process_group(
        'nccl'
    )
    time.sleep(30)

if __name__ == '__main__':
    main()
```

The only difference of the two code snippets is that server 2 will sleep 30 seconds and thus exit later. 


## Expected behavior

It should work. 

## Environment

```
Collecting environment information...
PyTorch version: 1.10.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 7.0.0-3~ubuntu0.18.04.1 (tags/RELEASE_700/final)
CMake version: version 3.19.6
Libc version: glibc-2.17

Python version: 3.6.13 |Anaconda, Inc.| (default, Jun  4 2021, 14:25:59)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.15.0-123-generic-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration:
GPU 0: Tesla P100-PCIE-12GB
GPU 1: Tesla P100-PCIE-12GB

Nvidia driver version: 455.32.00
cuDNN version: Probably one of the following:
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7.6.5
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn.so.7
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.10.0
[pip3] torchvision==0.10.0
[conda] blas                      1.0                         mkl
[conda] magma-cuda102             2.5.2                         1    pytorch
[conda] mkl                       2021.3.0           h06a4308_520
[conda] mkl-include               2021.3.0           h06a4308_520
[conda] mkl-service               2.4.0            py36h7f8727e_0
[conda] mkl_fft                   1.3.0            py36h42c9631_2
[conda] mkl_random                1.2.2            py36h51133e4_0
[conda] numpy                     1.19.5                   pypi_0    pypi
[conda] numpy-base                1.16.6           py36hdc34a94_3
[conda] torch                     1.10.0                   pypi_0    pypi
[conda] torchvision               0.10.0                   pypi_0    pypi
```

## Additional context

A possible solution is to shutdown the rdzv_handler in `_exit_barrier()`.