[v3] fix: zarr v2 compatibility fixes for Dask #2186

jhamman · 2024-09-14T22:50:33Z

port normalize_chunks from v2
add array.store property
default to append in create
rename zarr.store to zarr.storage

closes #1953
xref: dask/dask#11388

TODO:

Add unit tests and/or doctests in docstrings
Add docstrings and API docs for any new/modified user-facing classes and functions
New/modified features documented in docs/tutorial.rst
Changes documented in docs/release.rst
GitHub Actions have all passed
Test coverage is 100% (Codecov passes)

- port normalize_chunks from v2 - add array.store property - default to append in create

src/zarr/api/asynchronous.py

src/zarr/core/array.py

jhamman · 2024-09-14T22:59:22Z

src/zarr/core/chunk_grids.py

+    if all(isinstance(c, (tuple | list)) for c in chunks):
+        # take first chunk size for each dimension
+        chunks = (
+            c[0] for c in chunks
+        )  # TODO: check/error/warn for irregular chunks (e.g. if c[0] != c[1:-1])


except for this block, the rest of this function came from 2.x

will-moore · 2024-09-16T16:34:08Z

Hi @jhamman - thanks for looking into this!
I tried to remember where I got to when I was looking into this before...

Here is what we're doing to load remote data with dask and Zarr v2...

# dask==2023.5.0
# dask-image==2022.9.0
# zarr==2.18.2
import dask.array as da
from zarr import Array, Group, open
from zarr.storage import FSStore
import json

url = "https://uk1s3.embassy.ebi.ac.uk/idr/zarr/v0.4/idr0062A/6001240.zarr"
store = FSStore(url)
zattrs = json.loads(store.get(".zattrs"))
array_path = zattrs["multiscales"][0]["datasets"][0]["path"]
array_data = da.from_zarr(store, array_path)
print(array_data)
# dask.array<from-zarr, shape=(2, 236, 275, 271), dtype=uint16, chunksize=(1, 1, 275, 271), chunktype=numpy.ndarray>

With zarr v3 I'm not getting as far as opening with dask as I'm failing to access any remote data, using a couple of different approaches here (using zarr from this branch), both failing with exceptions noted below...

import zarr
from zarr.store import RemoteStore
import s3fs

url = "https://uk1s3.embassy.ebi.ac.uk/idr/share/ome2024-ngff-challenge/0.0.5/6001240.zarr"
store = RemoteStore(url)
# RuntimeError: Timeout context manager should be used inside a task
g = zarr.group(store = store)


s3 = s3fs.S3FileSystem(anon=True, client_kwargs={'endpoint_url': 'https://uk1s3.embassy.ebi.ac.uk'})
store = s3fs.S3Map(root='idr/share/ome2024-ngff-challenge/0.0.5/6001240.zarr', s3=s3, check=False)
# TypeError - only LocalStore or MemoryStore is supported
g = zarr.group(store = store)

Is there some other way that I can try to load remote data?

Many thanks,
Will

also fix failing ci

TomAugspurger · 2024-09-17T15:03:04Z

Thanks for trying this out @will-moore. Can you share the full traceback?

In the meantime, I would expect something like

 g = zarr.open_group(store=zarr.store.RemoteStore("s3://idr", endpoint_url="http://uk1s3.embassy.ebi.ac.uk", anon=True), path="zarr/v0.4/idr0062A/6001240.zarr")

Can you let us know if that works? It doesn't for me, but I'm not sure if some network / authentication setup is required.

…nto fix/dask-compat

tests/v3/test_codecs/test_codecs.py

src/zarr/api/asynchronous.py

will-moore · 2024-09-17T15:57:31Z

That sample you tried there is zarr v2 data, and it also seemed to need https rather than http!
So this works for me on another machine, with zarr==3.0.0a0, reading both zarr v2 and v3 data:

store=zarr.store.RemoteStore("s3://idr", endpoint_url="https://uk1s3.embassy.ebi.ac.uk", anon=True)
g = zarr.open_group(store=store, path="share/ome2024-ngff-challenge/idr0048/9846151.zarr")
print('zarr v3 group', g)

g2 = zarr.open_group(store=store, path="zarr/v0.4/idr0062A/6001240.zarr", zarr_version=2)
print('zarr v2', g2)

gives me:

zarr v3 group Group(_async_group=<AsyncGroup Remote fsspec store: S3FileSystem , idr/share/ome2024-ngff-challenge/idr0048/9846151.zarr>)
zarr v2 Group(_async_group=<AsyncGroup Remote fsspec store: S3FileSystem , idr/zarr/v0.4/idr0062A/6001240.zarr>)

