Databricks Agent

[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/

[Future-Outlier]

Just revise the code for the agent server only, please review, thanks a lot!

[codecov]

Codecov Report

Patch and project coverage have no change.

Comparison is base (ba571cd) 18.73% compared to head (f074740) 18.74%.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1797   +/-   ##
=======================================
  Coverage   18.73%   18.74%           
=======================================
  Files         332      332           
  Lines       31390    31387    -3     
  Branches     3081     3081           
=======================================
  Hits         5882     5882           
+ Misses      25424    25421    -3     
  Partials       84       84           

Files Changed	Coverage Δ
flytekit/extend/backend/agent_service.py	0.00% <ø> (ø)

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[Future-Outlier]

The error handler is really amazing!
Thanks a lot for helping me!

[Future-Outlier]

@pingsutw Just updated the test to the latest version:
#1802.
Please review.
Thanks a lot!