Repository: dask/dask
Pull request: Fixed flaky test-rearrange
Puts file cleanup in the task graph, rather than relying on Partd to do it for us. I'm not really happy with my approach, but wanted to try this on CI a few times.

Closes #5867 


Problem statement:
Flaky test test_rearrange
```
=================================== FAILURES ===================================

________________________ test_rearrange[processes-disk] ________________________

shuffle = 'disk', scheduler = 'processes'

    @pytest.mark.parametrize("shuffle", ["tasks", "disk"])

    @pytest.mark.parametrize("scheduler", ["threads", "processes"])

    def test_rearrange(shuffle, scheduler):

        df = pd.DataFrame({"x": np.random.random(10)})

        ddf = dd.from_pandas(df, npartitions=4)

        ddf2 = ddf.assign(_partitions=ddf.x % 4)

    

        result = rearrange_by_column(ddf2, "_partitions", max_branch=32, shuffle=shuffle)

        assert result.npartitions == ddf.npartitions

        assert set(ddf.dask).issubset(result.dask)

    

        # Every value in exactly one partition

>       a = result.compute(scheduler=scheduler)

dask/dataframe/tests/test_shuffle.py:272: 

_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

dask/base.py:166: in compute

    (result,) = compute(self, traverse=False, **kwargs)

dask/base.py:437: in compute

    results = schedule(dsk, keys, **kwargs)

dask/multiprocessing.py:212: in get

    result = get_async(

dask/local.py:486: in get_async

    raise_exception(exc, tb)

dask/local.py:316: in reraise

    raise exc

dask/local.py:222: in execute_task

    result = _execute_task(task, data)

dask/core.py:119: in _execute_task

    return func(*args2)

dask/dataframe/shuffle.py:669: in shuffle_group_3

    p.append(d, fsync=True)

../../../miniconda/envs/test-environment/lib/python3.8/site-packages/partd/encode.py:25: in append

    self.partd.append(data, **kwargs)

../../../miniconda/envs/test-environment/lib/python3.8/site-packages/partd/file.py:37: in append

    if lock: self.lock.acquire()

../../../miniconda/envs/test-environment/lib/python3.8/site-packages/locket/__init__.py:99: in acquire

    lock.acquire()

_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

>   self._file = open(self._path, "w")

E   FileNotFoundError: [Errno 2] No such file or directory: '/tmp/tmpn5wlil7v.partd/.lock'
```

Will look into this a bit. Both were on 3.8 IIRC.

Hints from issue discussion:
This seems to occur when the directory is removed prior between creating the `partd.File` object and attempting to append data to it. Roughly

```pytb
In [13]: import partd
In [14]: f = partd.File("data", )
In [16]: !rm -rf data

In [17]: f.append({"x": b"abc"})

---------------------------------------------------------------------------
FileNotFoundError                         Traceback (most recent call last)
<ipython-input-17-40cdad36fb6c> in <module>
----> 1 f.append({"x": b"abc"})

~/sandbox/partd/partd/file.py in append(self, data, lock, fsync, **kwargs)
     38     def append(self, data, lock=True, fsync=False, **kwargs):
     39         print(os.path.exists(self.filename('')))
---> 40         if lock: self.lock.acquire()
     41         try:
     42             for k, v in data.items():

~/sandbox/locket.py/locket/__init__.py in acquire(self)
    168     def acquire(self):
    169         self._lock.acquire(self._timeout, self._retry_period)
--> 170
    171     def release(self):
    172         self._lock.release()

~/sandbox/locket.py/locket/__init__.py in acquire(self, timeout, retry_period)
     98             for lock in self._locks:
     99                 lock.acquire(timeout, retry_period)
--> 100                 acquired_locks.append(lock)
    101         except:
    102             for acquired_lock in reversed(acquired_locks):

~/sandbox/locket.py/locket/__init__.py in acquire(self, timeout, retry_period)
    140         if self._file is None:
    141             self._file = open(self._path, "wb")
--> 142         if timeout is None and _lock_file_blocking_available:
    143             _lock_file_blocking(self._file)
    144         else:

FileNotFoundError: [Errno 2] No such file or directory: 'data/.lock'
```

It's not (yet) clear to my why that directory is being removed.
OK, I think what's happening is

1. Process A creates a `part.File` with no path. Partd makes a tmpdir, and marks it for removal when the `partd.File` is GCed
2. Process B reuses process A's partd.File's tmpdir for its own `partd.File`.
3. Process A's partd.File is GCed, triggering A's `part.File.drop`. This clears the tmpdir and removes it.
4. We have a race condition: If process B attempts to append data prior to A's `File.drop` being done, we get the exception. You can simulate this by inserting a `time.sleep` in `partd.File.drop`.

