Running tasks with a concurrency limit

pool/pool.py

"""Twenty report downloads, at most 3 at a time: five ways to run them, timed on a fake clock."""
import asyncio

from fake_clock import FakeClock
from reports import download_tasks


async def run_limited(tasks, k):
    """Runs every task, at most k at a time, and returns the results in input order."""
    results = [None] * len(tasks)
    next_index = 0

    async def worker():
        nonlocal next_index
        while next_index < len(tasks):
            i = next_index              # take the next task: no await between this
            next_index += 1             # read and this write, so no other worker gets i
            results[i] = await tasks[i]()

    await asyncio.gather(*(worker() for _ in range(min(k, len(tasks)))))
    return results


async def with_semaphore(tasks, k):
    """Starts a task for every download at once; each waits for one of k permits."""
    permits = asyncio.Semaphore(k)

    async def limited(task):
        async with permits:             # released even when the task raises
            return await task()

    return await asyncio.gather(*(limited(t) for t in tasks))


async def one_at_a_time(tasks, k):
    return [await t() for t in tasks]


async def all_at_once(tasks, k):
    return await asyncio.gather(*(t() for t in tasks))


async def in_batches(tasks, k):
    results = []
    for start in range(0, len(tasks), k):
        results += await asyncio.gather(*(t() for t in tasks[start:start + k]))
    return results


async def main():
    expected = [f"report-{i:02}.csv" for i in range(1, 21)]
    print("way to run them        finish   most at once   results in input order")
    for name, run in [("one at a time", one_at_a_time), ("all at once", all_at_once),
                      ("batches of 3", in_batches), ("pool of 3 workers", run_limited),
                      ("semaphore of 3", with_semaphore)]:
        clock = FakeClock()
        tasks = download_tasks(clock)
        results = await clock.run(run(tasks, 3))
        print(f"{name:20} {clock.now:6} s {clock.most:14}   {'yes' if results == expected else 'no'}")


asyncio.run(main())

way to run them        finish   most at once   results in input order
one at a time            66 s              1   yes
all at once               8 s             20   yes
batches of 3             37 s              3   yes
pool of 3 workers        25 s              3   yes
semaphore of 3           25 s              3   yes

pool/fake_clock.py

"""A fake clock for asyncio: simulated seconds that pass only while every task waits.

Tasks call `await clock.sleep(seconds)` where real code would wait for a download.
`await clock.run(main())` runs a program: each time every task is waiting, it moves the
time on to the earliest wake-up and wakes that sleeper. Nothing really waits, so the
times a program prints are the same in every run and on every computer.
"""
import asyncio
import heapq
import itertools

# Turns of the event loop that let every task run on until it waits again. The programs
# of this lesson need at most 6; 12 leave room for a longer chain of awaits.
SETTLE_TURNS = 12


class FakeClock:
    def __init__(self):
        self.now = 0
        self._sleepers = []                 # (wake-up time, order, future), earliest first
        self._order = itertools.count()

    async def sleep(self, seconds):
        wake = asyncio.get_running_loop().create_future()
        heapq.heappush(self._sleepers, (self.now + seconds, next(self._order), wake))
        await wake

    async def run(self, main):
        """Runs the coroutine main to its end, and returns its result or raises its error."""
        outcome = {}

        async def capture():
            try:
                outcome["result"] = await main
            except Exception as error:      # handed to the caller below
                outcome["error"] = error

        task = asyncio.ensure_future(capture())
        while True:
            for _ in range(SETTLE_TURNS):
                await asyncio.sleep(0)
            if task.done():
                if "error" in outcome:
                    raise outcome["error"]
                return outcome["result"]
            if not self._sleepers:
                task.cancel()
                raise RuntimeError(f"stuck at {self.now} s: every task waits, and none of them for the clock")
            self.now, _, wake = heapq.heappop(self._sleepers)
            wake.set_result(None)

pool/reports.py

"""Twenty synthetic report downloads: report i takes SECONDS[i] simulated seconds."""
SECONDS = [8, 3, 3, 6, 1, 5, 2, 1, 2, 1, 3, 5, 2, 4, 6, 1, 8, 2, 1, 2]


