JobSet: Neue API fĂĽr verteilte ML- und HPC-Anwendungen auf Kubernetes
Die neue Open-Source-API JobSet soll flexiblere und vielfältigere Konfigurationsoptionen für groß angelegte verteilte HPC- und ML-Anwendungsfälle liefern.
(Bild: Gorodenkoff / Shutterstock.com)
Kubernetes und dessen System für die Batch-Verarbeitung von Workloads sind grundsätzlich für verteiltes Machine-Learning-Modelltraining sowie High-Peformance-Compute-Anwendungen geeignet. Speziell größere Compute-Aufgaben wie etwa Large Language Models (LLMs), die angesichts begrenzter Speicherressourcen auf GPUs und TPUs über eine Vielzahl von Hosts hinweg verteilt werden müssen, profitieren vom containerisierten Einsatz auf Kubernetes. Bereits darauf ausgelegte Implementierungen wie die Job API oder der KubeFlow Operator lassen in der Praxis allerdings noch einige Konfigurationsoptionen vermissen – etwa die Kommunikation zwischen Pods, unterschiedliche Pod-Templates und Job-Gruppen betreffend. Einen anderen Ansatz für das Darstellen verteilter Jobs soll nun die neue Open-Source-API JobSet liefern.
JobSet setzt auf der Job API auf und erweitert sie
Auf der Job API aufbauend modelliert JobSet verteilte Batch-Workloads als eine Gruppe von Kubernetes-Jobs. Entwicklerinnen und Entwicklern eröffnet sich dadurch die Möglichkeit, verschiedene Pod-Templates unterschiedlichen Pod-Gruppen wie etwa Leader, Worker usw. zuzuweisen. Um auf deklarative Weise identische untergeordnete Jobs zu erstellen, die sich verteilt auf dedizierte Hardware-Beschleunigerbereiche (über Hochgeschwindigkeitsverbindungen vernetzte GPUs oder TPUs gleichen Typs) ausführen lassen, greift JobSet auf ReplicatedJob zurück – eine Vorlage, die eine bestimmte Anzahl von Jobreplikaten umfasst. Für die Kommunikation zwischen den Pods in den einzelnen Bereichen sieht JobSet einen Headless Service vor, der automatische Konfiguration und Lifecycle-Management sicherstellt.
(Bild: kubernetes.io)
JobSet erlaubt es zudem, Kind-Jobs innerhalb einer Topologie-Domäne explizit zuzuweisen – beispielsweise einem der dedizierten Hardware-Beschleunigerbereiche. Dadurch lassen sich unter anderem bestimmte Trainingsmethoden für ML-Modelle wie Distributed Data Parallel (DDP) umsetzen, bei denen nur jeweils eine Modell-Replica pro Hochgeschwindigkeits-Beschleunigerbereich ausgeführt wird und nur die Synchronisation der Replicas über die langsamere bereichsübergreifende Vernetzung erfolgt.
Darüber hinaus bietet JobSet konfigurierbare Erfolgs- und Fehlerrichtlinien. So können Entwicklerinnen und Entwickler beispielsweise eine Policy definieren, die festlegt, wie häufig ein JobSet nach einem Fehler maximal neu starten soll. Wird dann ein Job als fehlgeschlagen markiert, wird das gesamte JobSet neu erstellt, sodass der betreffende Workload ab dem letzten Prüfpunkt wieder aufgenommen werden kann.
Das Einsatzspektrum von JobSet und die wichtigsten bislang verfügbaren Funktionen fasst der Beitrag im Kubernetes-Blog zusammen. Anhand eines Beispiels für verteiltes ML-Training mit dem ML-Framework Jax demonstrieren die Autoren darin auch, wie sich JobSet für einen TPU-Multislice-Workload konfigurieren lässt. Das Entwicklungsteam der API will künftig weitere Funktionen ergänzen, die sich der Übersicht in der JobSet-Roadmap entnehmen lassen.
(map)