Add Kueue and Cluster Autoscaler on AKS example

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

@@ -7,13 +7,21 @@ 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 |
+
+> 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 +47,7 @@ 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 + `autoscale-queue` ClusterQueue + `autoscale` LocalQueue |
 
 ## Apply
 
@@ -61,6 +70,10 @@ 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
 ```
 
 > **⚠️ Choose one.** `20-single-queue.yaml` and `30-team-queues.yaml` are
@@ -237,6 +250,46 @@ 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: autoscale)
+   │
+   ▼
+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** `autoscale-queue` lists `cas-provisioning` under
+  `admissionChecks`, 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 cluster autoscaler running with `--enable-provisioning-requests` — this is
+  what consumes the ProvisioningRequest. Without it the request is created but
+  never satisfied.
+
+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.
+
 ### 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,96 @@
+# 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.
+#
+# Prerequisite: the cluster autoscaler must run with `--enable-provisioning-requests`.
+#
+# Apply after the namespace and (this file's own) resource flavor exist:
+#   kubectl apply -f 40-autoscale-queue.yaml
+#
+# Verify:
+#   kubectl get admissioncheck cas-provisioning
+#   kubectl get clusterqueue autoscale-queue
+#   kubectl -n ray get localqueue autoscale
+---
+# 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/v1beta1
+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/v1beta1
+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: autoscale-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.
+              # Size to the pool's max-count × per-node cores (e.g. 5 × 4 = 20).
+              nominalQuota: "100"
+            - name: memory
+              nominalQuota: 200Gi
+  admissionChecks:
+    - cas-provisioning
+---
+apiVersion: kueue.x-k8s.io/v1beta2
+kind: LocalQueue
+metadata:
+  name: autoscale
+  namespace: ray
+spec:
+  clusterQueue: autoscale-queue

examples/kueue-and-ray-on-aks/3-workloads/README.md

@@ -31,6 +31,11 @@ The shared base templates in [`_template/`](_template/) define the standard
 fields used by every workload. Per-example templates extend these with
 workload-specific pip packages, env vars, and resource counts.
 
+> **Exception:** [`cas-batch-job/`](cas-batch-job/) is a plain Kubernetes Job
+> (no Ray, no `envsubst`) applied directly with `kubectl apply -f`. It
+> demonstrates the cluster-autoscaler ProvisioningRequest path rather than a Ray
+> runtime, so it doesn't use the template system described here.
+
 ### Standard fields (all workloads)
 
 | Field | Value | Reason |
@@ -93,17 +98,29 @@ Ray Serve. Loads the LoRA adapter from blob storage on startup.
 - Access: `kubectl -n ray port-forward svc/${SERVICE_NAME}-serve-svc 8000:8000`
 - See [online-serving/](online-serving/)
 
+### 5. cas-batch-job/ — provision capacity on demand (cluster autoscaler)
+A plain Kubernetes batch Job (no Ray) that demonstrates the **ProvisioningRequest
+→ cluster autoscaler** path. Where the examples above admit against a fixed quota
+on an existing GPU node, this one drives the AKS cluster autoscaler to grow an
+autoscaling CPU pool *before* the Job is admitted — capacity provisioned
+just-in-time.
+
+- **CPU only** — 3 pods × 1800m on an autoscaling `scalepool`
+- Queue: `autoscale` (the ProvisioningRequest-gated queue from Module 2)
+- Requires: an autoscaling `scalepool` pool + CAS `--enable-provisioning-requests`
+- See [cas-batch-job/](cas-batch-job/)
+
 ## Comparison
 
-| | aurora-finetune | llm-training | batch-inference | online-serving |
-|--|--|--|--|--|
-| Kind | RayJob | RayJob | RayJob | RayService |
-| Kueue-admitted | ✓ | ✓ | ✓ | ✗ |
-| GPUs | 1 | 4 | 1 | 1 |
-| Queue label | `default` | `default` | `default` | — |
-| Deps | Aurora, torch, azure-storage-blob | LLaMA-Factory, azure-storage-blob | vLLM, azure-storage-blob | Aurora, azure-storage-blob |
-| Reads from blob | `aurora/data/` (init/truth) | `llm-pipeline/data/` (train.jsonl) | `llm-pipeline/data/` + `lora/` | `aurora/checkpoints/` (adapter) |
-| Writes to blob | `aurora/checkpoints/` | `llm-pipeline/lora/` | `llm-pipeline/inference/` | — |
+| | aurora-finetune | llm-training | batch-inference | online-serving | cas-batch-job |
+|--|--|--|--|--|--|
+| Kind | RayJob | RayJob | RayJob | RayService | Job |
+| Kueue-admitted | ✓ | ✓ | ✓ | ✗ | ✓ |
+| GPUs | 1 | 4 | 1 | 1 | 0 (CPU) |
+| Queue label | `default` | `default` | `default` | — | `autoscale` |
+| Deps | Aurora, torch, azure-storage-blob | LLaMA-Factory, azure-storage-blob | vLLM, azure-storage-blob | Aurora, azure-storage-blob | none (pause image) |
+| Reads from blob | `aurora/data/` (init/truth) | `llm-pipeline/data/` (train.jsonl) | `llm-pipeline/data/` + `lora/` | `aurora/checkpoints/` (adapter) | — |
+| Writes to blob | `aurora/checkpoints/` | `llm-pipeline/lora/` | `llm-pipeline/inference/` | — | — |
 
 ## Quick start
 

examples/kueue-and-ray-on-aks/3-workloads/cas-batch-job/README.md

@@ -0,0 +1,100 @@
+# cas-batch-job — provision capacity on demand with Kueue + cluster autoscaler
+
+A plain Kubernetes batch Job that demonstrates the **ProvisioningRequest → AKS
+cluster autoscaler (CAS)** path. The other Module 3 workloads admit against a
+fixed quota on an already-provisioned GPU node; this one instead asks CAS to
+*grow* an autoscaling CPU pool before the Job is admitted, so batch capacity is
+provisioned just-in-time.
+
+This isolates the Kueue↔CAS integration from the Ray runtime — it's the
+simplest way to see a *provision-first* scale-up end to end.
+
+## How it works
+
+```
+Job (suspend: true, queue-name: autoscale)
+   │
+   ▼
+Kueue creates a Workload ──► AdmissionCheck cas-provisioning
+   │                              │
+   │                              ▼
+   │                        ProvisioningRequest  ──►  cluster autoscaler
+   │                              │                     atomically scales
+   │                              ▼                     up `scalepool`
+   └──── admitted once nodes exist ◄── Provisioned=True
+   │
+   ▼
+suspend flips to false → 3 pods schedule on the new nodes → Job completes
+```
+
+The atomic `best-effort-atomic-scale-up.autoscaling.x-k8s.io` provisioning class
+adds the whole capacity block at once, avoiding the half-scheduled gang problem
+where some pods land and others sit Pending.
+
+## Prerequisites
+
+- Module 1 cluster deployed and Kueue running.
+- An **autoscaling** CPU pool named `scalepool`. If your Module 1 cluster
+  doesn't have one, add it:
+  ```bash
+  az aks nodepool add \
+    --resource-group <rg> --cluster-name <cluster> \
+    --name scalepool --mode User \
+    --node-vm-size Standard_D4s_v3 \
+    --enable-cluster-autoscaler --min-count 1 --max-count 5
+  ```
+- The cluster autoscaler must run with `--enable-provisioning-requests` (this is
+  what consumes the ProvisioningRequest — without it the request is created but
+  never satisfied).
+- The autoscale queue applied from Module 2:
+  ```bash
+  kubectl apply -f ../../2-kueue-queues/manifests/40-autoscale-queue.yaml
+  ```
+
+## Submit
+
+```bash
+kubectl apply -f manifests/job.yaml
+```
+
+The Job requests 3 pods at 1800m CPU each. On a pool starting at one 4-core node
+they cannot all fit, so Kueue asks CAS to scale up before admitting them.
+
+## Watch the flow
+
+```bash
+# The Job starts Suspended
+kubectl -n ray get job kueue-cas-job
+
+# Kueue creates a Workload, then a ProvisioningRequest
+kubectl -n ray get workloads
+kubectl -n ray get provisioningrequest
+
+# CAS provisions the nodes (Provisioned=True), the pool grows
+kubectl -n ray get provisioningrequest -o yaml | grep -A5 conditions
+kubectl get nodes -l agentpool=scalepool
+
+# Job runs to completion
+kubectl -n ray get job kueue-cas-job -w   # COMPLETIONS 3/3
+```
+
+Expected end state:
+
+```output
+NAME            STATUS     COMPLETIONS   DURATION   AGE
+kueue-cas-job   Complete   3/3           45s        6m
+```
+
+## Troubleshooting
+
+| Symptom | Cause | Fix |
+|---------|-------|-----|
+| Job stays `Suspended`, no ProvisioningRequest | Wrong queue-name label | Verify `kueue.x-k8s.io/queue-name: autoscale` matches the LocalQueue |
+| ProvisioningRequest created but `status.conditions` empty | CAS not consuming ProvisioningRequests | The autoscaler needs `--enable-provisioning-requests` |
+| Pods Pending after admission | Node label mismatch | Confirm `scalepool` nodes carry `agentpool=scalepool` (`kubectl get nodes --show-labels`) |
+
+## Clean up
+
+```bash
+kubectl delete -f manifests/job.yaml
+```

examples/kueue-and-ray-on-aks/3-workloads/cas-batch-job/manifests/job.yaml

@@ -0,0 +1,46 @@
+# A suspended batch Job routed through Kueue's autoscale queue.
+#
+# Unlike the Ray examples in this module, this is a plain Kubernetes Job — it
+# exists to demonstrate the ProvisioningRequest -> cluster autoscaler (CAS) path
+# without the Ray runtime. The Job requests more CPU than the autoscaling pool's
+# current nodes can satisfy, so Kueue's ProvisioningRequest AdmissionCheck asks
+# CAS to add nodes BEFORE admitting it. The flow:
+#
+#   suspend: true          -> Job created but no pods yet
+#   queue-name label       -> Kueue picks it up, creates a Workload
+#   AdmissionCheck         -> Kueue creates a ProvisioningRequest for CAS
+#   CAS atomic scale-up    -> scalepool grows to fit the 3 pods
+#   suspend flipped false  -> pods schedule onto the new nodes, Job runs
+#
+# 3 pods x 1800m CPU on 4-core nodes forces a scale-up beyond the single
+# starting node (~2 schedulable pods/node).
+#
+# Submit:
+#   kubectl apply -f manifests/job.yaml
+#
+# Watch the ProvisioningRequest and scale-up:
+#   kubectl -n ray get provisioningrequest
+#   kubectl get nodes -l agentpool=scalepool -w
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: kueue-cas-job
+  namespace: ray
+  labels:
+    kueue.x-k8s.io/queue-name: autoscale
+spec:
+  parallelism: 3
+  completions: 3
+  suspend: true
+  template:
+    spec:
+      nodeSelector:
+        agentpool: scalepool
+      containers:
+        - name: worker
+          image: mcr.microsoft.com/oss/v2/kubernetes/pause:3.6
+          resources:
+            requests:
+              cpu: "1800m"
+              memory: "256Mi"
+      restartPolicy: Never