Refcount problem for torch.distributed.Store objects defined in Python
Labels: oncall: distributed, module: c10d

### 🐛 Describe the bug

Consider the following Python code, call it `store_bug.py` (for a fully working version, see [here](https://gist.github.com/heiner/7ed5802eb9b3218a418262d542ecb827#file-store_bug-py)):

```python
import datetime
import sys

import torch

class MyStore(torch.distributed.Store):
    # Some store implementation.

def init(rank, world_size, use_workaround=False, *, _cache={}):
    store = MyStore()
    torch.distributed.init_process_group(
        backend="gloo",
        world_size=world_size,
        rank=rank,
        timeout=datetime.timedelta(seconds=3.0),
        store=store,
    )
    if use_workaround:
        _cache["store"] = store


def main():
    rank = int(sys.argv[1])
    world_size = int(sys.argv[2])
    use_workaround = sys.argv[3] if len(sys.argv) > 3 else False

    print(f"starting rank {rank}/{world_size} with use_workaround={use_workaround}")
    init(rank, world_size, use_workaround)

    g = torch.distributed.new_group(timeout=datetime.timedelta(seconds=2.0))

    payload = torch.tensor([3 * rank + 1], dtype=torch.int64)
    torch.distributed.all_reduce(payload, group=g)
    print(f"Rank {rank}, sum is {payload}")


if __name__ == "__main__":
    main()
```

Now, run the above code like this:

```sh
rm -rf /tmp/store_bug; for i in {0..1}; do python -u store_bug.py $i 2 & done && wait
```

This will result in

```sh
$ rm -rf /tmp/store_bug; for i in {0..1}; do python -u store_bug.py $i 2 & done && wait

[2] 69038
[3] 69039
starting rank 1/2 with use_workaround=False
starting rank 0/2 with use_workaround=False
Traceback (most recent call last):
  File "/Users/heiner/src/projects/pytorch_store_bug/store_bug.py", line 147, in <module>
    main()
  File "/Users/heiner/src/projects/pytorch_store_bug/store_bug.py", line 139, in main
    g = torch.distributed.new_group(timeout=datetime.timedelta(seconds=2.0))
  File "/Users/heiner/miniconda3/envs/dev/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 3520, in new_group
    pg = _new_process_group_helper(
  File "/Users/heiner/miniconda3/envs/dev/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 1009, in _new_process_group_helper
    backend_class = ProcessGroupGloo(backend_prefix_store, group_rank, group_size, timeout=timeout)
RuntimeError: fn INTERNAL ASSERT FAILED at "/Users/runner/work/pytorch/pytorch/pytorch/torch/csrc/distributed/c10d/init.cpp":137, please report a bug to PyTorch.
Traceback (most recent call last):
  File "/Users/heiner/src/projects/pytorch_store_bug/store_bug.py", line 147, in <module>
    main()
  File "/Users/heiner/src/projects/pytorch_store_bug/store_bug.py", line 139, in main
    g = torch.distributed.new_group(timeout=datetime.timedelta(seconds=2.0))
  File "/Users/heiner/miniconda3/envs/dev/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 3520, in new_group
    pg = _new_process_group_helper(
  File "/Users/heiner/miniconda3/envs/dev/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 1009, in _new_process_group_helper
    backend_class = ProcessGroupGloo(backend_prefix_store, group_rank, group_size, timeout=timeout)
RuntimeError: fn INTERNAL ASSERT FAILED at "/Users/runner/work/pytorch/pytorch/pytorch/torch/csrc/distributed/c10d/init.cpp":137, please report a bug to PyTorch.
[3]  + exit 1     python -u store_bug.py $i 2
[2]  + exit 1     python -u store_bug.py $i 2
```

In contrast, running the code such that `use_workaround` becomes true works:

```sh
$ rm -rf /tmp/store_bug; for i in {0..1}; do python -u store_bug.py $i 2 1 & done && wait

[2] 69067
[3] 69068
starting rank 1/2 with use_workaround=1
starting rank 0/2 with use_workaround=1
Rank 1, sum is tensor([5])
Rank 0, sum is tensor([5])
[3]  + done       python -u store_bug.py $i 2 1
[2]  + done       python -u store_bug.py $i 2 1
```

The difference, of course, is that the working version keeps an additional reference to the `store` object, making sure it stays alive.

This is a regression introduced in or around PyTorch 2.0, possibly in https://github.com/pytorch/pytorch/pull/91178. The [existing test](https://github.com/pytorch/pytorch/blob/a01e795a6d03dff853534c4c4a9f3522f3fd717d/test/distributed/test_store.py#L461) for this via `torch.distributed._test_python_store` doesn't catch this as it does not create a new group via the store.

To recreate at home:

```sh
git clone https://gist.github.com/7ed5802eb9b3218a418262d542ecb827.git storebug
cd storebug/
rm -rf /tmp/store_bug; for i in {0..1}; do python -u store_bug.py $i 2 1 & done && wait  # Works.
rm -rf /tmp/store_bug; for i in {0..1}; do python -u store_bug.py $i 2 & done && wait  # Breaks.
```

### Versions

Collecting environment information...
PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.6 (arm64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.102)
CMake version: version 3.24.2
Libc version: N/A

Python version: 3.9.11 (main, Mar 29 2022, 14:04:34)  [Clang 12.0.0 ] (64-bit runtime)
Python platform: macOS-12.6-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] flake8==4.0.1
[pip3] flake8-bugbear==22.3.23
[pip3] flake8-unused-arguments==0.0.10
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.5
[pip3] open-clip-torch==2.9.1
[pip3] pytorch-lightning==1.4.2
[pip3] torch==2.0.1
[pip3] torchaudio==0.13.1
[pip3] torchdiffeq==0.2.3
[pip3] torchmetrics==0.6.0
[pip3] torchvision==0.14.1
[conda] numpy                     1.23.5                   pypi_0    pypi
[conda] open-clip-torch           2.9.1                    pypi_0    pypi
[conda] pytorch-lightning         1.4.2                    pypi_0    pypi
[conda] torch                     1.11.0                   pypi_0    pypi
[conda] torchaudio                0.13.1                 py39_cpu    pytorch
[conda] torchdiffeq               0.2.3                    pypi_0    pypi
[conda] torchmetrics              0.6.0                    pypi_0    pypi
[conda] torchvision               0.14.1                   pypi_0    pypi

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu