#!/usr/bin/env python3
"""What does a GPU VM that can hold a 120B open model cost per month in each Azure region, at list prices?

The model is openai/gpt-oss-120b (Apache 2.0, 117B parameters, MXFP4 weights), read from the Hugging Face API
pinned to commit b5c939de8f754692c1647ca79fbf85e8c1e70f8a. The prices are the Azure Retail Prices API
(https://prices.azure.com/api/retail/prices, api-version 2023-01-01-preview), Linux pay-as-you-go rows, USD.

    python3 azure_gpu_cost_by_region.py --out today.csv --snapshot-out azure_vm_prices_today.csv
    python3 azure_gpu_cost_by_region.py --out azure_gpu_cost_by_region.csv --snapshot azure_vm_prices_2026-10-18.csv
    python3 azure_gpu_cost_by_region.py --out planted.csv --snapshot azure_vm_prices_2026-10-18.csv \
        --plant duplicate-row,zero-price,rename-sku,drop-sku,bad-unit,sp-above-price,weights-shrunk

The first form asks the live API (prices change: the numbers on the page are from the fetch time in the
snapshot). The second re-runs the page's numbers exactly from the saved raw answer. Needs Python 3 and the
duckdb module (tested on Python 3.12, DuckDB 1.5.6). It exits 1, and writes no CSV, if a check fails.

A month is 730 hours (8,760 / 12). Only the VM's compute meter is priced: no disks, network, egress, licenses
beyond the Linux image, staff or support, and no discount except the savings-plan rates the API itself lists.
"""
import argparse
import csv
import datetime
import io
import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

import duckdb

PRICES = "https://prices.azure.com/api/retail/prices"
API_VERSION = "2023-01-01-preview"
HF_MODEL = "openai/gpt-oss-120b"
HF_COMMIT = "b5c939de8f754692c1647ca79fbf85e8c1e70f8a"
HF_WEIGHTS_BYTES = 65_248_893_184    # sum of the root-level model-*.safetensors files at that commit
HOURS_PER_MONTH = 730
HEADROOM_GB = 10                      # room beyond the weights for the KV cache and the runtime: an assumption, change it
UA = "azure_gpu_cost_by_region.py (https://www.datameaning.com/guides/azure-gpu-cost-by-region/; info@datameaning.com)"

# The VM sizes priced. GPU count and memory per GPU are from Microsoft's size pages (read 2026-10-18).
# An armSkuName that the API returns and that is not a key here is counted as unmapped and fails the run.
SKUS = {
    "Standard_NC64as_T4_v3":    (4, 16, "Tesla T4", "ncast4v3-series"),
    "Standard_NC24ads_A100_v4": (1, 80, "A100 PCIe", "nca100v4-series"),
    "Standard_NC48ads_A100_v4": (2, 80, "A100 PCIe", "nca100v4-series"),
    "Standard_NC40ads_H100_v5": (1, 94, "H100 NVL", "ncadsh100v5-series"),
    "Standard_NC80adis_H100_v5": (2, 94, "H100 NVL", "ncadsh100v5-series"),
    "Standard_ND96isr_H100_v5": (8, 80, "H100 SXM", "ndh100v5-series"),
}
PLANTS = {"duplicate-row", "zero-price", "rename-sku", "drop-sku", "bad-unit", "sp-above-price", "weights-shrunk"}
RAW_FIELDS = ["fetched_at", "armSkuName", "armRegionName", "location", "skuName", "productName", "meterName", "type",
              "isPrimaryMeterRegion", "effectiveStartDate", "unitOfMeasure", "currencyCode", "retailPrice",
              "savingsPlan1y", "savingsPlan3y", "meterId"]
OUT_FIELDS = ["sku", "gpus", "gpu_memory_gb", "fits_gpt_oss_120b", "region", "location", "usd_per_hour",
              "usd_per_month_730h", "usd_per_month_savings_1y", "usd_per_month_savings_3y", "effective_start_date"]


def get_json(url, headers=None):
    """GET with one request at a time; honor retry-after on 429/503."""
    for attempt in range(6):
        req = urllib.request.Request(url, headers={"User-Agent": UA, **(headers or {})})
        try:
            with urllib.request.urlopen(req, timeout=60) as r:
                return json.load(r)
        except urllib.error.HTTPError as e:
            if e.code in (429, 503):
                wait = int(e.headers.get("retry-after", "5"))
                print(f"        HTTP {e.code}, retry-after {wait}s")
                time.sleep(wait)
                continue
            raise
    sys.exit(f"gave up on {url}")


