Весной этого года NVIDIA открыла исходный код KAI Scheduler — Kubernetes-нативного планировщика GPU-нагрузок, который раньше входил в состав платформы Run:ai и теперь распространяется под лицензией Apache 2.0. Интерес к этому проекту закономерен: планировщик давно работает в продакшене и решает ряд проблем, с которыми сталкивается любая команда, пытающаяся эффективно распределять GPU-ресурсы в кластере.
Мы в Orion soft изучили технические детали KAI Scheduler, чтобы понять, как он устроен изнутри, какие задачи закрывает и какие идеи могут быть полезны инженерам, работающим с Kubernetes, ML-нагрузками и распределёнными GPU-оркестраторами. Ниже — разбор архитектуры, базовых сущностей и цикла планирования.Преимущества KAI Scheduler
Управление AI-нагрузками на GPU и CPU сталкивается с рядом задач, которые традиционные планировщики ресурсов не всегда способны решать. KAI Scheduler был разработан специально для того, чтобы закрыть эти проблемы:
управление колеблющимся спросом на GPU;
сокращение времени ожидания доступа к вычислительным ресурсам;
гарантии ресурсов и предсказуемое выделение GPU;
бесшовная интеграция с AI-инструментами и фреймворками.
Управление колеблющимся спросом на GPU
AI-нагрузки могут меняться очень быстро. Например, вам может понадобиться всего один GPU для интерактивной работы (скажем, исследования данных), а затем — сразу н��сколько GPU для распределённого обучения или параллельных экспериментов. Традиционные планировщики плохо справляются с такой вариативностью.
KAI Scheduler пересчитывает fair-share и корректирует внутренние лимиты очередей, автоматически подстраиваясь под текущие требования нагрузки. Такой динамический подход позволяет эффективно распределять GPU без постоянного ручного вмешательства администраторов.
Сокращение времени ожидания доступа к вычислительным ресурсам
Для ML-инженеров время критично. Планировщик уменьшает время ожидания, сочетая gang scheduling, оптимальное распределение GPU-ресурсов и иерархическую систему очередей. Это позволяет отправлять большие наборы задач и не следить за ними вручную — задания стартуют сразу, как только появляются ресурсы, с учётом приоритетов и принципов справедливости.
Чтобы дополнительно оптимизировать использование ресурсов при колеблющемся спросе, планировщик применяет два подхода как для GPU-, так и для CPU-нагрузок:
Bin-packing и консолидация: повышают загрузку вычислительных ресурсов, уменьшая фрагментацию ресурсов — размещают задачи на частично занятых GPU и узлах, где есть свободные блоки ресурсов, а также устраняют «раздробленность» узлов за счёт перераспределения задач между ними.
Spreading: равномерно распределяет задачи по узлам, GPU и CPU, снижая нагрузку на каждый конкретный узел и повышая доступность ресурсов для каждого workload’а.
Гарантии ресурсов и выделение GPU
В общих кластерах исследователи иногда заранее бронируют больше GPU, чем им реально нужно, чтобы гарантировать доступ в течение дня. Из-за этого часть ресурсов простаивает, даже если у других команд остаются неиспользованные квоты.
KAI Scheduler решает эту проблему с помощью механизма гарантий ресурсов. Он гарантирует очередям их базовую квоту, а свободные ресурсы распределяет динамически в соответствии с fair-share и может передавать простаивающие ресурсы очередям, которым они нужнее. Такой подход предотвращает «захват» ресурсов и повышает общую эффективность кластера.
Бесшовная интеграция с AI-инструментами и фреймворками
Связать AI-нагрузки с разными фреймворками — задача непростая. Обычно командам приходится вручную настраивать связку рабочих нагрузок с такими инструментами, как Kubeflow, Ray, Argo или Training Operator. Эта сложность замедляет прототипирование.
KAI Scheduler упрощает процесс за счёт встроенного podgrouper — компонента, который автоматически определяет и подключает эти инструменты и фреймворки. Это снижает объём ручной конфигурации и ускоряет разработку.
Базовые сущности планировщика
В KAI Scheduler есть две основные сущности: группа подов и очередь.
Группа подов (PodGroup — отдельный CRD) — это атомарная единица планирования. Она представляет собой один или несколько взаимосвязанных подов, которые должны выполняться как единое целое — то есть с использованием gang scheduling. Такой подход критически важен для распределённых фреймворков обучения, таких как TensorFlow или PyTorch.
Ключевые атрибуты подгруппы:
Минимальное число участников (minimum members): минимальное количество подов, которые должны быть запланированы вместе. Если ресурсов не хватает, чтобы запланировать все необходимые поды, группа остаётся в состоянии ожидания.
Привязка к очереди (queue association): каждая группа подов связана с конкретной очередью, что встраивает её в общую стратегию распределения ресурсов.
Класс приоритета (priority class): определяет порядок планирования относительно других групп подов и влияет на приоритизацию задач.
Очереди — основа механизма обеспечения справедливого распределения ресурсов. Каждая очередь имеет свойства, которые определяют, как она получает ресурсы:
Квота (quota): базовый объём ресурсов, гарантированный очереди.
Вес сверх квоты (over-quota weight): задаёт, как распределяются дополнительные ресурсы выше базовой квоты между всеми очередями.
Лимит (limit): максимальный объём ресурсов, который может потреблять очередь.
Приоритет очереди (queue priority): определяет порядок планирования очередей относительно друг друга.
Архитектура и цикл планирования
В основе работы планировщика — постоянное получение состояния Kubernetes-кластера, вычисление оптимального распределения ресурсов и применение целевых действий планирования.
Процесс организован в следующие этапы:
Снэпшот состояния Kubernetes-к��астера
Деление ресурсов и вычисление fair share
Действия планировщика
Снэпшот кластера
Цикл планирования начинается с формирования снэпшота состояния Kubernetes-кластера. На этом этапе фиксируется текущее состояние узлов, групп подов и очередей.
Объекты Kubernetes преобразуются во внутренние структуры данных, что предотвращает несогласованность в середине цикла и гарантирует, что все решения принимаются на основе стабильного состояния кластера.
Деление ресурсов и вычисление fair share
После получения снэпшота алгоритм деления ресурсов вычисляет fair share для каждой очереди. Результатом является числовое значение, отражающее оптимальный объём ресурсов, который должна получить очередь для обеспечения максимальной справедливости между всеми очередями кластера.
Алгоритм деления состоит из таких этапов:
Deserved-quota: каждой очереди изначально выделяются ресурсы, равные её базовой квоте. Это гарантирует департаментам, проектам или командам минимальный объём ресурсов, на который они имеют право.
Over-quota: оставшиеся ресурсы распределяются между очередями, которым их не хватает, пропорционально их весу over-quota. Этот итеративный процесс уточняет fair share каждой очереди и динамически подстраивается под требования ворклоадов.
Действия планировщика
Когда fair share рассчитаны, планировщик выполняет целевые действия, чтобы привести текущее распределение ресурсов к оптимальному состоянию:
Аллокация (allocate): задачи в ожидании оцениваются на основе отношения allocated-to-fair-share. Задачи, подходящие под доступные ресурсы, сразу связываются с узлами; задачи, которым нужны ресурсы, находящиеся в процессе освобождения, cтавятся в очередь для последующего выделения (pipelined allocation).
Консолидация (consolidation): для тренировочных ворклоадов планировщик формирует упорядоченную очередь оставшихся задач в ожидании. Он проходит по ним и пытается выделить ресурсы, перезапуская поды на других узлах. Это уменьшает фрагментацию ресурсов и освобождает непрерывные блоки для ожидающих задач.
Возврат ресурсов (reclaim): чтобы соблюсти справедливое распределение ресурсов, планировщик выявляет очереди, которые потребляют больше своей fair share. Затем он удаляет выбранные поды согласно своим стратегиям, чтобы недополучающие очереди получили необходимые ресурсы. Это действие выполняет сам планировщик — а не автоматическое удаление подов kubelet при нехватке ресурсов.
Вытеснение задач с более низким приоритетом (preempt): внутри одной очереди задачи с более низким приоритетом могут быть вытеснены задачами с более высоким приоритетом, что предотвращает «голодание» критичных ворклоадов из-за конкуренции за ресурсы.
Пример сценария
Представим кластер из трёх узлов, каждый из которых оснащён восемью GPU — всего 24 GPU. В кластере выполняются два проекта: project-a и project-b, привязанные к очередям queue-a (очередь среднего приоритета) и queue-b (очередь высокого приоритета) соответственно. Предположим, что обе очереди находятся в пределах своей fair share. Текущее состояние показано на рисунке ниже.
Узел 1: Training Job 1 использует четыре GPU из очереди высокого приоритета, Training Job 2 — два GPU.
Узел 2: Training Job 3 использует шесть GPU.
Узел 3: интерактивная задача (например, Jupyter/отладочная сессия) использует пять GPU.
Теперь в очередь поступают две задачи, находящиеся в ожидании, схема показана на рисунке ниже:
Training Job A требует четыре смежных GPU на одном узле.
Training Job B требует три смежных GPU и отправлена в очередь высокого приоритета.
Аллокация
Планировщик упорядочивает ожидающие задачи по приоритету. В нашем случае Training Job B обрабатывается первой, потому что её очередь имеет более высокий уровень срочности. При этом сам процесс упорядочивания куда сложнее простой сортировки по приоритетам: планировщик использует дополнительные стратегии, чтобы одновременно сохранять эффективность и справедливое распределение ресурсов.
Для Training Job B подходит узел 3, так как на нём есть три смежных свободных GPU. Планировщик привязывает Training Job B к узлу 3, и этот узел становится полностью занятым.
Далее планировщик пытается выделить ресурсы для Training Job A. Но ни один из узлов не предоставляет четыре смежных свободных GPU.
Консолидация
Поскольку Training Job A не удалось запланировать на этапе выделения ресурсов, планировщик переходит к фазе консолидации. На этом этапе он проверяет узлы и пытается перераспределить уже работающие поды так, чтобы сформировать непрерывный блок GPU под требования Training Job A.
На узле 1, помимо Training Job 1 (которую нельзя перемещать из-за высокого приоритета), выполняется Training Job 2, занимающая два GPU. Планировщик переносит Training Job 2 с узла 1 на узел 2.
После этого перемещения количество свободных GPU на узел 1 увеличивается: вместо двух свободных GPU теперь доступно четыре смежных GPU.
С такой конфигурацией Training Job A можно разместить на узле 1 — требование о четырёх смежных GPU выполняется.
В этом сценарии этап консолидации оказался достаточным, и ни возврат ресурсов, ни вытеснение задач не потребовались. Однако если бы в кластере появилась ещё одна задача в ожидании от очереди, которая получает меньше своей справедливой доли ресурсов, и путь консолидации был бы невозможен, планировщик задействовал бы эти механизмы — возвращая ресурсы от очередей, использующих больше положенного, или вытесняя задачи с более низким приоритетом внутри той же очереди, чтобы восстановить баланс.
Обновление статуса
Статусы очередей обновляются, и весь цикл начинается заново, чтобы планировщик мог своевременно реагировать на новые рабочие нагрузки.
Заключение
KAI Scheduler — зрелый и активно используемый компонент: он работает в продакшене на базе NVIDIA Run:ai и прошёл проверку реальными рабочими нагрузками. Открытие его исходного кода делает технологию доступной для изучения, экспериментов и доработок под собственные задачи.
Если вы хотите глубже разобраться в механизмах распределения GPU и встраивать такие подходы в свои платформы, имеет смысл посмотреть код, попробовать планировщик в тестовом окружении и поделиться наблюдениями. Репозиторий открыт для вопросов, предложений и пулл-реквестов — опыт практического использования здесь особенно ценен.
Если тема окажется интересной, можем продолжить разбор и подробнее поговорить о практических сценариях использования KAI Scheduler — в том числе о том, как похожие подходы применяются в наших решениях. Будем рады вопросам и комментариям.
