# sender.py
import torch
from dora import Node
from dora.cuda import torch_to_ipc_buffer

        
def main():
    sender = Node(node_id="sender")
    tensor1 = torch.arange(6, dtype=torch.float32, device="cuda").reshape(2, 3)
    tensor2 = torch.arange(6, dtype=torch.float32, device="cuda").reshape(2, 3) + 6.0
    tensor = torch.stack([tensor1, tensor2], dim=0)
    ipc_buffer, metadata = torch_to_ipc_buffer(tensor)
    for _ in range(100):
        sender.send_output("tensor", ipc_buffer, metadata)

if __name__ == "__main__":
    main()

# receiver.py
import pyarrow as pa
import torch
from dora import Node
from dora.cuda import cudabuffer_to_torch, ipc_buffer_to_ipc_handle


def main():
    receiver = Node(node_id="receiver")
    ctx = pa.cuda.Context()
    tensor1 = torch.arange(6, dtype=torch.float32, device="cuda").reshape(2, 3)
    tensor2 = torch.arange(6, dtype=torch.float32, device="cuda").reshape(2, 3) + 6.0
    expected_tensor = torch.stack([tensor1, tensor2], dim=0)

    while True:
        event = receiver.next(timeout=1e-3)
        if event["type"] == "INPUT":
            ipc_handle = ipc_buffer_to_ipc_handle(event["value"])
            cudabuffer = ctx.open_ipc_buffer(ipc_handle)
            tensor = cudabuffer_to_torch(cudabuffer, event["metadata"]).clone()
            assert torch.allclose(tensor, expected_tensor), f"Received tensor: {tensor} does not match expected tensor: {expected_tensor}"

if __name__ == "__main__":
    main()

# mwe.yml
nodes:
  - id: sender
    path: /path/to/sender.py
    outputs:
      - tensor
  - id: receiver
    path: /path/to/receiver.py
    inputs:
      tensor:
        source: sender/tensor
        queue_size: 1