Conversation
|
Someone needs to test this on USDF before review & merge |
There was a problem hiding this comment.
Pull request overview
This PR aims to reduce repeated LSST Butler lookups by pre-resolving and caching deep_coadd dataset references (keyed by tract/patch/band) and caching skymap retrieval within LSSTDataset. It also extends the test Butler mocks to support query_datasets() and get(ref) to simulate the new access pattern.
Changes:
- Add an internal
deep_coaddref cache built viabutler.query_datasets()and switch patch retrieval tobutler.get(ref). - Add an
lru_cache-backed_get_skymap()helper to avoid repeated skymap fetches. - Extend
MockButlerwithquery_datasets()and a minimalMockDatasetRefso tests can cover the new ref-based retrieval.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
src/hyrax/datasets/lsst_dataset.py |
Adds deep_coadd ref pre-resolution/caching and cached skymap retrieval; updates patch requests to use resolved refs. |
tests/hyrax/mocks/lsst_butler_mocks.py |
Adds minimal query_datasets() + MockDatasetRef and enables MockButler.get(ref) to support the new dataset code path. |
| self._deep_coadd_refs = {} | ||
| if self._butler_config is not None: | ||
| self._build_deep_coadd_ref_cache() |
| for band in LSSTDataset.BANDS: | ||
| data_id = {"skymap": self.config["data_set"]["skymap"], "band": band} | ||
| refs = list( | ||
| butler.query_datasets( | ||
| "deep_coadd", | ||
| data_id=data_id, | ||
| collections=self._butler_config["collections"], | ||
| ) | ||
| ) | ||
| for ref in refs: | ||
| key = (ref.dataId["tract"], ref.dataId["patch"], ref.dataId["band"]) | ||
| if (key[0], key[1]) in needed_tract_patches: |
| @functools.lru_cache(maxsize=4) # noqa: B019 | ||
| def _get_skymap(self, skymap_name): | ||
| """Fetch and cache the immutable skymap object by name.""" | ||
| return self._get_butler_thread_safe().get("skyMap", {"skymap": skymap_name}) |
| for band in LSSTDataset.BANDS: | ||
| # Set up the data dict | ||
| butler_dict = { | ||
| "tract": tract_index, | ||
| "patch": patch_index, | ||
| "skymap": self.config["data_set"]["skymap"], | ||
| "band": band, | ||
| } | ||
|
|
||
| # pull from butler | ||
| image = self._get_butler_thread_safe().get("deep_coadd", butler_dict) | ||
| ref = self._deep_coadd_refs[(tract_index, patch_index, band)] | ||
| image = self._get_butler_thread_safe().get(ref) |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #933 +/- ##
==========================================
+ Coverage 62.94% 63.04% +0.09%
==========================================
Files 71 71
Lines 7357 7376 +19
==========================================
+ Hits 4631 4650 +19
Misses 2726 2726 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
I will try to test this at USDF soon to have some benchmarks. |
Motivation
Description
self._deep_coadd_refsand call_build_deep_coadd_ref_cache()when a Butler configuration is present inLSSTDataset._build_deep_coadd_ref_cache()to bulk-query the butler fordeep_coadddataset references per band and store them keyed by(tract, patch, band)for later retrieval._get_tract_patch()with a cached_get_skymap()helper decorated withfunctools.lru_cacheto avoid repeated lookups._request_patch()to use the pre-resolvedMockDatasetRef/butler refs viaself._deep_coadd_refs[(tract_index, patch_index, band)]and call the butler with the ref instead of constructing dict-based IDs.tests/hyrax/mocks/lsst_butler_mocks.pyby addingMockButler.query_datasets(), a minimalMockDatasetRefclass, and makingMockButler.get()acceptMockDatasetRefinstances so the tests can simulatequery_datasets+get(ref)behavior.Testing
tests/hyraxwithpytestto exercise the LSST dataset and mock butler paths, and the test suite passed.MockButler.query_datasets()andMockButler.get()in tests to verify resolving and fetchingdeep_coaddexposures, and those checks passed.Codex Task