Deallocate the GPU pool's idle node so later runs find the PyTorch image

tutorials/gpu-node-pools/README.md

@@ -61,11 +61,11 @@ Run the numbered scripts in order. Each one sources [00-variables.sh](scripts/00
 | --- | --- | --- |
 | `00-variables.sh` | Variables for this tutorial (the GPU node pool, the device plugin, Service Bus, the managed identity and the training job). Sourced by every script, never run directly. | Variables set. |
 | `functions.sh` | The functions the scripts share: reading the Service Bus endpoints, sending a training request, and following the jobs. Sourced by `04` to `06`. | Functions defined. |
-| `01-add-gpu-node-pool.sh` | Verifies the cluster exists with its OIDC issuer, enables the KEDA add-on if needed, merges the cluster credentials into `kubeconfig` (run it first), and adds the Spot GPU node pool with no node and the cluster autoscaler from 0 to 1 node. | The pool exists with `count: 0`, `scaleSetPriority: Spot` and the Spot label and taint AKS adds, and no node. |
+| `01-add-gpu-node-pool.sh` | Verifies the cluster exists with its OIDC issuer, enables the KEDA add-on if needed, merges the cluster credentials into `kubeconfig` (run it first), and adds the Spot GPU node pool with no node, the cluster autoscaler from 0 to 1 node, a managed OS disk and the `Deallocate` scale-down mode. | The pool exists with `count: 0`, `scaleSetPriority: Spot`, `scaleDownMode: Deallocate` and the Spot label and taint AKS adds, and no node. |
 | `02-deploy-device-plugin.sh` | Deploys the NVIDIA device plugin DaemonSet from Microsoft's guide, selecting the nodes AKS labels `accelerator=nvidia` and tolerating the GPU pool's taints. | The DaemonSet exists and schedules no pod until a GPU node joins. |
 | `03-create-resources.sh` | Creates (or reuses) the Service Bus namespace, the `training-requests` queue and the managed identity, grants the identity the **Azure Service Bus Data Owner** role on the namespace, federates it with the KEDA operator and with the training job's service account, and binds the operator to it. | The resources and the role assignment exist. The operator restart is skipped when the binding is already in place. |
 | `04-deploy-scaled-job.sh` | Creates the `gpu-training` namespace, the job's service account bound to the identity, the training script as a config map, the `TriggerAuthentication` and the `ScaledJob`, deriving the Service Bus endpoints at apply time. | The `ScaledJob` is `Ready`. |
-| `05-run-training-job.sh` | Sends a training request, waits for KEDA to start a job and for the cluster autoscaler to add a GPU node for its pod, prints the node's AKS labels, taints and GPU, and waits for the job to complete. | The job completes; its log says whether PyTorch found CUDA and how many steps it trained. |
+| `05-run-training-job.sh` | Sends a training request, waits for KEDA to start a job and for the cluster autoscaler to add a GPU node for its pod, or to start again the node an earlier run left deallocated, prints the node's AKS labels, taints and GPU, and waits for the job to complete. | The job completes; its log says whether PyTorch found CUDA and how many steps it trained. |
 | `06-simulate-eviction.sh` | Sends a longer training request and, once its pod reports training progress, simulates the eviction of its Spot node with `az vmss simulate-eviction`, follows the node until AKS deletes it, and waits for the job to complete on a new node. | The node is tainted for the eviction and deleted, and the job completes on another node. |
 | `07-cleanup.sh` | Deletes the two namespaces, the GPU node pool, the Service Bus namespace with the identity's role on it, and the identity's federation with the training job. | Everything the tutorial created is gone; the identity and the operator binding stay for the KEDA tutorials, which share them. |
 
@@ -81,10 +81,11 @@ cd gpu-node-pools/scripts
 ```
 
 > [!NOTE]
-> The first time you run the tutorial, the job's pod can take a few minutes to start running: the new GPU node has to download the pod's image, `pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime` (about 3 GB), which can take between 5 and 10 minutes depending on your network bandwidth.
+> The first run of a training job in a cluster waits several minutes before its pod starts: the new GPU node has to download the pod's image, `pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime` (about 3 GB), which takes between 5 and 10 minutes depending on your network bandwidth.
 > Until the download completes, the pod stays `Pending`, and `05-run-training-job.sh` keeps printing `(pods: Pending)`.
-> Every node the cluster autoscaler adds to the GPU pool starts without the image, so the wait comes back whenever a run finds the pool scaled back to zero, and in `06-simulate-eviction.sh` for the node that replaces the evicted one.
-> To check that the pod is downloading its image rather than stuck, run `kubectl get events --namespace gpu-training --watch` in another terminal: it shows `Pulling image` when the download starts and `Successfully pulled image` when it ends.
+> The next runs don't wait for it: once a job is done, the cluster autoscaler deallocates the idle node instead of deleting it, and starts it again, with the image on its disk, when the next job's pod asks for a GPU (see **A node that keeps its images** below).
+> The exception is the eviction in `06-simulate-eviction.sh`, which deletes the node: the new node that replaces it downloads the image again.
+> To check that the pod is downloading its image rather than stuck, run `kubectl get events --namespace gpu-training --watch` in another terminal: it shows `Pulling image` when the download starts and `Successfully pulled image` when it ends, or `Container image ... already present on machine` on a node that has the image.
 
 ### Kubernetes manifests
 
@@ -107,15 +108,17 @@ The job runs the official PyTorch image with CUDA (`pytorch/pytorch:2.5.1-cuda12
 
 ## How it works
 
-**A pool that scales from zero.** The GPU pool is created with `--node-count 0 --min-count 0 --max-count 1 --enable-cluster-autoscaler`. While no pod needs it, it has no node and costs nothing. The job's pod requests `nvidia.com/gpu: 1`, selects `workload: gpu` and tolerates the pool's `sku=gpu:NoSchedule` taint and the Spot taint, so no existing node can run it. The cluster autoscaler checks which pools could, finds the GPU pool, and adds a node to it. Once the job is done, it removes the node again after the node has been unneeded for a while: 10 minutes by default on AKS (the `scale-down-unneeded-time` of the [autoscaler profile](https://learn.microsoft.com/en-us/azure/aks/cluster-autoscaler#use-the-cluster-autoscaler-profile)), 30 seconds on the emulator.
+**A pool that scales from zero.** The GPU pool is created with `--node-count 0 --min-count 0 --max-count 1 --enable-cluster-autoscaler`. While no pod needs it, it has no node and costs nothing. The job's pod requests `nvidia.com/gpu: 1`, selects `workload: gpu` and tolerates the pool's `sku=gpu:NoSchedule` taint and the Spot taint, so no existing node can run it. The cluster autoscaler checks which pools could, finds the GPU pool, and adds a node to it. Once the job is done, it scales the pool back to zero after the node has been unneeded for a while: 10 minutes by default on AKS (the `scale-down-unneeded-time` of the [autoscaler profile](https://learn.microsoft.com/en-us/azure/aks/cluster-autoscaler#use-the-cluster-autoscaler-profile)), 30 seconds on the emulator.
+
+**A node that keeps its images.** The pool is also created with `--scale-down-mode Deallocate`, so the cluster autoscaler deallocates the idle node rather than deleting it, as [Scale-down Mode](https://learn.microsoft.com/en-us/azure/aks/scale-down-mode) describes. The deallocated node stays registered, `NotReady`, and its scale set instance stays stopped: it costs no compute, only its OS disk, which keeps the container images its pods pulled. When the next job's pod asks for a GPU, the autoscaler starts that node again before it adds any new node, and the pod starts without downloading its image. Scale-down Mode doesn't support ephemeral OS disks, hence `--node-osdisk-type Managed`. A Spot node in the stopped-deallocated state also still counts against your Spot vCPU quota, as the [Spot node pool guide](https://learn.microsoft.com/en-us/azure/aks/spot-node-pool) notes.
 
 **The labels and taints AKS sets.** AKS labels every node of an NVIDIA GPU size `accelerator=nvidia`, and every node of a Spot pool `kubernetes.azure.com/scalesetpriority=spot`, with the matching `NoSchedule` taint, whether or not the pool declares them. The device plugin DaemonSet selects the first, so it runs only on GPU nodes; the job tolerates the second, since only workloads that can handle an eviction belong on Spot nodes.
 
 **Why the device plugin.** AKS installs the NVIDIA driver on the nodes of a GPU size, but Kubernetes learns about their GPUs only from a device plugin. Until the plugin runs on a node, the node advertises no `nvidia.com/gpu`, and a pod that asks for one can't be scheduled there. This tutorial deploys the plugin's DaemonSet as Microsoft's guide does. AKS can also install it for you on [AKS-managed GPU node pools](https://learn.microsoft.com/en-us/azure/aks/aks-managed-gpu-nodes).
 
 **One job per message.** The `ScaledJob` declares an `azure-servicebus` trigger with `messageCount: "1"`: every polling interval, KEDA starts as many jobs as there are messages waiting, up to `maxReplicaCount`. Each job's pod receives one message and deletes it, so a request is processed once.
 
-**The eviction.** `az vmss simulate-eviction` evicts a Spot instance the way Azure does when it needs the capacity back. AKS taints the node `remediator.kubernetes.azure.com/unschedulable:NoSchedule`, evicts its pods and deletes the node with its scale set instance. The job's pod is gone, so the Job starts a replacement, which stays pending until the cluster autoscaler adds a new node, and completes there. The first pod had already received and deleted the message, so the replacement pod finds the queue empty and trains for the default time (`JOB_DEFAULT_TRAIN_SECONDS`). A production job would rather keep its message locked until it finishes, or checkpoint its training, so that a replacement pod picks up where the evicted one stopped.
+**The eviction.** `az vmss simulate-eviction` evicts a Spot instance the way Azure does when it needs the capacity back. AKS taints the node `remediator.kubernetes.azure.com/unschedulable:NoSchedule`, evicts its pods and deletes the node with its scale set instance: the pool's eviction policy, `Delete`, decides what happens to an evicted node, not its scale-down mode. The job's pod is gone, so the Job starts a replacement, which stays pending until the cluster autoscaler adds a new node, which downloads the image again, and completes there. The first pod had already received and deleted the message, so the replacement pod finds the queue empty and trains for the default time (`JOB_DEFAULT_TRAIN_SECONDS`). A production job would rather keep its message locked until it finishes, or checkpoint its training, so that a replacement pod picks up where the evicted one stopped.
 
 **How KEDA and the job authenticate.** The managed identity is federated with two service accounts: `keda-operator` in `kube-system`, so the operator can read the queue's message count, and the job's `gpu-training`, so the pod can receive a message. Both exchange their projected service account token for a Microsoft Entra token for the identity, and the **Azure Service Bus Data Owner** role on the namespace, granted on the identity's principal (object) id, authorizes them. No secret is involved, except the namespace key the scripts use on your machine to send the training requests over the HTTPS API.
 
@@ -126,6 +129,7 @@ The job runs the official PyTorch image with CUDA (`pytorch/pytorch:2.5.1-cuda12
 - [Use GPUs for compute-intensive workloads on AKS](https://learn.microsoft.com/en-us/azure/aks/use-nvidia-gpu)
 - [Add an Azure Spot node pool to an AKS cluster](https://learn.microsoft.com/en-us/azure/aks/spot-node-pool)
 - [Cluster autoscaling in AKS](https://learn.microsoft.com/en-us/azure/aks/cluster-autoscaler-overview)
+- [Use Scale-down Mode to delete or deallocate nodes in AKS](https://learn.microsoft.com/en-us/azure/aks/scale-down-mode)
 - [Simulate an eviction of an Azure Spot VM](https://learn.microsoft.com/en-us/azure/virtual-machine-scale-sets/use-spot#simulate-an-eviction)
 - [KEDA ScaledJobs](https://keda.sh/docs/2.20/concepts/scaling-jobs/) and the [Azure Service Bus scaler](https://keda.sh/docs/2.20/scalers/azure-service-bus/)
 - [Securely scale your applications using the KEDA add-on and workload identity](https://learn.microsoft.com/en-us/azure/aks/keda-workload-identity)

tutorials/gpu-node-pools/scripts/00-variables.sh

@@ -17,14 +17,21 @@ LOCATION="ItalyNorth"
 # nvidia.com/gpu, with 18 vCPUs. Spot VMs count against the region's Spot vCPU quota (lowPriorityCores) whatever their
 # VM family, and Azure starts one only when it has Spot capacity. Check the quota with:
 #   az vm list-usage --location italynorth --query "[?name.value=='lowPriorityCores']" --output table
-# The pool starts with no node: the cluster autoscaler adds one when a pod asks for a GPU and removes it when it is idle.
+# The pool starts with no node: the cluster autoscaler adds one when a pod asks for a GPU, and deallocates it when it
+# is idle instead of deleting it (scale-down mode Deallocate). A deallocated node keeps its OS disk, and with it the
+# images its pods pulled, and the autoscaler starts it again, before it adds any new node, when a pod asks for a GPU
+# again. Ephemeral OS disks don't support it, hence a managed OS disk. A Spot eviction still deletes the node: the
+# eviction policy decides what happens to an evicted node, the scale-down mode what happens to an idle one.
+# https://learn.microsoft.com/en-us/azure/aks/scale-down-mode
 GPU_NODE_POOL_NAME="gpupool"
 GPU_NODE_VM_SIZE="Standard_NV18ads_A10_v5"
 GPU_NODE_OS_SKU="AzureLinux"
+GPU_NODE_OS_DISK_TYPE="Managed"
 GPU_NODE_MIN_COUNT=0
 GPU_NODE_MAX_COUNT=1
 GPU_NODE_PRIORITY="Spot"
 GPU_NODE_EVICTION_POLICY="Delete"
+GPU_NODE_SCALE_DOWN_MODE="Deallocate"
 # -1 pays up to the on-demand price, so Azure evicts the node when it needs the capacity back, never for its price
 GPU_NODE_SPOT_MAX_PRICE=-1
 # The jobs select the pool's nodes by this label, and only pods that tolerate this taint run on them

tutorials/gpu-node-pools/scripts/01-add-gpu-node-pool.sh

@@ -1,13 +1,16 @@
 #!/bin/bash
 
 # Step 1: make sure the cluster runs the KEDA add-on, merge its credentials into kubeconfig, and add a Spot GPU node
-# pool that starts with no node and that the cluster autoscaler grows to one node when a pod asks for a GPU.
+# pool that starts with no node, that the cluster autoscaler grows to one node when a pod asks for a GPU, and whose
+# idle node the autoscaler deallocates, keeping the images on its disk for the next run, instead of deleting it.
 #
 # The cluster-creation scripts already pass --enable-keda, --enable-oidc-issuer and --enable-workload-identity, so on
-# such a cluster this script only adds the node pool. The pool's VM size and priority can't change after it is
-# created, so an existing pool with other values stops the script rather than being reused.
+# such a cluster this script only adds the node pool. The pool's VM size, priority and OS disk type can't change after
+# it is created, so an existing pool with other values stops the script rather than being reused; its scale-down mode
+# can, so an existing pool gets the one this tutorial uses.
 # https://learn.microsoft.com/en-us/azure/aks/use-nvidia-gpu
 # https://learn.microsoft.com/en-us/azure/aks/spot-node-pool
+# https://learn.microsoft.com/en-us/azure/aks/scale-down-mode
 
 # Variables
 source ./00-variables.sh
@@ -70,16 +73,36 @@ EXISTING_POOL=$(az aks nodepool show \
   --name "$GPU_NODE_POOL_NAME" \
   --cluster-name "$AKS_NAME" \
   --resource-group "$AKS_RESOURCE_GROUP_NAME" \
-  --query "[vmSize, scaleSetPriority]" \
+  --query "[vmSize, scaleSetPriority, osDiskType]" \
   --output tsv \
   --only-show-errors 2>/dev/null | tr -d '\r' | tr '\t\n' '  ')
 if [[ -n $EXISTING_POOL ]]; then
-  if [[ $EXISTING_POOL != "$GPU_NODE_VM_SIZE $GPU_NODE_PRIORITY " ]]; then
-    echo "The [$GPU_NODE_POOL_NAME] node pool exists with VM size and priority [$EXISTING_POOL], not [$GPU_NODE_VM_SIZE $GPU_NODE_PRIORITY]"
+  if [[ $EXISTING_POOL != "$GPU_NODE_VM_SIZE $GPU_NODE_PRIORITY $GPU_NODE_OS_DISK_TYPE " ]]; then
+    echo "The [$GPU_NODE_POOL_NAME] node pool exists with VM size, priority and OS disk type [$EXISTING_POOL], not [$GPU_NODE_VM_SIZE $GPU_NODE_PRIORITY $GPU_NODE_OS_DISK_TYPE]"
     echo "Delete it with 07-cleanup.sh, or change GPU_NODE_POOL_NAME in 00-variables.sh"
     exit 1
   fi
   echo "The [$GPU_NODE_POOL_NAME] node pool already exists"
+  # A pool added before this tutorial set the scale-down mode deletes its idle node, and with it the job's image
+  SCALE_DOWN_MODE=$(az aks nodepool show \
+    --name "$GPU_NODE_POOL_NAME" \
+    --cluster-name "$AKS_NAME" \
+    --resource-group "$AKS_RESOURCE_GROUP_NAME" \
+    --query scaleDownMode \
+    --output tsv \
+    --only-show-errors 2>/dev/null | tr -d '\r')
+  if [[ $SCALE_DOWN_MODE != "$GPU_NODE_SCALE_DOWN_MODE" ]]; then
+    echo "Setting the scale-down mode of the [$GPU_NODE_POOL_NAME] node pool to [$GPU_NODE_SCALE_DOWN_MODE]..."
+    if ! az aks nodepool update \
+      --name "$GPU_NODE_POOL_NAME" \
+      --cluster-name "$AKS_NAME" \
+      --resource-group "$AKS_RESOURCE_GROUP_NAME" \
+      --scale-down-mode "$GPU_NODE_SCALE_DOWN_MODE" \
+      --only-show-errors 1>/dev/null; then
+      echo "Failed to set the scale-down mode of the [$GPU_NODE_POOL_NAME] node pool"
+      exit 1
+    fi
+  fi
 else
   echo "Adding the [$GPU_NODE_POOL_NAME] Spot node pool of [$GPU_NODE_VM_SIZE] VMs, with $GPU_NODE_MIN_COUNT to $GPU_NODE_MAX_COUNT nodes, to the [$AKS_NAME] AKS cluster..."
   if ! az aks nodepool add \
@@ -89,9 +112,11 @@ else
     --mode User \
     --node-vm-size "$GPU_NODE_VM_SIZE" \
     --os-sku "$GPU_NODE_OS_SKU" \
+    --node-osdisk-type "$GPU_NODE_OS_DISK_TYPE" \
     --priority "$GPU_NODE_PRIORITY" \
     --eviction-policy "$GPU_NODE_EVICTION_POLICY" \
     --spot-max-price "$GPU_NODE_SPOT_MAX_PRICE" \
+    --scale-down-mode "$GPU_NODE_SCALE_DOWN_MODE" \
     --enable-cluster-autoscaler \
     --node-count "$GPU_NODE_MIN_COUNT" \
     --min-count "$GPU_NODE_MIN_COUNT" \
@@ -110,6 +135,6 @@ az aks nodepool show \
   --name "$GPU_NODE_POOL_NAME" \
   --cluster-name "$AKS_NAME" \
   --resource-group "$AKS_RESOURCE_GROUP_NAME" \
-  --query "{name: name, vmSize: vmSize, count: count, minCount: minCount, maxCount: maxCount, scaleSetPriority: scaleSetPriority, scaleSetEvictionPolicy: scaleSetEvictionPolicy, spotMaxPrice: spotMaxPrice, gpuProfile: gpuProfile, nodeLabels: nodeLabels, nodeTaints: nodeTaints}" \
+  --query "{name: name, vmSize: vmSize, osDiskType: osDiskType, count: count, minCount: minCount, maxCount: maxCount, scaleDownMode: scaleDownMode, scaleSetPriority: scaleSetPriority, scaleSetEvictionPolicy: scaleSetEvictionPolicy, spotMaxPrice: spotMaxPrice, gpuProfile: gpuProfile, nodeLabels: nodeLabels, nodeTaints: nodeTaints}" \
   --output json \
   --only-show-errors