Tasks which are obviously root tasks not considered `rootish`

[gjoseph92]

If a task does not have dependencies, it's obviously a root task.

However, is_rootish currently requires that the task also have > 2 * total_nthreads other tasks in its TaskGroup. (The 2x is its own problem: #7273.)

This means some tasks which are clearly root tasks don't go down root-task code paths.

The consequences of this aren't obvious. They may not be particularly noticeable as a user. See #7221 (comment) for discussion of practical implications. My main concern so far is that this could be hiding untested and possibly incorrect behavior. (So there might be practical implications, we just haven't noticed them yet.)

1. Many of our tests don't use root-task code paths, even when queuing is enabled, because they don't submit enough tasks. So whatever situations they're testing for (which presumably could happen with more tasks too) are un-tested with queuing on. At least 12 tests fail with queuing on when this problem is fixed; we don't know how many pass in both cases, but were not using the queuing code path and would not have caught a regression.
2. More codepaths make reasoning harder. Combined with the fact that not all tests may be running, it's harder to avoid incorrect behavior at edge cases. Specifically, the round-robin code path doesn't get much explicit testing, is something we've already wanted to remove Simplify decide_worker #6974, and does not interact well with worker-saturation Round-robin worker selection makes poor choices with worker-saturation > 1.0 #7197 — something CI didn't show until we tried changing the default.

[dcherian]

I think I'm running in to some version of this.

The behaviour changes drastically depending on whether a couple of xarray DataArrays are chunked at read-time. I have a reproducible workflow but it is complicated. I can try to minimize but maybe you can already tell what's wrong by looking at the graphs.

The computation involving these dataarrays is

u_e = (merged.UVEL*dyu*dzu/1.e4)
v_e = (merged.VVEL*dxu*dzu/1.e4)

UVEL, VVEL are chunked in the xr.open_dataset call, dzu always wraps a numpy array, the only difference between the two below is whether dxu and dyu are chunked.

unchunked

looks great, open_dataset tasks were identified as "root tasks" (is that what the diagonal lines mean) and memory usage was great

chunked

not good. Note that open_dataset is not identified as a root task, and memory use is much higher. The difference is open_dataset-DXU and open_dataset-DYU on the right side, i think the quivalents are array-24, array-48 in the previous graph

[mrocklin]

UVEL, VVEL are chunked in the xr.open_dataset call, dzu always wraps a numpy array, the only difference between the two below is whether dxu and dyu are chunked.

This is going to sound dumb, but can you leave them unchunked? 😉. Seriously though, is this you reporting a case that should theoretically be smoother but you're alright, or is this blocking you painfully?

@fjetter do you have anyone who can take a look at this meaningfully? I wouldn't be surprised if there was some special-case here, such as literals being fused in.

I have a reproducible workflow but it is complicated. I can try to minimize but maybe you can already tell what's wrong by looking at the graphs.

Thanks for the offer. Let's see what @fjetter comes back with. I wouldn't be surprised if minimizing ends up being necessary.

[fjetter]

First gut reaction without any meaningful thought: Is it easy and not too expensive for you to try #7531? Over there I played a bit with the rootish classification logic and our benchmarks looked fine. However, the logic I am introducing there is something we were not fully convinced of without more data.
Maybe this helps (or destroys) your workload which would be an interesting data point. Bonus points if we can run this on Coiled ;)

---

Can somebody help me understand what the difference between chunked and unchunked is? Does this mean chunks is explicitly provided in the call to open_dataset, i.e. the rechunking happening behind the scenes is different? Looking at the task screens it looks like the graphs contain the same prefixes and same task counts in both instances but I can't puzzle out the difference and I'm not sufficiently familiar with xarray.

[fjetter]

We can maybe run some diagnostics with the following snippet

def get_tg_info(dask_scheduler):
     return {
         tg.name: {
             'nthreads': dask_scheduler.total_nthreads,
             'len': len(tg),
             'len_dep': len(tg.dependencies),
             'sum_map': sum(map(len, tg.dependencies))
         }
         for tg in dask_scheduler.task_groups.values()
     }
client.run_on_scheduler(get_tg_info)

Can you run this once for the good and the bad example? This extracts the internal info that is being used to classify the tasks by rootish/non-rootish and we can see what exactly is leading to the wrong classification

[mrocklin]

My guess (uneducated) is that chunked xarray datasets use dask arrays while unchunked xarray datasets use numpy.