AWS Batch bekommt Quota-Management: Teams teilen sich GPU-Kapazität für SageMaker Training

Warum AWS das Quota-Management entwickelt hat

Die Unterstützung von AWS Batch für Amazon-SageMaker-Training-Jobs bot Kunden bereits eine vollständig verwaltete Möglichkeit, Trainings-Workloads für maschinelles Lernen in eine Warteschlange zu stellen, einzuplanen und bei Fehlern erneut auszuführen, ohne eine eigene Infrastruktur aufzubauen. Je mehr Teams hinzukommen, desto mehr Kommunikations-, Abstimmungs- und Konfigurationsaufwand entsteht für Administratoren bei der Koordination des Zugriffs auf gemeinsam genutzte GPU-Kapazität (Graphics Processing Unit), so AWS. Dazu gehören das Einrichten separater Warteschlangen, das Anpassen von Kapazitätsgrenzen und die Steuerung, welche Workloads welche Ressourcen erhalten.

Kunden mit Trainingsumgebungen für mehrere Teams hätten AWS zufolge wiederholt darum gebeten, Zuteilungen pro Team oder Projekt festlegen, ungenutzte Kapazität automatisch teilen und die Workloads mit der höchsten Priorität zuerst ausführen zu können. Das Quota-Management ist die Antwort des Unternehmens auf diesen Wunsch. AWS verknüpft es mit einem zentralen Ziel von AWS Batch: Kunden dabei zu helfen, die Auslastung ihrer Rechenressourcen zu steigern.

So funktionieren Quota-Shares, Verleihen und Preemption

Administratoren definieren für jede Einheit, etwa ein Team oder ein Projekt, einen Quota-Share. Jeder Share enthält Compute-Quoten, also Kapazitätsgrenzen, sowie eine Strategie zur gemeinsamen Nutzung ungenutzter Rechenleistung. Es gibt drei Strategien:

• RESERVE: Ungenutzte Rechenleistung wird zurückgehalten, bis die eigenen Jobs des Shares eintreffen.
• LEND: Ungenutzte Rechenleistung wird an andere Quota-Shares verliehen, die Arbeit zu erledigen haben.
• LEND_AND_BORROW: Der Share verleiht, wenn er untätig ist, und leiht sich Kapazität, wenn er Arbeit hat, die davon profitieren könnte.

Ein Quota-Share funktioniert wie eine virtuelle Warteschlange innerhalb einer Job-Queue. Bei der Planung bevorzugt AWS Batch Shares, die unter ihren Kapazitätsgrenzen liegen. Treffen Jobs in einem Share ein, der seine Kapazität verliehen hat, führt AWS Batch eine Cross-Share-Preemption durch und holt die Ressourcen zurück, sodass Administratoren die ausleihenden Jobs nicht selbst suchen und beenden müssen. Treffen Jobs in einem Share ein, der seine Grenze erreicht oder überschritten hat, prüft AWS Batch, ob ein anderer Share Kapazität zum Ausleihen hat.

Eine prioritätsbasierte Preemption innerhalb eines Shares lässt sich bei der Konfiguration des Shares aktivieren. Sie verdrängt Jobs mit niedrigerer Priorität im Status SCHEDULED, STARTING oder RUNNING, damit Jobs mit höherer Priorität am Anfang der Warteschlange schneller starten können.

Einrichtung mit dem SageMaker Python SDK

Für die Nutzung der Funktion braucht es eine Scheduling Policy mit einer quotaSharePolicy, ein Service Environment, das Kapazitätsgrenzen in Form konkreter Instanzen festlegt, und eine SageMaker-Training-Job-Queue, die mit beiden verbunden ist. Anschließend werden Quota-Shares der Job-Queue zugeordnet. Das AWS-Benutzerhandbuch beschreibt die Konfiguration über die AWS CLI und die Konsole, und ein Beispiel-Notebook zeigt ein durchgängiges Beispiel mit dem Modul aws_batch im SageMaker Python SDK v3.

Das Beispiel im Beitrag legt drei Quota-Shares für drei Projekte an, die jeweils auf eine einzelne ml.g5.xlarge-Instanz begrenzt sind:

• QS1 deckt Hintergrund-Jobs ab. Er nutzt LEND_AND_BORROW mit aktivierter Preemption innerhalb des Shares, was zu nicht dringlicher Arbeit mit Checkpointing passt.
• QS2 deckt Workloads mit mittlerer Priorität ab. Er nutzt LEND ohne Preemption innerhalb des Shares, für Teams, die lieber auf Kapazität warten, als zusätzliche Preemptions in Kauf zu nehmen. AWS Batch holt dennoch jede verliehene Kapazität zurück, die für die Ausführung seiner Jobs nötig ist.
• QS3 deckt Workloads ab, die schnell und mit minimalen Unterbrechungen starten müssen. Er nutzt RESERVE ohne Preemption innerhalb des Shares, sodass seine Jobs von AWS Batch nicht verdrängt werden, solange die Kapazitätsgrenze unverändert bleibt.

Diese Konfigurationen werden an eine Funktion create_quota_management_resources aus dem Beispiel-Notebook übergeben, die auch die Scheduling Policy, das Service Environment und die Job Queue einrichtet. Anschließend übermittelt der Beitrag einen einfachen Hello-World-ModelTrainer an die Queue, um die programmatische Übermittlung zu demonstrieren, die das Modul aws_batch inzwischen unterstützt.

Was das für Plattform-Teams bedeutet

Für Organisationen, die begrenzte GPU-Kapazität auf mehrere Teams verteilen, verlagert die Funktion die Zuteilungsregeln in den Scheduler, statt sie manuell abzustimmen. Teams können den Kompromiss wählen, der zu ihrer Arbeit passt: garantierte Kapazität ohne Preemption, das Verleihen ungenutzter Kapazität mit der Möglichkeit, sie zurückzuholen, oder das Ausleihen freier Kapazität für unterbrechbare Jobs. Voraussetzungen sind ein AWS-Konto mit Berechtigungen zur Verwaltung von AWS-Batch-Ressourcen sowie das SageMaker Python SDK v3.