Python: Luigiでデータパイプラインを作る

バッチ処理の実装にLuigiを使う機会があり、少し調べて整理しました。
irisデータセットをダウンロードしてきて、scikit-learnで学習したモデルをS3にアップロードする、簡単なサンプルも作ります。

Luigi

Luigiは、データパイプラインを記述するためのPythonフレームワークです。

github.com

特徴

タスク間の依存関係を定型的に定義できる
失敗した(アウトプットが生成できなかった)タスクから再開できる
スケジューリングはできない
- スタートキックは別の誰かが行う必要がある

ユースケース

処理間に依存関係があり、各処理が複雑な場合に役立ちます。

例えば、データの前処理やパラメータ最適化などを含む複雑な処理フローを構造化して実装したい
例えば、インプットとなるデータ量が膨大なので、適切に中間生成物を作りながら処理したい

一方で、ジョブのスケジューリングは別の機構(cronやAWS Data PipelineやAirFlowなど)が必要となります。

そうしたジョブスケジューラから呼び出され、依存関係があるまとまった処理を行うバッチ、という位置づけが良さそうです。

基本的な使い方

はじめにluigiをインストールします。

pip install luigi

最初に簡単な2つのタスクを定義して、動きを見ていきます。

Task1: intermediate/task1.txtにファイルを出力
Task2: intermediate/task2.txt(Task1の出力)を読み込み、output/task2.txtにファイル出力

フォルダ構造は以下のとおりです。main.pyが今回実装するプログラムです。python main.pyで実行できます。

.
├── Intermediate
├── main.py
└── output

main.pyの実装を示します。

import luigi
import time


class Task1(luigi.Task):
    task_namespace = 'tasks'

    def run(self):
        print("Task1: run")
        with self.output().open("w") as target:
            target.write(f"This file was generated by Task1 at {time.asctime()}.")

    def output(self):
        print("Task1: output")
        return luigi.LocalTarget("intermediate/task1.txt")


class Task2(luigi.Task):
    task_namespace = 'tasks'

    def requires(self):
        print("Task2: requires")
        return Task1()

    def run(self):
        print("Task2: run")
        with self.input().open("r") as intermediate, self.output().open("w") as target:
            task1_text = intermediate.read()

            target.write(f"{task1_text}\n")
            target.write(f"This file was generated by Task2 at {time.asctime()}.")

    def output(self):
        print("Task2: output")
        return luigi.LocalTarget("output/task2.txt")


if __name__ == '__main__':
    luigi.run(['tasks.Task2', '--workers', '1', '--local-scheduler'])

Luigiでは、入力・処理・出力の一連の処理をTaskクラス(あるいはその継承クラス)を継承することで定義します。
Taskクラスでは、requires、output、runの3つのメソッドを主に実装します。

requiresメソッドでは、前段のタスクを記述し、依存関係を定義する
outputメソッドでは、出力先を定義する
runメソッドでは、出力を生成する処理を記述する
- outputメソッドがFalse(存在しない)の場合のみ実行される

Task1では、依存するタスクが無いため、outputとrunのみを定義しています。

outputは、luigi.LocalTarget("intermediate/task1.txt")でintermediate/task1.txtファイルを出力先に指定
- LocalTarget以外にも、luigi.contrib.s3.S3Targetやluigi.contrib.bigquery.BigqueryTargetなどが提供されています (参考)
runは、出力先をouputメソッドで取得し、文字列(This file was generated by Task1 at Sat Mar 31 13:47:18 2018.)を書き込んでいます

Task2では、requires、output、putを定義してます。

requiresでは、依存するTask1を生成して返すことで、先にTask1を実行させるようにします
runでは、Task1の出力をinputメソッドで取得しています

上を実行すると、task1.txtとtask2.txtが生成されます。

.
├── Intermediate
│   └── task1.txt
├── input
├── main.py
└── output
    └── task2.txt

コンソールの表示からは以下の順番で実行されていることがわかります。また、Luigi Execution Summaryからは2つのタスクが正常に完了したことも確認できます。

DEBUG: Checking if tasks.Task2() is complete
Task2: output
Task2: requires
DEBUG: Checking if tasks.Task1() is complete
Task1: output
Task1: run
Task1: output
Task2: requires
Task1: output
Task2: run
Task2: requires
Task1: output
Task2: output
...
===== Luigi Execution Summary =====

