Skip to content

fix: race condition during merging #42

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alexander-held wants to merge 5 commits into
base: main
Choose a base branch
from
Open

fix: race condition during merging #42

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ca649e9
use distributed.as_completed instead
alexander-held Jan 20, 2026
07ff4e8
hide progress bars when subtasks complete
alexander-held Jan 20, 2026
578c11c
minor fixes
alexander-held Jan 20, 2026
cecb991
keep pbars visible
alexander-held Jan 20, 2026
9b8f3a3
flip order
alexander-held Jan 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 10 additions & 15 deletions util/_dask.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,13 +41,13 @@
from coffea.processor import Accumulatable, accumulate
from coffea.processor.executor import WorkItem
from coffea.util import rich_bar
from dask.distributed import Client
from dask.distributed import as_completed, Client
from dask.tokenize import tokenize
from rich.console import Group
from rich.live import Live
from rich.progress import Progress

from util._futures import DynamicAsCompleted, FutureLike
from util._futures import FutureLike

try:
from coffea.util import coffea_console
Expand Down Expand Up @@ -198,15 +198,10 @@ def dask_reduce(
)

failed_items: defaultdict[list[Failure]] = defaultdict(list)
dynac = DynamicAsCompleted(futures)
ac = as_completed(futures)

# in-dataset merging loop, we merge first within datasets to avoid large accumulators in memory
# some reasonable value for the batch_size:
# yield in batches of treereduction, and we want at least 1 item per batch
batch_size = min(
treereduction, max(int(len(futures) / 100), 1)
) # this is heuristic, can be tuned
for batch in dynac.iter_batches(batch_size=batch_size):
for batch in ac.batches():
for future in batch:
ds = key2ds[future.key]

Expand Down Expand Up @@ -281,9 +276,9 @@ def dask_reduce(
buf.clear()

# add back to the ac, recursively merge
dynac.add(work)
ac.add(work)

del dynac
del ac

# not needed anymore
pbar_merge_count.clear()
Expand Down Expand Up @@ -314,8 +309,8 @@ def dask_reduce(
# track how much to advance
pbar_merge_count = {}

dynac = DynamicAsCompleted(final_merge_futures.values())
for future in dynac:
ac = as_completed(final_merge_futures.values())
for future in ac:
if failed_future(future):
raise future.exception()

Expand Down Expand Up @@ -353,10 +348,10 @@ def dask_reduce(
buf.clear()

# add back to the ac, recursively merge
dynac.add(future)
ac.add(future)

# not needed anymore
del dynac
del ac
pbar_merge_count.clear()

pbars[_final_merge_sentinel].update(
Expand Down