NEXRAD issue - IndexError index out of bounds

[rabernat]

• xradar version: 0.6.4
• xarray version: 2024.7.0
• Python version: 3.12.4
• Operating System: Linux

Description

I have found a puzzling bug that only comes up in certain situations with Dask

What I Did

import xradar
import xarray as xr
import pooch

# download and open a NEXRAD2 file from S3
url = "https://noaa-nexrad-level2.s3.amazonaws.com/2024/09/01/FOP1/FOP120240901_000347_V06"
local_file = pooch.retrieve(url, known_hash=None)
ds = xr.open_dataset(local_file, group="sweep_0", engine="nexradlevel2")

# create a chunked version 
dsc = ds.chunk()
# load one variable - IMPORTANT - skipping this step makes the next line work
dsc.DBZH.load()
# load the entire dataset
dsc.load()
# - > IndexError: index 140 is out of bounds for axis 0 with size 38
# try all the variables
for v in dsc:
    print(v)
    try:
        dsc[v].load()  # also fails with dsc!
        print("ok")
    except Exception as e:
        print(e)
# DBZH
# ok
# ZDR
# index 212 is out of bounds for axis 0 with size 212
# PHIDP
# index 140 is out of bounds for axis 0 with size 130
# RHOHV
# index 140 is out of bounds for axis 0 with size 22
# CCORH
# index 140 is out of bounds for axis 0 with size 73
# sweep_mode
# ok
# sweep_number
# ok
# prt_mode
# ok
# follow_mode
# ok
# sweep_fixed_angle
# ok

Possibly related to #180.

Experience tells me this has something to do with Dask task tokenization.