Fabricでの処理中にサーバを再起動させる

TensorFlowで株価予想シリーズ を実装してるのですが、コーディングは Mac で行い、コードを GPU（TITAN X） が搭載された PC に rsync で同期してから実行しています。しかし、GPU を使ったプログラムを200回ほど実行すると cuda(7.5) のライブラリがエラーを出してそれ以降は全て失敗するようになります。そこでエラーを出すようになったら手動でサーバを再起動してたのですが、Fabric を使って自動で再起動するようにしました。

まずは全体のコードはこちら。

fabfile.py

import time
from fabric.api import run, cd, prefix, env
from fabric.tasks import execute


# -i を付けることでsshでログインした時に .bashrc に設定した cuda などのパス設定が有効になる
env.shell = '/bin/bash -l -i -c'
# サーバアドレス
env.hosts = ['10.0.1.17']
# ログインユーザー名（秘密鍵認証なのでパスワードはなし）
env.user = 'akiraak'


# ssh でアクセスしてコマンド実行
def run_jp():
    # プロジェクトのディレクトリに移動
    with cd('~/project/'):
        # virtualenv の設定を有効にする
        with prefix('. ~/tensorflow-env/bin/activate'):
            # コマンド実行
            run('time python run_jp.py', warn_only=True)
            # サーバ再起動（パスワード無しsudo設定を有効にしておくこと）
            run('sudo reboot')


if __name__ == "__main__":
    # サーバ上のコマンド実行と再起動を繰り返す
    while True:
        try:
            # ssh実行
            execute(run_jp)
        except Exception as e:
            print(e)
        # サーバが再起動するまでしばし待つ
        time.sleep(60)

Fabric はサーバに ssh でログインしてのコマンド実行を python のコードで記述できるものです。

Mac からこのコードを実行して、サーバを制御します。

run_jp() が ssh にアクセスした後にコマンドを実行する関数になります。通常だと以下のコマンドで処理を実行します。

$ fab run_jp

今回は「処理を実行」→「サーバ再起動」→「処理を実行」としたいので、fab コマンドは使わないで、普通に python コードを実行するかたちにします。

$ python fabfile.py

ここで run_jp() の実行と、再起後に60秒待つ処理を行います。

if __name__ == "__main__":
    # サーバ上のコマンド実行と再起動を繰り返す
    while True:
        try:
            # ssh実行
            execute(run_jp)
        except Exception as e:
            print(e)
        # サーバが再起動するまでしばし待つ
        time.sleep(60)

実際の処理を行う run_jp() では、コードがあるディレクトリに移動し(~/project/)、virtualenv の設定を有効(. ~/tensorflow-env/bin/activate)にしたのち、python run_jp.py を実行しています。

python run_jp.pyの中では別プロセスで複数の処理を実行してるのですが、前述のとおり途中で cuda がエラーを出すので一部の処理は失敗した状態で終わります。そこで最後に sudo reboot を行います。その後は if __name__ == "__main__": に戻り、sleepで再起動を待った後に、また同じコマンドを実行します。

# ssh でアクセスしてコマンド実行
def run_jp():
    # プロジェクトのディレクトリに移動
    with cd('~/project/'):
        # virtualenv の設定を有効にする
        with prefix('. ~/tensorflow-env/bin/activate'):
            # コマンド実行
            run('time python run_jp.py', warn_only=True)
            # サーバ再起動（パスワード無しsudo設定を有効にしておくこと）
            run('sudo reboot')