Airflowによるデータパイプラインの依存関係の管理

[Managing dependencies in data pipelines \| Databricks on AWS](https://docs.databricks.com/dev-tools/data-pipelines.html) [2021/3/17時点]の翻訳です。

データパイプラインに複雑な依存関係があることは頻繁に起こりえます。ワークフローシステムを利用することで、そのような依存関係を記述したり、パイプラインの実行をスケジュールすることができます。

# Apache Airflow

[Apache Airflow](https://airflow.apache.org/)はデータパイプラインを管理、スケジュールするためのソリューションです。Airflowはデータパイプラインを、オペレーションの有向非巡回グラフ(DAG)で表現します。ここでは、オペレーション間の論理的な依存関係をエッジで表現します。

AirflowはDatabricks、Airflow間の密なインテグレーションを提供します。Airflow-Databricksインテグレーションによって、Databricksによって提供される最適化されたSparkエンジンとAirflowのスケジューリング機能を活用することができます。

## Airflow-Databricksインテグレーションのインストール

Airflow-DatabricksインテグレーションはAirflowバージョン1.9.0で利用することができます。Airflow-Databricksインテグレーションをインストールするには、以下を実行します。

```bash:Bash
pip install "apache-airflow[databricks]"
```

[extras](https://airflow.apache.org/docs/apache-airflow/stable/installation.html?highlight=extras#airflow-extra-dependencies)(`celery`、`s3`、`password`など)をインストールするには、以下を実行します。

```bash:Bash
pip install "apache-airflow[databricks, celery, s3, password]"
```

## `DatabricksRunNowOperator`オペレーター

Airflow-Databricksインテグレーションは、処理のDAGにおけるノードとして[DatabricksRunNowOperator](https://airflow.apache.org/docs/apache-airflow-providers-databricks/stable/_api/airflow/providers/databricks/operators/databricks/index.html#airflow.providers.databricks.operators.databricks.DatabricksRunNowOperator)を提供します。このオペレーターは、Databricksジョブの[Run now](https://docs.databricks.com/dev-tools/api/latest/jobs.html#jobsjobsservicerunnow) APIエンドポイントに適合し、プログラムからノートブックやDBFSやS3にアップロードされたJARを実行することが可能になります。

## `DatabricksSubmitRunOperator`オペレーター

Airflow-Databricksインテグレーションは、処理のDAGにおけるノードとして[DatabricksSubmitRunOperator](https://airflow.apache.org/docs/apache-airflow-providers-databricks/stable/_api/airflow/providers/databricks/operators/databricks/index.html#airflow.providers.databricks.operators.databricks.DatabricksSubmitRunOperator)を提供します。このオペレーターは、Databricksジョブの[Runs submit](https://docs.databricks.com/dev-tools/api/latest/jobs.html#jobsjobsservicesubmitrun) APIエンドポイントに適合し、プログラムからノートブックやDBFSやS3にアップロードされたJARを実行することが可能になります。

## Databricks接続の設定

`DatabricksSubmitRunOperator`を使用するには、適切なAirflow接続設定に認証情報を提供する必要があります。`DatabricksSubmitRunOperator `で`databricks_conn_id`を指定しない場合、デフォルトでは、オペレーターは`databricks_default`という名前の接続設定から認証情報を探そうとします。

[Managing Connections](https://airflow.apache.org/docs/apache-airflow/stable/howto/connection.html)に記載されているように、Airflowのweb UIからAirflowの接続設定を行うこともできます。Databricksとの接続に関しては、**Login**フィールドを`token`にし、**Extra**フィールドに以下を設定します。

```json:JSON
{"token": "<personal access token>", "host":"<Databricks hostname>"}
```

ここで、`<personal access token>`は、Databricksで作成した[personal access token](https://docs.databricks.com/dev-tools/api/latest/authentication.html)であり、`<Databricks hostname>`はDatabricksがデプロイされているホスト名となります。

## サンプル

このサンプルでは、ローカルのマシンで実行されるシンプルなAirflowデプロイメントをどのように設定し、Databricksで実行されるように指定されたサンプルDAGをどのようにデプロイするのかを説明します。

### Airflowデータベースの初期化

Airflowが様々なメタデータを追跡するのに用いるSQLiteデータベースを初期化します。プロダクションのAirflowのデプロイメントにおいては、標準的なデータベースを用いてAirflowを設定するかもしれません。初期化を行うには以下を実行します。

```bash:Bash
airflow initdb
```

Airflowデプロイメントに対するデフォルト設定とSQLiteデータベースが`~/airflow`に初期化されます。

### DAG定義

DAG定義はPythonファイルであり、この例では`example_databricks_operator.py`という名前となります。この例では、一つの線形の依存関係を持つ二つのDatabricksジョブを実行します。最初のDatabricksジョブは`/Users/airflow@example.com/PrepareData`に配置されたノートブックを起動し、二つ目のジョブは`dbfs:/lib/etl-0.1.jar`に配置されているJARを実行します。サンプルのDAG定義は二つの`DatabricksSubmitRunOperator`タスクから構成され、最後に`set_downstream`メソッドを用いて依存関係を設定します。コードのスケルトンバージョンは以下のようなものになります。

```py:Python
notebook_task = DatabricksSubmitRunOperator(
    task_id='notebook_task',
    dag=dag,
    json=notebook_task_params)

spark_jar_task = DatabricksSubmitRunOperator(
    task_id='spark_jar_task',
    dag=dag,
    json=spark_jar_task_params)

notebook_task.set_downstream(spark_jar_task)
```

#### Airflowと必要なクラスのインポート

DAG定義のトップレベルでは、`airflow`、`DAG`、`DatabricksSubmitRunOperator`をインポートします。

```py:Python
import airflow

from airflow import DAG
from airflow.contrib.operators.databricks_operator import DatabricksSubmitRunOperator
```

#### グローバル引数の設定

次のセクションでは、DAGにおけるそれぞれのタスクに適用されるデフォルトの引数を設定します。

```py:Python
args = {
    'owner': 'airflow',
    'email': ['airflow@example.com'],
    'depends_on_past': False,
    'start_date': airflow.utils.dates.days_ago(0)
}
```

ここで興味深い引数が二つあります。`depends_on_past`と`start_date`です。`depends_on_past`を`true`に設定することで、以前のタスクが成功しない限り、このタスクは実行されなくなります。`start_date`引数は、最初のタスクをいつ実行するのかをスケジュールします。

#### DAGインスタンスの生成

DAGインスタンス作成のステートメントによって、DAGにユニークなIDを与え、デフォルトの引数を適用し、日次のスケジュールを設定します。

```py:Python
dag = DAG(dag_id='example_databricks_operator', default_args=args, schedule_interval='@daily')
```

次のステートメントでは、Sparkバージョン、ノードのタイプ、タスク実行に必要なクラスターのワーカー数を指定します。指定内容のスキーマはジョブの[Runs Submit](https://docs.databricks.com/dev-tools/api/latest/jobs.html#jobsjobsservicesubmitrun)エンドポイントの`new_cluster`フィールドに合致します。

```py:Python
new_cluster = {
    'spark_version': '6.0.x-scala2.11',
    'node_type_id': 'i3.xlarge',
    'aws_attributes': {
        'availability': 'ON_DEMAND'
    },
    'num_workers': 8
}
```

#### タスクをDAGに登録

`notebook_task`に対しては、`DatabricksSubmitRunOperator`インスタンスを作成します。

```py:Python
notebook_task_params = {
    'new_cluster': new_cluster,
    'notebook_task': {
    'notebook_path': '/Users/airflow@example.com/PrepareData',
  },
}
# Example of using the JSON parameter to initialize the operator.
notebook_task = DatabricksSubmitRunOperator(
  task_id='notebook_task',
  dag=dag,
  json=notebook_task_params)
```

このコードでは、JSONパラメーターとして`Runs Submit`エンドポイントに適合するPythonのディクショナリーを与えます。

`dbfs:/lib/etl-0.1.jar`に配置されているJARファイルを実行する`spark_jar_task`に対しては、`DatabricksSubmitRunOperator`のインスタンスを生成します。

```py:Python
# Example of using the named parameters of DatabricksSubmitRunOperator to initialize the operator.
spark_jar_task = DatabricksSubmitRunOperator(
  task_id='spark_jar_task',
  dag=dag,
  new_cluster=new_cluster,
  spark_jar_task={
    'main_class_name': 'com.example.ProcessData'
  },
  libraries=[
    {
      'jar': 'dbfs:/lib/etl-0.1.jar'
    }
  ]
)
```

後段の処理を実行するために`spark_jar_task`を設定するには、依存関係を登録するために`notebook_task`で`set_downstream`メソッドを使用します。

```py:Python
notebook_task.set_downstream(spark_jar_task)
```

ここで使用している`notebook_task`では、submit runエンドポイントに対するすべての仕様を`json`パラメーターで指定し、`spark_jar_task`では`DatabricksSubmitRunOperator`に対するsubmit runエンドポイントのトップレベルのキーをパラメーターとして直接記載したことに注意してください。オペレーターのインスタンスを生成する両方の方法は等価ですが、後者の方法では`spark_python_task`、`spark_submit_task`のように新たなトップレベルのフィールドを追加することはできません。詳細は[DatabricksSubmitRunOperator API](https://airflow.apache.org/docs/apache-airflow-providers-databricks/stable/operators.html)を参照ください。

### AirflowへのDAGのインストール、検証

AirflowにDAGをインストールするには、ディレクトリ`~/airflow/dags`を作成し、ディレクトリにDAG定義をコピーします。

AirflowがDAGを読み込んだことを検証するには、`list_dags`コマンドを実行します。

```bash:Bash
airflow list_dags

[2017-07-06 10:27:23,868] {__init__.py:57} INFO - Using executor SequentialExecutor
[2017-07-06 10:27:24,238] {models.py:168} INFO - Filling up the DagBag from /Users/<user>/airflow/dags

-------------------------------------------------------------------
DAGS
-------------------------------------------------------------------
example_bash_operator
example_branch_dop_operator_v3
example_branch_operator
example_databricks_operator
...
```

> **訳者注**
Airflow 2.1.0では`airflow dags list`というコマンドになります。

### Airflow UIでのDAGの可視化

Airflowのweb UIでDAGを可視化することができます。`airflow webserver`を実行し、`localhost:8080`に接続します。`example_databricks_operator`をクリックしDAGを可視化します。以下が例となります。
![](https://docs.databricks.com/_images/airflow-tutorial.png)

### Airflowへの接続設定

Databricksに対する接続の認証情報はDAG定義では指定されていません。デフォルトでは、`DatabricksSubmitRunOperator`は`databricks_conn_id`パラメーターを`databricks_default`に設定します。[Databricks接続の設定](#databricks接続の設定)に記載されているweb UIで、IDが`databricks_default`である接続設定を確認します。

### それぞれのタスクのテスト

`notebook_task`をテストするには、`airflow test example_databricks_operator notebook_task <YYYY-MM-DD>`を実行し、`spark_jar_task`に対しては、`airflow test example_databricks_operator spark_jar_task <YYYY-MM-DD>`を実行します。DAGのスケジュール処理を実行する場合には、`airflow scheduler`コマンドでスケジューラーデーモンプロセスを起動します。

スケジューラを開始した後は、web UIでDAGで処理待ちのランを参照することが可能になります。

> **訳者注**
Airflow 2.1.0では`airflow tasks test example_databricks_operator notebook_task 2021-07-02`というコマンドになります。

### Databricks 無料トライアル

[Databricks 無料トライアル](https://databricks.com/jp/try-databricks)