A Krea saját módszerrel tereli a GPU-s feladatokat Kubernetesben

A Kubernetes alapértelmezett ütemezőjét alakították a céljaikhoz

A Krea augusztus 7-én közzétett mérnöki bejegyzésében mutatta be, hogyan ütemezi GPU-s munkaterheléseit. A vállalat korábbi írása azt ismertette, miként tudnak bárhol következtetési feladatokat futtatni a Virtual Kubelet projekt segítségével. A mostani bejegyzés a rendszer másik felére, a klaszter erőforrásainak használatára koncentrál.

A Krea szerint több megközelítés módosítja a Kubernetes ütemezőjét, például konfigurációval vagy bővítményekkel, mások pedig teljesen lecserélik olyan megoldásokra, mint a Volcano, az Apache YuniKorn vagy a KAI-Scheduler. A vállalat egyik út mellett sem döntött. A bejegyzés szerint néhány, az alapértelmezett ütemezőre épített eszközzel el tudták érni a kívánt működést.

A cél az úgynevezett tartalékkapacitás kezelése. Amennyire lehet, a munkákat a klaszterben tartják, és csak akkor engedik ki őket a külső kapacitásba, amikor a klaszter megtelt, például kutatási feladatok foglalják le, vagy a következtetési igény kimeríti az erőforrásokat.

Dinamikus taint jelzi, mikor van szabad GPU

A Krea Virtual Kubelet-csomópontja rendkívül nagy mennyiségű erőforrást hirdet. Emiatt a Kubernetes NodeResourcesFit bővítménye aránytalanul magas pontszámot adhat ennek a csomópontnak, így a podok akkor is oda kerülhetnek, amikor a klaszterben még lenne szabad kapacitás. A vállalat kipróbálta a lágy csomópont-affinitást és a MostAllocated stratégiát is, de egyik megoldás sem felelt meg minden igénynek.

A rendszer ezért taintokat használ. A Virtual Kubelet-csomópont alapból rendelkezik egy NoSchedule tainttel, amelyet csak a már migrált munkaterhelések tolerálnak. A Krea szerint ezeknek a feladatoknak előzetesen több módosításon is át kell esniük, például a modelleket le kell tölteniük a külső szolgáltató oldalán, vagy kezelniük kell bizonyos működési eltéréseket.

Alacsony GPU-kihasználtság esetén egy második, senki által nem tolerált taint is megjelenik a csomóponton. Ezt a rendszeren belül futó háttérfeladat kezeli, amely a Prometheustól kérdezi le az elérhető GPU-k számát. A Krea a klaszter állapotát egy PromQL-lekérdezéssel írja le. Ebben például kizárják a legalacsonyabb prioritású adatfeldolgozási feladatokat, hogy a tanítási vagy következtetési munkák elsőbbséget élvezhessenek.

A lekérdezés eredményét időben átlagolják, majd csak akkor adják hozzá vagy távolítják el a taintet, ha az érték egy meghatározott ideig a küszöbérték felett vagy alatt marad. Ezzel igyekeznek elkerülni a folyamatos áthelyezést.

A Descheduler fokozatosan helyezi vissza a munkákat

Amikor újra felszabadulnak GPU-k a klaszterben, a feladatokat vissza kell migrálni. A Krea erre a Descheduler RemovePodsViolatingNodeTaints stratégiáját használja. Ez eltávolítja azokat a podokat, amelyek már megsértik a csomóponton érvényes NoSchedule taint szabályát.

A vállalat nem a NoExecute taintre épít, mert az azonnal kilakoltatná a csomóponton futó összes podot. Ez egy nagyobb következtetési terhelésnél egyszerre szakítaná meg a munkák jelentős részét. A Descheduler ezzel szemben beállítható minimális pod-életkorral, futási gyakorisággal és ciklusonkénti maximális kilakoltatási számmal.

A Krea szerint minden eltávolítás normál kilakoltatásnak számít, ezért érvényesülnek a Pod Disruption Budgetek. A munkaterhelések replikaszáma nem csökken nullára, és az előre beállított alsó határ alá sem kerül.

A Kueue kezeli a sorokat és a prioritásokat

A feladatok felvételéhez és sorrendjéhez a Krea a Kueue-t használja. Négy sort alakítottak ki. A training sor a klaszter GPU-inak szinte teljes készletét kezeli, és kölcsönözhet a többi sornak, illetve kölcsönözhet tőlük. Az inference egy kisebb GPU-készletet tart fenn azoknak a következtetési folyamatoknak, amelyek nem helyezhetők át a Virtual Kubelet-csomópontra.

A low-priority sorhoz nem tartozik saját GPU, csak más soroktól kölcsönözhet, és szükség esetén visszavehető. Az inference-vk ezzel szemben nagy GPU-készletet kezel, amelyet nem adhat kölcsön más soroknak, és más soroktól sem vehet fel erőforrást.

A prioritási rendszer azt veszi alapul, hogy a munkaterhelés hány GPU-t használ. Az 1 GPU-s feladat kapja a legalacsonyabb, a 8 GPU-s a legmagasabb prioritást. A Krea szerint ez segít megszüntetni a GPU-k feldarabolódását. Ha például összesen 8 GPU szabad két csomóponton, egy 8 GPU-t kérő pod a rendszer nélkül várakozó állapotban maradna. A prioritás használatával a rendszer eltávolítja az egyik csomópont feladatait, az új pod oda kerülhet, a korábban futó munkák pedig a másik csomópontra vándorolhatnak.

A megoldás a lekérdezés pontosságától is függ

A Krea szerint a rendszer működőképes, de nem tökéletes. Előfordul, hogy módosítani kell a PromQL-lekérdezést, amikor új logikát kell beépíteni. Ha a lekérdezés és a klaszter tényleges állapota eltér egymástól, a csomópont akkor is taintelt maradhat, amikor a klaszter valójában megtelt. A bejegyzés a korlátozások ismertetését ezen a ponton szakítja meg.

A bemutatott felépítés a Krea felhasználási céljára az alapértelmezett Kubernetes-ütemezőt egészíti ki dinamikus taintokkal, Deschedulerrel és Kueue-val. Így a vállalat a belső GPU-kapacitás kihasználására, a munkák fokozatos áthelyezésére és a külső kapacitás csak szükség esetén történő használatára építi a folyamatot.