Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix safe chunks validation #9513

Closed

Conversation

josephnowak
Copy link
Contributor

@josephnowak josephnowak commented Sep 18, 2024

This MR fixes multiple cases on the logic of the safe_chunks option, for more info please see the test added, which includes multiple scenarios.

I know that some of the issues are already closed, but if you see the comments none of them was successfully closed, but with this fix a proper error is going to be raised.

@josephnowak
Copy link
Contributor Author

Hi @shoyer, a long time ago I mentioned to you that it would be good to raise an error on certain scenarios where the dask and zarr chunks are not aligned, I was finally able to dedicate some time to this fix, and I think it would be useful to avoid many issues in the future related to incorrect writes on Zarr.

@@ -5496,24 +5496,26 @@ def test_encode_zarr_attr_value() -> None:

@requires_zarr
def test_extract_zarr_variable_encoding() -> None:
# The region is not useful in these cases, but I still think that it must be mandatory
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the region kwarg is required now, should we make it mandatory? Or at least add a TODO in the function if we want to push that off to another PR?

Copy link
Contributor Author

@josephnowak josephnowak Sep 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that ideally, the validation should be outside the extraction of the encoding, and that would avoid the use of the region inside the extract_zarr_variable_encoding function, I added a comment with that option in the zarr.py file but, I tried to reduce the amount of modifications to the code so I decided to not apply it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If needed I can add a TODO indicating that it would be good to apply that modification.

start = 0
if interval.start:
# If the start of the interval is not None or 0, it means that the data
# is being appended or updated, and in both cases it is mandatory that
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are we sure we don't need to distinguish between append mode vs. others? If we have an array of (5,5,5) and we write to location (4,10) in chunks of [(4,5),(6,10)], then that's only OK if we're appending — otherwise we could be writing (1,4) in another process?

(To the extent the proposed code dominates the existing code — i.e. the existing code still has this problem — I would still vote to merge even if we can't solve that on this iteration, albeit with a TODO)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you specify that case using Xarray? I'm not able to understand correctly the chunk sizes that you are indicating, but I think that it is not necessary to distinguish between append and others because Xarray internally represents all as a region write, and for the update region it is always going to apply the logic of check if the first chunk is aligned with the border_size + N * zarr_chunk_size, for any integer N >= 0 and border_size being the amount of data between interval.start and the corresponding chunk end.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Take the example above:

image

This seems safe or unsafe depending on whether we're appending:

  • If we're appending (mode="a"), then this is safe, because we can only append in a single process, so can only be writing the blue box
  • But if we're doing a region write without append (mode="r+"), then we can be writing to the yellow and blue outlined box in two separate processes, which would cause corruption. So when mode="r+", writing half a chunk isn't safe.

(tbc, you seem to understand this overall better than I do, so I suspect I'm wrong, but also that this is important enough that it's worth clarifying)

Copy link
Contributor Author

@josephnowak josephnowak Sep 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that case is already covered by the validation, if the first chunk of the array (the blue box) is not of the size of the region that needs to be updated it is also going to raise an error.

For example, if you see the test on the image (test_zarr_region_chunk_partial_offset) the first chunk of the array is going to be written from position 5 until position 15, which means that it is going to touch only half of the first and second Zarr chunk, this would be completely valid if there is no other chunk in the array, but as there is another one it means that the second Zarr chunk and, only the second one is going to suffer the issue of more than one Dask chunk being written to the same Zarr chunk in parallel.
image

border_size = zchunk - interval.start % zchunk
if dchunks[0] % zchunk != border_size:

the translation of that condition is that the first chunk of Dask must be always greater or equal to the portion of the data that needs to be updated on the first Zarr chunk (the border size), and if it is greater then it must be exactly aligned with the other Zarr chunks unless there is no other Dask chunk

Copy link
Contributor Author

@josephnowak josephnowak Sep 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I hope these examples help to better understand my explanation.

import xarray as xr

path = "C:/Users/josep/Desktop/test_zarr3"
arr = xr.DataArray(list(range(10)), dims=["a"], coords={"a": list(range(10))}, name="foo").chunk(a=3)
arr.to_zarr(path, mode="w")

try:
    arr.isel(a=slice(0, 3)).chunk(a=(2, 1)).to_zarr(path, region="auto")
except ValueError:
    print("Chunk error, the first chunk is not completely covering the border size")

# This is valid the border size is covered by the first chunk of Dask
arr.isel(a=slice(1, 4)).chunk(a=(2, 1)).to_zarr(path, region="auto")


try:
    arr.isel(a=slice(1, 5)).chunk(a=(3, 1)).to_zarr(path, region="auto")
except ValueError:
    print("Chunk error, the first chunk is covering the border size but it is not completely covering the second chunk")

# This is valid because there is a single dask chunk, so it is not possible to write multiple chunks in parallel 
arr.isel(a=slice(1, 5)).chunk(a=(4, )).to_zarr(path, region="auto")

try:
    arr.isel(a=slice(0, 5)).chunk(a=(3, 1, 1)).to_zarr(path, region="auto")
except ValueError:
    print("Chunk error, the first chunk is correct but the other two chunks are overlapping the same Zarr chunk")

# This is the simplest case, the first chunk is completely covering the first Zarr chunk
arr.isel(a=slice(0, 5)).chunk(a=(3, 2)).to_zarr(path, region="auto")


# The first chunk is covering the border size (2 elements) and also the second chunk (3 elements), so it is valid
arr.isel(a=slice(1, 8)).chunk(a=(5, 2)).to_zarr(path, region="auto")

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh ok, you are right it is hard to detect that alignment problem, and if that would be the meaning of the mode="a" it has sense, just that right now if I try to write on a region using that mode it raises the following error:

arr.isel(a=slice(1, 8)).chunk(a=(5, 2)).to_zarr(path, region="auto", mode="a")

ValueError: ``mode`` must be 'r+' when using ``region='auto'``, got 'a'

Do you think that we should allow to write on region using the "a" mode? if yes then I could proceed to change the validation code so it satisfy the condition that you specified before (partial writes for "a" and non-partial writes for "r+")

Copy link
Collaborator

@max-sixty max-sixty Sep 19, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that would be a nice addition, yes! Though would need to refresh my memory to be confident, and others may have views.

I would encourage landing this PR with the improvements, and then another PR with that change, unless that's difficult!

Thank you!

Copy link
Contributor Author

@josephnowak josephnowak Sep 19, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds fine, so would this PR be ready then? and just to clarify would you like that in the other PR I add the following features?

  1. If the mode is "r+" then only allows full chunks.
  2. If the mode is "a" allows writing using region.
  3. If the mode is "a" allows partial writes.
  4. Probably it would be good to also update the docs of the to_zarr method, honestly they do not indicate any of those behaviors.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds fine, so would this PR be ready then?

I added one comment above, otherwise I think so!

Copy link
Collaborator

@max-sixty max-sixty Sep 19, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1. If the mode is "r+" then only allows full chunks.

2. If the mode is "a" allows writing using region.

r+ can write using region — that's the default — as long as they are full chunks. (I think the 1&3 cover things, bullet 2 isn't needed)

