Asyncio Terminate Subprocess on wait_for timeout

Question

I have a long running process and its children (in this example it is stress) that I wish to terminate after some time. I am using asyncio.wait_for since it's what the documentation suggests, but while the timeout occurs and the asyncio.TimeoutError is raised, the process is still running. I'm running on Python 3.8.10.

Here's my code:

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))

Can someone suggest a way to kill this process after the timeout? Thanks in advance! :D

Answer 1

I ended up solving it by modifying the functions a bit. Before, my run() function returned the output of the command. By returning the process proc, I could monitor the timeout for the proc.communicate(). This is the portion that waits until the process is done. If the process takes longer than timeout, then I ask if it is done, by looking into proc.returncode. If it is other than None it has finished. If it is None, then I recursively kill every child and finally the parent process itself.

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))