Skip to content

Commit

Permalink
Fix calculation of MemoryShardsBuffer.bytes_read (#8289)
Browse files Browse the repository at this point in the history
  • Loading branch information
crusaderky authored Oct 19, 2023
1 parent cbc3a33 commit 7ea3bff
Showing 1 changed file with 5 additions and 8 deletions.
13 changes: 5 additions & 8 deletions distributed/shuffle/_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,16 +34,13 @@ def read(self, id: str) -> Any:
raise RuntimeError("Tried to read from file before done.")

with self.time("read"):
shards = self._shards.pop(id) # Raises KeyError
self.bytes_read += sum(map(sizeof, shards))
# Don't keep the serialized and the deserialized shards
# in memory at the same time
data = []
size = 0
shards = self._shards[id]
while shards:
shard = shards.pop()
data.append(self._deserialize(shard))
size += sizeof(shards)

if data:
self.bytes_read += size
return data
else:
raise KeyError(id)
return data

0 comments on commit 7ea3bff

Please sign in to comment.