register_dataset_accessor and joblib Parallel

[aaschwanden]

I am defining a custom method using register_dataset_accessor and want to use it in combination with joblib.Parallel.

Libraries
xarray: 2024.1.1
joblib: 1.3.2

import xarray as xr
from joblib import Parallel, delayed

@xr.register_dataset_accessor('custom')
class CustomDatasetMethods:
    def __init__(self, xarray_obj):
        self._obj = xarray_obj

    def function(self):
        return self._obj.sum()

d = {'coords': {'x': {'dims': ('x',),
   'attrs': {'units': 'm',
    'standard_name': 'projection_x_coordinate',
    'long_name': 'x coordinate of projection'},
   'data': [-285007.5, -284767.5, -284527.5, -284287.5]},
  'y': {'dims': ('y',),
   'attrs': {'units': 'm',
    'standard_name': 'projection_y_coordinate',
    'long_name': 'y coordinate of projection'},
   'data': [-1001392.5, -1001632.5, -1001872.5, -1002112.5]}},
 'attrs': {'GDAL_AREA_OR_POINT': 'Area',
  'motion_coordinates': 'map'},
 'dims': {'y': 4, 'x': 4},
 'data_vars': {'vx': {'dims': ('y', 'x'),
   'attrs': {'units': 'm/y',
    'standard_name': 'x_velocity',
    'quality_cutoff': 7.0,
    'map_scale_corrected': 1,
    'grid_mapping': 'Polar_Stereographic'},
   'data': [[0.24609112739562988,
     0.22844453155994415,
     0.21066682040691376,
     0.19329878687858582],
    [0.27507367730140686,
     0.254960834980011,
     0.2345907837152481,
     0.2146659940481186],
    [0.29915371537208557,
     0.27727755904197693,
     0.25520917773246765,
     0.23379768431186676],
    [0.31636136770248413,
     0.29350340366363525,
     0.2707109749317169,
     0.24898700416088104]]},
  'vy': {'dims': ('y', 'x'),
   'attrs': {'units': 'm/y',
    'standard_name': 'y_velocity',
    'quality_cutoff': 7.0,
    'map_scale_corrected': 1,
    'grid_mapping': 'Polar_Stereographic'},
   'data': [[-0.7410032749176025,
     -0.7680140733718872,
     -0.7932140827178955,
     -0.8151463270187378],
    [-0.7151581645011902,
     -0.7432965636253357,
     -0.7698800563812256,
     -0.7932386994361877],
    [-0.6926486492156982,
     -0.721246063709259,
     -0.7484308481216431,
     -0.7723526358604431],
    [-0.6764388680458069,
     -0.7047024369239807,
     -0.7315648794174194,
     -0.7550129890441895]]}}}

ds = xr.Dataset.from_dict(d)

def process(n, ds):
    return ds.custom.function() + n

This works just fine:

process(1, ds)

<xarray.Dataset>
Dimensions:  ()
Data variables:
    vx       float64 5.053
    vy       float64 -10.94

However, using Parallel

    delayed(process)(n, ds)
            for n in range(10)
        )

results in

    752 try:
    753     if self.status == TASK_ERROR:
--> 754         raise self._result
    755     return self._result
    756 finally:

AttributeError: 'Dataset' object has no attribute 'custom'

So it seems that the registered custom method is not passed on to Parallel.
How can I get custom methods to work in the Parallel case?