How to properly clean up shared cuda tensor in case of cuda error?

youkaichao1: <p>I’m trying to use pytorch cuda shared tensor to test device p2p capability:</p>
<pre data-code-wrap="python"><code class="lang-python"># run with `python test.py`
import json
import os
import sys
import tempfile
import time
from contextlib import contextmanager
from typing import Callable, Dict, List, Optional

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

import vllm.envs as envs
from vllm.logger import init_logger

logger = init_logger(__name__)


@contextmanager
def mute_output():
    with open(os.devnull, "w") as f:
        # sys.stderr = f
        # sys.stdout = f
        yield


def producer(i: int,
             init_method: str,
             consumer_done_event: mp.Event,
             result_quque: mp.Queue,
             cuda_visible_devices: Optional[str] = None,):
    if cuda_visible_devices is not None:
        os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
    with mute_output():
        dist.init_process_group(
            backend="gloo",
            init_method=init_method,
            world_size=2,
            rank=0,
        )
        # produce a tensor in GPU i
        data = torch.zeros((128, ), device=f"cuda:{i}")
        # get the information to reconstruct the shared tensor
        func, args = torch.multiprocessing.reductions.reduce_tensor(data)
        args = list(args)
        dist.broadcast_object_list([(func, args)], src=0)
        recv = [None]
        dist.broadcast_object_list(recv, src=1)
        success = recv[0]
        if success:
            dist.barrier()
            assert torch.all(data == 1).item()
        consumer_done_event.wait()
        result_quque.put(success)

def consumer(j: int,
             init_method: str,
             cuda_visible_devices: Optional[str] = None):
    if cuda_visible_devices is not None:
        os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
    with mute_output():
        dist.init_process_group(
            backend="gloo",
            init_method=init_method,
            world_size=2,
            rank=1,
        )
        torch.cuda.set_device(j)
        recv = [None]
        dist.broadcast_object_list(recv, src=0)
        func: Callable
        args: List
        func, args = recv[0]  # type: ignore
        # `args[6]` is the device id
        # by default pytorch will use `i` from the producer
        # here we need to set it to `j` to test P2P access
        args[6] = j
        success = False
        try:
            data = func(*args)
            data += 1
            success = True
        except Exception:
            pass
        dist.broadcast_object_list([success], src=1)
        if success:
            torch.cuda.synchronize()
            dist.barrier()
            assert torch.all(data == 1).item()
            del data
        del func, args, recv
        import gc
        gc.collect()
        torch.cuda.empty_cache()

def can_actually_p2p(i, j):
    """
    Usually, checking if P2P access is enabled can be done by
    `torch.cuda.can_device_access_peer(i, j)`. However, sometimes
    the driver might be broken, and `torch.cuda.can_device_access_peer(i, j)`
    returns `True` even if P2P access is not actually possible.
    See https://github.com/vllm-project/vllm/issues/2728 and
    https://forums.developer.nvidia.com/t/direct-gpu-gpu-communication-does-not-seem-to-work-properly/283264/10
    Therefore, we have to perform a real P2P access to check if it is actually
    possible.

    Note on p2p and cuda IPC:
    Usually, one process uses one GPU:
    GPU i --&gt; cuda context i --&gt; tensor i --&gt; process i

    We need to combine p2p and cuda IPC, so that:
    GPU i --&gt; cuda context i --&gt; tensor i --&gt; process i
                                 |shared|
    GPU j --&gt; cuda context j --&gt; tensor j --&gt; process j
    That is to say, process i creates a tensor in GPU i, passes IPC handle to
    process j, and process j accesses the tensor in GPU j. Any operation on the
    tensor in process j will be reflected in the tensor in process i, because
    they are the same memory segment.
    It is important to note that process j accesses the tensor in GPU j, not
    GPU i. That's why we need p2p access. # noqa
    """
    cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', None)
    # pass the CUDA_VISIBLE_DEVICES to the child process
    # to make sure they see the same set of GPUs

    # make sure the temp file is not the same across different calls
    temp_path = tempfile.mktemp() + str(time.time())
    # create an empty file
    with open(temp_path, "w"):
        pass
    init_method = f"file://{temp_path}"

    # make sure the processes are spawned
    smp = mp.get_context("spawn")
    consumer_done_event = smp.Event()
    result_quque = smp.Queue()
    pi = smp.Process(target=producer,
                     args=(i, init_method, consumer_done_event, result_quque, cuda_visible_devices, ))
    pj = smp.Process(target=consumer,
                     args=(j, init_method, cuda_visible_devices, ))
    pi.start()
    pj.start()
    pj.join()
    consumer_done_event.set()
    pi.join()
    return result_quque.get()

if __name__ == '__main__':
    print(can_actually_p2p(0, 1))
</code></pre>
<p>It works when two GPUs have p2p capacity, but will produce a warning message if not:</p>
<blockquote>
<p>[rank0]:[W CudaIPCTypes.cpp:16] Producer process has been terminated before all shared CUDA tensors released. See Note [Sharing CUDA tensors]</p>
</blockquote>
<p>This is strange, though. I already take care of the process termination order.</p>
<p>One hypothesis is, if the consumer errors, and because it is cuda error, the process cannot destruct cuda shared memory.</p>
<p>Any idea on how can I get rid of the warning message? <img src="https://discuss.pytorch.org/images/emoji/apple/pray.png?v=12" title=":pray:" class="emoji" alt=":pray:" loading="lazy" width="20" height="20"></p>