Add opt-in lazy Dask side inputs with worker-safe computation

[Gaurav598]

Summary

Add an opt-in --dask_lazy_side_inputs mode that consumes Dask Bag side inputs one partition at a time. Keep whole-bag evaluation as the default for compatibility. When either path runs inside a Dask worker, use Dask's worker_client() while waiting for nested computation so the worker task releases its scheduling slot.

Problem and root cause

DaskBagWindowedIterator originally evaluated list(self.bag) before yielding a value. This can materialize a large side input in one worker and exhaust its memory. Simply replacing it with sequential partition.compute() calls also introduces one scheduler round trip per partition and can block a worker task while its nested Dask tasks wait for a free slot. A one-worker, one-thread cluster is especially vulnerable to scheduler starvation.

Design

• The new flag is off by default. In opt-in mode, the iterator computes one partition when reached and preserves bag order and the existing window conversion.
• Outside a Dask worker, computation uses the configured Dask scheduler. Inside a worker, worker_client() temporarily releases the task slot during computation. Its context ends before values are yielded, including when a consumer stops early.
• Computation failures propagate unchanged. The code only catches ValueError from get_worker() to identify a non-worker context.
• The runner removes lazy_side_inputs from its options before constructing dask.distributed.Client.

This follows Dask's tasks-launching-tasks guidance. It also responds to the concerns raised on the earlier, closed PR #40283, including default-path safety, serial partition evaluation, and exception handling.

Testing

• Focused iterator tests: 5 passed. They cover default whole-bag behavior, partition-at-a-time consumption, empty partitions, repeated iteration, later-partition errors, order, and window conversion.
• Distributed tests with LocalCluster and Client: 4 passed. End-to-end Beam pipelines exercise AsIter, AsList, and AsSingleton with one worker/one thread and two workers. A worker-side later-partition error also propagates. Each distributed case runs in a separate process with a 45-second watchdog.
• Full dask_runner_test.py: 36 passed, 1 existing expected failure, 5 subtests passed. The run reported 22 warnings, including a missing pytest-timeout plugin configuration and Dask local dashboard port warnings.
• yapf --diff on the three changed files, python -m compileall -q on those files, and git diff --check: passed.

The tests ran with Beam's Python SDK and Dask 2026.8.0 installed in an isolated virtual environment. The distributed tests use Dask's distributed scheduler with in-process workers; they do not cover a remote multi-host cluster.

Compatibility and limitations

The default still materializes the entire side-input bag, so the issue's memory limitation remains for users who do not opt in. In opt-in mode, peak iterator materialization is bounded by the largest computed partition, but AsList must still materialize its complete view. Partitions are computed serially and may incur repeated graph work or scheduler round trips when views are consumed again or by multiple worker tasks. A shared or storage-backed side-input design would require separate design work.

Related to #40282

[github-actions]

Assigning reviewers:

R: @claudevdm for label python.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[Gaurav598]

R: @tvalentyn

Hi @tvalentyn, I reviewed your feedback on #40283, particularly the concerns around nested Dask computations, potential worker deadlocks, and changing the default behavior.

I've updated this PR to make partition-at-a-time evaluation opt-in, preserving the existing default behavior. Both paths now use worker_client() for nested computations inside Dask workers, and I've added distributed regression tests covering single-worker/single-thread and multi-worker execution.
The focused and distributed tests pass locally. The PR description also documents the remaining memory and performance trade-offs.

Would you be willing to review the updated approach when you have time? I'd appreciate your feedback. Thanks!

[github-actions]

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

[codecov]

Codecov Report

❌ Patch coverage is 0% with 21 lines in your changes missing coverage. Please review.
✅ Project coverage is 56.25%. Comparing base (76f6e6d) to head (2466846).
⚠️ Report is 140 commits behind head on master.

Files with missing lines	Patch %	Lines
...on/apache_beam/runners/dask/transform_evaluator.py	0.00%	17 Missing ⚠️
...dks/python/apache_beam/runners/dask/dask_runner.py	0.00%	4 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #40496      +/-   ##
============================================
+ Coverage     56.16%   56.25%   +0.08%     
  Complexity     2288     2288              
============================================
  Files          1124     1124              
  Lines        178336   178409      +73     
  Branches       1489     1489              
============================================
+ Hits         100170   100361     +191     
+ Misses        75645    75527     -118     
  Partials       2521     2521              

Flag	Coverage Δ
python	79.71% <0.00%> (+0.16%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Gaurav598]

Hi @claudevdm,

I've pushed a follow-up commit to address the failing Codecov patch coverage check.

The Dask-specific tests were passing, but the main Python coverage workflow skips them because Dask isn't installed in that environment. The new commit adds a Linux-only coverage upload to the existing Dask workflow, without changing its original test matrix or test commands.

The focused tests pass locally, with 17 of 21 changed production lines covered. The new Codecov result is still pending verification.

The latest GitHub Actions runs are currently blocked with action_required. Could a maintainer please review the workflow change and approve the pending runs if it looks appropriate?

Thanks!