#1797 Databricks Agent Pull request opened by Future-Outlier TL;DR As title. Type ☐ Bug Fix ☑︎ Feature ☐ Plugin Are all requirements met? ☑︎ Code completed ☐ Smoke tested ☑︎ Unit tests added ☐ Code documentation added ☐ Any pending items have an associated Issue Complete description I prefer using REST API because it's easier to maintain compared to the Databricks Python SDK. For the creation of clusters and jobs in Databricks, refer to the relevant sections in the Databricks REST APIs: https://docs.databricks.com/en/workflows/jobs/jobs-2.0-api.html#request-structure These sections indicate numerous arguments that can be passed. Let's discuss which ones to include in the Databricks agent. Here's a example to use the databricks agent:

SECRET_GROUP = "token-info"
SECRET_NAME = "token_secret"

@task(
    secret_requests=[Secret(group=SECRET_GROUP, 
                              key=SECRET_NAME,
                              mount_requirement=Secret.MountType.ENV_VAR,
                        )],
    task_config=DatabricksAgentTask(
        # this configuration is applied to the spark cluster
        spark_conf={
            "spark.driver.memory": "1000M",
            "spark.executor.memory": "1000M",
            "spark.executor.cores": "1",
            "spark.executor.instances": "2",
            "spark.driver.cores": "1",
        },
        executor_path="/databricks/python3/bin/python",
        applications_path="dbfs:/entrypoint.py",
        databricks_conf={
            "run_name": "flytekit databricks plugin example",
            "new_cluster": {
                "spark_version": "12.2.x-scala2.12",
                "node_type_id": "n2-highmem-4",
                "num_workers": 1,
            },
            "timeout_seconds": 3600,
            "max_retries": 1,
        },
        databricks_instance=os.environ["DATABRICKS_HOST"],
    ),
    container_image="flyteorg/flytekit:databricks-0.18.0-py3.7",
)
def hello_spark(partitions: int) -> float:
    print("Starting Spark with Partitions: {}".format(partitions))

    n = 100000 * partitions
    sess = flytekit.current_context().spark_session
    count = (
        sess.sparkContext.parallelize(range(1, n + 1), partitions).map(f).reduce(add)
    )
    pi_val = 4.0 * count / n
    print("Pi val is :{}".format(pi_val))
    return pi_val


def f(_):
    x = random.random() * 2 - 1
    y = random.random() * 2 - 1
    return 1 if x**2 + y**2 <= 1 else 0


@task(cache_version="1")
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

@workflow
def my_databricks_job(
    triggered_date: datetime.datetime = datetime.datetime.now(),
) -> float:
    """
    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.
    """
    pi = hello_spark(partitions=50)
    print_every_time(value_to_print=pi, date_triggered=triggered_date)
    return pi

Tracking Issue https://github.com/flyteorg/flyte/issues/ flyteorg/flytekit ✅ All checks have passed 30/30 successful checks