Dask annotations not working when connecting to an existing cluster #47

andreas-ntonas · 2022-11-01T23:54:31Z

I have noticed that Prefect does not respect the dask resources annotations when using an existing cluster

Minimal reproducible example:

I have setup a local distributed cluster like this:

First I spin up a local scheduler:

dask scheduler

And then I start 2 workers. The 1st one has 1 thread and 1 "annotation" resource available and the other one has 3 threads and no specified special resources available:

dask worker --nthreads '1' --resources "annotation=1"  tcp://192.168.1.13:8786 --name special_worker

&

dask worker --nthreads '3'  tcp://192.168.1.13:8786 --name only_regular_tasks

Suppose we have the following flow that gets executed 6 times and submits to Dask its tasks for execution:

#!/usr/bin/env python3
import dask
from prefect import flow, task
from prefect_dask.task_runners import DaskTaskRunner
from time import sleep

@task
def show_number(x: int):
    print(x)

@task
def say_hello():
    print("Hello!")
    sleep(4)

@flow(task_runner=DaskTaskRunner(address="tcp://192.168.1.13:8786"))
def example_flow(x):
    show_number.submit(x)
    with dask.annotate(resources={'annotation': 1}):
        say_hello.submit()

if __name__ == "__main__":
    for i in range(6):
        example_flow(i)

I would expect say_hello() tasks to execute only on the special_worker and the show_number() tasks to whichever available.

However this is not the case and I get:

On the 1st special worker:

dask worker --nthreads '1' --resources "annotation=1"  tcp://192.168.1.13:8786 --name special_worker
2022-11-02 01:40:21,707 - distributed.nanny - INFO -         Start Nanny at: 'tcp://192.168.1.13:40111'
2022-11-02 01:40:21,911 - distributed.worker - INFO -       Start worker at:   tcp://192.168.1.13:38027
2022-11-02 01:40:21,911 - distributed.worker - INFO -          Listening to:   tcp://192.168.1.13:38027
2022-11-02 01:40:21,911 - distributed.worker - INFO -           Worker name:             special_worker
2022-11-02 01:40:21,911 - distributed.worker - INFO -          dashboard at:         192.168.1.13:44211
2022-11-02 01:40:21,911 - distributed.worker - INFO - Waiting to connect to:    tcp://192.168.1.13:8786
2022-11-02 01:40:21,911 - distributed.worker - INFO - -------------------------------------------------
2022-11-02 01:40:21,911 - distributed.worker - INFO -               Threads:                          1
2022-11-02 01:40:21,911 - distributed.worker - INFO -                Memory:                   1.92 GiB
2022-11-02 01:40:21,911 - distributed.worker - INFO -       Local Directory: /tmp/dask-worker-space/worker-5zbq1zho
2022-11-02 01:40:21,911 - distributed.worker - INFO - -------------------------------------------------
2022-11-02 01:40:21,915 - distributed.worker - INFO -         Registered to:    tcp://192.168.1.13:8786
2022-11-02 01:40:21,915 - distributed.worker - INFO - -------------------------------------------------
2022-11-02 01:40:21,916 - distributed.core - INFO - Starting established connection
0
Hello!
Hello!
3
Hello!
Hello!

On the 2nd only_regular_tasks worker:

dask worker --nthreads '3'  tcp://192.168.1.13:8786 --name only_regular_tasks
2022-11-02 01:40:44,071 - distributed.nanny - INFO -         Start Nanny at: 'tcp://192.168.1.13:37631'
2022-11-02 01:40:44,271 - distributed.worker - INFO -       Start worker at:   tcp://192.168.1.13:41407
2022-11-02 01:40:44,271 - distributed.worker - INFO -          Listening to:   tcp://192.168.1.13:41407
2022-11-02 01:40:44,271 - distributed.worker - INFO -           Worker name:         only_regular_tasks
2022-11-02 01:40:44,271 - distributed.worker - INFO -          dashboard at:         192.168.1.13:43921
2022-11-02 01:40:44,271 - distributed.worker - INFO - Waiting to connect to:    tcp://192.168.1.13:8786
2022-11-02 01:40:44,271 - distributed.worker - INFO - -------------------------------------------------
2022-11-02 01:40:44,271 - distributed.worker - INFO -               Threads:                          3
2022-11-02 01:40:44,271 - distributed.worker - INFO -                Memory:                   5.76 GiB
2022-11-02 01:40:44,271 - distributed.worker - INFO -       Local Directory: /tmp/dask-worker-space/worker-b0x7pdsc
2022-11-02 01:40:44,271 - distributed.worker - INFO - -------------------------------------------------
2022-11-02 01:40:44,276 - distributed.worker - INFO -         Registered to:    tcp://192.168.1.13:8786
2022-11-02 01:40:44,276 - distributed.worker - INFO - -------------------------------------------------
2022-11-02 01:40:44,276 - distributed.core - INFO - Starting established connection
Hello!
1
2
Hello!
4
5

