Respect priority in earliest start time heuristic

[mrocklin]

Background

Today Dask decides where to place a task based on an "earliest start time" heuristic. For a given task it tries to find the worker on which it will start the soonest. It does this by computing two values:

• The amount of work currently on the worker (currently measured by occupancy)
• The amount of time it would take to transfer all dependencies not on the worker to that worker

This is what is computed by the worker_objective function

Problem

However, this is incorrect if the task that we're considering has higher priority than some of the tasks currently running on that worker. In that case looking at the occupancy of the worker is incorrect, because this task gets to cut in line.

But in general while we can count all of the work that is higher priority than this task, that might be somewhat expensive, especially in cases where there are lots of tasks on a worker (which is common). This might be the kind of thing where Cython can save us, but even then I'm not sure.

Proposed solutions

Let's write down a few possible solutions:

1. Brute force: we can look at all possible tasks in ws.processing and count up the amount of occupancy of tasks with higher priority
2. Ignore: We could ignore occupancy altogether, and just let work stealing take charge
3. Middle ground: We could randomly take a few tasks in ws.processing (maybe four?) and see where we stand among those four. If we're worse then all of them then great, we take the full brunt of occupancy. If we're better than all of them then we take 0%. If we're in the middle then we take 50% and so on.
4. fancy: we maintain some sort of t-digest per worker. This seems extreme, but we would only need to track like three quantile values for this to work well most of the time.
5. less fancy: maybe we track min/max/mean and blend between them?

3 and 5 seem like the most probable. Each has some concerns:

3. Sampling: I'm not sure how best to get these items. iter(seq(...)) is ordered these days, and so not a great sample. random.sample is somewhat expensive. %timeit random.sample(list(d), 2) takes 8us for me for a dict with 1000 items.
4. min/max/mean: Our priorities are hierarchical, and so mean (or any quantile) is a little wonky.

Importance

This is, I suspect, especially important in workloads where we have lots of rootish tasks, which are common. The variance among all of those tasks can easily swamp the signal that tasks should stay where their data is.

[fjetter]

A second effect where our occupation heuristic is off is that assigned tasks might share dependencies. Currently we double count this and inflate the occupancy. This is less important for root-ish tasks, though.

---

Regarding 5. the only thing I can really come up with to deal with the "wonkyness" of our priorities is to keep a sorted list/deque of our priorities. estimating quantiles would then be a bisect which should be relatively cheap but I don't know if this is cheap enough (bisect + insert on deque to maintain sort is ~1us on my machine)

[mrocklin]

I'm curious about inserting into a sorted deque. I'm not familiar with this as an option. Can you say more about what you're thinking here.

I did a tiny (/ probably wrong) micro-benchmark with a list, assuming 1000 tasks per worker.

In [1]: L = list(range(1000))

In [2]: import bisect

In [3]: %timeit bisect.bisect(L, 400)
198 ns ± 7.16 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

In [4]: %%timeit
   ...: L2 = L[:400] + [400] + L[400:]  # insert
   ...: L3 = L[:400] + L[401:]          # delete
   ...: 
   ...: 
7.84 µs ± 139 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

I think that 8us is probably too much to spend on this improvement

For comparison with the random solution

In [7]: %timeit random.choices(L, k=5)
978 ns ± 8.76 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

[fjetter]

