Open jacobtomlinson opened 2 years ago
I'm able to reproduce but not when I manually create the cluster and separately connect with client
I think what's happening here is that annotation is initially applied correctly to the dask layers in for mean (@rjzamora showed me a few tricks!). We can see that with the following:
In [42]: arr = da.random.random((10_000, 10_000), chunks=(1000, 1000))
...:
...: # Use the annotated workers for compute
...: with dask.annotate(resources={"FOO": 1}):
...: result = arr.mean()
In [43]: result.dask
Out[43]:
HighLevelGraph with 6 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x10de05e80>
0. random_sample-c79caaeb508f0375f8313ede30300804
1. mean_chunk-eeca2912866909c9659bb48928748826
2. mean_combine-partial-f96108f16300ce4e31953f304bf89e8a
3. mean_combine-partial-74c4e5e4ebe3ca4d290d1ee469004c03
4. mean_combine-partial-28ce6f709a0bde77d0d9e92006055c52
5. mean_agg-aggregate-a395f045817315d8e7bb1b32626d0da6
In [44]: from dask.utils_test import hlg_layer
In [45]: hlg_layer(result.dask, "random_sample").annotations
In [46]: hlg_layer(result.dask, "mean_").annotations
Out[46]: {'resources': {'FOO': 1}}
Note that I've removed the persist
calls. Above the annotations are correctly applied only to the mean_
prefix layers and not the random_sample
layer. When the persist()
calls are added back, we essentially have only one layer
In [50]: arr = da.random.random((10_000, 10_000), chunks=(1000, 1000)).persist()
...:
...: # Use the annotated workers for compute
...: with dask.annotate(resources={"FOO": 1}):
...: result = arr.mean().persist()
In [51]: result.dask
Out[51]:
HighLevelGraph with 1 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x1844bf1f0>
0. mean_agg-aggregate-188659ca8917dd3c92677f087bd12817
With only one layer and dependencies on random_sample
, I think the packing infrastructure will apply the annotation everywhere:
https://github.com/dask/dask/blob/5bd2bbabe4923869ef6d95585d93f36684b5e29c/dask/highlevelgraph.py#L1051-L1067
In the above link, self.get_all_external_keys()
will return a set of keys including: random_sample
and mean_*
and the annotation {FOO: 1}
will be applied globally
@gjoseph92 or @madsbk does this seem correct to you ? I think this is a bug or at least something we want to avoid
I'm surprised the annotation even made it to the scheduler without optimization.fuse.active: False
https://github.com/dask/dask/issues/7036. I'd expect the fusion that's turning a multi-layer HLG into a single materialized mean_agg-aggregate
layer would drop the annotations entirely, not apply them to every task; guess I'm wrong.
Regardless of whatever HLGs are doing, there's clearly a scheduler bug. That KeyError
should never be able to happen, no matter how messed up the annotations are. For release_resources
to fail in this way, I think that means consume_resources
was never called for that worker.
Taking a skim through the work-stealing code, I don't see anywhere in stealing.py
that worker resources get consumed when a task gets sent to them. We check that resources are available, but I don't see anywhere that we actually use them up. @fjetter is this correct? Does work-stealing not properly consume worker resources?
Even so, given that you have only 1 worker with the FOO
resource, how could a task requiring FOO
get stolen onto that worker—there's nowhere else it could run!
Haven't tried running your example yet, just reading code; we'll need to look into this more later.
That is surprising it's not getting removed given https://github.com/dask/dask/issues/7036 -- I had forgotten that was the case.
I don't think this is a work stealing issue as the layers are incorrectly annotated before it gets to the scheduler, right ? You can see the graph being incorrectly manipulated on the client side before getting to the scheduler (I did this with breakpoints in the client and in dask/highlevelgraph.py
Even so, given that you have only 1 worker with the FOO resource, how could a task requiring FOO get stolen onto that worker—there's nowhere else it could run!
I think this is just a result of me condensing the problem into an MRE. In my actual setup I had 2 workers with the annotation, and 10 without it.
Just a note that I was able to reproduce this, and the problem goes away when you move the persist
outside the dask.annotate(resources={'FOO': 1})
block. When persist
is in this block, the annotation is treated as a global annotation (not just a layer-specific annotation).
When
persist
is in this block, the annotation is treated as a global annotation
That's really interesting. Is that intentional, or maybe a HLG bug where during the optimization process, a final materialized HLG layer gets instantiated, which automatically picks up the config.set(annotations=...)
and assigns them to the whole materialized layer?
In my actual setup I had 2 workers with the annotation
This would make much more sense to explain my theory about a work-stealing bug.
I think there are two orthogonal problems here:
distributed.scheduler.release_resources
should never be failing with KeyError, regardless of what annotations dask/dask throws at it. This is a dask/distributed bug. It may happen to be triggered by whatever's going on with HLG annotations, but I believe it's a separate bug. And given @jacobtomlinson's update, I think I can see a way it could be caused by work-stealing.That's really interesting. Is that intentional, or maybe a HLG bug where during the optimization process, a final materialized HLG layer gets instantiated, which automatically picks up the config.set(annotations=...) and assigns them to the whole materialized layer?
This feels like a bug, but it actually seems to be the "expected" behavior (and I don't think it is related to optimization). When you call persist
, the graph eventually makes its way to _graph_to_futures
on the client. As Ben pointed out, the config
annotations (which will currently include resources={"FOO": 1}
) are merged into the global HLG annotations here. These global annotations are then packed with every layer in dask_distributed_pack.
What happened:
Tasks which are currently persisting cause annotated tasks to fail.
What you expected to happen:
Annotated and non-annotated tasks should run concurrently.
Minimal Complete Verifiable Example:
Optional: Create a new environment with latest Dask
Full traceback
```python-traceback distributed.scheduler - INFO - Clear task state distributed.scheduler - INFO - Scheduler at: tcp://10.51.100.43:41979 distributed.scheduler - INFO - dashboard at: :8787 distributed.nanny - INFO - Start Nanny at: 'tcp://10.51.100.43:42819' distributed.worker - INFO - Start worker at: tcp://10.51.100.43:43427 distributed.worker - INFO - Listening to: tcp://10.51.100.43:43427 distributed.worker - INFO - dashboard at: 10.51.100.43:33969 distributed.worker - INFO - Waiting to connect to: tcp://10.51.100.43:41979 distributed.worker - INFO - ------------------------------------------------- distributed.worker - INFO - Threads: 12 distributed.worker - INFO - Memory: 93.04 GiB distributed.worker - INFO - Local Directory: /home/jtomlinson/Scratch/dask-worker-space/worker-4eyrd3zl distributed.worker - INFO - ------------------------------------------------- distributed.scheduler - INFO - Register workerAnything else we need to know?:
If I add a
wait
to the data generation before moving to the annotated section things work. But this means that we have a synchronisation point here that isn't ideal.Environment:
Cluster Dump State: