kubeflow / kubeflow/sdk

Add DynamicAllocation support to SparkClient

Open
#489 4 comments 1 reaction 1 assignee Claimed by @KartikSuryavanshi View on GitHub
area/spark kind/feature
Dominant language
Python
Stars
148
Forks
262
Avg merge
1d 2h
Merged PRs (30d)
1

Description

### What you would like to be added?

Expose Spark dynamic allocation on `SparkClient.connect()` when creating a Spark Connect session (create mode), aligned with the Spark Operator `SparkConnect` CR field `spec.dynamicAllocation`.

**API**
- Add a `DynamicAllocation` dataclass in `kubeflow.spark.types` (and export it from `kubeflow.spark`) with fields matching the CRD:
- `enabled`
- `initial_executors`, `min_executors`, `max_executors`
- `shuffle_tracking_enabled`, `shuffle_tracking_timeout`
- Add a `dynamic_allocation: DynamicAllocation | None = None` parameter to `SparkClient.connect()` and thread it through the Kubernetes backend into `SparkConnect` CR creation.
**Behavior**
- When `dynamic_allocation.enabled` is `True`, do **not** set a fixed `spec.executor.instances` (scaling is owned by dynamic allocation).
- When dynamic allocation is disabled or unset, keep current behavior (`num_executors`, `Executor.num_instances`, defaults).
- Map SDK fields to `SparkV1alpha1DynamicAllocation` when building the CR.
**Docs & tests**
- Unit tests for CR building (enabled/disabled, min/max/initial, executor instances omitted when enabled).
- Docstring example on `SparkClient.connect()`.
- Optional example under `examples/spark/ `showing dynamic allocation usage.

Example
```python
from kubeflow.spark import DynamicAllocation, SparkClient

client = SparkClient()
spark = client.connect(
dynamic_allocation=DynamicAllocation(
enabled=True,
initial_executors=2,
min_executors=1,
max_executors=20,
shuffle_tracking_enabled=True,
),
resources_per_executor={"cpu": "2", "memory": "4Gi"},
)
```

### Why is this needed?

Spark on Kubernetes supports **dynamic allocation** (scale executors with workload). The Spark Operator already exposes this on `SparkConnect` via `spec.dynamicAllocation`, but the Kubeflow SDK today only supports a fixed executor count (`num_executors` / `Executor.num_instances`).

Without first-class SDK support, users must:
- hand-author `SparkConnect` YAML, or
- rely on `spark_conf` alone, which does not set operator-level `spec.dynamicAllocation` and is easy to get wrong (e.g. fixed `executor.instances` conflicting with scaling).
That matters for:
- **Variable workload**s (notebooks, ad-hoc queries, pipelines with uneven stages) where a static executor count wastes resources or under-provisions.
- **Cost and cluster efficiency** on shared Kubeflow clusters (scale down idle executors, cap with `max_executors`).
- **Parity with the operator CRD** so SDK-created sessions behave like operator-native manifests.

### Love this feature?

Give it a 👍 We prioritize the features with most 👍

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.