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

flyteplugins/dask/task.py

@@ -2,10 +2,10 @@ from dataclasses import dataclass, field
 from typing import Any, Dict, Optional
 
 import flyte
+from distributed import Client, WorkerPlugin
 from flyte import Resources
-from flyte._internal.runtime.resources_serde import get_proto_resources
-from flyte.extend import AsyncFunctionTaskTemplate, TaskPluginRegistry
-from flyte.models import SerializationContext
+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 google.protobuf.json_format import MessageToDict
 
@@ -55,6 +55,21 @@ class Dask:
     workers: WorkerGroup = field(default_factory=lambda: WorkerGroup())
 
 
+class DownloadCodeBundlePlugin(WorkerPlugin):
+    """
+    A Dask plugin to download and set up the code bundle on each worker.
+    """
+
+    def __init__(self, code_bundle: CodeBundle):
+        self.code_bundle = code_bundle
+
+    def setup(self, worker):
+        """
+        Runs on each worker as it is initialized.
+        """
+        download_code_bundle(self.code_bundle)
+
+
 @dataclass(kw_only=True)
 class DaskTask(AsyncFunctionTaskTemplate):
     """
@@ -65,12 +80,11 @@ class DaskTask(AsyncFunctionTaskTemplate):
     task_type: str = "dask"
 
     async def pre(self, *args, **kwargs) -> Dict[str, Any]:
-        from distributed import Client
-        from distributed.diagnostics.plugin import UploadDirectory
-
-        if flyte.ctx().is_in_cluster() and flyte.ctx().code_bundle:
+        ctx = flyte.ctx()
+        code_bundle = ctx.code_bundle
+        if ctx.is_in_cluster() and code_bundle:
             client = Client()
-            client.register_plugin(UploadDirectory(flyte.ctx().code_bundle.destination))
+            client.register_plugin(DownloadCodeBundlePlugin(code_bundle))
 
         return {}
 

{flyteplugins_dask-2.0.0b23.dist-info → flyteplugins_dask-2.0.0b24.dist-info}/METADATA

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: flyteplugins-dask
-Version: 2.0.0b23
+Version: 2.0.0b24
 Summary: Dask plugin for flyte
 Author-email: Kevin Su <pingsutw@users.noreply.github.com>
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: dask[distributed]>=2022.10.2
-Requires-Dist: flyte
+Requires-Dist: flyte>v2.0.0b23
 Requires-Dist: bokeh
 
 # Union Dask Plugin

flyteplugins_dask-2.0.0b24.dist-info/RECORD

@@ -0,0 +1,6 @@
+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,,

flyteplugins_dask-2.0.0b23.dist-info/RECORD

@@ -1,6 +0,0 @@
-flyteplugins/dask/__init__.py,sha256=UwmJjPRVqDyCTobQlT5a2g2oZJvK0vlDiv6rwugNs-g,112
-flyteplugins/dask/task.py,sha256=5jYtNY0wHE_EwoTYZDQ1vdlWjfrBzix3fxMz3Q4Rivo,3307
-flyteplugins_dask-2.0.0b23.dist-info/METADATA,sha256=pn5seXlxgSZGEPX_KVRp79UQ0ALQYo8G3UDg4jQYP7Y,938
-flyteplugins_dask-2.0.0b23.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-flyteplugins_dask-2.0.0b23.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
-flyteplugins_dask-2.0.0b23.dist-info/RECORD,,