databricksのdbxを試してみる（インストール〜サンプルプロジェクトのデプロイ） #Databricks

Notebook環境だけではMLOpsの実現は難しい

Databricksにはsparkによる計算基盤に加えnotebook環境、データ基盤、mlflowによる実験管理、モデルリポジトリが用意されており、機械学習モデル構築においては自己完結できる環境と言えます。
作成したNotebookもDatabricks上でスケジュールジョブ登録できるので、開発したnotebookをそのままproductionに回すことも可能ですが、MLOps的には以下の点で推奨されません。

処理がモジュール化されておらず
- 保守性が低い
- リファクタリングした際の単体テストが難しい
- コード再利用が難しい

従ってMLOpsの運用にあたっては、データサイエンティストが開発したnotebookをproduction用にモジュラーなコードに変換し、MLパイプラインとしてデプロイする必要があります。
DatabricksではコードをDatabricksにデプロイするツールとしてdbxが用意されています。

今回はこのdbxで用意されているサンプルプロジェクトの実行を試してみます。
なおdbx自体はまだメジャーバージョンが0で、メジャーアップデートになった際に、大幅な仕様変更が生じる可能性があります。今回の記事ではdbxバージョン0.87を想定しています。

インストール

$ pip install dbx==0.8.7

dbxをインストールすると、databricks-cliもインストールされます。

Databricks-CLIの設定

ターミナル上でdatabricks configure --tokenを実行します。ここでdatabricks hostとtokenを入力します。

