Open OriolAbril opened 5 years ago
rcParam for dask?
Some function to translate idata to dask
Some function to set dask rcParam -> also set if return dask graph or evaluate
Also, see summary func, it already contains the main calculations in one func.
For rhat, ess etc, we can add .compute based on rcparams
Possible reference: http://xarray.pydata.org/en/stable/examples/apply_ufunc_vectorize_1d.html
Currently, ArviZ does not try to parallelize its functions in order to improve performance. It would be great to combine the speed-ups from using Numba with parallel computation using Dask if desired. Citing xarray docs:
Both arguments about parallelization and about arrays not fitting in memory are relevant to ArviZ. Moreover, Dask may become a requirement for xarray, and writing our code to be dask compatible could save a lot of time if this eventually happens.
I propose to modify all functions using
wrap_xarray_ufunc
to have some kwargs passed toxr.apply_ufunc
so that if the datasets in the InferenceData object are dask arrays, its capabilities can be used to parallelize the code. In addition, allowing to pass kwargs also toxr.open_dataset
will allow to load InferenceData objects from netdcf files directly with dask (with chunks if desired for example).Below there are some experiments to see the possible benefits of this changes (and keep in mind I don't really know dask at all, I have read its tutorial and xarray docs on it):
which outputs:
As you can see, the speed-up is quite relevant directly with this vanilla solution. There is probably room for improvement using pure dask, however, I do not think this should be an immediate concern (or a concern at all, maybe an idea for next year's GSoC?), using dask via xarray should be simple enough and provide good speed-ups for large data.
Implementation ideas
As more or less said before, just summing up, we can use dask via xarray to handle parallelization with very little work involved (passing kwargs to
wrap_xarray_ufunc
and toxr.open_dataset
should be enough). However, some things must be taken into account:xr.apply_ufunc
core dims cannot be chunked, which means that chunking on chains or draws won't work with most of the functions..compute()
must be called on them. In xarray though, there is also the option of geting.values
attribute, which should also trigger the evaluation.I do not have a clear idea on how to pass these dask related kwargs for all relevant functions, for
ess
it can be quite straightforward, but for summary? Should there be an option to allow different dask kwargs for ess and for rhat? Also, should ArviZ return unevaluated dask espressions so the user can runcompute()
whenever he wants? Should there be an option to optionally compute things? For summary for instance it could be better to compute ess and rhat at the same time, and also, could summary work without evaluating the values? Or will it always trigger computation via.values
attribute?It would probably be great to eventually add some tests on InferenceData objects containing dask arrays too, but not too many to avoid segmentation fault like we would get if tests were not run on "eager" mode.