NNI 在执行并行任务时卡死在 WAITING 状态

Posted on In

在使用 NNI 自动跑实验，过了几个小时去看了一眼状态，结果发现有任务一直处于 WAITING 状态（如图所示），而实际上服务器的 GPU 并非处于全部被占用的状态。

经过查阅 issue 与查看源码，发现 nni 判定 WAITING 状态的任务在何时可以执行并将状态转变为 RUNNING 的条件是文件 /tmp/<user_name>/nni/script/gpu_metrics 中 gpuInfos 字段下各 GPU 的状态 activeProcessNum。由于服务器上有 GPU 实时监控软件在不断调用 nvidia-smi 程序，导致 nni 的检查 GPU 状态的程序一直卡在 nvidia-smi 处。

而 nni 中专门有个脚本可以用来检测 GPU 使用情况并更新 gpu_metrics 文件：/<python_path>/site-packages/nni/tools/gpu_tool/gpu_metrics_collector.py。查看代码可以看到：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def main(argv):
    metrics_output_dir = os.environ['METRIC_OUTPUT_DIR']

    cmd = 'nvidia-smi -q -x'.split()
    while(True):
        try:
            smi_output = subprocess.check_output(cmd)
        except Exception:
            traceback.print_exc()
            gen_empty_gpu_metric(metrics_output_dir)
            break
        parse_nvidia_smi_result(smi_output, metrics_output_dir)
        # TODO: change to sleep time configurable via arguments
        time.sleep(5)

因此，将环境变量 METRIC_OUTPUT_DIR 设定在 gpu_metrics 所在的目录，即可自动生成最新的 GPU 状态。在我这儿卡住的服务器上 kill 掉无响应的 nvidia-smi 程序，执行 METRIC_OUTPUT_DIR=/tmp/<user_name>/nni/script/ python3 -m nni.tools.gpu_tool.gpu_metrics_collector，成功地让一直卡在 WAITING 状态的程序继续运行，状态转为 RUNNING。

为了后续不被卡住，特意用了 crontab 定期执行一次杀掉 nvidia-smi 和执行 gpu_metrics_collector 的操作，一劳永逸。