dask · TomAugspurger · Mar 31, 2020 · Mar 4, 2020 · Mar 5, 2020 · Mar 5, 2020
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
@@ -1,6 +1,10 @@
+import contextlib
+import logging
 import math
+import shutil
 from operator import getitem
 import uuid
+import tempfile
 import warnings
 
 import tlz as toolz
@@ -17,6 +21,8 @@
 from ..utils import digit, insert, M
 from .utils import hash_object_dispatch, group_split_dispatch
 
+logger = logging.getLogger(__name__)
+
 
 def set_index(
     df,
@@ -340,6 +346,8 @@ def __reduce__(self):
     def __call__(self, *args, **kwargs):
         import partd
 
+        path = tempfile.mkdtemp(suffix=".partd", dir=self.tempdir)
+
         try:
             partd_compression = (
                 getattr(partd.compressed, self.compression)
@@ -353,10 +361,8 @@ def __call__(self, *args, **kwargs):
                     self.compression
                 )
             )
-        if self.tempdir:
-            file = partd.File(dir=self.tempdir)
-        else:
-            file = partd.File()
+        file = partd.File(path)
+        partd.file.cleanup_files.append(path)
         # Envelope partd file with compression, if set and available
         if partd_compression:
             file = partd_compression(file)
@@ -407,19 +413,33 @@ def rearrange_by_column_disk(df, column, npartitions=None, compute=False):
     dsk3 = {barrier_token: (barrier, list(dsk2))}
 
     # Collect groups
-    name = "shuffle-collect-" + token
+    name1 = "shuffle-collect-1" + token
     dsk4 = {
-        (name, i): (collect, p, i, df._meta, barrier_token) for i in range(npartitions)
+        (name1, i): (collect, p, i, df._meta, barrier_token) for i in range(npartitions)
     }
-
+    cleanup_token = "cleanup-" + always_new_token
+    barrier_token2 = "barrier2-" + always_new_token
+    # A task that depends on `cleanup-`, but has a small output
+    dsk5 = {(barrier_token2, i): (barrier, part) for i, part in enumerate(dsk4)}
+    # This indirectly depends on `cleanup-` and so runs after we're done using the disk
+    dsk6 = {cleanup_token: (cleanup_partd_files, p, list(dsk5))}
+
+    name = "shuffle-collect-2" + token
+    dsk7 = {(name, i): (_noop, (name1, i), cleanup_token) for i in range(npartitions)}
     divisions = (None,) * (npartitions + 1)
 
-    layer = toolz.merge(dsk1, dsk2, dsk3, dsk4)
+    layer = toolz.merge(dsk1, dsk2, dsk3, dsk4, dsk5, dsk6, dsk7)
     graph = HighLevelGraph.from_collections(name, layer, dependencies=dependencies)
-
     return DataFrame(graph, name, df._meta, divisions)
 
 
+def _noop(x, cleanup_token):
+    """
+    A task that does nothing.
+    """
+    return x
+
+
 def rearrange_by_column_tasks(
     df, column, max_branch=32, npartitions=None, ignore_index=False
 ):
@@ -611,10 +631,38 @@ def barrier(args):
     return 0
 
 
+def cleanup_partd_files(p, keys):
+    """
+    Cleanup the files in a partd.File dataset.
+
+    Parameters
+    ----------
+    p : partd.Interface
+        File or Encode wrapping a file should be OK.
+    keys: List
+        Just for scheduling purposes, not actually used.
+    """
+    import partd
+
+    if isinstance(p, partd.Encode):
+        maybe_file = p.partd
+    else:
+        maybe_file
+
+    if isinstance(maybe_file, partd.File):
+        path = maybe_file.path
+    else:
+        path = None
+
+    if path:
+        shutil.rmtree(path, ignore_errors=True)
+
+
 def collect(p, part, meta, barrier_token):
     """ Collect partitions from partd, yield dataframes """
-    res = p.get(part)
-    return res if len(res) > 0 else meta
+    with ensure_cleanup_on_exception(p):
+        res = p.get(part)
+        return res if len(res) > 0 else meta
 
 
 def set_partitions_pre(s, divisions):
@@ -683,10 +731,31 @@ def shuffle_group(df, col, stage, k, npartitions, ignore_index):
     return group_split_dispatch(df, c.astype(np.int64), k, ignore_index=ignore_index)
 
 
