【mlflow】作成済みのrun_idの情報（metrics, tags, etc...）を更新したい

ドキュメントみてたけど明示的に書いていない（はず）

mlflow.org

start_runの項目を確認すると、run_uuidの引数があるので「もしや・・・！」と思い試した。

# 色々情報はすっ飛ばす
# MLflowに記録
mlflow.set_tracking_uri('http://localhost:5000')
with mlflow.start_run():
    mlflow.sklearn.log_model(model, "ml_models")
    mlflow.set_tag('test_tag', 0)  # あとから更新したいタグ

引っ張ってくる

from mlflow.tracking.client import MlflowClient
from mlflow.entities.view_type import ViewType

query = 'tag.test_tag = "0"'
runs = MlflowClient().search_runs([0], query, ViewType.ACTIVE_ONLY)
_run = runs[0]
print(_run.info.run_uuid)  # 11923d8756ea4d21axxxxxxxxxxxxxxxxx

更新する

# MLflowに記録
mlflow.set_tracking_uri('http://localhost:5000')
with mlflow.start_run(run_uuid=_run.info.run_uuid):  # run_uuidを指定する
    mlflow.set_tag('test_tag', 1)  # 更新

これでできた（めっちゃ時間かかった....