feat: add GPU spot instance support for Azure node pools #1837 by zeel2104 · Pull Request #1870 · kaito-project/kaito

zeel2104 · 2026-03-24T04:36:08Z

Reason for Change:
GPU instances are expensive and have limited quota. This PR adds support for running KAITO workloads on Azure Spot GPU node pools to reduce costs in both dev and production environments.
Requirements

[X] added unit tests and e2e tests (if applicable).

Issue Fixed:

Fixes #1837

Notes for Reviewers:

-> Added a spotInstance.enabled flag to values.yaml (defaults to false, no breaking change for existing users)
-> When enabled, adds the kubernetes.azure.com/scalesetpriority=spot:NoSchedule toleration to the nvidia-device-plugin DaemonSet so it can be scheduled on Spot node pools
-> The workspace Deployment already supports tolerations dynamically via Values.tolerations so no changes needed there
-< To enable spot support, users simply set spotInstance.enabled: true in their values override
-> No existing Helm tests found in the chart. Validated via helm lint and helm template --set spotInstance.enabled=true dry-ruN

kaito-pr-agent · 2026-03-24T04:36:50Z

Title

Add Azure GPU Spot Instance Support

Description

Add spotInstance flag to Helm values.yaml.
Add Azure Spot toleration to device plugin.
Enable cost reduction for GPU workloads.
Default behavior remains unchanged with spot disabled.

Changes walkthrough 📝

Relevant files

Enhancement

nvidia-device-plugin-ds.yaml `Add Spot Tolerations to Device Plugin` charts/kaito/workspace/templates/nvidia-device-plugin-ds.yaml Added conditional toleration for Azure Spot priority. Wrapped toleration in `spotInstance.enabled` check.	+7/-1

Configuration changes

values.yaml `Add Spot Instance Configuration Flag` charts/kaito/workspace/values.yaml Added `spotInstance` configuration section. Set default enabled value to false.	+8/-1

Need help?
Type /help how to ... in the comments thread for any questions about PR-Agent usage.
Check out the documentation for more information.

kaito-pr-agent · 2026-03-24T04:38:02Z

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1837 - Partially compliant Compliant requirements: Add toleration for Azure Spot instances to nvidia-device-plugin-daemonset DaemonSet. Add `spotInstance.enabled` flag to `values.yaml`. Non-compliant requirements: Add unit tests and e2e tests (PR states no tests found/added). Requires further human verification: Add toleration for Azure Spot instances to workspace StatefulSet (PR claims existing support via `Values.tolerations`, but template not shown in diff).
⏱️ Estimated effort to review: 2 🔵🔵⚪⚪⚪
🧪 No relevant tests
🔒 No security concerns identified
⚡ Recommended focus areas for review StatefulSet Tolerations The PR description claims the workspace StatefulSet already supports tolerations via `Values.tolerations`, but the diff does not include the StatefulSet template. Verification is needed to ensure the StatefulSet actually applies these tolerations. {{- if .Values.spotInstance.enabled }} - key: kubernetes.azure.com/scalesetpriority operator: Equal value: spot effect: NoSchedule {{- end }} Test Coverage Gap The ticket requires unit and e2e tests, but the PR description states no tests were added. This requirement is not met. spotInstance: enabled: false Unrelated Change The diff shows a change to `observability.metrics.enabled` (line 154) which appears to be a whitespace or no-op change unrelated to the feature. enabled: false

zhuangqh · 2026-03-24T05:34:45Z

          operator: "Equal"
          value: "gpu"
          effect: "NoSchedule"
+        {{- if .Values.spotInstance.enabled }}


I think it doesnt hurt to add this toleration without toggle?

Agree at the end customer could use node labels/selectors. And Spot instances are legitimate target for running AI models.

Fei-Guo · 2026-03-24T23:17:03Z

If you check the issue, we would also need to add toleration in the sts pod template otherwise, the pod cannot use the spot gpu.

Fei-Guo · 2026-03-25T16:46:04Z

