Results from initial explorations in leveraging dask in yt

There are a number of possible ways for yt to leverage dask in order to simplify existing parallel and lazy operations but thus far, we have focused on two related areas in which dask may prove particularly useful:

1. data IO: reading data from disk
2. processing data after reading

Central to both of these areas is the chunk iterator that yt uses to index a dataset. The exact implementation depends both on the frontend and datatype, but in general a single chunk is a range of data indeces that may span datafiles or a subset of a single datafile. Thus, data can be read in by separate processes in parallel or sequentially on a single processor by iterating and aggregating over chunks separately. In the case of IO, data can also be subselected by selector objects (e.g., slices, regions) before aggregation across chunks into yt arrays. And in the case of data processing, some computations such as finding statistical measures of a data field (max, min, weighted means) can also process chunks separately and return intermediate results that are then aggregated.

In terms of dask, one straightforward conceptual approach with fairly minimal refactoring of yt to use the dask.delayed() decorator to construct graphs to execute across chunks. In pseudocode this could look like:

delayed_chunks = [dask.delayed(chunk_processor)(ch, *args) for ch in chunks]

where chunk_processor is some function that we want to apply to each chunk and *args are the arguments to that function. This list of delayed_chunks can be strung together with any number of subsequent operations to create a dask Task Graph. For example, if we want to find the minimum value of a field across chunks, we might construct a graph that first reads a chunk and then finds the min of each chunk:

delayed_chunks = [dask.delayed(read_one_chunk)(ch,field) for ch in chunks]
delayed_mins =  [dask.delayed(find_chunk_min)(dc) for dc in delayed_chunks]
min_val = min(dask.compute(*delayed_mins))

The final line is where the task graph actually executes, before which dask is only constructing a representation of the tasks. When we call dask.compute(), dask will distribute the tasks to the dask.distributed.Client if one is active, otherwise the execution continues sequentially on a single processor.

In the following, we first dive into data IO in more detail and describe a prototype particle reader that uses delayed dask.dataframes to read fields from a Gadget dataset. We then discuss calculating derived quantities of yt Data Objects using dask.

1. data IO:

yt reads particle and grid-based data by iterating across the chunks, with frontend-specific IO functions. For gridded data, each frontend implements a _read_fluid_selection (e.g., yt.frontend.amrvac.AMRVACIOHandler._read_fluid_selection) that iterates over chunks and returns a flat dictionary with numpy arrays concatenated across each chunk. The particle data, frontends must implement a similar function, _read_particle_fields, that typically gets invoked within the BaseIOHanlder._read_particle_selection function. In both cases, the read functions accept the chunks iterator, the fields to read and a selector object:

def _read_particle_fields(self, chunks, ptf, selector):

def _read_fluid_selection(self, chunks, selector, fields, size):

On reading a single chunk, the selector is applied so that we only return selected data from each chunk.

