Repository: pydata/xarray

Issue:
CI Failure in Xarray test suite post-Dask tokenization update
### What is your issue?

Recent changes in Dask's tokenization process (https://github.com/dask/dask/pull/10876) seem to have introduced unexpected behavior in Xarray's test suite. This has led to CI failures, specifically in tests related to tokenization.


- https://github.com/pydata/xarray/actions/runs/8069874717/job/22045898877

```python
---------- coverage: platform linux, python 3.12.2-final-0 -----------
Coverage XML written to file coverage.xml

=========================== short test summary info ============================
FAILED xarray/tests/test_dask.py::test_token_identical[obj0-<lambda>1] - AssertionError: assert 'bbd9679bdaf2...d3db65e29a72d' == '6352792990cf...e8004a9055314'
  
  - 6352792990cfe23adb7e8004a9055314
  + bbd9679bdaf284c371cd3db65e29a72d
FAILED xarray/tests/test_dask.py::test_token_identical[obj0-<lambda>2] - AssertionError: assert 'bbd9679bdaf2...d3db65e29a72d' == '6352792990cf...e8004a9055314'
  
  - 6352792990cfe23adb7e8004a9055314
  + bbd9679bdaf284c371cd3db65e29a72d
FAILED xarray/tests/test_dask.py::test_token_identical[obj1-<lambda>1] - AssertionError: assert 'c520b8516da8...0e9e0d02b79d0' == '9e2ab1c44990...6ac737226fa02'
  
  - 9e2ab1c44990adb4fb76ac737226fa02
  + c520b8516da8b6a98c10e9e0d02b79d0
FAILED xarray/tests/test_dask.py::test_token_identical[obj1-<lambda>2] - AssertionError: assert 'c520b8516da8...0e9e0d02b79d0' == '9e2ab1c44990...6ac737226fa02'
  
  - 9e2ab1c44990adb4fb76ac737226fa02
  + c520b8516da8b6a98c10e9e0d02b79d0
= 4 failed, 16293 passed, [628](https://github.com/pydata/xarray/actions/runs/8069874717/job/22045898877#step:9:629) skipped, 90 xfailed, 71 xpassed, 213 warnings in 472.07s (0:07:52) =
Error: Process completed with exit code 1.
```


previously, the following code snippet would pass, verifying the consistency of tokenization in Xarray objects:

```python
In [1]: import xarray as xr, numpy as np

In [2]: def make_da():
   ...:     da = xr.DataArray(
   ...:         np.ones((10, 20)),
   ...:         dims=["x", "y"],
   ...:         coords={"x": np.arange(10), "y": np.arange(100, 120)},
   ...:         name="a",
   ...:     ).chunk({"x": 4, "y": 5})
   ...:     da.x.attrs["long_name"] = "x"
   ...:     da.attrs["test"] = "test"
   ...:     da.coords["c2"] = 0.5
   ...:     da.coords["ndcoord"] = da.x * 2
   ...:     da.coords["cxy"] = (da.x * da.y).chunk({"x": 4, "y": 5})
   ...: 
   ...:     return da
   ...: 

In [3]: da = make_da()

In [4]: import dask.base

In [5]: assert dask.base.tokenize(da) == dask.base.tokenize(da.copy(deep=False))

In [6]: assert dask.base.tokenize(da) == dask.base.tokenize(da.copy(deep=True))

In [9]: dask.__version__
Out[9]: '2023.3.0'
```

However, post-update in Dask version '2024.2.1', the same code fails:

```python
In [55]: 
    ...: def make_da():
    ...:     da = xr.DataArray(
    ...:         np.ones((10, 20)),
    ...:         dims=["x", "y"],
    ...:         coords={"x": np.arange(10), "y": np.arange(100, 120)},
    ...:         name="a",
    ...:     ).chunk({"x": 4, "y": 5})
    ...:     da.x.attrs["long_name"] = "x"
    ...:     da.attrs["test"] = "test"
    ...:     da.coords["c2"] = 0.5
    ...:     da.coords["ndcoord"] = da.x * 2
    ...:     da.coords["cxy"] = (da.x * da.y).chunk({"x": 4, "y": 5})
    ...: 
    ...:     return da
    ...: 

In [56]: da = make_da()
```