I don't know how to fix this yet.
So why would process B be appending data before A's drop is called?  I would only expect any process to clear out the data directory after all of them are done operating on it.
That I don't know. Here's some debug output. These are printing either the task key & PID in `dask.local.execute_task`, or the `file` passed to `partd.File` and the PID.

```
** task ('zpartd-4320f2f849cb11eab8c6186590cd1c87',) 24524
** file None /var/folders/hz/f43khqfn7b1b1g8z_z6y3bsw0000gp/T/tmp46ml086h.partd 24524
task ('getitem-mod-assign-04375fd34e33c3cc6fbc296cec0cdf70', 0) 24523
task ('getitem-mod-assign-04375fd34e33c3cc6fbc296cec0cdf70', 1) 24525
task ('getitem-mod-assign-04375fd34e33c3cc6fbc296cec0cdf70', 2) 24526
file /var/folders/hz/f43khqfn7b1b1g8z_z6y3bsw0000gp/T/tmp46ml086h.partd /var/folders/hz/f43khqfn7b1b1g8z_z6y3bsw0000gp/T/tmp46ml086h.partd 24471
task ('assign-04375fd34e33c3cc6fbc296cec0cdf70', 1) 24527
task ('assign-04375fd34e33c3cc6fbc296cec0cdf70', 2) 24528
task ('assign-04375fd34e33c3cc6fbc296cec0cdf70', 0) 24529
file /var/folders/hz/f43khqfn7b1b1g8z_z6y3bsw0000gp/T/tmp46ml086h.partd /var/folders/hz/f43khqfn7b1b1g8z_z6y3bsw0000gp/T/tmp46ml086h.partd 24525
task ('shuffle-partition-4320f2f849cb11eab8c6186590cd1c87', 0) 24525
file /var/folders/hz/f43khqfn7b1b1g8z_z6y3bsw0000gp/T/tmp46ml086h.partd /var/folders/hz/f43khqfn7b1b1g8z_z6y3bsw0000gp/T/tmp46ml086h.partd 24524
** task ('shuffle-partition-4320f2f849cb11eab8c6186590cd1c87', 1) 24524
file /var/folders/hz/f43khqfn7b1b1g8z_z6y3bsw0000gp/T/tmp46ml086h.partd /var/folders/hz/f43khqfn7b1b1g8z_z6y3bsw0000gp/T/tmp46ml086h.partd 24530
task ('shuffle-partition-4320f2f849cb11eab8c6186590cd1c87', 2) 24530
partd.File.append          24525
partd.File.drop-start      24524
partd.File.append          24530
(Pdb) partd.File.drop-done 24524
partd.File.append          24524
```

I've prefixed some lines with `**`.

