ColabでwandbのAPI KEY入力を自動化する

Weights & Biases (wandb) をColab上で使う場合、wandb.login()でトークンを入力する形が多いように思う。こんなふうに。

しかし、いま僕は7つのA100インスタンスを同時にぶん回しており、インスタンス立ち上げのあといちいちログインで引っかかるとだるい。この入力プロセスを省略するためには、公式にもあるとおりで基本的には環境変数に入れた上でwandbのコマンドラインツール経由でログインすればいい。

ただ、このauthorizationのためにNotebook上に平文でAPI_KEYを入れる猛者がいそうなので、念のためColabの環境変数として設定する方法をここに残しておく。やり方は至極簡単。

Colabでは左の「シークレット」から環境変数等を格納するタブを開けるので、そこにとりあえずWANDB_API_KEYとか名前つけてvalueを入れてあげた上で、このNotebookからのアクセスを許可しておく。

この状態にしておけば、例のコードに示されているとおりuserdata.get(<name>)で対応する<value>を取得できるので、あとはこんな感じでwandbのコマンドラインを使って、

import wandb
from google.colab import userdata
wandb_api_key = userdata.get('WANDB_API_KEY')
!wandb login $wandb_api_key

とかしてあげればログインは完了。試しにwandb.init()とかで確認してみると、

今回はこれまでのように https://wandb.ai/authorize を開いてAPIキーを叩き込むようなことは一切なく実際にrunが立ち上がっているので、ログインは成功。あとは普通に使えばいい。