```python
In [57]: assert dask.base.tokenize(da) == dask.base.tokenize(da.copy(deep=False))
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[57], line 1
----> 1 assert dask.base.tokenize(da) == dask.base.tokenize(da.copy(deep=False))

AssertionError: 

In [58]: dask.base.tokenize(da)
Out[58]: 'bbd9679bdaf284c371cd3db65e29a72d'

In [59]: dask.base.tokenize(da.copy(deep=False))
Out[59]: '6352792990cfe23adb7e8004a9055314'

In [61]: dask.__version__
Out[61]: '2024.2.1'
```

additionally, a deeper dive into `dask.base.normalize_token()` across the two Dask versions revealed that the latest version includes additional state or metadata in tokenization that was not present in earlier versions. 


- old version 
```python
In [29]: dask.base.normalize_token((type(da), da._variable, da._coords, da._name))
Out[29]: 
('tuple',
 [xarray.core.dataarray.DataArray,
  ('tuple',
   [xarray.core.variable.Variable,
    ('tuple', ['x', 'y']),
    'xarray-<this-array>-14cc91345e4b75c769b9032d473f6f6e',
    ('list', [('tuple', ['test', 'test'])])]),
  ('list',
   [('tuple',
     ['c2',
      ('tuple',
       [xarray.core.variable.Variable,
        ('tuple', []),
        (0.5, dtype('float64')),
        ('list', [])])]),
    ('tuple',
     ['cxy',
      ('tuple',
       [xarray.core.variable.Variable,
        ('tuple', ['x', 'y']),
        'xarray-<this-array>-8e98950eca22c69d304f0a48bc6c2df9',
        ('list', [])])]),
    ('tuple',
     ['ndcoord',
      ('tuple',
       [xarray.core.variable.Variable,
        ('tuple', ['x']),
        'xarray-ndcoord-82411ea5e080aa9b9f554554befc2f39',
        ('list', [])])]),
    ('tuple',
     ['x',
      ('tuple',
       [xarray.core.variable.IndexVariable,
        ('tuple', ['x']),
        ['x',
         ('603944b9792513fa0c686bb494a66d96c667f879',
          dtype('int64'),
          (10,),
          (8,))],
        ('list', [('tuple', ['long_name', 'x'])])])]),
    ('tuple',
     ['y',
      ('tuple',
       [xarray.core.variable.IndexVariable,
        ('tuple', ['y']),
        ['y',
         ('fc411db876ae0f4734dac8b64152d5c6526a537a',
          dtype('int64'),
          (20,),
          (8,))],
        ('list', [])])])]),
  'a'])
```

- most recent version 

```python
In [44]: dask.base.normalize_token((type(da), da._variable, da._coords, da._name))
Out[44]: 
('tuple',
 [('7b61e7593a274e48', []),
  ('tuple',
   [('215b115b265c420c', []),
    ('tuple', ['x', 'y']),
    'xarray-<this-array>-980383b18aab94069bdb02e9e0956184',
    ('dict', [('tuple', ['test', 'test'])])]),
  ('dict',
   [('tuple',
     ['c2',
      ('tuple',
       [('__seen', 2),
        ('tuple', []),
        ('6825817183edbca7', ['48cb5e118059da42']),
        ('dict', [])])]),
    ('tuple',
     ['cxy',
      ('tuple',
       [('__seen', 2),
        ('tuple', ['x', 'y']),
        'xarray-<this-array>-6babb4e95665a53f34a3e337129d54b5',
        ('dict', [])])]),
    ('tuple',
     ['ndcoord',
      ('tuple',
       [('__seen', 2),
        ('tuple', ['x']),
        'xarray-ndcoord-8636fac37e5e6f4401eab2aef399f402',
        ('dict', [])])]),
    ('tuple',
     ['x',
      ('tuple',
       [('abc1995cae8530ae', []),
        ('tuple', ['x']),
        ['x', ('99b2df4006e7d28a', ['04673d65c892b5ba'])],
        ('dict', [('tuple', ['long_name', 'x'])])])]),
    ('tuple',
     ['y',
      ('tuple',
       [('__seen', 25),
        ('tuple', ['y']),
        ['y', ('88974ea603e15c49', ['a6c0f2053e85c87e'])],
        ('dict', [])])])]),
  'a'])
```


