JobSet: API для навчання на Kubernetes

Автори: Даніель Вега-Майре (Google), Абдулла Гарайбе (Google), Кевін Ханнон (Red Hat)

JobSet — це інноваційний API з відкритим кодом, призначений для представлення та оркестрації розподілених завдань. Його мета полягає в об’єднанні складнощів розподіленого навчання машин (ML) та робочих навантажень високопродуктивних обчислень (HPC) в єдиному, уніфікованому API, який працює на базі Kubernetes. У цій статті розглядається архітектура JobSet, технічні деталі та його трансформаційна роль у керуванні контейнеризованими навантаженнями на різноманітних обчислювальних островах.

Чому JobSet?

Недавні покращення в екосистемі пакетних завдань Kubernetes відіграли ключову роль у залученні інженерів ML до цієї платформи. Здатність платформи безшвидко планувати та управляти розподіленими навчальними завданнями робить її відмінним вибором для задоволення високих вимог великих ML та HPC навантажень.

Сучасні великі моделі ML, особливо великі мовні моделі (LLMs), часто перевищують обсяги пам’яті одного GPU або TPU. Як наслідок, навчання моделей розподіляється на десятки тисяч прискорювачів, розміщених на сотнях або тисячах вузлів. Ці контейнеризовані застосунки виконуються одночасно на кількох хостах, виконуючи складні розподілені операції за допомогою комунікаційних примітивів, таких як all-gather та all-reduce для синхронізації градієнтів між вузлами.

Вроджені переваги Kubernetes у оркестрації контейнерів, плануванні ресурсів та управлінні життєвим циклом роблять його ідеальною платформою для підтримки цих розподілених завдань. Його розширюваність дозволяє розробникам створювати кастомізовані API та контролери, хоча інноваційні випадки використання, такі як розподілене ML, все ще перевіряють межі наявних примітивів Kubernetes.

Історично склалося так, що оркестраційні фреймворки, такі як Kubeflow, надавали кастомні визначення, такі як PyTorchJob, TFJob та MPIJob, адаптовані до конкретних фреймворків. Проте ці специфічні рішення для фреймворків внесли фрагментацію в екосистему розподіленого навчання. JobSet вирішує цю проблему, використовуючи більш загальний Job API як будівельний блок і розширюючи його для врахування додаткових вимог — таких як мультишаблонні поди, групування завдань за мережевими топологіями та контрольована послідовність запуску — що є важливими для сучасних навантажень HPC та ML.

Як працює JobSet

JobSet розглядає розподілене навантаження як агрегат кількох Kubernetes Jobs. Ця стратегія дозволяє користувачам визначати різні шаблони подів для різних ролей, таких як лідер, працівники чи сервіси параметрів. Інтегруючи абстракцію ReplicatedJob — шаблон завдання, пов’язаний із зазначеною кількістю реплік — JobSet спрощує процес розгортання ідентичних дочірніх завдань на різних островах прискорювачів, усуваючи потребу в громіздкому сценарію або управлінні через Helm charts.

Ключові можливості JobSet включають:

• Репліковані завдання: Сучасні центри обробки даних зазвичай містять групи однорідних прискорювачів, з’єднаних через мережеві канали з високою пропускною здатністю. JobSet розподіляє великі розподілені завдання на менші, ідентичні завдання, що виконуються на островах прискорювачів. Цей підхід мінімізує міжострівну комунікацію та обіцяє покращення продуктивності в стратегіях розподіленого паралельного навчання даних.
• Автоматичне управління безголовими сервісами: Комунікація між подами є критично важливою для синхронізації навчальних навантажень. JobSet автоматично налаштовує безголові сервіси, забезпечуючи надійність комунікації між подами через імена хостів, що безперешкодно управляється протягом життєвого циклу завдання.
• Конфігуровані політики успіху та невдачі: Користувачі можуть тонко налаштувати поведінку JobSet за допомогою політик, що визначають, коли завдання вважається завершеним і скільки перезапусків дозволено у разі невдачі. Ця гнучкість дозволяє адміністраторам визначати умови успіху — наприклад, вимагати, щоб усі поди в певній групі успішно завершилися перед завершенням загального завдання — та автоматично відновлюватися після невдач.
• Ексклюзивне розміщення за топологічними доменами: Для завдань, чутливих до затримок, JobSet забезпечує ексклюзивне 1:1 призначення дочірніх завдань до топологічних доменів, таких як стійки або острови прискорювачів. Це оптимальне розміщення гарантує, що міжвузлова комунікація відбувається переважно через високошвидкісні з’єднання в межах одного домену, з мінімальною залежністю від мереж центрів обробки даних з нижчою пропускною спроможністю.
• Інтеграція з Kueue: JobSet підтримує подання через Kueue, що дозволяє перевищувати кількість кластерів, чергувати навантаження та забезпечувати багатокористувацький доступ. Ця інтеграція запобігає сценаріям часткового планування та потенційним блокуванням, тим самим підвищуючи загальну ефективність кластера.

Приклад використання: Розподілене ML навчання на кількох слайсах TPU з Jax

Використання архітектури мультислайсів TPU