def download_tasks(clock, seconds=SECONDS):
    """One task per report: a function that starts the download when called. Each task
    records how many downloads are running at once, in clock.running and clock.most."""
    clock.running = clock.most = 0

    def task(i):
        async def download():
            clock.running += 1
            clock.most = max(clock.most, clock.running)
            await clock.sleep(seconds[i])
            clock.running -= 1
            return f"report-{i + 1:02}.csv"
        return download

    return [task(i) for i in range(len(seconds))]

pool/pool.mjs

// Twenty report downloads, at most 3 at a time: four ways to run them, timed on a fake clock.
import { FakeClock } from './fake_clock.mjs';
import { downloadTasks } from './reports.mjs';

// Runs every task, at most k at a time, and resolves to the results in input order.
export async function runLimited(tasks, k) {
  const results = new Array(tasks.length);
  let next = 0;
  async function worker() {
    while (next < tasks.length) {
      const i = next++; // read and increment with no await between: no other worker gets i
      results[i] = await tasks[i]();
    }
  }
  const workers = Array.from({ length: Math.min(k, tasks.length) }, worker);
  await Promise.all(workers);
  return results;
}

async function oneAtATime(tasks) {
  const results = [];
  for (const t of tasks) results.push(await t());
  return results;
}

const allAtOnce = (tasks) => Promise.all(tasks.map((t) => t()));

async function inBatches(tasks, k) {
  const results = [];
  for (let start = 0; start < tasks.length; start += k) {
    results.push(...(await Promise.all(tasks.slice(start, start + k).map((t) => t()))));
  }
  return results;
}

const expected = Array.from({ length: 20 }, (_, i) => `report-${String(i + 1).padStart(2, '0')}.csv`);
console.log('way to run them        finish   most at once   results in input order');
for (const [name, run] of [['one at a time', oneAtATime], ['all at once', allAtOnce], ['batches of 3', inBatches], ['pool of 3 workers', runLimited]]) {
  const clock = new FakeClock();
  const tasks = downloadTasks(clock);
  const results = await clock.run(run(tasks, 3));
  const same = results.length === expected.length && results.every((r, i) => r === expected[i]);
  console.log(`${name.padEnd(20)} ${String(clock.now).padStart(6)} s ${String(clock.most).padStart(14)}   ${same ? 'yes' : 'no'}`);
}

way to run them        finish   most at once   results in input order
one at a time            66 s              1   yes
all at once               8 s             20   yes
batches of 3             37 s              3   yes
pool of 3 workers        25 s              3   yes

pool/fake_clock.mjs

// A fake clock for promises: simulated seconds that pass only while every task waits.
//
// Tasks call `await clock.sleep(seconds)` where real code would wait for a download.
// `await clock.run(promise)` runs a program: each time every task is waiting (no promise
// job is left to run), it moves the time on to the earliest wake-up and wakes that
// sleeper. Nothing really waits, so the times a program prints are the same in every run.

export class FakeClock {
  now = 0;
  #sleepers = []; // { at, order, wake }, kept earliest first
  #order = 0;

  sleep(seconds) {
    return new Promise((wake) => {
      this.#sleepers.push({ at: this.now + seconds, order: this.#order++, wake });
      this.#sleepers.sort((a, b) => a.at - b.at || a.order - b.order);
    });
  }

  async run(promise) {
    let done = false;
    promise.then(() => { done = true; }, () => { done = true; });
    for (;;) {
      // A timer fires only once every promise job has run: every task now waits.
      await new Promise((settled) => setTimeout(settled, 0));
      if (done) return promise;
      const next = this.#sleepers.shift();
      if (!next) throw new Error(`stuck at ${this.now} s: every task waits, and none of them for the clock`);
      this.now = next.at;
      next.wake();
    }
  }
}

pool/reports.mjs

// Twenty synthetic report downloads: report i takes SECONDS[i] simulated seconds.
export const SECONDS = [8, 3, 3, 6, 1, 5, 2, 1, 2, 1, 3, 5, 2, 4, 6, 1, 8, 2, 1, 2];

// One task per report: a function that starts the download when called. Each task
// records how many downloads are running at once, in clock.running and clock.most.
export function downloadTasks(clock, seconds = SECONDS) {
  clock.running = 0;
  clock.most = 0;
  return seconds.map((s, i) => async () => {
    clock.running += 1;
    clock.most = Math.max(clock.most, clock.running);
    await clock.sleep(s);
    clock.running -= 1;
    return `report-${String(i + 1).padStart(2, '0')}.csv`;
  });
}