3. If the mode is "a" allows partial writes.

4. Probably it would be good to also update the docs of the to_zarr method, honestly they do not indicate any of those behaviors.

That would be really good, we've fallen behind a bit on docs. Possibly an LLM can help if that makes improving the docs a bit more appealing :)

Copy link
Collaborator

@max-sixty max-sixty left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very nice! Thank you v much @josephnowak . Added a couple of questions and comments

@TomNicholas TomNicholas added the topic-zarr Related to zarr storage library label Sep 18, 2024
@josephnowak
Copy link
Contributor Author

By the way, probably it would be good to create a utility function on Xarray (or a parameter on the to_zarr method) that automatically aligns the Dask chunks of a dataset with the Zarr chunks, this could be inefficient in some scenarios but it can be quite useful to avoid dealing with the alignment of chunks logic or using a synchronizer.

for zchunk, dchunks, interval in zip(
enc_chunks_tuple, var_chunks, region, strict=True
):
if not safe_chunks or len(dchunks) <= 1:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re "so would this PR be ready then" — am I right in thinking the proposed code changes the behavior such that a single partial chunk can be written on mode="r+"? I don't think we want that yet (though am not confident, it seems unlikely to be a big problem, but still possibly confusing)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, the code is going to allow writing on a single partial chunk on mode="r+".
From my perspective and based on the docs of the to_zarr method, the new behavior would not be confusing, "“r+” means modify existing array values only (raise an error if any metadata or shapes would change)", and writing in partial chunk means modify the existing array values.

Copy link
Collaborator

@max-sixty max-sixty Sep 19, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But it breaks the rule we discussed above, re it being safe writing to different locations from distributed processes.

It's unlikely, because it's rare to write to a single partial chunk from distributed processes. But is it worth the complication, relative to forcing people to pass mode="a"?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@max-sixty I sent the following PR #9527 which contains all the changes requested, if you want I can add everything here or close this one, whatever you prefer.

@josephnowak josephnowak mentioned this pull request Sep 20, 2024
2 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
topic-zarr Related to zarr storage library
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants