AI21, Kueue로 GPU 스케줄링 자동화… 수동 개입 제로화 및 기아 시간 83% 단축

AI21은 약 10,000개 GPU가 포함된 단일 공유 GKE 클러스터를 관리하기 위해 Kubernetes용 배치 스케줄러인 Kueue를 도입하여 수동 GPU 자원 협상을 대체했습니다. 이전에는 내부 메신저 채널을 통해 접근 권한을 조율했으나, 이는 제로섬 게임, 유휴 용량 발생, 상당한 엔지니어링 오버헤드를 초래했습니다. Kueue 도입으로 AI21은 대기열 처리, 우선순위 지정, 정리 작업을 자동화하여 자원 할당을 비즈니스 우선순위에 맞추고 팀 간 공정한 공유를 보장했습니다.

Kubernetes 기본 기능의 한계 극복

공정성과 기술적 무결성 측면에서 PriorityClass와 ResourceQuota 같은 Kubernetes 기본 요소는 AI21의 요구사항을 충족하지 못했습니다. PriorityClass는 포드 단위에서 작동하여 갱 스케줄링 요구사항을 무시했고, ResourceQuota는 대기열 처리 대신 어드미션을 거부하여 유휴 용량 대여를 방지했습니다. 또한 Indexed Jobs는 올-오어-낫싱(all-or-nothing) 시맨틱이 부족하여 분산 학습 실행 시 일부 포드만 스케줄링될 경우 GPU 낭비 위험이 있었습니다. AI21은 단순성과 원활한 통합을 이유로 Volcano와 Apache YuniKorn 같은 대안 대신 Kueue를 선택했습니다.

구현 및 진화

도입은 지출, 선점 가능성, 우선순위를 위한 개발자 레이블을 도입한 v0.10으로 시작했습니다. 그러나 이 버전은 Admission Fair Sharing (AFS) 부재로 인한 공정성 격차와 클러스터 토폴로지를 인식하지 못한 Kueue가 멀티 노드 워크로드를 효율적으로 배치하지 못해 발생한 단편화 문제를 겪었습니다. AI21은 Kueue 팀과 협력하여 AFS를 도입하고 **Topology Aware Scheduling (TAS)**을 활성화하여 이러한 문제를 해결했습니다.

v0.15에서는 우선순위 인식 선점 가능 용량과 멀티 노드 베스트 에포트 레인을 포함하도록 스케줄링 로직을 재설계했습니다. 현재 시스템은 대기열 순서를 위해 과거 칩-시간 데이터를 사용하며, 단편화를 최소화하기 위해 LeastFreeCapacity 알고리즘을 적용합니다. 이러한 내부 변경에도 불구하고 개발자 인터페이스는 변경되지 않았으며, 라우팅 결정을 위해 동일한 세 가지 레이블이 유지되었습니다.

운영 성과

자동화된 스케줄링으로의 전환은 다음과 같은 상당한 운영 개선을 가져왔습니다:

• 수동 개입 주당 20회에서 0회로 감소.
• 단편화 15%에서 8%로 감소.
• 히어로 잡 기아 상태 시간 72시간에서 12시간으로 단축.
• 좀비 및 부분 할당 잡 완전히 제거.
• #gpu-resources 채널이 아카이브되어 GPU 분쟁에 대한 수동 중재 필요성이 종료됨.