Недавно компания Google Cloud опубликовала нашу историю об AI Hypercomputer, включая сокращение времени запуска высокоприоритетных рабочих нагрузок на 83%. Ниже приводится подробный технический разбор того, как мы этого добились с помощью Kueue на GKE.

Кратко

Совместное использование большого парка GPU разными командами — это постоянная проблема для всех, кто запускает рабочие нагрузки по обучению ИИ в Kubernetes. В этом блоге мы делимся примером из практики одного из наших общих кластеров GKE (около 10 000 GPU, используемых несколькими командами) о том, как мы перешли от ручной координации через внутренние каналы связи к Kueue — открытой системе очередей заданий для Kubernetes, которая автоматически принимает, расставляет приоритеты и вытесняет рабочие нагрузки. То, что начиналось как внедрение Kueue, переросло в тесное сотрудничество с командой разработчиков в Google, в ходе которого наши реальные производственные потребности напрямую влияли на дорожную карту Kueue. Мы расскажем о том, что мы пробовали, о результатах (драматическое падение таких показателей, как частота ручного вмешательства и время нахождения критически важного задания в «голодающем» состоянии), а также об уроках эксплуатации в продакшене, которые вернулись в виде улучшений в саму Kueue, включая новые функции, такие как Fair Sharing при допуске (AFS), которая теперь доступна всем пользователям.

Проблема совместного использования

Ни для кого не секрет, что GPU — это дефицитный ресурс. Каждый, кто занимается обучением моделей (что требует значительного использования GPU в течение длительного времени), знаком с этой проблемой не понаслышке. Их ограниченность означает, что управление распределением GPU между командами исследователей, MLOps, инженеров и alignment становится особенно острым во время обучения, вызывая сложные вопросы, такие как:

Как справедливо разделить конечный и дорогой ресурс между несколькими командами?
Кому отдается приоритет, когда спрос превышает емкость?
Стоит ли прерывать выполняющуюся рабочую нагрузку, чтобы освободить место для более приоритетной?
Как примирить команды с разными приоритетами и разными определениями «срочности»?

«До»: согласование ресурсов в командных каналах

Поначалу мы решали эти вопросы во внутреннем канале связи #gpu-resources. Когда кому-то требовался тип GPU, который был полностью занят, он спрашивал в канале и надеялся, что кто-то освободит мощности.

Такое согласование через каналы работало, пока в кластере был запас прочности. Как только загрузка зафиксировалась на уровне около 100% (что является целью для дорогого зарезервированного парка GPU), каждый входящий запрос превратился в разговор об эвакуации с нулевой суммой, и канал зашел в тупик.

При ручном согласовании трещины в этой системе проявились быстро:

Реальное инженерное время тратится впустую на сортировку и согласование
Несправедливое формирование очередей
Простаивающие мощности в одних сегментах, в то время как другие испытывали нехватку
Зависимость от специальных «магических команд» для освобождения достаточной емкости рабочих нагрузок

Нам нужно было создать что-то лучшее. Мы начали с разработки принципов проектирования, которые должны лечь в основу идеальной системы управления GPU.

Четыре столпа идеальной системы управления GPU

Справедливость: решения о планировании принимаются систематически и в соответствии с использованием учетных записей, а не по принципу «первым пришел — первым обслужен» или «кого громче слышно, тот и прав».
Порядок: очередь рабочих нагрузок отражает бизнес-приоритеты, а не просто временные метки.
Эффективность: разработчики избавлены от необходимости вручную удалять свои простаивающие или некритические рабочие нагрузки.
Техническая целостность: многоподовые (групповые) задания не запускаются до тех пор, пока все ресурсы не будут готовы, что предотвращает частичные развертывания.

От принципов проектирования к реализации: создание системы управления GPU

Наши строительные блоки

Для этого блога мы подробно рассмотрим один из наших общих кластеров GKE на Google Cloud (GCP) — тот, на котором мы развернули эту систему. В нем объединено порядка 10 000 GPU, используемых несколькими командами.

Мы намеренно запускаем его как единый общий пул, а не выделяем по кластеру на команду: Kubernetes позволяет легко развернуть множество кластеров, но сложно делиться емкостью между ними. Сохранение всего в одном пуле позволяет поддерживать высокую степень использования — и именно это создает нагрузку на то, как мы планируем внутри него.

Вот конфигурация этого кластера:

У нас есть три типа емкости машин — зарезервированная, спотовая и режим flex-start планировщика Dynamic Workload Scheduler, — которые нам необходимы для запуска трех различных типов исследовательских рабочих нагрузок обучения:

