Python Subprocess Concurrency with AsyncIO

Share
Python Subprocess Concurrency with AsyncIO

Here's a useful function I find myself rewriting quite often when I want to automate some IO-bound task that benefits from async concurrency (not quite true parallelism, since CPU-bound tasks would still ultimately run sequentially with this code).

import asyncio
import os
import time


MAX_SUBPROCESSES = 10000  # handle subprocesses in batches of this many
active_subprocesses = 0


async def run_command(cmd):
    """Run a subprocess asynchronously. Wait until fewer than MAX_SUBPROCESSES are running."""
    global active_subprocesses

    while active_subprocesses >= MAX_SUBPROCESSES:
        await asyncio.sleep(.1)

    active_subprocesses += 1

    proc = await asyncio.create_subprocess_shell(
        cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
        shell=True,
    )
    stdout, _ = await proc.communicate()

    active_subprocesses -= 1

    print(cmd, "finished")

    return stdout.decode()

Example use-case for getting smartctl information for all devices under /dev/:

import glob

async def main_parallel():
    # Create a run_command coro for each device under /dev
    device_map = {}
    coros = []
    for dev in glob.glob("/dev/*"):
        # here is where we actually utilize our handy run_command function
        coros.append(run_command(f"echo '{dev}'; smartctl -a {dev}"))

    # Run all the coros in parallel
    # Note that this usage of `async for` requires at least Python 3.11
    async for result in asyncio.as_completed(coros):
        result_lines = result.result().split("\n")
        dev = result_lines[0]
        device_map[dev] = "\n".join(result_lines[1:])

    return device_map


device_map = asyncio.run(main_parallel())

Here's the corresponding sequential method and a time comparison:

import time
import subprocess

def main_sequential():
    device_map = {}
    for dev in glob.glob("/dev/*"):
        try:
            device_map[dev] = subprocess.check_output(f"smartctl -a {dev}", shell=True)
        except subprocess.CalledProcessError:
            continue
        print(dev, "finished")
    return device_map

time_sequential_start = time.process_time()
main_sequential()
total_sequential_time = time.process_time() - time_sequential_start

time_parallel_start = time.process_time()
asyncio.run(main_parallel())
total_parallel_time = time.process_time() - time_parallel_start

print("Sequential time", total_sequential_time)
print("Parallel time", total_parallel_time)
print(f"Parallel execution is {1 - (total_parallel_time / total_sequential_time):.2%} faster")

Results:

Sequential time 0.18521522
Parallel time 0.15665329800000005
Parallel execution is 15.42% faster

Obviously this will only be beneficial in quite specific scenarios where IO-bound tasks are independent of one another (although of course a topological sort can be utilized before-hand if some tasks do depend on one another). All-the-same, I tend to end up using some version of this run_command function whenever I need to automate a task on Linux.