xarray.DataArray map_blocks failed to deserialize

[rsignell-usgs]

What happened

My xarray map_blocks call failed in pickling with a deserialization error:

distributed.protocol.core - CRITICAL - Failed to deserialize
...
TypeError: __init__() missing 1 required positional argument: 'code'

What you expected to happen:

The code should work (and it works if you close the client and just use Dask, not distributed).

Not-Quite-Minimal but Complete and Verifiable Example

import xarray as xr
from cartopy import crs as ccrs
import numpy as np
from dask.distributed import Client

client = Client()

nx = 10000
ny = 10000

x = (np.linspace(121940., 574180., nx))
y = (np.linspace(4250700., 4659150., ny))[::-1]

crs_from = ccrs.epsg(26917)

da = xr.DataArray(
    data=np.ones((ny,nx)),
    dims=["y", "x"],
    coords=dict(
        x=(["x"], x),
        y=(["y"], y))).chunk({'x':5120, 'y':5120})

crs_to = ccrs.PlateCarree()  

def xy_to_lonlat(da):
    x, y = np.meshgrid(da.x, da.y)
    ll = crs_to.transform_points(crs_from, x, y)
    lon = ll[:,:,0]
    lat = ll[:,:,1]
    da = da.assign_coords(dict(
        lon=(["y", "x"], lon),
        lat=(["y", "x"], lat)))
    return da

da2 = da.map_blocks(xy_to_lonlat).compute()

Anything else we need to know?:
May be related to dask/dask#8355 and/or #5495 ?

Environment:

• Dask version: 2021.11.0
• Python version: 3.8.10
• Operating System: Linux
• Install method (conda, pip, source): conda

[rsignell-usgs]

Just hopefully, I tried the reproducer code above again today with Dask 2021.11.2 and Python 3.9.7 and got the same error.

[ian-r-rose]

Thanks for bumping this @rsignell-usgs, this may also be related to #5001 . Are xarray and cartopy important for reproducing this?

[jrbourbeau]

Hrm, it looks like the issue is the crs_from = ccrs.epsg(26917) object, which is used in the map_blocks call isn't pickeable. For example, slightly modifying @rsignell-usgs code snippet to check that we can pickle crs_from

import xarray as xr
from cartopy import crs as ccrs
import numpy as np
from dask.distributed import Client

if __name__ == "__main__":

    client = Client()

    nx = 10000
    ny = 10000

    x = (np.linspace(121940., 574180., nx))
    y = (np.linspace(4250700., 4659150., ny))[::-1]

    crs_from = ccrs.epsg(26917)

    da = xr.DataArray(
        data=np.ones((ny,nx)),
        dims=["y", "x"],
        coords=dict(
            x=(["x"], x),
            y=(["y"], y))).chunk({'x':5120, 'y':5120})

    crs_to = ccrs.PlateCarree()  

    def xy_to_lonlat(da):
        x, y = np.meshgrid(da.x, da.y)
        ll = crs_to.transform_points(crs_from, x, y)
        lon = ll[:,:,0]
        lat = ll[:,:,1]
        da = da.assign_coords(dict(
            lon=(["y", "x"], lon),
            lat=(["y", "x"], lat)))
        return da
    
    # check that we can pickle crs_from
    from cloudpickle import loads, dumps
    print(f"{type(crs_from) = }")
    func_dumped = dumps(crs_from)
    print(f"{func_dumped = }")
    func_roundtrip = loads(func_dumped)

    # da2 = da.map_blocks(xy_to_lonlat).compute()

produces a similar error

Traceback (most recent call last):
  File "/Users/james/projects/dask/distributed/test-map-blocks.py", line 42, in <module>
    func_roundtrip = loads(func_dumped)
TypeError: __init__() missing 1 required positional argument: 'code'

[jrbourbeau]

@rsignell-usgs since crs_from is only used inside the xy_to_lonlat function, you could its definition inside xy_to_lonlat to help with pickling. I was able to run the following code snippet on my laptop successfully:

import xarray as xr
from cartopy import crs as ccrs
import numpy as np
from dask.distributed import Client

if __name__ == "__main__":

    client = Client()

    nx = 10000
    ny = 10000

    x = (np.linspace(121940., 574180., nx))
    y = (np.linspace(4250700., 4659150., ny))[::-1]

    da = xr.DataArray(
        data=np.ones((ny,nx)),
        dims=["y", "x"],
        coords=dict(
            x=(["x"], x),
            y=(["y"], y))).chunk({'x':5120, 'y':5120})

    crs_to = ccrs.PlateCarree()  

    def xy_to_lonlat(da):
        x, y = np.meshgrid(da.x, da.y)
        crs_from = ccrs.epsg(26917)
        ll = crs_to.transform_points(crs_from, x, y)
        lon = ll[:,:,0]
        lat = ll[:,:,1]
        da = da.assign_coords(dict(
            lon=(["y", "x"], lon),
            lat=(["y", "x"], lat)))
        return da

    da2 = da.map_blocks(xy_to_lonlat).compute()

[rsignell-usgs]

Brilliant! Gave me a fish and taught me a little more about fishing! Thanks @ian-r-rose and @jrbourbeau !