Поды для отладки: разработчики могут захотеть запустить отладку удаленно в реальной производственной среде, а не локально на своих ноутбуках
Задания обучения: пакетные рабочие нагрузки, которые могут масштабироваться от 1 до сотен заданий за секунду. Они могут выглядеть как задания с одним подом или как многоподовые задания, для запуска которых требуется несколько машин параллельно. Мы используем индексированные задания (Indexed Jobs) для управления параллелизмом.
Модели вывода: используются нашими рабочими нагрузками обучения с подкреплением и управляются как развертывания Kubernetes, масштабируемые вверх или вниз с помощью HPA.

Каждую из этих рабочих нагрузок можно дополнительно классифицировать по их способности к прерыванию:

Непрерываемые: прерывание их на середине выполнения — это автоматический красный свет, так как это дорого, неэффективно и потенциально опасно для критически важных по времени рабочих нагрузок.
Прерываемые: прерывание их на середине выполнения имеет незначительные или пренебрежимо малые последствия.

Поиск решения: оценка стандартного Kubernetes

Перед оценкой сторонних инструментов мы проверили, что входит в стандартную поставку Kubernetes. Вот что мы увидели:

Нам пришлось обратиться к другим опциям пакетных планировщиков, чтобы найти ту, которая соответствовала бы всем нашим четырем критериям. Поскольку GKE не поставляется с пакетным планировщиком, мы оценили три: Apache YuniKorn, Volcano и Kueue. В то время у Volcano было больше функций, но ее интеграция с нашей средой не была бесшовной, а YuniKorn не покрывала все наши сценарии использования. С точки зрения простоты и интеграции победила Kueue.

Работа с Kueue

Kueue — это набор API и контроллер для очередей заданий. Это менеджер на уровне заданий, который определяет, когда задание должно быть допущено к запуску (т. е. могут быть созданы поды) и когда оно должно быть остановлено (т. е. активные поды должны быть удалены).

Шаг 1: Проверка системных критериев

Сопоставив наши четыре столпа с функциями Kueue, мы увидели, что она может удовлетворить все наши требования:

Справедливость: Kueue предлагает два взаимодополняющих механизма: Preemption Fair Sharing (вытесняет рабочие нагрузки для балансировки использования кластера) и Admission Fair Sharing (AFS) (допускает в первую очередь рабочие нагрузки от команд, которые исторически использовали меньше ресурсов).
Порядок: Классы приоритета рабочих нагрузок (Workload Priority Classes) устраняют пробел, который мы обнаружили в стандартных PriorityClasses Kubernetes, управляя как порядком очереди, так и порядком вытеснения. Admission Fair Sharing добавляет еще один рычаг поверх упорядочивания по временным меткам.
Эффективность: основной ресурс Kueue, ClusterQueue, дает администраторам детальный контроль над поведением вытеснения и объединяет Fair Sharing + классы приоритета рабочих нагрузок. Kueue выполняет очистку рабочих нагрузок автоматически.
Техническая целостность: семантика «все или ничего» в Kueue органично интегрируется с нашими индексированными заданиями: многоподовые группы допускаются к работе только тогда, когда может работать каждый под.

Шаг 2: Определение требований к реализации

Kueue поддерживал основные функции, необходимые для нашей внутренней системы управления GPU. Тем не менее, для реализации проекта нам требовалось детализировать ряд дополнительных требований. Мы встретились с нашими ключевыми заинтересованными лицами — лидерами команд алгоритмических разработчиков, которые управляют использованием кластера, чтобы точно определить требуемое от системы поведение:

Шаг 3: Внедрение нашей первой архитектуры (v0.10)

С точки зрения администратора, в первой версии мы заложили три уровня емкости:

Гарантированная емкость (Guaranteed capacity): У каждой команды есть собственная квота для одноузловых рабочих нагрузок, общий пул для многоузловых рабочих нагрузок и еще один общий пул для одноузловых рабочих нагрузок.
Вытесняемая емкость (Preemptible capacity): Единый общий пул с честным распределением (fair-sharing) между тенантами, имеющий возможность заимствовать ресурсы у любой простаивающей гарантированной емкости. Именно это поддерживает высокую степень утилизации и помогает распределять ресурсы между командами.
Емкость по требованию (On-demand capacity): Для случаев, когда система переполнена и ни одна рабочая нагрузка не может быть безопасно вытеснена: разработчики могли динамически выделять новые машины. Также полезна для тестирования типов GPU, отсутствующих в зарезервированном пуле.

С точки зрения разработчика, мы хотели сделать маршрутизацию в очередь максимально простой. Достаточно выбрать метки, описывающие затраты на рабочую нагрузку (низкие/высокие), возможность вытеснения (false/true) и приоритет (низкий/средний/высокий/критический); эти ответы вместе с информацией о том, является ли рабочая нагрузка одноузловой или многоузловой, определяют, в какую именно очередь она в итоге попадет.

