Add get_dask_client

[ahuang11]

Summary

Adds get_dask_client.

This is intended to be called within tasks that run on workers, and is useful for operating on dask collections, such as a dask.DataFrame.

Without invoking this, workers in a task do not automatically get a client to connect to the full cluster. Therefore, it will attempt perform work within the worker itself serially, and potentially overwhelming the single worker.

Internally, this context manager is a simple wrapper around distributed.worker_client with separate_thread=False fixed.

With it:

Without it (a single worker takes on the entire job):

Why a separate util function, from Michael:
"""
I think we should document our required pattern for accessing the client. If we want, we can add a utility to the library that, as you noted, just sets the required keyword argument. The utility function is nice because:

• We can maintain it over time without our users having to change their invocations
• It is discoverable without reading documentation — otherwise I'd never guess that I need to pass separate_thread
• We can account for both sync and async invocations
  """

Relevant Issue(s)

Closes #26

Checklist

• Summarized PR's changes in CHANGELOG.md

[ahuang11]

I was able to get it to connect to existing client, but unable to get asynchronous working because if I make this is async, then users need to call async with get_dask_client

Also, I couldn't get sync_compatible working on it.

Thoughts?

[scharlottej13]

Thanks again for your work on this @ahuang11!

[ahuang11]

I think this is almost ready; just missing the @sync_compatible piece:

import dask
from prefect import flow, task
from prefect_dask import DaskTaskRunner, get_dask_client

@task
async def compute_task():
    with get_dask_client() as client:
        df = dask.datasets.timeseries("2000", "2001", partition_freq="4w")
        summary_df = await client.compute(df.describe())
    return summary_df

@flow(task_runner=DaskTaskRunner())
async def dask_flow():
    prefect_future = await compute_task.submit()
    return prefect_future.result()

await dask_flow()

results in:

/Users/andrew/mambaforge/envs/prefect-dask/lib/python3.9/site-packages/distributed/client.py:1405: RuntimeWarning: coroutine 'wait_for' was never awaited
  self.close()
RuntimeWarning: Enable tracemalloc to get the object allocation traceback
/Users/andrew/mambaforge/envs/prefect-dask/lib/python3.9/site-packages/distributed/client.py:1405: RuntimeWarning: coroutine 'Client._close' was never awaited
  self.close()
...
TypeError: cannot pickle 'coroutine' object

[zanieb]

This is looking good! I have some minor questions and comments.

[zanieb]

This seems contrary, can you clarify? Is the second bit a Dask bug?

[ahuang11]

Will have to try with dask alone.

[ahuang11]

import dask
from dask.distributed import Client

async with Client(asynchronous=True) as client:
    df = dask.datasets.timeseries("2000", "2001", partition_freq="4w")
    print(type(df))
    print(type(df.describe()))
    print(type(df.describe().compute())) # errors on this line here
    summary_df = df.describe().compute()

[ahuang11]

I guess this works

import dask
from dask.distributed import Client, wait

async with Client(asynchronous=True) as client:
    df = dask.datasets.timeseries("2000", "2001", partition_freq="4w")
    summary_df = await df.describe().compute(sync=False)[0].result()

[zanieb]

Thanks for dealing with all the comments!