Цей приклад демонструє конфігурацію JobSet, спрямовану на завдання мультислайсів TPU на 4 слайсах TPU v5e. Такі кластери, визначені архітектурою прискорювача, виграють від точного розміщення подів і розподілу ресурсів, що є важливим для навантажень, що використовують Jax. Jax пропонує рідну компіляцію Just-In-Time через OpenXLA, тоді як PyTorch/XLA залишається життєздатною альтернативою.

Нижче наведений приклад YAML-специфікації демонструє, як JobSet обробляє вимоги до планування, забезпечуючи, щоб кожне дочірнє завдання отримувало ексклюзивне використання призначених слайсів TPU з мінімальною ручною конфігурацією.

# Запустіть просте навантаження Jax на 
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice
  annotations:
    # Призначте ексклюзивне використання слайсів TPU
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 3
  replicatedJobs:
  - name: workers
    replicas: 4 # Встановіть кількість слайсів TPU
    template:
      spec:
        parallelism: 2 # Кількість ВМ на слайс TPU
        completions: 2 # Кількість ВМ на слайс TPU
        backoffLimit: 0
        template:
          spec:
            hostNetwork: true
            dnsPolicy: ClusterFirstWithHostNet
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
              cloud.google.com/gke-tpu-topology: 2x4
            containers:
            - name: jax-tpu
              image: python:3.8
              ports:
              - containerPort: 8471
              - containerPort: 8080
              securityContext:
                privileged: true
              command:
              - bash
              - -c
              - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  python -c 'import jax; print("Глобальна кількість пристроїв:", jax.device_count())'
                  sleep 60
              resources:
                limits:
                  google.com/tpu: 4

Глибший технічний аналіз

Дизайн JobSet сильно вплинений необхідністю точності в сучасних розподілених обчисленнях. Розкладаючи одне велике завдання на набір менших, взаємопов’язаних завдань, JobSet надає детальний контроль над розподілом ресурсів і мережею:

• Ізоляція та контроль ресурсів: Можливість визначати селектори вузлів та ексклюзивні політики топології дозволяє адміністраторам ізолювати навантаження в межах конкретних мережевих доменів. Це особливо важливо, коли навантаження є дуже чутливими до затримок в мережі та потребують оптимізованих каналів зв’язку з високою пропускною здатністю.
• Планування на основі політики: Завдяки конфігурованим політикам успіху та невдачі JobSet дозволяє оркестраційному механізму автоматично управляти повторними спробами та перезапусками завдань. Ці функції знижують експлуатаційні витрати, необхідні для моніторингу розподілених завдань, і підвищують стійкість конвеєрів навчання ML.
• Декларативна реплікація завдань: Введення шаблонів реплікації завдань через абстракцію ReplicatedJob значно зменшує складність управління великими розгортаннями. Це дозволяє розробникам налаштовувати однорідні репліки завдань без необхідності вручну дублювати код або конфігурації, при цьому враховуючи різні образи контейнерів, запити ресурсів та налаштування безпеки.

Експертні думки та вплив на галузь

Лідери в області хмарних обчислень та розподілених систем високо оцінюють JobSet за його потенціал стандартизувати оркестрацію розподілених завдань. Експерти зазначають, що абстрагуючи складнощі інфраструктури, JobSet може прискорити цикл розробки моделей ML та спростити операції HPC.

За словами кількох аналітиків галузі, уніфіковані API, такі як JobSet, можуть стати важливими у наступному поколінні досліджень штучного інтелекту, де розподілені робочі процеси потребують масштабованих, стійких та керованих оркестраційних шарів. Тісна інтеграція з Kubernetes і Kueue також відповідає ширшому зсуву в галузі до контейнерно-орієнтованих архітектур, що підкреслює портативність і сумісність між хмарами.

Майбутня робота, тенденції та залучення спільноти

Дорожня карта для JobSet включає кілька обіцяючих вдосконалень. Заплановані функції включають покращені протоколи міжподової комунікації, динамічне масштабування ресурсів, покращену спостережуваність та додаткову конфігурованість для гетерогенних навантажень. Постійні внески спільноти є ключем до його еволюції, і зацікавлені розробники запрошуються ознайомитися з дорожньою картою JobSet для детального списку майбутніх функцій.

Крім того, очікуються подальші вдосконалення в інтеграції JobSet з новими хмарними парадигмами, такими як безсерверні обчислення та обчислення на краю. Ці інтеграції, ймовірно, розширять його застосування за межі традиційних центрів обробки даних, роблячи його основою в оркестрації розподілених систем у гібридних середовищах.

Учасники та члени спільноти можуть долучитися до проекту через його репозиторій, поштовий список або на Slack-каналі. Відгуки, повідомлення про помилки та пропозиції щодо функцій вітаються, оскільки вони забезпечують, щоб JobSet залишався на передовій оркестрації розподілених навантажень.

Висновок

JobSet є значним кроком уперед в оркестрації великих розподілених навантажень. Забезпечуючи уніфікований API, який використовує рідні можливості Kubernetes, він заповнює розрив між навчанням ML, HPC та новими потребами розподілених систем. Завдяки своїм розширеним функціям, гнучким конфігураціям та потужній підтримці спільноти, JobSet готовий стати де-факто рішенням для оркестрації багатоядерних розподілених завдань у складних обчислювальних середовищах сьогоднішнього дня.