In order to construct a dask.delayed() Task Graph for IO, there are a number of changes required. First, dask uses pickle to serialize the functions and arguments that get distributed to workers, so all arguments to the _read functions above must be pickleable. A recent PR (PR #2934) added __getstate__ and __setstate__ methods to the geometric Cython selector objects, but the chunks can be more challenging. We touch on this further in the following section on the Particle IO prototype.

Second, the _read_particle_fields and _read_fluid_selection typically iterate over chunks internally, but it is more straightforward to construct a chunk-parallel dask process with a single-chunk read function. In the following Particle IO prototype, we avoid a large scale refactor by calling _read_particle_fields repeatedly with a list of a single-chunk, i.e.,_read_particle_fields([chunk],ptf,selector), but separating the chunk iteration and concatenation into a single chunk read function may improve maintainability.

Finally, one question on IO is what to return? The present read routines return a dictionary of numpy arrays. These arrays are typically pre-allocated though the recent PR #2416 removes particle pre-allocation. But in terms of dask, we could return a distritubed dask.array or dask.dataframe, allowing subsequent functions to compute in parallel. The dask.array needs to know chunk sizes a priori, which presents some difficulty. A dask.dataframe does not need to know the chunk sizes, just the expected data type for each column. In the following prototype, we construct a Task Graph that constructs a dask.dataframe from delayed chunk reads. In order to return the dictionary of flat numpy arrays expected by _read_particle_fields, prototype reduces the distributed dataframes to standard numpy arrays but in principle we could return the dask.dataframes, allowing subsequent calculations to leverage distributed calculations by using the distributed dataframes directly.

a particle IO prototype

The dxl yt-dask-experiment repo includes some initial chunking experiments with dask in the dask_chunking directory. That folder primarily experiments with using dask to do some manual IO on a gadget dataset. In this experiment, the gadget file reading is split into a series of dask-delayed operations that include file IO and selector application. The notebook here gives an overview of that attempt, so I won't go into details here, but will instead focus on an attempt to use dask for file IO natively within yt.

The focus of this attempt falls within BaseIOHandler._read_particle_selection(). In this function, yt currently pre-allocates arrays for all the particle fields that we're reading and then read in the fields for each chunk using the self._read_particle_fields(chunks, ptf, selector) generator.

One conceptually straightforward way to leverage dask is by building a dask.dataframe from delayed objects. Dask dataframes, unlike dask arrays, do not need to know the total size a priori, but we do need to specify a metadata dictionary to declare the column names and datatypes.

So to build a dask.dataframe, we can do the following (from within BaseIOHandler._read_particle_selection()):

        ptypes = list(ptf.keys())
        delayed_dfs = {}
        for ptype in ptypes:
            # build a dataframe from delayed for each particle type
            this_ptf = {ptype: ptf[ptype]}
            delayed_chunks = [               
                dask.delayed(self._read_single_ptype)(
                    ch, this_ptf, selector, ptype_meta[ptype]
                )
                for ch in chunks
            ]
            delayed_dfs[ptype] = ddf.from_delayed(delayed_chunks, meta=ptype_meta[ptype])

In this snippet, we build up a dictionary of dask dataframes organized by particle type (ptypes). First, let's focus on the innermost loop in the snippet, the actual application of the dask.delayed decorator:

dask.delayed(self._read_single_ptype)(
                    ch, this_ptf, selector, ptype_meta[ptype]
                )

This decorator wraps a _read_single_ptype function, which takes a single chunk, a single particle type dictionary (with multiple fields) and a column-datatype metadata dictionary and returns a normal pandas dataframe:

    def _read_single_ptype(self, chunk, ptf, selector, meta_dict):
        # read a single chunk and single particle type into a pandas dataframe so that 
        # we can use dask.dataframe.from_delayed! fields within a particle type should
        # have the same length?

        chunk_results = pd.DataFrame(meta_dict)
        # each particle type could be a different dataframe...
        for field_r, vals in self._read_particle_fields([chunk], ptf, selector):
            chunk_results[field_r[1]] = vals        
            

        return chunk_results

We can see this function just calls our normal self._read_particle_fields and pulls out the field values as usual. We are storing the fields in columns of our single chunk and single particle type in a pandas dataframe, chunk_results.

So for a single particle type, we build our delayed dataframe

            delayed_chunks = [               
                dask.delayed(self._read_single_ptype)(
                    ch, this_ptf, selector, ptype_meta[ptype]
                )
                for ch in chunks
            ]
            delayed_dfs[ptype] = ddf.from_delayed(delayed_chunks, meta=ptype_meta[ptype])

The delayed_dfs object is a dictionary with a delayed dataframe for each particle type. The reason we're organizing by particle type is the issue that different particle types may have a different number of records in a given chunk (otherwise we'd have to store a large number of null values). In order to return the expected in-memory dict that _read_particle_selection() should return, we can very easily pull all the records from across our chunks and particle typeswith:

        rv = {}
        for ptype in ptypes:
            for col in delayed_dfs[ptype].columns:
                rv[(ptype, col)] = delayed_dfs[ptype][col].values.compute()

So this is a fairly promising approach, particularly since the dataframes do not need to know the expected array size. And it does indeed work to read data from our particle front ends, with some modifications to the ParticleContainer class.

