diff --git a/examples/cuda-benchmark/demo_receiver.py b/examples/cuda-benchmark/demo_receiver.py index ea8609ae..0b2c86d6 100644 --- a/examples/cuda-benchmark/demo_receiver.py +++ b/examples/cuda-benchmark/demo_receiver.py @@ -1,7 +1,6 @@ #!/usr/bin/env python """TODO: Add docstring.""" - import os import time @@ -9,7 +8,7 @@ import numpy as np import pyarrow as pa import torch from dora import Node -from dora.cuda import cudabuffer_to_torch, ipc_buffer_to_ipc_handle +from dora.cuda import ipc_buffer_to_ipc_handle, open_ipc_handle from helper import record_results from tqdm import tqdm @@ -17,7 +16,6 @@ torch.tensor([], device="cuda") pa.array([]) -context = pa.cuda.Context() node = Node("node_2") current_size = 8 @@ -29,8 +27,6 @@ DEVICE = os.getenv("DEVICE", "cuda") NAME = f"dora torch {DEVICE}" -ctx = pa.cuda.Context() - print() print("Receiving 40MB packets using default dora-rs") @@ -49,13 +45,13 @@ while True: if event["metadata"]["device"] != "cuda": # BEFORE handle = event["value"].to_numpy() + scope = None torch_tensor = torch.tensor(handle, device="cuda") else: # AFTER - # storage needs to be spawned in the same file as where it's used. Don't ask me why. ipc_handle = ipc_buffer_to_ipc_handle(event["value"]) - cudabuffer = ctx.open_ipc_buffer(ipc_handle) - torch_tensor = cudabuffer_to_torch(cudabuffer, event["metadata"]) # on cuda + scope = open_ipc_handle(ipc_handle, event["metadata"]) + torch_tensor = scope.__enter__() else: break t_received = time.perf_counter_ns() @@ -73,6 +69,9 @@ while True: latencies = [] n += 1 + if scope: + scope.__exit__(None, None, None) + mean_cuda = np.array(latencies).mean() pbar.close() @@ -81,9 +80,9 @@ time.sleep(2) print() print("----") -print(f"Node communication duration with default dora-rs: {mean_cpu/1000:.1f}ms") -print(f"Node communication duration with dora CUDA->CUDA: {mean_cuda/1000:.1f}ms") +print(f"Node communication duration with default dora-rs: {mean_cpu / 1000:.1f}ms") +print(f"Node communication duration with dora CUDA->CUDA: {mean_cuda / 1000:.1f}ms") print("----") -print(f"Speed Up: {(mean_cpu)/(mean_cuda):.0f}") +print(f"Speed Up: {(mean_cpu) / (mean_cuda):.0f}") record_results(NAME, current_size, latencies) diff --git a/examples/cuda-benchmark/receiver.py b/examples/cuda-benchmark/receiver.py index 8115ed16..5c67a5a2 100644 --- a/examples/cuda-benchmark/receiver.py +++ b/examples/cuda-benchmark/receiver.py @@ -1,14 +1,13 @@ #!/usr/bin/env python """TODO: Add docstring.""" - import os import time import pyarrow as pa import torch from dora import Node -from dora.cuda import cudabuffer_to_torch, ipc_buffer_to_ipc_handle +from dora.cuda import ipc_buffer_to_ipc_handle, open_ipc_handle from helper import record_results from tqdm import tqdm @@ -17,7 +16,6 @@ torch.tensor([], device="cuda") pa.array([]) pbar = tqdm(total=100) -context = pa.cuda.Context() node = Node("node_2") @@ -29,8 +27,6 @@ DEVICE = os.getenv("DEVICE", "cuda") NAME = f"dora torch {DEVICE}" -ctx = pa.cuda.Context() - while True: event = node.next() if event["type"] == "INPUT": @@ -40,12 +36,12 @@ while True: # BEFORE handle = event["value"].to_numpy() torch_tensor = torch.tensor(handle, device="cuda") + scope = None else: # AFTER - # storage needs to be spawned in the same file as where it's used. Don't ask me why. ipc_handle = ipc_buffer_to_ipc_handle(event["value"]) - cudabuffer = ctx.open_ipc_buffer(ipc_handle) - torch_tensor = cudabuffer_to_torch(cudabuffer, event["metadata"]) # on cuda + scope = open_ipc_handle(ipc_handle, event["metadata"]) + torch_tensor = scope.__enter__() else: break t_received = time.perf_counter_ns() @@ -53,7 +49,6 @@ while True: if length != current_size: if n > 0: - pbar.close() pbar = tqdm(total=100) record_results(NAME, current_size, latencies) @@ -69,4 +64,7 @@ while True: n += 1 i += 1 + if scope: + scope.__exit__(None, None, None) + record_results(NAME, current_size, latencies)