[flyte-core] Flyte Connection

[pingsutw]

Tracking issue

flyteorg/flyte#3936

Why are the changes needed?

it will allow users to pass different API_KEY to the chatGPT agent from flytepropeller

What changes were proposed in this pull request?

• Add List[Secret] to the create/get/delete/executeSync request
• Read the secret in the base agent

How was this patch tested?

local/remote

Setup process

from flytekit import workflow, Secret
from flytekitplugins.chatgpt import ChatGPTTask

chatgpt_small_job = ChatGPTTask(
    name="3.5-turbo",
    openai_organization="org-P2rdnZQry4Fw7Ak3vSpXEIrx",
    openai_key=Secret("openai", "kevin"),
    chatgpt_config={
        "model": "gpt-3.5-turbo",
        "temperature": 0.7,
    },
)


@workflow
def my_chatgpt_job(message: str) -> str:
    message = chatgpt_small_job(message=message)
    return message

import datetime
import random
from operator import add

from click.testing import CliRunner

import flytekit
from flytekit import Resources, Secret, task, workflow, ImageSpec
from flytekit.clis.sdk_in_container import pyflyte
from flytekitplugins.spark import Databricks

image = ImageSpec(base_image="flyteorg/flytekit:spark-latest", registry="pingsutw")


@task(
    task_config=Databricks(
        # this configuration is applied to the spark cluster
        spark_conf={
            "spark.driver.memory": "600M",
            "spark.executor.memory": "600M",
            "spark.executor.cores": "1",
            "spark.executor.instances": "1",
            "spark.driver.cores": "1",
        },
        executor_path="/databricks/python3/bin/python",
        applications_path="dbfs:///FileStore/tables/entrypoint.py",
        databricks_conf={
            "run_name": "flytekit databricks plugin example",
            "job_id": "{{inputs.cluster_id}}",
            "new_cluster": {
                "spark_version": "13.3.x-scala2.12",
                "node_type_id": "m6i.large",  # TODO: test m6i.large, i3.xlarge
                "num_workers": 3,
                "aws_attributes": {
                    "availability": "SPOT_WITH_FALLBACK",
                    "instance_profile_arn": "arn:aws:iam::546011168256:instance-profile/databricks-demo",
                    "ebs_volume_type": "GENERAL_PURPOSE_SSD",
                    "ebs_volume_count": 1,
                    "ebs_volume_size": 100,
                    "first_on_demand": 1,
                },
                "policy_id": "0000FF579DC57A49"
            },
            # "existing_cluster_id": "1122-085119-brus1mft",
            "timeout_seconds": 3600,
            "max_retries": 1,
        },
        databricks_instance="dbc-429786b4-2d97.cloud.databricks.com",
    ),
    limits=Resources(mem="2000M"),
    container_image="pingsutw/databricks:v12",
    secret_requests=[Secret("databricks", "kevin")],
)
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 wf(
        triggered_date: datetime.datetime = datetime.datetime.now(),
) -> float:
    """
    Using the workflow is still as any other workflow. As image is a pr1operty of the task, the workflow does not care
    about how the image is configured.1
    """
    pi = hello_spark(partitions=50).with_overrides(name="bad cluster")
    print_every_time(value_to_print=pi, date_triggered=triggered_date)
    return pi

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

NA

Docs link

NA

[Future-Outlier]

Love it, will take a look

[Future-Outlier]

Works well to me in local.

[Future-Outlier]

I've failed in remote execution, but I think the reason is that I need more settings to let propeller can get my secret, right?
(which means kubectl create secret generic openai-remote\ --from-literal=hanru=xxx will not work in this case.

[Future-Outlier]

Coming, will do it today

[kumare3]

Looks nice, cc @EngHabu one comment.