flyteplugins-dask 2.0.0b11__py3-none-any.whl

flyteplugins/dask/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["Dask", "Scheduler", "WorkerGroup"]
+
+from flyteplugins.dask.task import Dask, Scheduler, WorkerGroup

flyteplugins/dask/task.py

@@ -0,0 +1,91 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+import flyte
+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 flyteidl.plugins.dask_pb2 import DaskJob, DaskScheduler, DaskWorkerGroup
+from google.protobuf.json_format import MessageToDict
+
+
+@dataclass
+class Scheduler:
+    """
+    Configuration for the scheduler pod
+
+    :param image: Custom image to use. If ``None``, will use the same image the task was registered with. Optional,
+        defaults to None. The image must have ``dask[distributed]`` installed and should have the same Python
+        environment as the rest of the cluster (job runner pod + worker pods).
+    :param resources: Resources to request for the scheduler pod. Optional, defaults to None.
+    """
+
+    image: Optional[str] = None
+    resources: Optional[Resources] = None
+
+
+@dataclass
+class WorkerGroup:
+    """
+    Configuration for a group of dask worker pods
+
+    :param number_of_workers: Number of workers to use. Optional, defaults to 1.
+    :param image: Custom image to use. If ``None``, will use the same image the task was registered with. Optional,
+        defaults to None. The image must have ``dask[distributed]`` installed. The provided image should have the
+        same Python environment as the job runner/driver as well as the scheduler.
+    :param resources: Resources to request for the worker pods. Optional, defaults to None.
+    """
+
+    number_of_workers: Optional[int] = 1
+    image: Optional[str] = None
+    resources: Optional[Resources] = None
+
+
+@dataclass
+class Dask:
+    """
+    Configuration for the dask task
+
+    :param scheduler: Configuration for the scheduler pod. Optional, defaults to ``Scheduler()``.
+    :param workers: Configuration for the pods of the default worker group. Optional, defaults to ``WorkerGroup()``.
+    """
+
+    scheduler: Scheduler = field(default_factory=lambda: Scheduler())
+    workers: WorkerGroup = field(default_factory=lambda: WorkerGroup())
+
+
+@dataclass(kw_only=True)
+class DaskTask(AsyncFunctionTaskTemplate):
+    """
+    Actual Plugin that transforms the local python code for execution within a spark context
+    """
+
+    plugin_config: Dask
+    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:
+            client = Client()
+            client.register_plugin(UploadDirectory(flyte.ctx().code_bundle.destination))
+
+        return {}
+
+    def custom_config(self, sctx: SerializationContext) -> Dict[str, Any]:
+        scheduler = self.plugin_config.scheduler
+        wg = self.plugin_config.workers
+
+        job = DaskJob(
+            scheduler=DaskScheduler(image=scheduler.image, resources=get_proto_resources(scheduler.resources)),
+            workers=DaskWorkerGroup(
+                number_of_workers=wg.number_of_workers, image=wg.image, resources=get_proto_resources(wg.resources)
+            ),
+        )
+
+        return MessageToDict(job)
+
+
+TaskPluginRegistry.register(Dask, DaskTask)

flyteplugins_dask-2.0.0b11.dist-info/METADATA

@@ -0,0 +1,24 @@
+Metadata-Version: 2.4
+Name: flyteplugins-dask
+Version: 2.0.0b11
+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: bokeh
+
+# Union Dask Plugin
+
+Flyte can execute `dask` jobs natively on a Kubernetes Cluster, which manages the virtual `dask` cluster's lifecycle
+(spin-up and tear down). It leverages the open-source Kubernetes Dask Operator and can be enabled without signing up
+for any service. This is like running a transient (ephemeral) `dask` cluster - a type of cluster spun up for a specific
+task and torn down after completion. This helps in making sure that the Python environment is the same on the job-runner
+(driver), scheduler and the workers.
+
+To install the plugin, run the following command:
+
+```bash
+pip install --pre flyteplugins-dask
+```

flyteplugins_dask-2.0.0b11.dist-info/RECORD

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

flyteplugins_dask-2.0.0b11.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flyteplugins_dask-2.0.0b11.dist-info/top_level.txt

@@ -0,0 +1 @@
+flyteplugins