stratum: A System Infrastructure for Massive Agent-Centric ML Workloads

Il paper presenta Stratum, un'infrastruttura di sistema unificata che supera le limitazioni dell'ecosistema ML Python per supportare la ricerca di pipeline guidata da agenti su larga scala, ottenendo un'accelerazione fino a 16,6 volte grazie alla decoupling dell'esecuzione dalla pianificazione e all'uso di un runtime ottimizzato in Rust.

L'essenziale

Immagina di dover costruire una casa. Fino a poco tempo fa, per farlo, avevi bisogno di un architetto umano (il data scientist) che disegnasse i piani, scegliesse i materiali e supervisionasse ogni mattone. Era un lavoro lento, faticoso e soggetto a errori.

Ora, grazie all'intelligenza artificiale (in particolare i grandi modelli linguistici o LLM), abbiamo assunto dei robot architetti (gli "agenti"). Questi robot sono incredibilmente veloci: possono disegnare migliaia di progetti diversi in pochi secondi, provare a costruirli, vedere quale funziona meglio e rifare i piani se qualcosa non va.

Il problema? Il nostro attuale "cantiere" (l'ecosistema Python che usiamo per l'IA) è fatto per essere guidato a mano da un umano, non per gestire migliaia di robot che lavorano tutti insieme. È come se avessimo un cantiere con un solo muletto e un solo operaio, ma avessimo assunto 10.000 robot che urlano ordini contemporaneamente. Risultato? Il cantiere va in tilt, i robot si scontrano, i materiali si accumulano ovunque e il lavoro procede lentissimamente.

Stratum è la soluzione proposta dagli autori di questo articolo. È un nuovo sistema operativo per il cantiere, progettato specificamente per gestire questi eserciti di robot.

Ecco come funziona, spiegato con metafore semplici:

1. Il Problema: Il "Cantiere Caotico"

Attualmente, quando un agente AI genera un nuovo piano (una "pipeline" di dati), il sistema lo esegue come se fosse un compito completamente nuovo, ignorando che 100 piani prima aveva già tagliato gli stessi mattoni o mescolato lo stesso cemento.

• Spreco: I robot fanno calcoli inutili (come tagliare lo stesso legno 50 volte).
• Ingolfamento: Il computer si blocca perché prova a fare troppe cose contemporaneamente senza un piano (memoria piena, CPU al 100% ma senza risultati).
• Lentezza: Tutto è scritto in Python, che è come usare un cucchiaio di legno per scavare una trincea: funziona, ma è lento rispetto a una pala di metallo.

2. La Soluzione: Stratum, il "Capocantiere Superintelligente"

Stratum non cerca di sostituire i robot (gli agenti AI), ma diventa il loro capocantiere perfetto. Ecco le sue tre armi segrete:

A. La Mappa Magica (Grafo di Esecuzione)

Invece di far partire ogni robot a caso, Stratum prende tutti i progetti che gli agenti hanno disegnato e li fonde in un'unica mappa gigante.

• Metafora: Immagina che invece di dare a ogni muratore un foglio di carta diverso, Stratum prende tutti i fogli, li sovrappone e vede che il "muro di mattoni rossi" è uguale in 50 progetti diversi. Invece di costruirlo 50 volte, lo costruisce una volta sola e lo usa per tutti. Questo si chiama ottimizzazione logica.

B. Il Motore in Acciaio (Backend Rust)

Python è comodo, ma lento. Stratum ha costruito un motore interno scritto in Rust (un linguaggio di programmazione super veloce e sicuro).

• Metafora: Se Python è un'auto con il motore a vapore, Rust è un razzo. Stratum prende i compiti pesanti (come mescolare il cemento o tagliare i metalli) e li passa al motore in acciaio, lasciando a Python solo il compito di dare gli ordini. Inoltre, questo motore non si blocca mai perché può gestire centinaia di robot che lavorano in parallelo senza impazzire.

C. Il Magazzino Intelligente (Caching)

Stratum tiene traccia di tutto ciò che è stato già fatto.

• Metafora: Se un robot ha già pulito e tagliato 1000 mattoni per un progetto che è stato scartato, Stratum li mette in un magazzino. Se il robot successivo ha bisogno degli stessi 1000 mattoni, Stratum non li fa tagliare di nuovo: li prende dal magazzino e li consegna istantaneamente. Questo fa risparmiare un tempo enorme.

3. Il Risultato: Velocità Esplosiva

