Creating A Spark Task

import datetime
import random
from operator import add

from truefoundry.deploy.v2.lib.patched_models import (
    Resources,
    SparkDriverConfig,
    SparkExecutorConfig,
    SparkExecutorFixedInstances,
    TaskPySparkBuild,
)
from truefoundry.workflow import PySparkTaskConfig, task, workflow, PythonTaskConfig, TaskPythonBuild, ExecutionConfig

@task(
    task_config=PySparkTaskConfig(
        image=TaskPySparkBuild(
            pip_packages=[
                "truefoundry[workflow,spark]",
                "setuptools"
            ],
            # Optional: specify a custom Spark base image instead of the default (public.ecr.aws/bitnami/spark)
            # container_image="your-registry.example.com/custom-spark:3.5.2",
        ),
        driver_config=SparkDriverConfig(
            resources=Resources(
                cpu_request=1.0,
                cpu_limit=2.0,
                memory_request=2048,
                memory_limit=2048,
                ephemeral_storage_request=2000,
                ephemeral_storage_limit=4000,
            )
        ), 
        # Executor configuration (required)
        executor_config=SparkExecutorConfig(
            instances=SparkExecutorFixedInstances(count=10),
            resources=Resources(
                cpu_request=1.0,
                cpu_limit=1.0,
                memory_request=1024,
                memory_limit=1024,
                ephemeral_storage_request=1000,
                ephemeral_storage_limit=2000,
            ),
        ),
        spark_conf={
            "spark.sql.adaptive.enabled": "true",
            "spark.sql.adaptive.coalescePartitions.enabled": "true",
            "spark.ui.enabled": "true",
            "spark.ui.proxyRedirectUri": "/",
        },
        # Replace with your k8s Service Account
        service_account="k8s-service-account",
    ), retries=3
)
def hello_spark(partitions: int) -> float:
    # generate a random number between 0 and 1
    import flytekit

    # Define map function inside task to avoid serialization issues when Flyte loads the task
    def _monte_carlo_point(_):
        x = random.random() * 2 - 1
        y = random.random() * 2 - 1
        return 1 if x**2 + y**2 <= 1 else 0

    print("Starting Spark with Partitions: {}".format(partitions))

    n = 1 * partitions
    sess = flytekit.current_context().spark_session
    count = (
        sess.sparkContext.parallelize(range(1, n + 1), partitions)
        .map(_monte_carlo_point)
        .reduce(add)
    )

    pi_val = 4.0 * count / n
    return pi_val

@task(
    task_config=PythonTaskConfig(
        image=TaskPythonBuild(
            python_version="3.9",
            pip_packages=["truefoundry[workflow,spark]"],
        ),
        resources=Resources(cpu_request=0.5),
    ),
)
def print_every_time(value_to_print: float, date_triggered: datetime.datetime) -> int:
    print("My printed value: {} @ {}".format(value_to_print, date_triggered))
    return 1

# Finally, define a workflow that connects your tasks in a sequence.
@workflow(
    # Schedule the workflow using a cron expression (runs every minute here)
    execution_configs=[
        ExecutionConfig(
            schedule="*/1 * * * *",
        )
    ],
    # Replace with your k8s Service Account
    service_account="k8s-service-account",
)
def my_spark(
    triggered_date: datetime.datetime = datetime.datetime(2020, 9, 11),
):
    """
    Using the workflow is still as any other workflow. As image is a property of the task, the workflow does not care
    about how the image is configured.
    """
    # task1 → task2 → task3 (sequential via data dependencies)
    pi = hello_spark(partitions=10000)
    print_every_time(value_to_print=pi, date_triggered=triggered_date)