AWS описала новую возможность планирования в AWS Batch под названием quota management в заметке на AWS Machine Learning Blog. Она опирается на поддержку задач SageMaker Training в AWS Batch, которая уже занимается очередями, планированием и повторными запусками для обучения ML-моделей.
Что она делает
Администратор делит мощность GPU на квоты - по одной на команду, проект или тип нагрузки. У каждой квоты есть лимит и стратегия на случай простоя:
- RESERVE держит мощность до прихода собственных задач квоты.
- LEND предлагает простаивающую мощность другим квотам, у которых есть ожидающая работа.
- LEND_AND_BORROW одалживает в простое и занимает, когда у квоты есть задачи, способные использовать дополнительную мощность.
Каждая квота работает как виртуальная очередь внутри очереди задач. Планировщик отдаёт предпочтение квотам, которые не дошли до своих лимитов. Если квота одолжила мощность, а её собственные задачи пришли, AWS Batch выполняет вытеснение между квотами (cross-share preemption) и забирает ресурсы обратно, и администратору не нужно искать задачи-заёмщики.
Есть и второй режим - вытеснение внутри квоты (in-share preemption). Когда он включён, задачи с низким приоритетом, которые запланированы, запускаются или уже выполняются, можно прервать, чтобы высокоприоритетные задачи из начала очереди стартовали раньше. Новый API UpdateServiceJob позволяет исследователю поднять приоритет уже отправленной задачи.
Настройка
Нужна политика планирования с политикой квот, сервисная среда, которая задаёт ограничения мощности по конкретным инстансам, и очередь задач SageMaker Training, связанная с обоими элементами. К очереди затем подключаются квоты. Задачи отправляются через модуль aws_batch в SageMaker Python SDK v3, а примеры для CLI и консоли авторы отсылают искать в руководстве пользователя.
Пример из заметки
AWS создаёт три квоты, и каждая ограничена одним инстансом ml.g5.xlarge:
- Квота для фоновой работы: она и одалживает, и занимает, вытеснение внутри квоты включено. Подходит для задач с контрольными точками, которые переносят прерывание.
- Квота среднего приоритета: она только одалживает, вытеснение внутри квоты выключено. Для команд, которые лучше подождут, чем будут раз за разом терять запуск.
- Квота только с резервированием: её задачи не вытесняются, пока лимит остаётся прежним.
В примере фоновые задачи занимают свободный инстанс квоты среднего приоритета. Когда в эту квоту приходит задача, мощность забирают у фоновой. Затем повышение приоритета ожидающей задачи вытесняет одну из выполняющихся в той же квоте.
Также в заметке сказано, что возможность работает с планами обучения SageMaker AI (training plans), поэтому зарезервированную мощность можно делить между командами.