Add Dask engine to dataset generation functions

[zmbc]

Add Dask engine to dataset generation functions

Description

• Category: feature
• JIRA issue: none

Successor to #349 -- using Dask instead of Modin. This simplifies things and hopefully gives us a shorter path to releasing distributed noising.

Testing

• all tests pass (pytest --runslow)

Added some tests that use Dask.

[aflaxman]

Can you also add something to the docs about how to install and use this, and a cautionary note that this is not important for the small data distributed with psp, but can be a big time saver for people working with the full dataset?

Also, I think I'm doing it wrong... without dask
E.g. without this I loaded Alabama in 1.29 hours, but with it, on a 32-core cluster node, it took 1.57 hours.

[zmbc]

@aflaxman How did you start your Dask cluster? And which dataset did you load?

[aflaxman]

I didn't start one! I just used

        df = psp.generate_decennial_census(source=full_data_path, state=name,
                                           config=my_config, verbose=False, engine='dask')

with a path to the full data that you sent me yesterday.

[zmbc]

I have a feeling you might be starting Dask with as many workers as there are physical CPUs on your cluster node, and then it is thrashing.

I'm in the process of testing this branch on the case study.

[zmbc]

Note: running into a fiddly bug with dtypes on a few columns. I don't think it will be overly complicated to fix.

[zmbc]

Blocked by #405 -- Dask currently can't handle writing our strange dtypes to Parquet files.

[aflaxman]

I didn't start one! I just used

        df = psp.generate_decennial_census(source=full_data_path, state=name,
                                           config=my_config, verbose=False, engine='dask')

with a path to the full data that you sent me yesterday.

When I added this code block (on a node with the srun from the commented line, after installing the packages from the other commented lines), I was able to load the full USA data in 40 minutes:

# pip install "dask[distributed]" --upgrade
# pip install --upgrade pyarrow

# srun -t 14-00:00:00 --mem=100G -c 32 -A proj_simscience -p long.q --pty bash

from dask.distributed import LocalCluster
cluster = LocalCluster(n_workers=10, threads_per_worker=1)

[rmudambi]

@zmbc can you target a (new) release branch rather than main? We probably want to release all of your PRs together as one release, and merging to main requires we release.