JobSet: Die neue Lösung für verteilte ML- und HPC-Workloads in Kubernetes

In der Welt der Kubernetes-Entwicklung gibt es spannende Neuigkeiten: JobSet wurde eingeführt, eine Open-Source-API, die speziell für die Verwaltung verteilter Jobs entwickelt wurde. Ziel von JobSet ist es, eine einheitliche API für das verteilte Training von Machine Learning (ML) und High-Performance-Computing (HPC) Workloads auf Kubernetes bereitzustellen.

Warum JobSet?

Die jüngsten Verbesserungen im Batch-Ökosystem von Kubernetes haben das Interesse von ML-Ingenieuren geweckt, die festgestellt haben, dass Kubernetes hervorragend für die Anforderungen an verteilte Trainings-Workloads geeignet ist.

Große ML-Modelle, insbesondere die sogenannten LLMs (Large Language Models), passen oft nicht in den Speicher der GPU- oder TPU-Chips eines einzelnen Hosts. Stattdessen werden sie über Zehntausende von Beschleunigerchips verteilt, die sich über Tausende von Hosts erstrecken.

In solchen Szenarien wird der Code für das Modelltraining häufig containerisiert und gleichzeitig auf all diesen Hosts ausgeführt. Dabei finden verteilte Berechnungen statt, bei denen sowohl die Modellparameter als auch die Trainingsdaten über die Zielbeschleunigerchips verteilt werden. Hierbei kommen kollektive Kommunikationsprimitive wie all-gather und all-reduce zum Einsatz, um Berechnungen durchzuführen und die Gradienten zwischen den Hosts zu synchronisieren.

Diese Merkmale machen Kubernetes zu einer ausgezeichneten Wahl für solche Workloads, da es die Lebenszyklen containerisierter Anwendungen effizient verwalten und terminieren kann. Es ist zudem sehr erweiterbar, was Entwicklern ermöglicht, eigene Kubernetes-APIs, Objekte und Controller zu definieren, um das Verhalten und den Lebenszyklus dieser Objekte zu steuern.

Allerdings können die bestehenden Kubernetes-Primitiven den sich weiterentwickelnden Techniken des verteilten ML-Trainings nicht mehr gerecht werden. Die Landschaft der Kubernetes-APIs für die Orchestrierung von verteiltem Training ist zudem fragmentiert, und jede der bestehenden Lösungen hat spezifische Einschränkungen, die sie für verteilte ML-Trainingssuboptimal machen.

Beispielsweise definiert der KubeFlow-Trainingsoperator benutzerdefinierte APIs für verschiedene Frameworks (z. B. PyTorchJob, TFJob, MPIJob etc.), jedoch sind diese Jobtypen speziell auf die jeweiligen Frameworks zugeschnitten, was zu unterschiedlichen Semantiken und Verhaltensweisen führt.

Die Job-API hat viele Lücken beim Ausführen von Batch-Workloads geschlossen, darunter den Indexierungsabschlussmodus, höhere Skalierbarkeit, Pod-Fehlermuster und Rücksetzrichtlinien für Pods. Allerdings benötigt die Ausführung von ML-Trainings- und HPC-Workloads mit der upstream Job-API zusätzliche Orchestrierung, um folgende Lücken zu schließen:

• Multi-Template-Pods: Die meisten HPC- oder ML-Trainingsjobs beinhalten mehr als einen Pod-Typ. Diese Pods gehören zum selben Workload, müssen jedoch unterschiedliche Container ausführen, verschiedene Ressourcen anfordern oder unterschiedliche Fehlermuster haben. Ein gängiges Beispiel ist das Treiber-Arbeiter-Muster.
• Job-Gruppen: Großangelegte Trainings-Workloads erstrecken sich über mehrere Netzwerk-Topologien und laufen beispielsweise über mehrere Racks. Solche Workloads sind latenzsensitiv und zielen darauf ab, die Kommunikation zu lokalisieren und den Verkehr über die Netzwerklatenz zu minimieren. Um dies zu ermöglichen, muss der Workload in Gruppen von Pods aufgeteilt werden, die jeweils einer Netzwerk-Topologie zugewiesen sind.
• Inter-Pod-Kommunikation: Ressourcen (z. B. headless Services) erstellen und verwalten, die notwendig sind, um die Kommunikation zwischen den Pods eines Jobs herzustellen.
• Startsequenz: Einige Jobs erfordern eine spezifische Startsequenz der Pods; manchmal wird erwartet, dass der Treiber zuerst startet (wie bei Ray oder Spark), in anderen Fällen müssen die Arbeiter bereit sein, bevor der Treiber startet (wie bei MPI).

JobSet zielt darauf ab, diese Lücken zu schließen, indem es die Job-API als Baustein verwendet, um eine reichhaltigere API für großangelegte verteilte HPC- und ML-Anwendungsfälle zu schaffen.

Die Einführung von JobSet ist ein bedeutender Schritt nach vorn für Entwickler und DevOps-Teams, die komplexe ML- und HPC-Workloads effizienter verwalten möchten. Bei ayedo sind wir stolz darauf, Kubernetes-Partner zu sein und freuen uns darauf, Ihnen bei der Umsetzung dieser neuen Möglichkeiten zu helfen.

Quelle: Kubernetes Blog