Migrate JPSS VIIRS, ATMS, and CrIS data sources to obstore - #1062
Merged
Conversation
Replaces s3fs with the shared obstore helpers across the three JPSS sources, which share the NOAA NESDIS S3 bucket layout and the granule-discovery-by-listing pattern: - JPSS (VIIRS): single-bucket store; day-directory listings memoized per instance once a day can no longer gain files, so repeated (time, variable) requests issue one LIST instead of one per pair; available() now verifies both data and geolocation prefixes via a sync obstore listing - JPSS_ATMS / JPSS_CRIS: one store per satellite bucket; missing day directories are detected via empty listings (obstore does not raise FileNotFoundError), and the CrIS SDR/GEO dual listings remain concurrent via asyncio.gather over obstore_list_prefix - Cache-key hashing is unchanged everywhere (sha256 of the historical bucket-prefixed path / s3://bucket/key URI), so warm caches populated before the migration remain valid - Tests: fake obspec store injection (Himawari pattern) for listing, discovery, availability, and a full VIIRS __call__ mock; no obstore internals patched
Contributor
Greptile SummaryThe PR migrates JPSS VIIRS, ATMS, and CrIS discovery and downloads from s3fs to per-bucket obstore stores while preserving cache keys and near-real-time listing behavior.
|
| Filename | Overview |
|---|---|
| earth2studio/data/jpss.py | Migrates VIIRS listing, availability checks, and cached downloads to obstore while retaining bucket-prefixed cache hashing. |
| earth2studio/data/jpss_atms.py | Adds per-satellite obstore instances, bucket-relative day listings, and obstore byte-range downloads without changing task-selection semantics. |
| earth2studio/data/jpss_cris.py | Migrates concurrent SDR/GEO discovery and paired downloads to per-bucket obstore stores while preserving granule matching. |
| test/data/test_jpss.py | Adds offline coverage for availability, listing memoization, closest-granule selection, complete VIIRS fetches, and warm-cache compatibility. |
| test/data/test_jpss_atms.py | Adds fake-store tests for ATMS task discovery, time-window filtering, and empty listings. |
| test/data/test_jpss_cris.py | Adds fake-store tests for CrIS SDR/GEO pairing, missing geolocation files, and empty listings. |
Reviews (1): Last reviewed commit: "Migrate JPSS VIIRS, ATMS, and CrIS data ..." | Re-trigger Greptile
Replaces the per-(FOV, channel) Python dict loop in _decode_bufr with numpy column assembly (repeat/tile over the arrays eccodes already returns) and one pyarrow Table per BUFR message, concatenated and converted to pandas once per file. This removes ~1M dict appends plus eleven sequential column astype passes for a typical +/-10m window. Decode semantics are preserved exactly and verified by exact-equality comparison (pandas assert_frame_equal, check_exact) of the full 931k-row DataFrame against the previous implementation on a real NOAA-20 granule set: identical values, dtypes, row order, and NaN handling (the historical fill test keeps NaN observations; invalid FOV timestamps coerce to NaT and are dropped, matching the old ValueError skip). Warm-cache decode+compile drops from ~1.8s to ~0.6s for 37 granules (931k rows); cold end-to-end median 2.2s -> 1.4s on top of the obstore migration (2.8s on main), i.e. ~2x for the source overall.
fetch_data and fetch_geolocation decoded HDF5 inline in async coroutines, blocking the event loop (and therefore all in-flight granule downloads) for the duration of each decode. The h5py blocks now run via asyncio.to_thread under a module-level _HDF5_LOCK (HDF5 is not thread-safe), following the himawari_ahi.py pattern. Decoded output verified identical (np.array_equal with equal_nan) on a live NOAA-20 M-band fetch. End-to-end time is unchanged at small request sizes (interleaved A/B within noise); the win is for large time/variable fan-outs where decode otherwise serializes downloads.
fetch() previously ran two serial phases: download every unique BUFR file, then decode them all. Each file now decodes in a worker thread (serialized by a module-level _ECCODES_LOCK — eccodes keeps global C state and is not thread-safe) as soon as its download lands, while other downloads continue. Decode work is also deduplicated to once per unique (uri, variable) instead of once per task, so overlapping tolerance windows no longer re-decode the same file. _compile_dataframe keeps all assembly semantics (task order, per-task window mask, drop_duplicates, empty-result shape, attrs). Output verified bit-identical (assert_frame_equal check_exact) against the 931k-row baseline. Cold end-to-end median 1.4s -> 0.9s.
CrIS is transfer-bound: granule HDF5 files are tens of MB and a typical request touches only a few files, so each file downloaded as a single whole-object GET at single-stream S3 throughput while the worker pool sat idle. _fetch_remote_file now reads large objects through a chunked helper: head for size, then 1 MiB byte-range GETs with up to 8 in flight per file, reassembled in order. Chunk size was tuned empirically on a real 16 MB SCRIF granule (1 MiB x 8 beat 2-8 MiB variants, which yield too few concurrent streams). Objects at or under one chunk take a single range GET. Retry/backoff, exception tuple, cache-write, and not-found translation (FileNotFoundError, mirroring obstore_read_range) are unchanged; the helper stays local to jpss_cris.py to avoid touching utils.py while another PR has it in flight. Byte parity verified via sha256 against the single-range path on a real granule; offline fake-store tests cover chunk reassembly, range coverage, the small-object path, and not-found. Cold end-to-end median ~8.2s -> ~4.4s (~1.9x).
6 tasks
NickGeneva
approved these changes
Aug 14, 2026
NickGeneva
left a comment
NickGeneva
left a comment
Collaborator
There was a problem hiding this comment.
Over all looks good, main item is the streaming util in the cris streaming, if that should be something thats more general util. Approving to unblock.
Address PR #1062 review feedback: - Move the CrIS-local _read_object_chunked helper to earth2studio/data/utils.py as obstore_read_chunked and give obstore_fetch_to_cache a chunked option, so CrIS granule downloads go through the shared save-to-cache util (cache file names unchanged) - Relocate the chunked-read unit tests next to the other obstore util tests in test_data_utils.py and cover the chunked ValueError guard - Clarify why VIIRS HDF5 decode runs via asyncio.to_thread despite _HDF5_LOCK, and why day-listing memoization waits an hour past the UTC day boundary
…ling-ladybug # Conflicts: # CHANGELOG.md
Member
Author
|
/ok to test 384a75a |
negin513
added a commit
to negin513/earth2studio
that referenced
this pull request
Aug 17, 2026
… series Benchmark harness and results comparing every optimized data source before/after the obstore migration series (NVIDIA#855-NVIDIA#1065): cold cache, maximum common variable set, sequential runs, per-run subprocess isolation. Includes claimed-vs-measured comparison against PR-body numbers and rows for the open PRs (NVIDIA#914, NVIDIA#1062, NVIDIA#1063, NVIDIA#1065).
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#2879, closes NVIDIA/physicsnemo-roadmap#2880, closes NVIDIA/physicsnemo-roadmap#2881.
Migrates the three JPSS data sources from s3fs to obstore as one batched effort — they share the NOAA NESDIS S3 bucket layout (
noaa-nesdis-{n20,n21,snpp}-pds) and the granule-discovery-by-listing pattern, so the migration and test-fixture work amortizes across all three. Continues the pattern established by GOES/GLM (#1042), Himawari AHI (#1043), and #1058.obstore_list_prefixwith per-instance memoization: previously every(time, variable)pair re-listed the same day directory (plus one more per geolocation fetch); now a completed day is listed once per product folder, while in-progress days are always re-listed so near-real-time polling still sees new granules.available()moves from sync s3fs to a sync obstore listing that verifies both the data and geolocation prefixes (one listed entry proves existence)._async_initbuilds one store per bucket. The_ls+FileNotFoundErrorhandling becomes an empty-listing check (obstore lists missing prefixes as empty). Download retry loop, backoff, and exception semantics are unchanged; only the transfer line moved toobstore_read_range.asyncio.gatherat the old jpss_cris.py:839) stays concurrent as two gatheredobstore_list_prefixcoroutines, and the entire two-levelFileNotFoundErrorfallback block is deleted — empty listings subsume it. Completed-day listings are memoized per bucket.Warm-cache compatibility: cache-key hashing is byte-identical everywhere — VIIRS hashes the same bucket-prefixed
bucket/keypaths s3fs used to return, and ATMS/CrIS rebuild the exacts3://bucket/keyURI strings that_cache_pathhashes — so local caches populated before the migration remain valid.Tests: each source gains offline tests that inject fake obspec stores (the
_FakeStorepattern from test_himawari_ahi.py) rather than patching obstore internals — covering granule discovery/time-window selection, SDR↔GEO pairing and missing-GEO skip (CrIS), empty-listing/data-gap branches,available(), and a full VIIRS__call__mock over synthetic HDF5 granules. Combined suite: 57 passed offline, plus the two liveavailable()network tests xpass through the new obstore path.earth2studio/data/utils.pyis untouched — all helpers needed already exist — so there is no overlap with #1058.Checklist
Performance
Cold-fetch medians (
cache=False, same host; interleaved A/B where noted):Variable-space coverage: ATMS is inherently the full space (the single
atmsvariable carries all 22 channels viasensor_index). VIIRS is covered for the full M-band lexicon (16 variables); I-band (5 variables, ~4× larger granules) and L2 EDR (14 variables) run the identical code path with different file sizes/counts and were not separately benchmarked. CrIS is covered at both the subset (3 of 2211 channels) and full-spectrum extremes.VIIRS full M-band reads as parity because 17 granule downloads dominate and S3 variance is high (worst rep 36s on a slow granule — bucket-side, both code paths affected). CrIS full-spectrum dilutes from 1.9× to ~1.5× because decode's share grows relative to transfer at 16M rows; its remaining warm-path costs are S3 listing latency (~60%) and the Planck inversion in
radiance_to_bt, both out of scope here.Beyond the migration itself, the follow-up commits:
astypepasses; warm decode+compile 1.8s → 0.6s for 37 granules. Verified bit-identical viaassert_frame_equal(check_exact=True)on the full 931k-row frame (values, dtypes, row order, NaN retention, NaT skips).asyncio.to_threadunder an HDF5 lock (himawari pattern). No measured end-to-end change at small request sizes (interleaved A/B within noise); prevents decode from serializing downloads on large fan-outs. Output verified identical.CrIS decode was measured at 0.14s warm and left alone — it is transfer-bound, which is why its win comes from the download path.