Snowpark Container ServicesでGPUを使うためのヒント

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

## Snowpark Container Servicesとは

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

https://docs.snowflake.com/ja/developer-guide/snowpark-container-services/overview  

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

### [Snowflake Documentation - サービス仕様リファレンス](https://docs.snowflake.com/ja/developer-guide/snowpark-container-services/specification-reference)  

https://docs.snowflake.com/ja/developer-guide/snowpark-container-services/specification-reference  

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](https://quickstarts.snowflake.com/guide/fine-tune-an-llm-in-snowpark-container-services-with-autotrain/index.html#0)  

https://quickstarts.snowflake.com/guide/fine-tune-an-llm-in-snowpark-container-services-with-autotrain/index.html#0  

QuickstartsというSnowflakeが公式で出しているハンズオンコンテンツの一つです。  
[RAPIDS](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/rapidsai/containers/rapidsai)とAutoTrainを使っています。確実にGPUを使えるコンテナの作り方が分からなかったので参考にしました。
このハンズオンではCUDA 11.8をベースイメージとして指定しています。  

### [Leveraging Snowpark Container Services for Advanced Q&A Retrieval: A Journey into GPU-Enhanced Semantic Search](https://medium.com/snowflake/leveraging-snowpark-container-services-for-advanced-q-a-retrieval-5bf5d2686fa6)  

https://medium.com/snowflake/leveraging-snowpark-container-services-for-advanced-q-a-retrieval-5bf5d2686fa6

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

### [Multi-GPU PyTorch Training in Snowflake](https://medium.com/@sikha.das_38205/multi-gpu-pytorch-training-in-snowflake-9ef1681867b5)  

https://medium.com/@sikha.das_38205/multi-gpu-pytorch-training-in-snowflake-9ef1681867b5

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

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

コンピュートプールを作成するとき、```INSTANCE_FAMILY```パラメータに```GPU_NV_S```か
```GPU_NV_M```を指定するだけです。  
```GPU_NV_L```は[英語版の公式ドキュメント](https://docs.snowflake.com/en/developer-guide/snowpark-container-services/working-with-compute-pool)によるとOn requestということです。そもそもクレジット消費量が怖くて使えないですが…  

https://docs.snowflake.com/ja/sql-reference/sql/create-compute-pool  

https://docs.snowflake.com/en/developer-guide/snowpark-container-services/working-with-compute-pool  

## サービスの作り方

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

```yaml
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をインストールすればなんでもよさそうです。  

![Jupyter Notebook with RAPIDS](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2917580/f50baa0f-cdf5-313e-df51-aa3b3388dba2.png)

### torchrunを使うとき

こちらについては未検証です。  
ですが、前述したMediumの記事の通りにやれば簡単にできそうな雰囲気はあります。  
記事執筆時点ではまだPrivate Preview (PrPr)なのですが、[ジョブの実行](https://docs.snowflake.com/ja/developer-guide/snowpark-container-services/working-with-jobs)を見据えて検証しておきたいですね。