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

feat/batch creation #2665

Open
wants to merge 54 commits into
base: main
Choose a base branch
from
Open

Conversation

d-v-b
Copy link
Contributor

@d-v-b d-v-b commented Jan 7, 2025

This PR adds a few routines for creating a collection of arrays and groups (i.e., a dict with path-like keys and ArrayMetadata / GroupMetadata values) in storage concurrently.

  • create_hierarchy takes a dict representation of a hierarchy, parses that dict to ensure that there are no implicit groups (creating group metadata documents as needed), then invokes create_nodes and yields the results
  • create_nodes concurrently writes metadata documents to storage, and yields the created AsyncArray / AsyncGroup instances.

I still need to wire up concurrency limits, and test them.

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)

@d-v-b d-v-b requested review from jhamman and dcherian January 7, 2025 13:27
@normanrz normanrz added this to the After 3.0.0 milestone Jan 7, 2025
@dstansby dstansby added the needs release notes Automatically applied to PRs which haven't added release notes label Jan 9, 2025
@github-actions github-actions bot removed the needs release notes Automatically applied to PRs which haven't added release notes label Jan 9, 2025
@d-v-b
Copy link
Contributor Author

d-v-b commented Jan 10, 2025

this is now working, so I would appreciate some feedback on the design.

The basic design is the same as what I outlined earlier in this PR: there are two new functions that take a dict[path, GroupMetadata | ArrayMetadata] like {'a': GroupMetadata(zarr_format=3), 'a/b': ArrayMetadata(...)} and concurrently persist those metadata documents to storage, resulting in a hierarchy on disk that looks like the dict.

approach

basically the same as concurrent group members listing, except we don't need any recursion. I'm scheduling writes and using as_completed to yield Arrays / Groups when they are available.

new functions

  • create_nodes is low-level and doesn't do any checking of its input, so it will happily create invalid hierarchies, e.g. nesting groups inside arrays, or mixing v2 and v3 metadata, and it won't create intermediate groups, either.

  • create_hierarchy is higher level, it parses the input, checking it for invalid hierarchies, and inserting implicit groups as needed.

  • Group.create_hierarchy is a new method on the Group / AsyncGroup classes that takes a hierarchy dict and creates the nodes specified in that dict at locations relative to the path of the group instance. the return value is dict[str, AsyncGroup | AsyncArray], but I guess it also doesn't have tor return anything, or it could be an async iterator, so that you can interact with the nodes as they are formed. This is flexible right now, but I think the iterator idea is nice.

  • _from_flat (names welcome) is a new function that creates a group entirely from a hierarchy dict + a store. that dict must specify a root group, otherwise an exception is raised. We could revise this to create a root group if one is not specified. Open to suggestions here.

Implicit groups

Partial hierarchies like {'a': GroupMetadata(), 'a/b/c': ArrayMetadata(...)} implicitly denote a group at a/b. When creating such a hierarchy, if we find an existing group at a/b, then we don't need to create a new one. So in the context of modeling a hierarchy, implicit groups are a little special -- by not specifying the properties of the group, the user / application is tolerant of any group being there. So I introduced a subclass of GroupMetadata called _ImplicitGroupMetadata, which can be inserted into a hierarchy dict to explicitly denote groups that don't need to be written if one already exists. _ImplicitGroupMetadata is just like GroupMetadata except it errors if you try to set any parameter except zarr_format.

streaming v2 vs v3 node creation

creating v3 arrays / groups requires writing 1 metadata document, but v2 requires 2. To get the most concurrency I await the write of each metadata document separately, which means that foo/.zattrs might resolve before foo/.zarray. So in the v2 case I only yield up an array / group when both documents were written.

Overlap with metadata consolidation logic

there's a lot of similarity between the stuff in this PR and routines used for consolidated metadata. it would be great to find ways to factor out some of the overlap areas