The notebook here shows the above approach in action, with notes on the required changes to the ParticleContainer class. To load actually execute in parallel, the only requirement is to spin up a dask Client:

from dask.distributed import Client 
import yt 
c = Client(threads_per_worker=3,n_workers=4)
ds = yt.load_sample("snapshot_033")
ad = ds.all_data()
d = ad[('PartType0','Density')

The dask.delayed and dask.compute will find and connect to the Client. If no client is present, the Task Graph will be executed sequentially. Here's a snapshot of the Dask Dashboard's Task Stream during a parallel particle read:

2. derived quantities

Calcluation of derived quantities in yt also uses the chunk iterator to return intermediate by-chunk results that are then aggregated. The base DerivedQuantity object's __call__ function is where the iteration occurs:

    def __call__(self, *args, **kwargs):
        """Calculate results for the derived quantity"""
        # create the index if it doesn't exist yet
        self.data_source.ds.index
        self.count_values(*args, **kwargs)
        chunks = self.data_source.chunks([], chunking_style="io")
        storage = {}
        for sto, ds in parallel_objects(chunks, -1, storage=storage):
            sto.result = self.process_chunk(ds, *args, **kwargs)
        # Now storage will have everything, and will be done via pickling, so
        # the units will be preserved.  (Credit to Nathan for this
        # idea/implementation.)
        values = [[] for i in range(self.num_vals)]
        for key in sorted(storage):
            for i in range(self.num_vals):
                values[i].append(storage[key][i])
        # These will be YTArrays
        values = [self.data_source.ds.arr(values[i]) for i in range(self.num_vals)]
        values = self.reduce_intermediate(values)
        return values

In principle, if a dask dataframe were to be returned by the chunk IO, then we could completely remove the consideration of intermediate values and reduction and simply use pandas-like operations to return values. More specifically, we could use dataframe aggregation directly, which specifies a by-chunk operation, a by-chunk aggregation operation and a chunk-aggregation operation. From the dask documentations, an example of manually computing a mean value for a dask dataframe, df, is

custom_mean = dd.Aggregation(
    'custom_mean',
    lambda s: (s.count(), s.sum()),
    lambda count, sum: (count.sum(), sum.sum()),
    lambda count, sum: sum / count,
)
df.groupby('g').agg(custom_mean)

So if we return dask dataframes from the IO, we could replace the derived quantities with custom Aggregation operations. We are currently working on a proof-of-concept prototype demonstrating this in action.

Working with unyt

yt 4.0 forward uses unyt to track and convert units – the base YTArray class is actually a model in unyt. If we are using dask for IO, there are situations where it may be advantageous to hold off computing the dask graph. For example, we may want to return a dask array to the user so that they can construct their own computations in parallel. This requires, however, some level of dask support in unyt.

In the notebook, working with unyt and dask, we demonstrate a limited prototype of a dask-unyt array. In this notebook, we create a custom dask collection by sublcassing the primary dask.array class and adding some unyt functionality. This custom class is handled automatically by the dask scheduler, so that if we have a large dask array with a dask client running and we create our new dask-unyt array, e.g.:

import dask.array as da
from dask.distributed import Client
client = Client(threads_per_worker=2, n_workers=2)

x = da.random.random((10000, 10000), chunks=(1000, 1000))
x_da = unyt_from_dask(x, unyt.m)

when we do operations like finding the minimum value across all the chunks:

x_da.min()

We are returned a standard unyt_array

unyt_array(3.03822589e-09, 'm')

that was calculated by processing each chunk of the array separately, as seen here in the Task Stream of the dask dashboard:

This notebook demonstrates a general and fairly straightforward way to build in dask support to unyt which can be used in conjuction with, for example, the prototype dask-enabled particle reader to return arrays with dask functionality preserved.

some final thoughts

In the above discussion, we've focused primarily on how to use dask within the yt chunking infrastructure. But it is also worth considering whether we can replace or simplify the chunking itself using dask. In the case of reading particle data, yt is primarily looping over datafiles, so it may be possible to read directly into a dask dataframe without the chunking architecture. This remains to be explored…