Где v0.10 не справилась: Справедливость и фрагментация

Проблема справедливости (The fairness gap)

На момент реализации этого решения функция Admission Fair Sharing (AFS) еще не существовала, и мы были ограничены использованием преемпшена Fair Sharing в нашей системе. Это помешало нам одновременно реализовать два требования для многоузловых систем:

Честное планирование между тенантами на основе чипов
Отсутствие вытеснения в этой квоте между разными командами, если только кто-то явно не указал «критический приоритет» для своей индексированной задачи

Следуя идеологии открытого исходного кода и вклада в сообщество, мы напрямую озвучили эту проблему команде Kueue; в ответ они оперативно выпустили AFS, которая обеспечивает два важнейших сценария поведения:

Вытесняет рабочие нагрузки команды, которая обделяет ресурсы других команд
Изменяет порядок рабочих нагрузок в очереди, чтобы в приоритете оказались команды с низким уровнем использования

Для нас это пример идеальной обратной связи, которая подпитывает любой сильный проект с открытым исходным кодом: выявление требований к использованию на основе реальных сценариев и передача этих знаний обратно в продукт.

Это также означало, что теперь у нас есть решение проблемы справедливости.

Проблема фрагментации (The fragmentation problem)

Мы столкнулись и со второй проблемой: иногда поды принимались к выполнению, хотя ни одна отдельная машина фактически не могла удовлетворить требования пода. Сумма наших номинальных квот равна общему количеству имеющихся у нас GPU, поэтому поначалу это выглядело как баг.

Вот сценарий, который это выявил: каждая машина имеет 8 GPU, и мы отправляем рабочую нагрузку, которой требуются все 8 на одном узле. Если емкость GPU распределена как 1+1+4+2 по четырем узлам, математика квот говорит, что свободно 8 GPU, но задаче на 8 GPU негде разместиться. С таким типом фрагментации мы и столкнулись.

Хотя GKE делает все возможное для оптимальной упаковки (bin-pack), Kueue не является планировщиком в нашей конфигурации и не знает топологию нашего кластера (8 GPU на машину).

Найденным решением стала функция Kueue, которую мы не использовали в v0.10: планирование с учетом топологии (Topology Aware Scheduling, TAS). Хотя TAS часто активируется для пакетного планирования (gang scheduling) — например, чтобы держать распределенные рабочие нагрузки физически близко друг к другу для снижения задержек и затрат — наша мотивация была иной.

У нас часто было больше 8 свободных GPU по всему кластеру, но не было ни одного узла с 8 свободными. Простой кастомный ресурс под названием Topology позволяет администраторам указать Kueue, какие топологии следует учитывать при принятии решений о планировании. Технически Kueue не является планировщиком, но при настройке TAS Kueue устанавливает nodeName для принятых рабочих нагрузок. Для наших целей мы получили новый планировщик.

Имея на руках AFS и TAS, мы вновь встретились с нашими ключевыми заинтересованными лицами — лидерами команд алгоритмических разработчиков. Исходный бэклог теперь можно было решить:

AFS обеспечила честное многоузловое совместное использование
А TAS означала, что Kueue может отклонять поды, которые не помещаются; планировать с помощью алгоритма LeastFreeCapacity, который в первую очередь размещает новые рабочие нагрузки на наиболее загруженных узлах для минимизации фрагментации; и прекратить принятие решений о вытечнении, которые фактически не освобождают емкость.

Теперь мы решили проблему справедливости и проблему фрагментации.

Но, как и на любой хорошей встрече, мы вышли из обсуждения с двумя новыми запросами от наших стейкхолдеров:

Вытеснение на основе приоритетов внутри пула вытесняемых одноузловых ресурсов. Одноузловая вытесняемая рабочая нагрузка с priority=low может быть вытеснена задачами со средним/высоким приоритетом в том же пуле.
Новый многоузловой пул с негарантированным качеством обслуживания (best-effort lane). Многоузловые рабочие нагрузки могут использовать его с priority=low; они вытесняются средними/высокими многоузловыми нагрузками при нехватке емкости и при этом сами вытесняют одноузловые вытесняемые рабочие нагрузки, когда это необходимо.

Учитывая то, что мы узнали об AFS и TAS, а также два новых запроса, мы переработали нашу реализацию Kueue до версии v0.15.

Ресохранение с включением AFS + TAS (v0.15)

Наша обновленная версия v0.15 сохранила структуру v0.10, но включила два изменения для поддержки необходимых сценариев поведения, выявленных при первом внедрении:

