HRRR Async Refactor

CHANGELOG.md

@@ -22,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Hemispheric centred bred vector perturbation now supports single/odd batch sizes
 - Refactored NCAR ERA5 source to have async structure
 - Refactored GFS and GFS_FX to have async structure
+- Refactored HRRR and HRRR_FX to have async structure
 - Expanded the data source protocol to also include async fetch functions for async
   data sources
 

earth2studio/data/base.py

@@ -116,7 +116,7 @@ async def fetch(  # type: ignore[override]
         lead_time: timedelta | list[timedelta] | LeadTimeArray,
         variable: str | list[str] | VariableArray,
     ) -> xr.DataArray:
-        """Async function to get data. Async data sources support.
+        """Async function to get data. Async forecast sources support this.
 
         Parameters
         ----------

earth2studio/data/gfs.py

@@ -235,6 +235,12 @@ async def fetch(
         if not self._cache:
             shutil.rmtree(self.cache)
 
+        # Close aiohttp client if s3fs
+        # https://github.com/fsspec/s3fs/issues/943
+        # https://github.com/zarr-developers/zarr-python/issues/2901
+        if isinstance(self.fs, s3fs.S3FileSystem):
+            s3fs.S3FileSystem.close_session(asyncio.get_event_loop(), self.fs.s3)
+
         return xr_array.isel(lead_time=0)
 
     async def _create_tasks(
@@ -343,7 +349,7 @@ async def fetch_array(
         xr.DataArray
             FS data array for given time and lead time
         """
-        logger.debug(f"Fetching GRS grib file: {grib_uri} {byte_offset}-{byte_length}")
+        logger.debug(f"Fetching GFS grib file: {grib_uri} {byte_offset}-{byte_length}")
         # Download the grib file to cache
         grib_file = await self._fetch_remote_file(
             grib_uri,
@@ -386,7 +392,6 @@ async def _fetch_index(self, index_uri: str) -> dict[str, tuple[int, int]]:
             Dictionary of GFS vairables (byte offset, byte length)
         """
         # Grab index file
-        # TODO: Change remote file to be more proper fetch
         index_file = await self._fetch_remote_file(index_uri)
         with open(index_file) as file:
             index_lines = [line.rstrip() for line in file]