As you can see the say_hello() tasks are getting executed on the 2nd worker that does not have the appropriate --resources "annotation=1" available

Relevant versions used in a Python 3.10 env: prefect==2.6.5, prefect-dask==0.2.1, dask==2022.10.2, distributed==2022.10.2

The text was updated successfully, but these errors were encountered:

zanieb · 2022-11-02T14:43:34Z

Ah this used to work and was the intended way to use annotations. I think this may be caused by PrefectHQ/prefect#6527 — it'll require further investigation though.

BitTheByte · 2022-12-12T03:13:18Z

@madkinsz I also encountered this while using dask_kubernetes to create dask cluster. some tasks have their correct annotations and others don't have any. it seems like there is a problem with the context manager

zanieb · 2022-12-21T19:59:25Z

We made an optimization to task submission but that means that tasks are submitted in the background now, which means that submission can actually occur outside of the context. I'm working on a fix.

BitTheByte · 2023-02-02T13:31:08Z

@madkinsz do you mind sharing the PR? I'd like to keep track of this issue and provide help if needed

tibuch · 2023-06-30T14:45:58Z

I just encountered this issue as well. I am using this config for the DaskTaskRunner:

cluster_class="dask_jobqueue.SLURMCluster",
    cluster_kwargs={
        "account": "...",
        "queue": "...",
        "cores": 5,
        "processes": 1,
        "memory": "5 GB",
        "walltime": "24:00:00",
        "worker_extra_args": [
            "--resources TEST=1",
        ],
    }

And then submit task with:

with dask.annotate(resources={"TEST": 1}):
    for i in range(10):
        task.submit(i)

This executes 5 tasks simultaneously on the same worker. As far as I can tell the limiting factor is the number of cores. Because if I set "cores": 1 only a single task is submitted at a time.

Jrodseth · 2023-08-01T17:42:53Z

Also experiencing this problem with temporary clusters on ECS. Workers can be started with resources configures as @tibuch showed with worker_extra_args in the cluster_kwargs e.g.

"worker_extra_args": [ "--resources", "MEMORY=4098"]

But using with dask.annotate(resources={"MEMORY": 4098}): doesn't prevent multiple tasks being submitted to workers.

chrisguidry · 2023-09-13T20:07:58Z

Hi @andreas-ntonas, thank you for the very detailed reproduction steps.

I'm working on reproducing this issue, but in my testing with the latest prefect and prefect-dask, I'm not able to reproduce this, and I believe I'm seeing the behavior you were expecting.

On the special worker:

$ dask worker --nthreads '1' --resources "annotation=1"  tcp://10.0.0.196:8786 --name special_worker
2023-09-13 15:57:23,709 - distributed.nanny - INFO -         Start Nanny at: 'tcp://10.0.0.196:35097'
2023-09-13 15:57:23,927 - distributed.worker - INFO -       Start worker at:     tcp://10.0.0.196:34437
2023-09-13 15:57:23,927 - distributed.worker - INFO -          Listening to:     tcp://10.0.0.196:34437
2023-09-13 15:57:23,927 - distributed.worker - INFO -           Worker name:             special_worker
2023-09-13 15:57:23,927 - distributed.worker - INFO -          dashboard at:           10.0.0.196:42291
2023-09-13 15:57:23,927 - distributed.worker - INFO - Waiting to connect to:      tcp://10.0.0.196:8786
2023-09-13 15:57:23,927 - distributed.worker - INFO - -------------------------------------------------
2023-09-13 15:57:23,927 - distributed.worker - INFO -               Threads:                          1
2023-09-13 15:57:23,927 - distributed.worker - INFO -                Memory:                   3.13 GiB
2023-09-13 15:57:23,927 - distributed.worker - INFO -       Local Directory: /tmp/dask-scratch-space/worker-uwkqglq5
2023-09-13 15:57:23,927 - distributed.worker - INFO - -------------------------------------------------
2023-09-13 15:57:23,932 - distributed.worker - INFO - Starting Worker plugin shuffle
2023-09-13 15:57:23,932 - distributed.worker - INFO -         Registered to:      tcp://10.0.0.196:8786
2023-09-13 15:57:23,932 - distributed.worker - INFO - -------------------------------------------------
2023-09-13 15:57:23,933 - distributed.core - INFO - Starting established connection to tcp://10.0.0.196:8786
Hello!
16:00:12.727 | INFO    | Task run 'say_hello-0' - Finished in state Completed()
Hello!
16:00:21.908 | INFO    | Task run 'say_hello-0' - Finished in state Completed()
Hello!
16:00:30.204 | INFO    | Task run 'say_hello-0' - Finished in state Completed()
Hello!
16:00:38.761 | INFO    | Task run 'say_hello-0' - Finished in state Completed()
Hello!
16:00:47.142 | INFO    | Task run 'say_hello-0' - Finished in state Completed()
Hello!
16:00:55.426 | INFO    | Task run 'say_hello-0' - Finished in state Completed()
Hello!
16:02:41.596 | INFO    | Task run 'say_hello-0' - Finished in state Completed()
Hello!
16:02:50.030 | INFO    | Task run 'say_hello-0' - Finished in state Completed()
Hello!
16:03:03.527 | INFO    | Task run 'say_hello-0' - Finished in state Completed()
Hello!
16:03:11.621 | INFO    | Task run 'say_hello-0' - Finished in state Completed()
Hello!
16:03:21.032 | INFO    | Task run 'say_hello-0' - Finished in state Completed()
Hello!
16:03:29.293 | INFO    | Task run 'say_hello-0' - Finished in state Completed()

On the regular worker:

$ dask worker --nthreads '3' tcp://10.0.0.196:8786 --name only_regular_tasks
2023-09-13 15:57:56,149 - distributed.nanny - INFO -         Start Nanny at: 'tcp://10.0.0.196:45717'
2023-09-13 15:57:56,398 - distributed.worker - INFO -       Start worker at:     tcp://10.0.0.196:34377
2023-09-13 15:57:56,398 - distributed.worker - INFO -          Listening to:     tcp://10.0.0.196:34377
2023-09-13 15:57:56,398 - distributed.worker - INFO -           Worker name:         only_regular_tasks
2023-09-13 15:57:56,398 - distributed.worker - INFO -          dashboard at:           10.0.0.196:42241
2023-09-13 15:57:56,398 - distributed.worker - INFO - Waiting to connect to:      tcp://10.0.0.196:8786
2023-09-13 15:57:56,398 - distributed.worker - INFO - -------------------------------------------------
2023-09-13 15:57:56,398 - distributed.worker - INFO -               Threads:                          3
2023-09-13 15:57:56,398 - distributed.worker - INFO -                Memory:                   9.40 GiB
2023-09-13 15:57:56,398 - distributed.worker - INFO -       Local Directory: /tmp/dask-scratch-space/worker-sz_hhz3t
2023-09-13 15:57:56,398 - distributed.worker - INFO - -------------------------------------------------
2023-09-13 15:57:56,402 - distributed.worker - INFO - Starting Worker plugin shuffle
2023-09-13 15:57:56,403 - distributed.worker - INFO -         Registered to:      tcp://10.0.0.196:8786
2023-09-13 15:57:56,403 - distributed.worker - INFO - -------------------------------------------------
2023-09-13 15:57:56,403 - distributed.core - INFO - Starting established connection to tcp://10.0.0.196:8786
0
16:00:08.917 | INFO    | Task run 'show_number-0' - Finished in state Completed()
1
16:00:17.455 | INFO    | Task run 'show_number-0' - Finished in state Completed()
2
16:00:26.501 | INFO    | Task run 'show_number-0' - Finished in state Completed()
3
16:00:34.491 | INFO    | Task run 'show_number-0' - Finished in state Completed()
4
16:00:43.324 | INFO    | Task run 'show_number-0' - Finished in state Completed()
5
16:00:51.269 | INFO    | Task run 'show_number-0' - Finished in state Completed()
0
16:02:32.554 | INFO    | Task run 'show_number-0' - Finished in state Completed()
1
16:02:46.101 | INFO    | Task run 'show_number-0' - Finished in state Completed()
2
16:02:59.269 | INFO    | Task run 'show_number-0' - Finished in state Completed()
3
16:03:07.540 | INFO    | Task run 'show_number-0' - Finished in state Completed()
4
16:03:16.945 | INFO    | Task run 'show_number-0' - Finished in state Completed()
5
16:03:25.199 | INFO    | Task run 'show_number-0' - Finished in state Completed()

