flyteplugins-spark 0.2.0b30__py3-none-any.whl

flyteplugins/spark/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["Spark"]
+
+from flyteplugins.spark.task import Spark

flyteplugins/spark/task.py

@@ -0,0 +1,86 @@
+import os
+import shutil
+import tempfile
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from flyte import PodTemplate
+from flyte._initialize import is_initialized
+from flyte.extend import AsyncFunctionTaskTemplate, TaskPluginRegistry
+from flyte.models import SerializationContext
+from flyteidl.plugins.spark_pb2 import SparkApplication, SparkJob
+from google.protobuf.json_format import MessageToDict
+
+
+@dataclass
+class Spark(object):
+    """
+    Use this to configure a SparkContext for a your task. Task's marked with this will automatically execute
+    natively onto K8s as a distributed execution of spark
+
+    Attributes:
+        spark_conf (Optional[Dict[str, str]]): Spark configuration dictionary.
+        hadoop_conf (Optional[Dict[str, str]]): Hadoop configuration dictionary.
+        executor_path (Optional[str]): Path to the Python binary for PySpark execution.
+        applications_path (Optional[str]): Path to the main application file.
+        driver_pod (Optional[PodTemplate]): Pod template for the driver pod.
+        executor_pod (Optional[PodTemplate]): Pod template for the executor pods.
+    """
+
+    spark_conf: Optional[Dict[str, str]] = None
+    hadoop_conf: Optional[Dict[str, str]] = None
+    executor_path: Optional[str] = None
+    applications_path: Optional[str] = None
+    driver_pod: Optional[PodTemplate] = None
+    executor_pod: Optional[PodTemplate] = None
+
+    def __post_init__(self):
+        if self.spark_conf is None:
+            self.spark_conf = {}
+
+        if self.hadoop_conf is None:
+            self.hadoop_conf = {}
+
+
+@dataclass(kw_only=True)
+class PysparkFunctionTask(AsyncFunctionTaskTemplate):
+    """
+    Actual Plugin that transforms the local python code for execution within a spark context
+    """
+
+    plugin_config: Spark
+    task_type: str = "spark"
+
+    async def pre(self, *args, **kwargs) -> Dict[str, Any]:
+        import pyspark as _pyspark
+
+        sess = _pyspark.sql.SparkSession.builder.appName("FlyteSpark").getOrCreate()
+
+        if is_initialized():
+            base_dir = tempfile.mkdtemp()
+            file_name = "flyte_wf"
+            file_format = "zip"
+            shutil.make_archive(f"{base_dir}/{file_name}", file_format, os.getcwd())
+            sess.sparkContext.addPyFile(f"{base_dir}/{file_name}.{file_format}")
+
+        return {"spark_session": sess}
+
+    def custom_config(self, sctx: SerializationContext) -> Dict[str, Any]:
+        driver_pod = self.plugin_config.driver_pod.to_k8s_pod() if self.plugin_config.driver_pod else None
+        executor_pod = self.plugin_config.executor_pod.to_k8s_pod() if self.plugin_config.executor_pod else None
+
+        job = SparkJob(
+            sparkConf=self.plugin_config.spark_conf,
+            hadoopConf=self.plugin_config.hadoop_conf,
+            mainApplicationFile=self.plugin_config.applications_path,
+            executorPath=self.plugin_config.executor_path,
+            mainClass="",
+            applicationType=SparkApplication.PYTHON,
+            driverPod=driver_pod,
+            executorPod=executor_pod,
+        )
+
+        return MessageToDict(job)
+
+
+TaskPluginRegistry.register(Spark, PysparkFunctionTask)

flyteplugins_spark-0.2.0b30.dist-info/METADATA

@@ -0,0 +1,18 @@
+Metadata-Version: 2.4
+Name: flyteplugins-spark
+Version: 0.2.0b30
+Summary: Spark plugin for flyte
+Author-email: Kevin Su <pingsutw@users.noreply.github.com>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: pyspark
+
+# Union Spark Plugin
+
+Union can execute Spark jobs natively on a Kubernetes Cluster, which manages a virtual cluster’s lifecycle, spin-up, and tear down. It leverages the open-sourced Spark On K8s Operator and can be enabled without signing up for any service. This is like running a transient spark cluster — a type of cluster spun up for a specific Spark job and torn down after completion.
+
+To install the plugin, run the following command:
+
+```bash
+pip install flyteplugins-spark
+```

flyteplugins_spark-0.2.0b30.dist-info/RECORD

@@ -0,0 +1,6 @@
+flyteplugins/spark/__init__.py,sha256=7q06oXtVhOuNW4Nv-YhJvt8ORqhXmMPNYa-m4IkougI,63
+flyteplugins/spark/task.py,sha256=3fnw8QzcROAqmSAZIJLPc90Q-AvENV1fVdHChMFWdhs,3174
+flyteplugins_spark-0.2.0b30.dist-info/METADATA,sha256=PmUrpy-01jPgkUUYkcYhljrjbErp2CVrmeAfGviv4YU,736
+flyteplugins_spark-0.2.0b30.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+flyteplugins_spark-0.2.0b30.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_spark-0.2.0b30.dist-info/RECORD,,

flyteplugins_spark-0.2.0b30.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_spark-0.2.0b30.dist-info/top_level.txt

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