Refactor zarr data copying to use store.set instead of fsspec operations#194
Conversation
| PROCESSED_VARIABLES_KEY = "processed_variables" | ||
|
|
||
|
|
||
| def _get_fs_and_path( |
There was a problem hiding this comment.
I moved this into here because now it is the only module making use of this function and we intend to deprecate/remove when we move onto refactoring progress tracking.
| zarr.core.sync.sync( | ||
| store.set( | ||
| key, | ||
| zarr.core.buffer.cpu.Buffer.from_bytes(data), |
There was a problem hiding this comment.
If there's any way we can avoid a copy when creating the Buffer that would be ideal. im not sure we can but check out the methods on it
There was a problem hiding this comment.
oh good thought, i'll take a look
There was a problem hiding this comment.
For posterity, we tracked this down and it looks like the CPU buffer will use np.frombuffer https://zarr.readthedocs.io/en/main/_modules/zarr/core/buffer/cpu.html#Buffer.from_bytes
which won't make a copy https://numpy.org/doc/2.1/reference/generated/numpy.frombuffer.html
| source = f"{tmp_store / relative_dir}/" | ||
| dest = f"{path}/{relative_dir}" | ||
| fsspec_apply(fs, "put", source, dest, recursive=True, auto_mkdir=True) | ||
| for file in tmp_store.glob(f"{relative_dir}**/*"): |
There was a problem hiding this comment.
this is likely going from parallel to not parallel. lets keep an eye on if this is impacting perf at all, but i think it should be fine
This PR refactors how we handle copying of data shards and metadata from the tmp store to another storage location. Instead of using
fsspec_applywe can use thesetmethod on our destination store. This also sets us up to be able to interface with anicechunkstore in the same way (where fsspec operations would not work).