Databricks Container Servicesを使って Runtime環境をカスタムする

はじめに

Databricksではクラスター作成時にRuntimeを選択することで、さまざまなライブラリーやバージョンの入ったRuntimeを選ぶことが出来ます。Runtimeに含まれるライブラリーやバージョンはこちらから確認出来ます。

また追加のライブラリーや設定ファイルなどは、Databricksライブラリーサービスやinitスクリプトを利用することでクラスター起動時に追加や設定が可能です。

• Databricksでライブラリをインストールする (@taka_yayoi)
• 自作ライブラリをインストールする - Databricks (by tktsqさん)
• Databricksで python モジュールをimportするには？
• Databricksクラスターノード初期化スクリプト (@taka_yayoi)

基本的にはベースのライブラリーには新機能への対応やBug fixなどのアップデートもあるため上記の方法で運用することがベストだと思います。

しかし毎回のクラスター起動時に大量のライブラリーなどをインストールしていては起動時間が遅くなってしまうのも事実です。またバージョンをロックダウンしたいケースなどもあるかと思います。

そこで、DatabricksではRuntime環境をカスタマイズして保存しておくことができます。これがDatabricks Container Serviceとなります。Databricksのクラスターインスタンス上ではDocker　Containerが動作するため、それを予め作成しておくことで起動時に読み込む事ができます。

今回はこのContainer Serviceを試してみたいと思います。

正式な設定方法については、こちらをご覧ください
https://learn.microsoft.com/ja-jp/azure/databricks/clusters/custom-containers

準備

• ローカルPCにDockerをインストールしておく
• クラスター環境からアクセスできるDockerレジストリを用意しておく（Amazon ECRや、 Azure Container Registry、Docker Hubなど）
• Databricks Container Serviceを有効にする。設定方法

コンテナ Baseイメージについて

コンテナのベース イメージは、https://hub.docker.com/u/databricksruntime の Docker Hub でホストされています。 これらのベースの生成に使用される Dockerfile は https://github.com/databricks/containers にありますのでベースの中身が知りたい方はご確認ください。

Standardのベースイメージとしては各LTSのバージョンが利用出来ますが、MLバージョンは用意されておりませんので、各自でBuildする必要があります。また、ベースイメージを用いずに１から作成することも可能です。詳しく知りたい方はこちらをご覧ください。（今回はベースイメージを利用します）

GPUをご利用の方は、GPUのライブラリーが含まれたベースイメージもあります。こちらの利用に関してはマニュアルをご覧ください。これらのベースの生成に使用される Dockerfile は https://github.com/databricks/containers にあります

Container Serviceの利用手順

1. ベースイメージを利用したDockerfileの作成

それでは、カスタムするためのDockerfileを作成していきます。ローカルPC上の任意のDirectory上にDockerfileを作成します。以下のようにベースイメージと追加するライブラリーなどを記述します。Dockerfileの記述方法は他に多数の解説記事があるのでそちらをご覧ください。（参考サイト)

Dockerfile

FROM databricksruntime/standard:11.3-LTS
RUN /databricks/python3/bin/pip install mlflow
RUN /databricks/python3/bin/pip install tensorflow

mlflow, tensorflowを入れたらサイズが3GBになってしまったので、試すのであればもっと別のライブラリーにした方が良かった。。

2. Build

作成したDockerfileを元に新規イメージをBuildします。

docker build --pull --rm -f "Dockerfile" -t standarddemo "." 

VSCodeに Docker Extensionが入っている場合は、右クリックでBuildできます。

3. Docker　Registoryにログイン

今回はDocker HUBを利用します。

docker login

VSCodeにDocker Extensionが入っている場合は、Registories のコンセント見たいなマークから接続できます。

4. イメージをPush

Docker HUBにイメージをプッシュします。tagはわかりやすいように今回は11.3-LTSとしておきます。

docker push <myregistry>/standarddemo:11.3-LTS 

VSCodeにDocker Extensionが入っている場合は、ImageをクリックしてPUSHで、namespaceやtagを選択して実行出来ます。

5. Databricks ClusterでContainerイメージを指定します。

新規クラスター作成にて、下部にある「高度なオプション」内の「Docker」タブを開き、「自分のDockerコンテナを使用する」にチェックを入れます。DockerイメージURLと認証（Docker HUBの場合はPublicのため認証なしなため、デフォルトで大丈夫です）。また<myregistry>箇所はPushしたレジストリー名に変更してください。

起動してpythonライブラリーを確認すると、mlflowやtensorflowが含まれているのが確認出来ました。

注意点

Containerイメージは自分でメンテナンスしないといけないため、ライブラリーのパッチやバージョンアップデートをセルフで行う必要があります。
またBaseイメージだけではいくつかのDatabricks機能が利用できませんでした。

• Gangulia
• クラスターのWeb Terminal
• MLflow Auto Logging
  など

Databricksの機能や新機能を利用したい場合は素直にDatabricksが用意しているRuntimeを使用しましょう。