Сегодня NVIDIA анонсировала открытый выпуск KAI Scheduler — решения для планирования GPU, основанного на Kubernetes, теперь доступного под лицензией Apache 2.0. Изначально разработанный в рамках платформы Run:ai, KAI Scheduler теперь доступен для сообщества и продолжает упаковываться в платформу NVIDIA Run:ai. Эта инициатива подчеркивает приверженность NVIDIA развитию как открытого кода, так и корпоративной AI инфраструктуры, содействуя активному и совместному сообществу, поощряя вклад, отзывы и инновации.
В этом посте представляем обзор технических деталей KAI Scheduler, подчеркиваем его ценность для IT и ML команд и объясняем цикл планирования и действия.
Преимущества KAI Scheduler
Управление загрузками AI на GPU и CPU представляет собой множество проблем, которые традиционные распределители ресурсов часто не могут решить. Планировщик был разработан для решения этих вопросов:
- Управление колеблющимся спросом на GPU
- Сокращение времени ожидания доступа к вычислениям
- Гарантии ресурсов или распределение GPU
- Бесшовное соединение инструментов и фреймворков AI
Управление колеблющимся спросом на GPU
Загрузки AI могут быстро изменяться. Например, может потребоваться только один GPU для интерактивной работы и затем внезапно понадобиться несколько GPU для распределенного обучения или нескольких экспериментов. Традиционные планировщики испытывают трудности с такой изменчивостью.
KAI Scheduler непрерывно пересчитывает значения справедливой доли и в реальном времени регулирует квоты и лимиты, автоматически соответствуя текущим требованиям загрузки. Этот динамический подход помогает обеспечить эффективное распределение GPU без постоянного ручного вмешательства администраторов.
Сокращение времени ожидания доступа к вычислениям
Для инженеров ML время имеет решающее значение. Планировщик сокращает время ожидания, комбинируя групповое планирование, совместное использование GPU и иерархическую систему очередей, позволяющую подавать партии заданий и уверенно отойти, зная, что задачи будут запущены, как только ресурсы станут доступными с учетом приоритетов и справедливости.
Чтобы оптимизировать использование ресурсов даже при изменении спроса, планировщик применяет две эффективные стратегии как для загрузок GPU, так и для CPU:
- Упаковка и консолидация: Максимизирует использование вычислений, борясь с фрагментацией ресурсов — упаковка мелких задач в частично использованные GPU и CPU — и устраняя фрагментацию узлов за счет перераспределения задач по узлам.
- Распределение: Равномерно распределяет нагрузки по узлам или GPU и CPU, минимизируя нагрузку на отдельный узел и максимизируя доступность ресурсов на нагрузку.
Гарантия ресурсов или распределение GPU
В совместных кластерах некоторые исследователи рано утром захватывают больше GPU, чем необходимо, чтобы гарантировать доступность на протяжении всего дня. Эта практика может привести к недоиспользованию ресурсов, даже когда у других команд еще есть неиспользованные квоты.
KAI Scheduler решает эту проблему, обеспечивая гарантии ресурсов. Он гарантирует, что команды практиков AI получают свои выделенные GPU, одновременно динамически перераспределяя неиспользуемые ресурсы для других загрузок. Этот подход предотвращает злоупотребление ресурсами и содействует общей эффективности кластера.
Бесшовное соединение инструментов и фреймворков AI
Соединение загрузок AI с различными фреймворками AI может быть сложной задачей. Ранее командам приходилось сталкиваться с множеством ручных конфигураций для объединения загрузок с инструментами, такими как Kubeflow, Ray, Argo и Training Operator. Эта сложность задерживала прототипирование.
KAI Scheduler решает эту задачу, имея встроенный подгруппер, который автоматически обнаруживает и соединяется с этими инструментами и фреймворками — снижая сложность конфигурации и ускоряя разработку.

Основные сущности планирования
Существует две основные сущности для KAI Scheduler: подгруппы и очереди.
Подгруппы
Подгруппы — это атомарные единицы для планирования и представляют собой одну или несколько взаимозависимых подов, которые должны выполняться как единое целое, также известное как групповое планирование. Эта концепция важна для распределенных фреймворков AI, таких как TensorFlow или PyTorch.
Вот ключевые атрибты подгруппы:
- Минимальное количество членов: Определяет минимальное количество подов, которые должны быть запланированы вместе. Если необходимые ресурсы недоступны для всех членов, подгруппа остается в ожидании.
- Ассоциация очереди: каждая подгруппа связана с конкретной очередью планирования, связывая ее с более широкой стратегией распределения ресурсов.
- Приоритетный класс: Определяет порядок планирования относительно других подгрупп, влияя на приоритезацию заданий.
Очереди
Очереди служат основными сущностями для обеспечения справедливости ресурсов. Каждая очередь имеет свои характеристики, которые направляют распределение ресурсов:
- Квота: Базовое распределение ресурсов, гарантированное очереди.
- Вес превышения квоты: Влияет на то, как дополнительные, избыточные ресурсы распределяются среди всех очередей помимо базовой квоты.
- Лимит: Определяет максимальные ресурсы, которые очередь может потреблять.
- Приоритет очереди: Определяет порядок планирования относительно других очередей, влияя на приоритезацию очередей.
Архитектура и цикл планирования
В своей основе планировщик работает, непрерывно фиксируя состояние кластера Kubernetes, вычисляя оптимальное распределение ресурсов и применяя целенаправленные действия по планированию.

