Generic tool for caching small metadata ops in datapipes

[coreyjadams]

PhysicsNeMo Pull Request

Initial draft of datapipes caching for small files and arrays and metadata

Caching datapipes is something @peterdsharpe and I had been kicking ideas around since the winter, Max suggested it recently as a major perf improvement on lustre for mesh datasets so I am testing a generic tool to implement his ideas. The idea is to deploy something that can work on almost all of our datasets, but is relatively not invasive.

I'll add more here with benchmarks and expect to refine this. Opening this up to share the design ideas. No need to review this yet.

Description

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

CODEOWNERS review map

Current for commit c62047f34dc7. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @coreyjadams — 20 file(s)

• docs/api/datapipes/physicsnemo.datapipes.rst
• examples/cfd/external_aerodynamics/unified_external_aero_recipe/conf/base.yaml
• examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/datasets.py
• examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/merge_global_data.py
• examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/train.py
• examples/cfd/external_aerodynamics/unified_external_aero_recipe/tests/test_merge_global_data_cache.py
• examples/cfd/external_aerodynamics/unified_external_aero_recipe/tests/test_reader_cache_wiring.py
• physicsnemo/datapipes/init.py
• physicsnemo/datapipes/caching-design.md
• physicsnemo/datapipes/caching.py
• physicsnemo/datapipes/readers/base.py
• physicsnemo/datapipes/readers/hdf5.py
• physicsnemo/datapipes/readers/mesh.py
• physicsnemo/datapipes/readers/numpy.py
• physicsnemo/datapipes/readers/tensorstore_zarr.py
• physicsnemo/datapipes/readers/vtk.py
• physicsnemo/datapipes/readers/zarr.py
• test/datapipes/core/test_caching.py
• test/datapipes/readers/test_mesh_reader_cache.py
• test/datapipes/readers/test_reader_cache_integrations.py

⏳ @megnvidia — 1 file(s)

• docs/api/datapipes/physicsnemo.datapipes.rst

⏳ @peterdsharpe — 6 file(s)

• examples/cfd/external_aerodynamics/unified_external_aero_recipe/conf/base.yaml
• examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/datasets.py
• examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/merge_global_data.py
• examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/train.py
• examples/cfd/external_aerodynamics/unified_external_aero_recipe/tests/test_merge_global_data_cache.py
• examples/cfd/external_aerodynamics/unified_external_aero_recipe/tests/test_reader_cache_wiring.py

Comment /codeowners-info to refresh.