I know it's been some time since the original report, would you be willing to test this again using these versions:

prefect==2.13.0
prefect-dask==0.2.5
dask==2023.9.1
distributed==2023.9.1

And then let me know if you're still having the same trouble?

chrisguidry · 2023-09-15T17:26:11Z

I'm going to close this one out as fixed, as we're unable to reproduce this on our end. There were a LOT of improvements to how tasks are handled between November and today, and I believe this was fixed as a side-effect.

If anyone is still having trouble with this, please open a new issue with repro steps.

BitTheByte · 2023-09-15T19:59:27Z

Hey @chrisguidry

Although this is fixed, there is a race condition somewhere as I see different tasks having other task annotations, I can't really provide a reproducible example due to the nature of the problem but it occur in flows that looks like the below

with annotate(priority=1000):
    task1.map(...)

with annotate(priority=-1000):
    task2.map(...)

chrisguidry · 2023-09-18T13:32:09Z

@BitTheByte thanks for the feedback, can you open a fresh issue with all the details you have available (prefect* and dask* versions, the behavior you're seeing and what you expect, etc)?

ghislainp mentioned this issue Apr 2, 2023

Trouble using get_dask_client with a Coiled cluster PrefectHQ/prefect#12971

Open

chrisguidry self-assigned this Sep 13, 2023

chrisguidry closed this as completed Sep 15, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Dask annotations not working when connecting to an existing cluster #47

Dask annotations not working when connecting to an existing cluster #47

andreas-ntonas commented Nov 1, 2022 •

edited

Loading

zanieb commented Nov 2, 2022

BitTheByte commented Dec 12, 2022

zanieb commented Dec 21, 2022

BitTheByte commented Feb 2, 2023

tibuch commented Jun 30, 2023

Jrodseth commented Aug 1, 2023

chrisguidry commented Sep 13, 2023

chrisguidry commented Sep 15, 2023

BitTheByte commented Sep 15, 2023

chrisguidry commented Sep 18, 2023

Dask annotations not working when connecting to an existing cluster #47

Dask annotations not working when connecting to an existing cluster #47

Comments

andreas-ntonas commented Nov 1, 2022 • edited Loading

I have noticed that Prefect does not respect the dask resources annotations when using an existing cluster

Minimal reproducible example:

I have setup a local distributed cluster like this:

On the 1st special worker:

On the 2nd only_regular_tasks worker:

zanieb commented Nov 2, 2022

BitTheByte commented Dec 12, 2022

zanieb commented Dec 21, 2022

BitTheByte commented Feb 2, 2023

tibuch commented Jun 30, 2023

Jrodseth commented Aug 1, 2023

chrisguidry commented Sep 13, 2023

chrisguidry commented Sep 15, 2023

BitTheByte commented Sep 15, 2023

chrisguidry commented Sep 18, 2023

andreas-ntonas commented Nov 1, 2022 •

edited

Loading