dagster partitionで機械学習の実験を管理する【競馬AI開発#65】

dagster の「パーティション（partition）」とは、dagster のアセットを日付や実験名などの「区切り」ごとに分けて実行・保存する仕組みです。パーティションを選んで実行すると、同じ名前なら上書きされ、別の名前なら前の結果の隣に並んで残ります。

本記事ではまず、パーティションの考え方を小さなデモで一通り見ます。後半では実例として、この【競馬AI開発】シリーズで作っている競馬予想AIの学習以降の成果物に実験名のパーティションを付けて、設定・モデル・評価結果・回収率を実験名で揃えて見比べられるようにします。

1. dagster のパーティションとは

dagster は、データ処理の流れを依存関係つきで管理・実行できる Python のツールです。

dagster では、テーブルや学習済みモデルのように、処理が作るデータを「アセット」と呼び、Python の関数として定義します。アセットの関数を実際に動かして結果を保存することを materialize（マテリアライズ）と言い、ふつう1つのアセットの結果は1つのファイルに保存されます。

パーティションは、この1つのアセットを、名前の付いた「区切り」ごとに分けて実行・保存する仕組みです。分けた1つ1つのこともパーティションと呼び、その名前を partition key と呼びます。

materialize するときは、どのパーティションを実行するかを1つ選び、その結果は partition key ごとに別々の場所へ保存されます。パーティションを付けても dagster の画面に並ぶアセットは1つのままで、partition key の数だけ増えるのは、保存される結果のほうです。

パーティションは日付で区切るのが定番の使い方ですが、区切る軸は名前を付けられるものなら何でもよく、地域ごとや顧客ごとにも区切れます。

今回は、実験名で区切ります。条件を変えて学習した結果を、実験名という区切りに1つずつ入れていく使い方です。ただ、実験名は「どんな実験をいくつ回すか」を試しながら決めていく場合が多いので、日付のように前もって全部を決めておくことができません。

そこで、「動的パーティション（dynamic partitions）」というものを使います。動的パーティションは、定義した時点ではパーティションが1つも無く、実行する段階でUI上から自由にpartition key名を決めて追加することができます。新しい実験を試すときは、その実験名でパーティションを追加してから実行する、というわけです。

1.1 動的パーティションを付けた最小のアセット

具体的なコードを見てみます。点数の一覧から合格点以上の点数だけを残すアセット passed_scores と、その人数を数えるアセット pass_count の2つです。合格点は、実行するときに値を変えられる設定（Config）として持たせています。

‍▼partition_demo.py

import dagster as dg

# 後から UI でパーティションを足せる、動的パーティションの定義
experiment = dg.DynamicPartitionsDefinition(name="experiment")


class ScoreConfig(dg.Config):
    passing_score: int = 60  # 合格点


@dg.asset(partitions_def=experiment)
def passed_scores(config: ScoreConfig) -> list[int]:
    # 例なので、点数はコードに直接書いている
    scores = [45, 52, 61, 68, 74, 83, 90]
    return [s for s in scores if s >= config.passing_score]


@dg.asset(partitions_def=experiment)
def pass_count(passed_scores: list[int]) -> int:
    return len(passed_scores)


defs = dg.Definitions(
    assets=[passed_scores, pass_count],
    # 戻り値を data/ 配下に保存する
    resources={"io_manager": dg.FilesystemIOManager(base_dir="data")},
)

やっていることは3つです。

1. パーティションを定義する。‍ dg.DynamicPartitionsDefinition が、動的パーティションの定義です。引数の name はこの定義自体の名前で、partition key とは別物です。定義した時点では、パーティションはまだ1つもありません。

2. アセットに付ける。‍ @dg.asset は、その関数が「データを作る関数（アセット）」だと dagster に伝えるデコレータです。その partitions_def に1の定義を渡すと、アセットにパーティションが付きます。pass_count の引数を passed_scores と書いているのは、「passed_scores の結果を使って pass_count を作る」というつながりを dagster に伝えるためです。

