Locks and chunked storage

[d-v-b]

dask.array.store takes an optional argument, lock, which (to my understanding) avoids write contention by forcing workers to request access to the write target before writing. But for chunked storage like zarr arrays, write contention happens at the level of individual chunks, not the entire array. So perhaps a lock for chunked writes should have the granularity of the chunk structure of the storage target, thereby allowing the scheduler to control access to individual chunks. Does this make sense?

The context for this question is my attempt to optimize storing multiple enormous dask arrays in zarr containers with a minimum amount of rechunking, which makes the locking mechanism attractive, as long as the lock happens at the right place.

cc @jakirkham in case you have ideas about this.

[mrocklin]

I don't know much about the store function (you're right that @jakirkham is maybe the right person there). But on the lock side you should be able to give the dask.distributed.Lock a unique identifier like the chunk ID and things should work ok. Dask locks are fairly cheap to create, even if you're making many of them.

[gjoseph92]

Is the chunk structure of your dask array the same as the chunk structure of your zarr array? If so, shouldn't each chunk only be written exactly once? Why would you need to prevent multiple tasks writing to the same chunk if there is only ever one task per chunk?

Having the chunk structures aligned generally seems like a good idea to me. But if they're not, what about https://zarr.readthedocs.io/en/stable/tutorial.html#parallel-computing-and-synchronization? I bet it would be only a few lines to write something equivalent to a zarr.sync.ProcessSynchronizer that used distributed.Locks internally.

[d-v-b]

Is the chunk structure of your dask array the same as the chunk structure of your zarr array?

Ideally, yes; typically, no. My dask arrays are really big (TB), so small chunks would be very inefficient for dask. But for storage (e.g., zarr), small chunks are ideal because we want to visualize the data with a chunk-aware visualization tool (neuroglancer). So I aim for a one-to-many relationship between dask chunks and zarr chunks, e.g. 512 x 512 x 512 on the dask side and 64 x 64 x 64 on the zarr side for uint8 data. This is all fine until I do something in dask that makes my chunks shrink (downsampling for making a multiresolution pyramid). If I downsample my 512-sized chunk by half 4 times, I end up with a 32-sized chunk, and that cannot be safely written to my 64-sized zarr chunks. Rechunking is a natural solution, but this moves data between workers and I would prefer to avoid this if possible. So I'm looking at locking as another approach.

Regarding the syncronizers in the zarr-python library, as I understand things the ProcessSynchronizer requires file system locking, which my storage backends (NFS and s3) don't support. So I need to do locking via the client instead. I actually think the distributed.MultiLock will work for me, but I need to test it.

[gjoseph92]

Rechunking is a natural solution, but this moves data between workers and I would prefer to avoid this if possible

You might want to try this again with the latest distributed version now that #4967 has been released. Some of these "dask is moving data around a lot" situations should happen a lot less.

I understand things the ProcessSynchronizer requires file system locking

Sorry, I was suggesting writing your own class with the same interface as ProcessSynchronizer—I'm assuming zarr will play nicely with that?

import distributed


class DistributedSynchronizer:
    """Provides synchronization using locks from Distributed.

    Parameters
    ----------
    prefix: string
        Prefix the names of all dask.Distributed locks with this.
    """
    def __init__(self, prefix: str = ""):
        self.prefix = ""

    def __getitem__(self, item):
        return distributed.Lock(self.prefix + str(item))

>>> synchronizer = DistributedSynchronizer()
>>> z = zarr.open_array('data/example', mode='w', shape=(10000, 10000),
...                     chunks=(1000, 1000), dtype='i4',
...                     synchronizer=synchronizer)

[mrocklin]

Rechunking is a natural solution, but this moves data between workers and I would prefer to avoid this if possible

You might want to try this again with the latest distributed version now that #4967 has been released. Some of these "dask is moving data around a lot" situations should happen a lot less.

Yeah, big +1 on this. I think that the recent memory scheduling things may have had a large positive impact on rechunking.

[mrocklin]

I actually think the distributed.MultiLock will work for me, but I need to test it.