still to do:

  • write some more tests (checking that implicit groups don't get written if a group already exists)
  • handle overwriting. I think the plan here is, if overwrite is False, then we do a check before any writing to ensure that there are no conflicts between the proposed hierarchy and the stuff that actually exists in storage. this check will involve more IO.

@d-v-b
Copy link
Contributor Author

d-v-b commented Jan 25, 2025

@TomNicholas you should have a look at some of these new functions / methods. I'd be happy to change things if you have some datatree conventions you'd like to suggest

This method takes a dictionary where the keys are the names of the arrays or groups
to create and the values are the metadata or objects representing the arrays or groups.

The method returns an asynchronous iterator over the created nodes.
Copy link
Contributor

Choose a reason for hiding this comment

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

A note about when to use this would be great

Copy link
Contributor Author

Choose a reason for hiding this comment

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

good point, would something like "use this method to create an entire tree of sub-groups and / or sub-arrays efficiently." suffice?

@@ -1407,6 +1395,42 @@ async def _members(
):
yield member

# TODO: find a better name for this. create_tree could work.
Copy link
Contributor

Choose a reason for hiding this comment

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

hierarchy is fine to me.

"""
ctx: asyncio.Semaphore | contextlib.nullcontext[None]

if semaphore is None:
Copy link
Contributor

@dcherian dcherian Jan 27, 2025

Choose a reason for hiding this comment

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

why do we need None as an option here? And why is it "no semaphore" instead of "default concurrency semaphore"?

Copy link
Contributor Author

@d-v-b d-v-b Jan 27, 2025

Choose a reason for hiding this comment

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

create_nodes is low-level, kind of dangerous, and the expectation is that the user of this function knows what they are doing. So it doesn't default with any concurrency limit. It also doesn't check if nodes already exist, or if the user wants to nest arrays inside arrays. A higher level function (create_hierarchy) does all that checking, and that's where the concurrency limit defaults to the value in the config.

Copy link
Contributor

Choose a reason for hiding this comment

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

so why not require the semaphore if it's an advanced user?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

the semaphore is only necessary if you want some concurrency limiting mechanism, but I think that's a special case. A default of None means that users of the function don't need to run from asyncio import Semaphore before creating some nodes. E.g., there are a lot of places in the tests where I would want to use create_nodes, and in basically none of those places would a semaphore be necessary.

# We will iterate over the dict again, but a full pass here ensures that the error message
# is comprehensive, and I think the performance cost will be negligible.
for k, v in data.items():
observed_zarr_formats[v.zarr_format].append(k)
Copy link
Contributor

Choose a reason for hiding this comment

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

can we short-circuit inside the loop and raise on the first instance of zarr_format that is not equal to next(data.items())[1].zarr_format?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

we definitely could, but we get a much nicer error message if we can traverse the entire proposed hierarchy and identify all the problematic nodes.

document stored at store_path.path / (.zgroup | .zarray). If no such document is found,
raise a KeyError.
"""
# TODO: consider first fetching array metadata, and only fetching group metadata when we don't
Copy link
Contributor

Choose a reason for hiding this comment

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

i think it's better to just grab all 3 at once honestly

"""


def create_array_metadata(
Copy link
Contributor

Choose a reason for hiding this comment

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

seems like this would duplicate logic form elsewhere? Is there nowhere else this function could be used?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

this should be duplicated, but sadly these functions don't exist in the codebase yet! I add some functions like this in another PR: #2761

@d-v-b
Copy link
Contributor Author

d-v-b commented Jan 28, 2025

question: should we export a sync version of create_hierarchy from the top-level zarr namespace?

@d-v-b d-v-b removed the needs release notes Automatically applied to PRs which haven't added release notes label Jan 28, 2025
@dcherian
Copy link
Contributor

question: should we export a sync version of create_hierarchy from the top-level zarr namespace?

Yes, this would be used in Xarray.

@github-actions github-actions bot added the needs release notes Automatically applied to PRs which haven't added release notes label Jan 28, 2025
@d-v-b
Copy link
Contributor Author

d-v-b commented Jan 28, 2025

this PR adds a few functions that have async implementations and sync wrappers, like create_hierarchy_a (async) and create_hierarchy. Instead of putting (async, sync) pairs in the same module with slightly different names, I wonder if we should split the group module into an async-only namespace and a sync namespace (that imports stuff from the async namespace)? Then we don't need to mangle function names.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New features or improvements needs release notes Automatically applied to PRs which haven't added release notes
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants