エムスリーテックブログ

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

機械学習パイプラインライブラリluigiのshell補完ツールを作った

こんにちは。Python謹製のshellを使って早3年、エムスリーエンジニアリングGの河合 (@vaaaaanquish) です。 この記事はエムスリー Advent Calendar 2019 の21日目の記事です。

エムスリー AIチームでは、機械学習パイプラインとして「luigi」及び、そのwrapperである「gokart」を利用した開発、運用を行なっています。 本記事は、エムスリーとluigiの繋がりと、私が作成したluigiのshell補完をサポートするmodule「luigi-completion」の概要、その使い方について示すものです。

f:id:vaaaaaanquish:20191216102634p:plain
luigi.readthedocs.ioより luigiロゴ

 

はじめに

前述の通り、エムスリーAIチームでは、機械学習プロジェクトにおけるパイプラインライブラリとして、luigiとgokartを利用しています*1*2

 

luigiは、spotifyがOSSとして作成、公開している、機械学習プロジェクト向けのパイプラインライブラリです。タスクの共通化や再現性の担保など、機械学習プロジェクトにおける様々なissueを解決するためのツールです。 github.com gokartは、エムスリーが独自に作成したluigi wrapperで、luigiに対してより簡易なタスク設計が行える拡張を追加していたり、連携モジュールであるthunderboltによるデータ管理をサポートしています。 github.com エムスリー AIチームでは、gokartの開発を通して得た知見や課題についてluigiにcommitしたり、READMEにリンクを掲載する等しており、それぞれ互いに機能を高め合う関係にあります。

  

luigiとgokartは、機械学習モデリングやモデル運用におけるタスク管理や再現性を強力にサポートしており、私達のチームとは切っても切り離せないツールとなっています。 実験を繰り返す機械学習モデリングの工程において、「パラメータ管理」の機能は最も重要です。これらは、パラメータ用のファイルや環境変数、CLIから実行する際の引数として設定する事が可能ですが、利用可能なパラメータは、luigi.Taskを継承したタスク、それらが持つ全てのパラメータとなり、膨大な数になっていまいます。それらを全て把握した上でCLI上で管理するのは困難です。現に私の持っているプロジェクトでは、importしているモジュール全て合わせて1000以上のパラメータがあり、強力な補完によるサポートが必要な状態になっていると言えます*3

  

luigi_completion

CLI上でluigiの補完をサポートする、luigi_completionというPythonパッケージを作成しました。

github.com

pypiよりInstallする事が可能です。

pip install luigi_completion

現在、luigi_completionは以下の機能をサポートしています。

  • mklcomp
    • プロジェクトのディレクトリに.luigi_completionファイルを作成するCLIコマンド
  • luigi_completion.bash, .zsh
    • .luigi_completionがあるディレクトリでpython, luigiコマンドの補完を拡張する

実際にluigiプロジェクトのrootディレクトリに移動し、mainとなるファイルを引数に指定する事で.luigi_completionを作成します*4

$ ls -a
hoge/    main.py

$ mklcomp main.py
[RUNNING] python main.py --help-all
[DONE] update completion

$ ls -a
hoge/    main.py    .luigi_completion

.luigi_completionファイルが出来ている事が確認できます。

 

実際に補完を有効にするためには、GitHub Repository内にあるbash、zshスクリプト*5をsourceで読み込みます。こちらは、各シェルのパスの通った場所に設置するか、.bashrcや.zshrcに書いておくと良いと思います。

$ source luigi_completion.bash

TAB補完に.luigi_completionに記載されたtaskとparameterが追加されます。

$ python main.py [TAB]
hoge.DownloadSample    hoge.TrainData    hoge.OutputFile   ...

$ python main.py hoge.TrainData [TAB]
--local-scheduler    --rerun    --from-date    --to-date    --iter    --sample-num   ...

f:id:vaaaaaanquish:20191216135918g:plain
TAB補完の様子

 

処理の中身

現在luigiがサポートしている機能は、補完先を.luigi_completionファイルとして保存する機能と、それらをshellの補完候補に追加する機能の2つです。

 

luigiの全てのタスク名とパラメータ、及びそれらのDescriptionを表示する方法として、--help-allというフラグを引数に入れ実行するという方法があります。

$ python main.py --help-all


