flyteplugins-dask 2.0.0b25__py3-none-any.whl → 2.0.0b28__py3-none-any.whl

flyteplugins/dask/task.py

@@ -1,8 +1,10 @@
+import asyncio
+import sys
 from dataclasses import dataclass, field
 from typing import Any, Dict, Optional
 
 import flyte
-from distributed import Client, WorkerPlugin
+from distributed import Client, SchedulerPlugin, WorkerPlugin
 from flyte import Resources
 from flyte.extend import AsyncFunctionTaskTemplate, TaskPluginRegistry, download_code_bundle, get_proto_resources
 from flyte.models import CodeBundle, SerializationContext
@@ -55,7 +57,20 @@ class Dask:
     workers: WorkerGroup = field(default_factory=lambda: WorkerGroup())
 
 
-class DownloadCodeBundlePlugin(WorkerPlugin):
+class DownloadCodeBundleSchedulerPlugin(SchedulerPlugin):
+    """
+    A Dask plugin to download and set up the code bundle on the scheduler.
+    """
+
+    def __init__(self, code_bundle: CodeBundle):
+        self.code_bundle = code_bundle
+
+    async def start(self, scheduler):
+        sys.path.insert(0, ".")
+        await download_code_bundle(self.code_bundle)
+
+
+class DownloadCodeBundleWorkerPlugin(WorkerPlugin):
     """
     A Dask plugin to download and set up the code bundle on each worker.
     """
@@ -67,7 +82,8 @@ class DownloadCodeBundlePlugin(WorkerPlugin):
         """
         Runs on each worker as it is initialized.
         """
-        download_code_bundle(self.code_bundle)
+        sys.path.insert(0, ".")
+        asyncio.run(download_code_bundle(self.code_bundle))
 
 
 @dataclass(kw_only=True)
@@ -85,7 +101,8 @@ class DaskTask(AsyncFunctionTaskTemplate):
         code_bundle = ctx.code_bundle
         if ctx.is_in_cluster() and code_bundle:
             client = Client()
-            client.register_plugin(DownloadCodeBundlePlugin(code_bundle))
+            client.register_plugin(DownloadCodeBundleWorkerPlugin(code_bundle))
+            client.register_plugin(DownloadCodeBundleSchedulerPlugin(code_bundle))
 
         return {}
 

{flyteplugins_dask-2.0.0b25.dist-info → flyteplugins_dask-2.0.0b28.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyteplugins-dask
-Version: 2.0.0b25
+Version: 2.0.0b28
 Summary: Dask plugin for flyte
 Author-email: Kevin Su <pingsutw@users.noreply.github.com>
 Requires-Python: >=3.10

flyteplugins_dask-2.0.0b28.dist-info/RECORD

@@ -0,0 +1,6 @@
+flyteplugins/dask/__init__.py,sha256=UwmJjPRVqDyCTobQlT5a2g2oZJvK0vlDiv6rwugNs-g,112
+flyteplugins/dask/task.py,sha256=2jmKzyhmsTXv8Ap8xNse3YATtnvYVXATzEEq1WP5L_8,4210
+flyteplugins_dask-2.0.0b28.dist-info/METADATA,sha256=HijZQu4OF3_hGZuHOaYQCTLIkfKYhFo-ePsAYT-71E4,948
+flyteplugins_dask-2.0.0b28.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+flyteplugins_dask-2.0.0b28.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_dask-2.0.0b28.dist-info/RECORD,,

flyteplugins_dask-2.0.0b25.dist-info/RECORD

@@ -1,6 +0,0 @@
-flyteplugins/dask/__init__.py,sha256=UwmJjPRVqDyCTobQlT5a2g2oZJvK0vlDiv6rwugNs-g,112
-flyteplugins/dask/task.py,sha256=8E8AhPYccIKrq1vADkLmvrufDmiGvFDRG3T96zsqOG0,3663
-flyteplugins_dask-2.0.0b25.dist-info/METADATA,sha256=5OidCmNiIFitaDLrf4eayrp5hIda_Ng8IqIlMdGDYN4,948
-flyteplugins_dask-2.0.0b25.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-flyteplugins_dask-2.0.0b25.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
-flyteplugins_dask-2.0.0b25.dist-info/RECORD,,