Így osztja meg az AI21 a közel 10 ezer GPU-t a csapatai között
Az AI21 egy mintegy 10 ezer GPU-t kezelő, több csapat által használt GKE-fürtön váltott kézi erőforrás-egyeztetésről automatizált feladatütemezésre. A rendszer központi eleme a nyílt forráskódú, Kubernetesre épülő Kueue, amely a munkafolyamatokat automatikusan sorolja, engedélyezi és szükség esetén megszakítja.
- Az AI21 mintegy 10 ezer GPU-t oszt meg több csapat között egy közös GKE-fürtön.
- A kézi, belső üzenetküldő csatornán zajló egyeztetést Kueue-alapú ütemezés váltotta fel.
- A Kueue feladat-szinten kezeli a prioritást, a sorban állást és a megszakítást.
- Az AI21 szerint a nagy prioritású munkafolyamatok indulásáig eltelt idő 83 százalékkal csökkent.
- A vállalat éles tapasztalatai hozzájárultak az Admission Fair Sharing fejlesztéséhez.
A GPU-k elosztása kézi egyeztetéssel akadt el
Az AI21 szerint a GPU-k megosztása különösen nehéz akkor, amikor kutatók, MLOps-csapatok, mérnökök és összehangolással foglalkozó munkatársak ugyanazért a korlátozott kapacitásért versenyeznek. A vállalatnál kezdetben egy belső üzenetküldő csatornán, a #gpu-resources felületen próbálták rendezni a kérelmeket.
Ha egy szükséges GPU-típus teljesen foglalt volt, a kérvényező abban reménykedhetett, hogy valaki felszabadít kapacitást. Amíg volt tartalék, ez működött, de a kihasználtság közel 100 százalékos szintje mellett minden új igény feladatok kiszorításáról szóló egyeztetéssé vált.
Az AI21 felsorolása szerint a kézi modell mérnöki időt vitt el a válogatással és az egyeztetéssel, igazságtalan sorban állást eredményezett, miközben egyes kapacitásrészek kihasználatlanok maradtak, más munkafolyamatok pedig erőforrásra vártak. A csapatok emellett alkalmi, úgynevezett „varázsparancsokra” támaszkodtak a szükséges kapacitás felszabadításához.
Négy alapelvre építették az új rendszert
Az AI21 négy szempontot határozott meg a GPU-kezeléshez. A méltányosság azt jelenti, hogy az ütemezési döntések a csapatok korábbi erőforrás-használatát is figyelembe vevő szabályok alapján születnek, nem pusztán érkezési sorrendben. Az elsőbbségi sorrend a munkafolyamatok üzleti fontosságát is kezeli.
A harmadik elv a hatékonyság: a fejlesztőknek ne kelljen kézzel törölniük a tétlen vagy kevésbé kritikus feladatokat. A negyedik a technikai integritás, amely szerint a több podból álló, úgynevezett gang feladatok csak akkor induljanak el, ha minden szükséges erőforrás rendelkezésre áll.
A vizsgált Google Cloud-fürt egyetlen közös készletként működik. Az AI21 ezt azért választotta, hogy magas maradjon a kihasználtság, miközben ugyanaz a döntés nagyobb nyomást helyez a fürtön belüli ütemezésre. A kapacitás lefoglalt, spot és a Dynamic Workload Scheduler flex-start módjához tartozó gépekből áll.
A fürt debug podokat, egy- és többgépes tréningfeladatokat, valamint megerősítéses tanuláshoz használt, Kubernetes-telepítésként kezelt következtetési modelleket futtat. Ezek között vannak megszakíthatatlan és megszakítható munkafolyamatok is.
A Kueue feladat-, nem pod-szinten dönt
Az AI21 először a Kubernetes beépített lehetőségeit vizsgálta. A PriorityClass képes kritikus podokat előnyben részesíteni, de podszinten működik, és nem ismeri a teljes feladat méretét. A ResourceQuota korlátokat szab, viszont nem biztosít kapacitásmegosztást a csapatok között, a Job és az Indexed Job pedig nem garantálja, hogy egy több podból álló feladat csak teljes egészében induljon el.
A vállalat az Apache YuniKorn, a Volcano és a Kueue megoldásait értékelte. A beszámoló szerint a Volcanónak akkor több funkciója volt, de a környezetbe való integrációja nem bizonyult zökkenőmentesnek, a YuniKorn pedig nem fedte le az összes felhasználási esetet. Egyszerűség és integráció szempontjából végül a Kueue mellett döntöttek.
A Kueue API-kból és vezérlőből álló rendszer, amely feladat-szinten határozza meg, mikor indulhat el egy munka, illetve mikor kell leállítani. A Workload Priority Classes a sorban állást és a megszakítási sorrendet kezeli. A Preemption Fair Sharing a fürthasználat kiegyensúlyozására szolgál, az Admission Fair Sharing pedig előnyben részesíti azokat a csapatokat, amelyek korábban kevesebb kapacitást használtak.
A Google-lel közösen fejlesztették tovább
Az AI21 a bevezetés során a Kueue mögött álló Google-csapattal is együttműködött. A vállalat éles környezetben szerzett tapasztalatai közvetlenül visszakerültek a Kueue fejlesztési ütemtervébe, többek között az Admission Fair Sharing funkcióba, amely az AI21 közlése szerint már minden felhasználó számára elérhető.
A Google Cloud korábbi beszámolója szerint a megközelítés a nagy prioritású munkafolyamatok indulásáig eltelt időt 83 százalékkal csökkentette. Az AI21 saját technikai összefoglalója emellett a kézi beavatkozási arány és a kritikus feladatok erőforrás-hiányos várakozási idejének jelentős visszaeséséről számol be.
A felhasználók számára ez azt jelenti, hogy a GPU-igények kezelése kevésbé függ a csapatok közötti kézi egyeztetéstől. A Kueue a prioritásokat, a méltányos megosztást és a megszakítható munkák kezelését szabályok szerint végzi, miközben a több gépet igénylő feladatok csak akkor indulnak, ha a teljes erőforrásigényük biztosítható.
AI21 Labs: From manual negotiation to automated scheduling: How AI21 manages its GPU fleet with Kueue | AI21
