sweep run のシグナル処理

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()

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