Вытесняемая емкость теперь учитывает приоритеты. Общий пул одноузловых вытесняемых ресурсов разделен по приоритету рабочих нагрузок: нагрузки с priority=low выполняются в собственном сегменте и вытесняются первыми, когда средним/высоким нагрузкам требуется место. Рядом с ним в архитектуру добавляется новый многоузловой пул best-effort: многоузловые рабочие нагрузки могут использовать его с priority=low, вытесняя при необходимости одноузловые вытесняемые рабочие нагрузки и в то же время вытесняясь средними/высокими многоузловыми нагрузками.
Прием и размещение стали умнее. Admission Fair Sharing (AFS) теперь перебалансирует порядок в очереди на основе исторических данных об использовании чипо-часов каждой командой в дополнение к приоритетам и временным меткам. Topology Aware Scheduling (TAS) отклоняет рабочие нагрузки, которые не помещаются на одном узле, и планирует размещение с помощью LeastFreeCapacity для поддержания плотной упаковки кластера.

Гарантированная емкость и емкость по требованию остались без изменений по сравнению с v0.10.

Дерево решений для разработчиков пополнилось новым сегментом:

С точки зрения пользователя, здесь нет ни новых меток, ни новых значений меток, поэтому интерфейс не изменился. Единственное, что изменилось — это то, как приоритеты влияют на маршрутизацию рабочей нагрузки: для многоузловых рабочих нагрузок и вытесняемых одноузловых нагрузок низкий приоритет теперь отделен от среднего и высокого.

«После внедрения»: Влияние развертывания Kueue

GPU больше не распределяются путем переговоров по внутренним каналам команд. Сегодня каждая рабочая нагрузка в AI21 — будь то одноузловой отладочный под, многоузловой обучающий запуск или развертывание инференса — автоматически направляется в нужную очередь. Kueue автоматически управляет процессом от начала до конца, не требуя ручных согласований: от постановки в очередь и приоритизации до очистки ресурсов и обеспечения справедливости использования GPU между командами.

Сегодня канал #gpu-resources уходит в архив. Руководители групп больше не выступают судьями в спорах за GPU, а исследователи тратят время на эксперименты, а не на запросы мощностей. Хорошая инфраструктура — это не просто часы работы GPU, это возможность для наших исследователей сосредоточиться на том, что они любят и умеют делать лучше всего: на науке. Вот чего мы добились в цифрах после перехода на Kueue:

Операционные достижения

Ручные вмешательства в неделю: 20 → 0
«Зомби»-задачи и задачи с частичным выделением ресурсов ликвидированы
Процент фрагментации: 15% → 8%
Время, впустую тратимое в командных чатах на согласование GPU в неделю: часы → ноль
Среднее время голодания приоритетной (hero) задачи: 72 часа → 12 часов

Выводы и заключение

Для всех, кто сталкивается с задачей совместного использования ресурсов GPU в кластере, вот основные уроки, которые мы усвоили при переходе от системы ручного согласования к автоматическому управлению с помощью Kueue:

Планируйте на уровне рабочей нагрузки (workload), а не на уровне пода. Встроенное вытеснение (preemption) в Kubernetes с радостью убьет один под из четырехподового процесса обучения и оставит остальные три впустую сжигать ресурсы GPU. Распределенное обучение требует пакетного (gang / all-or-nothing) приема задач.
Очередь — это не планировщик, пока она не понимает вашу топологию. Математика квот показывала, что свободно 8 GPU; из-за фрагментации (1+1+4+2 на четырех узлах) для задачи на 8 GPU не нашлось места. Планировщику с учетом топологии (Topology Aware Scheduling) удалось заставить Kueue отклонять задачи, которые не помещаются, и закреплять размещение на наиболее загруженных узлах в первую очередь, превратив очередь в нечто, что действительно занимается планированием.
Намеренно удерживайте загрузку парка близкой к 100%. Для зарезервированного парка GPU полная утилизация — это цель, а не повод для тревоги. Безопасность этого обеспечивается вытесняемой/оппортунистической линией, которая заимствует простаивающую емкость и возвращает ее при востребовании, поэтому «полный» не означает «заблокированный».

Как может подтвердить любая компания, обучающая модели, совместное использование GPU разными командами — сложная задача. Когда использование согласовывается асинхронно через внутренние каналы команд или жестко делится на отдельные кластеры для каждой команды, это оставляет желать лучшего с точки зрения эффективности времени и вычислений. Оба метода несут риски простоя GPU, несоответствия бизнес-приоритетам и прерывания задач.

Только после того, как мы перешли на использование Kueue для управления доступом, приоритетами и вытеснением в едином общем кластере нашей команды, емкость стала выделяться автоматически: справедливо между командами, с учетом бизнес-приоритетов и с возвратом ресурсов без ручного вмешательства.