Optimized HealDA data loaders to avoid GIL

[negin513]

Description

Migrates the HealDA ZarrLoader remote zarr reads and s3:// downloads in physicsnemo.core.filesystem from fsspec/s3fs to obstore. obstore's Rust-native client performs object fetches off the GIL with real connection pooling, so concurrent chunk reads no longer serialize on the Python interpreter.

Benchmark

Concurrent zarr chunk reads against public HRRR data on NOAA S3 (4 surface fields × 96 chunks = 384 GETs of ~45 KB, cold cache, 3 trials per backend), using the exact ZarrLoader code path:

Backend	Median wall clock
fsspec/s3fs (before)	12.4 s
obstore (after)	1.6 s (7.7× faster)

This is the many-small-GET regime that HealDA's per-timestep 2D/3D variable reads hit; for large-chunk stores both backends converge toward bandwidth-bound parity.

Changes

• ZarrLoader: remote URLs now open through zarr.storage.ObjectStore backed by obstore.store.from_url. fsspec-style storage_options keys (anon, key, secret, token, endpoint_url, region_name) are translated to obstore config names so existing call sites keep working; unknown keys pass through. Falls back to the fsspec path if obstore is not installed.
• core/filesystem.py: s3:// single-file and recursive downloads in _download_cached stream through obstore (8 MB chunks, layout-preserving recursive fetch, FileNotFoundError on empty prefixes). msc:// stays on fsspec; the checkpoint write path (_get_fs) is deliberately unchanged.
• Adds obstore>=0.6.0 dependency.
• New tests: test/datapipes/healda/test_zarr_loader_obstore.py (local baseline, remote-via-obstore with data verification, storage-options translation, read-only store) and four obstore download tests in test/core/test_filesystem.py against an in-memory store.

Notes for reviewers

• Credentials now resolve through obstore's chain (AWS env vars, profiles, IMDS) instead of s3fs/botocore — equivalent for standard setups, flagging in case anyone relies on exotic botocore config.
• The benchmark figure is committed under docs/img/ to embed here; happy to drop it from the tree if preferred.

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.

[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 290fdba7537b. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @CharlelieLrt — 1 file(s)

• physicsnemo/core/filesystem.py

⏳ @coreyjadams — 4 file(s)

• physicsnemo/core/filesystem.py
• pyproject.toml
• test/core/test_filesystem.py
• test/datapipes/healda/test_zarr_loader_obstore.py

⏳ @ktangsali — 2 file(s)

• physicsnemo/core/filesystem.py
• pyproject.toml

⏳ @pzharrington — 1 file(s)

• physicsnemo/experimental/datapipes/healda/loaders/zarr_loader.py

No CODEOWNER

• CHANGELOG.md
• docs/img/healda_obstore_benchmark.png
• uv.lock

Comment /codeowners-info to refresh.

[greptile-apps]

Greptile Summary

The PR migrates HealDA remote Zarr reads and cached S3 downloads from fsspec/s3fs to obstore to improve concurrent object-fetch performance.

• Adds an obstore-backed, read-only Zarr store with fsspec option translation.
• Adds streamed single-object and layout-preserving recursive S3 downloads.
• Adds obstore as a dependency and covers local, remote, option-translation, and download behavior with tests.

Security Review

The recursive S3 downloader does not constrain paths derived from object keys to its destination. A key containing parent-directory components can therefore cause writes outside the cache directory.

Important Files Changed

Filename	Overview
physicsnemo/core/filesystem.py	Replaces S3 downloads with obstore streaming, but recursive downloads permit destination escape through object-key path traversal.
physicsnemo/experimental/datapipes/healda/loaders/zarr_loader.py	Routes remote Zarr stores through a read-only obstore adapter while retaining the existing fsspec fallback when obstore is unavailable.
test/core/test_filesystem.py	Adds in-memory coverage for single and recursive downloads, but does not test hostile object keys or destination confinement.
test/datapipes/healda/test_zarr_loader_obstore.py	Covers local loading, obstore-backed remote loading, option translation, data correctness, and read-only behavior.
pyproject.toml	Adds the obstore runtime dependency.
uv.lock	Updates the lockfile to include obstore and its resolved metadata.

_{Reviews (1): Last reviewed commit: "Optimize HealDA data loaders with GIL-fr..." | Re-trigger Greptile}

[greptile-apps]

Recursive download path traversal

When a recursively listed S3 object key contains parent-directory components, rel is joined directly to destination without a containment check, causing the download to overwrite writable files outside the cache entry. Resolve each destination and reject it unless it remains beneath the recursive download root.

How this was verified: The public Package.get(..., recursive=True) path reaches these lines, where the raw object-key suffix is passed to Path(destination) / rel without sanitization or containment validation.

[pzharrington]

Is this used in docs somewhere?

[pzharrington]

HealDA changes look good to me, nice speedup