Documentation build fails with latest Dask release (2021.3.0)

[tomwhite]

>>>-------------------------------------------------------------------------
Exception in /home/runner/work/sgkit/sgkit/docs/getting_started.rst at block ending on line None
Specify :okexcept: as an option in the ipython:: block to suppress this message
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-55-0ac15a3c444c> in <module>
----> 1 sg.diversity(ds, merge=False)

~/work/sgkit/sgkit/sgkit/stats/popgen.py in diversity(ds, cohort_allele_count, merge)
     85         [1.        , 1.        ]])
     86     """
---> 87     ds = define_variable_if_absent(
     88         ds, variables.cohort_allele_count, cohort_allele_count, count_cohort_alleles
     89     )

~/work/sgkit/sgkit/sgkit/utils.py in define_variable_if_absent(ds, default_variable_name, variable_name, func)
    193             f"Variable '{variable_name}' with non-default name is missing and will not be automatically defined."
    194         )
--> 195     return func(ds)
    196 
    197 

~/work/sgkit/sgkit/sgkit/stats/aggregation.py in count_cohort_alleles(ds, call_allele_count, merge)
    303     AC, SC = da.asarray(ds.call_allele_count), da.asarray(ds.sample_cohort)
    304     n_cohorts = SC.max().compute() + 1  # 0-based indexing
--> 305     C = da.empty(n_cohorts, dtype=np.uint8)
    306 
    307     G = da.asarray(ds.call_genotype)

/opt/hostedtoolcache/Python/3.8.8/x64/lib/python3.8/site-packages/dask/array/wrap.py in wrap_func_shape_as_first_arg(func, *args, **kwargs)
     63         )
     64 
---> 65     parsed = _parse_wrap_args(func, args, kwargs, shape)
     66     shape = parsed["shape"]
     67     dtype = parsed["dtype"]

/opt/hostedtoolcache/Python/3.8.8/x64/lib/python3.8/site-packages/dask/array/wrap.py in _parse_wrap_args(func, args, kwargs, shape)
     33     dtype = np.dtype(dtype)
     34 
---> 35     chunks = normalize_chunks(chunks, shape, dtype=dtype)
     36 
     37     name = name or funcname(func) + "-" + tokenize(

/opt/hostedtoolcache/Python/3.8.8/x64/lib/python3.8/site-packages/dask/array/core.py in normalize_chunks(chunks, shape, limit, dtype, previous_chunks)
   2730 
   2731     if shape is not None:
-> 2732         chunks = tuple(c if c not in {None, -1} else s for c, s in zip(chunks, shape))
   2733 
   2734     if chunks and shape is not None:

/opt/hostedtoolcache/Python/3.8.8/x64/lib/python3.8/site-packages/dask/array/core.py in <genexpr>(.0)
   2730 
   2731     if shape is not None:
-> 2732         chunks = tuple(c if c not in {None, -1} else s for c, s in zip(chunks, shape))
   2733 
   2734     if chunks and shape is not None:

TypeError: unhashable type: 'DataArray'

[jeromekelleher]

Sigh...

[tomwhite]

Oddly, this is only happening with the documentation. That is, running the snippet of code in a python interpreter doesn't cause the problem:

import numpy as np
import sgkit as sg
import xarray as xr
ds = sg.simulate_genotype_call_dataset(n_variant=100, n_sample=50)

# Define a single cohort for all samples
ds["sample_cohort"] = xr.DataArray(np.full(ds.dims['samples'], 0), dims="samples")

# The diversity statistic is computed for every variant since no windows are defined
sg.diversity(ds, merge=False)

It turns out that when run as a part of the Sphinx build, the SC variable in count_cohort_alleles has the following form:

dask.array<array, shape=(50,), dtype=int64, chunksize=(50,), chunktype=xarray.DataArray>

whereas in the repl it is:

dask.array<array, shape=(50,), dtype=int64, chunksize=(50,), chunktype=numpy.ndarray>

And indeed, the ds.sample_cohort variable passed to da.asarray has type xarray.DataArray under Sphinx, and xarray.core.dataarray.DataArray in the repl. This, in turn, means that it is wrapped differently by Dask. Something changed between Dask 2021.2.0 and 2021.3.0 that surfaced the failure. (I'm not sure exactly what, but possibly dask/dask#6738.)

How does DataArray get a different type under Sphinx? It's because of this workaround:

# FIXME: Workaround for linking xarray module
# For some reason, intersphinx is not able to to link xarray objects.
# https://github.com/pydata/xarray/issues/4279
xarray.Dataset.__module__ = "xarray"
xarray.DataArray.__module__ = "xarray"

This can be fixed by not using the workaround for DataArray (Dataset is OK since it isn't wrapped in the same way using da.asarray).

I've also opened dask/dask#7335 for a longer-term fix.

[jeromekelleher]

Yeesh, sounds like you had a fun day working that out!