On my machine, I unfortunately get:

this error

    g = zarr.open_group(store=store, path="share/ome2024-ngff-challenge/idr0048/9846151.zarr")
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/zarr/api/synchronous.py", line 175, in open_group
    sync(
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/zarr/sync.py", line 92, in sync
    raise return_result
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/zarr/sync.py", line 51, in _runner
    return await coro
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/zarr/api/asynchronous.py", line 523, in open_group
    return await AsyncGroup.open(store_path, zarr_format=zarr_format)
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/zarr/group.py", line 150, in open
    zarr_json_bytes = await (store_path / ZARR_JSON).get()
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/zarr/store/core.py", line 35, in get
    return await self.store.get(self.path, prototype=prototype, byte_range=byte_range)
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/zarr/store/remote.py", line 98, in get
    await (
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/s3fs/core.py", line 1128, in _cat_file
    return await _error_wrapper(_call_and_read, retries=self.retries)
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/s3fs/core.py", line 145, in _error_wrapper
    raise err
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/s3fs/core.py", line 113, in _error_wrapper
    return await func(*args, **kwargs)
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/s3fs/core.py", line 1115, in _call_and_read
    resp = await self._call_s3(
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/s3fs/core.py", line 358, in _call_s3
    await self.set_session()
  File "/Users/wmoore/opt/anaconda3/envs/zarr_v3/lib/python3.10/site-packages/s3fs/core.py", line 519, in set_session
    self.session = aiobotocore.session.AioSession(**self.kwargs)
TypeError: AioSession.__init__() got an unexpected keyword argument '//'

Which is an issue that seems to affect just me, as I already reported at ome/ome2024-ngff-challenge#22 and I can't seem to fix on my machine, which is kinda annoying!
Might have to re-install anaconda, or try some other option...

Thanks for your help!

…nto fix/dask-compat

src/zarr/api/asynchronous.py

src/zarr/core/array.py

jhamman · 2024-09-29T22:22:56Z

src/zarr/core/array.py

+    parents = [
+        AsyncGroup(
+            metadata=GroupMetadata(zarr_format=node.metadata.zarr_format),
+            store_path=StorePath(store=node.store_path.store, path=""),
+        )
+    ]


@TomAugspurger - this was needed to enable zarr.create(store={}, path="a") such that the root zarr.json is created in addition to a/zarr.json.

tests/v3/test_codecs/test_codecs.py

src/zarr/core/chunk_grids.py

src/zarr/core/array.py

src/zarr/api/asynchronous.py

src/zarr/core/array.py

…nto fix/dask-compat

d-v-b

good to go i think

TomAugspurger

Looks good, just a couple minor things.

TomAugspurger · 2024-10-01T14:09:09Z

src/zarr/core/array.py

@@ -635,6 +639,8 @@ async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = F
            # To enable zarr.create(store, path="a/b/c"), we need to create all the intermediate groups.
            parents = _build_parents(self)

+            logger.debug("Ensure parents: %s", parents)


Do we want to keep this?

TomAugspurger · 2024-10-01T14:10:05Z

src/zarr/core/array.py

-    parents = []
+    store = node.store_path.store
+    path = node.store_path.path
+    if not path:


Do you know whether this branch is covered by an existing test? If not, it might be good to add one.

Yes, this path is highly exercised.

src/zarr/core/array.py

fix: zarr v2 compatability fixes

fe49f5f

- port normalize_chunks from v2 - add array.store property - default to append in create

jhamman commented Sep 14, 2024

View reviewed changes

jhamman mentioned this pull request Sep 15, 2024

[v3] creating a new array with path should create a root group node if it doesn't exist #2187

Closed

jhamman added 2 commits September 16, 2024 16:47

move zarr.store to zarr.storage

9a1580b

also fix failing ci

make chunks a tuple

0d89912

This was referenced Sep 17, 2024

Doc/v3 migration guide #2102

Open

Deprecate zarr.storage #2196

Closed

Merge branch 'v3' of https://github.com/zarr-developers/zarr-python i…

d78e384

…nto fix/dask-compat

jhamman commented Sep 17, 2024

View reviewed changes

tests/v3/test_codecs/test_codecs.py Outdated Show resolved Hide resolved

src/zarr/api/asynchronous.py Outdated Show resolved Hide resolved

src/zarr/api/asynchronous.py Outdated Show resolved Hide resolved

src/zarr/api/asynchronous.py Outdated Show resolved Hide resolved

Apply suggestions from code review

e534279

jhamman requested a review from d-v-b September 17, 2024 15:46

will-moore mentioned this pull request Sep 17, 2024

TypeError: AioSession.__init__() got an unexpected keyword argument '//' ome/ome2024-ngff-challenge#22

Open

jhamman mentioned this pull request Sep 18, 2024

Zarr Python 3 tracking issue pydata/xarray#9515

Closed

4 tasks

Merge branch 'v3' of https://github.com/zarr-developers/zarr-python i…

dea4a3d

…nto fix/dask-compat

will-moore mentioned this pull request Sep 19, 2024

Zarr v3 support ome/ome-zarr-py#383

Closed

jrbourbeau mentioned this pull request Sep 20, 2024

Zarr-Python 3 compatibility dask/dask#11388

Merged

3 tasks

jhamman added 2 commits September 21, 2024 20:54

Merge branch 'v3' of https://github.com/zarr-developers/zarr-python i…

7800f38

…nto fix/dask-compat

more merge conflict resolution

93b61fc

TomAugspurger mentioned this pull request Sep 25, 2024

Allow mode casting for Stores #2249

Merged

6 tasks

jhamman added 3 commits September 29, 2024 14:02

Merge branch 'v3' of https://github.com/zarr-developers/zarr-python i…

88afe52

…nto fix/dask-compat

fixups

fb6752d

fixup zipstore

0b1dedc

jhamman commented Sep 29, 2024

View reviewed changes

Apply suggestions from code review

322918a

d-v-b requested changes Sep 30, 2024

View reviewed changes

jhamman commented Sep 30, 2024

View reviewed changes

src/zarr/core/array.py Outdated Show resolved Hide resolved

Apply suggestions from code review

a95d54a

jhamman added 6 commits September 30, 2024 13:28

add test

128eb53

Merge branch 'v3' of https://github.com/zarr-developers/zarr-python i…

3c170ef

…nto fix/dask-compat

extend test

54ab9ef

clean up parents

77f2938

debug race condition

2295d76

more debug

5879d67

d-v-b approved these changes Oct 1, 2024

View reviewed changes

jhamman mentioned this pull request Oct 1, 2024

document storage classes and some developer apis #2279

Merged

6 tasks

TomAugspurger approved these changes Oct 1, 2024

View reviewed changes

jhamman commented Oct 1, 2024

View reviewed changes

src/zarr/core/array.py Outdated Show resolved Hide resolved

Update src/zarr/core/array.py

3940d22

jhamman merged commit f3a2e0a into v3 Oct 1, 2024
23 checks passed

dstansby deleted the fix/dask-compat branch October 13, 2024 07:31

dstansby restored the fix/dask-compat branch October 13, 2024 07:31

jhamman deleted the fix/dask-compat branch October 15, 2024 12:30

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[v3] fix: zarr v2 compatibility fixes for Dask #2186

[v3] fix: zarr v2 compatibility fixes for Dask #2186

jhamman commented Sep 14, 2024 •

edited

Loading

jhamman Sep 14, 2024

will-moore commented Sep 16, 2024

TomAugspurger commented Sep 17, 2024

will-moore commented Sep 17, 2024 •

edited

Loading

jhamman Sep 29, 2024

d-v-b left a comment

TomAugspurger left a comment

TomAugspurger Oct 1, 2024

TomAugspurger Oct 1, 2024

jhamman Oct 1, 2024

[v3] fix: zarr v2 compatibility fixes for Dask #2186

[v3] fix: zarr v2 compatibility fixes for Dask #2186

Conversation

jhamman commented Sep 14, 2024 • edited Loading

jhamman Sep 14, 2024

Choose a reason for hiding this comment

will-moore commented Sep 16, 2024

TomAugspurger commented Sep 17, 2024

will-moore commented Sep 17, 2024 • edited Loading

jhamman Sep 29, 2024

Choose a reason for hiding this comment

d-v-b left a comment

Choose a reason for hiding this comment

TomAugspurger left a comment

Choose a reason for hiding this comment

TomAugspurger Oct 1, 2024

Choose a reason for hiding this comment

TomAugspurger Oct 1, 2024

Choose a reason for hiding this comment

jhamman Oct 1, 2024

Choose a reason for hiding this comment

jhamman commented Sep 14, 2024 •

edited

Loading

will-moore commented Sep 17, 2024 •

edited

Loading