Add end-to-end GPU labs (Kueue + cluster autoscaler) to the Ray on AKS example

examples/kueue-and-ray-on-aks/2-kueue-queues/README.md

@@ -7,13 +7,22 @@ resource management for the Ray workloads. This module sets up:
 - ResourceFlavors that map to CPU and GPU node pools
 - Queue configurations that control how workloads are admitted to the cluster
 
-Two queue configurations are provided as **independent demos** — apply one or
-the other, not both:
+Three queue configurations are provided as **independent demos** — apply one,
+not several:
 
 | Configuration | File | What it demonstrates |
 |---------------|------|----------------------|
 | **Single queue** | `manifests/20-single-queue.yaml` | One ClusterQueue with admission backpressure — one workload runs, the next waits |
 | **Team queues** | `manifests/30-team-queues.yaml` | Two ClusterQueues in a shared cohort with borrowing and preemption |
+| **Autoscale queue** | `manifests/40-autoscale-queue.yaml` | A ProvisioningRequest AdmissionCheck that drives the AKS cluster autoscaler to provision capacity *before* admission |
+| **GPU autoscale queue** | `manifests/50-gpu-autoscale-queue.yaml` | The same provisioning gate applied to GPUs — atomic scale-up of a GPU pool, with GPU-aware quota |
+
+> The single and team queues admit against a **fixed** quota on the
+> already-provisioned GPU node. The autoscale queue is different: it pairs Kueue
+> with the cluster autoscaler so capacity is provisioned **on demand**. See
+> [Autoscale queue — provision on demand](#autoscale-queue--provision-on-demand)
+> below, and the [cas-batch-job](../3-workloads/cas-batch-job/) workload that
+> uses it.
 
 ## Prerequisites
 
@@ -39,6 +48,8 @@ the other, not both:
 | `manifests/10-resource-flavors.yaml` | `default` (any node) and `gpu` (NVIDIA accelerator nodes) ResourceFlavors |
 | `manifests/20-single-queue.yaml` | `cluster-queue` ClusterQueue + `default` LocalQueue |
 | `manifests/30-team-queues.yaml` | `team-a-cq` / `team-b-cq` ClusterQueues in `shared-cohort` + `team-a` / `team-b` LocalQueues |
+| `manifests/40-autoscale-queue.yaml` | `scalepool` ResourceFlavor + `cas-provisioning` AdmissionCheck + `cas-provreq-config` ProvisioningRequestConfig + `cas-cluster-queue` ClusterQueue + `cas-local-queue` LocalQueue (in its own `cas-kueue-demo` namespace) |
+| `manifests/50-gpu-autoscale-queue.yaml` | `gpu-a100` ResourceFlavor + `gpu-provisioning` AdmissionCheck + `gpu-provreq-config` ProvisioningRequestConfig + `gpu-cluster-queue` ClusterQueue + `gpu-local-queue` LocalQueue (in its own `gpu-lab` namespace) |
 
 ## Apply
 
@@ -61,6 +72,14 @@ kubectl apply -f manifests/20-single-queue.yaml
 
 #   Option B — Team queues (multi-tenant borrowing + preemption demo)
 kubectl apply -f manifests/30-team-queues.yaml
+
+#   Option C — Autoscale queue (provision capacity on demand via CAS)
+#   Requires an autoscaling `scalepool` pool — see the section below.
+kubectl apply -f manifests/40-autoscale-queue.yaml
+
+#   Option D — GPU autoscale queue (provision GPU capacity on demand via CAS)
+#   Requires an autoscaling GPU pool — see ../3-workloads/gpu-lab/.
+kubectl apply -f manifests/50-gpu-autoscale-queue.yaml
 ```
 
 > **⚠️ Choose one.** `20-single-queue.yaml` and `30-team-queues.yaml` are
@@ -237,6 +256,61 @@ Preemption policy:
 Then submit a 4-GPU RayJob to `team-b` — Kueue will preempt Team A down to 4
 GPUs and admit Team B's job. See Module 3 for ready-to-run workload examples.
 
+### Autoscale queue — provision on demand
+
+The single- and team-queue configs admit workloads against a **fixed** quota
+that assumes the GPU node already exists. The autoscale queue instead pairs
+Kueue with the AKS **cluster autoscaler (CAS)** so nodes are provisioned *before*
+a workload is admitted:
+
+```
+Workload (suspend: true, queue-name: cas-local-queue)
+   │
+   ▼
+Kueue ──► AdmissionCheck cas-provisioning ──► ProvisioningRequest ──► CAS
+   │                                                                    │
+   └──────── admitted once nodes exist ◄──── Provisioned=True ◄─────────┘
+```
+
+Three objects wire the gate together (all in `40-autoscale-queue.yaml`):
+
+- **ProvisioningRequestConfig** `cas-provreq-config` selects the
+  `best-effort-atomic-scale-up.autoscaling.x-k8s.io` provisioning class — CAS
+  adds the requested capacity as a single atomic increase (all-or-nothing),
+  which is what batch/gang workloads want.
+- **AdmissionCheck** `cas-provisioning` uses the
+  `kueue.x-k8s.io/provisioning-request` controller and points at that config.
+- **ClusterQueue** `cas-cluster-queue` gates admission on `cas-provisioning` via
+  `admissionChecksStrategy`, so every workload it admits first goes through the
+  provisioning gate. It admits onto the `scalepool` ResourceFlavor
+  (`agentpool: scalepool`).
+
+**Requirements:**
+
+- An autoscaling CPU pool named `scalepool`
+  (`az aks nodepool add ... --enable-cluster-autoscaler --min-count 1 --max-count 5`).
+
+The [cas-batch-job](../3-workloads/cas-batch-job/) workload in Module 3 submits
+a suspended Job through this queue and walks through the scale-up end to end.
+
+### GPU variant
+
+`manifests/50-gpu-autoscale-queue.yaml` applies the same gate to GPU capacity.
+Two things change: `managedResources` is `nvidia.com/gpu` rather than `cpu`, so
+only GPU-requesting podsets drive provisioning, and the ResourceFlavor carries
+a toleration for the `nvidia.com/gpu=present:NoSchedule` taint that AKS puts on
+GPU pools.
+
+The stakes are also higher. Atomic scale-up matters more when a half-placed
+gang holds A100s instead of vCPUs — see the
+[GPU labs](../3-workloads/gpu-lab/) for the full walkthrough, including
+multi-node training and queue contention.
+
+**Requirements:**
+
+- An autoscaling GPU pool named `gpupool`
+  (`az aks nodepool add ... --enable-cluster-autoscaler --min-count 0 --max-count 2`).
+
 ### Quota sizing
 
 The default quotas are sized for a single `Standard_ND96amsr_A100_v4` node

examples/kueue-and-ray-on-aks/2-kueue-queues/manifests/40-autoscale-queue.yaml

@@ -0,0 +1,111 @@
+# Autoscale queue — ProvisioningRequest AdmissionCheck driving the AKS cluster autoscaler.
+#
+# The single-queue and team-queue configs (20-/30-) admit workloads against a
+# FIXED quota that assumes the GPU node already exists. This configuration is
+# different: it pairs Kueue with the AKS cluster autoscaler (CAS) so capacity is
+# provisioned ON DEMAND, before the workload is admitted.
+#
+# How it fits together:
+#
+#   1. ProvisioningRequestConfig + AdmissionCheck delegate capacity provisioning
+#      to CAS. When Kueue needs capacity for a workload, the AdmissionCheck asks
+#      CAS — via a ProvisioningRequest — to atomically scale up the pool BEFORE
+#      the workload is admitted. The workload only starts once the nodes exist.
+#
+#   2. A ClusterQueue references that AdmissionCheck, so every workload routed
+#      through the queue goes through the provisioning gate first.
+#
+# provisioningClassName `best-effort-atomic-scale-up.autoscaling.x-k8s.io` tells
+# CAS to add all the requested capacity as a single atomic increase (all-or-
+# nothing), which is what batch/gang workloads want.
+#
+# This targets an AUTOSCALING CPU pool named `scalepool` (see the README for how
+# to add it with `az aks nodepool add --enable-cluster-autoscaler`). It is an
+# independent configuration — apply it INSTEAD OF 20-/30-, not alongside them.
+#
+# Unlike the Ray configs in this directory, this one is self-contained: it
+# creates its own `cas-kueue-demo` namespace and does not use the `ray`
+# namespace, because the workload it gates is a plain batch Job, not a RayJob.
+#
+# Apply:
+#   kubectl apply -f 40-autoscale-queue.yaml
+#
+# Verify:
+#   kubectl get admissioncheck cas-provisioning
+#   kubectl get clusterqueue cas-cluster-queue
+#   kubectl -n cas-kueue-demo get localqueue cas-local-queue
+---
+# Namespace for the cluster autoscaler + Kueue demo.
+apiVersion: v1
+kind: Namespace
+metadata:
+  name: cas-kueue-demo
+---
+# ResourceFlavor pinned to the autoscaling CPU pool. AKS stamps every node in a
+# pool with `agentpool: <pool-name>`, so this flavor only admits onto scalepool.
+apiVersion: kueue.x-k8s.io/v1beta2
+kind: ResourceFlavor
+metadata:
+  name: scalepool
+spec:
+  nodeLabels:
+    agentpool: scalepool
+---
+# AdmissionCheck that delegates capacity provisioning to CAS via ProvisioningRequest.
+apiVersion: kueue.x-k8s.io/v1beta2
+kind: AdmissionCheck
+metadata:
+  name: cas-provisioning
+spec:
+  controllerName: kueue.x-k8s.io/provisioning-request
+  parameters:
+    apiGroup: kueue.x-k8s.io
+    kind: ProvisioningRequestConfig
+    name: cas-provreq-config
+---
+apiVersion: kueue.x-k8s.io/v1beta2
+kind: ProvisioningRequestConfig
+metadata:
+  name: cas-provreq-config
+spec:
+  provisioningClassName: best-effort-atomic-scale-up.autoscaling.x-k8s.io
+  managedResources:
+    - cpu
+---
+apiVersion: kueue.x-k8s.io/v1beta2
+kind: ClusterQueue
+metadata:
+  name: cas-cluster-queue
+spec:
+  namespaceSelector: {}
+  queueingStrategy: BestEffortFIFO
+  resourceGroups:
+    - coveredResources:
+        - cpu
+        - memory
+      flavors:
+        - name: scalepool
+          resources:
+            - name: cpu
+              # Upper bound Kueue will admit before asking CAS to grow the pool.
+              # Set generously here so quota is never the limiting factor — the
+              # pool's --max-count is what actually caps the scale-up. Lower it
+              # to max-count × per-node cores (e.g. 5 × 4 = 20) if you want
+              # Kueue to enforce the ceiling too.
+              nominalQuota: "100"
+            - name: memory
+              nominalQuota: 200Gi
+  # Gate admission on the provisioning check, scoped to the scalepool flavor.
+  # This is the form the upstream Kueue provisioning guide uses.
+  admissionChecksStrategy:
+    admissionChecks:
+      - name: cas-provisioning
+        onFlavors: [scalepool]
+---
+apiVersion: kueue.x-k8s.io/v1beta2
+kind: LocalQueue
+metadata:
+  name: cas-local-queue
+  namespace: cas-kueue-demo
+spec:
+  clusterQueue: cas-cluster-queue

examples/kueue-and-ray-on-aks/2-kueue-queues/manifests/50-gpu-autoscale-queue.yaml

@@ -0,0 +1,120 @@
+# GPU autoscale queue — ProvisioningRequest AdmissionCheck driving the AKS
+# cluster autoscaler for GPU capacity.
+#
+# This is the GPU counterpart to 40-autoscale-queue.yaml. The mechanism is the
+# same (Kueue asks CAS for capacity via a ProvisioningRequest before admitting
+# the workload); what changes is what is scarce.
+#
+# With CPU, over-provisioning is cheap and a half-scheduled job is merely
+# wasteful. With GPUs it is neither. A100/H100 nodes are expensive and often
+# capacity-constrained, so two properties matter:
+#
+#   1. ATOMIC scale-up. A distributed training job that gets 6 of the 8 GPUs it
+#      asked for makes no progress, but still holds those 6 GPUs. The
+#      `best-effort-atomic-scale-up` class makes CAS add the whole increment or
+#      none of it, so a gang is never half-placed.
+#
+#   2. GPU-AWARE quota. `managedResources: [nvidia.com/gpu]` tells Kueue that
+#      the GPU is the resource CAS is being asked to provision. Workloads whose
+#      podsets request no GPU skip the provisioning gate entirely rather than
+#      waiting behind it.
+#
+# Apply:
+#   kubectl apply -f 50-gpu-autoscale-queue.yaml
+#
+# Verify:
+#   kubectl get admissioncheck gpu-provisioning
+#   kubectl get clusterqueue gpu-cluster-queue
+#   kubectl -n gpu-lab get localqueue gpu-local-queue
+---
+apiVersion: v1
+kind: Namespace
+metadata:
+  name: gpu-lab
+---
+# ResourceFlavor pinned to the autoscaling GPU pool.
+#
+# AKS stamps every node with `agentpool: <pool-name>`, which is the most precise
+# selector — it distinguishes an A100 pool from an H100 pool in the same cluster.
+#
+# The tolerations below are applied by Kueue to admitted pods. AKS taints GPU
+# pools created with `--node-taints nvidia.com/gpu=present:NoSchedule` so that
+# CPU work cannot land on expensive GPU nodes; the flavor has to tolerate that
+# taint or nothing will ever schedule.
+apiVersion: kueue.x-k8s.io/v1beta2
+kind: ResourceFlavor
+metadata:
+  name: gpu-a100
+spec:
+  nodeLabels:
+    agentpool: gpupool
+  tolerations:
+    - key: nvidia.com/gpu
+      operator: Equal
+      value: present
+      effect: NoSchedule
+---
+apiVersion: kueue.x-k8s.io/v1beta2
+kind: ProvisioningRequestConfig
+metadata:
+  name: gpu-provreq-config
+spec:
+  # All-or-nothing scale-up: CAS either provisions the entire gang's worth of
+  # GPU nodes or reports failure. This is what makes multi-node training safe.
+  provisioningClassName: best-effort-atomic-scale-up.autoscaling.x-k8s.io
+  # Only GPU requests drive provisioning. A podset that asks for CPU/memory but
+  # no GPU is considered ready immediately instead of blocking on CAS.
+  managedResources:
+    - nvidia.com/gpu
+---
+apiVersion: kueue.x-k8s.io/v1beta2
+kind: AdmissionCheck
+metadata:
+  name: gpu-provisioning
+spec:
+  controllerName: kueue.x-k8s.io/provisioning-request
+  parameters:
+    apiGroup: kueue.x-k8s.io
+    kind: ProvisioningRequestConfig
+    name: gpu-provreq-config
+---
+apiVersion: kueue.x-k8s.io/v1beta2
+kind: ClusterQueue
+metadata:
+  name: gpu-cluster-queue
+spec:
+  namespaceSelector: {}
+  queueingStrategy: BestEffortFIFO
+  resourceGroups:
+    - coveredResources:
+        - nvidia.com/gpu
+        - cpu
+        - memory
+      flavors:
+        - name: gpu-a100
+          resources:
+            # The GPU quota is the meaningful ceiling. Size it to
+            #   max-count x GPUs-per-node
+            # e.g. a 2-node pool of Standard_ND96amsr_A100_v4 (8 GPUs each)
+            # = 16. Kueue will not admit past this even if CAS could scale
+            # further, which is your guardrail against a runaway spend.
+            - name: nvidia.com/gpu
+              nominalQuota: "16"
+            # CPU and memory are sized generously: they are along for the ride
+            # and should not be what blocks admission.
+            - name: cpu
+              nominalQuota: "200"
+            - name: memory
+              nominalQuota: 3000Gi
+  admissionChecksStrategy:
+    admissionChecks:
+      - name: gpu-provisioning
+        onFlavors: [gpu-a100]
+---
+apiVersion: kueue.x-k8s.io/v1beta2
+kind: LocalQueue
+metadata:
+  name: gpu-local-queue
+  namespace: gpu-lab
+spec:
+  clusterQueue: gpu-cluster-queue