3. 保存先を決める。‍ resources の io_manager に、アセットの戻り値をファイルに保存する部品（IO Manager）として、dagster 標準の FilesystemIOManager を指定し、保存先を data フォルダにしています。パーティションを付けたアセットの結果は、この下の「data/<アセット名>/<partition key>」に保存されます。

ポイントは、関数の中身に partition key が一度も出てこないことです。どのパーティションとして実行されているかを関数は知らず、保存先を partition key ごとに分けるのは dagster 側で実行する時の仕事なわけです。

1.2 パーティションを追加して materialize する

ターミナルで次のコマンドを実行すると dagster の UI が起動するので、表示されたアドレス（デフォルトはlocalhost:3000）をブラウザで開きます。

dagster dev -f partition_demo.py

左のメニューの Lineage を開くと、passed_scores から pass_count へつながる2つのアセットが表示されます。

右上の Materialize all を押すと、「Partition selection」という見出しの画面が開きます。まだパーティションは1つも無いので、選択欄（Select a partition or create one）の下にある Add a partition を押し、Partition name に baseline と入力して Save を押します。

右下の Launch 1 run を押して実行すると、passed_scores と pass_count が baseline として順に materialize されます。

実行が終わると、partition_demo.py と同じフォルダにできた data フォルダの中は、次のようになります（右は各ファイルの中身）。

data/
├── passed_scores/
│   └── baseline      [61, 68, 74, 83, 90]
└── pass_count/
    └── baseline      5

アセット名のフォルダの中に、partition key と同じ baseline という名前のファイルが1つずつできています。

1.3 別の名前は並んで残り、同じ名前は上書きされる

次は、合格点を80点に上げた結果を、strict という別の名前で残します。

設定を変えて実行するときは、Lineage の画面で Materialize all の横の ▼ を押し、メニューから Open launchpad を選びます。開くのは、実行前に設定を書き換えられる Launchpad という画面で、詳しくは#59の記事で解説しています。

左上の Partition の欄の横にある Add new partition を押し、先ほどと同じ画面で strict を追加すると、strict が選ばれた状態になります。

設定欄の passing_score を80に書き換えて、右下の Materialize を押します。実行が終わると、data フォルダの中は次のようになります。

data/
├── passed_scores/
│   ├── baseline      [61, 68, 74, 83, 90]
│   └── strict        [83, 90]
└── pass_count/
    ├── baseline      5
    └── strict        2

baseline の結果はそのまま残り、その隣に strict の結果が並んでいます。

1つのアセットの中に partition key の数だけ結果が並び、同じ partition key のブロック同士が、上流の passed_scores から下流の pass_count へつながっています。pass_count を strict として作るときに渡されるのは、passed_scores の strict の結果です。dagster の公式ドキュメントでは、この「同じ名前同士でつながる」規則を次のように説明しています。

同じpartition keyでもう一度実行した時には、上書きされます。
例えば strict を選んで、passing_score を70に変えて実行すると、data フォルダの中は次のようになります。

data/
├── passed_scores/
│   ├── baseline      [61, 68, 74, 83, 90]
│   └── strict        [74, 83, 90]
└── pass_count/
    ├── baseline      5
    └── strict        3

ファイルは4つのままで、strict の中身だけが70点で計算し直した結果に置き換わっています。baseline には触れていません。

2. 競馬AIの実験管理をdagsterで出来るようにする

ここからは実例です。章1のパーティションを、本シリーズで開発している競馬予想AIの、学習以降のアセットに付けていきます。

これにより、学習の設定・予測モデル・確率補正モデル・評価結果・回収率が、実験名（partition）ごとに分かれて残り、実験名で揃えて見比べられるようになります。

2.1 学習以降のアセットに実験名のパーティションを付ける

パーティションを付けるのは、学習の設定（train_config）から回収率（return_rate_fixed_tansho）までの、学習以降のアセットです。特徴量やデータセットまでは今までどおり1本のままで、学習以降が実験名ごとに枝分かれします。

保存先のファイルで見ると、学習以降のアセットのフォルダの中に、実験名ごとのファイルが並ぶ形になります。

実験名で分けるのは学習の設定の違いで、データの期間は今までどおり全実験で共通です。

まず、実験名のパーティションを1つ定義します。