Specify worker ressource for dask.Delayed object with dask.annotate

[juergend7lytix]

I tried to execute a python function via the "dask.delayed" interface on different workers according to their resources via the "dask.annotate" context manager. According to the doc, this should be possible, if I'm reading it correctly.

MWE
I start a scheduler, two workers with different resources (namely CPU, GPU) in a background thread.
I execute a simple function two times with dask.delayed: I make two dask.delayed objects, where the first, should be executed on worker 1, the second on worker 2. Via get_worker() I print out, which worker is active.

import asyncio
import os
from dask.distributed import Scheduler, Worker, Client, Future, get_worker
from dask import delayed, annotate
import threading
from time import sleep

def add(a, b):
    sleep(b)
    print(os.getpid(), get_worker().name)
    return a + b

async def start_server():
    s = await Scheduler()
    os.environ["SCHEDULERADDRESS"] = s.address
    w1 = await Worker(s.address, resources={"GPU": 0, "CPU": 1}, name="worker1")
    w2 = await Worker(s.address, resources={"GPU": 1, "CPU": 0}, name="worker2")

def start_server_in_thread(loop):
    asyncio.set_event_loop(loop)
    loop.create_task(start_server())
    loop.run_forever()

def init_dask():
    t = threading.Thread(target=start_server_in_thread,
                         args=(asyncio.get_event_loop(),))
    t.start()
    sleep(1)
	
if __name__ == "__main__":
    init_dask()
    c = Client(os.environ["SCHEDULERADDRESS"])
    with annotate(resources={"CPU": 0, "GPU": 1}):
        d1 = delayed(add)(2, 2)
    with annotate(resources={"CPU": 1, "GPU": 0}):
        d2 = delayed(add)(4, 4)

    print(d1.compute())
    print(d2.compute())

When I submit the tasks, via client.submit, I get the expected behavior.

import asyncio
import os
from dask.distributed import Scheduler, Worker, Client, Future, get_worker
from dask import delayed, annotate
import threading
from time import sleep

def add(a, b):
    sleep(b)
    print(os.getpid(), get_worker().name)
    return a + b

async def start_server():
    s = await Scheduler()
    os.environ["SCHEDULERADDRESS"] = s.address
    w1 = await Worker(s.address, resources={"GPU": 0, "CPU": 1}, name="worker1")
    w2 = await Worker(s.address, resources={"GPU": 1, "CPU": 0}, name="worker2")

def start_server_in_thread(loop):
    asyncio.set_event_loop(loop)
    loop.create_task(start_server())
    loop.run_forever()

def init_dask():
    t = threading.Thread(target=start_server_in_thread,
                         args=(asyncio.get_event_loop(),))
    t.start()
    sleep(1)
	
if __name__ == "__main__":
    init_dask()
    c = Client(os.environ["SCHEDULERADDRESS"])
    with annotate(resources={"CPU": 0, "GPU": 1}):
        f1 = c.submit(add, 2, 2)
    with annotate(resources={"CPU": 1, "GPU": 0}):
        f2 = c.submit(add, 4, 4)
	
    print(f1.result())
    print(f2.result())

I'm using

• python 3.8
• dask 2021.02.0, on
• Arch linux, with kernel 5.10.16-arch1-1

\edit: Edited the minimal example, such that it also runs through, when executing it outside REPL.

[jrbourbeau]

Thanks for raising an issue @juergend7lytix. It looks like you're not creating a Client, which is what will cause compute() calls to run on the distributed cluster you created, in your first code snippet. Does adding a Client improve the situation?

Also FWIW when I tried running the second code snippet in an ipython session I get KeyError: 'SCHEDULERADDRESS':

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-1-7578deb10f5c> in <module>
     26 t.start()
     27
---> 28 c = Client(os.environ["SCHEDULERADDRESS"])
     29 with annotate(resources={"CPU": 0, "GPU": 1}):
     30     f1 = c.submit(add, 2, 2)

~/miniforge3/envs/dask/lib/python3.8/os.py in __getitem__(self, key)
    673         except KeyError:
    674             # raise KeyError with the original key value
--> 675             raise KeyError(key) from None
    676         return self.decodevalue(value)
    677

KeyError: 'SCHEDULERADDRESS'

[juergend7lytix]

Thank you for the answer @jrbourbeau !

Also with creating explicitly the client, the problem still persists.

In fact, I tried before the same snippet with an explicit creation of the client, but as I saw, that dask doesn't need it (because it probably just creates in the background a client, and uses a scheduler object, which is created in the same python session), I left it out for the minimal working example.

And regarding the problem you encounter with my minimal example: Thanks very much for that! I forget to make a slight delay after spinning up the server and then creating the client! :-) I ran the code somehow line-by-line in PyCharm, that's the reason why I didn't have any issue! I rearranged the code a little bit, now it "works" when executing it all at once.

[jrbourbeau]

Thanks for updating those examples @juergend7lytix! I think you may be running into #7036. I've opened up #7298 to fix the issue -- if you have time to try out the changes in #7298 for your use case that would be appreciated (though no obligation)

[ian-r-rose]

Agreeing with @jrbourbeau here, it looks like you have indeed run into #7036. By default dask will attempt to perform various graph optimization steps upon calling compute, during which time they can lose their annotations.

In addition to the fix that @jrbourbeau is working on, you can try to disable graph optimizations for the time being to get the expected result:

    print(d1.compute(optimize_graph=False))
    print(d2.compute(optimize_graph=False))

[juergend7lytix]

Thank you very much @jrbourbeau @ian-r-rose !
Indeed #7036 resolves my issue!

Adding optimize_graph=False indeed gives me the expected behaviour in the MWE.