+@contextlib.contextmanager
+def ensure_cleanup_on_exception(p):
+    """Ensure a partd.File is cleaned up.
+
+    We have several tasks referring to a `partd.File` instance. We want to
+    ensure that the file is cleaned up if and only if there's an exception
+    in the tasks using the `partd.File`.
+    """
+    try:
+        yield
+    except Exception:
+        # the function (e.g. shuffle_group_3) had an internal exception.
+        # We'll cleanup our temporary files and re-raise.
+        try:
+            p.drop()
+        except Exception:
+            logger.exception("ignoring exception in ensure_cleanup_on_exception")
+        raise
+
+
 def shuffle_group_3(df, col, npartitions, p):
-    g = df.groupby(col)
-    d = {i: g.get_group(i) for i in g.groups}
-    p.append(d, fsync=True)
+    with ensure_cleanup_on_exception(p):
+        g = df.groupby(col)
+        d = {i: g.get_group(i) for i in g.groups}
+        p.append(d, fsync=True)
 
 
 def set_index_post_scalar(df, index_name, drop, column_dtype):

diff --git a/dask/dataframe/tests/test_shuffle.py b/dask/dataframe/tests/test_shuffle.py
@@ -1,6 +1,8 @@
 import itertools
 import os
 import random
+import tempfile
+from unittest import mock
 
 import pandas as pd
 import pytest
@@ -24,7 +26,6 @@
     remove_nans,
 )
 from dask.dataframe.utils import assert_eq, make_meta
-from dask.compatibility import PY_VERSION
 
 
 dsk = {
@@ -277,6 +278,42 @@ def test_rearrange(shuffle, scheduler):
         assert sum(i in set(part._partitions) for part in parts) == 1
 
 
+def test_rearrange_cleanup():
+    df = pd.DataFrame({"x": np.random.random(10)})
+    ddf = dd.from_pandas(df, npartitions=4)
+    ddf2 = ddf.assign(_partitions=ddf.x % 4)
+
+    tmpdir = tempfile.mkdtemp()
+
+    with dask.config.set(temporay_directory=str(tmpdir)):
+        result = rearrange_by_column(ddf2, "_partitions", max_branch=32, shuffle="disk")
+        result.compute(scheduler="processes")
+
+    assert len(os.listdir(tmpdir)) == 0
+
+
+def test_rearrange_disk_cleanup_with_exception():
+    # ensure temporary files are cleaned up when there's an internal exception.
+    def mock_shuffle_group_3(df, col, npartitions, p):
+        raise ValueError("Mock exception!")
+
+    with mock.patch("dask.dataframe.shuffle.shuffle_group_3", new=mock_shuffle_group_3):
+        df = pd.DataFrame({"x": np.random.random(10)})
+        ddf = dd.from_pandas(df, npartitions=4)
+        ddf2 = ddf.assign(_partitions=ddf.x % 4)
+
+        tmpdir = tempfile.mkdtemp()
+
+        with dask.config.set(temporay_directory=str(tmpdir)):
+            with pytest.raises(ValueError, match="Mock exception!"):
+                result = rearrange_by_column(
+                    ddf2, "_partitions", max_branch=32, shuffle="disk"
+                )
+                result.compute(scheduler="processes")
+
+    assert len(os.listdir(tmpdir)) == 0
+
+
 def test_rearrange_by_column_with_narrow_divisions():
     from dask.dataframe.tests.test_multi import list_eq
 
@@ -779,34 +816,6 @@ def test_compute_divisions():
         compute_divisions(c)
 
 
-# TODO: Fix sporadic failure on Python 3.8 and remove this xfail mark
-@pytest.mark.xfail(PY_VERSION >= "3.8", reason="Flaky test", strict=False)
-def test_temporary_directory(tmpdir):
-    from multiprocessing.pool import Pool
-
-    df = pd.DataFrame(
-        {
-            "x": np.random.random(100),
-            "y": np.random.random(100),
-            "z": np.random.random(100),
-        }
-    )
-    ddf = dd.from_pandas(df, npartitions=10, name="x", sort=False)
-
-    # We use a pool to avoid a race condition between the pool close
-    # cleaning up files, and the assert below.
-    pool = Pool(4)
-    with pool:
-        with dask.config.set(
-            temporary_directory=str(tmpdir), scheduler="processes", pool=pool
-        ):
-            ddf2 = ddf.set_index("x", shuffle="disk")
-            ddf2.compute()
-            assert any(
-                fn.endswith(".partd") for fn in os.listdir(str(tmpdir))
-            ), os.listdir(str(tmpdir))
-
-
 def test_empty_partitions():
     # See https://github.com/dask/dask/issues/2408
     df = pd.DataFrame({"a": list(range(10))})