通过SLURM脚本运行时，如何让Python子进程等待输入并续跑？

1. 利用SLURM自定义信号触发续跑 #

SLURM支持向运行中作业发送自定义信号，你可以在Python代码里监听SIGUSR1这类非系统关键信号，修复许可证问题后通过SLURM命令发送信号，让程序退出等待状态续跑。

实现示例：

import signal
import time

resume_trigger = False

def signal_handler(signum, frame):
    global resume_trigger
    resume_trigger = True

# 注册信号监听
signal.signal(signal.SIGUSR1, signal_handler)

try:
    # 你的核心任务代码
    execute_main_task()
except LicenseRelatedError:  # 替换为实际捕获的许可证异常类型
    print("检测到许可证问题，等待修复后继续...")
    while not resume_trigger:
        time.sleep(10)  # 低间隔轮询，减少资源消耗
    # 从断点恢复执行
    resume_from_breakpoint()

修复完成后，在终端执行：scontrol signal SIGUSR1 <你的SLURM作业ID>，Python进程就会收到信号并续跑。

2. 借助SLURM作业的标准输入管道 #

SLURM作业默认保留标准输入通道，你可以让Python在捕获异常后等待特定输入指令，修复问题后直接通过SSH或SLURM命令向作业发送触发信号。

实现示例：

try:
    execute_main_task()
except LicenseRelatedError:
    print("许可证异常，请修复后输入'resume'继续...")
    while True:
        input_cmd = input().strip()
        if input_cmd == "resume":
            break
    resume_from_breakpoint()

修复后，找到作业所在节点和Python进程PID，执行：ssh <节点名称> "echo resume > /proc/<PID>/fd/0"，或者用SLURM命令：srun --jobid <作业ID> --pty echo resume发送触发指令。

3. 自动查询SLURM许可证状态重试 #

如果许可证是集群共享资源，可让Python自动调用SLURM工具查询许可证可用性，满足条件后自动续跑，无需手动干预。

实现示例：

import subprocess
import time

def check_license_available():
    # 替换为实际的许可证查询命令，比如检查指定许可证的剩余数量
    check_cmd = ["sacctmgr", "show", "license", "where", "Name=target_license", "--format=Used,Total"]
    result = subprocess.run(check_cmd, capture_output=True, text=True)
    if result.returncode != 0:
        return False
    # 解析输出获取已用/总数
    used, total = map(int, result.stdout.strip().split("\n")[-1].split())
    return used < total

try:
    execute_main_task()
except LicenseRelatedError:
    print("许可证不足，等待可用资源...")
    while not check_license_available():
        time.sleep(60)  # 每分钟检查一次
    resume_from_breakpoint()

4. 用文件锁替代内容轮询 #

比起轮询文件内容，检查文件锁的存在性更高效——修复完成后创建一个特定锁文件，Python通过判断文件是否存在来触发续跑，减少IO操作。

实现示例：

import os
import time

RESUME_LOCK = "/tmp/slurm_job_resume.lock"

try:
    execute_main_task()
except LicenseRelatedError:
    print("许可证问题，等待锁文件触发...")
    while not os.path.exists(RESUME_LOCK):
        time.sleep(10)
    # 续跑后删除锁文件，避免下次误触发
    os.remove(RESUME_LOCK)
    resume_from_breakpoint()

修复后执行touch /tmp/slurm_job_resume.lock即可触发程序续跑。

内容的提问来源于stack exchange，提问作者Lucas Moura Gomes