Skip to content

Add a tutorial that runs GPU training jobs on a Spot GPU node pool - #52

Open
paolosalvatori wants to merge 1 commit into
mainfrom
smf-1055-gpu-node-pools-tutorial
Open

paolosalvatori wants to merge 1 commit into
mainfrom
smf-1055-gpu-node-pools-tutorial

Conversation

@paolosalvatori

Copy link
Copy Markdown
Collaborator

Motivation

The repository has no GPU example.
On AKS, GPU workloads run on node pools of NVIDIA GPU VM sizes, and a Spot pool that scales from zero is the cheap way to run them, at the price of evictions.
LocalStack for Azure now emulates that flow: GPU nodes carry the labels and taints AKS gives them and advertise the GPUs of their size, the cluster autoscaler grows a pool from zero for a pod that asks for a GPU, and az vmss simulate-eviction evicts a Spot node as AKS does (see Run GPU workloads).
This tutorial walks through that flow with training jobs that KEDA starts for each message on a Service Bus queue, and runs unchanged on Azure, where the job trains on the GPU, and on the emulator, where it trains on the CPU.

Fixes SMF-1055

Changes

tutorials/gpu-node-pools/, in the shape of the other tutorials, for the cluster that scripts/01-user-assigned-managed-identity.sh creates:

  • 01-add-gpu-node-pool.sh adds gpupool: Spot Standard_NV18ads_A10_v5 VMs (eviction policy Delete, maximum price -1) on Azure Linux, the cluster autoscaler from 0 to 1 node, the label workload=gpu and the taint sku=gpu:NoSchedule. It enables the KEDA add-on if the cluster lacks it, and stops when a pool of that name exists with another VM size or priority, since neither can change.
  • 02-deploy-device-plugin.sh deploys the NVIDIA device plugin DaemonSet of Microsoft's manual install guide (nvcr.io/nvidia/k8s-device-plugin:v0.18.0), plus a node selector for the accelerator=nvidia label AKS gives GPU nodes and a toleration for the Spot taint.
  • 03-create-resources.sh creates the Service Bus namespace and the training-requests queue, and the managed identity the KEDA tutorials share (local-keda-uami-test) unless one of them already did. The identity gets Azure Service Bus Data Owner on the namespace, granted on its principal id, and is federated with the KEDA operator and with the job's service account. The operator is bound to it, and restarted only when the binding changes.
  • 04-deploy-scaled-job.sh deploys the job's service account, the training script as a config map, the TriggerAuthentication and the ScaledJob: one job per message, whose pod asks for nvidia.com/gpu: 1, selects the pool's nodes and tolerates its taints. It derives the Service Bus endpoints when it applies them, so the same manifests serve both targets: on the emulator the job connects through a custom endpoint address, which makes the Python SDK use AMQP over WebSockets and so needs websocket-client, and the scaler reads the HTTPS API on the gateway port.
  • 05-run-training-job.sh sends a training request and follows its job: the autoscaler adds a GPU node, and the script prints the node's labels, taints and GPUs, then the job's log.
  • 06-simulate-eviction.sh sends a 10-minute training request, runs az vmss simulate-eviction on the node of the job's pod, follows the node until AKS deletes it, and waits for the job to complete on a new node.
  • 07-cleanup.sh deletes what the tutorial created, and keeps the identity the KEDA tutorials share.
  • src/train.py receives one message with DefaultAzureCredential, which signs in with the pod's workload identity, and trains a small PyTorch model on cuda when PyTorch finds a GPU, on cpu otherwise.
  • The tutorial's README.md, and its row in the root README.md.

The scripts follow the repository's conventions (#!/bin/bash, 00-variables.sh, manifests filled in with yq, a check after every command that can fail) and need nothing beyond bash 3.2.

