def recv_obj(self, src, group=None):
"""Receives an object from a source `src`."""
if group is None:
group = self.main_group
size = torch.tensor(1, dtype=torch.int32)
dist.irecv(size, src=src, group=group).wait()
data = torch.empty(size=(size,), dtype=torch.int8)
dist.irecv(data, src=src, group=group).wait()
buf = data.numpy().tobytes()
return serial.restricted_loads(buf)
def broadcast_obj(self, obj, src, group=None):
"""Broadcasts a given object to all parties."""
if group is None:
group = self.main_group
if self.rank == src:
assert obj is not None, "src party must provide obj for broadcast"
buf = pickle.dumps(obj)
size = torch.tensor(len(buf), dtype=torch.int32)
arr = torch.from_numpy(numpy.copy(numpy.frombuffer(buf, dtype=numpy.int8)))
dist.broadcast(size, src, group=group)
dist.broadcast(arr, src, group=group)
else:
size = torch.tensor(1, dtype=torch.int32)
dist.broadcast(size, src, group=group)
data = torch.empty(size=(size,), dtype=torch.int8)
dist.broadcast(data, src, group=group)
buf = data.numpy().tobytes()
obj = serial.restricted_loads(buf)
return obj