Xbatcher Caching Feature

This notebook demonstrates the new caching feature added to xbatcher’s BatchGenerator. This feature allows you to cache batches, potentially improving performance for repeated access to the same batches.

Introduction

The caching feature in xbatcher’s BatchGenerator allows you to store generated batches in a cache, which can significantly speed up subsequent accesses to the same batches. This is particularly useful in scenarios where you need to iterate over the same dataset multiple times.

The cache is pluggable, meaning you can use any dict-like object to store the cache. This flexibility allows for various storage backends, including local storage, distributed storage systems, or cloud storage solutions.

Installation

To use the caching feature, you’ll need to have xbatcher installed, along with zarr for serialization. If you haven’t already, you can install these using pip:

python -m pip install xbatcher zarr

or

using conda:

conda install -c conda-forge xbatcher zarr

Basic Usage

Let’s start with a basic example of how to use the caching feature:

import tempfile

import xarray as xr
import zarr

import xbatcher

# create a cache using Zarr's DirectoryStore
directory = f'{tempfile.mkdtemp()}/xbatcher-cache'
print(directory)
cache = zarr.storage.DirectoryStore(directory)

/tmp/tmphgywdjbm/xbatcher-cache

In this example, we’re using a local directory to store the cache, but you could use any zarr-compatible store, such as S3, Redis, etc.

# load a sample dataset
ds = xr.tutorial.open_dataset('air_temperature', chunks={})
ds

<xarray.Dataset> Size: 31MB
Dimensions:  (lat: 25, time: 2920, lon: 53)
Coordinates:
  * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
  * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
  * time     (time) datetime64[ns] 23kB 2013-01-01 ... 2014-12-31T18:00:00
Data variables:
    air      (time, lat, lon) float64 31MB dask.array<chunksize=(2920, 25, 53), meta=np.ndarray>
Attributes:
    Conventions:  COARDS
    title:        4x daily NMC reanalysis (1948)
    description:  Data is from NMC initialized reanalysis\n(4x/day).  These a...
    platform:     Model
    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...

xarray.Dataset

• Dimensions:
  • lat: 25
  • time: 2920
  • lon: 53
• Coordinates: (3)
  • lat
    (lat)
    float32
    75.0 72.5 70.0 ... 20.0 17.5 15.0

    standard_name :

    latitude

    long_name :

    Latitude

    units :

    degrees_north

    axis :

    Y

    array([75. , 72.5, 70. , 67.5, 65. , 62.5, 60. , 57.5, 55. , 52.5, 50. , 47.5,
           45. , 42.5, 40. , 37.5, 35. , 32.5, 30. , 27.5, 25. , 22.5, 20. , 17.5,
           15. ], dtype=float32)

  • lon
    (lon)
    float32
    200.0 202.5 205.0 ... 327.5 330.0

    standard_name :

    longitude

    long_name :

    Longitude

    units :

    degrees_east

    axis :

    X

    array([200. , 202.5, 205. , 207.5, 210. , 212.5, 215. , 217.5, 220. , 222.5,
           225. , 227.5, 230. , 232.5, 235. , 237.5, 240. , 242.5, 245. , 247.5,
           250. , 252.5, 255. , 257.5, 260. , 262.5, 265. , 267.5, 270. , 272.5,
           275. , 277.5, 280. , 282.5, 285. , 287.5, 290. , 292.5, 295. , 297.5,
           300. , 302.5, 305. , 307.5, 310. , 312.5, 315. , 317.5, 320. , 322.5,
           325. , 327.5, 330. ], dtype=float32)

  • time
    (time)
    datetime64[ns]
    2013-01-01 ... 2014-12-31T18:00:00

    standard_name :

    time

    long_name :

    Time

    array(['2013-01-01T00:00:00.000000000', '2013-01-01T06:00:00.000000000',
           '2013-01-01T12:00:00.000000000', ..., '2014-12-31T06:00:00.000000000',
           '2014-12-31T12:00:00.000000000', '2014-12-31T18:00:00.000000000'],
          shape=(2920,), dtype='datetime64[ns]')

