AI21 Labs ha adottato Google Cloud AI Hypercomputer per automatizzare la gestione dei propri workload AI su Google Kubernetes Engine (GKE). Il risultato è una riduzione dell'83% dei tempi di attesa per i job ad alta priorità, passati da 72 ore a 12, e l'eliminazione degli interventi manuali di scheduling, che erano circa 20 a settimana.
Il contesto: cluster GPU condivisi e scheduling manuale
AI21 sviluppa modelli linguistici, tra cui la famiglia Jamba, e prodotti per l'ottimizzazione di agenti AI. I suoi workload di addestramento e produzione richiedono capacità GPU molto elevata. L'azienda utilizza un cluster GKE condiviso che aggrega migliaia di istanze A3 con NVIDIA H100 e A3 Ultra con NVIDIA H200, consentendo a diversi team di attingere alla stessa capacità invece di restare vincolati a slice dedicate.
Prima dell'automazione, la capacità veniva negoziata manualmente in un canale di chat. Quando l'utilizzo del cluster si avvicinava al 100%, ogni richiesta diventava una trattativa. I job ad alta priorità, spesso multi-nodo e critici per i progetti di modello, potevano restare bloccati fino a 72 ore in attesa di capacità contigua sufficiente.
Due problemi distinti: contenzione e frammentazione
La scarsità di risorse generava due problemi separati. Il primo era la contenzione: decidere quale team avesse accesso al calcolo successivo. Il secondo era la frammentazione: capacità tecnicamente libera ma distribuita in frammenti troppo piccoli per un job grande, un tipico problema di packing che la negoziazione manuale non poteva risolvere.
Senza un meccanismo di ammissione completo o nulla, il cluster poteva entrare in uno stato di stallo, con macchine che tenevano risorse senza produrre lavoro utile finché qualcuno non interveniva manualmente.
La scelta di Kueue
Il team di AI21 ha valutato scheduler open source come Apache YuniKorn, Volcano e Kueue. Kueue è stato scelto per semplicità e integrazione: funziona con Kubernetes standard, non richiede la sostituzione dei componenti core dello scheduler e non impone di riscrivere le specifiche dei job.
- Integrazione nativa con GKE, inclusi tipi di capacità come Spot VM e Dynamic Workload Scheduler.
- Possibilità di passare automaticamente alla capacità elastica quando la capacità riservata è piena.
- Mantenimento di un livello di controllo adeguato per workload AI intensivi, con accesso vicino a GPU e driver senza gestire istanze nude.
Feedback loop open source con Google Cloud
Grazie alla partnership con Google Cloud, AI21 è diventata design partner del team Kueue. Ha condiviso un documento di requisiti per comportamenti necessari in produzione ma non ancora disponibili, come un ordinamento equo di ammissione per job multi-nodo senza preemption. Il team Kueue ha sviluppato Admission Fair Sharing, che riordina la coda di ammissione a favore dei team che storicamente hanno usato meno capacità, senza interrompere i job già in esecuzione.
AI21 ha inoltre abilitato Topology Aware Scheduling, che rende Kueue consapevole della topologia fisica del cluster e rifiuta di ammettere job che non possono essere eseguiti su un singolo nodo, evitando collocazioni senza capacità reale.
Risultati operativi
- Interventi manuali di scheduling: da circa 20 a settimana a zero.
- Attesa dei job ad alta priorità: da 72 ore a 12 ore.
- Frammentazione del cluster: dal 15% all'8%.
- Eliminazione dei job parzialmente ammessi senza capacità sufficiente per completare l'esecuzione.
Il costo totale non è cambiato in modo significativo, perché la flotta riservata viene mantenuta deliberatamente vicino al 100% di utilizzo. Il miglioramento principale riguarda il tempo delle persone: i responsabili di team non devono più arbitrare dispute di capacità e i ricercatori non devono più attendere risposte in chat.
Perché è utile per PMI e MSP
Il caso mostra come l'automazione dello scheduling possa ridurre il toil operativo in ambienti AI o container ad alta densità. Per le PMI che adottano workload AI, e per gli MSP che gestiscono piattaforme cloud, l'automazione della capacità riduce interventi manuali, migliora la prevedibilità dei job critici e aiuta a mantenere alta l'utilizzazione delle risorse senza aumentare i costi infrastrutturali.
Il passaggio chiave non è solo aggiungere più GPU, ma automatizzare l'ammissione, la priorità e la gestione della frammentazione in modo trasparente e ripetibile.
