#330 Databricks connector Pull request opened by pingsutw • Add databricks example (tested locally, Flyte will create a databricks job from the local desktop) • Add databricks connector & task • Add a new entrypoint for databricks job since databricks job cannot run the click command • Update

pre

method in the spark plugin, so the code bundle is copied to the executors if it doesn’t already exist. • Update AsyncConnectorExecutorMixin to render TaskTemplate, upload input, and download output if the raw output path is a remote path [Screenshot 2025-11-26 at 1 15 39 PM](https://private-user-images.githubusercontent.com/37936015/519405135-aa85e709-0cda-4f31-b72a-2aa65aef68b4.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NjQ3MDc4ODYsIm5iZiI6MTc2NDcwNzU4NiwicGF0aCI6Ii8zNzkzNjAxNS81MTk0MDUxMzUtYWE4NWU3MDktMGNkYS00ZjMxLWI3MmEtMmFhNjVhZWY2OGI0LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNTEyMDIlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjUxMjAyVDIwMzMwNlomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTYzNGJlNmQ3NDU0NzViNGFlMDMzYTc5OTMzMjExNDY1NmMxZWQ4NjE5YzliODg1MDg2MDZmNDI4NWIzYTcwMjAmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0In0.Y6z__19k5nyvKjrjCE8sZZZsC4H_gcRYujwNwhBhzok) import random from operator import add from flyteplugins.connectors.databricks.task import Databricks import flyte.remote from flyte.storage import S3 image = ( # https://hub.docker.com/r/databricksruntime/standard/tags flyte.Image.from_base("databricksruntime/standard:16.4-LTS") .clone(name="spark", registry="ghcr.io/flyteorg", registry_secret="docker-g") .with_apt_packages("git", "vim") .with_env_vars({"UV_PYTHON": "/databricks/python3/bin/python"}) .with_pip_packages( "git+https://github.com/flyteorg/flyte-sdk.git@62a1cf3bf4ba091beff2804263664e083038dd6c#subdirectory=plugins/connectors", pre=True, ) .with_pip_packages( "git+https://github.com/flyteorg/flyte-sdk.git@62a1cf3bf4ba091beff2804263664e083038dd6c#subdirectory=plugins/spark", pre=True, ) .with_local_v2() .with_pip_packages("nest-asyncio", "aiohttp", "click==8.1.6") .with_env_vars({"AWS_REGION": "us-west-1", "LOG_LEVEL": 10}) ) task_env = flyte.TaskEnvironment( name="get_pi", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image ) databricks_conf = Databricks( spark_conf={ "spark.driver.memory": "2000M", "spark.executor.memory": "1000M", "spark.executor.cores": "1", "spark.executor.instances": "2", "spark.driver.cores": "1", "spark.kubernetes.file.upload.path": "/opt/spark/work-dir", "spark.jars": "https://storage.googleapis.com/hadoop-lib/gcs/gcs-connector-hadoop3-latest.jar,https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/3.2.2/hadoop-aws-3.2.2.jar,https://repo1.maven.org/maven2/com/amazonaws/aws-java-sdk-bundle/1.12.262/aws-java-sdk-bundle-1.12.262.jar", }, executor_path="/databricks/python3/bin/python", databricks_conf={ "run_name": "flytekit databricks plugin example", # "new_cluster": { # "spark_version": "13.3.x-scala2.12", # "autoscale": { # "min_workers": 1, # "max_workers": 1, # }, # "node_type_id": "m6i.large", # "num_workers": 1, # "aws_attributes": { # "availability": "SPOT_WITH_FALLBACK", # "instance_profile_arn": "arnawsiam:339713193121instance-profile/databricks-demo", # "ebs_volume_type": "GENERAL_PURPOSE_SSD", # "ebs_volume_count": 1, # "ebs_volume_size": 100, # "first_on_demand": 1, # }, # }, "existing_cluster_id": "1113-204018-tb9vr2fm", "timeout_seconds": 3600, "max_retries": 1, }, databricks_instance="dbc-6f73e2e6-19e4.cloud.databricks.com", databricks_token="DATABRICKS_TOKEN", ) databricks_env = flyte.TaskEnvironment( name="databricks_env", resources=flyte.Resources(cpu=(1, 2), memory=("3000Mi", "5000Mi")), plugin_config=databricks_conf, image=image, depends_on=[task_env], ) def f(_): x = random.random() * 2 - 1 y = random.random() * 2 - 1 return 1 if x**2 + y**2 <= 1 else 0 def get_pi(count: int, partitions: int) -> float: return 4.0 * count / partitions @databricks_env.task async def hello_databricks_nested() -> float: partitions = 3 n = 1 * partitions spark = flyte.ctx().data["spark_session"] count = spark.sparkContext.parallelize(range(1, n + 1), partitions).map(f).reduce(add) print("getttttt pi") res = get_pi(count, partitions) print(f"result: {res}") return res async def test_main(): import flyte.storage as storage proto_str = b"".join( [c async for c in storage.get_stream(path="s3://my-v2-connector/a209f3771e880f2ea8eb44a24ca572cf/inputs.pb")] ) print(proto_str) if name == "__main__": flyte.init_from_config(storage=S3().auto()) run = flyte.with_runcontext(mode="local", raw_data_path="s3://my-v2-connector/").run(hello_databricks_nested) print(run) # print("run name:", run.name) # print("run urlllll:", run.url) # asyncio.run(test_main()) flyteorg/flyte-sdk