<#6667 Add dask plugin proto> Pull request opened ...
# flyte-github
a
#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