• Data variables: (1)
  • air
    (time, lat, lon)
    float64
    dask.array<chunksize=(2920, 25, 53), meta=np.ndarray>

    long_name :

    4xDaily Air temperature at sigma level 995

    units :

    degK

    precision :

    2

    GRIB_id :

    11

    GRIB_name :

    TMP

    var_desc :

    Air temperature

    dataset :

    NMC Reanalysis

    level_desc :

    Surface

    statistic :

    Individual Obs

    parent_stat :

    Other

    actual_range :

    [185.16 322.1 ]

    Array Chunk Bytes 29.52 MiB 29.52 MiB Shape (2920, 25, 53) (2920, 25, 53) Dask graph 1 chunks in 2 graph layers Data type float64 numpy.ndarray

• Indexes: (3)
  • lat
    PandasIndex

    PandasIndex(Index([75.0, 72.5, 70.0, 67.5, 65.0, 62.5, 60.0, 57.5, 55.0, 52.5, 50.0, 47.5,
           45.0, 42.5, 40.0, 37.5, 35.0, 32.5, 30.0, 27.5, 25.0, 22.5, 20.0, 17.5,
           15.0],
          dtype='float32', name='lat'))

  • lon
    PandasIndex

    PandasIndex(Index([200.0, 202.5, 205.0, 207.5, 210.0, 212.5, 215.0, 217.5, 220.0, 222.5,
           225.0, 227.5, 230.0, 232.5, 235.0, 237.5, 240.0, 242.5, 245.0, 247.5,
           250.0, 252.5, 255.0, 257.5, 260.0, 262.5, 265.0, 267.5, 270.0, 272.5,
           275.0, 277.5, 280.0, 282.5, 285.0, 287.5, 290.0, 292.5, 295.0, 297.5,
           300.0, 302.5, 305.0, 307.5, 310.0, 312.5, 315.0, 317.5, 320.0, 322.5,
           325.0, 327.5, 330.0],
          dtype='float32', name='lon'))

  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2013-01-01 00:00:00', '2013-01-01 06:00:00',
                   '2013-01-01 12:00:00', '2013-01-01 18:00:00',
                   '2013-01-02 00:00:00', '2013-01-02 06:00:00',
                   '2013-01-02 12:00:00', '2013-01-02 18:00:00',
                   '2013-01-03 00:00:00', '2013-01-03 06:00:00',
                   ...
                   '2014-12-29 12:00:00', '2014-12-29 18:00:00',
                   '2014-12-30 00:00:00', '2014-12-30 06:00:00',
                   '2014-12-30 12:00:00', '2014-12-30 18:00:00',
                   '2014-12-31 00:00:00', '2014-12-31 06:00:00',
                   '2014-12-31 12:00:00', '2014-12-31 18:00:00'],
                  dtype='datetime64[ns]', name='time', length=2920, freq=None))

• Attributes: (5)

  Conventions :

  COARDS

  title :

  4x daily NMC reanalysis (1948)

  description :

  Data is from NMC initialized reanalysis (4x/day). These are the 0.9950 sigma level values.

  platform :

  Model

  references :

  http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanalysis.html

# create a BatchGenerator with caching enabled
gen = xbatcher.BatchGenerator(ds, input_dims={'lat': 10, 'lon': 10}, cache=cache)

Performance Comparison

Let’s compare the performance with and without caching:

import time


def time_iteration(gen):
    start = time.time()
    for batch in gen:
        pass
    end = time.time()
    return end - start

directory = f'{tempfile.mkdtemp()}/xbatcher-cache'
cache = zarr.storage.DirectoryStore(directory)

# Without cache
gen_no_cache = xbatcher.BatchGenerator(ds, input_dims={'lat': 10, 'lon': 10})
time_no_cache = time_iteration(gen_no_cache)
print(f'Time without cache: {time_no_cache:.2f} seconds')

Time without cache: 0.14 seconds

# With cache
gen_with_cache = xbatcher.BatchGenerator(
    ds, input_dims={'lat': 10, 'lon': 10}, cache=cache
)
time_first_run = time_iteration(gen_with_cache)
print(f'Time with cache (first run): {time_first_run:.2f} seconds')


time_second_run = time_iteration(gen_with_cache)
print(f'Time with cache (second run): {time_second_run:.2f} seconds')

Time with cache (first run): 0.18 seconds

