[Core feature] Add Databricks Serverless Compute Support to Databricks Connector
- Dominant language
- Go
- Stars
- 7.5k
- Forks
- 886
- Avg merge
- 1d 12h
- Merged PRs (30d)
- 120
Description
### Motivation: Why do you think this is important?
Add comprehensive support for **Databricks Serverless Compute** in the flytekit databricks connector, enabling Flyte tasks to run on auto-scaling, on-demand compute without cluster management.
## Motivation
Classic Databricks compute requires **10-15 minutes for cluster startup**, which significantly impacts:
- Developer productivity and iteration speed
- Time-to-production for ML experiments
- Resource costs (paying for idle cluster time)
Databricks Serverless provides:
- **Sub-minute job startup times**
- Automatic scaling based on workload
- Zero cluster management overhead
- Better cost efficiency (pay only for the compute used)
## Proposed Changes
### 1. Serverless Compute Detection and Job Format
**Auto-detect** serverless vs classic compute based on configuration:
| Configuration | Compute Mode |
|--------------|--------------|
| `existing_cluster_id` | Classic |
| `new_cluster` | Classic |
| `environment_key` (no cluster) | **Serverless** |
| `environments` array (no cluster) | **Serverless** |
Generate correct Databricks Jobs API 2.1 format for serverless (multi-task format with `tasks` array).
### 2. SparkSession Preservation
Databricks Serverless provides a pre-configured SparkSession via Spark Connect. However, Flytekit's `fast_execute_task_cmd` clears Python module state, losing the SparkSession.
**Solution**:
- Custom serverless-compatible entrypoint (`entrypoint_serverless.py`)
- Direct task execution to preserve module state
- SparkSession stored in `sys.modules` and `builtins` for reliable retrieval
- New `get_spark()` helper function for simplified access
### 3. S3 Access via Databricks Service Credentials
Serverless doesn't have AWS instance metadata. Solution:
- Use `dbutils.credentials.getServiceCredentialsProvider()` to obtain credentials
- New `databricks_service_credential_provider` configuration option
- Automatic credential injection as environment variables
### 4. Notebook Task Support
Enable scheduling Databricks notebooks from Flyte workflows:
- New `notebook_path` and `notebook_base_parameters` configuration options
- Generate `notebook_task` instead of `spark_python_task` when notebook_path is set
- Support both workspace and git-sourced notebooks
## API Examples
### Serverless Spark Task
```python
from flytekitplugins.spark import DatabricksV2, get_spark
@task(
task_config=DatabricksV2(
databricks_conf={
"run_name": "my-serverless-job",
"environment_key": "default",
"environments": [{
"environment_key": "default",
"spec": {
"client": "4",
"dependencies": ["pandas>=2.0.0"],
}
}],
"git_source": {
"git_url": "https://github.com/org/repo",
"git_provider": "gitHub",
"git_branch": "main",
},
},
databricks_service_credential_provider="my-s3-credential",
),
container_image="my-image:tag",
)
def serverless_spark_task(n: int) -> int:
spark = get_spark() # Simply call get_spark()
return spark.range(n).count()
```
### Notebook Task
```python
@task(
task_config=DatabricksV2(
databricks_conf={...},
notebook_path="/Workspace/Users/user@example.com/my-notebook",
notebook_base_parameters={"param1": "value1"},
),
)
def run_notebook() -> None:
pass # Notebook tasks return None
```
## Files Changed
| File | Change |
|------|--------|
| `connector.py` | Serverless detection, multi-task job format, notebook task support |
| `task.py` | New config options, documentation |
| `__init__.py` | Export `get_spark()` helper |
| `entrypoint_serverless.py` | **NEW** - Serverless-compatible entrypoint |
| `test_connector.py` | Comprehensive serverless tests |
## Testing
- Unit tests for serverless detection and job spec generation
- Unit tests for notebook task generation
- Integration tests with Databricks Serverless Compute
- Verified with:
- PySpark DataFrame operations
- SparkSQL queries
- Window functions and aggregations
- Notebook execution
## Results
| Metric | Before (Classic) | After (Serverless) |
|--------|------------------|---------------------|
| Job startup | 10-15 minutes | < 1 minute |
| Cluster management | Manual | Zero |
| Cost model | Per-cluster | Per-job |
## References
- Databricks Serverless: https://docs.databricks.com/en/compute/serverless.html
- Databricks Service Credentials: https://docs.databricks.com/en/dev-tools/unity-catalog/service-credentials.html
- Spark Connect: https://spark.apache.org/docs/latest/spark-connect-overview.html
FYI: @kumare3
### Goal: What should the final outcome look like, ideally?
All of the above changes have been tested, and the working connector version has been deployed to our environment. Let us know what you guys think about it. I can raise a PR for review.
### Describe alternatives you've considered
N/A
### Propose: Link/Inline OR Additional context
The proposal is to extend the Flyte Databricks connector to support Serverless APIs and to submit jobs in the Databricks console; details are provided in the Motivation section.
### Are you sure this issue hasn't been raised already?
- [x] Yes
### Have you read the Code of Conduct?
- [x] Yes
Contributor guide
Assessment
This issue has not been assessed yet.