Seceded tasks don't release resources if they fail

[crusaderky]

The long-running -> error worker transition does not release worker resources.

On a related note, there's an unused method transition_long_running_rescheduled that should be cleaned up.

[crusaderky]

Also, there's nothing that ever removes keys from the long_running set.

[RaphaelRobidas]

I launch tasks from tasks and secede in the initial tasks while their subtasks complete. It seems like the resources of the initial tasks are not released, which can lead to a situation where no work is getting done since every thread is waiting for resources to be freed up. This issue is quite similar to yours.

[fjetter]

xref #5882 this is not directly related but also connected to the handling of long-running state

[crusaderky]

@RaphaelRobidas to my knowledge resources are released upon successful completion. Do you have a reproducer that shows otherwise?

[RaphaelRobidas]

@crusaderky yes, but seceding during the task is not successful completion and does not release resources. In my understanding, seceding is meant to free up the current worker and allow it to execute other tasks. Logically, the resources should be temporarily released as well.

See the example below: the jobs will never complete and the whole cluster will get stuck doing nothing. Is this a related issue or am I using resources incorrectly?

import time

import dask
from dask.distributed import (
    Client,
    LocalCluster,
    get_client,
    secede,
    rejoin,
    as_completed,
)


def subjob(arg):
    print("Starting subjob")

    time.sleep(5)

    print("Subjob done")
    return 0


def job(arg):
    print("Starting job")

    time.sleep(1)  # Wait until all jobs are submitted
    client = get_client()
    futs = [client.submit(subjob, i, resources={"CPU": 1}) for i in range(5)]

    secede()
    rets = client.gather(futs)
    rejoin()

    time.sleep(5)
    print("Job done")

    return 0


if __name__ == "__main__":
    with dask.config.set({"distributed.worker.resources.CPU": 8}):
        cluster = LocalCluster(n_workers=1, processes=False)

        client = Client(cluster)

        futures = [
            client.submit(
                job,
                i,
                resources={"CPU": 1},
            )
            for i in range(16)
        ]

        rets = client.gather(futures)

        print("This will never be reached")

[crusaderky]

Logically, the resources should be temporarily released as well.

It makes sense for resources to be treated the same as threads by secede()/rejoin().
However, this detail was never thought through. I suspect that implementing it in rejoin() would be a major complication and I doubt that many users would benefit from it.
So i'm oriented towards just documenting that resources won't be released by secede().

Is this a related issue or am I using resources incorrectly?

The only case where you would use a "CPU" resource is if you needed to mix multithreaded and singlethreaded tasks. If all tasks use the same amount of CPU, then you should not use a resource for it - just set the number of threads.

[RaphaelRobidas]

Yes, this is actually my use case, I am calling subprocesses which use either one or multiple CPU cores and want to ensure that no oversubscribing occurs. I guess a workaround would be to request the amount of additional CPU cores for the subtasks and limiting the number of threads in the cluster to the number of cores. This wouldn't work for multiple independent machines as in a real cluster, however.

[RaphaelRobidas]

My workaround ended up not working very well at all. I need the ability to release resources when seceding for my uses, so I dug a bit in the code. I found that it is possible to manually release and reattribute resources when seceding/rejoining:

def job():
    [...]
    secede()
    with release_resources(CPU=1):
        res = client.gather(futures)
    rejoin()

I have defined the context release_resources, which modifies the available_resources of the current worker. The context exit blocks until the resources can be reserved for the current thread again. This is fine for my usage and works as expected, as far as I can tell.

Is there an obvious problem with this approach and is this something that would be worth merging into the main branch? I can prepare a PR so you can have a look.

[crusaderky]

Please never, ever bypass Worker.handle_stimulus and modify the WorkerState directly.

[RaphaelRobidas]

Thanks for the help, I adjusted the code to make use of Worker.handle_stimulus. I will submit a PR for those who might also benefit from it.