# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Simulate the spot gang.

A job that needs all of its 100 spot nodes at once survives an hour only if
none of them is interrupted (simulated). Then checkpoints: the expected work
lost per interruption, and the expected completion time of a 10-hour job,
with all nodes on spot and with only the workers on spot (closed form).
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--nodes", type=int, default=100, help="Nodes the job needs at once.")
p.add_argument("--p", type=float, default=0.05, help="Probability that one node is interrupted in an hour.")
p.add_argument("--job-hours", type=float, default=10, help="Work in the job, in hours on all nodes.")
p.add_argument("--checkpoints", type=lambda s: [float(v) for v in s.split(",")], default=[0.02, 0.05, 0.1, 0.2, 0.5, 1],
               help="Checkpoint intervals in hours, comma-separated.")
p.add_argument("--restart-hours", type=float, default=0.1, help="Time to replace a node and restore from the checkpoint.")
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()
rng = np.random.default_rng(args.seed)

survive = np.mean(np.all(rng.random((100_000, args.nodes)) > args.p, axis=1))
print(f"survive one hour with {args.nodes} spot nodes: {survive:.4f}   (0.95^{args.nodes} = {(1 - args.p) ** args.nodes:.4f})")

def completion(interval, spot_nodes):
    """Expected hours to finish job_hours of work, checkpointing every `interval` hours. Interruptions arrive
    at rate lam across the gang; a segment of length tau restarts until it runs clean, which takes
    (e^(lam*tau) - 1) * (1/lam + restart) hours on average. Lost work per interruption is what ran before it."""
    lam = -np.log(1 - args.p) * spot_nodes                  # interruptions per hour across the gang
    segments = args.job_hours / interval
    per_segment = (np.exp(lam * interval) - 1) * (1 / lam + args.restart_hours)
    lost = 1 / lam - interval / (np.exp(lam * interval) - 1)  # mean work lost per interruption
    return float(segments * per_segment), float(lost)

rows = []
for interval in args.checkpoints:
    for label, spot in (("all nodes on spot", args.nodes), ("workers on spot, coordinator on demand", args.nodes - 1)):
        t, l = completion(interval, spot)
        rows.append({"checkpoint_h": interval, "setup": label, "completion_h": t, "lost_h": l})
    print(f"checkpoint every {interval:>4g} h: completion {rows[-2]['completion_h']:6.1f} h (all spot)   "
          f"{rows[-1]['completion_h']:6.1f} h (coordinator on demand); {args.job_hours:g} h of work")
interrupt_rate = -np.log(1 - args.p) * args.nodes
print(f"interruptions per hour across the gang: {interrupt_rate:.2f}; lost work per interruption: "
      + ", ".join(f"{r['lost_h'] * 60:.1f} min at {r['checkpoint_h'] * 60:g} min" for r in rows[::2]))

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"checkpoints": rows, "summary": [{"survive": float(survive), "formula": (1 - args.p) ** args.nodes,
                   "interrupts_per_h": interrupt_rate}]}}, f)