chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,529 @@
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
|
||||
import ray
|
||||
import ray.experimental.internal_kv as internal_kv
|
||||
from ray.util.collective import (
|
||||
allreduce,
|
||||
broadcast,
|
||||
create_collective_group,
|
||||
init_collective_group,
|
||||
)
|
||||
from ray.util.collective.backend_registry import (
|
||||
_global_registry,
|
||||
register_collective_backend,
|
||||
)
|
||||
from ray.util.collective.collective_group.base_collective_group import BaseGroup
|
||||
from ray.util.collective.types import (
|
||||
AllGatherOptions,
|
||||
AllReduceOptions,
|
||||
BarrierOptions,
|
||||
BroadcastOptions,
|
||||
RecvOptions,
|
||||
ReduceOp,
|
||||
ReduceOptions,
|
||||
ReduceScatterOptions,
|
||||
SendOptions,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
|
||||
def _unregister_collective_backend(name: str) -> None:
|
||||
"""Helper function to unregister a backend for testing purposes."""
|
||||
upper_name = name.upper()
|
||||
if upper_name in _global_registry._map:
|
||||
del _global_registry._map[upper_name]
|
||||
|
||||
|
||||
def get_data_key(group_name: str, rank: int, op_name: str):
|
||||
return f"collective_mock_{group_name}_{op_name}_rank_{rank}"
|
||||
|
||||
|
||||
def get_barrier_key(group_name: str, barrier_id: int):
|
||||
return f"collective_mock_{group_name}_barrier_{barrier_id}"
|
||||
|
||||
|
||||
class MockInternalKVGroup(BaseGroup):
|
||||
def __init__(self, world_size: int, rank: int, group_name: str):
|
||||
super().__init__(world_size, rank, group_name)
|
||||
self._barrier_counter = 0
|
||||
|
||||
@classmethod
|
||||
def backend(cls):
|
||||
return "MOCK"
|
||||
|
||||
@classmethod
|
||||
def check_backend_availability(cls) -> bool:
|
||||
return True
|
||||
|
||||
def _check_tensor_input(self, tensor):
|
||||
assert isinstance(tensor, list) and len(tensor) == 1
|
||||
t = tensor[0]
|
||||
if isinstance(t, np.ndarray):
|
||||
return t
|
||||
try:
|
||||
import torch
|
||||
|
||||
if isinstance(t, torch.Tensor):
|
||||
return t
|
||||
except ImportError:
|
||||
pass
|
||||
raise ValueError(
|
||||
f"MockInternalKVGroup only only accepts numpy.ndarray or torch.Tensor, received {type(t)}"
|
||||
)
|
||||
|
||||
def _serialize_tensor(self, tensor):
|
||||
if isinstance(tensor, np.ndarray):
|
||||
return tensor.tobytes(), tensor.shape, tensor.dtype
|
||||
try:
|
||||
import torch
|
||||
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
return (
|
||||
tensor.cpu().numpy().tobytes(),
|
||||
tensor.shape,
|
||||
tensor.cpu().numpy().dtype,
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
raise ValueError(f"Unsupported tensor type: {type(tensor)}")
|
||||
|
||||
def _deserialize_tensor(self, data: bytes, shape, dtype, target_tensor):
|
||||
if isinstance(target_tensor, np.ndarray):
|
||||
np_array = np.frombuffer(data, dtype=dtype).reshape(shape)
|
||||
target_tensor[:] = np_array
|
||||
else:
|
||||
try:
|
||||
import torch
|
||||
|
||||
if isinstance(target_tensor, torch.Tensor):
|
||||
np_array = np.frombuffer(data, dtype=dtype).reshape(shape)
|
||||
target_tensor.copy_(torch.from_numpy(np_array))
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
def broadcast(self, tensor, broadcast_options=BroadcastOptions()):
|
||||
tensor = self._check_tensor_input(tensor)
|
||||
root_rank = broadcast_options.root_rank
|
||||
|
||||
data_key = get_data_key(self._group_name, root_rank, "broadcast")
|
||||
|
||||
if self._rank == root_rank:
|
||||
data, shape, dtype = self._serialize_tensor(tensor)
|
||||
internal_kv._internal_kv_put(data_key, data)
|
||||
internal_kv._internal_kv_put(f"{data_key}_shape", str(shape))
|
||||
internal_kv._internal_kv_put(f"{data_key}_dtype", dtype.name)
|
||||
else:
|
||||
deadline_s = time.time() + 30.0
|
||||
while True:
|
||||
data = internal_kv._internal_kv_get(data_key)
|
||||
if data is not None:
|
||||
break
|
||||
if time.time() > deadline_s:
|
||||
raise TimeoutError(
|
||||
f"Timed out waiting for broadcast data from rank {root_rank}"
|
||||
)
|
||||
time.sleep(0.01)
|
||||
|
||||
deadline_s = time.time() + 30.0
|
||||
while True:
|
||||
shape_data = internal_kv._internal_kv_get(f"{data_key}_shape")
|
||||
dtype_data = internal_kv._internal_kv_get(f"{data_key}_dtype")
|
||||
if shape_data is not None and dtype_data is not None:
|
||||
break
|
||||
if time.time() > deadline_s:
|
||||
raise TimeoutError(
|
||||
f"Timed out waiting for broadcast metadata from rank {root_rank}"
|
||||
)
|
||||
time.sleep(0.01)
|
||||
|
||||
shape_str = shape_data.decode()
|
||||
shape = eval(shape_str)
|
||||
dtype_name = dtype_data.decode()
|
||||
dtype = np.dtype(dtype_name)
|
||||
|
||||
self._deserialize_tensor(data, shape, dtype, tensor)
|
||||
|
||||
def allreduce(self, tensor, allreduce_options=AllReduceOptions()):
|
||||
tensor = self._check_tensor_input(tensor)
|
||||
reduce_op = allreduce_options.reduceOp
|
||||
|
||||
data_key = get_data_key(self._group_name, self._rank, "allreduce")
|
||||
done_key = (
|
||||
f"collective_mock_{self._group_name}_allreduce_done_rank_{self._rank}"
|
||||
)
|
||||
|
||||
data, shape, dtype = self._serialize_tensor(tensor)
|
||||
internal_kv._internal_kv_put(data_key, data)
|
||||
internal_kv._internal_kv_put(f"{data_key}_shape", str(shape))
|
||||
internal_kv._internal_kv_put(f"{data_key}_dtype", dtype.name)
|
||||
internal_kv._internal_kv_put(done_key, b"1")
|
||||
|
||||
deadline_s = time.time() + 30.0
|
||||
while True:
|
||||
all_done = True
|
||||
for r in range(self._world_size):
|
||||
key = f"collective_mock_{self._group_name}_allreduce_done_rank_{r}"
|
||||
if internal_kv._internal_kv_get(key) is None:
|
||||
all_done = False
|
||||
break
|
||||
if all_done:
|
||||
break
|
||||
if time.time() > deadline_s:
|
||||
raise TimeoutError(
|
||||
"Timed out waiting for allreduce data from all ranks"
|
||||
)
|
||||
time.sleep(0.01)
|
||||
|
||||
result = None
|
||||
|
||||
for r in range(self._world_size):
|
||||
rank_data_key = get_data_key(self._group_name, r, "allreduce")
|
||||
rank_data = internal_kv._internal_kv_get(rank_data_key)
|
||||
rank_shape_data = internal_kv._internal_kv_get(f"{rank_data_key}_shape")
|
||||
rank_dtype_data = internal_kv._internal_kv_get(f"{rank_data_key}_dtype")
|
||||
|
||||
rank_shape_str = rank_shape_data.decode()
|
||||
rank_shape = eval(rank_shape_str)
|
||||
rank_dtype_name = rank_dtype_data.decode()
|
||||
rank_dtype = np.dtype(rank_dtype_name)
|
||||
|
||||
if isinstance(tensor, np.ndarray):
|
||||
rank_tensor = np.frombuffer(rank_data, dtype=rank_dtype).reshape(
|
||||
rank_shape
|
||||
)
|
||||
else:
|
||||
import torch
|
||||
|
||||
rank_np = np.frombuffer(rank_data, dtype=rank_dtype).reshape(rank_shape)
|
||||
rank_tensor = torch.from_numpy(rank_np)
|
||||
|
||||
if result is None:
|
||||
result = (
|
||||
rank_tensor.copy()
|
||||
if isinstance(rank_tensor, np.ndarray)
|
||||
else rank_tensor.clone()
|
||||
)
|
||||
else:
|
||||
if reduce_op == ReduceOp.SUM:
|
||||
result += rank_tensor
|
||||
elif reduce_op == ReduceOp.PRODUCT:
|
||||
result *= rank_tensor
|
||||
elif reduce_op == ReduceOp.MAX:
|
||||
if isinstance(result, np.ndarray):
|
||||
result = np.maximum(result, rank_tensor)
|
||||
else:
|
||||
import torch
|
||||
|
||||
result = torch.maximum(result, rank_tensor)
|
||||
elif reduce_op == ReduceOp.MIN:
|
||||
if isinstance(result, np.ndarray):
|
||||
result = np.minimum(result, rank_tensor)
|
||||
else:
|
||||
import torch
|
||||
|
||||
result = torch.minimum(result, rank_tensor)
|
||||
|
||||
if isinstance(tensor, np.ndarray):
|
||||
tensor[:] = result
|
||||
else:
|
||||
import torch
|
||||
|
||||
if isinstance(result, np.ndarray):
|
||||
tensor.copy_(torch.from_numpy(result))
|
||||
else:
|
||||
tensor.copy_(result)
|
||||
|
||||
def barrier(self, barrier_options=BarrierOptions()):
|
||||
barrier_id = self._barrier_counter
|
||||
barrier_key = get_barrier_key(self._group_name, barrier_id)
|
||||
rank_key = f"{barrier_key}_rank_{self._rank}"
|
||||
|
||||
internal_kv._internal_kv_put(rank_key, b"1")
|
||||
|
||||
deadline_s = time.time() + 30.0
|
||||
while True:
|
||||
all_arrived = True
|
||||
for r in range(self._world_size):
|
||||
key = f"{barrier_key}_rank_{r}"
|
||||
if internal_kv._internal_kv_get(key) is None:
|
||||
all_arrived = False
|
||||
break
|
||||
if all_arrived:
|
||||
break
|
||||
if time.time() > deadline_s:
|
||||
raise TimeoutError("Timed out waiting for barrier")
|
||||
time.sleep(0.01)
|
||||
|
||||
self._barrier_counter += 1
|
||||
|
||||
def reduce(self, tensor, reduce_options=ReduceOptions()):
|
||||
raise NotImplementedError("reduce is not implemented in MockInternalKVGroup")
|
||||
|
||||
def allgather(self, tensor_list, tensor, allgather_options=AllGatherOptions()):
|
||||
raise NotImplementedError("allgather is not implemented in MockInternalKVGroup")
|
||||
|
||||
def reducescatter(
|
||||
self, tensor, tensor_list, reducescatter_options=ReduceScatterOptions()
|
||||
):
|
||||
raise NotImplementedError(
|
||||
"reducescatter is not implemented in MockInternalKVGroup"
|
||||
)
|
||||
|
||||
def send(self, tensor, send_options: SendOptions):
|
||||
raise NotImplementedError("send is not implemented in MockInternalKVGroup")
|
||||
|
||||
def recv(self, tensor, recv_options: RecvOptions):
|
||||
raise NotImplementedError("recv is not implemented in MockInternalKVGroup")
|
||||
|
||||
|
||||
def test_mock_backend_create_group():
|
||||
"""Test using create_collective_group (driver-managed approach).
|
||||
|
||||
In this approach:
|
||||
- Driver calls create_collective_group() to declare the group
|
||||
- Workers only need to register the backend
|
||||
- Workers do NOT call init_collective_group()
|
||||
- The group is automatically initialized when workers call collective ops
|
||||
"""
|
||||
|
||||
ray.init()
|
||||
register_collective_backend("MOCK", MockInternalKVGroup)
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self, rank):
|
||||
self.rank = rank
|
||||
|
||||
def setup(self):
|
||||
from ray.util.collective.backend_registry import register_collective_backend
|
||||
|
||||
register_collective_backend("MOCK", MockInternalKVGroup)
|
||||
|
||||
def broadcast_test(self):
|
||||
if self.rank == 0:
|
||||
tensor = np.array([42.0, 43.0, 44.0], dtype=np.float32)
|
||||
else:
|
||||
tensor = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||||
|
||||
broadcast(tensor, src_rank=0)
|
||||
return tensor.tolist()
|
||||
|
||||
def allreduce_test(self):
|
||||
tensor = np.array([float(self.rank + 1)], dtype=np.float32)
|
||||
allreduce(tensor, op=ReduceOp.SUM)
|
||||
return tensor.item()
|
||||
|
||||
actors = [Worker.remote(rank=i) for i in range(3)]
|
||||
|
||||
create_collective_group(
|
||||
actors=actors,
|
||||
world_size=3,
|
||||
ranks=[0, 1, 2],
|
||||
backend="MOCK",
|
||||
group_name="default",
|
||||
)
|
||||
|
||||
ray.get([a.setup.remote() for a in actors])
|
||||
|
||||
results = ray.get([a.broadcast_test.remote() for a in actors])
|
||||
expected = [[42.0, 43.0, 44.0]] * 3
|
||||
if results == expected:
|
||||
print("Broadcast test passed!")
|
||||
else:
|
||||
print(f"Broadcast test failed! Expected {expected}, got {results}")
|
||||
|
||||
results = ray.get([a.allreduce_test.remote() for a in actors])
|
||||
|
||||
if results == [6.0, 6.0, 6.0]:
|
||||
print("AllReduce test passed!")
|
||||
else:
|
||||
print(f"AllReduce test failed! Expected [6.0, 6.0, 6.0], got {results}")
|
||||
|
||||
ray.shutdown()
|
||||
_unregister_collective_backend("MOCK")
|
||||
print("test_mock_backend_create_group completed!")
|
||||
|
||||
|
||||
def test_mock_backend_init_group():
|
||||
"""Test using init_collective_group (worker-managed approach).
|
||||
|
||||
In this approach:
|
||||
- Workers call init_collective_group() inside their setup method
|
||||
- Driver does NOT call create_collective_group()
|
||||
- Each worker explicitly initializes its own group membership
|
||||
"""
|
||||
|
||||
ray.init()
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self, rank):
|
||||
self.rank = rank
|
||||
|
||||
def setup(self, world_size):
|
||||
from ray.util.collective.backend_registry import register_collective_backend
|
||||
|
||||
register_collective_backend("MOCK", MockInternalKVGroup)
|
||||
|
||||
init_collective_group(
|
||||
world_size=world_size,
|
||||
rank=self.rank,
|
||||
backend="MOCK",
|
||||
group_name="default",
|
||||
)
|
||||
|
||||
def broadcast_test(self):
|
||||
if self.rank == 0:
|
||||
tensor = np.array([42.0, 43.0, 44.0], dtype=np.float32)
|
||||
else:
|
||||
tensor = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||||
|
||||
broadcast(tensor, src_rank=0)
|
||||
return tensor.tolist()
|
||||
|
||||
def allreduce_test(self):
|
||||
tensor = np.array([float(self.rank + 1)], dtype=np.float32)
|
||||
allreduce(tensor, op=ReduceOp.SUM)
|
||||
return tensor.item()
|
||||
|
||||
actors = [Worker.remote(rank=i) for i in range(3)]
|
||||
|
||||
# Do NOT call create_collective_group here
|
||||
ray.get([a.setup.remote(3) for a in actors])
|
||||
|
||||
results = ray.get([a.broadcast_test.remote() for a in actors])
|
||||
expected = [[42.0, 43.0, 44.0]] * 3
|
||||
if results == expected:
|
||||
print("Broadcast test passed!")
|
||||
else:
|
||||
print(f"Broadcast test failed! Expected {expected}, got {results}")
|
||||
|
||||
results = ray.get([a.allreduce_test.remote() for a in actors])
|
||||
|
||||
if results == [6.0, 6.0, 6.0]:
|
||||
print("AllReduce test passed!")
|
||||
else:
|
||||
print(f"AllReduce test failed! Expected [6.0, 6.0, 6.0], got {results}")
|
||||
|
||||
ray.shutdown()
|
||||
_unregister_collective_backend("MOCK")
|
||||
print("test_mock_backend_init_group completed!")
|
||||
|
||||
|
||||
def test_mock_backend_worker_not_registered():
|
||||
"""Test error handling when backend is not registered in worker.
|
||||
|
||||
This test uses create_collective_group (driver-managed approach).
|
||||
The driver registers the backend, but workers do not.
|
||||
When workers try to call collective ops, they should fail.
|
||||
|
||||
Note: We use world_size=1 to avoid "Unhandled error" messages
|
||||
from multiple workers failing simultaneously.
|
||||
"""
|
||||
|
||||
ray.init()
|
||||
register_collective_backend("MOCK", MockInternalKVGroup)
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self, rank):
|
||||
self.rank = rank
|
||||
|
||||
def broadcast_test(self):
|
||||
tensor = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||||
broadcast(tensor, src_rank=0)
|
||||
return tensor.tolist()
|
||||
|
||||
# Use single actor to avoid multiple "Unhandled error" messages
|
||||
actors = [Worker.remote(rank=0)]
|
||||
|
||||
create_collective_group(
|
||||
actors=actors,
|
||||
world_size=1,
|
||||
ranks=[0],
|
||||
backend="MOCK",
|
||||
group_name="default",
|
||||
)
|
||||
|
||||
test_passed = False
|
||||
try:
|
||||
ray.get([a.broadcast_test.remote() for a in actors])
|
||||
print("ERROR: Should have raised an exception for missing registration!")
|
||||
except Exception as e:
|
||||
if "not registered" in str(e) or "not initialized" in str(e):
|
||||
print(
|
||||
"Test passed! Correctly raised error for missing worker registration."
|
||||
)
|
||||
test_passed = True
|
||||
else:
|
||||
print(f"ERROR: Unexpected error: {e}")
|
||||
|
||||
ray.shutdown()
|
||||
_unregister_collective_backend("MOCK")
|
||||
|
||||
if not test_passed:
|
||||
print("Test failed!")
|
||||
|
||||
|
||||
def test_mock_backend_driver_not_registered():
|
||||
"""Test error handling when backend is not registered on driver.
|
||||
|
||||
This test uses create_collective_group, but the driver doesn't
|
||||
register the backend first, so it should fail immediately.
|
||||
"""
|
||||
|
||||
ray.init()
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self, rank):
|
||||
self.rank = rank
|
||||
|
||||
actors = [Worker.remote(rank=i) for i in range(2)]
|
||||
|
||||
try:
|
||||
create_collective_group(
|
||||
actors=actors,
|
||||
world_size=2,
|
||||
ranks=[0, 1],
|
||||
backend="MOCK",
|
||||
group_name="default",
|
||||
)
|
||||
print("ERROR: Should have raised an exception for missing registration!")
|
||||
except Exception as e:
|
||||
if "not registered" in str(e):
|
||||
print(
|
||||
"Test passed! Correctly raised error for missing driver registration."
|
||||
)
|
||||
else:
|
||||
print(f"ERROR: Unexpected error: {e}")
|
||||
|
||||
ray.shutdown()
|
||||
_unregister_collective_backend("MOCK")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print("Test 1: create_collective_group approach (driver-managed)")
|
||||
print("=" * 60)
|
||||
test_mock_backend_create_group()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 2: init_collective_group approach (worker-managed)")
|
||||
print("=" * 60)
|
||||
test_mock_backend_init_group()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 3: Error handling - worker not registered")
|
||||
print("=" * 60)
|
||||
test_mock_backend_worker_not_registered()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 4: Error handling - driver not registered")
|
||||
print("=" * 60)
|
||||
test_mock_backend_driver_not_registered()
|
||||
@@ -0,0 +1,37 @@
|
||||
import cupy as cp
|
||||
|
||||
import ray
|
||||
import ray.util.collective as collective
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
self.send = cp.ones((4,), dtype=cp.float32)
|
||||
self.recv = cp.zeros((4,), dtype=cp.float32)
|
||||
|
||||
def setup(self, world_size, rank):
|
||||
collective.init_collective_group(world_size, rank, "nccl", "default")
|
||||
return True
|
||||
|
||||
def compute(self):
|
||||
collective.allreduce(self.send, "default")
|
||||
return self.send
|
||||
|
||||
def destroy(self):
|
||||
collective.destroy_group()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
send = cp.ones((4,), dtype=cp.float32)
|
||||
ray.init(num_gpus=2)
|
||||
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))
|
||||
_ = ray.get(init_rets)
|
||||
results = ray.get([w.compute.remote() for w in workers])
|
||||
ray.shutdown()
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
import cupy as cp
|
||||
|
||||
import ray
|
||||
import ray.util.collective as collective
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
self.send = cp.ones((4,), dtype=cp.float32)
|
||||
|
||||
def compute(self):
|
||||
collective.allreduce(self.send, "177")
|
||||
return self.send
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ray.init(num_gpus=2)
|
||||
|
||||
num_workers = 2
|
||||
workers = []
|
||||
for i in range(num_workers):
|
||||
w = Worker.remote()
|
||||
workers.append(w)
|
||||
_options = {
|
||||
"group_name": "177",
|
||||
"world_size": 2,
|
||||
"ranks": [0, 1],
|
||||
"backend": "nccl",
|
||||
}
|
||||
collective.create_collective_group(workers, **_options)
|
||||
results = ray.get([w.compute.remote() for w in workers])
|
||||
ray.shutdown()
|
||||
@@ -0,0 +1,43 @@
|
||||
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
|
||||
|
||||
self.recv = cp.zeros((4,), dtype=cp.float32)
|
||||
|
||||
def setup(self, world_size, rank):
|
||||
collective.init_collective_group(world_size, rank, "nccl", "177")
|
||||
return True
|
||||
|
||||
def compute(self):
|
||||
collective.allreduce_multigpu([self.send1, self.send2], "177")
|
||||
return [self.send1, self.send2], self.send1.device, self.send2.device
|
||||
|
||||
def destroy(self):
|
||||
collective.destroy_collective_group("177")
|
||||
|
||||
|
||||
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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user