#233 Fix fast registration for Dask Pull request opened by pingsutw Add a worker plugin to download the code bundle in the worker nodes import asyncio import typing from pathlib import Path from distributed import Client from flyteplugins.dask import Dask, Scheduler, WorkerGroup from task_functions import hello import flyte.remote import flyte.storage from flyte import Resources # image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("flyteplugins-dask") image = ( flyte.Image.from_debian_base(python_version=(3, 12)) .with_pip_packages("dask[distributed]", "bokeh>=3.1.0") .with_source_folder(Path(file).parent.parent.parent / "plugins/dask", "./dask") .with_env_vars({"PYTHONPATH": "./dask/src:${PYTHONPATH}", "hello": "world2"}) ) 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() def _test(x: int) -> int: hello() return x + 1 client.upload_file("./task_functions.py") # from distributed.diagnostics.plugin import UploadFile # client.register_plugin(UploadFile("./task_functions.py")) futures = client.map(_test, range(n)) res = client.gather(futures) await t return res if name == "__main__": flyte.init_from_config(root_dir=Path(file).parent) 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}") flyteorg/flyte-sdk