flyteplugins-dask 2.0.0b24__py3-none-any.whl → 2.0.0b26__py3-none-any.whl

flyteplugins/dask/task.py

@@ -1,12 +1,14 @@
+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
-from flyteidl.plugins.dask_pb2 import DaskJob, DaskScheduler, DaskWorkerGroup
+from flyteidl2.plugins.dask_pb2 import DaskJob, DaskScheduler, DaskWorkerGroup
 from google.protobuf.json_format import MessageToDict
 
 
@@ -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)
@@ -78,13 +94,15 @@ class DaskTask(AsyncFunctionTaskTemplate):
 
     plugin_config: Dask
     task_type: str = "dask"
+    debuggable: bool = True
 
     async def pre(self, *args, **kwargs) -> Dict[str, Any]:
         ctx = flyte.ctx()
         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.0b24.dist-info → flyteplugins_dask-2.0.0b26.dist-info}/METADATA

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

flyteplugins_dask-2.0.0b26.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.0b26.dist-info/METADATA,sha256=1gRWI9D_6tklLf8wFBrcvhRtSnQv_9AW1qj54ufs9Kk,948
+flyteplugins_dask-2.0.0b26.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+flyteplugins_dask-2.0.0b26.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_dask-2.0.0b26.dist-info/RECORD,,

flyteplugins_dask-2.0.0b24.dist-info/RECORD

@@ -1,6 +0,0 @@
-flyteplugins/dask/__init__.py,sha256=UwmJjPRVqDyCTobQlT5a2g2oZJvK0vlDiv6rwugNs-g,112
-flyteplugins/dask/task.py,sha256=giAaSDZb0cIedzJ2ij8LRq7u8WvtB75l2lgLyeKwKQI,3634
-flyteplugins_dask-2.0.0b24.dist-info/METADATA,sha256=IDPjmHMoQ8iCirFQE3-eFb_n18ISB9J-J4ti7DYU8Cg,948
-flyteplugins_dask-2.0.0b24.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-flyteplugins_dask-2.0.0b24.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
-flyteplugins_dask-2.0.0b24.dist-info/RECORD,,