PG-Stromをコンテナで動かしてみた

去年くらいからVTJのプロジェクトとしては「爆速DB」というプロジェクトで、DBの高速化の技術について検証をしています。

PG-Strom はPostgreSQLの拡張モジュールであり、データ分析やバッチ処理のために SQL ワークロードの GPU アクセラレーションを可能にするという技術です。GitHubでオープンソースで開発されており、インストール方法などは公式のドキュメントサイトにまとめられています。

コンテナで動かしてみようと思ったきっかけ

GitHubのIssueに、「コンテナでPG-Stromが動いたらいいよね」と言ったようなリクエストがされていました。

Official Docker image for pg-storm · Issue #548 · heterodb/pg-strom · GitHub

これをみた時点ではコンテナで動かす利点は思いつかなかったものの、試しに動かしてみることにしました。

必要なもの

試すにあたり、次のものが事前に必要です。

NVIDIAのGPU

Pascal世代以降のGPUが必要です。すべての機能を使うにはNVIDIA Teslaのリストにあるモデルが必要ですが、試しに触ってみるだけであればPascal世代以降のGPUであれば動作します。

CUDA

これも必要です。ただ、今回はDockerコンテナーでCUDAを動かすので、ホストOSで必要なのはNVIDIA Driverだけです。NVIDIA Container Toolkitのセットアップ手順に従って、GPUドライバーをインストーするか、CUDAのインストールの時にdnf install cudaの代わりにdnf install cuda-driversのようにインストールすると、互換性のあるドライバーがインストールされます（詳細はこちらへ）。

Docker

今回はNVIDIA Container Toolkitを使ってCUDAをコンテナーで動かします。NVIDIA Container Toolkitはかつて「NVIDIA Docker」などと呼ばれていたものの、現在の名称です。名前の通りDockerランタイムをエンジンとして使う技術でしたが、現在はDocker以外でも動作します。

今回はDockerを使います。

オフィシャルのCUDAイメージ

NVIDIAがコンテナでCUDAを使いやすくするため、オフィシャルのCUDAイメージをDocker Hubで公開してくれています。

hub.docker.com

現在はいろいろなOSベースのイメージが提供されています。

Ubuntu 22.04
Ubuntu 20.04
Ubuntu 18.04
Red Hat Universal Base Images (UBI) 8
Red Hat Universal Base Images (UBI) 7
Rocky Linux 8
CentOS 7

どうやって動かすか

さてでは実際にどう動かすかなのですが、次のリポジトリーにDockerfileやそれを使ったカスタムイメージの作成方法、簡単な使い方をまとめています。

github.com

動作確認に使っているクエリーについては、次を参考にしました。

www.cybertec-postgresql.com

ちょっとした検証

実行計画の比較

次のように実行して、GPU側で処理が実行されるか実行計画を確認してみます。

SET pg_strom.enabled = on;
EXPLAIN ANALYZE SELECT count(*)
FROM   t_test1
WHERE sqrt(x) > 0
GROUP BY y;

実行すると次のような結果になるはずです。

                                                                        QUERY PLAN                            
                                             
--------------------------------------------------------------------------------------------------------------
---------------------------------------------
 GroupAggregate  (cost=100031.53..100031.56 rows=1 width=109) (actual time=2009.044..2013.132 rows=1 loops=1)
   Group Key: y
   ->  Sort  (cost=100031.53..100031.53 rows=2 width=109) (actual time=2009.028..2013.115 rows=3 loops=1)
         Sort Key: y
         Sort Method: quicksort  Memory: 25kB
         ->  Gather  (cost=100031.31..100031.52 rows=2 width=109) (actual time=1930.220..2012.977 rows=3 loops
=1)
               Workers Planned: 2
               Workers Launched: 2
               ->  Parallel Custom Scan (GpuPreAgg) on t_test1  (cost=99031.31..99031.32 rows=1 width=109) (ac
tual time=1874.588..1874.596 rows=1 loops=3)
                     Reduction: GroupBy (Global+Local [nrooms: 1974])
                     Group keys: y
                     Outer Scan: t_test1  (cost=2833.33..98814.29 rows=694445 width=101) (actual time=177.739.
.4320.970 rows=5000000 loops=1)
                     Outer Scan Filter: (sqrt((x)::double precision) > '0'::double precision)
 Planning Time: 1.079 ms
 Execution Time: 2071.091 ms
(15 rows)

一行目をoffにすると、CPUで実行する場合の実行計画を確認できます。CPUとGPUの実行計画を比較すると、PostgreSQL標準ではないGpuPreAggを使ってスキャンしていることがわかります。

                                                                      QUERY PLAN                              
                                         
--------------------------------------------------------------------------------------------------------------
-----------------------------------------
 Finalize GroupAggregate  (cost=187989.65..187989.90 rows=1 width=109) (actual time=3287.826..3291.186 rows=1 
loops=1)
   Group Key: y
   ->  Gather Merge  (cost=187989.65..187989.88 rows=2 width=109) (actual time=3287.817..3291.177 rows=3 loops
=1)
         Workers Planned: 2
         Workers Launched: 2
         ->  Sort  (cost=186989.62..186989.63 rows=1 width=109) (actual time=3284.143..3284.144 rows=1 loops=3
)
               Sort Key: y
               Sort Method: quicksort  Memory: 25kB
               Worker 0:  Sort Method: quicksort  Memory: 25kB
               Worker 1:  Sort Method: quicksort  Memory: 25kB
               ->  Partial HashAggregate  (cost=186989.60..186989.61 rows=1 width=109) (actual time=3284.104..
3284.105 rows=1 loops=3)
                     Group Key: y
                     Batches: 1  Memory Usage: 24kB
                     Worker 0:  Batches: 1  Memory Usage: 24kB
                     Worker 1:  Batches: 1  Memory Usage: 24kB
                     ->  Parallel Seq Scan on t_test1  (cost=0.00..183517.38 rows=694445 width=101) (actual ti
me=0.037..1013.665 rows=1666667 loops=3)
                           Filter: (sqrt((x)::double precision) > '0'::double precision)
 Planning Time: 0.109 ms
 Execution Time: 3291.247 ms
(19 rows)

CPUとGPUでの比較

実行時間は\timingをつけて実行しています。クエリ例はちょっと複雑になるので\timing周りは省略しています。このデータはテーブルに含まれる行数をカウントした結果を出力しています。物理的なテーブルサイズはおおよそ1GBです。

この程度だとあまり差は出てこないのですが、それでもGPUで処理を実行した方が速いことが確認できました。

//CPU
testdb2=# SELECT count(*)
testdb2-#   FROM   t_test1
testdb2-#   WHERE sqrt(x) > 0
testdb2-# GROUP BY y;
  count  
---------
 5000000
(1 row)

Time: 3130.320 ms (00:03.130)

//GPU
testdb2=# SELECT count(*)
testdb2-#   FROM   t_test1
testdb2-#   WHERE sqrt(x) > 0
testdb2-# GROUP BY y;
  count  |                                                                                                    
  
---------+----------------------------------------------------------------------------------------------------
--
 5000000 | a                                                                                                  
 
(1 row)

Time: 2957.607 ms (00:02.958)

この結果は物理環境でPG-Stromを動かした時とほぼ一緒の結果になっています。参考までに以前物理マシンでPG-Stromを動かし、もう少し大きいテーブルで同じテストを実行した時の結果を載せておきます。