Skip to content

Commit a735a62

Browse files
rohitrshRohit Sharmasamhita-alla
authored
Add Databricks Serverless Compute Support (#3392)
* feat(spark): Add Databricks Serverless Compute support Signed-off-by: Rohit Sharma <rohitrsh@gmail.com> * refactor(spark): Address review comments on serverless support - Extract _build_notebook_job_spec() and _build_python_file_job_spec() - Move compute validation to beginning of _get_databricks_job_spec() - Add API docs link and config format to _configure_serverless docstring - Create shared utils.py with is_serverless_config() for reuse in task.py Signed-off-by: Rohit Sharma <rohitrsh@gmail.com> * refactor(spark): Address review comments - value checks, serverless flag, lint fixes - Switch validation to value-based .get() checks so None values are rejected - Combine serverless config flag with is_databricks check in task.py - Fix ruff, ruff-format, and pydoclint violations - Update tests for new serverless detection behavior - Update pydoclint-errors-baseline.txt (fixed DatabricksV2 docstring) Signed-off-by: Rohit Sharma <rohitrsh@gmail.com> --------- Signed-off-by: Rohit Sharma <rohitrsh@gmail.com> Co-authored-by: Rohit Sharma <rohsharma@expediagroup.com> Co-authored-by: Samhita Alla <aallasamhita@gmail.com>
1 parent c0921e9 commit a735a62

7 files changed

Lines changed: 1195 additions & 78 deletions

File tree

‎Dockerfile.connector‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,15 @@ ARG VERSION
77

88
RUN apt-get update && apt-get install build-essential -y \
99
&& pip install uv
10-
11-
RUN uv pip install --system --no-cache-dir -U flytekit[connector]==$VERSION \
10+
# Pin pendulum<3.0: Apache Airflow (via flytekitplugins-airflow) imports
11+
# pendulum.tz.timezone() at module load time (airflow/settings.py).
12+
# Pendulum 3.x changed the tz API, causing the connector to crash on startup:
13+
# airflow/settings.py → TIMEZONE = pendulum.tz.timezone("UTC") → AttributeError
14+
# Without this pin, uv resolves to pendulum 3.x which breaks the import chain:
15+
# pyflyte serve connector → load_implicit_plugins → airflow → pendulum → crash
16+
RUN uv pip install --system --no-cache-dir -U \
17+
"pendulum>=2.0.0,<3.0" \
18+
flytekit[connector]==$VERSION \
1219
flytekitplugins-airflow==$VERSION \
1320
flytekitplugins-bigquery==$VERSION \
1421
flytekitplugins-k8sdataservice==$VERSION \

‎plugins/flytekit-spark/flytekitplugins/spark/connector.py‎

Lines changed: 209 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,13 @@
1515
from flytekit.models.literals import LiteralMap
1616
from flytekit.models.task import TaskTemplate
1717

18+
from .utils import is_serverless_config as _is_serverless_config
19+
1820
aiohttp = lazy_module("aiohttp")
1921

2022
DATABRICKS_API_ENDPOINT = "/api/2.1/jobs"
2123
DEFAULT_DATABRICKS_INSTANCE_ENV_KEY = "FLYTE_DATABRICKS_INSTANCE"
24+
DEFAULT_DATABRICKS_SERVICE_CREDENTIAL_PROVIDER_ENV_KEY = "FLYTE_DATABRICKS_SERVICE_CREDENTIAL_PROVIDER"
2225

2326

2427
@dataclass
@@ -27,38 +30,219 @@ class DatabricksJobMetadata(ResourceMeta):
2730
run_id: str
2831

2932

33+
def _configure_serverless(databricks_job: dict, envs: dict) -> str:
34+
"""
35+
Configure serverless compute settings and return the environment_key to use.
36+
37+
Databricks serverless requires the ``environments`` array in the job submission.
38+
This function ensures the array exists and injects Flyte environment variables
39+
into the matching environment's ``spec.environment_vars``.
40+
41+
Reference: https://docs.databricks.com/api/workspace/jobs/submit
42+
43+
Expected ``environments`` format::
44+
45+
"environments": [
46+
{
47+
"environment_key": "<key>",
48+
"spec": {
49+
"client": "1",
50+
"dependencies": ["pandas==2.0.0"],
51+
"environment_vars": {"KEY": "VALUE"}
52+
}
53+
}
54+
]
55+
56+
Tasks reference an environment via their own ``environment_key`` field,
57+
analogous to how ``job_cluster_key`` links a task to a shared cluster.
58+
59+
Args:
60+
databricks_job (dict): The databricks job configuration dict.
61+
envs (dict): Environment variables to inject into the environment spec.
62+
63+
Returns:
64+
str: The environment_key to use in the task definition.
65+
"""
66+
environment_key = databricks_job.get("environment_key", "default")
67+
environments = databricks_job.get("environments", [])
68+
69+
# Check if environment already exists in the array
70+
env_exists = any(env.get("environment_key") == environment_key for env in environments)
71+
72+
if not env_exists:
73+
# Create the environment entry - Databricks serverless requires environments
74+
# to be defined in the job submission (not externally pre-configured)
75+
new_env = {
76+
"environment_key": environment_key,
77+
"spec": {
78+
"client": "1", # Required: Databricks serverless client version
79+
},
80+
}
81+
environments.append(new_env)
82+
databricks_job["environments"] = environments
83+
84+
# Inject Flyte environment variables into the environment spec
85+
for env in environments:
86+
if env.get("environment_key") == environment_key:
87+
spec = env.setdefault("spec", {})
88+
existing_env_vars = spec.get("environment_vars", {})
89+
# Merge Flyte env vars with any existing ones (Flyte vars take precedence)
90+
merged_env_vars = {**existing_env_vars, **{k: v for k, v in envs.items()}}
91+
spec["environment_vars"] = merged_env_vars
92+
break
93+
94+
# Remove environment_key from top level (it's now in the task definition)
95+
databricks_job.pop("environment_key", None)
96+
97+
return environment_key
98+
99+
100+
def _configure_classic_cluster(databricks_job: dict, custom: dict, container: typing.Any, envs: dict) -> None:
101+
"""
102+
Configure classic compute (existing cluster or new cluster).
103+
104+
Args:
105+
databricks_job (dict): The databricks job configuration dict.
106+
custom (dict): The custom config from task template.
107+
container (typing.Any): The container config from task template.
108+
envs (dict): Environment variables to inject.
109+
"""
110+
if databricks_job.get("existing_cluster_id") is not None:
111+
# Using an existing cluster, no additional configuration needed
112+
return
113+
114+
new_cluster = databricks_job.get("new_cluster")
115+
if new_cluster is None:
116+
return
117+
118+
if not new_cluster.get("docker_image"):
119+
new_cluster["docker_image"] = {"url": container.image}
120+
if not new_cluster.get("spark_conf"):
121+
new_cluster["spark_conf"] = custom.get("sparkConf", {})
122+
if not new_cluster.get("spark_env_vars"):
123+
new_cluster["spark_env_vars"] = {k: v for k, v in envs.items()}
124+
else:
125+
new_cluster["spark_env_vars"].update({k: v for k, v in envs.items()})
126+
127+
128+
def _build_notebook_job_spec(
129+
databricks_job: dict, custom: dict, container: typing.Any, envs: dict, is_serverless: bool
130+
) -> dict:
131+
"""Build the Databricks job spec for a notebook task."""
132+
notebook_path = custom["notebookPath"]
133+
notebook_base_parameters = custom.get("notebookBaseParameters", {})
134+
135+
notebook_task = {"notebook_path": notebook_path}
136+
if notebook_base_parameters:
137+
notebook_task["base_parameters"] = notebook_base_parameters
138+
139+
user_git_source = databricks_job.get("git_source")
140+
if user_git_source:
141+
notebook_task["source"] = "GIT"
142+
143+
if is_serverless:
144+
environment_key = _configure_serverless(databricks_job, envs)
145+
task_def = {
146+
"task_key": "flyte_notebook_task",
147+
"notebook_task": notebook_task,
148+
"environment_key": environment_key,
149+
}
150+
databricks_job["tasks"] = [task_def]
151+
else:
152+
_configure_classic_cluster(databricks_job, custom, container, envs)
153+
databricks_job["notebook_task"] = notebook_task
154+
155+
databricks_job.pop("git_source", None)
156+
if user_git_source:
157+
databricks_job["git_source"] = user_git_source
158+
159+
return databricks_job
160+
161+
162+
def _build_python_file_job_spec(
163+
databricks_job: dict, custom: dict, container: typing.Any, envs: dict, is_serverless: bool
164+
) -> dict:
165+
"""Build the Databricks job spec for a python file (spark_python_task)."""
166+
user_git_source = databricks_job.get("git_source")
167+
user_python_file = databricks_job.get("python_file")
168+
169+
default_git_source = {
170+
"git_url": "https://github.com/flyteorg/flytetools",
171+
"git_provider": "gitHub",
172+
"git_commit": "572298df1f971fb58c258398bd70a6372f811c96",
173+
}
174+
default_classic_python_file = "flytekitplugins/databricks/entrypoint.py"
175+
default_serverless_python_file = "flytekitplugins/databricks/entrypoint_serverless.py"
176+
177+
if is_serverless:
178+
git_source = user_git_source or default_git_source
179+
python_file = user_python_file or default_serverless_python_file
180+
181+
environment_key = _configure_serverless(databricks_job, envs)
182+
183+
parameters = list(container.args) if container.args else []
184+
185+
service_credential_provider = custom.get(
186+
"databricksServiceCredentialProvider", os.getenv(DEFAULT_DATABRICKS_SERVICE_CREDENTIAL_PROVIDER_ENV_KEY)
187+
)
188+
if service_credential_provider:
189+
parameters.append(f"--flyte-credential-provider={service_credential_provider}")
190+
191+
spark_python_task = {
192+
"python_file": python_file,
193+
"source": "GIT",
194+
"parameters": parameters,
195+
}
196+
197+
task_def = {
198+
"task_key": "flyte_task",
199+
"spark_python_task": spark_python_task,
200+
"environment_key": environment_key,
201+
}
202+
203+
databricks_job["tasks"] = [task_def]
204+
else:
205+
git_source = user_git_source or default_git_source
206+
python_file = user_python_file or default_classic_python_file
207+
208+
spark_python_task = {
209+
"python_file": python_file,
210+
"source": "GIT",
211+
"parameters": container.args,
212+
}
213+
214+
_configure_classic_cluster(databricks_job, custom, container, envs)
215+
databricks_job["spark_python_task"] = spark_python_task
216+
217+
databricks_job.pop("git_source", None)
218+
databricks_job.pop("python_file", None)
219+
databricks_job["git_source"] = git_source
220+
221+
return databricks_job
222+
223+
30224
def _get_databricks_job_spec(task_template: TaskTemplate) -> dict:
31225
custom = task_template.custom
32226
container = task_template.container
33227
envs = task_template.container.env
34228
envs[FLYTE_FAIL_ON_ERROR] = "true"
35229
databricks_job = custom["databricksConf"]
36-
if databricks_job.get("existing_cluster_id") is None:
37-
new_cluster = databricks_job.get("new_cluster")
38-
if new_cluster is None:
39-
raise ValueError("Either existing_cluster_id or new_cluster must be specified")
40-
if not new_cluster.get("docker_image"):
41-
new_cluster["docker_image"] = {"url": container.image}
42-
if not new_cluster.get("spark_conf"):
43-
new_cluster["spark_conf"] = custom.get("sparkConf", {})
44-
if not new_cluster.get("spark_env_vars"):
45-
new_cluster["spark_env_vars"] = {k: v for k, v in envs.items()}
46-
else:
47-
new_cluster["spark_env_vars"].update({k: v for k, v in envs.items()})
48-
# https://docs.databricks.com/api/workspace/jobs/submit
49-
databricks_job["spark_python_task"] = {
50-
"python_file": "flytekitplugins/databricks/entrypoint.py",
51-
"source": "GIT",
52-
"parameters": container.args,
53-
}
54-
databricks_job["git_source"] = {
55-
"git_url": "https://github.com/flyteorg/flytetools",
56-
"git_provider": "gitHub",
57-
# https://github.com/flyteorg/flytetools/commit/572298df1f971fb58c258398bd70a6372f811c96
58-
"git_commit": "572298df1f971fb58c258398bd70a6372f811c96",
59-
}
60230

61-
return databricks_job
231+
has_cluster = databricks_job.get("existing_cluster_id") is not None or databricks_job.get("new_cluster") is not None
232+
has_serverless = bool(databricks_job.get("environment_key") or databricks_job.get("environments"))
233+
if not has_cluster and not has_serverless:
234+
raise ValueError(
235+
"No compute configuration found in databricks_conf. "
236+
"Provide one of: 'existing_cluster_id' (classic), 'new_cluster' (classic), "
237+
"'environment_key' (serverless), or 'environments' (serverless)."
238+
)
239+
240+
is_serverless = _is_serverless_config(databricks_job)
241+
242+
if custom.get("notebookPath"):
243+
return _build_notebook_job_spec(databricks_job, custom, container, envs, is_serverless)
244+
245+
return _build_python_file_job_spec(databricks_job, custom, container, envs, is_serverless)
62246

63247

64248
class DatabricksConnector(AsyncConnectorBase):

0 commit comments

Comments
 (0)