KAI Scheduler

KAI Scheduler (Kubernetes AI Scheduler) est un scheduler Kubernetes open source développé par NVIDIA, extrait de la plateforme Run:ai et publié en avril 2025. Il remplace le scheduler par défaut de Kubernetes pour les workloads GPU : gang scheduling, files d'attente hiérarchiques avec fairness dynamique, partage GPU fractionnel, et consolidation des ressources.

Idéal pour : clusters GPU partagés entre équipes, entraînements distribués (PyTorch, Ray, Horovod), inférence multi-locataire, optimisation de l'utilisation GPU.

Informations essentielles

Origine : NVIDIA (extrait de Run:ai)  ·  Licence : Apache 2.0  ·  Statut : CNCF Sandbox (2025)

Liens : GitHub  ·  Releases  ·  Annonce NVIDIA

Support : Projet CNCF Sandbox en développement actif - pas de LTS défini. Consulter les release notes avant toute mise à jour, les CRDs peuvent évoluer entre versions mineures.

Stack par défaut

Concepts clés

Queues et hiérarchie

Les queues sont des ressources Kubernetes (kind: Queue) représentant des unités organisationnelles (équipes, départements). Chaque queue définit :

• quota : ressources garanties (GPU, CPU, mémoire) - -1 = illimité
• limit : plafond dur de consommation - -1 = pas de limite
• overQuotaWeight : poids relatif pour la distribution des ressources sur-quota

KAI Scheduler requiert au minimum deux niveaux : département (parent) → équipe (feuille). Les queues feuilles n'ont pas d'enfants - les workloads y sont soumis. Deux queues sont créées automatiquement à l'installation : default-parent-queue et default-queue.

Gang scheduling

Un PodGroup définit un minMember : le nombre minimum de pods qui doivent être schedulables simultanément. Le scheduler alloue tous les pods ou aucun - indispensable pour les entraînements distribués où des pods partiellement schedulés brûlent du budget GPU sans avancer.

Le PodGrouper crée automatiquement les PodGroups en suivant les ownerReferences des pods (PyTorchJob, RayCluster, MPIJob, ArgoWorkflow, etc.) - aucune annotation manuelle nécessaire pour les frameworks supportés.

Partage GPU

Préemption et priorités

Valeur < 100 = préemptible, ≥ 100 = non préemptible. Classes pré-installées :

Prérequis

• Kubernetes ≥ 1.24
• Helm ≥ 3.0
• NVIDIA GPU Operator (requis pour tout workload GPU)
• Droits cluster-admin (installation des CRDs et RBAC)
• Accès réseau à ghcr.io (chart OCI et images)

Installation

helm upgrade -i kai-scheduler \
  oci://ghcr.io/nvidia/kai-scheduler/kai-scheduler \
  -n kai-scheduler \
  --create-namespace \
  --version v0.14.2

Options d'installation fréquentes

# Activer le partage GPU (fractions)
--set global.gpuSharing=true

# Environnement CDI (GB200/GB300)
--set binder.cdiEnabled=true

# OpenShift avec GPU Operator < v25.10.0
--set admission.gpuPodRuntimeClassName=null

Vérifier l'installation

# Pods du namespace kai-scheduler
kubectl get pods -n kai-scheduler
# Attendu : admission, binder, kai-operator, pod-grouper,
#           podgroup-controller, queue-controller, scheduler

# CRDs installés
kubectl get crd | grep kai

Configuration des queues

# queues.yaml
---
apiVersion: scheduling.run.ai/v2
kind: Queue
metadata:
  name: departement-ml
spec:
  resources:
    gpu:
      quota: 16
      limit: 32
      overQuotaWeight: 1
    cpu:
      quota: -1
      limit: -1
      overQuotaWeight: 1
    memory:
      quota: -1
      limit: -1
      overQuotaWeight: 1
---
apiVersion: scheduling.run.ai/v2
kind: Queue
metadata:
  name: equipe-training
spec:
  parentQueue: departement-ml
  resources:
    gpu:
      quota: 10
      limit: 20
      overQuotaWeight: 2
    cpu:
      quota: -1
      limit: -1
      overQuotaWeight: 1
    memory:
      quota: -1
      limit: -1
      overQuotaWeight: 1
---
apiVersion: scheduling.run.ai/v2
kind: Queue
metadata:
  name: equipe-inference
spec:
  parentQueue: departement-ml
  resources:
    gpu:
      quota: 4
      limit: 16
      overQuotaWeight: 1
    cpu:
      quota: -1
      limit: -1
      overQuotaWeight: 1
    memory:
      quota: -1
      limit: -1
      overQuotaWeight: 1

kubectl apply -f queues.yaml
kubectl get queues

Soumettre des workloads

Chaque workload nécessite deux éléments :

1. Label queue : kai.scheduler/queue: <nom-queue>
2. Scheduler : spec.schedulerName: kai-scheduler

Pod GPU entier

apiVersion: v1
kind: Pod
metadata:
  name: training-job
  labels:
    kai.scheduler/queue: equipe-training
spec:
  schedulerName: kai-scheduler
  priorityClassName: train
  containers:
  - name: trainer
    image: nvcr.io/nvidia/pytorch:24.01-py3
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

