Skip to content

Commit

Permalink
Note what struct lines are packing/unpacking
Browse files Browse the repository at this point in the history
  • Loading branch information
jakirkham committed Apr 21, 2020
1 parent 98d82dd commit c59f95d
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions distributed/comm/ucx.py
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,11 @@ async def write(
]

# Send meta data

# Send # of frames (uint64)
await self.ep.send(struct.pack("Q", nframes))
# Send which frames are CUDA (bool) and
# how large each frame is (uint64)
await self.ep.send(
struct.pack(nframes * "?" + nframes * "Q", *cuda_frames, *sizes)
)
Expand Down Expand Up @@ -222,11 +226,15 @@ async def read(self, deserializers=("cuda", "dask", "pickle", "error")):

try:
# Recv meta data

# Recv # of frames (uint64)
nframes_fmt = "Q"
nframes = host_array(struct.calcsize(nframes_fmt))
await self.ep.recv(nframes)
(nframes,) = struct.unpack(nframes_fmt, nframes)

# Recv which frames are CUDA (bool) and
# how large each frame is (uint64)
header_fmt = nframes * "?" + nframes * "Q"
header = host_array(struct.calcsize(header_fmt))
await self.ep.recv(header)
Expand Down

0 comments on commit c59f95d

Please sign in to comment.