Gli autori hanno testato questo sistema con un agente AI reale (chiamato AIDE).

• Senza Stratum: Il sistema impiegava molto tempo, si bloccava spesso e usava male le risorse.
• Con Stratum: Il sistema è diventato 16,6 volte più veloce.

È come se, invece di impiegare 16 ore per completare un lavoro, ci volesse solo un'ora, con meno sprechi di energia e meno stress per il computer.

In Sintesi

Stratum è l'infrastruttura che mancava per permettere all'Intelligenza Artificiale di fare il suo vero lavoro: esplorare migliaia di possibilità per trovare la soluzione migliore, senza impantanarsi nella lentezza dei vecchi sistemi informatici.

Mentre prima dovevamo guidare l'AI passo dopo passo, ora possiamo lasciarle il volante e darle un motore potente e un navigatore intelligente che le permette di correre alla massima velocità senza mai uscire dalla strada.

Sommario tecnico

Ecco un riassunto tecnico dettagliato del paper "stratum: A System Infrastructure for Massive Agent-Centric ML Workloads [Vision]", presentato in italiano.

1. Il Problema: Il Mismatch tra Agenti ML e Infrastruttura Esistente

Il lavoro identifica una crescente discrepanza tra l'evoluzione dei modelli linguistici (LLM) e l'infrastruttura di sistema attuale per il Machine Learning (ML).

• Nuovo Carico di Lavoro (Agentic Pipeline Search): Gli agenti ML, guidati da LLM, stanno automatizzando la creazione, la validazione e l'ottimizzazione di pipeline di data science. Questi agenti operano in modo esplorativo, generando e eseguendo migliaia di varianti di pipeline Python (basate su librerie come Pandas e scikit-learn) per profilare dati, generare codice e rifinare iperparametri.
• Limitazioni dell'Ecosistema Python: L'ecosistema ML attuale è progettato per flussi di lavoro interattivi e sequenziali guidati da umani. È vincolato dal modello di esecuzione interpretato di Python, dall'isolamento a livello di libreria e dalla mancanza di supporto runtime per l'esecuzione massiva di pipeline.
• Inefficienze degli Agenti: Gli agenti attuali mancano di supporto runtime dedicato, portando a:
  • Esecuzioni sovrapposte e ridondanti.
  • Fallimenti per esaurimento della memoria (OOM).
  • Sottoutilizzazione dell'hardware (CPU/GPU).
  • Mancanza di ottimizzazione olistica tra le diverse fasi della pipeline.
• Inadeguatezza dei Sistemi Esistenti: I sistemi ML ad alte prestazioni proposti dalla comunità di ricerca (es. SystemML, DAPHNE) spesso richiedono linguaggi specifici (DSL) o modelli di programmazione specializzati, rendendoli difficili da integrare con le librerie Python standard su cui si basano gli agenti.

2. Metodologia e Architettura di Stratum

Per colmare questo divario, gli autori propongono stratum, un'infrastruttura di sistema unificata che disaccoppia l'esecuzione della pipeline dalla pianificazione e dal ragionamento dell'agente.

Principi di Design Chiave:

1. Supporto Trasparente: Integrazione senza interruzioni con librerie ML arbitrarie (scikit-learn, Pandas, PyTorch, ecc.) senza richiedere la riscrittura degli operatori.
2. Semantica degli Operatori e Valutazione Pigra (Lazy Evaluation): Rappresentazione delle pipeline come Grafi Aciclici Diretti (DAG) per permettere ottimizzazioni a livello di compilatore.
3. Runtime Eterogeneo: Esecuzione efficiente su backend diversi (CPU, GPU, distribuiti) con un runtime nativo in Rust.

Componenti Architetturali (Figura 3):

• Astrazione Dichiarativa: Stratum utilizza skrub come punto di ingresso. Trasforma codice Python arbitrario (inclusi UDF e trasformazioni Pandas) in operatori semantici espliciti (DataOps), formando un DAG valutato in modo pigro.
• Ottimizzatore Logico:
  • Raccolta Metadati: Estrae informazioni sugli operatori (tipo, libreria sorgente, proprietà strutturali, caratteristiche dei dati).
  • Rewrite: Applica riscritture basate su regole (es. spinta dei predicati, condivisione della lettura, eliminazione delle sottoespressioni comuni - CSE, folding delle costanti) e riscritture specifiche per le API (es. riordinamento operazioni Pandas per minimizzare le copie).
  • Abbassamento degli Operatori (Operator Lowering): Scompone operatori di alto livello (es. TableVectorizer o cross-validation) in operatori atomici più fini per un'ottimizzazione granulare.
• Selezione degli Operatori: Mantiene implementazioni fisiche multiple per ogni operatore logico. Utilizza un'euristica basata sui costi (memoria e tempo) per scegliere l'implementazione migliore (es. SVD di scikit-learn vs. Frequent Directions vs. kernel nativi Rust).
• Backend Rust: Sviluppa un runtime in Rust per gli operatori frequenti. Questo permette:
  • Controllo esplicito della memoria e accesso ai dati senza copie (zero-copy).
  • Rilascio del GIL (Global Interpreter Lock) di Python per l'esecuzione concorrente.
  • Sfruttamento del parallelismo dati nativo (tramite Rayon).
• Ottimizzazioni Runtime:
  • Parallelismo: Gestione intelligente del parallelismo intra-operatore e inter-operatore per evitare l'oversubscription delle risorse.
  • Riuso degli Intermedi: Caching dei risultati degli operatori (hashing basato su input e specifica) per riutilizzare dati e modelli parziali tra iterazioni successive della ricerca degli agenti.

3. Contributi Chiave

1. Analisi dei Carichi di Lavoro Agentic: Identificazione e caratterizzazione dei pattern di esecuzione tipici degli agenti ML (alta ridondanza, stato stateless, generazione massiva di pipeline) e dimostrazione delle inefficienze delle soluzioni attuali.
2. Progettazione di Stratum: Un'architettura sistemica che integra ottimizzazioni logiche, selezione di operatori, un backend in Rust e caching, specificamente mirata a supportare la ricerca di pipeline su larga scala.
3. Integrazione con l'Ecosistema Python: La capacità di operare su librerie esistenti senza richiedere agli agenti di imparare nuovi linguaggi o DSL.
4. Co-Design Agente-Sistema: Visione di agenti "consapevoli del sistema" che possono annotare le pipeline con metadati per guidare ottimizzazioni specifiche (es. selezione di operatori a bassa fedeltà nelle fasi esplorative iniziali).

4. Risultati Sperimentali

Gli autori hanno valutato un prototipo di stratum utilizzando un carico di lavoro reale generato dall'agente AIDE su un dataset immobiliare del Regno Unito (Kaggle).

• Ambiente: Nodo singolo con CPU AMD EPYC (48 core virtuali) e 256 GB di RAM.
• Confronto:
  • Base: Esecuzione sequenziale di pipeline (come fa AIDE standard).
  • Base_par: Esecuzione parallela tramite multiprocessing (AIDE standard).
  • Stratum: Prototipo con tutte le ottimizzazioni attivate.
• Performance:
  • Stratum è 16.6 volte più veloce rispetto all'esecuzione sequenziale di base.
  • Stratum è 7.8 volte più veloce rispetto all'esecuzione parallela di base (che soffre di overhead di serializzazione e consumo di memoria elevato).
• Analisi di Ablazione:
  • L'ottimizzazione logica (CSE, rewrite) da sola offre un speedup di 2.2x.
  • La selezione degli operatori (sostituzione di Python con Rust/Polars) aggiunge un ulteriore 4.5x.
  • Il parallelismo inter-operatore contribuisce ulteriormente, portando al risultato finale.

5. Significato e Impatto

Il paper "stratum" è significativo perché:

• Ridefinisce l'infrastruttura ML: Propone un cambio di paradigma da sistemi ottimizzati per task isolati a sistemi progettati per la ricerca esplorativa massiva guidata da agenti.
• Ponte tra AI e Sistemi: Risolve il problema dell'integrazione tra l'alta astrazione delle librerie Python (usate dagli LLM) e le esigenze di performance dei sistemi distribuiti.
• Scalabilità: Dimostra che è possibile scalare l'uso degli agenti ML in ambienti enterprise senza costi proibitivi di risorse, rendendo la ricerca automatica di pipeline (AutoML avanzato) praticabile su larga scala.
• Futuro della Ricerca: Apre la strada a sistemi multi-tenant, agenti consapevoli del sistema e motori di inferenza ottimizzati specificamente per le interazioni con gli agenti AI.

In sintesi, stratum non è solo un ottimizzatore, ma un'infrastruttura fondamentale per abilitare la prossima generazione di automazione nel Machine Learning, trasformando la ricerca di pipeline da un processo lento e costoso in un'operazione efficiente e scalabile.