def fetch_prices(snapshot_out):
    """Page through the Retail Prices API for the SKUs above and save the raw rows."""
    flt = "serviceName eq 'Virtual Machines' and priceType eq 'Consumption' and (" + \
          " or ".join(f"armSkuName eq '{s}'" for s in SKUS) + ")"
    url = PRICES + "?" + urllib.parse.urlencode({"api-version": API_VERSION, "currencyCode": "USD", "$filter": flt},
                                                quote_via=urllib.parse.quote)
    fetched_at = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    rows, pages = [], 0
    while url:
        d = get_json(url)
        pages += 1
        rows += d["Items"]
        url = d.get("NextPageLink")
        if url:
            time.sleep(1)
    print(f"source  {PRICES}?api-version={API_VERSION} fetched {fetched_at}: {len(rows)} rows in {pages} page(s)")
    with open(snapshot_out, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=RAW_FIELDS, lineterminator="\n")
        w.writeheader()
        for i in rows:
            sp = {s["term"]: s["retailPrice"] for s in i.get("savingsPlan", [])}
            w.writerow({"fetched_at": fetched_at, "savingsPlan1y": sp.get("1 Year", ""),
                        "savingsPlan3y": sp.get("3 Years", ""),
                        **{k: i.get(k, "") for k in RAW_FIELDS if k not in ("fetched_at", "savingsPlan1y",
                                                                           "savingsPlan3y")}})
    return snapshot_out


def model_weights():
    """Weights size of the pinned model commit: the root-level model-*.safetensors files."""
    d = get_json(f"https://huggingface.co/api/models/{HF_MODEL}/revision/{HF_COMMIT}?blobs=true")
    files = [s for s in d["siblings"] if "/" not in s["rfilename"] and s["rfilename"].startswith("model-")
             and s["rfilename"].endswith(".safetensors")]
    total = sum(s["size"] for s in files)
    lic = d.get("cardData", {}).get("license")
    print(f"model   {HF_MODEL} @ {d['sha']} license {lic}: {len(files)} weight files, {total} bytes "
          f"({total / 1e9:.1f} GB)")
    return d["sha"], total


