Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds PrefectDistributedClient to prefect-dask #13537

Merged
merged 13 commits into from
May 24, 2024
Merged

Conversation

desertaxle
Copy link
Member

@desertaxle desertaxle commented May 23, 2024

This PR adds a PrefectDistributedClient to prefect-dask. The client is a drop in replacement for the client from distributed for submitting work to a Dask cluster. This client will power the new task runner, but splitting out this client allows users to take advantage of Prefect features without changing their existing Dask code.

Example

An example flow with tasks submitted to Dask via the PrefectDistributedClient

from prefect.flows import flow
from prefect.tasks import task
from prefect_dask import PrefectDistributedClient

@task
def square(x):
    return x**2


@task
def neg(x):
    return -x


@task
def summation(args):
    return sum(args)


@flow
def my_flow():
    with PrefectDistributedClient() as client:
        square_futures = client.map(square, range(10))
        neg_futures = client.map(neg, square_futures)
        total_future = client.submit(
            summation,
            args=neg_futures,
        )
        return total_future.result()


if __name__ == "__main__":
   my_flow()

Checklist

  • This pull request references any related issue by including "closes <link to issue>"
    • If no issue exists and your change is not a small fix, please create an issue first.
  • If this pull request adds new functionality, it includes unit tests that cover the changes
  • This pull request includes a label categorizing the change e.g. maintenance, fix, feature, enhancement, docs.

For documentation changes:

  • This pull request includes redirect settings in netlify.toml for files that are removed or renamed.

For new functions or classes in the Python SDK:

  • This pull request includes helpful docstrings.
  • If a new Python file was added, this pull request contains a stub page in the Python SDK docs and an entry in mkdocs.yml navigation.

@desertaxle desertaxle marked this pull request as ready for review May 24, 2024 03:19
@desertaxle desertaxle requested review from zzstoatzz and a team as code owners May 24, 2024 03:19
Comment on lines +438 to +441
with hydrated_context(self.context):
with get_client(sync_client=True) as client:
self._client = client
self._is_started = True
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had to change this order to ensure the client gets the correct API URL from the current settings.

@desertaxle desertaxle requested a review from cicdw May 24, 2024 14:10
Copy link
Member

@cicdw cicdw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@desertaxle desertaxle merged commit 1566e18 into main May 24, 2024
30 checks passed
@desertaxle desertaxle deleted the udpated-dask-task-runner branch May 24, 2024 15:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants