Gang Scheduling & Batch Reservation — ProvisioningRequest, Kueue, Dynamic Workload Scheduler

Tại sao gang scheduling quan trọng

Suppose distributed training job yêu cầu 8 GPUs:

Scenario 1: NO gang scheduling (Kubernetes default)
─────────────────────────────────────────────────

t=0: Pod 0 scheduled on GPU node → starts immediately
t=1: Pod 1 scheduled on GPU node → starts immediately
t=2-10: Pods 2-7 pending (GPU nodes saturated)

Training framework (e.g., PyTorch DDP) tries to initialize:
→ 2 processes connected, 6 missing
→ Timeout after 30min → entire job fails + all 2 running processes killed
→ Waste compute + start over

Scenario 2: WITH gang scheduling (ProvisioningRequest)
──────────────────────────────────────────────────────

t=0: Pod requests 8 GPUs total (gang of 8)
t=0-5: Kubernetes cluster autoscaler receives signal:
       "need 8 GPUs simultaneously (not incrementally)"
t=5: Cluster autoscaler provisions 1 new node with 8 GPUs (1 operation)
t=6: All 8 Pods scheduled together → training starts immediately
t=10: Training completes successfully

Result: No partial scheduling → no timeout → training succeeds

Mental model: Gang scheduling ensures atomic resource allocation — either all Pods get resources, or none do. Prevents partial setup → timeout failures.

---

ProvisioningRequest: Batch reservation API

Cơ chế

ProvisioningRequest là custom resource (CRD) nói với cluster autoscaler:

"I have a batch of N Pods with total requirement R.
 Don't schedule them incrementally.
 Only schedule when you can provision all R at once."

Setup: ProvisioningRequest manifest

apiVersion: autoscaling.gke.io/v1beta1
kind: ProvisioningRequest
metadata:
  name: training-batch-1
  namespace: ml-workloads
spec:
  # Define resource class (optional, for quota management)
  parameters:
    key: "gpu-training-class"
  
  # What this ProvisioningRequest needs
  podTemplateSpec:
    metadata:
      labels:
        batch: training-1
    spec:
      terminationGracePeriodSeconds: 300
      # Actual workload spec (e.g., training job)
  
  # How many replicas of this Pod template
  podSetTemplates:
  - name: gpu-pods
    count: 8
    podTemplateSpec:
      spec:
        containers:
        - name: training
          image: training:latest
          resources:
            limits:
              nvidia.com/gpu: 1
        tolerations:
        - key: nvidia.com/gpu
          value: "true"
          effect: NoSchedule

Kubernetes cluster autoscaler:

1. Receives ProvisioningRequest
2. Calculates total need: 8 Pods × 1 GPU = 8 GPUs
3. Checks available capacity:
   - Current nodes have 0 free GPUs
   - Must provision 1 new node (8 GPUs)
4. If provision succeeds:
   → All 8 Pods scheduled together
5. If timeout (provision fails after 1 hour):
   → ProvisioningRequest marked FAILED
   → Pods remain PENDING

Dynamic Workload Scheduler vs Cluster Autoscaler

GKE provides two gang-scheduling mechanisms:

1. Cluster Autoscaler + ProvisioningRequest (resource-aware)

ProvisioningRequest → Cluster autoscaler → provision new nodes

Best for: Large batch jobs (training, data processing) where you need new GPU nodes.

2. Dynamic Workload Scheduler (delay-tolerant)

Pod → DWS → queue waiting for resources
     → automatically schedule when resources free up

Best for: Batch jobs that can wait (deadline not immediate), prefer not to provision new nodes.

Using Kueue for job management

Kueue is higher-level job queueing system (sits above ProvisioningRequest):

# 1. Define queue with resource quota
apiVersion: kueue.x-k8s.io/v1beta1
kind: Queue
metadata:
  name: gpu-training-queue
spec:
  namespaceSelector:
    matchLabels:
      queue: training
  resources:
  - name: gpu
    flavors:
    - name: nvidia-h100
      resources:
      - name: count
        nominalQuota: 32  # Max 32 GPUs in queue

---
# 2. Define Job using Kueue integration
apiVersion: batch/v1
kind: Job
metadata:
  name: training-job-batch
