Python: mlflow.start_run(nested=True) は使い方に注意しよう

今回は MLflow Tracking のすごーく細かい話。ソースコードを読んでいて、ハマる人もいるかもなと思ったので書いておく。結論から先に書くと、MLflow Tracking には次のような注意点がある。

MLflow Tracking で標準的に使う API はマルチスレッドで Run が同時並行に作られることを想定していない
同時並行に作れそうなmlflow.start_run(nested=True) は、あくまで Run を入れ子にするときだけ使える
この点に気をつけないと MLflow Tracking で記録されるデータが壊れる

使った環境は次の通り。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.14.6
BuildVersion:   18G5033
$ python -V
Python 3.7.7
$ pip list | grep -i mlflow              
mlflow                    1.8.0

下準備

あらかじめ MLflow をインストールして Python のインタプリタを起動しておく。

$ pip install mlflow
$ python

そして MLflow をインポートしておく。

>>> import mlflow

注意点の解説

MLflow Tracking では、新しく Run を作るときに mlflow.start_run() という関数を使う。

>>> mlflow.start_run()
<ActiveRun: >

この関数は、すでに同じ Python プロセスで呼ばれていると、再度呼び出したときに例外となるよう作られている。

>>> mlflow.start_run()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/amedama/.virtualenvs/py37/lib/python3.7/site-packages/mlflow/tracking/fluent.py", line 112, in start_run
    _active_run_stack[0].info.run_id))
Exception: Run with UUID a14a8f9bf213473aa4a7a437ac811077 is already active. To start a new run, first end the current run with mlflow.end_run(). To start a nested run, call start_run with nested=True

ただし、オプションに nested=True を指定すると、エラーにならず新たに Run を作ることができる。

>>> mlflow.start_run(nested=True)
<ActiveRun: >

ただ、上記で作られた Run は MLflow のモジュールの中でグローバル変数として用意されたスタック構造で管理されているに過ぎない。スタック構造は隠しオブジェクトだけど、あえて中身を見せるとこんな感じ。

>>> from mlflow.tracking import fluent
>>> fluent._active_run_stack
[<ActiveRun: >, <ActiveRun: >]

この状況で mlflow.end_run() を呼び出すと、スタック構造からひとつずつ Run を表すオブジェクトが POP されるという寸法。

>>> mlflow.end_run()
>>> fluent._active_run_stack
[<ActiveRun: >]
>>> mlflow.end_run()
>>> fluent._active_run_stack
[]

つまり何が言いたいかというと、mlflow.start_run(nested=True) は、あくまで Run の中で別の Run を「入れ子」にすることだけを想定している。要するに、以下のようなコード。

>>> # この使い方はセーフ (Run を入れ子にする)
>>> with mlflow.start_run():
...     with mlflow.start_run(nested=True):
...         ...  # 何か時間のかかる実験
...

「入れ子」ではなく、同時並行に Run を作る用途で使ってしまうとスタック構造が壊れる。たとえば、次のようなコードを書くとレースコンディションを生む。具体的には、今実行しているのとは関係ない Run にメトリックやパラメータが記録される。

>>> # この使い方はアウト (スタックが壊れる)
>>> import threading
>>> def f():
...     with mlflow.start_run(nested=True):
...         ...  # 何か時間のかかる実験
... 
>>> for _ in range(10):
...     t = threading.Thread(target=f, daemon=True)
...     t.start()
...

それでも同時に Run を実行したい！というときはプロセスを分けよう。プロセスが違えばグローバル変数のいるメモリ空間も分かれるので問題ないはず。そもそも、一般的な Python の処理系には GIL (Global Interpreter Lock) があるので、I/O を並行処理するときしかマルチスレッドが意味を成さない。

ちなみに、ここらへんの実装は以下にある。

github.com

上記を読むと mlflow.tracking.client.MlflowClient を直接使えばマルチスレッドで複数の Run を同時に使うこともできそう。ただ、得られる嬉しさが手間に見合わないだろうなという感じがする。

いじょう。