/home/docs/checkouts/readthedocs.org/user_builds/xbatcher/conda/latest/lib/python3.10/site-packages/xarray/core/dataset.py:2292: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]
/home/docs/checkouts/readthedocs.org/user_builds/xbatcher/conda/latest/lib/python3.10/site-packages/xarray/core/dataset.py:2292: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]
/home/docs/checkouts/readthedocs.org/user_builds/xbatcher/conda/latest/lib/python3.10/site-packages/xarray/core/dataset.py:2292: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]
/home/docs/checkouts/readthedocs.org/user_builds/xbatcher/conda/latest/lib/python3.10/site-packages/xarray/core/dataset.py:2292: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]
/home/docs/checkouts/readthedocs.org/user_builds/xbatcher/conda/latest/lib/python3.10/site-packages/xarray/core/dataset.py:2292: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]
/home/docs/checkouts/readthedocs.org/user_builds/xbatcher/conda/latest/lib/python3.10/site-packages/xarray/core/dataset.py:2292: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]
/home/docs/checkouts/readthedocs.org/user_builds/xbatcher/conda/latest/lib/python3.10/site-packages/xarray/core/dataset.py:2292: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]
/home/docs/checkouts/readthedocs.org/user_builds/xbatcher/conda/latest/lib/python3.10/site-packages/xarray/core/dataset.py:2292: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]
/home/docs/checkouts/readthedocs.org/user_builds/xbatcher/conda/latest/lib/python3.10/site-packages/xarray/core/dataset.py:2292: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]
/home/docs/checkouts/readthedocs.org/user_builds/xbatcher/conda/latest/lib/python3.10/site-packages/xarray/core/dataset.py:2292: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]

Time with cache (second run): 0.07 seconds

You should see that the second run with cache is significantly faster than both the first run and the run without cache.

Advanced Usage

Custom Cache Preprocessing

You can also specify a custom preprocessing function to be applied to batches before they are cached:

# create a cache using Zarr's DirectoryStore
directory = f'{tempfile.mkdtemp()}/xbatcher-cache'
cache = zarr.storage.DirectoryStore(directory)


def preprocess_batch(batch):
    # example: add a new variable to each batch
    batch['new_var'] = batch['air'] * 2
    return batch


gen_with_preprocess = xbatcher.BatchGenerator(
    ds,
    input_dims={'lat': 10, 'lon': 10},
    cache=cache,
    cache_preprocess=preprocess_batch,
)

# Now, each cached batch will include the 'new_var' variable
for batch in gen_with_preprocess:
    print(batch)
    break

<xarray.Dataset> Size: 5MB
Dimensions:  (lat: 10, time: 2920, lon: 10)
Coordinates:
  * lat      (lat) float32 40B 75.0 72.5 70.0 67.5 65.0 62.5 60.0 57.5 55.0 52.5
  * lon      (lon) float32 40B 200.0 202.5 205.0 207.5 ... 217.5 220.0 222.5
  * time     (time) datetime64[ns] 23kB 2013-01-01 ... 2014-12-31T18:00:00
Data variables:
    air      (time, lat, lon) float64 2MB 241.2 242.5 243.5 ... 281.8 282.3
    new_var  (time, lat, lon) float64 2MB 482.4 485.0 487.0 ... 563.6 564.6
Attributes:
    Conventions:  COARDS
    title:        4x daily NMC reanalysis (1948)
    description:  Data is from NMC initialized reanalysis\n(4x/day).  These a...
    platform:     Model
    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...

/home/docs/checkouts/readthedocs.org/user_builds/xbatcher/conda/latest/lib/python3.10/site-packages/xarray/core/dataset.py:2292: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]

Using Different Storage Backends

While we’ve been using a local directory for caching, you can use any dict-like that is compatible with zarr. For example, you could use an S3 bucket as the cache storage backend:

import s3fs
import zarr

# Set up S3 filesystem (you'll need appropriate credentials)
s3 = s3fs.S3FileSystem(anon=False)
cache = s3.get_mapper('s3://my-bucket/my-cache.zarr')

# Use this cache with BatchGenerator
gen_s3 = xbatcher.BatchGenerator(ds, input_dims={'lat': 10, 'lon': 10}, cache=cache)

Considerations and Best Practices

• Storage Space: Be mindful of the storage space required for your cache, especially when working with large datasets.

• Cache Invalidation: The current implementation doesn’t handle cache invalidation. If your source data changes, you’ll need to manually clear or update the cache.

• Performance Tradeoffs: While caching can significantly speed up repeated access to the same data, the initial caching process may be slower than processing without a cache. Consider your use case to determine if caching is beneficial.

• Storage Backend: Choose a storage backend that’s appropriate for your use case. Local storage might be fastest for single-machine applications, while distributed or cloud storage might be necessary for cluster computing or cloud-based workflows.