I would have expected a vanilla dask.distributed.Lock per chunk would have done the job.

with dask.distributed.Lock(chunk_id):
    # write chunk 

[d-v-b]

the vanilla lock works in the special case where each worker needs to lock just 1 chunk, but in the more general case when a worker attempts to write a region that spans multiple chunks, that worker must iteratively lock the chunk files spanned by the region, which introduces a potential deadlock due to the iteration, and I suspect the MultiLock addresses this (but I could be wrong).

[mrocklin]

Ah, got it. That makes sense.

[d-v-b]

And I should add that I really like the idea of using the existing synchronizer api in zarr. I will look into implementing that.

[mrocklin]

I do agree with @gjoseph92 though that rechunking seems maybe better. I would be curious to learn more about what is pushing you away from that path.

[gjoseph92]

worker must iteratively lock the chunk files spanned by the region, which introduces a potential deadlock due to the iteration, and I suspect the MultiLock addresses this (but I could be wrong).

I think the code I posted handles these issues, since I believe that's what the Zarr synchronization API is all about. It looks like you just have to tell Zarr how to get the lock instance to use for a particular chunk of the dataset, and it manages the rest.

[jakirkham]

Ideally would probably try to rechunk to match a Dask Array to the chunks that the Zarr Array would have. This effectively causes Dask to figure out how to efficiently compute and write things out. Maybe it's worth discussing why that does not work for your use case?

[d-v-b]

I hope this illustrates why I'm trying to avoid rechunking.

I'm working with task graphs that look like this. A big (500 MB) piece of data enters at the bottom, then gets downscaled and saved recursively. The input chunks need to be big because I have a lot of them. I'm working on datasets that exceed 10 TB, so chunks need to be large to avoid painful task graph creation times.

Each downscaling step reduces the chunk size, which is bad when the dask chunk size dips below my zarr chunk size. Now my workers can't do chunk-aligned writes. If I solve this problem with rechunking, the above task graph transforms to this:

As you can see, rechunking introduces data dependencies between the end leaves of the task graph and breaks the nice parallelism we had before. I suspect this is the cause of some memory struggles my workers are experiencing (and these issues persisted and / or got worse on the latest version of dask).

I suspect that locking zarr chunks lets me keep the first graph. Although locking zarr chunks would break parallelism, it should do so in a way that only increases runtime and not memory use. Hence my interest.

[mrocklin]

You're right that rechunking will introduce transfers which might cause memory issues. Recent versions of dask (2021-07 or later) greatly minimize the effect of these issues

…

On Tue, Jul 20, 2021 at 9:59 AM Davis Bennett ***@***.***> wrote: I hope this illustrates why I'm trying to avoid rechunking. I'm working with task graphs that look like this. A big (500 MB) piece of data enters at the bottom, then gets downscaled and saved recursively. The input chunks need to be big because I have a *lot* of them. I'm working on datasets that exceed 10 TB, so chunks need to be large to avoid painful task graph creation times. [image: image] <https://user-images.githubusercontent.com/3805136/126364246-e2eab385-a9e6-403d-b67a-342004f23e90.png> Each downscaling step reduces the chunk size, which is bad when the dask chunk size dips below my zarr chunk size. Now my workers can't do chunk-aligned writes. If I solve this problem with rechunking, the above task graph transforms to this: [image: image] <https://user-images.githubusercontent.com/3805136/126364465-5298c077-3722-4a2f-975e-c763bbbcdf62.png> As you can see, rechunking introduces data dependencies between the end leaves of the task graph and breaks the nice parallelism we had before. I suspect this is the cause of some memory struggles my workers are experiencing (and these issues persisted and / or got worse on the latest version of dask). I suspect that locking zarr chunks lets me keep the first graph. Although locking zarr chunks would break parallelism, it should do so in a way that only increases runtime and not memory use. Hence my interest. — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#5083 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTEOSG73UKEPYTOI2ETTYWTNVANCNFSM5AQOCA3Q> .

[d-v-b]

I found that my memory issues got worse, not better, after 2021.07.0