エムスリーテックブログ

エムスリー(m3)のエンジニア・開発メンバーによる技術ブログです

Genericやらoverloadやらを使って、MLパイプラインツールgokartを型安全にしてみた

gokart python

AI・機械学習チームの北川(@kitagry)です。 このブログはサテライトオフィスのメンバーで投稿されるブログリレー１日目の記事になります。1

サテライトオフィスブログリレーとある通り、４月から京都オフィスに在籍しています。 最近京都オフィスメンバーで居酒屋に行ったりしたのですが、やはり関西めっちゃ安い。。 みんなおいでよと思っているところです。

本ブログではエムスリー株式会社がOSSとして公開しているgokartというPythonツールの型を強化した話について書きます。 Python3.5から型ヒントの機能が追加され、現在最新のPython3.13に至るまでに徐々に機能が追加されています。 最近はGILの無効化など高速化に関する話題が上りがちですが、その一方で型についても徐々に充実してきています。 このブログではPythonで利用可能な型について言及しながら、gokartに型を導入した話について書きます。

本ブログは次のブログの続きでもあるのですが、このブログ単体でも読むことが出来ます。 mypy pluginが気になる方は是非次の記事も読んでみてください。

www.m3tech.blog

文脈とは何も関係ない、ただただデスクで寝ている猫

gokartについて

gokartはAI・機械学習チーム(以降AIチーム)で良く使われている機械学習パイプラインツールです。 gokartを使うことによって次のメリットがあります。

  • タスクのモジュール化
  • キャッシュの利用
  • スケジューラによる並列処理

ここではgokartを使ったことがない方のために簡単に使い方を書きます。

gokartはgokart.TaskOnKartというクラスを継承して利用します。 ユーザーが主に意識するのは requires と run メソッドの２つになります。

  • requires: 依存関係を列挙するメソッド
  • run: タスク自体の実行部分を記述するメソッド

次のコードが実際にgokartを利用したタスクの書き方の一例になります。

from gokart import TaskOnKart

class Hoge(TaskOnKart):
    def requires(self):
        """自身のtaskへの依存関係を列挙する"""
        return dict(fuga=Fuga(), piyo=Piyo())

    def run(self):
        """このタスクの実行部分"""
        # requiresで書いた依存からデータを取得する
        fuga = self.load('fuga')
        piyo = self.load('piyo')

        result = fuga + piyo

        # 自身のタスクの結果はdumpメソッドで保存する
        self.dump(result)

このようにgokartではrequiresで書いた依存関係をrunメソッドの中で利用して後続の処理を書くことが出来ます。 最後にはdumpメソッドによって、最終的な結果を保存することが出来ます。

さて、では問題です。

上記の依存関係であるfugaとpiyoは何が返ってくるタスクで、Hoge自体は何をdumpしているクラスなのでしょう？

答えはintやfloat型かもしれませんし、もしかしたらstr型かもしれません。 これが何かを知るためには実際にFugaクラスやPiyoクラスを見に行く必要があります。 これらはコードを書くときはまだ記憶があるかもしれませんが、コードを読むときにはかなり負荷がかかります。 このような問題を解決するために、gokartに型を付けて可読性を上げることにしました。

Taskに型をつける

まず手始めに取り組んだのはgokartのタスク自体に型をつけることです。 最初に利用方法について書きます。

class Hoge(TaskOnKart[str]):
    def run(self):
        self.dump(1)  # mypy error
        self.dump('hello')  # mypy ok

dumpで別の型を入れると怒られている図

上記のように、タスクの定義を見るだけでこのタスクの結果の型を見られるようにしました。 この機能の実現にはGenericクラスを用いて実現しています。 Generic型は型をパラメータ化することで、コードの再利用性を高め、利用側の型安全性を確保するためのものになります。 gokartのTaskOnKartの実装は次のようになっています。

class TaskOnKart(luigi.Task, Generic[T]):

    def dump(self, obj: T) -> None: ...

この状態ではT型は不定なのですが、ユーザーがHoge(TaskOnKart[str])と指定することによって、T型がstr型に固定されます。 そのため、dumpメソッドにはstr型以外のものを引数に入れるとmypyに怒られるという仕組みになっています。 もっと詳細な実装方法を見たい人は実際のPRを見てみてください。

この型を定義することで間違った型をdumpすることを防げます。 また、副次的な効果としてTaskの定義を見ただけで何の結果をdumpするタスクなのかが一目でわかります。

依存に型をつける

前章ではタスク自体に型をつけるというものでした。 これによって、そのタスク内においては型がついている状態になりました。 続いて、タスク間の型情報のやり取りについて考えてみます。

そのためにまずは、gokartのmypyプラグインが必要不可欠になります。 mypyプラグインの有効化についてはこちらをご参照ください。

github.com

またmypyプラグインについて知りたい場合は、同じく京都オフィスの山本のブログに詳しくまとまっているのでそちらをご参照ください。

www.m3tech.blog

mypyプラグインを利用することによってGeneric型と組み合わせて外部から注入するタスクの型を指定することが出来ます。 次の例は、strを返すタスクのみに絞るコードです。

class Hoge(TaskOnKart):
    str_task: TaskOnKart[str] = gokart.TaskInstanceParameter()

Hoge(str_task=StrTask())  # mypy: ok
Hoge(str_task=IntTask())  # mypy: ng

このようにTaskの型とmypyプラグインを組み合わせることによって、外部から注入するときに型を指定することが出来ました。 しかし、まだloadメソッドには型をつけることが出来ませんでした。

