Smarter stealing with dependencies

[hendrikmakait]

• Based on (and blocked by) No longer double count transfer cost in stealing #7036

• Tests added / passed
• Passes pre-commit run --all-files

[github-actions]

Unit Test Results

See test report for an extended history of previous test failures. This is useful for diagnosing flaky tests.

       15 files  ±  0         15 suites  ±0   6h 5m 45s ⏱️ - 34m 8s
  3 146 tests +11    3 060 ✔️ +12    85 💤 ±0  1 ❌  - 1 
23 287 runs  +88  22 381 ✔️ +92  905 💤  - 3  1 ❌  - 1 

For more details on these failures, see this check.

Results for commit b5afd77. ± Comparison against base commit 8f36aa5.

♻️ This comment has been updated with latest results.

[hendrikmakait]

For better readability/lookup in VS Code, I made the in-line comments parametrization IDs. There are no changes to the actual test params.

[hendrikmakait]

Everything below is code for enabling the test cases above

[hendrikmakait]

A/B test: https://github.com/coiled/coiled-runtime/actions/runs/3068174742

N=2, so be cautious with the results. Notably, test_vorticity looks much better, yet test_basic_sum much worse. The remainder of the changes might be noise.

[fjetter]

@hendrikmakait you set the worker-saturation parameter as well https://github.com/coiled/coiled-runtime/blob/hendrik/AB_stealing_objective/AB_environments/AB_sample.dask.yaml which explains the huge impact

[hendrikmakait]

@hendrikmakait you set the worker-saturation parameter as well https://github.com/coiled/coiled-runtime/blob/hendrik/AB_stealing_objective/AB_environments/AB_sample.dask.yaml which explains the huge impact

🤦

[fjetter]

FYI I retriggered the computation already with updated info

[hendrikmakait]

A/B test: https://github.com/coiled/coiled-runtime/actions/runs/3135725432 (ignore AB_sample)
n: 7
baseline: 2022.9.1

With n==7, there appear to be no noticeable effects of the change.

[hendrikmakait]

I could not reproduce the flake of test_balance_prefers_busier_with_dependency locally in 1000 runs.
A possible cause might be described in this comment:

distributed/distributed/tests/test_steal.py

Lines 751 to 756 in 288305f

	# NOTE: There is a timing issue that workers may already start executing
	# tasks before we call balance, i.e. the workers will reject the
	# stealing request and we end up with a different end result.
	# Particularly tests with many input tasks are more likely to fail since
	# the test setup takes longer and allows the workers more time to
	# schedule a task on the threadpool

[hendrikmakait]

I could not reproduce the flake of test_balance_prefers_busier_with_dependency locally in 1000 runs. A possible cause might be described in this comment:

distributed/distributed/tests/test_steal.py

Lines 751 to 756 in 288305f

	# NOTE: There is a timing issue that workers may already start executing
	# tasks before we call balance, i.e. the workers will reject the
	# stealing request and we end up with a different end result.
	# Particularly tests with many input tasks are more likely to fail since
	# the test setup takes longer and allows the workers more time to
	# schedule a task on the threadpool

The culprit is worker-saturation (see #7085). To reproduce this locally, make sure to enable root task withholding by setting worker-saturation as hinted at by Tests / test (ubuntu-latest, 3.9, queue, ci1). It's easy to miss that one.

[hendrikmakait]

CI flakes:

• test_secede_rejoin_busy flaky #5275

[hendrikmakait]

It looks like test_balance flakes after merging with main (likely due to AMM being enabled). I'm investigating and filing a follow-up PR.