@zeel2104, the toleration we add is specific for azure. Please add a guard in both yaml and controller to add the toleration for azure cloud provider only.

HoussemDellai · 2026-03-25T17:11:23Z

@zeel2104, the toleration we add is specific for azure. Please add a guard in both yaml and controller to add the toleration for azure cloud provider only.

Oh good point. KAITO is meant to be multi-provider and the Spot tolerations might be different across different providers. At the end may be it is good to add an input variable in values.yaml for 'custom'/provider-specific tolerations.

zeel2104 · 2026-03-25T21:50:24Z

@Fei-Guo
Updated the PR to guard the Azure spot toleration by cloud provider in both the Helm chart and controller-generated workloads.

Validation:

go test ./pkg/workspace/inference ./pkg/workspace/tuning
helm template checks:
- cloudProviderName=azure, spotInstance.enabled=true -> toleration present
- cloudProviderName=aws, spotInstance.enabled=true -> toleration absent
- cloudProviderName=azure, spotInstance.enabled=false -> toleration absent
  Let me know if any modifications need to be done

andyzhangx

this PR makes it possible to run KAITO workspace on existing spot node pool, that's BYO node scenario, I think you need to mention this in PR description.

Moreover, shall we consider adding spot node pool support in node auto-provisioning scenario?

andyzhangx · 2026-04-03T11:26:35Z

+# Set enabled: true to allow workloads to be scheduled on Azure Spot nodes.
+# This adds the required toleration for Azure Spot node pools.
+spotInstance:
+  enabled: false


just set default as true since it would be easier to covered in e2e test and no harm if it's supporting spot by default?

Copilot

Pull request overview

Adds Azure Spot GPU support by introducing the required tolerations so KAITO components can run on Azure Spot-tainted node pools, reducing GPU cost/pressure for dev and production clusters.

Changes:

Adds Azure Spot toleration constants and cloud-provider detection helper.
Updates inference/tuning pod generation to include Azure Spot toleration.
Adds Helm spotInstance.enabled value to conditionally tolerate Spot taint for the NVIDIA device plugin DaemonSet.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
pkg/workspace/tuning/preset_tuning.go	Builds default tolerations (now includes Azure Spot) and switches output dir path handling.
pkg/workspace/tuning/preset_tuning_test.go	Adds unit test coverage for default tolerations behavior.
pkg/workspace/inference/template_inference.go	Uses computed default tolerations when generating template inference manifests.
pkg/workspace/inference/preset_inferences.go	Builds default tolerations (now includes Azure Spot) for inference pods.
pkg/workspace/inference/preset_inferences_test.go	Adds unit test coverage for default tolerations behavior.
pkg/utils/consts/consts.go	Introduces constants for Azure Spot taint key/value.
pkg/utils/common.go	Adds `IsAzureCloudProvider()` helper based on `CLOUD_PROVIDER`.
charts/kaito/workspace/values.yaml	Adds `spotInstance.enabled` Helm value with documentation.
charts/kaito/workspace/templates/nvidia-device-plugin-ds.yaml	Conditionally adds Azure Spot toleration to the device plugin DaemonSet.

Comments suppressed due to low confidence (1)

pkg/workspace/tuning/preset_tuning.go:103

PrepareOutputDir intends to constrain outputDir under /mnt, but the current prefix check is not path-boundary safe (e.g., /mntx passes HasPrefix("/mnt")). This can allow paths outside the intended base directory. Use a boundary-aware check (e.g., require /mnt/ prefix) or compute a path.Rel to ensure the cleaned path does not escape the base dir.

	cleanPath := outputDir
	if !strings.HasPrefix(cleanPath, DefaultBaseDir) {
		cleanPath = path.Join(DefaultBaseDir, outputDir)
	}
	cleanPath = path.Clean(cleanPath)
	if cleanPath == DefaultBaseDir || !strings.HasPrefix(cleanPath, DefaultBaseDir) {
		klog.InfoS("Invalid output_dir specified: '%s', must be a directory. Using default output_dir: %s", outputDir, DefaultOutputVolumePath)
		return DefaultOutputVolumePath, fmt.Errorf("invalid output_dir specified: '%s', must be a directory", outputDir)
	}

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-03T11:31:04Z

