Kedroにおける特徴量のバージョニング方法の一案

（異常に脱字が多かったので見直しました 12/28)

最近検証を兼ねて、Kedroを使ってMLモデル開発におけるワークフロー管理をしています。
Kedroはワークフロー管理ツールの一つであり、python関数をnodeというクラスでラップしてつなぎ合わせることでPipelineを構成すること、nodeの入出力は全てData Catalogと呼ばれるメモリやストレージを抽象化したファイルシステムクラスで管理されることが特徴です。io自体は別個のDataSetというモジュールを用意して対応しています。

kedroの利用にあたって学習用に生成した特徴量と推論用に生成した特徴量を、中間生成物を含めて丸ごと別々に保存したいと考えましたが、どうもその分けて保存する方法についてストレートな解決方法が見つからなかったため、今回試行した方法をこちらに残します。

前提

kedroが自動生成してくれるプロジェクトテンプレートに従ってファイル構成を切ったとします。
kedroのバージョンは0.17.0です。

Data Catalogでのデータ管理

Data Catalogの用意の仕方はpython code上で定義する方法とyamlで記述し、hooksを仕掛けてkedro実行時にそのyamlをベースに生成する方法があります。今回は後者に基づくので後者の例だけ記述します。

Kedroは実行時に自動的にルートディレクトリからconf/*/catalog*という名前のyamlを見つけて読みに行きます。
yamlに書かれた設定を元に、Data Catalogが作られます。

titanic_train:
  type: pandas.CSVDataSet
  filepath: s3://competitions/titanic/train.csv
  credentials: minio

outputs:
  type: pandas.CSVDataSet
  filepath: s3://competitions/titanic/outputs.csv
  credentials: minio
  versioned: true

現行のkedroのバージョニング

DataSetのバージョニング

DataSetにはバージョニングの方法が提供されています。
AbstractVersionedDataSetというクラスを継承してDataSetクラスを実装し、かつcatalog.ymlにおいてversionedをtrueにすることでバージョニングが有効になります。
ただ困ったことに、この方法ではtimestampしかバージョン名に使えません。（おそらく）
コードでDataCatalogを設定する場合にはこの限りではないのですが、少なくともyamlの記述からバージョンのフォーマットは変更できないようです。
さらにCLIを使って実行する場合はコードで生成したDataCatalogを注入する手段は現状ないようです。
-> HOOKSにおいて直接言語で定義したDataCatalogインスタンスを呼び出せば良いようです。

時間以外でバージョニングする動機

私は生成した特徴量を時間以外のバージョン名でバージョニングしたい（trainとtestのように分けたい）と思っています。
これのために取れる手段としてはcatalog.ymlに新しくtest用のデータセットを記述し、さらにpipelineも別途test用の物を生成する必要があります。
pipeline定義時にはnodeが定義されている必要があり、nodeの定義時にバージョンを読み取って参照するデータセットを変更する処理が必要になります。読めないコードの始まりです。
一方で当然ながら特徴量生成のロジック自体は学習時と推論時で同じである必要があるし、コード上でも共通化したいです。
これの回避を目的としてデータを時間以外の何かしらの値でバージョニング（というよりはタグ付与）がしたいです。

TemplatedConfigLoaderによる設定への値の注入

0.17.0から実装された機能です。
yamlにプレイスホルダーを設定でき、値をコードから注入できます。外部にプレイスホルダーへのマッピングを記述したyamlを用意することでも注入できます。
これを使うことでcatalog.ymlにおけるfilepathを直接変更し、バージョニングを実現します。
バージョンごとに保存するパスを変更することでバージョニングを実現しようという腹です。
また、これによって、全てのデータセットが一律のバージョン名であることを強制されていましたが、データセットに応じてバージョン名を変更することができます。

TemplatedConfigLoaderはhydraの簡易版ぐらいの機能です。
hydraとはyamlの構成管理のツールの一つです。コマンドラインからの値の注入やyamlの構造化などを実現できるようにするOSSです。
hydraはhydraで独特の記法を要求され、その書き方がkedroと相性がちょっと悪いのでkedroでテンプレート機能を用意してもらえるのは便利ですね。

Data Catalogの用意の流れ

DataCatalogが用意される時の処理を、プロジェクトテンプレートからプロジェクトを生成した場合について記述します。これによってどこにTemplatedConfigLoaderを適用すべきかを理解します。

プロジェクトテンプレートよりcliなどを用意した場合、kedroの実行セッションはKedroSessionと呼ばれるクラスによって生成されます。
細かいところは正しく説明しきれないので無視しますが、この時<project_name>/src/配下にあるhooks.pyにあるProjectHooksのインスタンスを参照して、記述されている各hooksを実行します。

このうちDataCatalogに関わるhooksはregister_config_loaderとregister_catalogの二つです。
前者がConfigLoaderを用意し、後者がDataCatalogを用意します。
また、register_catalogにおいてはregister_config_loaderを利用して読み込まれたcatalog.ymlに基づいてDataCatalogを生成します。

以上からregister_config_loaderにおけるConfigLoaderをTemplatedConfigLoaderに差し替えれば、保存するパスを動的に変更できそうなことがわかります。

実装

TemplatedConfigLoaderをregister_config_loaderに実装すれば、catalog.ymlを動的に変更できそうなことはわかりました。
実際にはCLIなどでバージョン名を操作したいので、その入力を受け付けられるようにします。
なので
- ProjectHooksクラスにバージョンに対応する変数を設定する
- register_config_loaderにおいてTemplatedConfigLoaderによって設定を読み込むようにする。この時catalog.ymlに値を注入するように変更する
- cli.pyでコマンドラインからバージョン名を受け取り、セッション生成前にProjectHooksにバージョン名を渡すようにする
- catalog.ymlをTemplatedConfigLoaderから値を受け取れるようにプレイスホルダーを用意する
ということをします。

ProjectHooksクラスにバージョンに対応する変数を設定する

class ProjectHooks:
    _mode: str = ''

    @classmethod
    def set_mode(cls, mode: str):
        cls._mode = mode

Sessionではここで生成されたhooksのインスタンスを呼び出すという仕様になってます。シングルトンパターンを採用されてないので、呼び出されても同一のインスタンスであることが保証できません。なのでバージョン名に対してクラス変数を設定しています。（変数名がmodeなのは実際のところ、自分は実行目的に合わせて利用するデータを変更したいだけであり正しくバージョニングしたいわけではないからです。）

register_config_loaderの変更

ProjectHooksに次のインスタンスメソッドを実装します。

    @hook_impl
    def register_config_loader(self, conf_paths: Iterable[str]) -> TemplatedConfigLoader:
        return TemplatedConfigLoader(conf_paths, globals_dict=dict(mode=self._mode))

これによって、modeというプレイスホルダーを含んでいればそこにself._modeの値を代入するようになります。

cli.pyでコマンドラインからバージョン名を受け取り、セッション生成前にProjectHooksにバージョン名を渡すようにする

@click.option(
    "--run-mode", type=click.Choice(['train', 'inference'], case_sensitive=True), default="train"
)
def run(
    ...,
    run_mode
):
   ...
   from .hooks import project_hooks
   project_hooks.set_mode(run_mode)

(importは書く上で見やすいところに置いただけなので好きなところにおいてください）

catalog.ymlをTemplatedConfigLoaderから値を受け取れるようにプレイスホルダーを用意する

outputs:
  type: pandas.CSVDataSet
  filepath: s3://competitions/titanic/${mode}/outputs.csv
  credentials: minio

実行

上記の変更で、kedro run --run-mode=trainとした時、ディレクトリtrainとしてきられています。

versionedで制御できなくなってしまいましたが、まぁこの実装の時点で使わない方針なのでいいと思います。

まとめ

catalog.ymlを直接テンプレート化することでバージョニングを実装する方法を説明しました。
モデルについてのみ変更したい場合はまた別途モデル用のプレイスホルダーを用意して何らかの値を与えればよいです。
他の手段として、自前でVersionedDataSetを実装して、yamlの値で変更できるようにすることも考えましたが、結局プレイスホルダーを用意しないと動的に変更できないことからこの形にしました。
今後変更によって利用できなくなるかもしれませんが、一旦はこれでやっていこうと思います。
使っていて結構クリティカルに足りない機能な気がしたので、今後機能追加もありそうですね。（個人の感想です）
issueやPRを覗いてみると同じようなトピックがWIPで進んでいるのでそのうち追加されそうです。
雑に確認したので動かなければ教えて下さい。