[Jupyter ユーザ向け] mlflow で実験時の notebook も一緒に保存しよう

TL;DR

• Jupyter Notebook / Lab を使ってる人向け
• notebook で実験するときに実験前後で notebook を mlflow に上げておくと、再現性に困った時助かる
• notebook のアップロードは単純にはできなくて少し面倒

背景

機械学習の実験などで Jupyter notebook や Jupyter lab を使っている方は多いと思います。インタラクティブに扱えて便利な反面、グローバルな状態が見えづらく、セルの実行順序で結果が変わるなど、注意深く扱わないと再現性が著しく落ちてしまうというデメリットもあります。
全ての notebook が Run All Cells できちんと再現できればよいですが、現実的にはそうでない場合もあり、特に複数人で notebook を共有する場合などは難しく、実行順序の再現はある種の「職人技」が必要になります。
この記事は、「実験したあとに notebook でセルをいじってしまって、どれが元だったか分からない」、「データサイエンティストから受け取った notebook で同じ結果が再現できない」などを経験したことがある方向けに、実験した瞬間の notebook を mlflow に保存することで保険にしようというものです。

理想はこのような問題が起きないように皆が気をつけることなので、上記のような経験がないあなたはそのままで良いと思います！残念ながら私はありました;(

対策

非常に単純ですが、単に実験の前後で notebook ファイルを mlflow にアップロードして、もし何かあったときに見返せるようにする、というものです。容量もモデル自体と比べれば無視できるレベルのものが多いかと思います。
mlflow については本記事では特に説明しませんが、artifact を置けるような実験管理ツールであれば他のものでも問題ないかと思います。

logger 側

具体的に、私のチームでは次のような mlflow (など)のラッパを作って、実験の start 時に前後で notebook を S3 にアップロードしています。(本記事に関係ないコードも少し混ざってます)

class MLLogger:
    ...

    @contextlib.contextmanager
    def start(self, experiment_name: str, run_name: str) -> Generator:
        mlflow.set_experiment(experiment_name)
        with mlflow.start_run(run_name=run_name):
            self.log_running()
            self.log_user(self._user)
            self.log_notebook()  # <- ここで notebook をアップロード
            try:
                yield
                self.log_notebook()  # <- ここで notebook をアップロード
                self.log_success()
            except (Exception, KeyboardInterrupt):
                self.log_failure()
                raise
    ...

    def log_notebook(self) -> None:
        '''
        notebook を mlflow に保存します。
        '''
        notebook = get_notebook_path()  # notebook のパスを取得
        self.log_artifact(notebook)  # mlflow.log_artifact(file) を呼んでるだけ

notebook のパスの取得

これが単純には行かず、遠回りな方法で解決しました。こちらの StackOverflow を参考にさせていただきました。

Jupyter notebook 内では __file__ が None になります(ナント...)。そのため、ファイル名取得のために現在のカーネルのIDから一旦外に出てリクエストを自身に投げ、session 一覧から notebook を取ってくるということをします。

from notebook import notebookapp
import urllib
import json
import os
import ipykernel


def get_notebook_path():
    '''
    mlflow に notebook をそのまま保存する際に notebook のパスを取得するのに使います。
    単純にファイル名が取れないためかなりややこしいことになっています。
    使用するには jupyterlab が token ありで起動していないといけません。
    '''
    connection_file = os.path.basename(ipykernel.get_connection_file())
    kernel_id = connection_file.split('-', 1)[1].split('.')[0]

    for server in notebookapp.list_running_servers():
        try:
            if server['token'] == '' and not server['password']:  # No token and no password, ahem...
                req = urllib.request.urlopen(server['url'] + 'api/sessions')
            elif server['token'] != '':
                req = urllib.request.urlopen(server['url'] + 'api/sessions?token=' + server['token'])
            else:
                continue
            sessions = json.load(req)
            for sess in sessions:
                if sess['kernel']['id'] == kernel_id:
                    return os.path.join(server['notebook_dir'], sess['notebook']['path'])
        except Exception:
            raise
    raise EnvironmentError('JupyterLab の token を指定するか、パスワードを無効化してください。')

これを使うには Jupyter が token 付きで立ち上がっている必要があります。
適当に token 付きで Jupyter を立ち上げるには、次のようにコマンドで指定するか、設定ファイルに書いておきましょう。

# 立ち上げるときに指定する場合
$ jupyter lab --notebook-dir . --NotebookApp.token='tekitou_na_token'

# 設定ファイルに書いておく場合
$ echo "c.NotebookApp.token = 'tekitou_na_token'" >> ~/.jupyter/jupyter_notebook_config.py

これで実際に notebook が mlflow の artifact から確認できれば完了です!

まとめ

実験開始時・終了時の Jupyter notebook を mlflow に上げておくことで、実験実行時の notebook を保存する方法を紹介しました。
notebook のパスの取得は少しやっかいでした。もし良い方法を知ってる方がいたらコメントお願いします
理想はこんなもの必要がない状態を作ることですが、既に mlflow などの実験管理ツールを導入していれば少し足すだけで実現できるので、やっておくといざ困ったときに助けになるかもしれません。