python手写守护进程（daemon）

2021.10.31 Updates

发现在用 gpu 跑实验程序的时候，运行报错之后，此 daemon 的线程函数 run 里用 os.system 执行 shell 的语句并没有除实验程序的报错、退出而结束，而是卡住了，进程没有结束，gpu 也一直占住，不知道什么原因。

面前的解决方法是，在代码里封装一个 main 函数，并手动处理 exception、退出：

# import traceback

def main():
    # blabla...
    pass

if "__main__" == __name__:
    # 封装虽丑陋但有用
    try:
        main()
    except:
        traceback.print_exc()
        exit(1)

当然也可以手动 kill 进程。

---

服务器的用法是提交一个 pbs 脚本（感觉就是 shell 脚本，只是开头多了一些 pbs 的指令）去排队，轮到的时候会执行 pbs 脚本的内容。

本文试图（部分）对抗服务器的排队机制：思路是让 pbs 脚本执行一个 daemon，让它挂着动态监测、调另外的 shell，并且就算目前所有 shell 都运行完， daemon 也保持一段等待时间才退出。在 shell 里面写要跑的命令，且可以多次动态加入，那只要 daemon 已开始、未退出，就可以随时加跑新任务，而不用重新排队。

要点：

• 固定输入 shell 的名字，如 zombie.sh；固定输入 shell 的摆放路径，即探测路径，如 ~/pool/。
  • （2022.1.11）现在把探测路径改为当前路径（即 ./）。当运行多个 daemons 时，在 pbs 脚本（其实排队的是 pbs 脚本，而由该 pbs 脚本运行此 daemon）里指定 daemon ID，并创建相应的一个独立 shell pool 目录，以隔离不同 daemons 的探测目录。
• 输入 shell 要指定一个 flag，表示已经写完命令，可以执行，防止未敲完命令就被拿去执行。如约定 shell 最后一定以 echo DONE 命令结束，且只出现一次，则当且仅当检测到存在 zombie.sh，且里面有句 echo DONE，daemon 才拿它去执行。
• 为固定 shell 名字且避免命名冲突，执行 shell 之前改掉文件名先，如加一个序号；运行结束之后删掉相应 shell 文件，保持简洁。
• deamon 死循环监测（中间小 sleep 一阵，如 30s），且在所有 shell 都运行完之后，还保持一段等待时间（如半个钟），不然又要重新排；但又不好等太久，别被人投诉或者被管理员发现了。
• 每个 shell 文件都开一个线程执行，全部执行完之前 daemon 都不能停（无视等待时间超时），且每个 shell 执行完之后，都更新一下等待时间的开始计时时间（即 tic）。

Code

import argparse
import os
import os.path as osp
import re
import shutil
import subprocess
import threading
import time


# parser = argparse.ArgumentParser(description='daemon')
# parser.add_argument('-i', '--id', '--daemon-id', dest="daemon_id", type=int, default=0,
#                     help="seperate the script pools for multiple daemons")
# args = parser.parse_args()

DETECT_PATH = "."
# if args.daemon_id > 0:
#     DETECT_PATH += str(args.daemon_id)
INPUT_SHELL = osp.join(DETECT_PATH, "zombie.sh")
RENAME_TEMPLATE = osp.join(DETECT_PATH, "zombie.{}.sh")
TIMEOUT = 60 * 60  # seconds
SLEEP = 7  # seconds


fin_flag = re.compile(r"\necho\s+\"*DONE\"*")
tic, mutex_tic = time.time(), threading.Lock()
n_running, mutex_run = 0, threading.Lock()
logger = open("log.txt", "w")
logger.write("DAEMON START: {}\n".format(time.strftime(
    "%Y-%m-%d %H:%M:%S", time.localtime())))


def run(tid):
    global tic, n_running, logger

    # starting log
    _timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    _info_str = "START #{} at {}".format(tid, _timestamp)
    print(_info_str)
    logger.write(_info_str + "\n")
    logger.flush()

    # running counting +1
    with mutex_run:
        n_running += 1

    new_f = RENAME_TEMPLATE.format(tid)
    os.system("bash {}".format(new_f))
    os.remove(new_f)

    # update global tic
    with mutex_tic:
        tic = max(tic, time.time())

    # running counting -1
    with mutex_run:
        n_running -= 1

    # ending log
    _timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    _info_str = "END #{} at {}".format(tid, _timestamp)
    print(_info_str)
    logger.write(_info_str + "\n")
    logger.flush()


n_thread = 0
while True:
    if osp.exists(INPUT_SHELL):
        try:
            with open(INPUT_SHELL, "r") as _f:
                _txt = _f.read()
            if len(fin_flag.findall(_txt)) > 0:
                # rename shell -> avoid name conflict
                new_f = RENAME_TEMPLATE.format(n_thread)
                os.rename(INPUT_SHELL, new_f)

                # run
                t = threading.Thread(target=run, args=[n_thread])
                t.start()
                # t.join()  # NOT now, do it later

                # update tic
                mutex_tic.acquire()
                tic = max(tic, time.time())
                mutex_tic.release()

                n_thread += 1
        except:
            pass

    toc = time.time()
    if (n_running < 1) and (toc - tic > TIMEOUT):
        logger.write("* TIME OUT\n")
        break

    time.sleep(SLEEP)

# for t in thread_list:
#     t.join()

logger.write("DAEMON END: {}\n".format(time.strftime(
    "%Y-%m-%d %H:%M:%S", time.localtime())))
logger.flush()
logger.close()

References

1. 多线程
2. Python 多线程
3. Python 日期和时间
4. 由bibtex生成引用文献字符串
5. 【Python学习笔记】（八）多线程与并行：_thread模块、threading模块、Queue模块；os模块、subprocess模块、multiprocessing.Process模块
6. 编写健壮的 Shell 脚本