複数の実行中プロセスをテキストファイルで簡易的に把握するPythonのデコレータ

例えば複数のジョブをサーバに投げている時、いま自分がどのジョブを投げているのかを忘れてしまうことがある。
自分は深層学習の実験用スクリプトを投げた時に、どの実験を投げているのか結構忘れる。

もちろんmlflowなどの実験管理ツールを使えば、どのモデルの学習がどれくらい進んでいるかはある程度把握できるのだけど、「学習用スクリプトをどんなコマンドでいつどのサーバで投げたか」あたりの情報はあんまり管理できない(たぶん)。
あとローカルサーバでいいとはいえ、確認のために何か立ち上げたりするのは面倒。

なので、学習用スクリプトが投げられた時に、main()をラップすることで、どのプロセスがいつどう立ち上がったかを把握できるデコレータを書いた。

以下のような感じで使える、memoize_runnnig_processesが今回作ったデコレータ。
memoizeという名前がいいのかは微妙。

hoge.py

@memoize_runnnig_processes
def main():
    exp = HogeExperiment()
    exp.run()
    
if __name__ == "__main__":
    main()

これをやると、カレントディレクトリに.runnning-processes.txtというテキストファイルに実行中プロセスの情報を書き込む、その中身はこんな感じになる。

.runnning-processes.txt

server1	2021年5月1日 3時17分36.110417秒	hoge.py arg1
server2	2021年5月1日 9時52分0.868838秒	huga.py arg1 arg2 arg3

状態の管理は簡易的にテキストファイルでやることにした。
mmapとかいろいろな手段があると思うけど、NFS上の複数サーバ構成とかも統一的に扱いたい場合はファイルベースでやり取りするのがいいんじゃないかという気がする。

コードは以下の通り、普通のデコレータ。
排他ロックをとって、テキストファイルの各行を生きているプロセスに対応させている。Pythonのfinallyで、途中で投げているスクリプトが死んでも、きっちりテキストファイルの該当行を消してから死ぬようにした。

memoize_runnnig_processes.py

import os, sys, fcntl
from datetime import datetime
from pathlib import Path


def memoize_runnnig_processes(func):
    memory = Path("./.runnning-processes.txt")
    memory.touch()

    nodename = os.uname().nodename
    now = datetime.today().strftime('%Y年%-m月%-d日 %-H時%-M分%-S.%f秒')
    args = " ".join(sys.argv)
    l = f"{nodename}\t{now}\t{args}"

    def inner(*args, **kwargs):
        # add process status
        try:
            with memory.open("r+") as f:
                fcntl.flock(f, fcntl.LOCK_EX)

                lines = set(line.strip() for line in f.readlines() if line.strip())
                lines.add(l)
                lines = sorted(lines)

                f.truncate(0)
                f.seek(os.SEEK_SET)
                f.write("\n".join(lines))
                f.flush()

            return func(*args, **kwargs)

        # remove process status
        finally:
            with memory.open("r+") as f:
                fcntl.flock(f, fcntl.LOCK_EX)

                lines = set(line.strip() for line in f.readlines() if line.strip())
                lines.discard(l)
                lines = sorted(lines)

                f.truncate(0)
                f.seek(os.SEEK_SET)
                f.write("\n".join(lines))
                f.flush()

    return inner

.runnning-processes.txtは事前に作っておいてもいいけど、touchすればファイルがない場合は作成、みたいな挙動を実現できて便利。

発展として、Pythonのデコレータで書くんじゃなくてコマンドを作っちゃうみたいな方針もあると思う。
memo python hoge.pyみたいな感じで動かせるやつ。時間あったらやろうと思う。

もっといいやり方もあると思うので、何かあったら教えてください🙏