Az NVIDIA Workload Map (korábbi nevén Karta) egy nyílt forráskódú Kubernetes CRD és egy Go-könyvtár, amely deklaratív módon írja le a különböző, elosztott AI-munkaterhelések szerkezetét. A megoldás célja, hogy a külső ütemezők, vezérlők és platformok ugyanazt az újrahasznosítható struktúra-leírást használhassák, így csökken a keretrendszerenként szükséges adapterek fejlesztési terhe.
Miért van szükség erre?
AI-feladatok Kubernetesen ritkán jelentik egyetlen pod futtatását: egy elosztott tanítási feladat vagy többcsomópontos inferencia szolgáltatás több szerepet, replikaszámot, pod-sablont és státuszinformációt tartalmaz. A különböző keretrendszerek (például JobSet, LeaderWorkerSet, PyTorchJob, RayCluster) ezeket eltérő CRD-sémákkal reprezentálják, ezért egy platformnak egyedi adaptereket kellett írnia minden támogatott típushoz. Az adapterek karbantartása és fejlesztése így többszörös munkát jelent minden új keretrendszer bevezetésekor.
A szerződéses példa egy JobSet manifesztum: egy leader Job és két worker Job, mindegyik pod 4 GPU-t kér. A JobSet összesített GPU-igénye így 12 GPU. Kubectl azonban a JobSet, a benne lévő jobok és a podok külön erőforrásként jeleníti meg, ezért a platformoknak külön értelmezniük kell a szerepeket és az összesített erőforrásigényeket.
Hogyan működik a Workload Map?
A Workload Map erőforrásként (Karta kind) és egy Go könyvtárként érkezik. Minden munkaterhelés-leírás (workload definition) meghatározza, hogy egy adott CRD-ben hol találhatók a pod-template-ek, a replikaszámok és a státuszmezők. A könyvtár egy élő objektumból egységes, fa-struktúrájú nézetet alakít ki, így a fogyasztók egyszerűen:
- kinyerhetik a pod-sablonokat, replikaszámokat, státuszt és metaadatokat,
- módosíthatják a pod-specet, címkéket és annotációkat minden példányon,
- megérthetik a hierarchiát (például JobSet → replicated Jobs → podok).
A Workload Map katalógusában előre elkészített definíciók találhatók többek között a következőkhöz: JobSet, PyTorchJob, RayCluster, RayJob, RayService, KServe, Knative, MPIJob, NVIDIA NIM, LeaderWorkerSet, Milvus, Grove PodCliqueSet és NVIDIA Dynamo. Új típusok YAML definícióval adhatók hozzá.
A JobSet példa — hogyan írható le egy definíció
A példa Karta-erőforrás gyökérkomponense meghatározza a JobSet saját típusát és hogy a státuszt hogyan kell olvasni. A státuszok átképezése (statusMappings) a Workload Map azon rétege, amely a natív mezőket egy közös, fogyasztók által használt fázis-sémára fordítja. A JobSet nem ad meg egyszerű "Running" feltételt, így a definíció egy jq-kifejezéssel állapítja meg a futó fázist (.status.replicatedJobsStatus alapján: van-e kész és aktív replika, és nincs-e hiba).
A childComponents rész specifikálja a belső szerkezetet: a .spec.replicatedJobs egy lista, amely minden szerephez egy bejegyzést tartalmaz. A podTemplateSpecPath a pod-sablon helyére mutat, a scaleDefinition pedig kiszámolja az egyes szerepekhez tartozó podok számát (.replicas * .template.spec.parallelism). Az instanceIdPath adja a szerepek azonosítóját (például "leader" és "worker").
Ezzel a deklaratív definícióval a keretrendszer-specifikus szerkezet egy újrahasznosítható leírásba kerül, amelyet több fogyasztó is használhat.
Hogyan használják különböző eszközök ugyanazt a leírást?
-
Workload Map CLI: a kliens eszköz betölti a definíciót, feloldja az élő objektumot és annak podjait, és fává rendezi azokat. A CLI telepítése Homebrew-val ismertetett parancssorokkal történik, és nem követeli meg a Workload Map klaszterbe telepítését. A demo JobSet esetén a kimenet megjeleníti a leader és worker szerepeket, valamint az összesített 12 GPU-igényt.
-
KAI Scheduler: az opcionális gang-scheduling instrukciók a definícióban adhatók meg (optimizationInstructions → gangScheduling). A KAI Scheduler ezeket a csoportosítási információkat használja a gang-admission és elhelyezési döntésekhez, és Workload Map-re támaszkodik a keretrendszer nélküli típusoknál.
-
NVIDIA Run:ai: egy API-hívással regisztrálható a munkaterhelés-típus a katalógus definíciójával, további kódfeladat nélkül. A platform ezután a leírást használja a JobSet egységként történő kezelése és láthatósága érdekében, például úgy, hogy a felhasználó számára egyetlen, 12 GPU-s munkaterhelésként jeleníti meg azt.
A CLI, a KAI Scheduler és az NVIDIA Run:ai ugyanazt a definíciót használja különböző célokra: ellenőrzésre, ütemezésre és platform-láthatóságra.
Elfogadás és használat
A Workload Map az Apache 2.0 licenc alatt érhető el és éles környezetben használja az NVIDIA Run:ai. A deklaratív definíciók támogatják a láthatóságot, a státuszértelmezést, a beküldést és a módosítást a támogatott Kubernetes munkaterhelés-típusok esetén. Többgyűrűs ütemezésben is alkalmazzák: Elotl’s Nova például a Workload Map-et használja a munkaterhelés szerkezetének és erőforrásigényének értelmezésére multi-cluster döntések során.
Hogyan kezdjem el?
A Workload Map elérhető az NVIDIA DSX AI Factory GitHub-szervezetén keresztül. Kezdésként érdemes megnézni a workload catalogot, hogy találjon-e már létező definíciót a használt frameworkhöz, vagy egy meglévőt adaptálni egy egyedi munkaterheléshez. A közösségi részvételhez csatlakozhat a NVIDIA DSX AI Factory Slack munkaterületéhez is.
A Workload Map célja, hogy a munkaterhelés-ismeretet a fordított kód helyett deklaratív, felülvizsgálható és klaszterbe alkalmazható manifestekbe helyezze, így csökkentve az egyedi adapterek sorozatos fejlesztésének szükségességét.