+	if utils.IsAzureCloudProvider() {
+		tolerations = append(tolerations, corev1.Toleration{
+			Effect:   corev1.TaintEffectNoSchedule,
+			Key:      consts.SpotInstanceKey,
+			Operator: corev1.TolerationOpEqual,
+			Value:    consts.SpotInstanceValue,
+		})
+	}


Spot toleration is added for all Azure clusters purely based on CLOUD_PROVIDER, regardless of whether the user opted into Spot support. This can unintentionally allow inference pods to land on Spot GPU node pools when both Spot and on-demand pools exist, which conflicts with the Helm-level spotInstance.enabled flag and the PR description’s “opt-in” behavior. Consider gating this behind an explicit setting (e.g., env var wired from Helm values, a Workspace spec field, or a feature gate) rather than cloud-provider detection alone.

Copilot · 2026-04-03T11:31:04Z

+	if utils.IsAzureCloudProvider() {
+		tolerations = append(tolerations, corev1.Toleration{
+			Effect:   corev1.TaintEffectNoSchedule,
+			Key:      consts.SpotInstanceKey,
+			Operator: corev1.TolerationOpEqual,
+			Value:    consts.SpotInstanceValue,
+		})
+	}


Spot toleration is added for all Azure clusters purely based on CLOUD_PROVIDER, not an explicit opt-in. This can broaden scheduling to Spot GPU node pools even when users did not enable Spot support (and Helm’s spotInstance.enabled currently only affects the device plugin DaemonSet). Consider guarding this behind an explicit configuration flag that’s plumbed into the controller/workspace generation path.

Copilot · 2026-04-03T11:31:04Z

+# Set enabled: true to allow workloads to be scheduled on Azure Spot nodes.
+# This adds the required toleration for Azure Spot node pools.


The spotInstance.enabled comment says it “allow[s] workloads to be scheduled on Azure Spot nodes”, but this chart value only toggles an extra toleration on the nvidia-device-plugin DaemonSet. Consider clarifying the wording to reflect what it actually affects (device plugin scheduling), and/or document how workspace pods opt into Spot toleration (if that’s intended to be separate).

Suggested change

# Set enabled: true to allow workloads to be scheduled on Azure Spot nodes.

# This adds the required toleration for Azure Spot node pools.

# Set enabled: true to add the required toleration to the NVIDIA device plugin

# DaemonSet so it can run on Azure Spot node pools. Workspace pods must opt

# into Azure Spot toleration separately if they also need to run on Spot nodes.

Copilot · 2026-04-03T11:31:05Z