$ databricks configure --token
Databricks Host (should begin with https://):

databricks host
databricks hostはdatabricksのワークスペースにログインした際のURLになります。
通常であれば?o=の直前のスラッシュまでがワークスペースURLになります。下記の例でのワークスペースURLはhttps://xxx.net/になります。

URLをコピーしたら、ターミナル上のDatabricks Host:にペーストしてエンターを押しましょう。
token
続いてはtokenの入力が求められます。

Token:

tokenを取得するには、まず「User Settings」にアクセスします。
「Access Token」の画面で「Generate new token」をクリック
利用目的が分かるようにコメントを入力
「Generate」をクリック
tokenが生成されるのでコピーする

ターミナルに戻り、tokenをペーストしエンターを押します。特に結果出力などはありませんが、設定は完了しており、設定ファイル.databrickscfgが作成されています。
Macであれば~/.databrickscfgにあります。
Windowsであれば%USERPROFILE%/.databrickscfgです。
なおプロファイル名は[DEFAULT]になってますが、複数のプロファイルを作成したい場合は

databricks configure --token --profile <profile-name>

とすると[DEFAULT]以外のプロファイル名で接続設定を作成することができます。

サンプルプロジェクトの作成

いよいよサンプルプロジェクトの作成です。ターミナル上で

$ dbx init

を実行します。ここからはいくつか設定事項を質問されます。

project_name [cicd-sample-project]: 
version [0.0.1]: 
Select cloud:
1 - AWS
2 - Azure
3 - Google Cloud
Choose from 1, 2, 3 [1]: 
Select cicd_tool:
1 - GitHub Actions
2 - Azure DevOps
3 - GitLab
4 - None
Choose from 1, 2, 3, 4 [1]: 
project_slug [cicd_sample_project]: 
workspace_dir [/Shared/dbx/cicd_sample_project]: 
artifact_location [dbfs:/Shared/dbx/projects/cicd_sample_project]: 
profile [DEFAULT]:

project_name
- プロジェクトパッケージ名
version
- プロジェクトバージョン
cloud
- databricksが稼働しているクラウド環境。AWS, Azure Google Cloudから選択。デフォルトはAWS。
cicd_tool
- GitHub Actions, Azure DevOps, GitLab, Noneから選択。
project_slug
- URLやディレクトリに使用されるプロジェクト名。
workspace_dir
- databricks上のワークスペースディレクトリ
artifact_location
- アーティファクトの保存先
profile
- databricks configureで設定したプロファイル。基本はDEFAULTで大丈夫ですが、databricks configure -- profile <profile-name>で、任意のプロファイル名を指定している場合はそのプロファイル名を指定します。

今回、私のdatabricks環境はAzureなのでクラウド環境はAzureを選択し、それ以外はデフォルトのままにしておきます。
作成されたサンプルプロジェクトのフォルダ構造は以下のようになります。

cicd-sample-project
├── README.md
├── cicd_sample_project
│   ├── __init__.py
│   ├── common.py
│   └── tasks
│       ├── __init__.py
│       ├── sample_etl_task.py
│       └── sample_ml_task.py
├── conf
│   ├── deployment.yml
│   └── tasks
│       ├── sample_etl_config.yml
│       └── sample_ml_config.yml
├── notebooks
│   └── sample_notebook.py
├── pyproject.toml
├── setup.py
└── tests
    ├── entrypoint.py
    ├── integration
    │   └── e2e_test.py
    └── unit
        ├── conftest.py
        └── sample_test.py

細かい中身については次回以降に触れますが、実行プログラムはtasksディレクトリ内のsample_etl_task.pyとsample_ml_task.pyでこれらのプログラムで参照するパラメータはconf/tasks内にあるyamlファイルで定義しています。

少し中身の探検

dbxは開発したプログラムをdatabricks上のworkflowsとしてデプロイが可能です。workflowsへのデプロイはconf/deployment.ymlで定義します。なおこのyamlファイルのフォーマットはdatabricksのJobs APIに準拠しているようです。

ここではsample_etl_task.pyを単体で実行する場合の定義を見てみます。

deployment.yml

environments:
  default:
    workflows:
      - name: "cicd-sample-project-sample-etl"
        tasks:
          - task_key: "main"
            <<: *basic-static-cluster
            python_wheel_task:
              package_name: "cicd_sample_project"
              entry_point: "etl" # take a look at the setup.py entry_points section for details on how to define an entrypoint
              parameters: ["--conf-file", "file:fuse://conf/tasks/sample_etl_config.yml"]

sample_etl_task.pyで参照するパラメータyamlファイルのパスはparametersで指定されています。
この定義ではプログラムをwheelパッケージとして参照します。entry_pointキーで指定している値"etl"はsetup.pyのentry_pointで定義されている必要があります。

setup.py

setup(
    name="cicd_sample_project",
    packages=find_packages(exclude=["tests", "tests.*"]),
    setup_requires=["setuptools","wheel"],
    install_requires=PACKAGE_REQUIREMENTS,
    extras_require={"local": LOCAL_REQUIREMENTS, "test": TEST_REQUIREMENTS},
    entry_points = {
        "console_scripts": [
            "etl = cicd_sample_project.tasks.sample_etl_task:entrypoint",
            "ml = cicd_sample_project.tasks.sample_ml_task:entrypoint",
    ]},
    version=__version__,
    description="",
    author="",
)

setup.pyではetlのentry pointはsample_etl_task.pyのentrypoint関数を指定しており、deployment.ymlで実行entry_pointにetlを指定することで、sample_etl_task.pyのentrypoint関数が実行されます。

sample_etl_task.py

class SampleETLTask(Task):
    def _write_data(self):
        db = self.conf["output"].get("database", "default")
        table = self.conf["output"]["table"]
        self.logger.info(f"Writing housing dataset to {db}.{table}")
        _data: pd.DataFrame = fetch_california_housing(as_frame=True).frame
        df = self.spark.createDataFrame(_data)
        df.write.format("delta").mode("overwrite").saveAsTable(f"{db}.{table}")
        self.logger.info("Dataset successfully written")

    def launch(self):
        self.logger.info("Launching sample ETL task")
        self._write_data()
        self.logger.info("Sample ETL task finished!")

# if you're using python_wheel_task, you'll need the entrypoint function to be used in setup.py
def entrypoint():  # pragma: no cover
    task = SampleETLTask()
    task.launch()

流れをまとめると以下のとおり。

sample_etl_task.py内でメイン実行関数entrypointを定義
entrypoint関数をsetup.py内でetlという名前でentry_pointsに登録
deployment.ymlはsetup.pyのentry_pointsを参照。etlを指定することでsample_etl_tasklpyのentrypoint関数をworkflowsのジョブに指定。

デプロイ

いよいよdatabricksにデプロイしてみましょう。デプロイコマンドは以下のとおり。

dbx deploy --workflow <workflow-name>

<workflow-name>はdeployment.ymlで定義されているnameキーを指定します。今回は上の例で挙げているcicd-sample-project-sample-etlをデプロイしてみます。

dbx deploy --workflow cicd-sample-project-sample-etl

以下のような出力が出たらデプロイは成功しています。

[dbx][2023-01-10 23:22:35.542] ✨ Deployment for environment default finished successfully

ではdatabricksワークスペースでworklowsのジョブが作成されたか確認してみましょう。
databricksワークスペースの画面左側にあるメニューから「workflows」をクリックしましょう。
「Jobs」タブにdbx deployで指定したのジョブが作成されていれば成功です。

ジョブの実行はこのワークスペースのUI上からも実行できますが、ローカルPCのターミナルからdbx launchコマンドでも実行可能です。

dbx launch <workflow-name>

先ほどデプロイした｀cicd-sample-project-sample-etl`を実行してみましょう。

dbx launch cicd-sample-project-sample-etl

以下のような出力がされればジョブ実行リクエストが正常にサブミットされています。

[dbx][2023-01-10 23:37:26.779] Workflow successfully launched in the non-tracking mode 🚀. Please check Databricks UI for job status 👀

再びdatabricksのワークスペースでworkflows画面でジョブを見ると、ジョブが実行されていることが確認できます。

databricksのdbxを試してみる（インストール〜サンプルプロジェクトのデプロイ）