Skip to content

Commit

Permalink
a future must be linked to a client
Browse files Browse the repository at this point in the history
  • Loading branch information
sanderegg committed Dec 4, 2024
1 parent f5d9f45 commit e5708ec
Showing 1 changed file with 9 additions and 3 deletions.
12 changes: 9 additions & 3 deletions services/director-v2/tests/unit/test_modules_dask_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -613,7 +613,9 @@ def fake_sidecar_fct(
)
assert published_computation_task[0].node_id in image_params.fake_tasks
# creating a new future shows that it is not done????
assert not distributed.Future(published_computation_task[0].job_id).done()
assert not distributed.Future(
published_computation_task[0].job_id, client=dask_client.backend.client
).done()

# as the task is published on the dask-scheduler when sending, it shall still be published on the dask scheduler
list_of_persisted_datasets = await dask_client.backend.client.list_datasets() # type: ignore
Expand All @@ -636,7 +638,9 @@ def fake_sidecar_fct(
assert isinstance(task_result, TaskOutputData)
assert task_result.get("some_output_key") == 123
# try to create another future and this one is already done
assert distributed.Future(published_computation_task[0].job_id).done()
assert distributed.Future(
published_computation_task[0].job_id, client=dask_client.backend.client
).done()


async def test_abort_computation_tasks(
Expand Down Expand Up @@ -1022,7 +1026,9 @@ def fake_remote_fct(

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

await _assert_wait_for_task_status(
Expand Down

0 comments on commit e5708ec

Please sign in to comment.