Dask/Distributed Performance Regression

[seanlaw]

Somewhere between Dask/Distributed version 2022.5.0 and 2022.5.1 (and subsequently 2022.5.2), STUMPY's dask-related unit tests (e.g., https://github.com/TDAmeritrade/stumpy/blob/main/tests/test_stumped.py) are around 10x slower. This is also causing our Github Actions workflows to timeout.

We'll need to track this down and find out what may have changed in either dask and/or distributed.

[seanlaw]

@jrbourbeau @mrocklin @crusaderky It'll take me a few days to formulate a better reproducible example (beyond switching from 2022.5.0 to 2022.5.1) but would any of you happen to have any thoughts or ideas of what might be causing this? I've tried looking at the distributed changelogs (https://github.com/dask/distributed/blob/main/docs/source/changelog.rst) but didn't notice anything obvious.

[mrocklin]

Nothing is immediately obvious to me personally. @crusaderky might be the best bet.

Alternatively, is git bisect a feasible option here? (how easy is it to reproduce the bad behavior?)

[seanlaw]

Thanks @mrocklin. I'm going to find some time and see if I can isolate a minimum reproducible example

[mrocklin]

Sounds good. git bisect may genuinely be faster is the stumpy test suite is easy to run. Happy to walk you through using it if you're unfamiliar. It's a useful tool.

[seanlaw]

I had heard about git bisect before but this is my first time using it and it felt like magic! Thanks for your recommendation @mrocklin 😄

I was able to isolate the problem in distributed. This was the first "bad" commit:

commit 36e99465bdd80be283440ad823cb8231ccb471f5
Author: Florian Jetter <fjetter@users.noreply.github.com>
Date:   Wed May 18 11:45:38 2022 +0200

    Ensure multiple clients can cancel their key without interference (#6016)

 distributed/client.py                  |  2 ++
 distributed/scheduler.py               | 39 +++++++++++++++++------------
 distributed/tests/test_as_completed.py |  8 +++---
 distributed/tests/test_client.py       | 45 ++++++++++++++++++++++------------
 distributed/tests/test_worker.py       |  4 +--
 5 files changed, 61 insertions(+), 37 deletions(-)

So, to reproduce the problem, I started with a fresh miniconda environment and then did:

conda install -y -c conda-forge stumpy
cd dask
git checkout 298112115f63bd8846527bfafae2886999f0f4ec
python -m pip install .
cd ../distributed
git checkout 36e99465bdd80be283440ad823cb8231ccb471f5
python -m pip install .

Then, I ran this test.py script:

#!/usr/bin/env python

import numpy as np
from stumpy import stumped
from dask.distributed import Client, LocalCluster
import dask, distributed
import time


if __name__ == "__main__":
    print(dask.__version__, distributed.__version__)
    T_A = np.random.uniform(-1000, 1000, [8]).astype(np.float64)
    T_B = np.random.uniform(-1000, 1000, [64]).astype(np.float64)
    m = 3
    cluster = LocalCluster(n_workers=2, threads_per_worker=2)

    start = time.time()
    for i in range(10):
        with Client(cluster) as dask_client:
            stumped(dask_client, T_A, m, T_B, ignore_trivial=False)
    print(time.time() - start)
    cluster.close()

On my Apple M1 Macbook Air, this took around 26 seconds. However, if I revert distributed back by one commit (a "good" commit):

git checkout 32313c0150b5a5527c81829b5de4911444204975

test.py now runs in around 6 seconds, which is ~4x faster!

@mrocklin Would it be more appropriate to create a new issue in the distributed repo?

[mrocklin]

I'll defer that to @fjetter

[mrocklin]

Also, as a heads-up, today was a German holiday, so Florian was out.

[seanlaw]

Also, as a heads-up, today was a German holiday, so Florian was out.

Thank you for your thoughtfulness @mrocklin ! I understand and no rush

[fjetter]

I'll have a look

[fjetter]

The commit you are referencing is PR dask/distributed#6016
and the regression introduced is us now waiting explicitly for a second upon a cancellation request if a key is no longer available dask/distributed#6517
This wait period is hit by https://github.com/fjetter/stumpy/blob/89487ae3991c50ead0e927c90a19ed6ace32df7e/stumpy/stumped.py#L266-L282
I assume you're cancelling some futures multiple times and we're hitting this wait period on two calls (every stumped call requires about 2s more than before). I haven't tried to find out what futures you're cancelling multiple times.

I think the simplest fix for you is to no longer cancel things explicitly, see #617