この記事は備忘録であり、Snowpark Container Servicesの最適な設定を示したものではありません。

Snowpark Container Servicesとは

Snowpark Container ServicesはSnowflakeのインフラ上でKubernetesクラスタを構築できる機能です。
略してSPCSと呼ばれることもあります。

参考になるドキュメントや記事など

Snowflake Documentation - サービス仕様リファレンス

Snowflakeの公式ドキュメントです。GPUを使うにはここに記載されている通りにcontainer.resources.requiresおよびcontainer.resources.limitsを指定する必要があります。

nvidia.com/gpu: GPUs が必要な場合は、それらをリクエストしなければならず、同じ数量に対して limit も指定する必要があります。コンテナーが GPU の容量に対するリクエストと制限を指定しない場合は、 GPUs にアクセスできません。

Snowflake Quickstarts - Fine-Tuning an LLM in Snowpark Container Services with AutoTrain

QuickstartsというSnowflakeが公式で出しているハンズオンコンテンツの一つです。
RAPIDSとAutoTrainを使っています。確実にGPUを使えるコンテナの作り方が分からなかったので参考にしました。
このハンズオンではCUDA 11.8をベースイメージとして指定しています。

Leveraging Snowpark Container Services for Advanced Q&A Retrieval: A Journey into GPU-Enhanced Semantic Search

Mediumに掲載された記事です。上述のハンズオンよりシンプルで分かりやすいです。
プレーンなPython 3.10にSentence TransformersとPyTorchをインストールしただけのものをベースイメージにしています。おそらくPyTorchがあればGPUを使えるのかな？
GitHubでソースコードが公開されています。
個人的には、Docker Composeの使い方が参考になりました。これまで逐一手動でSnowflakeのイメージリポジトリ用にタグ付けしていました。

Multi-GPU PyTorch Training in Snowflake

Mediumに掲載された記事です。torchrunでマルチGPUを使ってトレーニングする流れを紹介しています。
汎用的で実用的。PyTorchを使った機械学習に慣れている方ならこの記事を眺めると「なんだ普通に使えるじゃん」となるんじゃないでしょうか？

コンピュートプールの作り方

コンピュートプールを作成するとき、INSTANCE_FAMILYパラメータにGPU_NV_Sか
GPU_NV_Mを指定するだけです。
GPU_NV_Lは英語版の公式ドキュメントによるとOn requestということです。そもそもクレジット消費量が怖くて使えないですが…

サービスの作り方

サービス仕様を定義するYAMLのcontainer.resources.requiresとcontainer.resources.limitsにnvidia.com/gpuを指定する必要があります。また、それらの個数は同じ値を指定しないといけないそうです。以下が記載例です。

spec:
  containers:
  - name: resource-test-gpu
    image: ...
    resources:
      requests:
        nvidia.com/gpu: 1
      limits:
        nvidia.com/gpu: 1

コンテナの作り方

Jupyter Notebookを使うとき

RAPIDSをベースイメージにしてPyTorchとJupyter Notebookをインストールするだけで問題なく動きました。
他の記事を確認した感じではPyTorchをインストールすればなんでもよさそうです。

torchrunを使うとき

こちらについては未検証です。
ですが、前述したMediumの記事の通りにやれば簡単にできそうな雰囲気はあります。
記事執筆時点ではまだPrivate Preview (PrPr)なのですが、ジョブの実行を見据えて検証しておきたいですね。