Tests

  • Static: bash -n on every script with bash 5 and with GNU bash 3.2.57, the version macOS ships; shfmt -i 2 -d; src/train.py parses; the Mermaid diagram validates; the repository's pre-commit hooks (gitleaks passed, zizmor has no workflow to check).
  • LocalStack for Azure, the published image: localstack/localstack-azure:latest as pushed on 2026-11-04 at 14:42 UTC (sha256:31573c66...), the first latest with GPU node pool support, on a fresh emulator routed with lstk az start-interception. The emulator ran with docker run as the samples workflow's Windows job starts it, with the external service ports from 4512, because another program held port 4511 on the test machine and lstk start checks the whole range. scripts/01-user-assigned-managed-identity.sh created the cluster in 420 s, then the scripts as committed ran in order, with 01 to 03 run twice: every step exited 0, and the second runs changed nothing.
    • 05: the cluster autoscaler added a GPU node; the job received its message and trained on cpu (CUDA available: False), 563 s in all, most of it pulling the PyTorch image.
    • 06: AKS tainted the node remediator.kubernetes.azure.com/unschedulable:NoSchedule and deleted it 32 s after the eviction request; the replacement pod found the queue empty and completed on a new node.
    • 07 deleted everything it lists.
  • An earlier run on the emulator built from source with the same AKS code passed the same sequence.
  • Azure, italynorth, twice; the second run used the scripts as committed. The cluster came from az aks create with --enable-keda --enable-oidc-issuer --enable-workload-identity and the names of the repository's script, because scripts/01-user-assigned-managed-identity.sh needs the object id of a Microsoft Entra admin group, which the test service principal can't look up. Every step exited 0, and the second runs of 01 to 03 changed nothing.
    • 05: the cluster autoscaler added a Spot Standard_NV18ads_A10_v5 node, the device plugin advertised one nvidia.com/gpu, and the job trained on cuda on an NVIDIA A10-12Q, 254 s in all.
    • 06: the same eviction taint, the node deleted 158 s after the eviction request, and the replacement pod completed on a new node, on cuda.
    • The node used 18 of the subscription's 32 Spot vCPUs in the region, while the Standard NVADSA10v5 Family quota is 0, which Spot doesn't draw on. 07 and the deletion of the resource group left no resource behind and no Spot vCPU in use.
  • Parity: the pool 01 prints and the node 05 prints (labels, taints and GPUs) are identical on the emulator and on Azure, and so are the eviction taints.
  • Not run: on macOS beyond the bash 3.2 parse, from Git Bash on Windows, and on Azure with the repository's cluster script.

🤖 Generated with Claude Code

The tutorial adds a Spot NVIDIA GPU node pool that scales from zero to the cluster the repository's scripts create, and runs training jobs on it: KEDA starts a job for each training request on a Service Bus queue, the cluster autoscaler adds a GPU node for its pod, and the pod trains a small PyTorch model, signing in with Microsoft Entra Workload ID.
A second run simulates the eviction of the Spot node with az vmss simulate-eviction, and the job completes on the node that replaces it.
The scripts run unchanged on Azure, where the job trains on the GPU, and on LocalStack for Azure, whose GPU nodes advertise their GPUs without one behind them, so the job trains on the CPU.
The root README lists the tutorial.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Copilot AI balanced review requested due to automatic review settings November 4, 2026 18:43

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

The eviction script can evict a pod before training starts, falsely reporting recovery without interrupting training; it must wait for training progress and confirm the pod is still running on the selected node before eviction.

1 open issue
What changed in this PR

Adds the repository’s first GPU tutorial, showing how Service Bus requests trigger KEDA training jobs on an AKS Spot GPU pool that scales from zero. The same setup supports Azure GPU training and LocalStack CPU execution.

Changes:

  • Adds scripts and manifests for the GPU pool, NVIDIA device plugin, workload identity, and KEDA jobs.
  • Demonstrates training, simulated Spot eviction, replacement-job completion, and cleanup.
  • Documents setup, architecture, and Azure versus emulator behavior.
File Description
tutorials/​gpu-node-pools/​src/​train.py Receives a request and trains with PyTorch on GPU or CPU.
tutorials/​gpu-node-pools/​scripts/​triggerauthentication.yml Configures KEDA workload identity authentication.
tutorials/​gpu-node-pools/​scripts/​serviceaccount.yml Defines the training job’s service account.
tutorials/​gpu-node-pools/​scripts/​scaledjob.yml Defines GPU-requesting, queue-triggered training jobs.
tutorials/​gpu-node-pools/​scripts/​nvidia-device-plugin.yml Defines the NVIDIA device plugin DaemonSet.
tutorials/​gpu-node-pools/​scripts/​functions.sh Provides endpoint, messaging, monitoring, and logging helpers.
tutorials/​gpu-node-pools/​scripts/​07-cleanup.sh Removes tutorial resources while preserving the shared identity.
tutorials/​gpu-node-pools/​scripts/​06-simulate-eviction.sh Simulates Spot eviction and checks completion on another node.
tutorials/​gpu-node-pools/​scripts/​05-run-training-job.sh Sends a request and follows training to completion.
tutorials/​gpu-node-pools/​scripts/​04-deploy-scaled-job.sh Populates and deploys training resources and KEDA configuration.
tutorials/​gpu-node-pools/​scripts/​03-create-resources.sh Creates messaging resources and configures identity permissions.
tutorials/​gpu-node-pools/​scripts/​02-deploy-device-plugin.sh Deploys the NVIDIA device plugin.
tutorials/​gpu-node-pools/​scripts/​01-add-gpu-node-pool.sh Enables KEDA if needed and adds the autoscaling Spot GPU pool.
tutorials/​gpu-node-pools/​scripts/​00-variables.sh Centralizes tutorial settings and timeouts.
tutorials/​gpu-node-pools/​README.md Explains prerequisites, execution, architecture, and limitations.
README.md Adds the GPU tutorial to the repository’s tutorial index.

Review effort: Balanced

echo "No pod of the [$JOB] job runs after $JOB_START_TIMEOUT seconds"
exit 1
fi
echo "The pod of the [$JOB] job runs on the [$EVICTED_NODE] node"
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants