xarray.DataArray map_blocks failed to deserialize

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 https://github.com/dask/dask/issues/8355 and/or https://github.com/dask/distributed/issues/5495 ?

Environment:

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

[rsignell-usgs]

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

[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()