Optimize and migrate MRMS and NClimGridDaily data sources on obstore - #1061
Merged
Conversation
Replaces the s3fs filesystem with an obstore S3 store using the shared byte-range helpers. Day-directory listings go through obstore_list_prefix with per-instance memoization: past (complete) days are cached so repeated fetches over the same day share one LIST request, while the current UTC day is always re-listed. Gzip decompression and the cache write now run in a worker thread instead of blocking the event loop. Cache file names hash the same bucket-prefixed s3:// URI as before, so warm caches remain valid.
Replaces the synchronous s3fs streaming reads with obstore. Each monthly NetCDF (~58 MB) is now downloaded once into the cache via obstore_fetch_to_cache and shared by every (day, variable) slice of that month, instead of every slice independently seeking through the remote file over fsspec. Months whose per-slice caches are already warm skip the download entirely. Benchmarked at ~1.2x faster medians for multi-day requests with far fewer network-latency outliers (16 s worst case dropped to 6 s).
MRMS fetch time was dominated by latlons(), which decodes full 2-D coordinate arrays (~8 s per 3500x7000 file) even though the data source only keeps the 1-D axes of the regular CONUS lat/lon grid. Compute the axes from the grid-definition header keys instead (~0.1 ms, verified identical to within 2e-6 degrees on a 0.01-degree grid), keeping latlons() as a fallback for non-regular grids. The decode also moves into a worker thread so concurrent (time, product) tasks no longer serialize on the event loop. Full-lexicon fetch of 2 times x 2 products: ~32 s -> ~5 s.
for more information, see https://pre-commit.ci
Contributor
Greptile SummaryThe PR migrates MRMS and NClimGridDaily remote reads from s3fs to obstore and accelerates MRMS coordinate decoding.
|
| Filename | Overview |
|---|---|
| earth2studio/data/mrms.py | Migrates discovery and reads to obstore and moves optimized GRIB decoding off the event loop; no actionable defect was established for supported products. |
| earth2studio/data/nclimgrid.py | Adds monthly local downloads, but stores invocation-specific paths on the instance, allowing concurrent fetches to overwrite one another. |
| test/data/test_nclimgrid.py | Updates the mocked download seam but does not cover overlapping async fetches on one instance. |
| CHANGELOG.md | Documents the MRMS and NClimGridDaily obstore migrations and MRMS decoding optimization. |
Reviews (1): Last reviewed commit: "Compute MRMS lat/lon axes from grib head..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
This was referenced Aug 13, 2026
negin513
changed the title
negin513
changed the title
NickGeneva
reviewed
Aug 13, 2026
NickGeneva
reviewed
Aug 13, 2026
NickGeneva
approved these changes
Aug 13, 2026
Member
Author
|
/ok to test |
@negin513, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/ |
Member
Author
|
/okay to test 18552ec |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Earth2Studio Pull Request
Description
Closes NVIDIA/physicsnemo-roadmap#2882, closes NVIDIA/physicsnemo-roadmap#2885.
Migrates the MRMS and NClimGridDaily data sources from s3fs to obstore, continuing the migration series (GOES #1024, Himawari AHI #1043, ISD/IBTrACS/CFS-reforecast/OPERA #1058):
obstore_list_prefixwith per-instance memoization: past (complete) days are cached so repeated fetches over the same day share one LIST request, while the current UTC day is always re-listed (the GOES pattern). Gzip decompression and cache writes now run in a worker thread. Cache file names hash the same bucket-prefixeds3://URIs as before, so warm caches remain valid.latlons(), which decodes full 2-D coordinate arrays (~8 s per 3500×7000 file) even though the source only keeps the 1-D axes of the regular CONUS grid. The axes are now computed from grid-definition header keys (~0.1 ms, identical to within 2×10⁻⁶ degrees on the 0.01° grid), withlatlons()kept as a fallback for non-regular grids. The decode also moves into a worker thread so concurrent (time, product) tasks no longer serialize on the event loop.obstore_fetch_to_cache, shared by every (day, variable) slice of that month. Months whose per-slice caches are already warm skip the download entirely.Performance
Full lexicon,
cache=False, 3 reps, same host and time window:Nearly all of the MRMS win comes from the header-key axis computation and threaded decode; the obstore transport itself was a wash there (gunzip is 0.01 s, download 0.6 s —
latlons()was ~8 s of every ~9 s task). Output verified: identical shapes, no NaN change, axes matchlatlons()to float rounding.Checklist
test_mrms.py,test_nclimgrid.py; mocked NClimGrid test updated to the new download seam).Dependencies
No new dependencies; obstore and obspec are already core dependencies.