Skip to content

Commit

Permalink
fixup prefill logic that could miss adding pkgs
Browse files Browse the repository at this point in the history
  • Loading branch information
antheas committed Aug 11, 2024
1 parent dc55b6f commit a5b19ba
Showing 1 changed file with 16 additions and 13 deletions.
29 changes: 16 additions & 13 deletions src/rechunk/alg.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ def prefill_layers(
curr = []
l_upd = np.zeros(n_segments, dtype=np.bool)
l_size = 0
while todo:
while True:
# We will fill layers in two steps:
# If the layer is emptly, we will insert the largest package
# in todo.
Expand All @@ -87,23 +87,26 @@ def prefill_layers(
# There will be packages left over in the end, which will be handled
# differently.

if not curr:
p = max(todo, key=lambda p: p.size)
todo.pop(p)
curr.append(p)
l_upd |= upd_matrix[p.index]
l_size += p.size
elif l_size > fill_size:
layers.append(curr)
logger.info(
f"Layer {dedi+len(layers):2d}: {l_size / 1e9:.3f} GB with {len(curr):3d} packages."
)
if len(layers) >= max_layers:
if l_size > fill_size or not todo:
if curr:
# Since this also gets hit with not todo
# curr might be empty, avoid creating a layer
layers.append(curr)
logger.info(
f"Layer {dedi+len(layers):2d}: {l_size / 1e9:.3f} GB with {len(curr):3d} packages."
)
if len(layers) >= max_layers or not todo:
break
curr = []
l_upd = np.zeros(n_segments, dtype=np.bool)
l_size = 0
pbar.update(1)
elif not curr:
p = max(todo, key=lambda p: p.size)
todo.pop(p)
curr.append(p)
l_upd |= upd_matrix[p.index]
l_size += p.size
else:
# Calculate the bandwidth increase for each package
# and select the one with the smallest increase
Expand Down

0 comments on commit a5b19ba

Please sign in to comment.