spec:
  # Kueue scheduling class
  queue: gpu-training-queue
  
  parallelism: 8
  completions: 8
  
  template:
    spec:
      containers:
      - name: trainer
        image: training:latest
        resources:
          limits:
            nvidia.com/gpu: 1
      # Kueue automatically adds
      # ProvisioningRequest + tolerations

Kueue benefits:

1. Quota management: enforce max GPU per team/namespace
2. Job fairness: queue-based scheduling, avoid priority inversion
3. Automatic ProvisioningRequest: no manual CRD creation
4. Resource flavor abstraction: "gpu" instead of hardcoding "nvidia-h100"

---

Failure modes and recovery

Partial scheduling (without gang scheduling)

8-GPU training job, no ProvisioningRequest:

t=0: Pod-0 → scheduled, starts
t=1: Pod-1 → scheduled, starts
t=2: Pod-2 → scheduled, starts
t=3: Pod-3 → PENDING (node limit)
t=4: Pod-4 → PENDING
...
t=30min: training framework timeout
          "waiting for 8 processes, got 3"
          → kill all 3 processes
          → pod crashes

Result: 30min of compute waste + job restart

Timeout behavior

ProvisioningRequest has timeout window (default 1 hour):

t=0: ProvisioningRequest created
t=0-5min: Cluster autoscaler tries to provision nodes
t=25min: Provision fails (not enough quota, or zones saturated)
t=60min: ProvisioningRequest timeout → status = FAILED
         → Pods remain PENDING
         
Option 1: retry (delete + recreate ProvisioningRequest)
Option 2: reduce batch size
Option 3: request more quota

Zone failures

Cluster autoscaler provisions nodes in preferred zone, but if zone fails:

Node pool zones: us-central1-{a,b,c}
ProvisioningRequest requests 8 GPUs

Autoscaler tries:
t=0: provision in us-central1-a → zone saturated (quota)
t=5: provision in us-central1-b → success, 4 GPUs
t=10: provision in us-central1-c → success, 4 GPUs

But if topology-sensitive (e.g., multi-host TPU):
→ 4 TPUs in zone-a, 4 in zone-b
→ Inter-zone latency kills all-reduce bandwidth
→ Training crashes

→ Solution: strict zone affinity (all VMs same zone)
           or dynamic scheduling (avoid zone-lock)

Cost implications of gang scheduling

Cost model 1: without gang scheduling
─────────────────────────────────────
Batch job requests 32 GPUs
Cluster autoscaler provisions incrementally (4 GPUs/min):
- t=0-1min: 4 GPUs ready, training starts (1 Pod)
- t=1-2min: 4 more GPUs, training waiting for sync (4 Pods)
- ...
- t=7-8min: all 32 GPUs ready (but partially utilized for 8 min)
- t=8min-3hour: training at full speed

Cost: 3.5 hours × 32 GPUs = 112 GPU-hours

Cost model 2: with gang scheduling (ProvisioningRequest)
──────────────────────────────────────────────────────
- t=0-5min: cluster autoscaler provisions all 32 GPUs (in bulk)
- t=5min: training starts
- t=5min-3hour: training at full speed

Cost: 3.08 hours × 32 GPUs = 98.5 GPU-hours

Savings: 13.5 GPU-hours (12%) by avoiding partial utilization

BUT if timeout happens (5 min provisioning fails):
→ ProvisioningRequest fails, batch must retry
→ Lost cost: 5 min × 32 GPUs = 2.67 GPU-hours (debugging, retry overhead)
→ Break-even if timeout rate < 13.5 / 2.67 = 5% (rare)

---

Integration with training frameworks

PyTorch DDP + ProvisioningRequest

import torch
import torch.distributed as dist

# Assuming ProvisioningRequest + Kueue already allocated all 8 GPUs

