matplotlibを使ったpythonジョブでsshサーバーを抜けたあとでもエラーが出ないようにbackendをうまく切り替える

エラーになる条件

以下のmatplotlibで図を保存するpythonコードを考える。

check.py

import time
import matplotlib
import matplotlib.pyplot as plt

# あとでsshサーバーを抜けるための猶予
time.sleep(10)

print(matplotlib.get_backend())

fig = plt.figure(figsize=(3.0, 3.0))
plt.plot([0, 1, 2, 3], [0, 1, 2, 3])
plt.savefig("hoge.png")

このコードを

$ python3 check.py
$ python3 check.py $ 

のように普通にフォアグラウンド、バックグラウンドで実行しても問題なく図は保存される。
しかしながら、sshサーバー上で次のようにバックグラウンドジョブで投げたあとにサーバーを抜けるとtkinterのdisplay周りのエラーで落ちる(sshの-Xや-Yオプション関係なく)。

$ nohup python3 check.py > log 2>&1 & # sleepで止まっている間にサーバを抜ける

log

TkAgg

Traceback (most recent call last):
  File "check.py", line 9, in <module>
    fig = plt.figure(figsize=(3.0, 3.0))
...
    window = tk.Tk(className="matplotlib")
  File "$path/lib/python3.7/tkinter/__init__.py", line 2023, in __init__
    self.tk = _tkinter.create(screenName, baseName, className, interactive, wantobjects, useTk, sync, use)
_tkinter.TclError: couldn't connect to display "localhost:10.0"

これは今のbackendにTkAggのようなGUIを使うbackendを指定していることが原因なので、コード内で明示的にmatplotlib.use('Agg')を指定するか、

$ python3 -c "import matplotlib ;print(matplotlib.get_configdir())"

で確認できる場所にあるmatplotlibrcのbackendをAggにするか、plt.ioff()を使いインタラクティブモードをoffにするかなどをして回避する必要がある。

いずれにしてもサーバーを抜けるようなバックグラウンドジョブを実行するときだけbackendを書き換える方法では反映し忘れたり、戻し忘れたりするのでnohupを使いバックグラウンドジョブを実行するときだけbackendをAggに切り替えられるようにしたい。

解決策

nohupを指定したバックグラウンドジョブを実行しても特に新たな環境変数は作成されないので自分で環境変数を適当に指定する。
このときにexportコマンドで指定してしまうとログアウトするまで環境変数が残ってしまうので上記のように実行時にのみ反映されるようにする。

$ BG=1 nohup python3 check2.py &
or
$ alias nohup=`BG=1 nohup` # bashrcやzshrcに追加
$ nohup python3 check2.py &

pythonコードは次のように環境変数BGの有無によってbackendをAggに変更する処理を入れる。

check2.py

import os
import time
import matplotlib

# 環境変数BGがある場合のみ実行される
if os.getenv('BG') != None:
    matplotlib.use('Agg')  # backendを明示的に指定する場合は import matplotlib.pyplot より前に書く

import matplotlib.pyplot as plt

time.sleep(10)

print(matplotlib.get_backend())

fig = plt.figure(figsize=(3.0, 3.0))
plt.plot([0, 1, 2, 3], [0, 1, 2, 3])
plt.savefig("hoge.png")

BGという環境変数が存在しなければos.getenv('BG')でNoneが返るので、存在するときだけ上記のmatplotlib.use('Agg')が実行され、サーバーを抜けたあとでも正常に図が保存できる。なお、plt.show()があるようなコードではbackendがAggの場合はエラーになるので上記の環境変数かmatplotlib.get_backend()結果を使い実行させないようにする必要がある。