Skip to content

Commit 586218a

Browse files
committed
a future must be linked to a client
1 parent 2e67bb7 commit 586218a

File tree

1 file changed

+9
-3
lines changed

1 file changed

+9
-3
lines changed

services/director-v2/tests/unit/test_modules_dask_client.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -683,7 +683,9 @@ def fake_sidecar_fct(
683683
)
684684
assert published_computation_task[0].node_id in image_params.fake_tasks
685685
# creating a new future shows that it is not done????
686-
assert not distributed.Future(published_computation_task[0].job_id).done()
686+
assert not distributed.Future(
687+
published_computation_task[0].job_id, client=dask_client.backend.client
688+
).done()
687689

688690
# as the task is published on the dask-scheduler when sending, it shall still be published on the dask scheduler
689691
list_of_persisted_datasets = await dask_client.backend.client.list_datasets() # type: ignore
@@ -706,7 +708,9 @@ def fake_sidecar_fct(
706708
assert isinstance(task_result, TaskOutputData)
707709
assert task_result.get("some_output_key") == 123
708710
# try to create another future and this one is already done
709-
assert distributed.Future(published_computation_task[0].job_id).done()
711+
assert distributed.Future(
712+
published_computation_task[0].job_id, client=dask_client.backend.client
713+
).done()
710714

711715

712716
async def test_abort_computation_tasks(
@@ -1111,7 +1115,9 @@ def fake_remote_fct(
11111115

11121116
assert published_computation_task[0].node_id in cpu_image.fake_tasks
11131117
# let's get a dask future for the task here so dask will not remove the task from the scheduler at the end
1114-
computation_future = distributed.Future(key=published_computation_task[0].job_id)
1118+
computation_future = distributed.Future(
1119+
key=published_computation_task[0].job_id, client=dask_client.backend.client
1120+
)
11151121
assert computation_future
11161122

11171123
await _assert_wait_for_task_status(

0 commit comments

Comments
 (0)