def main():
    # Initialize distributed group
    # (assumes all 8 Pods have environment variables set)
    rank = int(os.environ['RANK'])
    world_size = int(os.environ['WORLD_SIZE'])
    
    dist.init_process_group(
        backend='nccl',
        rank=rank,
        world_size=world_size,
        timeout=timedelta(minutes=30)
    )
    
    # Now all 8 processes connected synchronously
    # (ProvisioningRequest ensures all 8 Pods started ~simultaneously)
    
    model = MyModel()
    model = DistributedDataParallel(model)
    
    # Training
    for epoch in range(100):
        train_epoch()
        dist.all_reduce(loss)  # Synchronized across 8 GPUs

JAX multi-host training + ProvisioningRequest

import jax
import jax.numpy as jnp
from jax.distributed import initialize, shutdown

def main():
    # Initialize JAX distributed
    # Pod spec must include REPLICA_COUNT, REPLICA_INDEX env vars
    initialize()
    
    # devices = jax.devices()  # Should see all 8 GPUs (across all Pods)
    
    # Training loop with collective ops
    @jax.jit
    def training_step(state, batch):
        loss, grads = compute_loss_and_grads(state, batch)
        # Synchronize gradients across all devices
        loss = jax.lax.all_reduce(loss, 'sum')
        return update_state(state, grads), loss

---

Dynamic Workload Scheduler (DWS)

Alternative to ProvisioningRequest (queue-based scheduling):

apiVersion: v1
kind: Pod
metadata:
  name: batch-job-1
spec:
  schedulingGates:
  - name: queued-provisioning.gke.io/quota-gate
  
  containers:
  - name: job
    image: training:latest
    resources:
      limits:
        nvidia.com/gpu: 8

# Kubernetes behavior:
# 1. Pod created with scheduling gate
# 2. Scheduler marks Pod as "gated" (not scheduled yet)
# 3. DWS monitors quota availability
# 4. When 8 GPUs available in single node → remove gate
# 5. Scheduler schedules Pod immediately

# Advantage: no need to provision new nodes (use idle resources)

Use case: Batch jobs with flexible deadline, shared cluster.

---

Cost optimization strategies

Strategy 1: Spot VMs + fault tolerance

apiVersion: autoscaling.gke.io/v1beta1
kind: ProvisioningRequest
spec:
  podTemplateSpec:
    spec:
      tolerations:
      - key: cloud.google.com/gke-preemptible
        operator: Equal
        value: "true"
        effect: NoExecute
  # ProvisioningRequest provisions Spot VMs (cheaper)
  # Training job must checkpoint + resume on preemption

Savings: 60-80% on GPU cost, but added retry logic complexity.

Strategy 2: Zone-local scheduling (latency + cost)

apiVersion: autoscaling.gke.io/v1beta1
kind: ProvisioningRequest
spec:
  podTemplateSpec:
    spec:
      affinity:
        nodeAffinity:
          preferredDuringSchedulingIgnoredDuringExecution:
          - weight: 100
            preference:
              matchExpressions:
              - key: topology.kubernetes.io/zone
                operator: In
                values: ["us-central1-a"]  # Prefer single zone
  # Autoscaler provisions in us-central1-a
  # → lower inter-GPU latency
  # → potentially cheaper (fewer cross-zone transfers)

Strategy 3: Reserved capacity

# For predictable high-volume GPU workloads:
# Reserve GPUs via Commitment (monthly/yearly discount)

gcloud compute commitments create gpu-commit \
  --region=us-central1 \
  --plan=one-year \
  --resources=accelerators=nvidia-tesla-h100,count=64

# Then use ProvisioningRequest on reserved nodes (auto-matched)
# Savings: 30% vs on-demand

---

Mental model: Resource atomicity

Key insight:

Without gang scheduling:   With gang scheduling:
─────────────────────      ─────────────────────
Pod 0: [GPU0]              Pods 0-7: requested atomically
Pod 1: [GPU1]              Cluster autoscaler: provision all-or-nothing
Pod 2: [GPU2]              Result: 0 or 8 Pods, never 1-7
Pod 3: pending
...
Partial state → framework timeout → restart

Gang scheduling ensures no partially-initialized states — either job is ready to run, or it's waiting (not starting partial training).

---

References

• GKE ProvisioningRequest
• GKE Deploy Batch System with Kueue
• GKE Best Practices for Batch
• Google Cloud Blog: Using KubeRay and Kueue
• Kueue Documentation
• Kubernetes Scheduling Gates