WandB で作成済みの Run の情報を後から更新したい

Weight & Biases (wandb) は MLOps のためのプラットフォームであり，機械学習実験の管理や学習ログの可視化機能を提供しています

W&B を使う中で，過去の実験 (Run) を整理するために後からタグをつけたりグループ化したい場面があるかと思います．
W&B では GUI から過去の Run の情報を更新できない (?) みたいなので，
今回は Python で API を叩いて更新する方法を記事にします

Python から Wandb API を叩く

まず，クライアントライブラリをインストールします

    pip install wandb

CLIで wandb login するか 環境変数WANDB_API_KEYにAPIキーを設定して認証を済ませておきます (ここまでは設定済みかと思います)

認証が終わったら API にアクセスするためのインスタンスを作成します

    from wandb import Api

    # W&B APIを初期化
    api = Api()

Run クラスのオブジェクトを取得する

過去の Run を更新するために Python のオブジェクトとして取得します

Runオブジェクトを取得する場合
api.runから Run のパスを指定して取得します．パスは Run を開いたときのURL https://wandb.ai/<entity>/<project>/runs/<run_id>で確認できます．あるいはブラウザから Run を開いて Overview タブ > Run path でも確認できます

    project = 'project_name'
    entity = 'user_name'  # W&Bアカウントのユーザー名またはチーム名
    run_id = 'run_id'
    
    run = api.run(path=f"{entity}/{project}/{run_id}")

Runsオブジェクトを取得する場合
過去の Run から特定の条件にマッチするものを絞って取得することも可能です
config.dataset が CIFAR10 である Run を取得する場合は次のようにします

    project = 'project_name'
    entity = 'user_name'  # W&Bアカウントのユーザー名またはチーム名

    filters = {"config.dataset": "cifar10"}
    runs = api.runs(path=f"{entity}/{project}", filters=filters)

    # runsはiterable
    for run in runs:
        print(run.name)

Run の各種設定を更新する

Run オブジェクトのプロパティを直接更新すれば良いです．run.update()を呼び出せば更新した情報がブラウザ上のデータと同期されます

    import time
    
    for run in runs:
        # configの更新 (nameのみ置き換わる)
        run.config.update({"name": "updated"})
        # tagの更新 (全てのタグが置き換わるので注意)
        run.tags = ["updated"]
        # groupの更新
        run.group = "updated"
    
        # データを置き換えたら同期する
        run.update()

        # 念の為APIサーバーに負荷をかけないようにしておく
        time.sleep(5)

参考