# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Whole-query restart against task retry.

A job of n tasks, each of duration T hours, all in parallel, with an
independent failure rate lambda per task-hour. Whole-query restart reruns
everything on any failure; task retry reruns only the task that failed.
"""
import argparse, json, math, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--tasks", type=int, default=1000, help="n: tasks in the job.")
p.add_argument("--task-hours", type=float, default=0.1, help="T: duration of one task, in hours.")
p.add_argument("--exposures", type=lambda s: [float(v) for v in s.split(",")], default=[0.5, 1.0, 2.0, 3.0, 4.0],
               help="Values of nT*lambda to simulate, comma-separated.")
p.add_argument("--runs", type=int, default=200, help="Simulated jobs per value.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)


def attempts_until_success(n, T, lam):
    p_task_ok = math.exp(-lam * T)            # probability that one task does not fail
    tries = 0
    while True:
        tries += 1
        if all(random.random() < p_task_ok for _ in range(n)): return tries


def task_retry_extra(n, T, lam):
    """Task-hours lost when only failed tasks are retried (each until it succeeds)."""
    lost = 0.0
    for _ in range(n):
        while (fail_at := random.expovariate(lam)) < T:   # the attempt dies after fail_at hours
            lost += fail_at
    return lost


n, T = args.tasks, args.task_hours
rows = []
for x in args.exposures:
    lam = x / (n * T)
    attempts = sum(attempts_until_success(n, T, lam) for _ in range(args.runs)) / args.runs
    extra = sum(task_retry_extra(n, T, lam) for _ in range(args.runs)) / args.runs / (n * T)
    rows.append({"exposure": x, "simulated": attempts, "e^x": math.exp(x), "task_retry_extra_pct": extra * 100})
    print(f"nT*lam = {x:<4}  whole-query attempts: {attempts:6.2f}  (e^x = {math.exp(x):6.2f})   "
          f"task retry: +{extra:.2%} task-time")

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