シグナル処理と sweep の run

このページでは、W&B Sweeps がシステムシグナルとプロセスの終了コードをどのように処理するかについて詳しく説明します。この情報は、SLURM、EC2 Spot、Google Cloud のプリエンプト可能な VM などのプリエンプト可能な環境で sweep を安定して実行するのに役立ちます。以下のセクションでは、キーボードから run を正常に中断する方法を説明するとともに、run をキューに入れ直す際の動作を理解し、予測するための詳細情報を提供します。このページは、プリエンプト可能なインフラストラクチャー上で sweep を実行するユーザーや、run のライフサイクルとクリーンアップをきめ細かく制御する必要があるユーザーを対象としています。プリエンプトされた run を W&B がキューに入れ直す仕組みの詳細については、プリエンプト可能な Sweeps run を再開するを参照してください。

終了ステータスとシグナル

W&B は、トレーニングプロセスの終了ステータスに基づいて、run をキューに入れ直すかどうか、および run の状態をどのように記録するかを決定します。 終了コードの規約:

終了コード 0: W&B は run が正常に完了したとみなし、キューに入れ直しません。
0 以外の終了コード: W&B は run を失敗またはプリエンプトされたものとして扱います。mark_preempting() を使用すると、W&B は run をキューに入れ直し、別のエージェント (または再起動後の同じエージェント) が run を再開できるようにします。

これは、プロセスがシグナルハンドラから終了した場合、例外によって終了した場合、明示的な sys.exit() 呼び出しによって終了した場合のいずれにも当てはまります。プリエンプト可能な環境やクラスター環境では、この規約を理解し、それを前提に実装することが重要です。プロセスが捕捉可能なシグナルによって終了する場合は、ハンドラを実行できます。ハンドラでは、run をキューに入れ直したい場合に wandb.run.mark_preempting() を呼び出し、クリーンアップ (チェックポイントの保存など) を行ってから、0 以外のコードで終了します。シグナルによる終了では、sys.exit(128 + signum) とするのが一般的な慣例です。W&B はその終了コードを記録し、同じキューに入れ直すルールが適用されます。一方、オペレーティングシステムのカーネルが SIGKILL でプロセスを強制終了した場合、プロセスは終了フックを実行できません。そのため W&B は最終サマリーを書き込めず、run が crashed または killed と表示されることがあります。その場合でも、エージェントは次の run を開始します。

古い run とサーバー側のタイムアウト

W&B は、終了コードと run のアクティビティの両方に基づいて run の状態を判断します。run が終了せず、新しいメトリクスの送信も約 5 分間ない場合、W&B はその run を CRASHED としてマークします。これは、トレーニングプロセスが応答しなくなった場合や、ログするのを停止した場合、または正常に終了せずに停止した場合 (SIGKILL を送信した場合など) に発生することがあります。run の状態を実際の状況と一致させるには、一定の間隔でメトリクスをログするか、明示的な終了コードで終了するようにしてください。

捕捉可能なシグナルとプリエンプション

プリエンプト可能な環境で送信されるシグナルのほとんどは捕捉可能です。つまり、トレーニングスクリプトでシグナルをインターセプトし、正常にシャットダウンできます。トレーニングスクリプトには独自のシグナルハンドラを登録できます。システムが捕捉可能なシグナルを送信すると、登録したハンドラが実行されます。W&B がすでに受信したメトリクスは保持されます。エージェントはプロセスの終了を検出すると、次の run を開始します。 ベストプラクティス:

ハンドラはできるだけ早い段階 (たとえば、メインのトレーニングループに入る前) で登録してください。
プリエンプション後に run をキューに入れ直す場合は、ハンドラ内で wandb.run.mark_preempting() を呼び出してください。その後、クリーンアップ (たとえば、チェックポイントの保存) を実行し、0 以外の終了コードで終了します。

次の例では、SIGUSR1 (クラスターでよく使われるプリエンプションシグナル) と SIGTERM のハンドラを登録します。SIGINT は対話的な操作 (たとえば、ターミナルからの手動キャンセル) のために残しておきます。ハンドラは wandb.run.mark_preempting() を呼び出し、128 + signum を終了コードとして終了します。

import signal
import sys
import wandb


def signal_handler(signum, frame):
    if wandb.run is not None:
        # オプション: モデル チェックポイントの保存やバッファのフラッシュなどを行います。
        print(f"Preempted with signal: {signal.Signals(signum).name}.")
        wandb.run.mark_preempting()
    sys.exit(128 + signum)


def train():
    signal.signal(signal.SIGUSR1, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    with wandb.init() as run:
        config = wandb.config
        for epoch in range(100):
            # トレーニングステップ。必要に応じて wandb.log(...) を呼び出します
            pass


if __name__ == "__main__":
    train()

`SIGKILL` (捕捉不可)

SIGKILL はオペレーティングシステムのカーネルから送信されるシグナルで、捕捉することも無視することもできません。プロセスは即座に終了するため、ハンドラや atexit コールバックが実行されることはありません。W&B はその run の最終サマリーを書き込めません。エージェントは復旧して sweep を続行しますが、その run のデータは不完全になります。SIGKILL は最後の手段としてのみ使用してください。正常なシャットダウンが必要な場合は、SIGTERM または SIGINT を使用することをお勧めします。

エージェントから子プロセスへのシグナル転送

wandb agent CLI を使用すると、エージェントはトレーニングスクリプトを子プロセスとして実行します。エージェントを中断しても (Ctrl+C を押した場合や、スケジューラーがジョブに SIGTERM を送信した場合など)、デフォルトでは子プロセス (トレーニングプロセス) にはシグナルが届きません。そのため、トレーニングスクリプトはハンドラを実行することも、mark_preempting() を呼び出すこともできません。詳細については、wandb GitHub issue #3667 を参照してください。子プロセスを正常にシャットダウンさせ、ハンドラ内で wandb.run.mark_preempting() を呼び出せるようにするには、--forward-signals を指定して CLI エージェントを実行します。

wandb agent --forward-signals entity/project/sweep_ID

W&B は、Python API の wandb.agent() ではシグナル転送をサポートしていません。この方法では、トレーニング関数は別の子プロセスとしてではなくスレッド内で実行されるため、同じ転送動作は適用されません。転送が有効な状態で CLI エージェントが SIGINT または SIGTERM を受信すると、エージェントはそのシグナルを子プロセスに中継します。これにより、トレーニングスクリプトのハンドラが実行され、wandb.run.mark_preempting() を呼び出したうえで、必要に応じて 0 以外の終了コードを指定して wandb.finish() を呼び出し、0 以外のコードで終了できます。エージェントプロセスで Ctrl+C を 2 回押すと、デフォルトではエージェントは SIGTERM を受信します。--forward-signals を指定すると、エージェントが SIGINT を子プロセスに転送できるため、ハンドラが実行されます。詳細については、wandb agent の CLI リファレンスを参照してください。

SLURM などのプリエンプト可能なクラスター

このセクションでは、SLURM、EC2 Spot、Google Cloud のプリエンプト可能な VM などのクラスターで、プリエンプションが発生しても run を継続できるように sweep を設定する方法を説明します。プリエンプション発生時には、トレーニングプロセスがシグナルを受信し、run をプリエンプト中としてマークしたうえで、0 以外の終了コードで終了する必要があります。これにより、W&B が run をキューに入れ直します。その後、新しいエージェント (またはジョブがキューに入れ直された後の同じエージェント) が run を再開できます。 トレーニングプロセスがシグナルを受信できるようにする:

スケジューラーがエージェントにシグナルを送る場合: wandb agent --forward-signals でエージェントを実行します。これにより、スケジューラー (またはユーザー) がエージェントにシグナルを送ると、エージェントがそのシグナルを子プロセスに転送します。子プロセスのハンドラでは、wandb.run.mark_preempting()、0 以外のコードを指定した wandb.finish(exit_code=...)、および sys.exit(128 + signum) (または別の 0 以外の終了コード) を呼び出せます。
スケジューラーが (エージェントではなく) 起動スクリプトにシグナルを送る場合: 起動スクリプトからトレーニングプロセスへプリエンプションシグナルを直接送信するようにします。たとえば、トレーニングスクリプトが自身のプロセス ID をファイルに書き込み、起動スクリプトがクラスターのシグナル (たとえば SIGUSR1) をトラップして kill -SIGUSR1 $(cat $PID_FILE) を実行することで、トレーニングプロセスのハンドラが実行されます。

トレーニングスクリプト内: クラスターが使用するシグナル (たとえば SIGTERM や SIGUSR1) のハンドラを登録します。ハンドラ内では、run が実行中であれば wandb.run.mark_preempting() を呼び出してから、0 以外の終了コードで run を終了し、sys.exit(128 + signum) (または別の 0 以外のコード) を呼び出します。これにより、W&B が run をキューに入れ直します。W&B が run をキューに入れ直す条件や、mark_preempting() との関係について詳しくは、プリエンプト可能な Sweeps run を再開するを参照してください。 sweep の状態: エージェントを起動する前に wandb sweep entity/project/sweep_ID --resume を実行して sweep を再開モードにし、キューに入れ直された run が割り当てられるようにします。 複数エージェントの調整: 多数のエージェントを同時に実行する場合 (SLURM のアレイジョブなど)、プリエンプトされた同じ run を複数のエージェントが奪い合う競合が発生することがあります。これは既知の制限事項です。回避策として、エージェントの起動タイミングをずらすか、ロックなどの外部の調整メカニズムを使用してください。マルチ GPU の SLURM ジョブで wandb.agent() を 1 つのプロセスからのみ呼び出す必要がある場合は、SLURM で sweep を実行するにはどうすればよいですか? を参照してください。

`wandb sweep --cancel`

キャンセルは OS シグナルを直接送信する場合とは動作が異なります。このセクションでは、--cancel コマンドがシグナルや子プロセスとどのように関係するかを説明します。sweep のキャンセルには、OS シグナルではなく W&B API を使用します。wandb sweep --cancel entity/project/sweep_ID のようなコマンドを実行してください。サーバーがエージェントに終了を指示すると、エージェントは実行中の子プロセスを終了させてから停止します。キャンセルが反映されるまでに、エージェントの API ポーリング間隔程度の短い遅延が発生する場合があります。キャンセルすると、run に SIGKILL が送信されます。そのため、子プロセスはユーザー定義のシグナルハンドラを実行できません。Sweeps UI の Cancel コントロールを使用した場合も同様です。--cancel は、sweep 全体を停止してキャンセル済みとしてマークしたい場合に使用してください。現在の run を正常にシャットダウンするには、run に捕捉可能なシグナルを送信します (または CLI エージェントで --forward-signals を使用し、エージェントにシグナルを送信します) 。sweep を正常に完了させるには、--cancel ではなく wandb sweep --stop を使用してください。一時停止、再開、停止、キャンセルの各オプションの詳細については、sweep を管理するを参照してください。

エージェントへのシグナルと run へのシグナルの違い

エージェントへのシグナル送信とトレーニング run へのシグナル送信の違いを理解しておくと、孤立プロセスの発生や予期しない動作を防げます。子のトレーニングプロセスではなくエージェントプロセスにシグナルを送信すると、エージェントは終了しても、子プロセスが孤立プロセスとして実行され続ける場合があります。この孤立プロセスがターミナルに出力し続けることもあり、Enter キーを押すまでシェルに新しいプロンプトが表示されない場合もあります。 CLI エージェントで --forward-signals を使用しない限り、エージェントを停止しても子のトレーニングプロセスが停止するとは限りません。エージェントが終了したかどうかは、プロンプトが表示されたかどうかで判断せず、ps -p [AGENT-PID] や pgrep -f "wandb agent" などの OS コマンドで確認してください。

リファレンス: `mark_preempting()` と最終的な run の状態

次の表は、mark_preempting() を呼び出すタイミングとプロセスの終了のしかたによって、run の状態がどのように決まるかをまとめたものです。この表は、wandb agent CLI を使用し、トレーニングプログラムをそのサブプロセスとして実行することを前提としています。

シナリオ	`mark_preempting()` なし	シグナルハンドラが `mark_preempting()` を呼び出し、ゼロ以外で終了	`init()` の直後に常に `mark_preempting()` を呼び出す
run が終了コード 0 で正常に完了	FINISHED	FINISHED	FINISHED
run がゼロ以外の終了コードで失敗	FAILED	FAILED	PREEMPTED
run が `SIGKILL` を受信	約 5 分後に CRASHED	約 5 分後に CRASHED (捕捉不可)	約 5 分後に PREEMPTED
run が `SIGINT` を受信	KILLED	PREEMPTED (`SIGINT` ハンドラがある場合)	PREEMPTED
run が別のシグナル (`SIGTERM` や `SIGUSR1` など) を受信	約 5 分後に CRASHED	PREEMPTED (対応するハンドラがある場合)	約 5 分後に PREEMPTED

mark_preempting() をシグナルハンドラ内でのみ呼び出す場合、SIGKILL のようにハンドラが実行されないケースには対応できません。 wandb.init() の直後に常に mark_preempting() を呼び出すと、W&B はあらゆる失敗をプリエンプションとして扱う可能性があります。その結果、バグや誤った設定が原因の失敗であっても、run が繰り返しキューに入れ直されることがあります。プリエンプションシグナルが明確に定義されている環境では、init() の後に無条件で呼び出すのではなく、mark_preempting() を呼び出してゼロ以外で終了するシグナルハンドラを使用するのが一般的です。

​終了ステータスとシグナル

​古い run とサーバー側のタイムアウト

​捕捉可能なシグナルとプリエンプション

​SIGKILL (捕捉不可)

​エージェントから子プロセスへのシグナル転送

​SLURM などのプリエンプト可能なクラスター

​wandb sweep --cancel

​エージェントへのシグナルと run へのシグナルの違い

​リファレンス: mark_preempting() と最終的な run の状態