Problem with using mp.Queue with CUDA Tensors on AGX ORIN

Robotics & Edge Computing Jetson Systems Jetson AGX Orin
,
1

Hello,

In a multiprocessing code I am trying to transfer a GPU tensor between two processes. I use torch.multiprocessing.Queue to do the transfers. I have provided a basic code below that causes the problem. Either Queue is not accepting CUDA tensors or it cannot share GPU memory. It only happens on AGX ORIN device. When I try the same code on a workstation with multiple GPUs and CPUs it doesn’t happen and the code works.

On Orin Jetpack 5.1.2 is installed.
On the workstation I use pytorch2206 container from Nvidia.

My question is why is this error appearing on Orin and not on a Workstation? How can I make it work on Orin?
(I can share other details that you see relevant)

Thanks in advance,
Cem

import torch

def another_dummy_thread(batch_queue):
    while True:
        a = batch_queue.get()
        print('the tensor is', a)

def dummy_thread(batch_queue):
    while True:
        a = torch.tensor([1, 2], device='cuda')
        batch_queue.put(a)

if __name__ == "__main__":
    torch.multiprocessing.set_start_method("spawn")
    dummy_2_dummy_queue = torch.multiprocessing.Queue()

    another_dummy_thread_ = torch.multiprocessing.Process(target= another_dummy_thread, args=(dummy_2_dummy_queue,))
    dummy_thread_ = torch.multiprocessing.Process(target=dummy_thread, args=(dummy_2_dummy_queue,))

    dummy_thread_.start()
    another_dummy_thread_.start()

    dummy_thread_.join()
    another_dummy_thread_.join()

The Error:

Traceback (most recent call last):
File “/usr/lib/python3.8/multiprocessing/queues.py”, line 239, in _feed
obj = _ForkingPickler.dumps(obj)
File “/usr/lib/python3.8/multiprocessing/reduction.py”, line 51, in dumps
cls(buf, protocol).dump(obj)
File “/usr/local/lib/python3.8/dist-packages/torch/multiprocessing/reductions.py”, line 249, in reduce_tensor
event_sync_required) = storage.share_cuda()
File “/usr/local/lib/python3.8/dist-packages/torch/storage.py”, line 623, in share_cuda
return self._storage.share_cuda(*args, **kwargs)
RuntimeError: CUDA error: operation not supported

3

Hi,

This is not supported.

The feature is implemented with CUDA IPC which is not available on the Jetson platform.
For more details, please check the below topic:

Thanks.