Scheduled 2 tasks of which:
* 2 ran successfully:
    - 1 tasks.Task1()
    - 1 tasks.Task2()

This progress looks :) because there were no failed tasks or missing external dependencies

===== Luigi Execution Summary =====

ここで、Task2に問題がありtask2.txtのみが出力されなかった場合を想定して、task2.txtのみを削除します。main.pyを再実行すると、Task2のみが実行されます。

===== Luigi Execution Summary =====

Scheduled 2 tasks of which:
* 1 present dependencies were encountered:
    - 1 tasks.Task1()
* 1 ran successfully:
    - 1 tasks.Task2()

This progress looks :) because there were no failed tasks or missing external dependencies

===== Luigi Execution Summary =====

一方で、task1.txtのみを削除して再実行しても、task2.txtは存在するため何も行われません(最初から再作成されたりしません)。

===== Luigi Execution Summary =====

Scheduled 1 tasks of which:
* 1 present dependencies were encountered:
    - 1 tasks.Task2()

Did not run any tasks
This progress looks :) because there were no failed tasks or missing external dependencies

===== Luigi Execution Summary =====

機械学習やS3と組み合わせたサンプル

応用例として、Irisデータセットを使って機械学習した結果をS3にアップロードする3つのタスクを定義します。

iris_tasks.DownloadCsvTask: Irisデータセット(CSVファイル)をダウンロードして、intermediate/original.csvに出力する
iris_tasks.CreateModelTask: original.csvを読み込み、SVMで機械学習してモデルを作り、intermediate/model.pklに出力する
iris_tasks.UploadS3Task: model.pklを読み込み、S3にアップロードする

設定ファイルはフォルダ構成は以下のとおりです。

.
├── intermediate
├── luigi.cfg
├── main.py
└── output

新たにluigi.cfgという設定ファイルを追加しています。
Luigiでは、実行時のカレントパスに"luigi.cfg"を配置するか、または、"LUIGI_CONFIG_PATH"でファイルパスを指定することで、実行時に参照できるパラメータを設定できます。

luigi.cfgは以下のように記述しています。
タスクの名前(例えば[iris_tasks.DownloadCsvTask])を指定し、パラメータ(iris_tasks.DownloadCsvTaskの場合は、source_urlとintermediate_original_path)を渡します。
[s3]はS3のアップロードで必要な認証情報を渡しています。このあと触れるS3Targetでは、[s3]で渡された認証情報を使ってS3にアクセスします。

[iris_tasks.DownloadCsvTask]
source_url=https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data
intermediate_original_path=intermediate/original.csv

[iris_tasks.CreateModelTask]
intermediate_model_path=intermediate/model.pkl

[iris_tasks.UploadS3Task]
s3_model_path=s3://test-s3-path/model.pkl

[s3]
aws_access_key_id=XXXXXXXXXXXXXXXXXXXX
aws_secret_access_key=XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

それでは実装を示します。こちらもpython main.pyで実行できます。

import luigi
from luigi.contrib.s3 import S3Target # pip install boto
import requests
import pandas
import pickle
import shutil
from sklearn import svm


# irisデータをダウンロードしてCSVファイルに出力
class DownloadCsvTask(luigi.Task):
    task_namespace = 'iris_tasks'

    source_url = luigi.Parameter()
    intermediate_original_path = luigi.Parameter()

    def run(self):
        with self.output().open("w") as target:
            response = requests.get(self.source_url)
            target.write(response.text)

    def output(self):
        return luigi.LocalTarget(self.intermediate_original_path)


# irisの分類モデル(SVM)を作り、pickleでシリアライズしてファイル出力
class CreateModelTask(luigi.Task):
    task_namespace = 'iris_tasks'

    intermediate_model_path = luigi.Parameter()

    def requires(self):
        return DownloadCsvTask()

    def run(self):
        with open(self.output().path, "wb") as target:
            # self.output().open("wb") ではエラーが発生した
            # https://github.com/spotify/luigi/issues/1647

            data = pandas.read_csv(self.input().path, header=None)
            x = data.iloc[:, [0, 1, 2, 3]]
            y = data.iloc[:, 4].replace({'Iris-setosa': 0, 'Iris-versicolor': 1, 'Iris-virginica': 2})

            model = svm.SVC()
            model.fit(x, y)

            pickle.dump(model, target)

    def output(self):
        # バイナリファイルを出力する場合は、formatにNopFormatを指定
        return luigi.LocalTarget(self.intermediate_model_path, format=luigi.format.NopFormat)


