#94 Add serverless-compatible entrypoint for Databricks Pull request opened by rohitrsh Tracking issue flyteorg/flyte#6911 Why are the changes needed? Databricks Serverless Compute has restrictions that make the existing classic entrypoint (

entrypoint.py

) incompatible: | Restriction | Classic Behavior | Serverless Impact | | ------------------------ | --------------------------- | --------------------------------------------- | | No /root access | os.chdir("/root") | PermissionError blocks all tasks | | No instance metadata | AWS creds from EC2 metadata | S3 access fails no data I/O | | No spark global variable | Injected by cluster runtime | SparkSession unavailable in python_file tasks | | Spark Connect only | Full Spark (JVM) | RDD APIs fail, must use DataFrame/Connect | What changes were proposed in this pull request? Add

flyteplugins/databricks/entrypoint_serverless.py

a serverless-compatible entrypoint that complements the existing

entrypoint.py

(classic compute). The classic entrypoint remains unchanged. The serverless entrypoint: 1. Uses

/tmp

as working directory (always writable in serverless) 2. Configures AWS credentials via

dbutils.credentials.getServiceCredentialsProvider()

3. Pre-initializes SparkSession and stores it in

sys.modules

+

builtins

for reliable retrieval by flytekit's

current_context().spark_session

4. Executes tasks directly instead of

pyflyte-fast-execute

subprocess to preserve Python module state (and the SparkSession) The flytekit-spark connector (in a companion PR) will reference this file by default for serverless jobs, using the same pattern as classic: # Both use the same flytetools repo; only the python_file differs default_classic_python_file = "flytekitplugins/databricks/entrypoint.py" default_serverless_python_file = "flytekitplugins/databricks/entrypoint_serverless.py" New file

flyteplugins/databricks/entrypoint_serverless.py

(~440 lines) | Function | Purpose | | --------------------------------------- | ----------------------------------------------------------------------- | | main() | Orchestrates the full entrypoint flow | | setup_aws_credentials_from_databricks() | Gets AWS creds from Databricks service credentials | | setup_spark_session() | Pre-creates SparkSession and stores in sys.modules + builtins | | setup_environment() | Sets working directory to /tmp/flyte, sets env markers | | execute_flyte_command_inprocess() | Routes pyflyte-fast-execute / pyflyte-execute commands | | execute_flyte_task_directly() | Downloads tarball, imports task module, calls execute_task_cmd directly | | parse_credential_provider_from_args() | Parses --flyte-credential-provider=<name> from args | | download_and_extract_distribution() | Downloads and extracts task tarball from S3 | Repository structure after

flytetools/
└── flyteplugins/
    └── databricks/
        ├── entrypoint.py              # Existing  classic compute (unchanged)
        └── entrypoint_serverless.py   # NEW  serverless compute

How was this patch tested? • Verified

entrypoint_serverless.py

is correctly placed alongside

entrypoint.py

• No changes to existing

entrypoint.py

• Manual test: Ran Flyte Spark tasks on Databricks Serverless using this entrypoint • Verified AWS credential setup succeeds with a valid service credential provider • Verified SparkSession is available via

flytekit.current_context().spark_session

• Verified the entrypoint handles both

pyflyte-fast-execute

and

pyflyte-execute

commands Setup process No setup needed. This adds a single new file no dependencies or configuration changes. Check all the applicable boxes • I updated the documentation accordingly. • All new and existing tests passed. • All commits are signed-off. Related PRs • Companion PR (flytekit): [link to flytekit PR] Adds serverless support to the flytekit-spark connector (To be update later) Docs link • Databricks Serverless Compute • Databricks Serverless Limitations flyteorg/flytetools