Can you pass Future objects to apply_ufunc

[alessioarena]

Hi,

I have an array to pass to a function mapped using apply_ufunc across a DataArray.
This is relatively small, but gets replicated for every dask task and end up filling up the scheduler memory.

To address that I tried to pre-scatter the array to the dask cluster, and therefore provide the Future to apply_ufunc instead of the actual array.
This is done using the kwargs argument, but it seems the computation is failing and I don't really know why that is the case.
My gut feeling is that apply_ufunc recognises that I'm passing a dask object and tries to do something with that.

Here is a simplified code of what I'm trying to do:

    problem_array = ....

    current_client = dask.distributed.get_client()
    future_array = current_client.scatter(problem_array, broadcast=True)
    
    # update kwargs with the array
    kwargs.update({'array':future_array, 'freq':freq, 'aggregation':aggregation})
          
    # run _proc_daily as a ufunc
    resampled = _xr.apply_ufunc(
        _proc_daily,
        data,
        input_core_dims=[[apply_axis]],
        output_core_dims=[[apply_axis]],
        exclude_dims=set([apply_axis]),
        vectorize=True,
        kwargs=kwargs,
        dask='parallelized',
        output_dtypes="float64",
        dask_gufunc_kwargs={'output_sizes': {apply_axis: len(problem_array)}}
    )

and the function I'm applying:

def _proc_daily(data, array, freq, aggregation, **kwargs):
    array = array.result()
    ...
    return

Unfortunately it seems that using "parallelized" causes exceptions not to be raised back, and the scheduler keeps resubmitting the failed tasks over and over (I did check the various dask configurations in that respect)

Did anyone ever tried to do something like this?

[crusaderky]

Closing as duplicate of #6803