Allow retire_workers to run concurrently

[fjetter]

Closes #8053

The added test is rather superficial but it is simple and conveys what we want to work. There is also a little heavier test already available that does a similar thing

distributed/distributed/tests/test_active_memory_manager.py

Lines 1230 to 1269 in 9a9d2c3

	@pytest.mark.slow
	@pytest.mark.parametrize("use_ReduceReplicas", [False, True])
	@gen_cluster(
	client=True,
	nthreads=[("", 1)] * 10,
	config={
	"distributed.scheduler.active-memory-manager.start": True,
	"distributed.scheduler.active-memory-manager.interval": 0.1,
	"distributed.scheduler.active-memory-manager.measure": "managed",
	"distributed.scheduler.active-memory-manager.policies": [
	{"class": "distributed.active_memory_manager.ReduceReplicas"},
	],
	},
	scheduler_kwargs={"transition_counter_max": 500_000},
	worker_kwargs={"transition_counter_max": 500_000},
	)
	async def test_RetireWorker_stress(c, s, *workers, use_ReduceReplicas):
	"""It is safe to retire the best part of a cluster in the middle of a computation"""
	if not use_ReduceReplicas:
	s.extensions["amm"].policies.clear()
	addrs = list(s.workers)
	random.shuffle(addrs)
	print(f"Removing all workers except {addrs[-1]}")
	# Note: Scheduler._lock effectively prevents multiple calls to retire_workers from
	# running at the same time. However, the lock only exists for the benefit of legacy
	# (non-AMM) rebalance() and replicate() methods. Once the lock is removed, these
	# calls will become parallel and the test *should* continue working.
	tasks = [asyncio.create_task(tensordot_stress(c))]
	await asyncio.sleep(1)
	tasks.append(asyncio.create_task(c.retire_workers(addrs[0:2])))
	await asyncio.sleep(1)
	tasks.append(asyncio.create_task(c.retire_workers(addrs[2:5])))
	await asyncio.sleep(1)
	tasks.append(asyncio.create_task(c.retire_workers(addrs[5:9])))
	await asyncio.gather(*tasks)
	assert set(s.workers) == {addrs[9]}

running those locally didn't reveal any problems. if CI is not sadder than usual, I'm fine with this

Regarding the conditions: This will prohibit running replicate/rebalance while retire_workers is running and will prohibit running anything else while replicate/rebalance is running.
The alternative is to nuke these APIs entirely but the way I see it this complexity is not too bad. I considered factoring this out but this abstraction would look aweful since the semantics we need here are hard to generalize. This just felt unnecessary. Happy to discuss either option.

[fjetter]

FWIW I believe we also have pretty good consensus to nuke replicate/rebalance. Happy either way (from my POV this was actually less work but I'm happy if somebody comes around deleting code)

[github-actions]

Unit Test Results

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

       20 files   -        1         20 suites   - 1   9h 36m 3s ⏱️ - 1h 16m 47s
  3 821 tests +       6    3 707 ✔️ +       3     107 💤 ±  0    7 ❌ +3 
35 347 runs   - 1 516  33 624 ✔️  - 1 427  1 712 💤  - 92  11 ❌ +3 

For more details on these failures, see this check.

Results for commit 9b63ce6. ± Comparison against base commit b14f6ae.

♻️ This comment has been updated with latest results.

[crusaderky]

I like this, but I think the synchronisation primitive is complex and hard to read enough to warrant a bit of encapsulation. I'm refactoring it now.

[crusaderky]

ready for review (assuming green tests)

[crusaderky]

@fjetter all review comments have been addressed; ready for review again

[fjetter]

LGTM (I cannot approve since I'm the original author)