# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Stranded resources and DRF.

(a) How many (1 CPU, 16 GB) tasks fit on a (32 CPU, 256 GB) machine, and how
much CPU is stranded. (b) Dominant resource fairness on 9 CPU and 18 GB with
tenant A's tasks needing (1 CPU, 4 GB) and tenant B's (3 CPU, 1 GB): one task
at a time to the tenant with the smallest dominant share.
"""
import argparse, json

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

machine, task = (32, 256), (1, 16)
fit = min(machine[0] // task[0], machine[1] // task[1]); stranded = machine[0] - fit * task[0]
cap = (9, 18); need = {"A": (1, 4), "B": (3, 1)}; used = {"A": [0, 0], "B": [0, 0]}; tasks = {"A": 0, "B": 0}; steps = []
def dom(t): return max(used[t][0] / cap[0], used[t][1] / cap[1])
while True:
    t = min(need, key=dom)                                                # tenant with the smallest dominant share
    total = [used["A"][r] + used["B"][r] for r in (0, 1)]                 # the machine is shared: check what both tenants use
    if total[0] + need[t][0] > cap[0] or total[1] + need[t][1] > cap[1]: break       # simple stop: the next task does not fit
    used[t][0] += need[t][0]; used[t][1] += need[t][1]; tasks[t] += 1
    steps.append({"step": len(steps) + 1, "tenant": t, "a_dominant": dom("A"), "b_dominant": dom("B")})
cpu = used["A"][0] + used["B"][0]; mem = used["A"][1] + used["B"][1]
summary = {"fit": fit, "stranded_cpu": stranded, "tasks_a": tasks["A"], "tasks_b": tasks["B"], "cpu_used": cpu, "mem_used": mem,
           "mem_stranded": cap[1] - mem, "cpu_stranded": cap[0] - cpu}
print(f"(a) {fit} tasks fit; {stranded} of {machine[0]} CPUs stranded")
print(f"(b) DRF: {tasks}; CPU {cpu} of {cap[0]}, memory {mem} of {cap[1]} GB ({cap[1] - mem} GB stranded)")

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