# pklファイルをS3のバケットにアップロード
class UploadS3Task(luigi.Task):
    task_namespace = 'iris_tasks'

    s3_model_path = luigi.Parameter()

    def requires(self):
        return CreateModelTask()

    def run(self):
        with self.output().open("w") as target:
            # 存在するファイルをアップロードする場合は、tmp_pathにファイルをコピー
            shutil.copy(self.input().path, target.tmp_path)
        
    def output(self):
        # S3Targetでアップロードできる(botoのインストールが必要)
        return S3Target(self.s3_model_path)


if __name__ == '__main__':
    luigi.run(['iris_tasks.UploadS3Task', '--workers', '1', '--local-scheduler'])

最初のDownloadCsvTaskでは、irisのデータをダウンロードしてCSVファイルでローカルに保存します。

luigi.cfgで渡されたパラメータですが、変数名 = luigi.Parameter()とすることで、変数名と一致するパラメータの値が得られます。ここでは、[iris_tasks.DownloadCsvTask]からsource_urlとintermediate_original_pathの値を取り出しています。

class DownloadCsvTask(luigi.Task):
    task_namespace = 'iris_tasks'

    source_url = luigi.Parameter()
    intermediate_original_path = luigi.Parameter()

    def run(self):
        with self.output().open("w") as target:
            response = requests.get(self.source_url)
            target.write(response.text)

    def output(self):
        return luigi.LocalTarget(self.intermediate_original_path)

2番目にCreateModelTaskが実行されます。前段のDownloadCsvTaskからoriginal.csvを受け取り、分類モデル(SVM)を作り、pickleでシリアライズしてファイル出力しています。

これまでと同様LocalTargetを使っていますが、バイナリで出力する場合は、formatオプションにNorFormatを指定します。ちなみに、runメソッドでファイルをオープンするときは、open(self.output().path, "wb")としてます(バージョンの問題なのか self.output().open("wb")ではだめでした)。

class CreateModelTask(luigi.Task):
    task_namespace = 'iris_tasks'

    intermediate_model_path = luigi.Parameter()

    def requires(self):
        return DownloadCsvTask()

    def run(self):
        with open(self.output().path, "wb") as target:
            # self.output().open("wb") ではエラーが発生した
            # https://github.com/spotify/luigi/issues/1647

            data = pandas.read_csv(self.input().path, header=None)
            x = data.iloc[:, [0, 1, 2, 3]]
            y = data.iloc[:, 4].replace({'Iris-setosa': 0, 'Iris-versicolor': 1, 'Iris-virginica': 2})

            model = svm.SVC()
            model.fit(x, y)

            pickle.dump(model, target)

    def output(self):
        # バイナリファイルを出力する場合は、formatにNopFormatを指定
        return luigi.LocalTarget(self.intermediate_model_path, format=luigi.format.NopFormat)

最後のUploadS3Taskでは、作成したpklファイルをS3のパスにアップロードします。

S3に出力する場合は、S3Targetを使うと自然に出力できます。
botoも事前にインストールしておく必要があります(内部的にはS3Clientでアクセスしています)。また、luigi.cfgの[s3]で渡したaws_access_key_idとaws_secret_access_keyの値が認証情報として使われます。

class UploadS3Task(luigi.Task):
    task_namespace = 'iris_tasks'

    s3_model_path = luigi.Parameter()

    def requires(self):
        return CreateModelTask()

    def run(self):
        with self.output().open("w") as target:
            # 存在するファイルをアップロードする場合は、tmp_pathにファイルをコピー
            shutil.copy(self.input().path, target.tmp_path)
        
    def output(self):
        # S3Targetでアップロードできる(botoのインストールが必要)
        return S3Target(self.s3_model_path)

まとめ

ジョブスケジューラを使うまででもないけど、データの受け渡しを伴う処理だったり、時間がかかる処理を分けたい場合に使えそうですね。