Process `24524` created the initial `partd.File` without a path during `zpartd-...`, so partd made it and marked it for deletion when that `File` goes out of scope. Other processes start using that directory. Then process `24524` moved on to `shuffle-partition-...`, which (I think?) will cause the `File` to go out of scope and eventually get GCed. That's logged with a `partd.File.drop-start`, and finishes with a `partd.File.drop-done` (I added a 3 second `time.sleep` in there). In the meantime, a few other processes did appends to that directory.
So, IMO, the easiest fix here is to have the initial task creating the file pass a `path` to `partd.File` (which disables partd.File's auto-cleanup), and *perhaps* add a task at the end that cleans stuff up?
Hrm, but if we fail partway through the computation then we won't clean things up?

cc also @jcrist , who I think had a hand in the multi-processing Partd implementation (although no pressure)
@jcrist I think that you volunteered to take a look at this.  Is this still something that you can do?
I dove into this this morning, and wasn't able to reproduce on either 3.7 or 3.8. Tom eventually was able to reproduce locally, but only after a bit of work - running the same reproducer on my machine doesn't trigger anything even after 500 runs. I'm not sure what's going on here, but without being able to reproduce the issue I'm not sure what I can do to help.
I'm getting this error when running compute fairly consistently at the moment.

I'm currently playing around with my worker/thread distribution on my LocalCluster, and this mainly seems to crop up when I have higher process count vs lower (8 vs 2)
@KrishanBhasin are you interested in working on a fix? My best understanding is in https://github.com/dask/dask/issues/5867#issuecomment-583479468, but I haven't come up with an actual solution yet.

Commit message: hrmm

Patch to review:
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
index 1d29d0b4220..153704cb67c 100644
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -1,6 +1,9 @@
+import logging
 import math
+import shutil
 from operator import getitem
 import uuid
+import tempfile
 import warnings
 
 import toolz
@@ -17,6 +20,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 +345,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 +360,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 +412,26 @@ 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
+    dsk5 = {cleanup_token: (cleanup_partd_files, p, list(dsk4))}
 
+    name = "shuffle-collect-2" + token
+    dsk6 = {(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)
     graph = HighLevelGraph.from_collections(name, layer, dependencies=dependencies)
-
     return DataFrame(graph, name, df._meta, divisions)
 
 
+def noop(x, cleanup_token):
+    return x
+
+
 def rearrange_by_column_tasks(
     df, column, max_branch=32, npartitions=None, ignore_index=False
 ):
@@ -611,6 +623,33 @@ 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.
+    exc_only : bool, default True
+        Whether to cleanup if and only if there's an exception.
+    """
+    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)
@@ -684,9 +723,16 @@ def shuffle_group(df, col, stage, k, npartitions, ignore_index):
 
 
 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)
+    try:
+        g = df.groupby(col)
+        d = {i: g.get_group(i) for i in g.groups}
+        p.append(d, fsync=True)
+    except Exception:
+        try:
+            p.drop()
+        except Exception:
+            logger.exception("ignoring exception dask.dataframe.shuffle.collect")
+        raise
 
 
 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
index 8be133504a0..f5ed39f10be 100644
--- a/dask/dataframe/tests/test_shuffle.py
+++ b/dask/dataframe/tests/test_shuffle.py
@@ -1,6 +1,7 @@
 import itertools
 import os
 import random
+import tempfile
 
 import pandas as pd
 import pytest
@@ -257,8 +258,8 @@ def test_shuffle_sort(shuffle):
     assert_eq(ddf2.loc[2:3], df2.loc[2:3])
 
 
-@pytest.mark.parametrize("shuffle", ["tasks", "disk"])
-@pytest.mark.parametrize("scheduler", ["threads", "processes"])
+@pytest.mark.parametrize("shuffle", ["disk"])
+@pytest.mark.parametrize("scheduler", ["processes"])
 def test_rearrange(shuffle, scheduler):
     df = pd.DataFrame({"x": np.random.random(10)})
     ddf = dd.from_pandas(df, npartitions=4)
@@ -269,7 +270,7 @@ def test_rearrange(shuffle, scheduler):
     assert set(ddf.dask).issubset(result.dask)
 
     # Every value in exactly one partition
-    a = result.compute(scheduler=scheduler)
+    a = result.compute(scheduler=scheduler, optimize_graph=False)
     get = dask.base.get_scheduler(scheduler=scheduler)
     parts = get(result.dask, result.__dask_keys__())
 
@@ -277,6 +278,20 @@ 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_by_column_with_narrow_divisions():
     from dask.dataframe.tests.test_multi import list_eq
 


Code review comments:
diff_hunk: @@ -257,8 +258,8 @@ def test_shuffle_sort(shuffle):
     assert_eq(ddf2.loc[2:3], df2.loc[2:3])
 
 
-@pytest.mark.parametrize("shuffle", ["tasks", "disk"])
-@pytest.mark.parametrize("scheduler", ["threads", "processes"])
+@pytest.mark.parametrize("shuffle", ["disk"])
+@pytest.mark.parametrize("scheduler", ["processes"])
path: dask/dataframe/tests/test_shuffle.py
text: @user1:
Why did we drop the extra parameters here?

@author:
Sorry, just debugging.

diff_hunk: @@ -407,19 +412,26 @@ 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
+    dsk5 = {cleanup_token: (cleanup_partd_files, p, list(dsk4))}
 
+    name = "shuffle-collect-2" + token
+    dsk6 = {(name, i): (noop, (name1, i), cleanup_token) for i in range(npartitions)}
path: dask/dataframe/shuffle.py
text: @user1:
`(toolz.first, [(name1, i), cleanup_token])`?

diff_hunk: @@ -407,19 +412,26 @@ 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
+    dsk5 = {cleanup_token: (cleanup_partd_files, p, list(dsk4))}
path: dask/dataframe/shuffle.py
text: @user1:
I think that you're bringing back all of the results into a single task with `list(dsk4)`.

If we do this approach then we might have to have an intermediary set of tasks that take in the results of dsk4 and return None for each.  In this way we maintain the dependencies, but don't move the data around.

@author:
Mmm OK. Do you know, does `dsk3` at https://github.com/dask/dask/pull/5977/files#diff-83ae5352ddc87bd80c831102addd9b1eL407 have a similar problem? Perhaps that result isn't as large.

@user1:
I don't think so, my guess is that the output tasks in dsk2 write things to disk and then return `None`.  

The output tasks of `dsk4` read from disk and return dataframes, so this is more of a concern.

Instead, I think that the solution is to have a `dsk4b` that maps `lambda x: None` across the outputs of dsk4.