Update usage of get_worker() in tests

[pentschev]

In dask/distributed#7580 get_worker was modified to return the worker of a task, thus it cannot be used by client.run, and we must now use dask_worker as the first argument to client.run to obtain the worker.

[jacobtomlinson]

Looks like this test also needs updating

dask-cuda/dask_cuda/tests/test_proxify_host_file.py

Line 418 in ec2de78

async def test_worker_force_spill_to_disk():

[wence-]

@madsbk there are usages of get_worker inside the explicit comms implementation. Are these likely to be safe because they're run in a context where there is a task?

[pentschev]

Are these likely to be safe because they're run in a context where there is a task?

I guess you mean

dask-cuda/dask_cuda/explicit_comms/comms.py

Line 77 in ec2de78

worker: Any = get_worker()

which is called from

dask-cuda/dask_cuda/explicit_comms/comms.py

Lines 92 to 93 in ec2de78

	def _run_coroutine_on_worker(sessionId, coroutine, args):
	session_state = worker_state(sessionId)

and submitted from

dask-cuda/dask_cuda/explicit_comms/comms.py

Lines 228 to 235 in ec2de78

	ret = self.client.submit(
	_run_coroutine_on_worker,
	self.sessionId,
	coroutine,
	args,
	workers=[worker],
	pure=False,
	)

The above should be fine as _run_coroutine_on_worker and is executed by client.submit.

[madsbk]

LGTM, thanks Peter

[wence-]

The above should be fine as _run_coroutine_on_worker and is executed by client.submit.

884a595 says "Using client.submit seems not allowed to resolve get_worker, but client.run does."

So this seems in conflict?

[pentschev]

884a595 says "Using client.submit seems not allowed to resolve get_worker, but client.run does."

So this seems in conflict?

I meant that for async functions only, it does work with non-async. I've also commented about that behavior in dask/distributed#7696 (comment) , but I currently don't know if it's possible to resolve, although it shouldn't be a problem for us currently, as we don't use it in actual code, just in that one test.

[wence-]

Ah, ok, thanks for the explanation and digging!

[pentschev]

Thanks everyone for reviews!

[pentschev]

/merge

[pentschev]

For completeness in this thread, regarding the substitution of client.submit by client.run, this is what Distributed devs answered in dask/distributed#7696 (comment):

I looked at 884a595 and argue this "workaround" is how it is supposed to be. client.submit is intended to run user functions (on a thread) while client.run is intended to run a function/coroutine on the worker event loop, in the main thread, outside of our task scheduling. If you are actually inspecting the worker in your test, that's how you should do it