Pod GPU fractionnel (time-slicing)

apiVersion: v1
kind: Pod
metadata:
  name: inference-light
  labels:
    kai.scheduler/queue: equipe-inference
  annotations:
    gpu-fraction: "0.5"   # 50% d'un GPU
spec:
  schedulerName: kai-scheduler
  priorityClassName: inference
  containers:
  - name: server
    image: nvcr.io/nvidia/tritonserver:24.01-py3
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

RayCluster avec gang scheduling

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: training-distributed
  labels:
    kai.scheduler/queue: equipe-training
spec:
  headGroupSpec:
    template:
      spec:
        schedulerName: kai-scheduler
        priorityClassName: train
        containers:
        - name: ray-head
          resources:
            requests:
              cpu: "2"
              memory: "4Gi"
  workerGroupSpecs:
  - replicas: 4
    minReplicas: 4
    maxReplicas: 4
    template:
      spec:
        schedulerName: kai-scheduler
        containers:
        - name: ray-worker
          resources:
            requests:
              nvidia.com/gpu: "1"
              memory: "8Gi"

Partage GPU - configuration GPU Operator

Le partage fractionnel (time-slicing) nécessite une ConfigMap dans le namespace gpu-operator :

apiVersion: v1
kind: ConfigMap
metadata:
  name: time-slicing-config
  namespace: gpu-operator
data:
  any: |-
    version: v1
    flags:
      migStrategy: none
    sharing:
      timeSlicing:
        replicas: 4
        failRequestsGreaterThanOne: false

Sans isolation mémoire : un pod peut provoquer un OOM sur les autres pods partageant le même GPU. Utiliser MIG (GPU Ampere+) pour une isolation complète en production.

Monitoring

KAI Scheduler expose des métriques Prometheus depuis chaque composant.

# État des queues (utilisation vs quota)
kubectl get queues -o wide

# État des PodGroups
kubectl get podgroups -A

# Config et statut des composants
kubectl get config -n kai-scheduler

# Logs du scheduler
kubectl logs -n kai-scheduler -l app=kai-scheduler --tail=100

# Logs du binder (binding asynchrone)
kubectl logs -n kai-scheduler -l app=binder --tail=100

Un Prometheus interne peut être activé via le Helm chart. Avec Prometheus Operator existant, les ServiceMonitor sont créés automatiquement.

Troubleshooting

Workload bloqué en Pending

# Vérifier l'état du PodGroup
kubectl get podgroup -n <namespace>
kubectl describe podgroup <nom> -n <namespace>

# Événements du pod
kubectl describe pod <nom-pod> -n <namespace>

# Vérifier la queue et ses ressources
kubectl describe queue <nom-queue>

GPU non alloué

# Vérifier que le GPU Operator est actif
kubectl get pods -n gpu-operator

# Vérifier les ressources GPU disponibles sur les nœuds
kubectl describe nodes | grep -A10 "Allocatable"
kubectl get nodes -o custom-columns="NAME:.metadata.name,GPU:.status.allocatable.nvidia\.com/gpu"

Admission Controller en erreur

kubectl logs -n kai-scheduler -l app=admission --tail=100

Commandes utiles

# État des queues
kubectl get queues
kubectl get queues -o wide

# PodGroups actifs
kubectl get podgroups -A

# Pods dans le namespace de réservation GPU
kubectl get pods -n kai-resource-reservation

# Vérifier le scheduler actif d'un pod
kubectl get pod <nom> -o jsonpath='{.spec.schedulerName}'

# Ressources GPU par nœud
kubectl get nodes -o custom-columns="NAME:.metadata.name,GPU:.status.allocatable.nvidia\.com/gpu"

# Pods GPU en cours d'exécution
kubectl get pods -A -o wide | grep nvidia

Mise à jour

KAI Scheduler se met à jour via Helm en bumping la version du chart.

# Vérifier la version installée
helm list -n kai-scheduler

# Mettre à jour le chart
helm upgrade kai-scheduler \
  oci://ghcr.io/nvidia/kai-scheduler/kai-scheduler \
  -n kai-scheduler \
  --version v0.14.2   # remplacer par la version cible

Attention aux CRDs : Helm ne met pas à jour les CRDs automatiquement. En cas de changement de CRD entre versions (indiqué dans les release notes), les appliquer manuellement avant le helm upgrade :

# Vérifier les CRDs installés
kubectl get crd | grep kai

# Appliquer les nouveaux CRDs depuis le source
kubectl apply -f https://raw.githubusercontent.com/NVIDIA/KAI-Scheduler/<version>/manifests/crds/

Consulter les release notes avant chaque mise à jour.

Ressources

• GitHub : https://github.com/NVIDIA/KAI-Scheduler
• Releases : https://github.com/NVIDIA/KAI-Scheduler/releases
• Quickstart : https://github.com/NVIDIA/KAI-Scheduler/blob/main/docs/quickstart/README.md
• Annonce NVIDIA : https://developer.nvidia.com/blog/nvidia-open-sources-runai-scheduler-to-foster-community-collaboration/
• Intégration Ray : https://docs.ray.io/en/latest/cluster/kubernetes/k8s-ecosystem/kai-scheduler.html