def load(snapshot, plant):
    with open(snapshot, "rb") as f:
        file_lines = f.read().count(b"\n") - 1
    con = duckdb.connect()
    con.sql(f"CREATE TABLE raw AS SELECT row_number() OVER () AS file_row, * FROM read_csv('{snapshot}', header=true, "
            f"all_varchar=true)")
    con.sql("CREATE TABLE px AS SELECT file_row, fetched_at, armSkuName AS sku, armRegionName AS region, location, "
            "skuName, productName, meterName, type, isPrimaryMeterRegion AS primary_meter, "
            "effectiveStartDate AS eff, unitOfMeasure AS uom, currencyCode AS cur, "
            "TRY_CAST(retailPrice AS DOUBLE) AS price, TRY_CAST(savingsPlan1y AS DOUBLE) AS sp1, "
            "TRY_CAST(savingsPlan3y AS DOUBLE) AS sp3 FROM raw")
    lin = "type = 'Consumption' AND skuName NOT LIKE '%Spot%' AND skuName NOT LIKE '%Low Priority%' " \
          "AND productName NOT LIKE '% Windows'"
    if "duplicate-row" in plant:
        con.sql(f"INSERT INTO px SELECT * FROM px WHERE {lin} ORDER BY file_row LIMIT 1")
        print("PLANTED: one Linux price row loaded twice")
    if "zero-price" in plant:
        con.sql(f"UPDATE px SET price = 0 WHERE file_row IN (SELECT file_row FROM px WHERE {lin} ORDER BY file_row "
                f"LIMIT 1)")
        print("PLANTED: one Linux price set to 0")
    if "rename-sku" in plant:
        con.sql(f"UPDATE px SET sku = 'Standard_NC24ads_A100_v4_x' WHERE file_row IN (SELECT file_row FROM px "
                f"WHERE sku = 'Standard_NC24ads_A100_v4' ORDER BY file_row LIMIT 1)")
        print("PLANTED: one NC24ads A100 v4 row renamed to a size that is not in SKUS")
    if "drop-sku" in plant:
        con.sql("DELETE FROM px WHERE sku = 'Standard_NC80adis_H100_v5'")
        print("PLANTED: every NC80adis H100 v5 row deleted")
    if "bad-unit" in plant:
        con.sql(f"UPDATE px SET uom = '1 Month' WHERE file_row IN (SELECT file_row FROM px WHERE {lin} ORDER BY "
                f"file_row LIMIT 1)")
        print("PLANTED: one Linux price quoted per '1 Month' instead of '1 Hour'")
    if "sp-above-price" in plant:
        con.sql(f"UPDATE px SET sp3 = price * 1.5 WHERE file_row IN (SELECT file_row FROM px WHERE {lin} AND sp3 IS "
                f"NOT NULL AND price > 0 ORDER BY file_row LIMIT 1)")  # not the row zero-price hit
        print("PLANTED: one 3-year savings-plan rate set 50% above the pay-as-you-go price")
    return con, file_lines, lin


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True, help="CSV to write")
    g = ap.add_mutually_exclusive_group(required=True)
    g.add_argument("--snapshot", help="re-run from a saved raw price file")
    g.add_argument("--snapshot-out", help="ask the live API and save the raw rows here")
    ap.add_argument("--plant", default="", help=f"comma list of test faults: {', '.join(sorted(PLANTS))}")
    a = ap.parse_args()
    plant = [x for x in a.plant.split(",") if x]
    if set(plant) - PLANTS:
        ap.error(f"unknown --plant value(s) {sorted(set(plant) - PLANTS)}; use {sorted(PLANTS)}")
    if plant and a.snapshot_out:
        ap.error("plant faults only run from a --snapshot, so the live API is never touched with a fault")
    print(f"duckdb {duckdb.__version__}, python {sys.version.split()[0]}")

    snapshot = a.snapshot or fetch_prices(a.snapshot_out)
    sha, weights = model_weights()
    if "weights-shrunk" in plant:
        weights //= 2
        print("PLANTED: the weights size halved, as if only half the files were summed")
    need_gb = weights / 1e9 + HEADROOM_GB
    print(f"need    {weights / 1e9:.1f} GB of weights + {HEADROOM_GB} GB headroom = {need_gb:.1f} GB of GPU memory")

    con, file_lines, lin = load(snapshot, plant)
    fetched = con.sql("SELECT MIN(fetched_at), MAX(fetched_at) FROM px").fetchone()
    print(f"prices  {snapshot}: {file_lines} data lines, fetched {fetched[0]}")
    n = con.sql("SELECT COUNT(*) FROM px").fetchone()[0]

    classes = con.sql(f"""SELECT CASE WHEN type <> 'Consumption' THEN 'other type'
                                      WHEN skuName LIKE '%Spot%' THEN 'spot'
                                      WHEN skuName LIKE '%Low Priority%' THEN 'low priority'
                                      WHEN productName LIKE '% Windows' THEN 'windows'
                                      ELSE 'linux pay-as-you-go' END AS cls, COUNT(*) FROM px GROUP BY 1 ORDER BY 1""").fetchall()
    print("rows by meter class: " + ", ".join(f"{c} {k}" for c, k in classes))
    nonprimary = con.sql(f"SELECT COUNT(*) FROM px WHERE {lin} AND lower(primary_meter) <> 'true'").fetchone()[0]
    print(f"Linux pay-as-you-go rows from a non-primary meter region: {nonprimary} (kept: they are the only Linux row "
          f"for their region and size if the checks below pass)")

    mapped = ", ".join(repr(k) for k in SKUS)
    unmapped = dict(con.sql(f"SELECT sku, COUNT(*) FROM px WHERE sku NOT IN ({mapped}) GROUP BY 1 ORDER BY 1").fetchall())
    present = {k for (k,) in con.sql(f"SELECT DISTINCT sku FROM px WHERE {lin}").fetchall()}
    dups = con.sql(f"SELECT COUNT(*) FROM (SELECT 1 FROM px WHERE {lin} GROUP BY sku, region HAVING COUNT(*) > 1)").fetchone()[0]
    bad_price = con.sql(f"SELECT COUNT(*) FROM px WHERE {lin} AND (price IS NULL OR price <= 0)").fetchone()[0]
    bad_unit = con.sql(f"SELECT COUNT(*) FROM px WHERE {lin} AND (uom <> '1 Hour' OR cur <> 'USD')").fetchone()[0]
    bad_sp = con.sql(f"SELECT COUNT(*) FROM px WHERE {lin} AND ((sp1 IS NOT NULL AND sp1 > price) OR "
                     f"(sp3 IS NOT NULL AND (sp3 > price OR (sp1 IS NOT NULL AND sp3 > sp1))))").fetchone()[0]
    sp_same = con.sql(f"SELECT COUNT(*) FROM px WHERE {lin} AND sp1 = price AND sp3 = price").fetchone()[0]
    print(f"Linux rows whose 1-year and 3-year savings-plan rates both equal the pay-as-you-go price (no discount): "
          f"{sp_same}")

    # one row per size and region: Linux pay-as-you-go only
    con.sql(f"""CREATE TABLE lin AS SELECT sku, region, location, price, sp1, sp3, eff FROM px WHERE {lin}
                AND sku IN ({mapped}) AND price > 0 AND uom = '1 Hour'""")
    con.sql("CREATE TABLE sku_t (sku VARCHAR, gpus INT, gb INT, model VARCHAR)")
    for k, (n_gpu, gb, model, _) in SKUS.items():
        con.sql(f"INSERT INTO sku_t VALUES ('{k}', {n_gpu}, {gb}, '{model}')")
    rows = con.sql(f"""SELECT l.sku, t.gpus, t.gb, (t.gpus * t.gb >= {need_gb}) AS fits, l.region, l.location, l.price,
                              round(l.price * {HOURS_PER_MONTH}, 2) AS m, round(l.sp1 * {HOURS_PER_MONTH}, 2) AS m1,
                              round(l.sp3 * {HOURS_PER_MONTH}, 2) AS m3, substr(l.eff, 1, 10) AS eff
                       FROM lin l JOIN sku_t t USING (sku) ORDER BY fits DESC, t.gpus, l.sku, m, l.region""").fetchall()

    # the tables are held back and printed only if every check below passes
    real_stdout, tables = sys.stdout, io.StringIO()
    sys.stdout = tables
    # summary per size
    summ = con.sql(f"""SELECT l.sku, t.gpus, t.gb, t.gpus * t.gb AS tot, (t.gpus * t.gb >= {need_gb}) AS fits,
                              COUNT(*) AS regions, round(MIN(price) * {HOURS_PER_MONTH}, 0) AS lo,
                              round(quantile_cont(price, 0.5) * {HOURS_PER_MONTH}, 0) AS med,
                              round(MAX(price) * {HOURS_PER_MONTH}, 0) AS hi,
                              round(MAX(price) / MIN(price), 2) AS ratio,
                              first(region ORDER BY price, region) AS cheapest, first(region ORDER BY price DESC, region) AS dearest
                       FROM lin l JOIN sku_t t USING (sku) GROUP BY l.sku, t.gpus, t.gb ORDER BY tot, l.sku""").fetchall()

    print(f"\nLinux pay-as-you-go, USD per month at {HOURS_PER_MONTH} hours, by VM size")
    head = ["size", "gpus", "gpu_gb", "fits", "regions", "min", "median", "max", "max/min", "cheapest", "dearest"]
    print("  ".join(h.rjust(w) for h, w in zip(head, [26, 4, 6, 5, 7, 7, 7, 7, 7, 18, 18])))
    for sku, g, gb, tot, fits, nr, lo, med, hi, ratio, ch, de in summ:
        cells = [sku, str(g), str(g * gb), "yes" if fits else "no", str(nr), f"{lo:.0f}", f"{med:.0f}", f"{hi:.0f}",
                 f"{ratio:.2f}", ch, de]
        print("  ".join(c.rjust(w) for c, w in zip(cells, [26, 4, 6, 5, 7, 7, 7, 7, 7, 18, 18])))

    fit_skus = [s[0] for s in summ if s[4]]
    one_gpu = [s[0] for s in summ if s[4] and s[1] == 1]   # the sizes with one GPU that fit
    regions = sorted({r[4] for r in rows if r[0] in one_gpu})
    by = {(r[0], r[4]): r for r in rows}
    print(f"\nUSD per month by region, one-GPU sizes that fit ({HOURS_PER_MONTH} h; payg = pay-as-you-go, "
          f"1y and 3y = savings-plan rate; '-' = no price or no rate listed)")
    short = {s: s.replace("Standard_", "").replace("_", " ") for s in one_gpu}
    print("".ljust(20) + "".join(short[s].center(30) for s in one_gpu))
    print("region".ljust(20) + "".join("payg".rjust(10) + "1y".rjust(10) + "3y".rjust(10) for _ in one_gpu))
    for reg in regions:
        cells = []
        for s in one_gpu:
            r = by.get((s, reg))
            cells += [f"{r[k]:,.2f}" if r and r[k] is not None else "-" for k in (7, 8, 9)]
        print(reg.ljust(20) + "".join(c.rjust(10) for c in cells))
    sp = con.sql(f"""SELECT COUNT(*) FILTER (WHERE sp1 IS NOT NULL) AS with1, COUNT(*) FILTER (WHERE sp3 IS NOT NULL) AS with3,
                            COUNT(*) AS n FROM lin l JOIN sku_t t USING (sku) WHERE t.gpus * t.gb >= {need_gb}""").fetchone()
    print(f"\nrows for sizes that fit: {sp[2]}; with a 1-year savings-plan rate {sp[0]}; with a 3-year rate {sp[1]}")
    sys.stdout = real_stdout
    print()

    failures = []
    if n != file_lines:   # duplicate-row
        failures.append(f"loaded {n} price rows but the file has {file_lines} data lines ({n - file_lines:+d})")
    if weights != HF_WEIGHTS_BYTES:   # weights-shrunk
        failures.append(f"weights are {weights} bytes, the pinned commit has {HF_WEIGHTS_BYTES}")
    if unmapped:   # rename-sku
        failures.append("rows with a size that is not in SKUS: " + ", ".join(f"'{k}' x {v}" for k, v in unmapped.items()))
    for k in SKUS:   # drop-sku
        if k not in present:
            failures.append(f"SKUS key '{k}' matches no Linux pay-as-you-go row")
    if dups:   # duplicate-row
        failures.append(f"{dups} size-and-region pairs have more than one Linux pay-as-you-go row")
    if bad_price:   # zero-price
        failures.append(f"{bad_price} Linux rows have a price that is missing or not above zero")
    if bad_unit:   # bad-unit
        failures.append(f"{bad_unit} Linux rows are not priced in USD per '1 Hour'")
    if bad_sp:   # sp-above-price
        failures.append(f"{bad_sp} Linux rows have a savings-plan rate above the pay-as-you-go price "
                        f"(or a 3-year rate above the 1-year rate)")
    if not fit_skus:   # weights-shrunk does not trip this; a too-large headroom would
        failures.append("no VM size in SKUS can hold the weights")
    unfit = [s[0] for s in summ if not s[4]]
    if "Standard_NC64as_T4_v3" not in unfit:   # weights-shrunk: 4 x 16 GB = 64 GB would then fit
        failures.append("Standard_NC64as_T4_v3 (4 x 16 GB) was expected not to fit the weights, but it does")
    if len(rows) != con.sql("SELECT COUNT(*) FROM lin").fetchone()[0]:   # no run plants this: a join that drops rows
        failures.append("the output has a different row count than the Linux price rows")

    if failures:
        print("SELF-CHECK FAILED: no CSV written")
        for f in failures:
            print("  -", f)
        sys.exit(1)
    print(tables.getvalue())
    print(f"SELF-CHECK PASSED: {n} price rows equal the {file_lines} data lines of the file, the weights equal the pinned "
          f"{HF_WEIGHTS_BYTES} bytes, every size is mapped, every SKUS key matches a Linux row, one Linux row per size "
          f"and region, every price is above zero in USD per hour, no savings-plan rate is above its price, "
          f"{len(unfit)} size(s) do not fit")
    with open(a.out, "w", newline="") as f:
        w = csv.writer(f, lineterminator="\n")
        w.writerow(OUT_FIELDS)
        for sku, g, gb, fits, reg, loc, ph, m, m1, m3, eff in rows:
            w.writerow([sku, g, gb, "yes" if fits else "no", reg, loc, f"{ph:.6f}".rstrip("0").rstrip("."), f"{m:.2f}",
                        "" if m1 is None else f"{m1:.2f}", "" if m3 is None else f"{m3:.2f}", eff])
    print(f"wrote {a.out}")


if __name__ == "__main__":
    main()
