Repository navigation
Asynchronous Disk Access in Workers #4424
Description
Currently reading from or writing to disk blocks the event loop
distributed/distributed/worker.py
Lines 1948 to 1964 in 74e8dc6
| def put_key_in_memory(self, ts, value, transition=True): | |
| if ts.key in self.data: | |
| ts.state = "memory" | |
| return | |
| if ts.key in self.actors: | |
| self.actors[ts.key] = value | |
| else: | |
| start = time() | |
| self.data[ts.key] = value | |
| ts.state = "memory" | |
| stop = time() | |
| if stop - start > 0.020: | |
| ts.startstops.append( | |
| {"action": "disk-write", "start": start, "stop": stop} | |
| ) |
This can cause workers to become unresponsive, especially in systems with very slow disk access. Ideally Disk I/O would happen concurrently. There are a couple of ways to do this.
Offload to separate thread
We could move all manipulation of the Worker.data MutableMapping to a separate thread, such as we do with the offload function, which we use today for deserialization.
However, if we do this then we need to do it for all access to Worker.data including seemingly innocuous checks like if key in self.data which may become annoying.
Handle Disk logic directly in the worker
We could also break apart the MutableMapping abstraction, and unpack the zict logic directly into the Worker code. This would allow us to keep a lot of the fast access in the event loop, while treating disk access specially. It would also open the door for more performance improvements, like trying to schedule tasks for data that is currently in memory rather than data that is currently on disk. In general if we want to improve out-of-memory handling in Dask we'll eventually need to break this abstraction.
However, breaking this abstraction comes at considerable cost. First, it means that there is more to manage in a monolithic Worker codebase (zict has tricky logic that we haven't really had to touch or maintain in years). Second, it means that we'll have to find a way that still lets other groups like RAPIDS extend the storage hierarchy (they have device->host->disk rather than just host->disk).
Activity
fjetter commentedon Jan 13, 2021
Is it a possibility to use a bigger/different interface than
MutableMapping? I'm not too worried about "not using zict" any longer but I am concerned about the monolithic worker problem you mentioned and I think we should keep some level of abstraction for this, let alone for easier testing once the logic continues to grow with the out-of-memory topic you mentioned.What I would also hope to address is the potential data duplication caused by spilling #3756 if we are reworking this area. I never followed up on that topic but fixing it would also require breaking up that abstraction to some extend.
mrocklin commentedon Jan 13, 2021
quasiben commentedon Jan 14, 2021
cc @madsbk who has also recently been doing work around spilling
gjoseph92 commentedon Apr 30, 2021
I'm curious about the ripple effects this could be having on performance. See ocean-transport/coiled_collaboration#9, where the same workload is running with and without
dask.config.set({"distributed.worker.memory.spill": None}),doing much better with spilling turned off. Copying the links for convenience:EDIT: the spill config wasn't actually being applied to workers: ocean-transport/coiled_collaboration#9 (comment). So whatever's causing the difference is unrelated. However, I think the observations of blocking IO might still be relevant?
Spilling offThe interesting thing to me is that with spilling on, the task stream has lots of blank space in it—workers are "idle" (not doing user work) a lot of the time. With spilling off, work is fully packed, and there's no idle time. I think ideally, spilling should not cause more idle time in the task stream. The red transfer times might be longer, and we'd see more yellow
disk-read-tasks, but we'd hope other useful work could happen concurrently, instead of blank space.I notice a pattern looking at the task stream in the spilling case. You can kind of see white/"idle" patches on one worker that match up with data-transfer waits on another worker. This might just be a Rorschach interpretation. But it seems to indicate that the sender of the data is blocked from running other tasks until the data it's sending has been un-spilled from disk. (This makes sense, since if the event loop is blocked, how could new tasks be submitted to run?)
(@jrbourbeau: here's another case where the "inverted task stream" / "what is my working doing when it's not working" visualization we were talking about could be useful)
I'm skeptical that disk IO could be taking 40-60sec in some of these cases (getting that number from the corresponding transfer durations). That seems really slow. But, in the administrative worker profile, we can see that ~1200 seconds are spent in
f.read()fromzict/file.py, called fromget_data. 120sec per worker is too low to add up to the total transfer time we see, but is still a lot. It would be helpful to know the distribution of thesef.read()times.Basically it seems that transferring spilled data locks up the sender of the data for part of the time. Blocking the worker IO loop to me seems consequential. Based on my incomplete understanding of the system—please correct what I'm getting wrong here—I think this could be the fallout (assuming IO is blocking for ~seconds at a time):
Semi-related thing I noticed: in
get_data, we're deserializing data stored on disk, then immediately re-serializing it to send over the comm. It would be cool if, in theget_datacase, we could just wrap the raw bytes from disk inSerialized. Even cooler if we could avoid slurping all the data from disk into memory first, and instead stream it. That would also allow us to remove the worker paused throttle, since transferring spilled data then wouldn't increase the memory usage of a worker (beyond some small buffer size). I wouldn't be surprised if that throttle is a bottleneck as well when workers come under memory pressure.fjetter commentedon Apr 30, 2021
@gjoseph92
You might want to look into the crick statistics for
get-data-load-duration, seedistributed/distributed/worker.py
Lines 1366 to 1380 in 233ec88
That's the time spent loading the data before the sender submits it. That's 100% disk IO.
Similarly, there is the
get-data-send-durationdistributed/distributed/worker.py
Lines 1383 to 1397 in 233ec88
That should at least help understanding where the time is spent on the sender. In particular if no task was queued for execution before this is initiated, the worker will idle until this is over, that's correct.
fjetter commentedon Apr 30, 2021
I've thought about the suggest of @gjoseph92 w.r.t a streaming like interface and I believe this aligns pretty well with the target of avoiding data duplication. I tried to put my idea about how this interface could look like into some code and would be glad to receive feedback. To be entirely honest, this proposal would not allow for direct from-disk-streaming but I think this could be easily added
https://gist.github.com/fjetter/25fae963c70c9b756b591213244af96a
This proposal keeps the mutable mapping and still allows for fast
is inor.keys()but introduces an objectDataProxywhich would act as a file-like object and would only allow data access via open/close which allows us to track if it is currently in use or not. If it is not in use, it would allow for spilling the data. Instead of referring to spilling directly, I figured we might want to have multiple layers, e.g.in-memory,in-memory-compressedanddisk/remote(I remember the GPU folks having a concept of multi layered memory?) and would suggest to implement this transition in another class I called "Processor" which would allow for pluggable storage. This draft is not functional but there is enough meat to deliver the idea, I hope.The interface itself would be write async and whether or not the
Dataimplements actual async access is then an implementation detailjrbourbeau commentedon Apr 30, 2021
cc'ing @madsbk who has worked on spilling over in
dask-cudaand has experience with using data proxy objects (though to be clear, I've not yet looked deeply into how @madsbk and @fjetter proposals may overlap)madsbk commentedon May 3, 2021
@gjoseph92 I think you are right in your observations. I made some similar ones, which turned in to the JIT-unspilling approach for GPU -> CPU memory spilling: rapidsai/dask-cuda#451. Notice, this approach delegates the de-serialization to the worker-thread executing the task. Also it handles the de-serializing & re-serializing issue in
get_data().I have summarized the shortcomings of the current spilling here: #4568 (comment), which are all fixed in Dask-CUDA.
Our plan is to generalize the approach to include CPU -> Disk spilling.
... and a small commercial, I will be talking about this at the Dask Distributed Summit 2021 :)
madsbk commentedon May 3, 2021
I think there is a significant overlap. If I am not mistaken, the interface @fjetter propose is very similar to the
ProxifyHostFilein Dask-CUDA. The main difference being that Dask-CUDA leaks the proxy objects into the executing task. However, there is no reason why a worker couldn't un-proxify the input arguments before executing tasks.fjetter commentedon May 3, 2021
There are at least similar ideas. The execution is quite different but it is also not a fair comparison since my code is not actually working :)
Considering that the goal of this issue is to allow for async data access, this would require us to change the interface for the worker anyhow and I would consider un-proxification a hard requirement for this feature since generic user functions should not be required to deal with our internal abstraction.
I could not 100% follow the implementation of the cuda proxies but it looks reasonably unspecific to GPUs which is good. From what I understand, most of the complexity is introduced since the implementation deduplicates by
id(at least this is what I understand from theProxiesTally) instead of merely bykeyand I'm wondering what the reason for this decision is.I assume that I would need to put in locks for multithreading as well and the two approaches, assuming un-proxification, would probably end up looking fairly similar indeed.
I'm wondering how much of the "how do we want to spill" discussion should/must leak into the "what is a good asyncio interface" discussion.
mrocklin commentedon May 3, 2021
I will put forward three options here for asyncio interfaces. I'm not recommending these interfaces in particular; instead my objective is to highlight a choice we have.
AsyncMutableMapping
We could make an async extension of the MutableMapping Interface with
async def getitemandasync def setitemmethods. This would be a relatively small change, but would stop the stalling that we currently see in workersHowever, it still keeps the keys in memory / disk / gpu memory opaque. This is good because it's a separate concern and extensible, but also somewhat bad because we can't use this informaiton for scheduling decisions or diagnostics.
No abstraction at all
Alternativley we could just fully bake in ideas for memory / disk / gpu memory into Dask Workers. After several years there have only been two serious implementations of this abstraction, the default memory/disk layout and dask-cuda's DeviceHostDisk layout. Let's just absorb these two options into the
Workerand be done with it.This closes the door for easy addition of future layers to the memory hierarchy (object stores? compressed-in-memory? other accelerators?) but opens the door for smarter scheduling. In particular there are often workloads where we should be computing on data that we have close at hand in main memory first, but instead we read in data from disk, and compute on that data while shoving the easy-to-access in-memory data to disk, only to have to read it later.
Storage Layer abstraction (middle ground)
Maybe in the middle, we create an abstraction for how to asynchronously get and set data into some storage system, and the Worker has a stack of these things with associated size limitations.
The stack-of-layers/zict.Buffer/zict.LRU logic is hard-coded into the Worker while the logic of each storage-layer is abstracted into a separate class.
mrocklin commentedon May 3, 2021
Hrm, or perhaps that's too simple. I now realize that the
DeviceHostDiskabstraction isn't just a stack of storage layers where one overflows to the other, but that there is some more complex logic about what to put where.gjoseph92 commentedon May 3, 2021
As another option (or an extension of the Storage Layer abstraction), could we define an interface for storage layers to give us enough information to make such smarter scheduling decisions? What are the questions this would need to answer? It sounds like we want to know "which keys are fast to load". Maybe also "how much memory will this key take if I load it", "give me the ranking of (these keys|all keys) in 'fastness'"?
mrocklin commentedon May 4, 2021
Perhaps it would make sense to just start with adding async methods to the MutableMappings in zict so that we can avoid blocking the event loop. This is probably an easy(ish) change and would relieve short term pressure.
mrocklin commentedon May 4, 2021
It would also avoid us having to come to some grand design here
23 remaining items
crusaderky commentedon Mar 30, 2023
Delivery timeline
This adds a new public API for SpillBuffer alternatives (e.g. Dask-CUDA),
AsyncBufferProto, which however is not going to be final until the next step. This interface is in addition to the legacyManualEvictProto.AsyncBufferProto.Officially document the new
AsyncBufferProto.AsyncBufferProtoAPIzict.AsyncBufferclass in order to increase code reuseManualEvictPrototoAsyncBufferProtoManualEvictProto. This will remove a lot of dead code from dask/distributed. As I don't really think anybody besides Dask-CUDA overrides the SpillBuffer with a spillable interface, I think it's safe to do so as soon as the new Dask-CUDA release is out, without waiting for several months?New public API
dask/distributed will indefinitely continue to support passing a generic
MutableMapping[str, object]instead of the SpillBuffer when you initialise the Worker.I would like to deprecate
ManualEvictProto(which is an extension toMutableMapping[str, object]) and instead create a new APIAsyncBufferProto, which crucially is not a MutableMapping.This is the new API, downstream of #7686 + #7421:
CC @madsbk is this feasible to implement on your side?
madsbk commentedon Mar 30, 2023
Yes, it looks feasible. For dask-cuda's JIT-unspill, this is almost trivial. Dask-cuda's default spilling requires a bit more work but definitely feasible.
crusaderky commentedon Apr 3, 2023
Status update:
dask/zict#92 is currently undergoing peer review.
I've run a first batch of performance tests.
The good news:
The change works exactly as intended. The event loop is not blocked anymore while spilling/unspilling is in effect.
If a task does not depend on spilled-out data, it will be completely unaffected by any ongoing spilling/unspilling activity.
The looks-bad-but-it's-actually-good news:
In all use cases where there was spilling and unspilling caused by a lot of independent tasks, this change will cause workers to reach the paused state a lot faster - effectively resulting in higher memory consumption. The reason is that, in main, once you pass the
targetthreshold you have effectively a mini-pause at the end of every task, until as much data as the task's output has been spilled out from the bottom of the LRU. With this change, there's nothing slowing tasks until you hit the pause threshold - which means that memory will actually get used.We may need to tweak the pause threshold down a bit - maybe 5% -or people will start experiencing OOM restarts whereas before they weren't, simply because their heap (across all threads) is more than 15% total memory.
The bad news:
This change in behaviour uncovered a very, very ugly bug, which completely destroys performance: #7742.
I've worked around it by introducing an artificial delay in Worker.execute - effectively, an undeclared mini-pause when you start a new task if there's spilling activity going on - which degrades performance but produces a memory profile more similar to main.
The second piece of bad news is that it turns out that the current coiled-runtime test bench is not very impacted by a choked event loop. The two tests in test_spill are not:
test_spill.py::test_spillinghas trivial tasks and does not perform any meaningful amounts of network transferstest_spill.py::test_dot_product_spillhas very heavy tasks, that are long to execute and to transfer, which means that while a lot of time is spent with a blocked event loop, it doesn't impact runtime much:There's only one test, which features fast-running and a lot of memory generation:
test_join.py::test_set_index. I'm still running A/B tests, but preliminary results show how the test now runs in half the time compared to main.test_spilling [uncompressible]
[1] hack in the test code, as shown in the last video in #7742, which waits for workers to unpause
[2] hack in the worker code, as described above, which prevents more tasks from running if the're current spill activity
[3] these are 10.5 seconds over 10 threads, so 1.05 wall seconds
[4] You can see how
otherdrops to nothing - that's time with a blocked event loop![5] This includes the hack slowdown
[6] The increase is caused by workers getting paused
The test has the same end-to-end runtime in the three use cases; this is unsurprising as it's dominated by raw disk throughput.
test_dot_product_spill [uncompressible]
Again, note how other has dropped to zero; there is also a substantial drop in unknown time, which is where workers are waiting on network transfers from other workers or the scheduler.
The zict-offload time shows how you frequently have a task that needs to unspill its dependencies, while either another task is doing the same or older tasks are being spilled.
test_set_index
I've failed to reproduce this test in a jupyter notebook so far - which I need to extract the metrics.
As a side note, I've found the metrics above exceptionally useful, and I'd definitely want to have them saved in the coiled-runtime CI artifacts.
fjetter commentedon Apr 3, 2023
How come? Did you fail to reproduce the runtime difference or are you not able to run the test?
fjetter commentedon Apr 3, 2023
For these tests I'd be interested to see the raw disk IO rates. Did the AsyncBuffer achieve higher peak rates or how does this look like? If you could share the grafana link, that'd be helpful (the benchmark stuff should be publicly available in grafana)
fjetter commentedon Apr 3, 2023
I assume this refers to
set_indexw/ tasks?crusaderky commentedon Apr 3, 2023
I've copied the test verbatim (including the call to dask.config.set to choose the shuffle engine) but it completes in ~20 seconds, without ever getting even close to spilling, instead of ~200s with heavy spilling. Could not figure out what I'm doing wrong.
Yes
crusaderky commentedon Apr 3, 2023
They're roughly the same. In the new version, disk write and task execution are pipelined, whereas before they were interleaved; however in
test_spillingspecifically task execution takes a very modest amount of time so you can't tell the impact of the difference.fjetter commentedon Apr 4, 2023
I could only confirm a runtime decrease of about 33%. Still nice but a little less impressive than what your measurements are showing. I can also only confirm this speedup if we are truly under very heavy memory pressure in the case of memory_multiplier=1 and persist=True, i.e. this is a very unhealthy place for the cluster to be in. Without persisting data I can even see the opposite effect that the async spilling implementation performs worse (haven't collected the data)
crusaderky commentedon Apr 6, 2023
This issue has been postponed indefinitely following disappointing A/B test results.