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

導入

Databricks Runtime(DBR)のバージョン15.1が4月末に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サポートがありました。

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

Step1. 準備

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

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

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

# 保存先の作成
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/という場所へ保管します。

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

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

# 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上でインストールする場合は以下のようになります。

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

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

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

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

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

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

まとめ

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