usage: task.py [--TestNotificationsTask-raise-in-complete]
               [--email-force-send] [--email-format EMAIL_FORMAT]
               [--email-method EMAIL_METHOD] [--email-prefix EMAIL_PREFIX]
               [--email-receiver EMAIL_RECEIVER] [--email-sender EMAIL_SENDER]
               [--smtp-host SMTP_HOST]
               [--smtp-local-hostname SMTP_LOCAL_HOSTNAME] [--smtp-no-tls]
               [--smtp-password SMTP_PASSWORD] [--smtp-port SMTP_PORT]
               [--smtp-ssl] [--smtp-timeout SMTP_TIMEOUT]

---[中略]---

               [--sample.SecondTask-number SAMPLE.SECONDTASK_NUMBER]
               [--sample.SecondTask-param SAMPLE.SECONDTASK_PARAM]
               [Required root task]

positional arguments:
  Required root task    Task family to run. Is not optional.

optional arguments:
  --TestNotificationsTask-raise-in-complete
                        If true, fail in complete() instead of run()
  --email-force-send    Send e-mail even from a tty
  --email-format EMAIL_FORMAT
                        Format type for sent e-mails Choices: {none, html,
                        plain}

---[以下略]---

--help-allはPythonを一度起動し、全ての外部moduleをimportしてluigiタスクとパラメータを探しています。そのため、1回結果を参照するのに10秒以上かかる事もあります。また、プロジェクト毎に補完対象が変化するため、プロジェクトに合わせて実行を行う必要があります。

 
mklcompコマンドは、この--help-allの結果をparseし、テンポラリとして.luigi_completionというファイルに変換するコマンドです。作成した.luigi_completionファイルは以下のように、補完候補となるmodule名とparameterがそれぞれ入っています。

module:sample.SecondTask
module:sample.SampleTaskparameter:--RangeBase-reverse
parameter:--wait-interval
parameter:--scheduler-state-path

---[以下略]---

mklcompコマンドでアップデートも出来ますし、不要と考えられる補完候補を削除する事も可能です。 luigi_completion.bash、luigi_completion.zshスクリプトは、この.luigi_completionがディレクトリ内にある場合に、.luigi_completionをparseして補完候補に追加します。 このような構成にする事で、プロジェクト毎のtaskやパラメータの違いを吸収しつつ、素早い補完に対応しています。

 

おわりに

本記事では、エムスリーAIチームとluigiの繋がり、およびluigiのshell補完をサポートするmodule「luigi-completion」の概要と、その使い方について記載しました。

luigiは、ここ数年で大きく成長したプロダクトの1つであると認識しています。一方で、機能の多さや粗削りな部分もまだまだ残るライブラリです。また、gokartやthunderboltなど、エムスリーが作成している周辺ライブラリにもまだまだな所が多くあります。

本記事で紹介しているluigi-completionも、同僚のふとした会話から私が1日で作成したものであるのが現状です。冒頭の通り、私は普段xonshというshellを利用して、Pythonプロセスを経由した非同期なluigiタスク補完を長らく利用していたのですが、チーム内に以下のような投稿を試しにしてみた所、bashやzshが圧倒的な利用率でした。その結果、勢いで作成したのがluigi-completionです。

f:id:vaaaaaanquish:20191216121920p:plain

Mac OSのdefault shellがzshになるという話もありますし、今後もこういったツールを徐々に作成、改修して、luigiに対しても貢献して行ければと思っています。

もちろんOSSですので、皆さんのコミットメントもお待ちしております。

 

We are hiring !

最後に、手前味噌ではありますが、エムスリーでは、こういったOSSの公開や技術検討、議論とそれらの公開が「技術力向上」として評価されるようになっています。 記事のブックマークやシェア、各ライブラリへのStar、commit、issue投稿、使ってみた感想の投稿に加えて、エムスリー株式会社採用への応募も心よりお待ちしております。宜しくお願いします。

open.talentio.com jobs.m3.com

 

*1:機械学習プロジェクト向けPipelineライブラリgokartを用いた開発と運用 - エムスリーテックブログ

*2:機械学習で便利なluigiフレームワークの紹介 - エムスリーテックブログ

*3:パラメータを共通化して設定可能なluigiの機能によって実際に触るパラメータは100に満たない状態ですがそれでも多い

*4:mainファイルではluigi.run()が走るようになっている想定です

*5:https://github.com/vaaaaanquish/luigi_completion/tree/master/bin