Databricksランタイム15.xで対応されたクラスタライブラリのrequirements.txtサポートを試す

# 導入

Databricks Runtime(DBR)のバージョン15.1が4月末にGAになっていました。

https://docs.databricks.com/ja/release-notes/product/2024/april.html#databricks-runtime-151-is-ga

15.1は15.0の破壊的変更(pandasのバージョン問題)をリカバーするためにリリースされたもののようで、合わせて15.0のサポートは2024/5月末に終了となります。
※ 私も上記のpandasバージョン問題で利用を躊躇っていたのですが、これで移行ができそうです。

DBRが15.x系にアップグレードするにあたって、他にも様々な変更があります。
個人的に大きなもので言うと、
- Pythonのバージョンが3.10 -> 3.11にアップグレード
- GPUを使うMLランタイムににおいてCUDAのバージョンが11.8 -> 12.1へアップグレード

などです。

他にも様々なものがあるのですが、その中で地味にありがたいクラスタライブラリのrequirements.txtサポートがありました。

https://docs.databricks.com/ja/release-notes/runtime/15.1.html#cluster-libraries-now-support-requirementstxt-files

こちらの機能を少し試してみます。

# Step1. 準備

試験用のrequirements.txtを準備します。

今回は、CUDA 12.1対応のpytorch 2.3.0をローカルのWheelファイルからインストールするrequirements.txtを作成します。
余談ですが、2024/5月上旬時点でベータとして公開されているDBR15.2は[pytorch 2.2.2](https://docs.databricks.com/ja/release-notes/runtime/15.2ml.html)が提供されるようです。(15.1は2.1.2)

ノートブックを作成し、必要なwheelファイルをダウンロードする処理を記述・実行。
※ torchvisonなどのインストールを割愛しているので注意。

```python
# 保存先の作成
dbutils.fs.mkdirs("file:/tmp/torch/whl/")

!wget https://download.pytorch.org/whl/cu121/torch-2.3.0%2Bcu121-cp311-cp311-linux_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_nvjitlink_cu12-12.1.105-py3-none-manylinux1_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/cu121/nvidia_nccl_cu12-2.20.5-py3-none-manylinux2014_x86_64.whl -P /tmp/torch/whl

!wget https://download.pytorch.org/whl/triton-2.3.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl -P /tmp/torch/whl
```

ローカルにダウンロードしたwhlファイルをUnity Catalog Volumesにコピーします。
今回はtrainingカタログのllmスキーマにpackagesというボリュームを事前に作成しておき、```/Volumes/training/llm/packages/torch/```という場所へ保管します。

```python
dbutils.fs.cp("file:/tmp/torch/whl/", "/Volumes/training/llm/packages/torch/", recurse=True)
```

保管したwhlファイルの一覧から```requirementx.txt```を作成します。

```python
# whlファイルの一覧を取得

whl_files = dbutils.fs.ls("/Volumes/training/llm/packages/torch/")
whl_files = [f.path.lstrip("dbfs:") for f in whl_files if f.path.endswith(".whl")]
whl_files_str = "\n".join(whl_files)

# requirements.txtを作成
dbutils.fs.put("/Volumes/training/llm/packages/torch/requirements.txt", whl_files_str, overwrite=True)

print(whl_files_str)
```

これで準備完了です。

ちなみに、作成したrequirements.txtを使ってNotebook上でインストールする場合は以下のようになります。

```python
%pip install -U -r /Volumes/training/llm/packages/torch/requirements.txt
```

# Step2. クラスタライブラリの設定

作成したrequirements.txtをクラスタライブラリに設定します。

クラスタメニューから、クラスタライブラリを設定したクラスタを開き、「ライブラリ」タブを開きます。対象クラスタのDBRは15.1以降としてください。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/59d1ac91-e3d6-08fd-fe50-d125fbc61c20.png)

「新規をインストール」ボタンを押して、ライブラリをインストール。
ボリュームを選択し、requirements.txtのパスを指定してください。
※　メッセージにもrequirements.txtを選択できる旨記載されていますね。DBR14.3以前は表示されていませんでした。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/0a012e43-8610-71f9-1152-4dedfd1df4d9.png)

これでクラスタが起動すると、合わせてrequirements.txtを使ってパッケージ類がインストールされます。
クラスタが利用可能になるまでの待ち時間は伸びますが、開発/実行環境の標準化には便利になると思います。
(よりちゃんとやる場合はコンテナイメージを作るべきだと思いますが、一部パッケージのインストールだけの利用であればこれで十分かと思います)

# まとめ

DBR 15.xから対応されたクラスタライブラリのrequirements.txtサポートを試してみました。
ちょっとした改善かもしれませんが、痒いところに手が届く感じでこじんてきにありがたいです。