The power of two choices

# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The power of two choices.

100,000 tasks thrown at 100,000 machines: each task goes to one random
machine, to the less loaded of two random machines, or to the least loaded of
three. The maximum load, over several seeds.
"""
import argparse, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=100_000, help="Tasks and machines.")
p.add_argument("--seeds", type=int, default=5, help="Runs per policy.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def run(choices):
    load = [0] * args.n
    for _ in range(args.n):
        picks = [random.randrange(args.n) for _ in range(choices)]
        m = min(picks, key=lambda i: load[i]); load[m] += 1
    return max(load)
rows = []
for c in (1, 2, 3):
    maxes = []
    for s in range(args.seeds): random.seed(s); maxes.append(run(c))
    rows.append({"choices": c, "max_load_min": min(maxes), "max_load_max": max(maxes), "max_loads": ", ".join(map(str, maxes))})
for r in rows: print(f"{r['choices']} choice{'s' if r['choices'] > 1 else ' '}: maximum load {r['max_loads']}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"policies": rows}}, f)