AI21 Labs ersetzt Slack-Verhandlungen um GPUs durch Kueue auf einem GKE-Cluster mit 10.000 GPUs

Vom Chat-Kanal zur automatisierten Ressourcenplanung

AI21 zufolge begann die gemeinsame GPU-Nutzung als Gespräch. Die Teams nutzten einen internen Messaging-Kanal namens #gpu-resources: Wer einen GPU-Typ brauchte, der vollständig ausgelastet war, fragte im Kanal nach und hoffte, dass ein anderes Team Kapazität freigibt.

Solange der Cluster Reserven hatte, funktionierte dieses Vorgehen. Sobald die Auslastung dauerhaft nahe 100 % lag, was AI21 für eine teure, reservierte GPU-Flotte als Ziel beschreibt, wurde jede Anfrage zu dem, was das Unternehmen ein Nullsummenspiel um Verdrängung nennt. Als Folgen nennt AI21 verschwendete Ingenieurzeit für Triage und Verhandlungen, unfaire Warteschlangen, brachliegende Kapazität in einigen Bereichen bei gleichzeitigem Mangel in anderen sowie die Abhängigkeit von ad hoc eingesetzten „Zauberbefehlen“, um Kapazität freizuräumen.

Die Fallstudie bezieht sich auf einen gemeinsamen GKE-Cluster auf Google Cloud. AI21 betreibt ihn bewusst als einen einzigen Pool statt als je einen Cluster pro Team, weil sich in Kubernetes zwar leicht viele Cluster anlegen lassen, Kapazität zwischen ihnen aber nur schwer teilen lässt.

Vier Designprinzipien und warum reines Kubernetes nicht reichte

AI21 definierte vier Säulen für ein ideales GPU-Management-System:

• Fairness: Scheduling-Entscheidungen richten sich nach der Nutzung der Konten, nicht nach dem Prinzip „Wer zuerst kommt“ oder „Wer am lautesten ist“.
• Ordnung: Die Warteschlange spiegelt Geschäftsprioritäten wider, nicht nur Zeitstempel.
• Effizienz: Entwickler müssen ungenutzte oder nicht kritische Workloads nicht mehr manuell löschen.
• Technische Integrität: Jobs mit mehreren Pods (Gang-Jobs) starten erst, wenn alle Ressourcen bereitstehen.

Zunächst prüfte das Unternehmen, was Kubernetes mitbringt. PriorityClass arbeitet auf Pod-Ebene und kennt keine Gang-Größe, sodass ganze Läufe mit niedrigerer Priorität verdrängt werden können, selbst wenn nur eine von mehreren GPUs benötigt wird. ResourceQuota erlaubt kein Ausleihen zwischen Teams bei freier Kapazität, und Pods über dem Kontingent werden bei der Zulassung abgelehnt, statt in eine Warteschlange zu kommen. LimitRange fungiert als Leitplanke, nicht als Scheduler. Job und Indexed Job kennen kein Alles-oder-nichts-Prinzip: Werden drei von vier Pods eingeplant und der vierte passt nicht mehr hinein, laufen die ersten drei trotzdem weiter und verbrennen GPU-Zeit.

Da GKE keinen Batch-Scheduler mitbringt, evaluierte AI21 Apache YuniKorn, Volcano und Kueue. Volcano habe damals mehr Funktionen geboten, sich aber nicht nahtlos in die eigene Umgebung integrieren lassen, und YuniKorn habe nicht alle Anwendungsfälle abgedeckt. Kueue überzeugte durch Einfachheit und Integration.

Wie Kueue die Anforderungen abbildet

AI21 beschreibt Kueue als Manager auf Job-Ebene, der entscheidet, wann ein Job zum Start zugelassen wird und wann er enden sollte. Laut dem Unternehmen erfüllte es die vier Säulen wie folgt:

• Fairness: Preemption Fair Sharing gleicht die Cluster-Nutzung durch Verdrängung von Workloads aus, während Admission Fair Sharing (AFS) Workloads von Teams zuerst zulässt, die historisch weniger genutzt haben.
• Ordnung: Workload Priority Classes steuern sowohl die Reihenfolge in der Warteschlange als auch die Verdrängungsreihenfolge.
• Effizienz: Die ClusterQueue gibt Admins fein abgestufte Kontrolle über die Verdrängung und integriert Fair Sharing und Workload Priority Classes, zudem räumt Kueue Workloads automatisch auf.
• Technische Integrität: Alles-oder-nichts-Semantik funktioniert mit Indexed Jobs, sodass Multi-Pod-Gangs nur zugelassen werden, wenn jeder Pod laufen kann.

Der Cluster verarbeitet drei Kapazitätstypen (reserviert, Spot und der Flex-Start-Modus des Dynamic Workload Scheduler) und drei Arten von Workloads: Debug-Pods, Trainingsjobs und Inferenzmodelle, die von Reinforcement-Learning-Jobs genutzt werden. Nach Gesprächen mit den Algorithm Developer Leads, die die Clusternutzung verantworten, definierte AI21 mehrere Pools: ein garantiertes Kontingent pro Team, das nie verdrängt wird, gemeinsame garantierte Pools für Single-Node- und Multi-Node-Jobs, einen opportunistischen Pool für Best-Effort-Arbeiten wie Debug-Pods und einen On-Demand-Pool für sonstige Fälle. Laut AI21 ist der opportunistische Pool der Hebel, der die Auslastung hoch hält.

Eine Zusammenarbeit, die in Kueue zurückfloss

AI21 zufolge entwickelte sich die Einführung von Kueue zu einer engen Zusammenarbeit mit dem dahinterstehenden Google-Team, wobei die Anforderungen aus dem Produktivbetrieb in die Roadmap von Kueue einflossen. Dazu gehören Funktionen wie Admission Fair Sharing, das nach Angaben des Unternehmens inzwischen allen Nutzern zur Verfügung steht. Der vorliegende Quelltext bricht bei der Beschreibung des ersten Design-Rollouts mit der Bezeichnung v0.10 ab, weitere Ergebnisse und Erkenntnisse sind daher hier nicht abgedeckt.

Für Teams, die große Trainings-Workloads auf gemeinsam genutzten Kubernetes-Clustern betreiben, liefert der Fall ein konkretes Beispiel dafür, wie sich ad hoc geführte Verhandlungen durch richtliniengesteuerte Zulassung, Priorisierung und Verdrängung ersetzen lassen.