chore: import upstream snapshot with attribution

python/ray/util/collective/examples/nccl_p2p_example_multigpu.py

@@ -0,0 +1,53 @@
+import cupy as cp
+from cupy.cuda import Device
+
+import ray
+import ray.util.collective as collective
+
+
+@ray.remote(num_gpus=2)
+class Worker:
+    def __init__(self):
+        with Device(0):
+            self.send1 = cp.ones((4,), dtype=cp.float32)
+        with Device(1):
+            self.send2 = cp.ones((4,), dtype=cp.float32) * 2
+
+        with Device(0):
+            self.recv1 = cp.zeros((4,), dtype=cp.float32)
+        with Device(1):
+            self.recv2 = cp.zeros((4,), dtype=cp.float32)
+        self.rank = -1
+
+    def setup(self, world_size, rank):
+        self.rank = rank
+        collective.init_collective_group(world_size, rank, "nccl", "8")
+        return True
+
+    def compute(self):
+        if self.rank == 0:
+            with Device(0):
+                collective.send_multigpu(self.send1 * 2, 1, 1, "8")
+        else:
+            # with Device(1):
+            collective.recv_multigpu(self.recv2, 0, 0, "8")
+        return self.recv2
+
+    def destroy(self):
+        collective.destroy_collective_group("8")
+
+
+if __name__ == "__main__":
+    ray.init(address="auto")
+    num_workers = 2
+    workers = []
+    init_rets = []
+    for i in range(num_workers):
+        w = Worker.remote()
+        workers.append(w)
+        init_rets.append(w.setup.remote(num_workers, i))
+    a = ray.get(init_rets)
+    results = ray.get([w.compute.remote() for w in workers])
+    print(results)
+    ray.get([w.destroy.remote() for w in workers])
+    ray.shutdown()