class Hoge(TaskOnKart):
    str_task: TaskOnKart[str] = gokart.TaskInstanceParameter()

    def requires(self):
        return dict(str=self.str_task)

    def run(self):
        s = self.load('str')  # Any型になってしまう。

次の章からはこのloadメソッドに型をつけることの難しさについて解説しつつ、どのように型をつけるようにしたかについて説明します。

Pythonの型の限界

loadメソッドに型をつけるにあたって、まずはrequiresの型について考えてみます。 requiresでは複数の依存を表現するためにdict型を使うことが多いです。 複数依存に対応するため、requiresメソッドの返り値はdict[str, TaskOnKart[Any]]型になります。 これだと、loadでdictのキーを指定して取得するときにどの型のTaskOnKartが返ってくるのかがわかりません。

そこで、各key-valueの型を対応させるような型であるTypedDictを利用すればよいのではということが浮かびます。 TypedDictを利用することによって、"str"がkeyのときにはTaskOnKart[str]が返ってくることがわかります。

class RequiresDict(TypedDict):
    str: TaskOnKart[str]
    int: TaskOnKart[int]

class Hoge(TaskOnKart):
    ...
    def requires(self) -> RequiresDict:
        return dict(str=self.str_task, int=self.int_task)

では、引数に"str"文字列が来たことをloadメソッドにわたすにはどうすればよいでしょうか？

def load(self, key: Literal["str"]) -> RequiresDict[Literal["str"]]:  # どう書けば良いか？

実は、このようなTypedDictの特定の値の型を返すような記法はPythonには用意されていません。 TypedDict側の実装ではmypyプラグインを用意して型をうまく推論しているようです。 そのため、ユーザー側でどのようなkeyの型からvalue型を返すかを利用する型を書くことは出来ません。 gokartでも出来なくはないのでしょうが、現在はこの方法で型を推測する方法はないです。

それでもloadメソッドに型をつけたい

そこで、loadに型をつけるために次のような文法を導入しました。

class Hoge(TaskOnKart):
    str_task: TaskOnKart[str] = gokart.TaskInstanceParameter()

    def requires(self):
        return dict(str=self.str_task)

    def run(self):
        s = self.load(self.str_task)  # str型だと推測出来る

今までとの違いはloadメソッドへの引数がTaskOnKartのインスタンスになっている点です。 これによってloadへの型の宣言が容易になります。 このようにdictのkeyから型を引っ張ってくるのは現在のPythonでは難しいところを、実際のタスクを渡すことによって型を推測出来ます。

loadに型が付いた

実際のgokartの実装は次のような形になっています。 これは、Genericsによってタスク自体に型が付いているため実現できた書き方になっています。

class TaskOnKart:
    def load(self, task: TaskOnKart[K]) -> K:
        ...

今までの文法も引き続き使えるようにoverloadデコレータを利用しています。 overloadデコレータは引数の型によって、レスポンス型を変えたいときなどに用いると便利です。

    @overload
    def load(self, key: TaskOnKart[K]) -> K: ...

    @overload
    def load(self, key: str) -> Any: ...

    def load(self, key: Union[str, TaskOnKart[K]]) -> Any:

詳しい実装は是非該当のPRを御覧ください。

終わりに

今回の一連のgokartの型周りを使ってみたい方は是非gokartのバージョンを1.6.0以上にあげて使ってみてください。

最後に、最初に上げたgokartの例を今回の型を利用して書いてみましょう。

class Hoge(TaskOnKart[int]):
    def requires(self):
        return dict(fuga=self.fuga_task(), piyo=self.piyo_task())

    def run(self):
        fuga = self.load(self.fuga_task())
        piyo = self.load(self.piyo_task())

        result = fuga + piyo

        self.dump(result)

    def fuga_task(self) -> TaskOnKart[int]:
        return Fuga()  # fugaがint以外の返り値の場合mypyが怒ってくれる

    def piyo_task(self) -> TaskOnKart[int]:
        return Piyo()  # piyoがint以外の返り値の場合mypyが怒ってくれる

多少冗長になった気はしますが、そのかわりにFugaとPiyoがintを返すタスクであり、Hoge自身もintを返すタスクなんだなと言うことがFugaやPiyoの実装を読まずとも分かるかと思います。 また、FugaやPiyoの返り値がintから別のものになったときは、mypyがエラーを返してくれるため実行するまでミスに気づかないということも少なくなります。

まとめ

型は実行前にバグの原因を知らせてくれるテストのような存在だと思っています。 また、型をちゃんと付けることによってエディタの補完が効いたり、可読性が上がったりとメリットがたくさんあります。 Pythonの型の表現の幅が広がってきたことで今回のような改善が出来ました。 今後もPythonの型が進化して、今回断念した場所に型をつけられるようになれば嬉しいなと思います。

We are hiring

実は、今回の一連のgokartへの型強化実装のほとんどが京都オフィスメンバーの２人で実装されました。 東京オフィスメンバーだけでなく、地方メンバーも今回のgokartのようなOSSや外部登壇・カンファレンスなどで活躍している様子が届けられれば幸いです。

www.m3tech.blog

www.m3tech.blog

地方でも活躍できるエムスリーに応募してみるのはいかがでしょうか？

jobs.m3.com

  1. タイトル一緒に考えて欲しいって言ったらいろんなタイトル案が出てきたのでここに供養します。Re:ゼロから始める型付け生活 ~ Dataframe型入出力 ~・Python世界なのに型を持った俺が成り上がる話・いまさら型を付けてくださいなんて言っても遅すぎる。Python界最下級ジョブ型付け師の俺が、狂った動的言語世界で安全スローライフする話・型上手の若君・Re:ゼロから始めるPython型付け生活↩