reading gadget files with dask natively in yt¶

In this notebook, we show a simple application of the dask chunking experiments for particle reading within yt itself.

Several changes are necessary for this to work:

1. we need to modify how the BaseIOHandler reads in particle data from files.
2. modify the ParticleContainer so that our chunks can be effectively pickled
3. use yt's on-disk parameter storage option

1. modifying the BaseIOHandler¶

One way to use dask effectively is to send stack each chunk's read in a delayed dataframe. In this approach, we have a delayed dataframe for each particle type we're reading where each dataframe is partitioned by dask over the chunks we supply. Dask delayed dataframes don't require the size a priori, which simplifies building the dataframe but we do need to supply a metadata dictionary to designate the expected datatype for each column. The modified _read_particle_selection is not yet on a committed branch, so here's the function:

def _read_particle_selection(self, chunks, selector, fields):
        rv = {}
        ind = {}
        # We first need a set of masks for each particle type
        ptf = defaultdict(list)  # ptype -> on-disk fields to read
        fsize = defaultdict(lambda: 0)  # ptype -> size of return value
        psize = defaultdict(lambda: 0)  # ptype -> particle count on disk
        field_maps = defaultdict(list)  # ptype -> fields (including unions)
        chunks = list(chunks)
        unions = self.ds.particle_unions
        # What we need is a mapping from particle types to return types
        for field in fields:
            ftype, fname = field
            fsize[field] = 0
            # We should add a check for p.fparticle_unions or something here
            if ftype in unions:
                for pt in unions[ftype]:
                    ptf[pt].append(fname)
                    field_maps[pt, fname].append(field)
            else:
                ptf[ftype].append(fname)
                field_maps[field].append(field)
        # Now we have our full listing

        # an experimental dask approach

        # build the meta dictionary for each chunk's dataframe so that empty chunks
        # don't cause problems.
        ptype_meta = {}
        for ptype, flds in ptf.items():
            meta_dict = {}
            for fld in flds:
                meta_dict[fld] = pd.Series([], dtype=np.float64)
            ptype_meta[ptype] = meta_dict

        # build the delayed chunk reader (still has parallel issues...)
        # particle_field_args = self._read_particle_field_args()
        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 = [
                # all these things need to be pickleable...
                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])

        # up to here, everything is delayed w dask, need to read into into memory across
        # chunks to return rv: 
        rv = {}
        for ptype in ptypes:
            for col in delayed_dfs[ptype].columns:
                rv[(ptype, col)] = delayed_dfs[ptype][col].values.compute()

        # but if we returned the delayed dataframes, could do things like this:
        #   delayed_dfs['PartType0'].Density.mean().compute()

        return rv

    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

2. modifying the ParticleContainer¶

A single chunk is composed of various particle continers, and there are several modifications we need to make to the ParticleContainer related to pickleability. So first off, the ParticleContainer on yt's master branch currently cannot be pickled as it inherits a __reduce__ method from the Dataset class that requires the class to be a registered method, which the ParticleContainer is not. The particleIndexRefactor branch at github.com/chrishavlin/yt does just that by overriding the __reduce__ method in ParticleContainer.

The other important bit in the particleIndexRefactor branch is that it adds an explicit base_selector argument to __init__ rather than pulling out the selector object from the base_region argument. The reason for this is that when the arguments are unpickled, the base_region object will not have an initialized index. And so when base_region.selector is accessed, it first triggers an index-read. Reading the index for every dask process is very slow. But since the majority of the selector objects can now be pickled directly, by adding it as an argument we can avoid that index read.

3. on-disk parameter storage option¶

When yt pickles a Dataset object, it uses a __reduce__ method that stores the object arguments used to build the Dataset and then re-initializes the dataset on unpickling. Those arguments are kept in the ParameterStore object, which under normal operation is an in-memory cache of parameters. Since the parameter store is in memory, it's not available to the separate dask tasks. But there is a yt config option that stores these parameters in an on-disk csv instead. Setting StoreParameterFiles = True in the yt config will allow the separate dask tasks to read the paramter file and reconstitute the dataset as needed.

Single Processor Read¶

Ok, so first let's use a single processor, in which case nothing changes in terms of user experience:

Let's check out some chunk info:

So we have 12 chunks for this dataset. Ok, let's read in a field:

parallel read¶

OK, so now if we spin up a dask client, the delayed dataframes that we build in _read_particle_selection will automatically be split up across our chunks! So let's spin up a client, using 4 workers and 3 threads per worker, so each chunk will get its own process:

and read in a couple different fields (so that we're not just reading the cached field):

That initial density read takes a bit of extra time (due to some internal dask initialization?) but we can see subsequent reads of other fields are much faster. The CPU time are fairly close to the single processor read, but we can see the wall time is 400-500 seconds slower for the parallel read, which can be attributed to the extra communication overhead between processes.

Here's a screenshot of the Task Stream during the Temperature read: