Asyncio 在 wait_for 超时时终止子进程

Asyncio Terminate Subprocess on wait_for timeout

我有一个很长的 运行 进程及其子进程(在这个例子中是 stress),我希望在一段时间后终止。我正在使用 asyncio.wait_for,因为它是文档所建议的,但是当发生超时并且引发 asyncio.TimeoutError 时,该过程仍然是 运行。我 运行 Python 3.8.10.

这是我的代码:

import asyncio

async def run(cmd):
    print("Running: ", cmd)

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE)

    stdout, stderr = await proc.communicate()
    stdout = stdout.decode('UTF-8')
    stderr = stderr.decode('UTF-8')
    
    return (stdout, stderr, proc.pid, proc.returncode)

async def run_with_timeout(cmd, timeout=20):
    task = asyncio.wait_for(run(cmd), timeout=timeout)
    
    try:
        output = await task
        stdout, _, pid, _ = output
        return str(stdout).strip()
    except asyncio.TimeoutError as e:
        print("Terminating Process '{0}' (timed out)".format(cmd))
        
asyncio.run(run_with_timeout(['stress', '--cpu', '2'], timeout=5))

有人可以建议一种在超时后终止此进程的方法吗?提前致谢! :D

我最终通过稍微修改函数解决了它。之前,我的 run() 函数返回命令的输出。通过返回进程 proc,我可以监控 proc.communicate() 的超时。这是等待过程完成的部分。如果该过程花费的时间超过 timeout,那么我会通过查看 proc.returncode 来询问是否已完成。如果它不是 None,它就完成了。如果是 None,那么我递归地杀死每个子进程,最后杀死父进程本身。

async def run(cmd):
    print("Running: ", cmd)

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE)

    return proc

async def run_with_timeout(cmd, timeout=20):
    proc = await run(cmd)

    try:
        output = await asyncio.wait_for(proc.communicate(), timeout=timeout)
        stdout, _ = output
        return str(stdout).strip()
    except asyncio.TimeoutError:
        if proc.returncode is None:
            parent = psutil.Process(proc.pid)
            for child in parent.children(recursive=True): 
                child.terminate()
            parent.terminate()
            print("Terminating Process '{0}' (timed out)".format(cmd))
        
        
asyncio.run(run_with_timeout(['stress', '--cpu', '2'], timeout=5))