diff --git a/dask_cuda/explicit_comms/dataframe/shuffle.py b/dask_cuda/explicit_comms/dataframe/shuffle.py index ca69156dd..6b4150041 100644 --- a/dask_cuda/explicit_comms/dataframe/shuffle.py +++ b/dask_cuda/explicit_comms/dataframe/shuffle.py @@ -1,8 +1,11 @@ from __future__ import annotations import asyncio +import contextlib import functools import inspect +import threading +import uuid from collections import defaultdict from math import ceil from operator import getitem @@ -16,18 +19,55 @@ from dask.base import tokenize from dask.dataframe.core import DataFrame, Series, _concat as dd_concat, new_dd_object from dask.dataframe.shuffle import group_split_dispatch, hash_object_dispatch -from distributed import wait +from distributed import get_worker, wait from distributed.protocol import nested_deserialize, to_serialize from distributed.worker import Worker from .. import comms +try: + from tqdm import tqdm +except ImportError: + tqdm = lambda x: x + T = TypeVar("T") Proxify = Callable[[T], T] +_WORKER_CACHE = {} +_WORKER_CACHE_LOCK = threading.RLock() + + +@contextlib.contextmanager +def get_worker_cache(name): + with _WORKER_CACHE_LOCK: + yield _get_worker_cache(name) + + +def _get_worker_cache(name): + """Utility to get the `name` element of the cache + dictionary for the current worker. If executed + by anything other than a distributed Dask worker, + we will use the global `_WORKER_CACHE` variable. + """ + try: + worker = get_worker() + except ValueError: + # There is no dask.distributed worker. + # Assume client/worker are same process + global _WORKER_CACHE # pylint: disable=global-variable-not-assigned + if name not in _WORKER_CACHE: + _WORKER_CACHE[name] = {} + return _WORKER_CACHE[name] + if not hasattr(worker, "worker_cache"): + worker.worker_cache = {} + if name not in worker.worker_cache: + worker.worker_cache[name] = {} + return worker.worker_cache[name] + + def get_proxify(worker: Worker) -> Proxify: """Get function to proxify objects""" from dask_cuda.proxify_host_file import ProxifyHostFile @@ -328,6 +368,8 @@ async def shuffle_task( ignore_index: bool, num_rounds: int, batchsize: int, + parquet_dir: str | None, + final_task: bool, ) -> Dict[int, DataFrame]: """Explicit-comms shuffle task @@ -371,6 +413,8 @@ async def shuffle_task( assert stage.keys() == rank_to_inkeys[myrank] no_comm_postprocess = get_no_comm_postprocess(stage, num_rounds, batchsize, proxify) + fns = [] + append_files = True # Whether to keep files open between batches out_part_id_to_dataframe_list: Dict[int, List[DataFrame]] = defaultdict(list) for _ in range(num_rounds): partitions = create_partitions( @@ -386,6 +430,50 @@ async def shuffle_task( out_part_id_to_dataframe_list, ) + if parquet_dir: + import cudf + + out_part_ids = list(out_part_id_to_dataframe_list.keys()) + for out_part_id in out_part_ids: + if append_files: + writers = _get_worker_cache("writers") + try: + writer = writers[out_part_id] + except KeyError: + fn = f"{parquet_dir}/part.{out_part_id}.parquet" + fns.append(fn) + writer = cudf.io.parquet.ParquetWriter(fn, index=False) + writers[out_part_id] = writer + dfs = out_part_id_to_dataframe_list.pop(out_part_id) + dfs = [df for df in dfs if len(dfs) > 0] + for df in dfs: + writer.write_table(df) + del dfs + else: + dfs = out_part_id_to_dataframe_list.pop(out_part_id) + id = str(uuid.uuid4())[:8] + fn = f"{parquet_dir}/part.{out_part_id}.{id}.parquet" + fns.append(fn) + dfs = [df for df in dfs if len(dfs) > 0] + if len(dfs) > 1: + with cudf.io.parquet.ParquetWriter(fn, index=False) as writer: + for df in dfs: + writer.write_table(df) + elif dfs: + dfs[0].to_parquet(fn, index=False) + del dfs + await asyncio.sleep(0) + + if parquet_dir: + if append_files: + if final_task: + for out_part_id in list(writers.keys()): + writers.pop(out_part_id).close() + await asyncio.sleep(0) + del writers + return {i: fn for i, fn in enumerate(fns)} + return {i: fn for i, fn in enumerate(fns)} + # Finally, we concatenate the output dataframes into the final output partitions ret = {} while out_part_id_to_dataframe_list: @@ -518,6 +606,7 @@ def shuffle( ignore_index, num_rounds, batchsize, + True, ) wait(list(shuffle_result.values())) @@ -547,6 +636,89 @@ def shuffle( return ret +def shuffle_to_parquet( + full_df: DataFrame, + column_names: List[str], + parquet_dir: str, + npartitions: Optional[int] = None, + ignore_index: bool = False, + batchsize: int = 2, + pre_shuffle: Optional[int] = None, + overwrite: bool = False, +) -> None: + from dask_cuda.explicit_comms.dataframe.utils import ( + _clean_worker_storage, + _prepare_dir, + ) + + c = comms.default_comms() + + # Assume we are writing to local worker storage + if overwrite: + wait(c.client.run(_clean_worker_storage, parquet_dir)) + wait(c.client.run(_prepare_dir, parquet_dir)) + + # The ranks of the output workers + ranks = list(range(len(c.worker_addresses))) + + # By default, we preserve number of partitions + if npartitions is None: + npartitions = full_df.npartitions + + # Find the output partition IDs for each worker + div = npartitions // len(ranks) + rank_to_out_part_ids: Dict[int, Set[int]] = {} # rank -> set of partition id + for i, rank in enumerate(ranks): + rank_to_out_part_ids[rank] = set(range(div * i, div * (i + 1))) + for rank, i in zip(ranks, range(div * len(ranks), npartitions)): + rank_to_out_part_ids[rank].add(i) + + parts_per_batch = len(ranks) * batchsize + num_rounds = ceil(full_df.npartitions / parts_per_batch) + for stage in tqdm(range(num_rounds)): + offset = parts_per_batch * stage + df = full_df.partitions[offset : offset + parts_per_batch] + + # Execute pre-shuffle function on each batch + if callable(pre_shuffle): + df = pre_shuffle(df) + + df = df.persist() # Make sure optimizations are apply on the existing graph + wait([df]) # Make sure all keys has been materialized on workers + name = ( + "explicit-comms-shuffle-" + f"{tokenize(df, column_names, npartitions, ignore_index)}" + ) + + # Stage all keys of `df` on the workers and cancel them, which makes it possible + # for the shuffle to free memory as the partitions of `df` are consumed. + # See CommsContext.stage_keys() for a description of staging. + rank_to_inkeys = c.stage_keys(name=name, keys=df.__dask_keys__()) + max_num_inkeys = max(len(k) for k in rank_to_inkeys.values()) + c.client.cancel(df) + + # Run a shuffle task on each worker + shuffle_result = {} + for rank in ranks: + shuffle_result[rank] = c.submit( + c.worker_addresses[rank], + shuffle_task, + name, + rank_to_inkeys, + rank_to_out_part_ids, + column_names, + npartitions, + ignore_index, + 1, + max_num_inkeys, + parquet_dir, + stage == (num_rounds - 1), + ) + wait(list(shuffle_result.values())) + + return + + def _use_explicit_comms() -> bool: """Is explicit-comms and available?""" if dask.config.get("explicit-comms", False): diff --git a/dask_cuda/explicit_comms/dataframe/utils.py b/dask_cuda/explicit_comms/dataframe/utils.py new file mode 100644 index 000000000..9d1783b3e --- /dev/null +++ b/dask_cuda/explicit_comms/dataframe/utils.py @@ -0,0 +1,140 @@ +import glob +import os +import pickle +from collections import defaultdict + +from dask.blockwise import BlockIndex +from distributed import wait +from distributed.protocol import dask_deserialize, dask_serialize + +from dask_cuda.explicit_comms import comms + + +class LazyLoad: + def __init__(self, path, index, **kwargs): + self.path = path + self.index = index + self.kwargs = kwargs + + def pre_serialize(self): + """Make the unloaded partition serializable""" + return self.load() + + def load(self): + """Load the partition into memory""" + import cudf + + fn = glob.glob(f"{self.path}/*.{self.index}.parquet") + return cudf.read_parquet(fn, **self.kwargs) + + +@dask_serialize.register(LazyLoad) +def _serialize_unloaded(obj): + return None, [pickle.dumps(obj.pre_serialize())] + + +@dask_deserialize.register(LazyLoad) +def _deserialize_unloaded(header, frames): + return pickle.loads(frames[0]) + + +def _prepare_dir(dirpath: str): + os.makedirs(dirpath, exist_ok=True) + + +def _clean_worker_storage(dirpath: str): + import shutil + + if os.path.isdir(dirpath): + shutil.rmtree(dirpath) + + +def _write_partition(part, dirpath, index, token=None): + if token is None: + fn = f"{dirpath}/part.{index[0]}.parquet" + else: + fn = f"{dirpath}/part.{token}.{index[0]}.parquet" + part.to_parquet(fn) + return index + + +def _get_partition(dirpath, index): + return LazyLoad(dirpath, index) + + +def _get_metadata(dirpath, index): + import glob + + import pyarrow.parquet as pq + + import cudf + + fn = glob.glob(f"{dirpath}/*.{index}.parquet")[0] + return cudf.DataFrame.from_arrow( + pq.ParquetFile(fn).schema.to_arrow_schema().empty_table() + ) + + +def _load_partition(data): + if isinstance(data, LazyLoad): + data = data.load() + return data + + +def to_worker_storage(df, dirpath, shuffle_on=None, overwrite=False, **kwargs): + + if shuffle_on: + from dask_cuda.explicit_comms.dataframe.shuffle import shuffle_to_parquet + + if not isinstance(shuffle_on, list): + shuffle_on = [shuffle_on] + return shuffle_to_parquet( + df, shuffle_on, dirpath, overwrite=overwrite, **kwargs + ) + + c = comms.default_comms() + if overwrite: + wait(c.client.run(_clean_worker_storage, dirpath)) + wait(c.client.run(_prepare_dir, dirpath)) + df.map_partitions( + _write_partition, + dirpath, + BlockIndex((df.npartitions,)), + **kwargs, + ).compute() + + +def from_worker_storage(dirpath): + import dask_cudf + + c = comms.default_comms() + + def get_indices(path): + return {int(fn.split(".")[-2]) for fn in glob.glob(path + "/*.parquet")} + + worker_indices = c.client.run(get_indices, dirpath) + + summary = defaultdict(list) + for worker, indices in worker_indices.items(): + for index in indices: + summary[index].append(worker) + + assignments = {} + futures = [] + meta = None + for i, (worker, indices) in enumerate(summary.items()): + assignments[worker] = indices[i % len(indices)] + futures.append( + c.client.submit(_get_partition, dirpath, i, workers=[assignments[i]]) + ) + if meta is None: + meta = c.client.submit(_get_metadata, dirpath, i, workers=[assignments[i]]) + wait(meta) + meta = meta.result() + wait(futures) + + return dask_cudf.from_delayed(futures, meta=meta, verify_meta=False).map_partitions( + _load_partition, + meta=meta, + enforce_metadata=False, + )