Cc @dcherian / @crusaderky for visibility 

Discussion:

Patch:
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
index c00fe1a9..aeb6b221 100644
--- a/xarray/core/dataarray.py
+++ b/xarray/core/dataarray.py
@@ -1070,7 +1070,7 @@ class DataArray(
         dataset[self.name] = self.variable
         return dataset
 
-    def __dask_tokenize__(self):
+    def __dask_tokenize__(self) -> object:
         from dask.base import normalize_token
 
         return normalize_token((type(self), self._variable, self._coords, self._name))
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
index 884e302b..e1fd9e02 100644
--- a/xarray/core/dataset.py
+++ b/xarray/core/dataset.py
@@ -694,7 +694,7 @@ class Dataset(
             data_vars, coords
         )
 
-        self._attrs = dict(attrs) if attrs is not None else None
+        self._attrs = dict(attrs) if attrs else None
         self._close = None
         self._encoding = None
         self._variables = variables
@@ -739,7 +739,7 @@ class Dataset(
 
     @attrs.setter
     def attrs(self, value: Mapping[Any, Any]) -> None:
-        self._attrs = dict(value)
+        self._attrs = dict(value) if value else None
 
     @property
     def encoding(self) -> dict[Any, Any]:
@@ -856,11 +856,11 @@ class Dataset(
 
         return self
 
-    def __dask_tokenize__(self):
+    def __dask_tokenize__(self) -> object:
         from dask.base import normalize_token
 
         return normalize_token(
-            (type(self), self._variables, self._coord_names, self._attrs)
+            (type(self), self._variables, self._coord_names, self._attrs or None)
         )
 
     def __dask_graph__(self):
diff --git a/xarray/core/variable.py b/xarray/core/variable.py
index cd0c022d..315c4636 100644
--- a/xarray/core/variable.py
+++ b/xarray/core/variable.py
@@ -2592,11 +2592,13 @@ class IndexVariable(Variable):
         if not isinstance(self._data, PandasIndexingAdapter):
             self._data = PandasIndexingAdapter(self._data)
 
-    def __dask_tokenize__(self):
+    def __dask_tokenize__(self) -> object:
         from dask.base import normalize_token
 
         # Don't waste time converting pd.Index to np.ndarray
-        return normalize_token((type(self), self._dims, self._data.array, self._attrs))
+        return normalize_token(
+            (type(self), self._dims, self._data.array, self._attrs or None)
+        )
 
     def load(self):
         # data is already loaded into memory for IndexVariable
diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py
index 29722690..fd209bc2 100644
--- a/xarray/namedarray/core.py
+++ b/xarray/namedarray/core.py
@@ -511,7 +511,7 @@ class NamedArray(NamedArrayAggregations, Generic[_ShapeType_co, _DType_co]):
 
     @attrs.setter
     def attrs(self, value: Mapping[Any, Any]) -> None:
-        self._attrs = dict(value)
+        self._attrs = dict(value) if value else None
 
     def _check_shape(self, new_data: duckarray[Any, _DType_co]) -> None:
         if new_data.shape != self.shape:
@@ -570,13 +570,12 @@ class NamedArray(NamedArrayAggregations, Generic[_ShapeType_co, _DType_co]):
             return real(self)
         return self._new(data=self._data.real)
 
-    def __dask_tokenize__(self) -> Hashable:
+    def __dask_tokenize__(self) -> object:
         # Use v.data, instead of v._data, in order to cope with the wrappers
         # around NetCDF and the like
         from dask.base import normalize_token
 
-        s, d, a, attrs = type(self), self._dims, self.data, self.attrs
-        return normalize_token((s, d, a, attrs))  # type: ignore[no-any-return]
+        return normalize_token((type(self), self._dims, self.data, self._attrs or None))
 
     def __dask_graph__(self) -> Graph | None:
         if is_duck_dask_array(self._data):
diff --git a/xarray/namedarray/utils.py b/xarray/namedarray/utils.py
index 0326a617..b82a80b5 100644
--- a/xarray/namedarray/utils.py
+++ b/xarray/namedarray/utils.py
@@ -218,7 +218,7 @@ class ReprObject:
     def __hash__(self) -> int:
         return hash((type(self), self._value))
 
-    def __dask_tokenize__(self) -> Hashable:
+    def __dask_tokenize__(self) -> object:
         from dask.base import normalize_token
 
-        return normalize_token((type(self), self._value))  # type: ignore[no-any-return]
+        return normalize_token((type(self), self._value))


Test patch:
diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py
index 07bf773c..517fc0c2 100644
--- a/xarray/tests/test_dask.py
+++ b/xarray/tests/test_dask.py
@@ -299,17 +299,6 @@ class TestVariable(DaskTestCase):
         self.assertLazyAndAllClose(u + 1, v)
         self.assertLazyAndAllClose(u + 1, v2)
 
-    def test_tokenize_empty_attrs(self) -> None:
-        # Issue #6970
-        assert self.eager_var._attrs is None
-        expected = dask.base.tokenize(self.eager_var)
-        assert self.eager_var.attrs == self.eager_var._attrs == {}
-        assert (
-            expected
-            == dask.base.tokenize(self.eager_var)
-            == dask.base.tokenize(self.lazy_var.compute())
-        )
-
     @requires_pint
     def test_tokenize_duck_dask_array(self):
         import pint
@@ -1573,6 +1562,30 @@ def test_token_identical(obj, transform):
     )
 
 
+@pytest.mark.parametrize(
+    "obj",
+    [
+        make_ds(),  # Dataset
+        make_ds().variables["c2"],  # Variable
+        make_ds().variables["x"],  # IndexVariable
+    ],
+)
+def test_tokenize_empty_attrs(obj):
+    """Issues #6970 and #8788"""
+    obj.attrs = {}
+    assert obj._attrs is None
+    a = dask.base.tokenize(obj)
+
+    assert obj.attrs == {}
+    assert obj._attrs == {}  # attrs getter changed None to dict
+    b = dask.base.tokenize(obj)
+    assert a == b
+
+    obj2 = obj.copy()
+    c = dask.base.tokenize(obj2)
+    assert a == c
+
+
 def test_recursive_token():
     """Test that tokenization is invoked recursively, and doesn't just rely on the
     output of str()
diff --git a/xarray/tests/test_sparse.py b/xarray/tests/test_sparse.py
index 289149bd..09c12818 100644
--- a/xarray/tests/test_sparse.py
+++ b/xarray/tests/test_sparse.py
@@ -878,10 +878,6 @@ def test_dask_token():
     import dask
 
     s = sparse.COO.from_numpy(np.array([0, 0, 1, 2]))
-
-    # https://github.com/pydata/sparse/issues/300
-    s.__dask_tokenize__ = lambda: dask.base.normalize_token(s.__dict__)
-
     a = DataArray(s)
     t1 = dask.base.tokenize(a)
     t2 = dask.base.tokenize(a)