def collective_fn(input_tensor, output_tensor, context):
pygloo.allreduce(
context, gloo_util.get_tensor_ptr(input_tensor),
gloo_util.get_tensor_ptr(output_tensor),
gloo_util.get_tensor_n_elements(input_tensor),
gloo_util.get_gloo_tensor_dtype(input_tensor),
gloo_util.get_gloo_reduce_op(allreduce_options.reduceOp))
def test_redis(rank, world_size, redis_ip_address, redis_port, redis_password):
'''
rank # Rank of this process within list of participating processes
world_size # Number of participating processes
'''
context = pygloo.rendezvous.Context(rank, world_size)
attr = pygloo.transport.tcp.attr("localhost")
# Perform rendezvous for TCP pairs
dev = pygloo.transport.tcp.CreateDevice(attr)
redisStore = pygloo.rendezvous.RedisStore(redis_ip_address, redis_port)
redisStore.authorize(redis_password)
store = pygloo.rendezvous.PrefixStore("default", redisStore)
context.connectFullMesh(store, dev)
print("Using RedisStore rendezvous, connect successful!!")
sendbuf = np.array([[1, 2, 3], [1, 2, 3]], dtype=np.float32)
recvbuf = np.zeros_like(sendbuf, dtype=np.float32)
sendptr = sendbuf.ctypes.data
recvptr = recvbuf.ctypes.data
data_size = sendbuf.size if isinstance(
sendbuf, np.ndarray) else sendbuf.numpy().size
datatype = pygloo.glooDataType_t.glooFloat32
op = pygloo.ReduceOp.SUM
algorithm = pygloo.allreduceAlgorithm.RING
pygloo.allreduce(context, sendptr, recvptr, data_size, datatype, op,
algorithm)
print(f"rank {rank} sends {sendbuf}, receives {recvbuf}")
def test_multiGroup(rank, world_size, redis_ip_address, redis_port,
redis_password):
'''
test the multiGroup without prefixStore
rank # Rank of this process within list of participating processes
world_size # Number of participating processes
'''
groups = [f"multiGroup{i}" for i in range(3)]
contexts = {}
for group_name in groups:
context = pygloo.rendezvous.Context(rank, world_size)
attr = pygloo.transport.tcp.attr("localhost")
# Perform rendezvous for TCP pairs
dev = pygloo.transport.tcp.CreateDevice(attr)
redisStore = pygloo.rendezvous.RedisStore(redis_ip_address, redis_port)
redisStore.authorize(redis_password)
context.connectFullMesh(redisStore, dev)
if rank == 0:
keys = []
keys += [f"rank_{i}" for i in range(world_size)]
keys += [f"{i}" for i in range(world_size)]
redisStore.delKeys(keys)
contexts[group_name] = context
print("Using RedisStore rendezvous, connect successful!!")
for group_name in groups:
context = contexts[group_name]
sendbuf = np.array([[1, 2, 3], [1, 2, 3]], dtype=np.float32)
recvbuf = np.zeros_like(sendbuf, dtype=np.float32)
sendptr = sendbuf.ctypes.data
recvptr = recvbuf.ctypes.data
data_size = sendbuf.size if isinstance(
sendbuf, np.ndarray) else sendbuf.numpy().size
datatype = pygloo.glooDataType_t.glooFloat32
op = pygloo.ReduceOp.SUM
algorithm = pygloo.allreduceAlgorithm.RING
pygloo.allreduce(context, sendptr, recvptr, data_size, datatype, op,
algorithm)
print(f"rank {rank} sends {sendbuf}, receives {recvbuf}")
def test_barrier(rank, world_size, fileStore_path):
'''
rank # Rank of this process within list of participating processes
world_size # Number of participating processes
'''
if rank == 0:
if os.path.exists(fileStore_path):
shutil.rmtree(fileStore_path)
os.makedirs(fileStore_path)
else:
time.sleep(0.5)
context = pygloo.rendezvous.Context(rank, world_size)
attr = pygloo.transport.tcp.attr("localhost")
# Perform rendezvous for TCP pairs
dev = pygloo.transport.tcp.CreateDevice(attr)
fileStore = pygloo.rendezvous.FileStore(fileStore_path)
store = pygloo.rendezvous.PrefixStore(str(world_size), fileStore)
context.connectFullMesh(store, dev)
sendbuf = np.array([[1, 2, 3], [1, 2, 3]], dtype=np.float32)
recvbuf = np.zeros_like(sendbuf, dtype=np.float32)
sendptr = sendbuf.ctypes.data
recvptr = recvbuf.ctypes.data
# sendbuf = torch.Tensor([[1,2,3],[1,2,3]]).float()
# recvbuf = torch.zeros_like(sendbuf)
# sendptr = sendbuf.data_ptr()
# recvptr = recvbuf.data_ptr()
data_size = sendbuf.size if isinstance(
sendbuf, np.ndarray) else sendbuf.numpy().size
datatype = pygloo.glooDataType_t.glooFloat32
op = pygloo.ReduceOp.SUM
algorithm = pygloo.allreduceAlgorithm.RING
pygloo.allreduce(context, sendptr, recvptr, data_size, datatype, op,
algorithm)
pygloo.barrier(context)
print(f"rank {rank} sends {sendbuf}, receives {recvbuf}")