P-MOSS: Scheduling Main-Memory Indexes Over NUMA Servers Using Next Token Prediction

Riassunto Tecnico: P-MOSS

Definizione del Problema

I moderni server multi-core che utilizzano architetture NUMA (Non-Uniform Memory Access) e Chiplet esibiscono una significativa eterogeneità hardware. Questa complessità crea uno vasto spazio di progettazione in cui le prestazioni degli indici in memoria principale (ad esempio, i B+-Tree) possono variare fino a un ordine di grandezza a seconda di due fattori:

1. Partizionamento dei Dati: La posizione delle partizioni dell'indice attraverso i nodi NUMA (controller di memoria).
2. Scheduling dei Core: L'assegnazione dei core logici per eseguire query mirate a specifiche partizioni dell'indice.

Questa combinazione è definita scheduling spaziale. Gli approoli tradizionali si affidano spesso a euristiche statiche (ad esempio, round-robin, interleaving) o ai default del sistema operativo, che non riescono ad adattarsi alle diverse topologie degli hardware moderni (Intel, AMD, NVIDIA, ARM) e ai vari pattern di carico di lavoro (read-heavy, write-heavy, scan). Gli studi esistenti si concentrano spesso sullo scheduling temporale (quando eseguire una query) o affrontano il posizionamento dei dati per tabelle relazionali, ignorando spesso l'aspetto critico dello scheduling spaziale tra core e dati negli indici in memoria principale. Inoltre, le soluzioni basate su euristiche faticano a generalizzare tra diversi vendor di CPU e architetture a causa della natura combinatoria del problema dello scheduling.

Metodologia: P-MOSS

Il documento presenta P-MOSS (Performance MOnitoring Unit-driven Spatial Query Scheduling), un framework appreso che sfrutta i principi dei Large Language Models (LLM) per ottimizzare lo scheduling spaziale.

Paradigma Centrale: Next Token Prediction (NTP)

P-MOSS riformula il problema dello scheduling spaziale come un compito di Next Token Prediction (Predizione del Prossimo Token). Invece di predire la parola successiva in una frase, il sistema predice il prossimo Core ID in una sequenza di decisioni di scheduling.

• Sequenza: La sequenza rappresenta la policy di scheduling per un indice, dove ogni "token" corrisponde all'assegnazione di uno slice dell'indice (una partizione logica basata sull'intervallo di chiavi) a un core CPU specifico.
• Contesto: A differenza dei token del linguaggio naturale, i Core ID mancano di una regolarità sintattica intrinseca. P-MOSS affronta questo problema arricchendo i "token" con il contesto derivato dalle statistiche delle performance hardware.

Architettura e Pipeline di Addestramento

P-MOSS utilizza un'architettura Decision Transformer (DT), addestrata tramite Reinforcement Learning (RL) Offline. Il sistema opera in tre fasi distinte:

1. Pre-addestramento (Generative Pre-training):

   • Fonte dei Dati: Un dataset offline è curato da hardware diversificati (Intel, AMD, NVIDIA, IBM), varie configurazioni NUMA (1NPS, 4NPS, Chiplet) e molteplici carichi di lavoro (varianti YCSB).
   • Processo: Il sistema interroga le Performance Monitoring Units (PMU) dell'hardware durante l'esecuzione delle query per generare "Hardware Snapshots". Questi snapshot catturano statistiche a basso livello (conteggi delle istruzioni, cache miss, traffico di memoria) senza contabilità a livello software.
   • Tokenizzazione: Lo snapshot hardware, la policy di scheduling corrente e le metriche di performance vengono tokenizzati in token State (Stato), Action (Azione), Return-to-Go (RTG) e Meta.
   • Obiettivo: Il DT apprende pattern generalizzabili di scheduling efficiente attraverso ambienti diversi.

2. Inferenza (Esecuzione Online):

   • Il modello "base" addestrato predice le policy di scheduling per l'hardware e il carico di lavoro target.
   • P-MOSS avvia il processo con una policy casuale o precedente, monitora le statistiche PMU in tempo reale e le fornisce al DT per inferire la mappatura ottimale degli slice dell'indice ai core.
   • Il sistema impone queste decisioni migrando gli slice dell'indice ai corrispondenti nodi NUMA e vincolando l'esecuzione delle query a core specifici.

3. Post-addestramento (Fine-tuning):

   • Mentre il sistema opera sull'hardware target, raccoglie un "dataset di fine-tuning" specifico per quell'ambiente.
   • Il modello base viene aggiornato incrementalmente (offline) per creare un modello "assistente", adattandosi alle caratteristiche uniche del server e del carico di lavoro specifico.

Componenti Tecnici Chiave

• Hardware Profiling (PMU): P-MOSS si affida esclusivamente a statistiche hardware a basso livello raccolte tramite la PMU (ad esempio, LLC misses, traffico del controller di memoria). Ciò evita l'overhead e l'imprecisione della contabilità a livello software.
• Strategia di Tokenizzazione:
  • State Token: Rappresenta lo stato hardware attuale (core occupati, statistiche cache/memoria) tramite token View, Position e Machine.
  • Action Token: Il Core ID predetto.
  • RTG Token: La ricompensa futura desiderata (throughput target) meno le ricompense attuali, guidando il modello verso risultati ad alte prestazioni.
  • Meta Token: Contesto di sistema (numero di socket, topologia NUMA, vendor della CPU).
• Offline RL: Decoupling (disaccoppiamento) del processo di apprendimento dal percorso critico di esecuzione delle query, garantendo che il DBMS non soffra di degradazione delle prestazioni durante la fase di apprendimento.

Contributi Chiave

1. Framework P-MOSS: Il primo framework appreso per lo scheduling spaziale delle query in indici in memoria principale su moderni server NUMA e Chiplet, dando priorità sia al posizionamento dei dati che allo scheduling dei core.
2. Ottimizzazione guidata da PMU: Un approccio innovativo all'ottimizzazione del kernel del DBMS che si basa esclusivamente su statistiche hardware a basso livello (PMU), garantendo l'adattabilità tra diversi vendor di hardware.
3. Formulazione NTP: La formulazione dello scheduling delle query come un compito di Next Token Prediction, consentendo l'uso di Offline RL e Decision Transformer per apprendere le policy di scheduling senza interazione online con il DBMS.
4. Validazione Empirica: Valutazione estesa su una vasta gamma di CPU (Intel, AMD, ARM, IBM), architetture NUMA (1-4 socket, 2-8 nodi) e carichi di lavoro.

Risultati Sperimentali

La valutazione è stata condotta su un B+-Tree in memoria principale con 30M fino a 1B di record su cinque diverse configurazioni di server (AMD Milan, Intel Skylake X, Intel Sandy Bridge, NVIDIA Grace Hopper, IBM Power).

• Miglioramento del Throughput: P-MOSS raggiunge un miglioramento fino a 6× nel throughput delle query rispetto alle tradizionali policy di scheduling.
• Confronto con i Baseline:
  • Rispetto ai default del sistema operativo (OS:D, OS:I) e alle strategie NUMA-aware (SE:N, SN:N), P-MOSS mostra un aumento medio di 1.92× a 3.66×.
  • Rispetto ad altre strategie "Shared Nothing" basate su apprendimento o euristiche (Grouped, Mixed, Spread, Random), P-MOSS supera costantemente le altre, con miglioramenti che vanno da 1.27% a 3.95% a seconda dell'hardware e del carico di lavoro specifici.
• Generalizzazione: P-MOSS dimostra forti capacità di generalizzazione, performando bene su hardware non visto (IBM Power) e distribuzioni di dati non viste (dataset OSM) senza dover riaddestrare da zero, sfruttando il suo modello base pre-addestrato.
• Overhead: L'overhead per la raccolta delle statistiche hardware tramite PMU è minimo (~0.73% di degradazione delle prestazioni, e l'uso di istruzioni SIMD per l'elaborazione di tali statistiche migliora le prestazioni del sistema del 60.37% durante la fase di profiling).

Significato e Rivendicazioni

Il documento afferma che P-MOSS dimostra la fattibilità dell'applicazione dei concetti fondamentali degli LLM, specificamente la Next Token Prediction, ai compiti di ottimizzazione dei database. Trattando lo scheduling come un problema di modellazione di sequenze, il framework può navigare la complessità combinatoria degli spazi di progettazione dell'hardware moderno in modo più efficace rispetto alle euristiche statiche.

Gli autori sottolineano che:

• Lo scheduling spaziale è critico per massimizzare le prestazioni su hardware eterogenei moderni.
• Le statistiche delle performance hardware (PMU) forniscono una rappresentazione delle feature robusta e a basso overhead per l'ottimizzazione del database guidata dal ML.
• L'Offline RL permette il disaccoppiamento dell'apprendimento dall'esecuzione, consentendo un adattamento sicuro e scalabile a diversi ambienti.
• L'approccio non è limitato allo scheduling spaziale ma offre un blueprint per generalizzare le tecniche ispirate agli LLM a sfide di ottimizzazione del DBMS più ampie.