-	cleanPath = filepath.Clean(cleanPath)
+	cleanPath = path.Clean(cleanPath)
 	if cleanPath == DefaultBaseDir || !strings.HasPrefix(cleanPath, DefaultBaseDir) {
 		klog.InfoS("Invalid output_dir specified: '%s', must be a directory. Using default output_dir: %s", outputDir, DefaultOutputVolumePath)


klog.InfoS does not apply printf-style formatting; the %s placeholders will be logged literally. Switch to structured logging (message + key/value pairs) or use klog.Infof if formatting is desired.

Suggested change

klog.InfoS("Invalid output_dir specified: '%s', must be a directory. Using default output_dir: %s", outputDir, DefaultOutputVolumePath)

klog.InfoS("Invalid output_dir specified, using default output_dir", "outputDir", outputDir, "defaultOutputDir", DefaultOutputVolumePath)

feat: add GPU spot instance support for Azure node pools

cd11618

zeel2104 requested review from Fei-Guo, chewong and zhuangqh as code owners March 24, 2026 04:36

github-project-automation Bot added this to KAITO Roadmap Mar 24, 2026

kaito-pr-agent Bot added the Review effort 2/5 label Mar 24, 2026

zhuangqh reviewed Mar 24, 2026

View reviewed changes

zeel2104 and others added 2 commits March 25, 2026 12:33

fix: add spot toleration to workspace StatefulSet and tuning pod spec

2b49d90

Merge branch 'main' into feature/gpu-spot-instance-support

d298615

fix: guard Azure spot tolerations by cloud provider

c83a98a

andyzhangx reviewed Apr 3, 2026

View reviewed changes

andyzhangx requested a review from Copilot April 3, 2026 11:26

Copilot started reviewing on behalf of andyzhangx April 3, 2026 11:27 View session

Copilot AI reviewed Apr 3, 2026

View reviewed changes

Fei-Guo approved these changes Apr 8, 2026

View reviewed changes

Merge branch 'main' into feature/gpu-spot-instance-support

cd8c80d

Fei-Guo had a problem deploying to e2e-test April 8, 2026 03:21 — with GitHub Actions Failure

Fei-Guo temporarily deployed to unit-tests April 8, 2026 03:21 — with GitHub Actions Inactive

Fei-Guo merged commit 4dc9815 into kaito-project:main Apr 8, 2026
16 of 19 checks passed

github-project-automation Bot moved this to Done in KAITO Roadmap Apr 8, 2026

Navigation Menu

Uh oh!

feat: add GPU spot instance support for Azure node pools #1837 - #1870

feat: add GPU spot instance support for Azure node pools #1837#1870
Fei-Guo merged 5 commits into
kaito-project:mainfrom
zeel2104:feature/gpu-spot-instance-support

zeel2104 commented Mar 24, 2026

kaito-pr-agent Bot commented Mar 24, 2026

kaito-pr-agent Bot commented Mar 24, 2026

zhuangqh Mar 24, 2026

HoussemDellai Mar 25, 2026

Fei-Guo commented Mar 24, 2026

Fei-Guo commented Mar 25, 2026

HoussemDellai commented Mar 25, 2026

zeel2104 commented Mar 25, 2026

andyzhangx left a comment

andyzhangx Apr 3, 2026

zhuangqh Apr 8, 2026

Copilot AI left a comment

Copilot AI Apr 3, 2026

Copilot AI Apr 3, 2026

Copilot AI Apr 3, 2026

Copilot AI Apr 3, 2026

Uh oh!

Labels

6 participants

		# Set enabled: true to allow workloads to be scheduled on Azure Spot nodes.
		# This adds the required toleration for Azure Spot node pools.

-# Set enabled: true to allow workloads to be scheduled on Azure Spot nodes.
-# This adds the required toleration for Azure Spot node pools.
+# Set enabled: true to add the required toleration to the NVIDIA device plugin
+# DaemonSet so it can run on Azure Spot node pools. Workspace pods must opt
+# into Azure Spot toleration separately if they also need to run on Spot nodes.

	klog.InfoS("Invalid output_dir specified: '%s', must be a directory. Using default output_dir: %s", outputDir, DefaultOutputVolumePath)
	klog.InfoS("Invalid output_dir specified, using default output_dir", "outputDir", outputDir, "defaultOutputDir", DefaultOutputVolumePath)

Uh oh!

Conversation

zeel2104 commented Mar 24, 2026

kaito-pr-agent Bot commented Mar 24, 2026

Title

Description

Changes walkthrough 📝

kaito-pr-agent Bot commented Mar 24, 2026

PR Reviewer Guide 🔍

zhuangqh Mar 24, 2026

Choose a reason for hiding this comment

HoussemDellai Mar 25, 2026

Choose a reason for hiding this comment

Fei-Guo commented Mar 24, 2026

Fei-Guo commented Mar 25, 2026

HoussemDellai commented Mar 25, 2026

zeel2104 commented Mar 25, 2026

andyzhangx left a comment

Choose a reason for hiding this comment

andyzhangx Apr 3, 2026

Choose a reason for hiding this comment

zhuangqh Apr 8, 2026

Choose a reason for hiding this comment

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Copilot AI Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

Labels

6 participants