Repository: dask/dask

Issue:
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.


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.
I'd love to but I doubt I have the required knowledge to take a good stab at this.

Resolving patch:
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- 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):


Test patch:
diff --git a/dask/dataframe/tests/test_shuffle.py b/dask/dataframe/tests/test_shuffle.py
--- 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))})