Fix for accelerated DAG GPU microbenchmarks (ray-project#46260)

The GPU microbenchmark code became out-of-sync with the accelerated DAG library. This PR fixes the microbenchmark code. Signed-off-by: Jack Humphries <[email protected]>
harborn · Jun 27, 2024 · 1e2599f · 1e2599f
1 parent 692c0ec
commit 1e2599f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/release/microbenchmark/experimental/accelerated_dag_gpu_microbenchmark.py b/release/microbenchmark/experimental/accelerated_dag_gpu_microbenchmark.py
@@ -113,7 +113,7 @@ def exec_ray_dag(
  TorchTensorType(
  "auto" if dynamic_shape else SHAPE,
  "auto" if dynamic_shape else DTYPE,
- transport="nccl" if use_nccl else None,
+ transport="nccl" if use_nccl else "auto",
  )
  )
 
@@ -124,9 +124,9 @@ def exec_ray_dag(
 
  def _run():
  i = np.random.randint(100)
- output_channel = dag.execute(i)
+ ref = dag.execute(i)
  # TODO(swang): Replace with fake ObjectRef.
- result = output_channel.read()
+ result = ray.get(ref)
  assert result == (i, SHAPE, DTYPE)
 
  else: