Primi passi con l’accelerazione GPU di Polars: query 13 volte più veloci

Capire l’integrazione con NVIDIA RAPIDS cuDF

La domanda di elaborazione e analisi dei dati più veloce è cresciuta esponenzialmente negli ultimi anni, man mano che i dataset sono arrivati a centinaia di milioni o persino miliardi di righe. Per affrontare questa sfida, la data science ha visto un passaggio crescente dall’elaborazione basata su CPU al calcolo accelerato da GPU. È qui che entra in gioco NVIDIA RAPIDS cuDF.

NVIDIA RAPIDS è una suite open-source di librerie che abilita la data science e l’analisi accelerate da GPU. Al suo nucleo, RAPIDS è progettato per semplificare i flussi di lavoro sui dati sfruttando il massiccio parallelismo delle GPU NVIDIA per accelerare compiti tipicamente vincolati alla CPU. 

All’interno di RAPIDS, cuDF è la libreria responsabile delle operazioni sui DataFrame. cuDF estende l’astrazione familiare del DataFrame alla memoria della GPU, consentendo l’integrazione nei flussi di lavoro senza ampie modifiche al codice. La libreria supporta tutte le principali operazioni che i data scientist eseguono sui DataFrame: filtro, aggregazione, join e ordinamento, tutto alimentato dalla GPU per la velocità.

NVIDIA e Polars hanno appena annunciato l’open beta del motore GPU di Polars per Python, basato su RAPIDS cuDF. Questa nuova funzionalità segna un passo avanti significativo per l’elaborazione dati ad alte prestazioni, offrendo fino a 13 volte flussi di lavoro più rapidi agli utenti Polars con GPU NVIDIA!

Progettazione di alto livello dell’API di Polars, incluso un motore GPU. Fonte immagine: Polars

Con l’aggiunta del motore GPU, gli utenti di Polars possono decidere su quale motore eseguire i loro carichi di lavoro. L’ottimizzatore di Polars supporta tutti i motori, determinando dinamicamente quali query possono essere eseguite su GPU o CPU.

Il motore GPU di Polars: una svolta

Ora hai un po’ di contesto su Polars e sulla libreria NVIDIA RAPIDS cuDF, su come si integrano e come sono rese disponibili tramite l’API Python di Polars. Ma cosa significa tutto questo per sviluppatori e data scientist?

Ecco un riepilogo di ciò che abilita il motore GPU di Polars:

• Integrazione con la Lazy API di Polars in Python: puoi semplicemente passare engine="gpu" all’operazione collect() della Lazy API di Polars, abilitando l’elaborazione su GPU senza modifiche significative al codice.
• Elaborazione interattiva di grandi dataset: il motore GPU è pensato per rendere interattiva l’elaborazione di centinaia di milioni di righe. Con una singola GPU puoi gestire volumi di dati enormi che richiederebbero soluzioni più complesse e lente.
• Ottimizzato per l’efficienza: il motore GPU di Polars sfrutta appieno l’ottimizzatore di Polars per garantire un’esecuzione il più efficiente possibile, riducendo tempi e uso di memoria.
• Fallback su CPU senza interruzioni: nei casi in cui specifiche query non siano supportate su GPU, il motore include un fallback su CPU. In questo modo il flusso di lavoro non si interrompe, tornando automaticamente alla CPU quando necessario.

E i risultati finora? Vediamo. 

Il motore GPU di Polars accelera l’elaborazione dei dati fino a 13 volte. Fonte immagine: Polars

Il benchmark sopra dimostra le impressionanti capacità del motore GPU di Polars, accelerando le query fino a 13 volte.

Tuttavia, è importante notare che usare Polars su GPU non è sempre garantito essere più veloce rispetto alla CPU, in particolare per query semplici che non sono computazionalmente intensive. In questi casi, le prestazioni sono spesso limitate dalla velocità di lettura dei dati dal disco: quando le operazioni di input/output (I/O) sono il collo di bottiglia, l’accelerazione GPU offre un vantaggio minore.

Detto ciò, il motore GPU offre accelerazioni significative per operazioni dati più complesse, rendendolo uno strumento prezioso per i team che gestiscono grandi dataset o query con join, group-by e processamento di stringhe.

L’interazione tra CPU e GPU

Una delle caratteristiche più interessanti di questa nuova integrazione è la sua flessibilità. Come accennato, l’ottimizzatore può passare tra esecuzione su CPU e su GPU in base alla complessità delle tue query. Ma come avviene esattamente questa transizione?

Polars senza accelerazione GPU

Quando esegui Polars sulla CPU, il motore segue una pipeline strutturata e altamente ottimizzata per eseguire le operazioni sui dati. Il processo prevede approssimativamente i seguenti passaggi:

1. DSL (domain-specific language): il motore Polars crea innanzitutto una traccia di query strutturata dal tuo codice, simile all’albero sintattico di un linguaggio di programmazione. Questa traccia definisce la sequenza di operazioni da eseguire, come filtrare, raggruppare e aggregare.
2. IR (intermediate representation): successivamente, Polars converte la traccia di query in un piano di esecuzione dettagliato, assicurandosi che tutte le operazioni siano corrette e che gli schemi dei dati combacino. Questo passaggio getta le basi di come Polars elaborerà i dati.
3. Ottimizzatore: l’ottimizzatore migliora il piano di esecuzione riordinando le operazioni e rimuovendo passaggi non necessari. Ad esempio, i filtri possono essere spostati prima di join o aggregazioni per ridurre prima possibile la dimensione del dataset. Questa ottimizzazione garantisce che il piano sia il più efficiente possibile prima dell’esecuzione.
4. IR ottimizzata: dopo l’ottimizzazione, il piano migliorato è finalizzato e pronto per essere elaborato dal motore Polars. Questo piano è il blueprint di come Polars eseguirà ogni operazione sui dati.
5. Motore in-memory di Polars (esecuzione su CPU): infine, il motore in-memory di Polars esegue il piano ottimizzato interamente sulla CPU. Nonostante le ottimizzazioni ad alte prestazioni di Polars, questo passaggio può richiedere tempo con dataset di grandi dimensioni, poiché le CPU elaborano i dati in modo sequenziale o con parallelismo limitato.

Polars con accelerazione GPU

Quando abiliti l’accelerazione GPU passando engine="gpu" al metodo .collect(), il processo in superficie resta essenzialmente lo stesso, ma con controlli e ottimizzazioni aggiuntivi dietro le quinte per determinare quanto carico può essere trasferito sulla GPU tramite cuDF:

1. DSL (domain-specific language): come in modalità CPU, Polars inizia creando una traccia di query strutturata dal tuo codice. Questa traccia è la stessa sia che tu stia eseguendo su CPU sia su GPU, e cattura le operazioni che vuoi eseguire.
2. IR (intermediate representation): Polars converte la traccia in un piano di esecuzione, verificando la correttezza e la validità degli schemi dei dati. In questa fase l’IR resta neutrale, poiché il motore non ha ancora deciso se le operazioni saranno gestite dalla CPU o dalla GPU.
3. Ottimizzatore: l’ottimizzatore lavora come per l’esecuzione su CPU, riordinando le operazioni per la massima efficienza e rimuovendo i passaggi ridondanti. A questo punto, il piano è ancora ottimizzato per l’esecuzione generale.
4. IR ottimizzata: il piano ottimizzato è finalizzato come per l’esecuzione su CPU. Tuttavia, con l’accelerazione GPU attivata, il motore verifica ora se parti di questo piano possono essere spostate sulla GPU.
5. Callback cuDF (verifica GPU): è qui che avviene la magia! Il callback di cuDF determina se la GPU può gestire il piano di esecuzione. Alcune operazioni si prestano naturalmente all’accelerazione GPU, altre meno. Il callback modifica il piano per usare la GPU nelle operazioni supportate.
6. Motore in-memory di Polars (esecuzione su GPU e CPU): infine, il motore in-memory di Polars esegue il piano. La GPU gestisce le query che possono essere accelerate, mentre le query non gestibili dalla GPU tornano alla CPU.

Questo modello di esecuzione ibrido—dove GPU e CPU sono usate in tandem—rende Polars estremamente flessibile. Benefici dell’accelerazione GPU senza rinunciare alla compatibilità per le query non supportate su GPU. Il post di rilascio di Polars lo spiega in maggior dettaglio tecnico.

Primi passi con il motore GPU di Polars

Finalmente, la parte più entusiasmante! In questa sezione ti guiderò nei passaggi necessari per configurare l’ambiente e iniziare a sfruttare l’elaborazione dati accelerata dalla GPU.

Il modo più semplice per partire è seguire questo notebook Google Colab. 

Colab offre un uso GPU gratuito ma limitato, particolarmente utile se non hai accesso immediato a una GPU. Se usi il notebook fornito, partiamo dal punto 2 qui sotto, quindi puoi saltare i passaggi precedenti.

Prerequisiti

Prima di installare il motore GPU di Polars, assicurati che il tuo sistema soddisfi i requisiti per NVIDIA RAPIDS cuDF. Puoi consultare le specifiche di sistema necessarie, inclusa la compatibilità GPU e i requisiti dei driver, nella pagina di documentazione di RAPIDS.

1. Crea un ambiente virtuale (consigliato)

Per iniziare con il motore GPU di Polars, è buona pratica creare un ambiente virtuale per isolare le dipendenze del progetto. 

In questo esempio uso conda, ma sentiti libero di utilizzare il tuo gestore di pacchetti preferito.

Oltre a Python, puoi includere JupyterLab nell’ambiente per l’esplorazione e lo sviluppo interattivi dei dati, particolarmente utile per eseguire piccoli snippet di codice e analizzare i dataset al volo.

conda create -n polars-gpu -c conda-forge python=3.11 jupyterlab

Dopo aver creato l’ambiente virtuale, attivalo con:

conda activate polars-gpu

2. Installa il motore GPU di Polars

Ora abilitiamo l’accelerazione GPU installando Polars e il motore GPU. Questo imposterà anche cuDF e altre dipendenze necessarie per usare le GPU NVIDIA:

pip install -U polars[gpu] --extra-index-url=https://pypi.nvidia.com

Una volta installato il motore GPU di Polars, usare Polars su GPU sarà simile a lavorare con la CPU, ma con risultati molto più rapidi per flussi di lavoro complessi o con molti dati.

Come spiegato in precedenza, il motore di Polars ottimizza automaticamente il piano di esecuzione, sfruttando appieno la GPU NVIDIA per velocizzare le operazioni e ridurre al minimo l’uso di memoria. Per le query che non supportano l’accelerazione GPU, il motore GPU di Polars esegue un fallback su CPU senza interruzioni, garantendo la continuità del flusso di lavoro.

Mettiamo il motore all’opera con alcuni semplici esempi!

Esecuzione di alcune query con il motore GPU di Polars

Per questa dimostrazione useremo un dataset da 22 GB di transazioni finanziarie simulate da Kaggle. NVIDIA ospita il dataset su un bucket Google Cloud Storage (GCS), garantendo velocità di download elevate.

Per iniziare, scarica il dataset:

!wget https://storage.googleapis.com/rapidsai/polars-demo/transactions.parquet -O transactions.parquet

Poiché questo dataset proviene da Kaggle, è soggetto a una licenza e termini d’uso specifici per i dataset Kaggle.

Carichiamo il dataset con Polars e ispezioniamo lo schema:

import polars as pl
from polars.testing import assert_frame_equal

transactions = pl.scan_parquet("transactions.parquet")
transactions.collect_schema()

Ecco l’output:

Schema([('CUST_ID', String),
        ('START_DATE', Date),
        ('END_DATE', Date),
        ('TRANS_ID', String),
        ('DATE', Date),
        ('YEAR', Int64),
        ('MONTH', Int64),
        ('DAY', Int64),
        ('EXP_TYPE', String),
        ('AMOUNT', Float64)])

Ora calcoliamo il volume totale delle transazioni sommando la colonna AMOUNT. Prima proviamo senza accelerazione GPU:

transactions.select(pl.col("AMOUNT").sum()).collect()

Output:

AMOUNT
f64
3.6183e9

È un volume totale molto alto! Eseguiamo la stessa query sulla GPU:

transactions.select(pl.col("AMOUNT").sum()).collect(engine="gpu")

Output:

AMOUNT
f64
3.6183e9

Per operazioni semplici come questa, CPU e GPU producono lo stesso risultato con velocità simili, perché la query non è abbastanza intensiva da trarre grande vantaggio dalla GPU. La GPU però brillerà con query più complesse.

Passiamo a una query più complessa. In questa query raggruppiamo le transazioni per ID cliente (CUST_ID), sommiamo gli importi totali delle transazioni e poi ordiniamo i risultati per i maggiori spender. 

Per prima cosa, eseguiamola sulla CPU:

%%time
res_cpu = (
   transactions
   .group_by("CUST_ID")
   .agg(pl.col("AMOUNT").sum())
   .sort(by="AMOUNT", descending=True)
   .head()
   .collect()
)
res_cpu

Ecco l’output:

CPU times: user 4.63 s, sys: 3.75 s, total: 8.38 s
Wall time: 6.04 s
CUST_ID
AMOUNT
str
f64
"CA9UYOQ5DA"
2.0290e6
"CJUK2MTM5Q"
1.8115e6
"CYXX1NBIKL"
1.8082e6
"C6ILEYAYQ9"
1.7961e6
"CCNBC305GI"
1.7274e6

Ora eseguiamo la stessa query sulla GPU:

%%time
res_gpu = (
   transactions
   .group_by("CUST_ID")
   .agg(pl.col("AMOUNT").sum())
   .sort(by="AMOUNT", descending=True)
   .head()
   .collect(engine=”gpu”)
)
res_gpu

Output:

CPU times: user 347 ms, sys: 0 ns, total: 347 ms
Wall time: 353 ms
shape: (5, 2)
CUST_ID
AMOUNT
str
f64
"CA9UYOQ5DA"
2.0290e6
"CJUK2MTM5Q"
1.8115e6
"CYXX1NBIKL"
1.8082e6
"C6ILEYAYQ9"
1.7961e6
"CCNBC305GI"
1.7274e6

Come si vede, usare la GPU per query più complesse offre guadagni notevoli, riducendo il tempo di esecuzione da 6,04 secondi su CPU a soli 353 millisecondi su GPU!

Questo esempio dimostra la potente spinta prestazionale che il motore GPU di Polars offre per le operazioni su larga scala. 

Puoi trovare esempi più avanzati nel relativo notebook Colab.

Conclusione

Il motore GPU di Polars, basato su NVIDIA RAPIDS cuDF, porta notevoli miglioramenti di velocità, con elaborazioni dei dati fino a 13 volte più rapide per operazioni complesse. Per dataset su larga scala, Polars offre un chiaro vantaggio rispetto alle librerie di DataFrame tradizionali.

Anche se non tutte le query trarranno lo stesso beneficio dall’accelerazione GPU, Polars resta uno strumento potente per chi lavora con grandi quantità di dati. La facilità di integrazione, unita al modello di esecuzione ibrido CPU-GPU, rende Polars un forte candidato per i flussi di lavoro moderni sui dati.

Se vuoi potenziare le tue competenze di manipolazione dei dati, soprattutto con Pandas, ti consiglio vivamente questi corsi:

• Data Manipulation with Pandas
• Streamlined Data Ingestion with Pandas
• Writing Efficient Code with Pandas

Queste risorse ti daranno basi solide che si integrano con le potenti capacità di strumenti come Polars!