‏AI21 چطور ۱۰ هزار GPU رو با Kueue بین تیم‌ها تقسیم کرد

تیم AI21 تو این مقاله تعریف می‌کنه که چطور یه کلاستر مشترک GKE با حدود ۱۰ هزار GPU رو بین چند تیم تحقیق و مهندسی تقسیم کرده. اولش همه‌چی تو یه چنل داخلی به اسم #gpu-resources حل می‌شد: هر کی GPU می‌خواست، می‌پرسید و امیدوار بود یکی جا خالی کنه. تا وقتی کلاستر جای خالی داشت این روش جواب می‌داد، ولی وقتی مصرف همیشه نزدیک ۱۰۰ درصد موند، هر درخواست تبدیل شد به یه دعوا سر اینکه کار کی باید بیرون بره.

نویسنده‌ها چهار تا اصل برای سیستم ایده‌آل تعریف کردن: عدالت (بر اساس مصرف قبلی، نه اولویت با کسی که بلندتر داد می‌زنه)، ترتیب (بر اساس اولویت کاری، نه فقط زمان ثبت)، بهره‌وری (لازم نباشه کسی دستی کارهاشو پاک کنه) و یکپارچگی فنی، یعنی job چندپادی (gang) تا همه منابعش آماده نشده شروع نشه. بین YuniKorn، Volcano و Kueue، آخر سر Kueue به خاطر سادگی و ادغام راحت‌تر انتخاب شد.

Kueue یه job queueing system بومی Kubernetes هست، یعنی تصمیم می‌گیره یه job کِی اجازه شروع بگیره (پادهاش ساخته بشن) و کِی متوقف بشه. نسخه اول پیاده‌سازی (v0.10) سه لایه ظرفیت داشت: ظرفیت تضمینی برای هر تیم، یه pool قابل preempt (یعنی کارهایی که می‌شه وسط اجرا قطعشون کرد) که GPUهای بیکار رو قرض می‌گیره، و ظرفیت on-demand برای وقتی همه‌چی پره. توسعه‌دهنده فقط چند تا label مثل هزینه، قابلیت preempt و اولویت می‌زنه و صف درست خودش انتخاب می‌شه.

دو تا مشکل پیش اومد. اول اینکه عدالت بین تیم‌ها برای کارهای چندنودی بدون preemption ممکن نبود. AI21 اینو به تیم Kueue گفت و اونا قابلیت Admission Fair Sharing (AFS) رو ساختن، یعنی تیمی که قبلاً کمتر مصرف کرده جلوتر تو صف قرار می‌گیره. مشکل دوم fragmentation بود: مثلاً ۸ GPU آزاد بود ولی پخش شده به شکل ۱+۱+۴+۲ روی چهار نود، و jobی که ۸ GPU روی یه نود می‌خواست جایی برای اجرا نداشت.

راه‌حل مشکل دوم Topology Aware Scheduling (TAS) بود. با TAS، Kueue می‌دونه هر ماشین ۸ تا GPU داره، کاری که جا نمی‌شه رو قبول نمی‌کنه و با الگوریتم LeastFreeCapacity کارها رو اول روی شلوغ‌ترین نودها می‌چینه تا GPUها پخش‌وپلا نشن. نسخه v0.15 با AFS و TAS ساخته شد و یه مسیر best-effort با اولویت low هم برای کارهای چندنودی اضافه شد، بدون اینکه labelهای کاربر عوض بشه.

نتیجه به گفته AI21: چنل #gpu-resources آرشیو شده و پژوهشگرها به‌جای چونه زدن سر GPU روی آزمایش‌هاشون تمرکز می‌کنن. درس اصلی‌شون اینه که زمان‌بندی باید در سطح workload باشه نه پاد، چون preemption خود Kubernetes ممکنه یه پاد از یه training چهارپادی رو بکشه و سه تای دیگه بیخودی GPU بسوزونن.

نکات کلیدی:

• کلاستر مشترک GKE با حدود ۱۰ هزار GPU، به‌جای یه کلاستر جدا برای هر تیم
• دخالت دستی هفتگی از ۲۰ بار به صفر رسید
• fragmentation از ۱۵ درصد به ۸ درصد کم شد
• میانگین انتظار کارهای خیلی مهم از ۷۲ ساعت به ۱۲ ساعت رسید
• قابلیت Admission Fair Sharing در Kueue از نیاز واقعی AI21 ساخته شد و الان برای همه در دسترسه
• TAS با الگوریتم LeastFreeCapacity جلوی پذیرش کارهایی رو می‌گیره که روی یه نود جا نمی‌شن