Deques are implemented as a linked list afaik (or at least similar). That gives us the opportunity to have (almost; I don't know the exact implementation) constant time insert performance. Lists have an O(N) insert performance since they are implemented using an array.

This is also a "wrong benchmark" but it's hard to benchmark the insert properly. If we subtract the copy performance very naively we're back to ns performance for the insert. bisect also works on deques.

In [1]: from collections import deque

In [2]: D = deque(range(1_000))

In [3]: %timeit D.copy()
3.61 µs ± 23.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [4]: %%timeit
   ...: D_c = D.copy()
   ...: D_c.insert(123, 123)
   ...:
   ...:
3.68 µs ± 15.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [5]: import bisect

In [6]: %timeit bisect.bisect(D, 123)
196 ns ± 1.84 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

[fjetter]

Very naive cumsum occupancy could then look like the following but I don't have a good recipe for removal, therefore I think we occasionally would need to recalculate the cumsum and live with inaccurate numbers for a while

def task_to_processing(ts, ws):
    ix = bisect.bisect(ws._processing_deque, ts.priority)
    max_clusters = 4
    cluster = min(int(ix / len(ws._processing_deque) * max_clusters), max_clusters - 1)
    ws._processing_deque.insert(ix, ts.priority)
    for cluster_to_update in range(cluster, max_clusters):
        ws._occupancies[cluster_to_update] += ts.duration

Mathematically speaking this update is wrong, I believe, but if we occasionally recalculate the cumsum entirely we might be just getting away with it. Maybe even remove the update entirely and just periodically recalc the cumsum

[mrocklin]

So given this complexity I started thinking "Gosh, it probably makes sense to just go random here instead" so I did a tiny micro-benchmark

In [1]: import random

In [2]: L = list(range(1000))

In [3]: %timeit random.choices(L, k=5)
1.34 µs ± 63.4 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

Great, one microsecond seems doable.

But then I realized "No, this happens for every valid worker on this task as we compute the worker objective function, and so this could easily add up to 100us".

Now I'm leaning more towards tracking a distribution of priorities, like mean and standard deviation, assuming that there is a nice way to do that in a streaming population where elements can leave the set.

Thoughts?

[fjetter]

Now I'm leaning more towards tracking a distribution of priorities, like mean and standard deviation, assuming that there is a nice way to do that in a streaming population where elements can leave the set.

I assume you are talking about some abstract "mean of priority" to estimate whether a task has the potential to cut the line? Or are you talking about mean/stddev of expected starttime? I'm currently a bit confused tbh

assuming that there is a nice way to do that in a streaming population where elements can leave the set.

I'm not sure if we can find anything which is significantly simpler than my proposed sorted deque + bisect but I'm glad to be educated about this topic if I'm wrong

[fjetter]

FWIW the mean/quantile of expected starttime is what I called cumsum occupancy above (If that's not clear, I'll need to find a pen and paper :P)

[mrocklin]

I was thinking of tracking the mean/std of the priority of the tasks in the processing state.

Then, we would determine start time by assuming that all processing tasks had equal duration, and looking at how the priority of the new task compared to the mean/std, and taking an appropriate fraction of the total occupancy.

I'm pretty comfortable making pretty broad assumptions here. I think that in most cases where this makes a difference in scheduling we'll find that the priority is either at the beginning or at the end of the distribution. Exact placement within the distribution probably isn't useful. My hope is that by making simplifying assumptions like this that we'll be able to simplify and accelerate things considerably.

[fjetter]

Then, we would determine start time by assuming that all processing tasks had equal duration, and looking at how the priority of the new task compared to the mean/std, and taking an appropriate fraction of the total occupancy.

Got it. Still, that requires us to define a mean of a tuple of things (tie breakers are not even integers, are they?). Is it possible for us to ignore some/most of the priority hierarchies? Calculating a mean (even with removal) for integers is trivial; I believe it is also possible for stddev but I'll need to do a bit more research on this. Maybe we can calculate a mean + stddev for every hierarchy and ignore all non-integer levels?

If we need to take the full priority tuple, I'll stick to my bisection proposal to calculate quantiles instead.

[mrocklin]

Most of the priority stack is non-informative for most comparisons. A mean/std per level would probably be fine. It'll get a little interesting though to do comparisons.

[mrocklin]

Interestingly, this problem could also go away if we got very aggressive with speculative task assignment, and if we tried to assign tasks in a depth first manner. The order in which we would assign tasks to workers would more closely match their current priority, and so the current assumption that all tasks on a worker are higher priority would more likely be true.

This fails the first time we're unable to speculatively assign something though (which, in a non-trivial situation will probably happen quickly).

I don't currently think we should go down the path of STA as a cure for all ailments, but I thought I'd bring it up here for completeness.

[gjoseph92]

Just thinking about a couple properties we might be able to take advantage of:

1. Tasks should be assigned to workers in roughly priority order, right? So just appending priorities to a deque as they get assigned would already be somewhat sorted.
2. Priorities are discrete and unique, so once we have both priority N and N+1 on a worker, we know that no task can ever be assigned which would fall between them in priority. This, combined with point 1, might allow us to efficiently keep a much smaller list of priorities per worker, since we could skip over any runs of consecutive priorities from the perspective of figuring out where a new priority falls within the list.

Basically maybe we could maintain a run-length encoding (kinda) of priorities on each worker? And if the latest task being assigned is 1 greater than the previous priority (common case, for rootish tasks at least, which should make up the vast majority of tasks in ws.processing), then this would be O(1).

Still, that requires us to define a mean of a tuple of things

Obviously, assessing "1 greater than the previous" would require this as well.

[gjoseph92]

In fact I think we can narrow this more. Doesn't this problem only really happen with downstream tasks that need to beat out root-ish tasks? I believe root tasks are the only case where workers can get thousands/large numbers of tasks assigned. Once all root tasks are done, shouldn't len(ws.processing) be O(ws.nthreads), because new tasks can't get assigned without old ones completing a dependency of that task (besides work stealing / poor decide_worker assignment)? And furthermore, once we're in this state, I'd think the cases where a new task is runnable and higher-priority than other tasks that are also runnable on the worker are relatively rare, and possibly not as important from a scheduling perspective.

So perhaps we could track the occupancy from root-ish tasks separately for each worker, and if the task to assign is not root-ish, then subtract that occupancy from the total worker occupancy? In doing so, we're assuming that non-root tasks are higher priority than all root tasks, which is approximately true, though false for those tasks' dependencies (though acting like it's true from the perspective of worker assignment still feels correct).

I think there's also a danger in a perfect "Respect priority in earliest start time heuristic", which is that it would tend to assign downstream tasks to the wrong workers. Consider the graph

2 4    6 8
| |    | |
1 3    5 7

 A      B
^ ideal worker

with 2 workers, each with 2 threads. 1,3 go to A, 5,7 go to B. Worker B completes all its tasks first, and gets 6 and 8 to run. Then worker A finishes, and the scheduler has to assign 1 and 2. It gives 1 to worker A. But it gives 2 to worker B, since it thinks 2 can cut in line in front of 6 and 8, which looks faster than waiting for 2 to finish. (This would happen if the estimated runtime of the downstream taskgroup was higher than the transfer time of the root task.)

Whereas if we just take the "ignore all root-ish tasks when estimating downstream task start time" approach, it basically levels the playing field across all workers, and the data-transfer metric should take over (I think).

[mrocklin]

So perhaps we could track the occupancy from root-ish tasks separately for each worker, and if the task to assign is not root-ish, then subtract that occupancy from the total worker occupancy?

Hrm, that does indeed seem simple and practical. I have a slight hesitation to encode rootish-ness more deeply throughout the codebase but 🤷‍♂️ it worked pretty well last time :) If we can't think of general purpose cases where this is important then I like this idea more than anything that has been proposed before.

[mrocklin]

As a lesser alternative I'll simplify my mean/std thought and suggest just min/max, and assume uniform distribution between those two, which might be simpler. I'm not proposing this above rootish handling. I'm just mentioning this here for completeness.