# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Fan-out as a batch size.

Measures the launch overhead F of a Python subprocess (and, if Docker is
available, of `docker run --rm alpine true`), computes the useful share
W / (W + F) for tasks of 1 s to 600 s, and the task duration at which overhead
falls below 10%. Then runs 1,000 tiny tasks with one subprocess each, and as
10 batches of 100.
"""
import argparse, json, os, shutil, statistics, subprocess, sys, tempfile, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--tasks", type=int, default=1000, help="Tiny tasks.")
p.add_argument("--batches", type=int, default=10, help="Batches in the batched run.")
p.add_argument("--docker", default="auto", choices=["auto", "yes", "no"], help="Also time a container start.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

TASK = "import sys; sum(range(1000))"
def launch_times(cmd, n, env=None):
    ts = []
    for _ in range(n):
        t = time.perf_counter(); subprocess.run(cmd, check=True, capture_output=True, env=env); ts.append(time.perf_counter() - t)
    return statistics.median(ts)
launchers = [{"launcher": "Python subprocess", "F_s": launch_times([sys.executable, "-c", TASK], 30)}]
if args.docker != "no" and shutil.which("docker"):
    cfg = tempfile.mkdtemp(); open(os.path.join(cfg, "config.json"), "w").write("{}")      # no credential helper
    env = {**os.environ, "DOCKER_CONFIG": cfg}
    try:
        subprocess.run(["docker", "pull", "-q", "alpine:3.20"], check=True, capture_output=True, env=env, timeout=300)
        launchers.append({"launcher": "docker run --rm alpine true", "F_s": launch_times(["docker", "run", "--rm", "alpine:3.20", "true"], 10, env)})
    except Exception as e:
        print("docker unavailable:", str(e)[:80])
shares = [{"launcher": l["launcher"], "W_s": W, "useful_pct": W / (W + l["F_s"]) * 100} for l in launchers for W in (1, 10, 60, 600)]
for l in launchers: l["W_for_10pct_s"] = 9 * l["F_s"]                 # F / (W + F) < 0.1  <=>  W > 9F

t = time.perf_counter()
for _ in range(args.tasks): subprocess.run([sys.executable, "-c", TASK], check=True)
one_each = time.perf_counter() - t
per = args.tasks // args.batches
t = time.perf_counter()
for _ in range(args.batches): subprocess.run([sys.executable, "-c", f"for _ in range({per}): sum(range(1000))"], check=True)
batched = time.perf_counter() - t
runs = [{"run": f"{args.tasks} tasks, one subprocess each", "seconds": one_each},
        {"run": f"{args.batches} batches of {per}", "seconds": batched}]
summary = {"speedup": one_each / batched}
for l in launchers: print(f"{l['launcher']:<30} F = {l['F_s'] * 1000:7.1f} ms; overhead below 10% for tasks longer than {l['W_for_10pct_s']:.2f} s")
for s in shares: print(f"  {s['launcher']:<30} W = {s['W_s']:>4} s: useful share {s['useful_pct']:6.2f}%")
for r in runs: print(f"{r['run']:<36} {r['seconds']:6.2f} s")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"launchers": launchers, "shares": shares, "runs": runs}}, f)