Процесс организован в следующие фазы:
- Снимок кластера
- Деление ресурсов и вычисление справедливой доли
- Действия по планированию
Снимок кластера
Цикл планирования начинается с захвата полного снимка кластера Kubernetes. В этой фазе записывается текущее состояние узлов, подгрупп и очередей.
Объекты Kubernetes преобразуются в внутренние структуры данных, предотвращая несоответствия в середине цикла и обеспечивая, чтобы все решения по планированию основывались на стабильном представлении о кластере.
Деление ресурсов и вычисление справедливой доли
Имея точный снимок, алгоритм деления ресурсов рассчитывает справедливую долю для каждой очереди планирования. Результат этого алгоритма — числовое значение, представляющее оптимальное количество ресурсов, которое очередь должна получить для максимизации справедливости среди всех очередей в кластере.
Алгоритм деления имеет следующие фазы:
- Заслуженная квота: Каждой очереди планирования изначально выделяются ресурсы, равные ее базовой квоте. Это гарантирует, что каждый департамент, проект или команда получают минимальные ресурсы, на которые они вправе претендовать.
- Избыточная квота: Оставшиеся ресурсы распределяются между очередями, которые не обеспечены, пропорционально их весу превышения квоты. Этот итеративный процесс тонко настраивает справедливую долю для каждой очереди, динамически адаптируя к требованиям загрузки.
Действия по планированию
После вычисления справедливых долей планировщик применяет ряд целенаправленных действий для согласования текущих распределений с вычисленным оптимальным состоянием:
- Распределение: Ожидающие задания оцениваются на основе их коэффициента распределения к справедливой доле. Задания, которые подходят под доступные ресурсы, немедленно связываются, в то время как те, которые требуют ресурсов, которые в данный момент освобождаются, ставятся в очередь для последовательного распределения.
- Консолидация: Для учебных загрузок планировщик создает упорядоченную очередь оставшихся ожидающих учебных заданий. Он итеративно проходит их и пытается их распределить, перемещая текущие назначенные поды на другой узел. Этот процесс минимизирует фрагментацию ресурсов, освобождая смежные блоки для ожидающих заданий.
- Рекламация: Для обеспечения справедливости планировщик определяет очереди, потребляющие больше ресурсов, чем их справедливая доля. Затем он изгоняет выбранные задания на основе четко определенных стратегий — обеспечивая, чтобы недостаточно удовлетворенные очереди получили необходимые ресурсы.
- Прерывание: Dentro de la misma cola, los trabajos de menor prioridad pueden ser interrumpidos en favor de trabajos pendientes de alta prioridad, asegurando que las cargas de trabajo críticas no se vean afectadas debido a la competencia por recursos.
Пример сценария
Представьте кластер из трех узлов, каждый из которых оснащен восемью GPU, всего 24 GPU. Запущены два проекта: project-a и project-b. Они связаны с очередью- а (очередь среднего приоритета) и очередью-b (очередь высокого приоритета) соответственно. Предположим, что все очереди находятся в пределах своей справедливой доли. Рисунок 3 демонстрирует текущее состояние.

- Узел 1: Учебное задание 1 использует четыре GPU из очереди высокого приоритета, а учебное задание 2 использует два GPU.
- Узел 2: Учебное задание 3 использует шесть GPU.
- Узел 3: Интерактивное задание, которое может быть связано с исследованием данных или отладкой скриптов, использует 5 GPU.
Теперь есть два ожидающих задания, поданных в очередь (Рисунок 4):
- Учебное задание A требует четырех смежных GPU на одном узле.
- Учебное задание B требует трех смежных GPU, поданных в очередь высокого приоритета.

Фаза распределения
Планировщик упорядочивает ожидающие задания по приоритету. В данном случае сначала обрабатывается Учебное задание B, так как его очередь имеет более высокий рейтинг по срочности. Конечно, процесс упорядочивания заданий более сложен, чем простая сортировка по приоритету. Планировщик использует передовые стратегии для обеспечения справедливости и эффективности.

- Для Учебного задания B Узел 3 подходит, так как у него есть три смежных свободных GPU. Планировщик связывает Учебное задание B с Узлом 3, и Узел 3 полностью заполняется.
- Затем планировщик пытается распределить Учебное задание A. Однако ни один из узлов не предоставляет четырех смежных свободных GPU.
Фаза консолидации
Поскольку Учебное задание A не удалось запланировать в ходе действия распределения, планировщик переходит в фазу консолидации. Он проверяет узлы, чтобы выяснить, может ли он перераспределить работающие поды для формирования смежного блока для Учебного задания A.

- На Узле 1, помимо Учебного задания 1 (которое должно оставаться из-за высокого приоритета), есть Учебное задание 2, занимающее два GPU. Планировщик перемещает Учебное задание 2 с Узла 1 на Узел 2.
- После этого перемещения консолидации количество свободных GPU на Узле 1 увеличивается с двух до четырех смежных GPU.
- С учетом нового расположения Учебное задание A теперь может быть распределено на Узел 1, удовлетворяя его требованию в четыре смежных GPU.
В этом сценарии действие консолидации было достаточным для того, чтобы ни рекламация, ни прерывание не требовались. Тем не менее, если бы было еще одно ожидающее задание из очереди ниже своей справедливой доли без доступного пути консолидации, тогда действия рекламации или прерывания планировщика вступили бы в силу — рекламация ресурсов от переполненных очередей или от заданий низкого приоритета в той же очереди для восстановления баланса.