#6667 Add dask plugin proto
Pull request opened by
pingsutw
## Tracking issue
NA
## What changes were proposed in this pull request?
Add dask plugin proto
## How was this patch tested?
import asyncio
import typing
from distributed import Client
from flyteplugins.dask import Dask, Scheduler, WorkerGroup
import flyte.remote
import flyte.storage
from flyte import Resources
image = (
flyte.Image.from_debian_base(python_version=(3, 12))
.with_env_vars({"SETUPTOOLS_SCM_PRETEND_VERSION_FOR_FLYTEIDL2": "3.0.0"})
.with_apt_packages("git")
.with_pip_packages("flyteplugins-dask")
.with_uv_project(pyproject_file="/Users/kevin/git/cloud/flyte2/gen/python/pyproject.toml")
)
dask_config = Dask(
scheduler=Scheduler(resources=Resources(cpu="1", memory="1Gi")),
workers=WorkerGroup(number_of_workers=4, resources=Resources(cpu="1", memory="1Gi")),
)
task_env = flyte.TaskEnvironment(
name="hello_dask", resources=Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image
)
dask_env = flyte.TaskEnvironment(
name="dask_env",
plugin_config=dask_config,
image=image,
resources=Resources(cpu="1", memory="1Gi"),
depends_on=[task_env],
)
@task_env.task()
async def hello_dask():
await asyncio.sleep(5)
print("Hello from the Dask task!")
@dask_env.task
async def hello_dask_nested(n: int = 3) -> typing.List[int]:
print("running dask task")
t = asyncio.create_task(hello_dask())
client = Client()
futures = client.map(lambda x: x + 1, range(n))
res = client.gather(futures)
await t
return res
if
name == "__main__":
flyte.init_from_config()
run = flyte.run(hello_dask_nested, n=3)
print("run name:", run.name)
print("run url:", run.url)
run.wait(run)
action_details = flyte.remote.ActionDetails.get(run_name=run.name, name="a0")
for log in action_details.pb2.attempts[-1].log_info:
print(f"{log.name}: {log.uri}")
### Labels
Please add one or more of the following labels to categorize your PR:
•
added: For new features.
•
changed: For changes in existing functionality.
•
deprecated: For soon-to-be-removed features.
•
removed: For features being removed.
•
fixed: For any bug fixed.
•
security: In case of vulnerabilities
This is important to improve the readability of release notes.
### Setup process
### Screenshots
## Check all the applicable boxes
• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.
## Related PRs
NA
## Docs link
NA
flyteorg/flyte