([\s\S]*?)<\/section>/);
+ assert.ok(m, "chapter exists");
+ return { html, body: m[1] };
+}
+
+test("the Heatwave window module registers a mount function and keeps the house style", async () => {
+ const src = await fs.readFile(path.join(DOCS, "heatwave-lab.js"), "utf8");
+ assert.ok(!/[—–]/.test(src), "no em or en dash");
+ assert.ok(src.length < 40 * 1024, "module stays small: " + src.length);
+ globalThis.window = globalThis;
+ await import(path.join(DOCS, "heatwave-lab.js"));
+ assert.equal(typeof globalThis.GPUStackHeatwave.mount, "function");
+});
+
+test("the Heatwave data file has every quantile state and labels its data", async () => {
+ const d = await readJson("e003-heatwave-lab-v1.json");
+ assert.equal(d.schema, "gpu-stack.e003-heatwave-lab.v1");
+ for (const key of ["packing", "outage", "cost", "audit", "forecast", "thermal", "ablation", "ablation_sim", "finetune"]) {
+ assert.ok(d[key] && d[key].status, key + " block");
+ }
+ assert.equal(d.packing.status, "ready");
+ assert.equal(d.packing.data_kind, "simulated");
+ assert.equal(d.packing.quantile_levels.length, 21);
+ for (const q of d.packing.quantile_levels) {
+ const run = d.packing.trace.runs["timestral-5M/q" + String(q)];
+ assert.ok(run, "trace for level " + q);
+ assert.equal(run.site_power_kw.length, d.packing.trace.runs.reactive.site_power_kw.length);
+ assert.ok(run.matches_committed_part, "replay matches the committed eval part for " + q);
+ assert.ok(d.packing.policies["timestral-5M/q" + String(q)], "summary for level " + q);
+ }
+ assert.ok(d.packing.policies.reactive);
+ assert.equal(d.outage.data_kind, "simulated");
+ assert.equal(d.outage.versions.v2.status, "ready", "outage v2 is in");
+ const raw = await fs.readFile(path.join(DOCS, "data", "e003-heatwave-lab-v1.json"), "utf8");
+ assert.ok(raw.length < 1.5 * 1024 * 1024, "data file under 1.5 MB");
+});
+
+test("the story data file is the lab file without the replay traces", async () => {
+ const lab = await readJson("e003-heatwave-lab-v1.json");
+ const story = await readJson("e003-heatwave-story-v1.json");
+ assert.equal(story.schema, "gpu-stack.e003-heatwave-story.v1");
+ for (const key of Object.keys(lab)) {
+ if (key === "schema" || key === "packing") continue;
+ assert.deepEqual(story[key], lab[key], key + " matches");
+ }
+ assert.equal(story.packing.trace.runs, undefined);
+ assert.deepEqual(story.packing.policies, lab.packing.policies);
+ const raw = await fs.readFile(path.join(DOCS, "data", "e003-heatwave-story-v1.json"), "utf8");
+ assert.ok(raw.length < 150 * 1024, "story data stays small: " + raw.length);
+});
+
+test("pending blocks say so and ready blocks carry a source and a data label", async () => {
+ const d = await readJson("e003-heatwave-lab-v1.json");
+ for (const key of ["forecast", "thermal", "ablation", "ablation_sim", "finetune"]) {
+ assert.ok(["ready", "pending"].includes(d[key].status), key);
+ }
+ assert.equal(d.ablation.status, "ready");
+ for (const [name, ds] of Object.entries(d.ablation.datasets)) {
+ assert.ok(["public", "measured", "simulated"].includes(ds.data_kind) || /simulated/.test(ds.data_kind), name + " data kind");
+ assert.ok(ds.methods["timestral-209M"].none.wql > 0, name);
+ for (const cfg of Object.values(ds.methods["timestral-209M"])) assert.ok(!cfg.groups.includes("L"), "no config that reads realised values");
+ }
+ assert.ok(d.ablation.source.sha256 && d.ablation.source.file.endsWith("ablation_real.json"));
+});
+
+test("the Story chapter cites its numbers and every number matches the data file", async () => {
+ const { body } = await chapter();
+ const story = await readJson("e003-heatwave-story-v1.json");
+ const byId = new Map(story.story_numbers.map((n) => [n.id, n]));
+ assert.ok(byId.size === story.story_numbers.length, "unique ids");
+ assert.ok(!/[—–]/.test(plain(body)), "no em dash in copy");
+ assert.ok(!plain(body).includes("!"), "no exclamation point");
+ const ALLOW = new Set(["1", "5", "10", "21", "99", "90", "2026"]);
+ const paragraphs = body.match(/[^]*?<\/p>/g) || [];
+ let cited = 0;
+ for (const p of paragraphs) {
+ const text = plain(p);
+ const ids = [...p.matchAll(//g)].flatMap((m) => m[1].trim().split(/\s+/));
+ const covered = ids.map((id) => {
+ assert.ok(byId.has(id), "unknown number id " + id);
+ const n = byId.get(id);
+ assert.ok(text.includes(n.text), `paragraph must contain "${n.text}" (${id}): ${text.slice(0, 70)}`);
+ cited += 1;
+ return n.text;
+ }).join(" | ");
+ for (const tok of text.replace(/\b[HB]\d{3}\b/g, "").match(/\d[\d,]*(?:\.\d+)?/g) || []) {
+ const t = tok.replace(/,$/, "");
+ assert.ok(ALLOW.has(t) || covered.includes(t), `number ${t} has no source id in: ${text.slice(0, 70)}`);
+ }
+ }
+ assert.ok(cited >= 40, "the chapter cites its numbers: " + cited);
+});
+
+test("claims the prose makes about intervals and counts hold in the data", async () => {
+ const d = await readJson("e003-heatwave-story-v1.json");
+ assert.equal(d.forecast.n_series_both_sizes_lower, d.forecast.n_series, "Timestral lower than every textbook method on every series");
+ const u = d.outage.versions.v2.policies.ts28m.useful_gpu_h_vs_reactive;
+ assert.ok(u.lo < 0 && u.hi > 0, "useful GPU-hours difference is inside its interval");
+ const u3 = d.outage.versions.v3.policies.ts28m.useful_gpu_h_vs_reactive;
+ assert.ok(u3.lo < 0 && u3.hi > 0, "v3 useful GPU-hours difference is inside its interval");
+ const l3 = d.outage.versions.v3.policies.ts28m.lost_gpu_h_vs_reactive;
+ assert.ok(l3.hi < 0, "v3 lost GPU-hours are lower than reacting");
+ const lv = d.packing.learned_vs_reactive;
+ assert.ok(lv["10"].vs_reactive.lo < 0 && lv["10"].vs_reactive.hi > 0, "the $10 difference spans zero");
+ assert.ok(lv["40"].vs_reactive.lo > 0 && lv["160"].vs_reactive.lo > 0 && lv["2.5"].vs_reactive.hi < 0);
+ const cfg = (ds, c) => d.ablation.datasets[ds].methods["timestral-209M"][c].vs_none;
+ assert.equal(cfg("caiso_rt_lmp_5min", "DA"), "lower");
+ assert.equal(cfg("caiso_rt_lmp_5min", "WX"), "overlap");
+ assert.equal(cfg("laptop_power_1s_run2", "S"), "lower");
+ assert.equal(cfg("laptop_power_10hz_run2", "S"), "higher");
+ assert.equal(cfg("gb_carbon_30min", "F"), "lower");
+ assert.ok(d.outage.versions.v2.policies.ts28m.recall_in_time > d.outage.versions.v1.policies.ts28m.recall_in_time);
+});
+
+test("a pending marker in the chapter exists only while its result is pending", async () => {
+ const { body } = await chapter();
+ const d = await readJson("e003-heatwave-story-v1.json");
+ for (const key of ["ablation_sim", "finetune"]) {
+ const marker = new RegExp("HEATWAVE PENDING: [a-z_ ]*\\b" + key + "\\b").test(body);
+ const ready = d[key].status === "ready";
+ assert.ok(!(marker && ready), key + " landed: replace its pending paragraph with its numbers from data");
+ if (!marker) assert.ok(ready, key + ": pending paragraph removed too early");
+ }
+});
+
+test("the chapter has its figures and the window exists", async () => {
+ const { html, body } = await chapter();
+ assert.ok(body.includes('data-figure="packing-tradeoff"') && body.includes('data-figure="known-future"'));
+ assert.ok(html.includes('src="figures/packing-tradeoff.js"') && html.includes('src="figures/known-future.js"'));
+ assert.ok(html.includes('data-window="heatwave"'), "window exists");
+});
+
+test("glossary entries for Heatwave are edited and their numbers are sourced", async () => {
+ const ex = await readJson("explainers.json");
+ const story = await readJson("e003-heatwave-story-v1.json");
+ const all = story.story_numbers.map((n) => n.text).join(" | ");
+ for (const id of ["quantile-forecast", "timestral", "thermal-throttling", "power-cap", "carbon-intensity"]) {
+ const e = ex.entries.find((x) => x.id === id);
+ assert.ok(e, id);
+ assert.equal(e.review, "edited", id);
+ for (const tok of e.body.match(/\d[\d,]*(?:\.\d+)?/g) || []) {
+ assert.ok(["1", "5", "10", "21", "99", "2026"].includes(tok) || all.includes(tok), id + ": unsourced number " + tok);
+ }
+ }
+ const visuals = new Set(ex.entries.filter((e) => ["quantile-forecast", "timestral", "thermal-throttling", "power-cap", "carbon-intensity"].includes(e.id)).map((e) => e.visual));
+ assert.equal(visuals.size, 5, "each Heatwave term has its own visual");
+});
+
+test("the Heatwave icon is its own and is used on the desktop, the title bar and the taskbar", async () => {
+ const css = await fs.readFile(path.join(DOCS, "styles", "70-pixel-icons.css"), "utf8");
+ const grab = (cls) => (css.match(new RegExp("\\." + cls + " \\{\\s*background-image: url\\(([^)]*)\\)")) || [])[1];
+ assert.ok(grab("px-heat") && grab("px-graph"));
+ assert.notEqual(grab("px-heat"), grab("px-graph"));
+ const html = await fs.readFile(path.join(DOCS, "index.html"), "utf8");
+ assert.equal((html.match(/px-heat/g) || []).length, 3);
+ assert.ok(!/px-graph"><\/span>\s*Heatwave/.test(html));
+});
diff --git a/tests/test_heatwave_lab_data.py b/tests/test_heatwave_lab_data.py
new file mode 100644
index 0000000..457b49c
--- /dev/null
+++ b/tests/test_heatwave_lab_data.py
@@ -0,0 +1,61 @@
+"""The Heatwave site data is rebuilt from committed result files.
+
+The e003 blocks (packing, outage, cost, audit) come from files inside this
+repo, so they can be rebuilt here and compared with the committed site data.
+The blocks that read the Heatwave repo are not rebuilt here; the test only
+checks that they report a pending state when that repo is not given.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+SCRIPT = ROOT / "scripts" / "build_heatwave_lab_data.py"
+LAB = ROOT / "docs" / "data" / "e003-heatwave-lab-v1.json"
+STORY = ROOT / "docs" / "data" / "e003-heatwave-story-v1.json"
+
+
+def _builder():
+ spec = importlib.util.spec_from_file_location("build_heatwave_lab_data", SCRIPT)
+ mod = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(mod)
+ return mod
+
+
+def _roundtrip(x):
+ return json.loads(json.JSONEncoder(allow_nan=False).encode(x))
+
+
+def test_e003_blocks_match_committed_data():
+ b = _builder()
+ lab = json.loads(LAB.read_text())
+ for name, fn in [("packing", b.packing_block), ("outage", b.outage_block), ("cost", b.cost_block), ("audit", b.audit_block)]:
+ assert _roundtrip(fn()) == lab[name], f"{name} block differs from {LAB.name}; rebuild with the script"
+
+
+def test_story_file_is_lab_file_without_replay_traces():
+ lab = json.loads(LAB.read_text())
+ story = json.loads(STORY.read_text())
+ assert "runs" in lab["packing"]["trace"]
+ assert "runs" not in story["packing"]["trace"]
+ for key in ("outage", "cost", "audit", "ablation", "story_numbers"):
+ assert story[key] == lab[key]
+
+
+def test_heatwave_repo_blocks_are_pending_without_the_repo():
+ b = _builder()
+ blocks = b.heatwave_blocks(None)
+ assert blocks, "expected the Heatwave repo blocks to be present"
+ for name, block in blocks.items():
+ assert block.get("status") == "pending", f"{name} should be pending when no Heatwave repo is given"
+
+
+def test_story_numbers_have_unique_ids_and_sources():
+ lab = json.loads(LAB.read_text())
+ ids = [n["id"] for n in lab["story_numbers"]]
+ assert len(ids) == len(set(ids))
+ for n in lab["story_numbers"]:
+ assert n["text"] and n["source"], n["id"]
diff --git a/tests/test_import_registry.py b/tests/test_import_registry.py
index 5e46e93..d948578 100644
--- a/tests/test_import_registry.py
+++ b/tests/test_import_registry.py
@@ -1,12 +1,12 @@
"""Pins the registry's published size and proves it can rebuild from empty.
-PUBLISHED_SNAPSHOT records the exact counts the project advertises — 1267
-variables, 701 equations, 583 root inputs, and the rest. Any scope change
+PUBLISHED_SNAPSHOT records the exact counts the project advertises — 1319
+variables, 735 equations, 604 root inputs, and the rest. Any scope change
that adds or removes a variable moves these numbers, and that is the
point: the change must be seen and the snapshot updated deliberately. If a
change legitimately moves the numbers, update the expectations here — never
delete the assertions. The second test resets the registry to zero and
-bootstraps it back, proving the full graph is reconstructible from code
+bootstraps it back, proving the full graph is rebuildable from code
alone.
"""
@@ -16,19 +16,19 @@
PUBLISHED_SNAPSHOT = {
- "systems": 16,
- "variables": 1267,
+ "systems": 17,
+ "variables": 1319,
"constants": 18,
- "equations": 701,
- "root_inputs": 583,
- "leaves": 241,
- "topological_order_length": 1267,
- "with_sp_units": 1249,
- "with_references": 1249,
- "equations_with_references": 701,
- "equations_with_unit_check": 642,
- "root_kind": 583,
- "derived_kind": 666,
+ "equations": 735,
+ "root_inputs": 604,
+ "leaves": 252,
+ "topological_order_length": 1319,
+ "with_sp_units": 1301,
+ "with_references": 1301,
+ "equations_with_references": 735,
+ "equations_with_unit_check": 672,
+ "root_kind": 604,
+ "derived_kind": 697,
"measured_kind": 0,
"definitional_kind": 18,
}
diff --git a/tests/test_power_forecast_scope.py b/tests/test_power_forecast_scope.py
new file mode 100644
index 0000000..cfd5ad1
--- /dev/null
+++ b/tests/test_power_forecast_scope.py
@@ -0,0 +1,641 @@
+"""Tests for the forecast-aware power, thermal, and cost scope.
+
+The scope puts the equations behind forecast-driven operation into the map:
+planning site power at a forecast quantile, stranded power, a lumped RC
+thermal model, the value of an early checkpoint, and time-shifted cost and
+carbon per token. These tests cover each group three ways. Metadata: every
+variable has a unit and a reference, every equation has a reference, and the
+unit-checked set is pinned. Numbers: each equation is evaluated through the
+resolver at inputs chosen so the answer can be derived by hand, and the
+expected value is computed in the test from those same inputs. Behavior: the
+limiting cases (never reaches the throttle threshold, no recall, no shifting)
+come out where the algebra says they must. A unit check that has been
+deliberately broken must raise, so the checks are known to bite.
+
+Inputs in these tests are arbitrary numbers picked for arithmetic, not
+measurements of any system. The one sampling test uses simulated (synthetic)
+draws from a normal distribution to check the exceedance rule 1 - q.
+"""
+
+import math
+
+import numpy as np
+import pytest
+import sympy as sp
+
+import gpu_stack
+from gpu_stack import Registry, resolve
+from gpu_stack.core.units import (
+ HOUR, JOULE, KELVIN, KILOGRAM, KWH, SECOND, WATT, UnitError,
+ check_equation_units, infer_expr_units_with_coefficient,
+)
+from gpu_stack.scopes import (
+ SCOPE_DESCRIPTIONS, SCOPE_MODULES, power_forecast as pf,
+)
+from gpu_stack.scopes.economics_capex_refs import USD
+
+
+# Equations that carry a unit check. Pinned on purpose: the one piecewise
+# relation is not unit-checked (the checker cannot infer a Piecewise), and
+# its closed-form branch is checked on its own; inequalities state bounds.
+UNCHECKED_POWER_FORECAST_RELATIONS = {
+ "forecast.eq.rc.time_to_throttle",
+ "forecast.ineq.plan.quantile_level_below_one",
+ "forecast.ineq.rc.threshold_above_supply",
+ "forecast.ineq.ckpt.lead_time_covers_checkpoint",
+}
+
+
+def value(variable_name, **assignments):
+ """Resolve a variable with assignments keyed by registered variable name."""
+ result = resolve(variable_name, assignments=assignments)
+ assert result.value is not None, result
+ return result.value
+
+
+def constraint(variable_name, equation, **assignments):
+ result = resolve(variable_name, assignments=assignments)
+ return next(c for c in result.constraints if c.equation == equation)
+
+
+# ---------------------------------------------------------------------------
+# Registration and metadata
+# ---------------------------------------------------------------------------
+
+def test_scope_is_registered_after_economics():
+ assert "power_forecast" in SCOPE_MODULES
+ assert SCOPE_MODULES.index("power_forecast") > SCOPE_MODULES.index("economics")
+ assert SCOPE_DESCRIPTIONS["power_forecast"]
+ assert "power_forecast" in Registry.systems
+
+
+def test_every_variable_and_equation_has_units_and_references():
+ assert pf.POWER_FORECAST_VARIABLES
+ assert [v.name for v in pf.POWER_FORECAST_VARIABLES if v.sp_units is None] == []
+ assert [v.name for v in pf.POWER_FORECAST_VARIABLES if not v.references] == []
+ assert [e.name for e in pf.POWER_FORECAST_EQUATIONS if not e.references] == []
+ assert all(v.scope == "power_forecast" for v in pf.POWER_FORECAST_VARIABLES)
+
+
+def test_unit_checked_equations_are_explicitly_curated():
+ unchecked = {
+ e.name
+ for e in pf.POWER_FORECAST_EQUATIONS
+ if not getattr(e, "_check_units_flag", False)
+ }
+ assert unchecked == UNCHECKED_POWER_FORECAST_RELATIONS
+
+
+def test_names_follow_the_forecast_prefix_and_are_unique():
+ names = [v.name for v in pf.POWER_FORECAST_VARIABLES]
+ assert len(names) == len(set(names))
+ assert all(n.startswith("forecast.") for n in names)
+ eq_names = [e.name for e in pf.POWER_FORECAST_EQUATIONS]
+ assert len(eq_names) == len(set(eq_names))
+ assert all(n.startswith("forecast.eq.") or n.startswith("forecast.ineq.") for n in eq_names)
+
+
+def test_sources_are_cited_with_links():
+ sources = {
+ r
+ for v in pf.POWER_FORECAST_VARIABLES
+ for r in v.references
+ if r.kind in {"paper", "textbook"}
+ }
+ # Gneiting & Raftery, Koenker, Koenker & Bassett, Fan et al., Incropera,
+ # Young, Daly, Wiesner et al., Radovanovic et al.
+ assert len(sources) == 9
+ for ref in sources:
+ if ref.kind == "paper":
+ assert ref.url and ref.year
+
+
+def test_forecaster_inputs_point_at_the_experiment_and_carry_no_default():
+ for name in (
+ "forecast.plan.site_power_planned",
+ "forecast.plan.site_power_expected",
+ "forecast.plan.overshoot_power_mean",
+ "forecast.rc.heat_power",
+ "forecast.ckpt.lead_time",
+ "forecast.ckpt.recall",
+ "forecast.shift.price_shifted",
+ "forecast.shift.carbon_intensity_shifted",
+ ):
+ v = Registry.variables[name]
+ assert v.is_root_input, name
+ assert any("e003-heatwave" in r.citation for r in v.references), name
+
+
+def test_the_scope_adds_no_cycles_and_no_duplicate_of_existing_variables():
+ assert gpu_stack.find_cycles() == []
+ reused = {
+ "thermal.dc.pue",
+ "thermal.coolant.t_inlet",
+ "cluster.rel.checkpoint_time",
+ "cluster.rel.optimal_checkpoint_interval",
+ "econ.power.price_kwh",
+ "econ.carbon.intensity_kg_per_kwh",
+ "econ.carbon.price_per_tonne",
+ "econ.job.dc_power",
+ "econ.cost.per_token",
+ "training.tokens_per_sec",
+ }
+ used = {
+ v.name
+ for e in pf.POWER_FORECAST_EQUATIONS
+ for v in e.variables_in_relation()
+ }
+ assert reused <= used
+
+
+def test_representative_units():
+ assert pf.site_power_cap.sp_units == WATT
+ assert pf.overcap_energy.sp_units == KWH
+ assert pf.overcap_penalty_cost.sp_units == USD
+ assert pf.rc_resistance.sp_units == KELVIN / WATT
+ assert pf.rc_capacitance.sp_units == JOULE / KELVIN
+ assert pf.rc_time_constant.sp_units == SECOND
+ assert pf.rc_steady_temperature.sp_units == KELVIN
+ assert pf.lost_work_unwarned.sp_units == HOUR
+ assert pf.facility_energy_per_token.sp_units == JOULE
+ assert pf.carbon_per_token.sp_units == KILOGRAM
+ assert pf.price_paid.sp_units == USD / KWH
+
+
+def _units_agree(equation, rhs):
+ symbol_units = {
+ v.symbol: v.sp_units for v in Registry.variables.values() if v.sp_units is not None
+ }
+ lhs_u = infer_expr_units_with_coefficient(equation.lhs, symbol_units, equation.name)
+ rhs_u = infer_expr_units_with_coefficient(rhs, symbol_units, equation.name)
+ check_equation_units(lhs_u, rhs_u, equation.name)
+
+
+def test_unit_checks_pass_as_written_and_raise_when_broken():
+ for equation in pf.POWER_FORECAST_EQUATIONS:
+ if getattr(equation, "_check_units_flag", False):
+ _units_agree(equation, equation.rhs)
+
+ # Dropping or doubling a conversion factor must be caught.
+ broken = {
+ pf.eq_overcap_energy: pf.eq_overcap_energy.rhs * sp.Integer(3_600_000),
+ pf.eq_lost_work_unwarned: pf.eq_lost_work_unwarned.rhs * sp.Integer(3600),
+ pf.eq_energy_cost_per_token: pf.eq_energy_cost_per_token.rhs * sp.Integer(3_600_000),
+ }
+ # A time constant swapped for a resistance changes the exponent's dimension.
+ broken[pf.eq_rc_temperature_at_time] = pf.eq_rc_temperature_at_time.rhs.subs(
+ pf.rc_time_constant.symbol, pf.rc_resistance.symbol
+ )
+ for equation, bad_rhs in broken.items():
+ with pytest.raises(UnitError):
+ _units_agree(equation, bad_rhs)
+
+
+# ---------------------------------------------------------------------------
+# Group 1: planning at a quantile, packing, exceedance, over-cap energy
+# ---------------------------------------------------------------------------
+
+def test_headroom_and_packable_gpus():
+ cap, planned = 10.0e6, 8.0e6
+ gpu_it_power, pue = 700.0, 1.25
+ headroom = cap - planned
+ assert value(
+ "forecast.plan.headroom",
+ **{"forecast.plan.site_power_cap": cap, "forecast.plan.site_power_planned": planned},
+ ) == pytest.approx(headroom)
+ per_gpu_site = gpu_it_power * pue
+ assert value(
+ "forecast.plan.gpu_site_power_planned",
+ **{"forecast.plan.gpu_power_planned": gpu_it_power, "thermal.dc.pue": pue},
+ ) == pytest.approx(per_gpu_site)
+ assert value(
+ "forecast.plan.n_gpus_packable",
+ **{
+ "forecast.plan.site_power_cap": cap,
+ "forecast.plan.site_power_planned": planned,
+ "forecast.plan.gpu_power_planned": gpu_it_power,
+ "thermal.dc.pue": pue,
+ },
+ ) == pytest.approx(headroom / per_gpu_site)
+
+
+def test_plan_over_the_cap_gives_negative_headroom_and_gpus_to_remove():
+ inputs = {
+ "forecast.plan.site_power_cap": 5.0e6,
+ "forecast.plan.site_power_planned": 6.0e6,
+ "forecast.plan.gpu_power_planned": 800.0,
+ "thermal.dc.pue": 1.0,
+ }
+ assert value("forecast.plan.headroom", **inputs) == pytest.approx(-1.0e6)
+ assert value("forecast.plan.n_gpus_packable", **inputs) == pytest.approx(-1.0e6 / 800.0)
+
+
+@pytest.mark.parametrize("q", [0.5, 0.9, 0.95, 0.99])
+def test_exceedance_probability_is_one_minus_q(q):
+ assert value(
+ "forecast.plan.exceedance_probability", **{"forecast.plan.quantile_level": q}
+ ) == pytest.approx(1.0 - q)
+
+
+def test_calibrated_quantile_exceedance_matches_one_minus_q_on_synthetic_draws():
+ # Simulated (synthetic) site power draws, not a measurement.
+ rng = np.random.default_rng(0)
+ draws = rng.normal(loc=8.0e6, scale=0.5e6, size=200_000)
+ for q in (0.5, 0.9, 0.99):
+ planned = float(np.quantile(draws, q))
+ empirical = float(np.mean(draws > planned))
+ predicted = value(
+ "forecast.plan.exceedance_probability", **{"forecast.plan.quantile_level": q}
+ )
+ assert empirical == pytest.approx(predicted, abs=0.004)
+
+
+def test_overcap_energy_and_penalty():
+ q, overshoot, horizon, price = 0.9, 5.0e5, 3600.0, 0.2
+ p_exceed = 1.0 - q
+ energy_kwh = p_exceed * overshoot * horizon / 3.6e6
+ inputs = {
+ "forecast.plan.quantile_level": q,
+ "forecast.plan.overshoot_power_mean": overshoot,
+ "forecast.plan.horizon": horizon,
+ "forecast.plan.overcap_price": price,
+ }
+ assert value("forecast.plan.overcap_energy", **inputs) == pytest.approx(energy_kwh)
+ assert value("forecast.plan.overcap_penalty_cost", **inputs) == pytest.approx(
+ energy_kwh * price
+ )
+
+
+def test_raising_the_quantile_lowers_expected_overcap_penalty():
+ base = {
+ "forecast.plan.overshoot_power_mean": 5.0e5,
+ "forecast.plan.horizon": 7200.0,
+ "forecast.plan.overcap_price": 0.1,
+ }
+ costs = [
+ value("forecast.plan.overcap_penalty_cost", **base, **{"forecast.plan.quantile_level": q})
+ for q in (0.5, 0.9, 0.99)
+ ]
+ assert costs[0] > costs[1] > costs[2] > 0
+
+
+def test_quantile_level_must_lie_strictly_between_zero_and_one():
+ for q, ok in ((0.5, True), (0.0, False), (1.0, False)):
+ above = constraint(
+ "forecast.plan.exceedance_probability",
+ "domain.forecast.plan.quantile_level.positive",
+ **{"forecast.plan.quantile_level": q},
+ )
+ below = constraint(
+ "forecast.plan.exceedance_probability",
+ "forecast.ineq.plan.quantile_level_below_one",
+ **{"forecast.plan.quantile_level": q},
+ )
+ assert (above.satisfied and below.satisfied) is ok
+
+
+# ---------------------------------------------------------------------------
+# Group 2: stranded power
+# ---------------------------------------------------------------------------
+
+def test_stranded_power_utilization_and_reserve():
+ cap, expected, planned = 10.0e6, 6.0e6, 8.0e6
+ inputs = {
+ "forecast.plan.site_power_cap": cap,
+ "forecast.plan.site_power_expected": expected,
+ "forecast.plan.site_power_planned": planned,
+ }
+ assert value("forecast.plan.stranded_power", **inputs) == pytest.approx(cap - expected)
+ assert value("forecast.plan.cap_utilization", **inputs) == pytest.approx(expected / cap)
+ assert value("forecast.plan.uncertainty_reserve", **inputs) == pytest.approx(planned - expected)
+
+
+def test_stranded_power_plus_utilization_account_for_the_whole_cap():
+ inputs = {
+ "forecast.plan.site_power_cap": 12.0e6,
+ "forecast.plan.site_power_expected": 9.0e6,
+ }
+ stranded = value("forecast.plan.stranded_power", **inputs)
+ utilization = value("forecast.plan.cap_utilization", **inputs)
+ assert stranded / inputs["forecast.plan.site_power_cap"] == pytest.approx(1.0 - utilization)
+
+
+# ---------------------------------------------------------------------------
+# Group 3: lumped RC thermal model
+# ---------------------------------------------------------------------------
+
+RC = {
+ "forecast.rc.resistance": 0.05, # K/W
+ "forecast.rc.capacitance": 2000.0, # J/K
+ "forecast.rc.heat_power": 700.0, # W
+ "thermal.coolant.t_inlet": 300.0, # K
+ "forecast.rc.initial_temperature": 310.0,
+ "forecast.rc.throttle_temperature": 330.0,
+ "forecast.rc.elapsed_time": 60.0,
+}
+
+
+def rc(**overrides):
+ return {**RC, **overrides}
+
+
+def test_rc_time_constant_and_steady_state():
+ tau = RC["forecast.rc.resistance"] * RC["forecast.rc.capacitance"]
+ steady = RC["thermal.coolant.t_inlet"] + RC["forecast.rc.heat_power"] * RC["forecast.rc.resistance"]
+ assert value("forecast.rc.time_constant", **RC) == pytest.approx(tau)
+ assert value("forecast.rc.steady_temperature", **RC) == pytest.approx(steady)
+
+
+def test_rc_temperature_after_time_matches_the_closed_form():
+ tau = RC["forecast.rc.resistance"] * RC["forecast.rc.capacitance"]
+ steady = RC["thermal.coolant.t_inlet"] + RC["forecast.rc.heat_power"] * RC["forecast.rc.resistance"]
+ t0 = RC["forecast.rc.initial_temperature"]
+ for t in (0.0, 30.0, tau, 5 * tau, 50 * tau):
+ expected = steady + (t0 - steady) * math.exp(-t / tau)
+ assert value(
+ "forecast.rc.temperature_at_time", **rc(**{"forecast.rc.elapsed_time": t})
+ ) == pytest.approx(expected)
+
+
+def test_rc_temperature_limits_and_one_time_constant():
+ tau = RC["forecast.rc.resistance"] * RC["forecast.rc.capacitance"]
+ steady = RC["thermal.coolant.t_inlet"] + RC["forecast.rc.heat_power"] * RC["forecast.rc.resistance"]
+ t0 = RC["forecast.rc.initial_temperature"]
+
+ def at(t):
+ return value("forecast.rc.temperature_at_time", **rc(**{"forecast.rc.elapsed_time": t}))
+
+ assert at(0.0) == pytest.approx(t0)
+ assert at(100 * tau) == pytest.approx(steady)
+ # After one time constant the gap to steady state has shrunk to 1/e.
+ assert at(tau) - steady == pytest.approx((t0 - steady) / math.e)
+
+
+def test_rc_time_to_throttle_is_the_time_the_closed_form_reaches_the_threshold():
+ tau = RC["forecast.rc.resistance"] * RC["forecast.rc.capacitance"]
+ steady = RC["thermal.coolant.t_inlet"] + RC["forecast.rc.heat_power"] * RC["forecast.rc.resistance"]
+ t0, thr = RC["forecast.rc.initial_temperature"], RC["forecast.rc.throttle_temperature"]
+ expected = tau * math.log((steady - t0) / (steady - thr))
+ t_hit = value("forecast.rc.time_to_throttle", **RC)
+ assert t_hit == pytest.approx(expected)
+ # Round trip: evaluated at that time, the temperature equals the threshold.
+ assert value(
+ "forecast.rc.temperature_at_time", **rc(**{"forecast.rc.elapsed_time": t_hit})
+ ) == pytest.approx(thr)
+
+
+def test_rc_time_to_throttle_is_infinite_when_the_threshold_is_never_reached():
+ # Steady state 325 K sits below the 330 K threshold.
+ low_power = rc(**{"forecast.rc.heat_power": 500.0})
+ assert value("forecast.rc.time_to_throttle", **low_power) == sp.oo
+ # Steady state exactly at the threshold is approached but not reached.
+ at_threshold = rc(**{"forecast.rc.heat_power": 600.0})
+ assert value("forecast.rc.time_to_throttle", **at_threshold) == sp.oo
+
+
+def test_rc_time_to_throttle_is_zero_when_already_at_or_above_the_threshold():
+ for t0 in (330.0, 335.0):
+ assert value(
+ "forecast.rc.time_to_throttle", **rc(**{"forecast.rc.initial_temperature": t0})
+ ) == 0
+
+
+def test_rc_cooling_margin_sign_matches_whether_the_threshold_is_reached():
+ sustainable = (RC["forecast.rc.throttle_temperature"] - RC["thermal.coolant.t_inlet"]) / RC["forecast.rc.resistance"]
+ assert value("forecast.rc.sustainable_power", **RC) == pytest.approx(sustainable)
+ for power in (500.0, 599.0, 601.0, 700.0):
+ margin = value("forecast.rc.cooling_margin", **rc(**{"forecast.rc.heat_power": power}))
+ assert margin == pytest.approx(sustainable - power)
+ t_hit = value("forecast.rc.time_to_throttle", **rc(**{"forecast.rc.heat_power": power}))
+ assert (margin > 0) == (t_hit == sp.oo)
+
+
+def test_rc_matches_the_package_junction_temperature_in_the_thermal_scope():
+ # One package: R = theta_ja and Q = package power, so the RC steady state
+ # equals the thermal scope's junction temperature.
+ theta_jc, theta_ca, power, supply = 0.015, 0.025, 1000.0, 305.0
+ theta_ja = theta_jc + theta_ca
+ junction = value(
+ "thermal.t_junction",
+ **{
+ "thermal.coolant.t_inlet": supply,
+ "thermal.theta_ca": theta_ca,
+ "thermal.theta_jc": theta_jc,
+ "gpu.power.total": power,
+ },
+ )
+ steady = value(
+ "forecast.rc.steady_temperature",
+ **{
+ "forecast.rc.resistance": theta_ja,
+ "forecast.rc.heat_power": power,
+ "thermal.coolant.t_inlet": supply,
+ },
+ )
+ assert steady == pytest.approx(junction)
+
+
+def test_rc_threshold_must_lie_above_supply_temperature():
+ ok = constraint(
+ "forecast.rc.sustainable_power", "forecast.ineq.rc.threshold_above_supply", **RC
+ )
+ bad = constraint(
+ "forecast.rc.sustainable_power",
+ "forecast.ineq.rc.threshold_above_supply",
+ **rc(**{"forecast.rc.throttle_temperature": 295.0}),
+ )
+ assert ok.satisfied is True
+ assert bad.satisfied is False
+
+
+# ---------------------------------------------------------------------------
+# Group 4: early checkpoint value
+# ---------------------------------------------------------------------------
+
+CKPT = {
+ "cluster.rel.optimal_checkpoint_interval": 3600.0, # s
+ "cluster.rel.checkpoint_time": 120.0, # s
+ "forecast.ckpt.gpus_affected": 1024,
+ "forecast.ckpt.recall": 0.8,
+ "forecast.ckpt.lead_time": 300.0,
+}
+
+
+def test_lost_work_without_and_with_warning():
+ unwarned = 0.5 * CKPT["cluster.rel.optimal_checkpoint_interval"] * CKPT["forecast.ckpt.gpus_affected"] / 3600.0
+ warned = CKPT["cluster.rel.checkpoint_time"] * CKPT["forecast.ckpt.gpus_affected"] / 3600.0
+ assert value("forecast.ckpt.lost_work_unwarned", **CKPT) == pytest.approx(unwarned)
+ assert value("forecast.ckpt.lost_work_warned", **CKPT) == pytest.approx(warned)
+ # 0.5 h on 1024 GPUs is 512 GPU-hours; one 120 s write is 1024/30 GPU-hours.
+ assert unwarned == pytest.approx(512.0)
+ assert warned == pytest.approx(1024.0 / 30.0)
+
+
+def test_expected_loss_and_gpu_hours_saved_weight_the_cases_by_recall():
+ unwarned = 0.5 * 3600.0 * 1024 / 3600.0
+ warned = 120.0 * 1024 / 3600.0
+ r = CKPT["forecast.ckpt.recall"]
+ expected = r * warned + (1 - r) * unwarned
+ assert value("forecast.ckpt.lost_work_expected", **CKPT) == pytest.approx(expected)
+ assert value("forecast.ckpt.gpu_hours_saved", **CKPT) == pytest.approx(unwarned - expected)
+ assert value("forecast.ckpt.gpu_hours_saved", **CKPT) == pytest.approx(r * (unwarned - warned))
+
+
+def test_recall_zero_saves_nothing_and_recall_one_saves_the_full_difference():
+ unwarned = 0.5 * 3600.0 * 1024 / 3600.0
+ warned = 120.0 * 1024 / 3600.0
+ none = {**CKPT, "forecast.ckpt.recall": 0.0}
+ full = {**CKPT, "forecast.ckpt.recall": 1.0}
+ assert value("forecast.ckpt.gpu_hours_saved", **none) == pytest.approx(0.0)
+ assert value("forecast.ckpt.gpu_hours_saved", **full) == pytest.approx(unwarned - warned)
+
+
+def test_gpu_hours_saved_scale_with_job_size_and_interval():
+ base = value("forecast.ckpt.gpu_hours_saved", **CKPT)
+ double_gpus = value("forecast.ckpt.gpu_hours_saved", **{**CKPT, "forecast.ckpt.gpus_affected": 2048})
+ assert double_gpus == pytest.approx(2 * base)
+ longer = value(
+ "forecast.ckpt.gpu_hours_saved",
+ **{**CKPT, "cluster.rel.optimal_checkpoint_interval": 7200.0},
+ )
+ assert longer > base
+
+
+def test_lead_time_must_cover_the_checkpoint_write():
+ ok = constraint(
+ "forecast.ckpt.gpu_hours_saved", "forecast.ineq.ckpt.lead_time_covers_checkpoint", **CKPT
+ )
+ short = constraint(
+ "forecast.ckpt.gpu_hours_saved",
+ "forecast.ineq.ckpt.lead_time_covers_checkpoint",
+ **{**CKPT, "forecast.ckpt.lead_time": 60.0},
+ )
+ exact = constraint(
+ "forecast.ckpt.gpu_hours_saved",
+ "forecast.ineq.ckpt.lead_time_covers_checkpoint",
+ **{**CKPT, "forecast.ckpt.lead_time": 120.0},
+ )
+ assert ok.satisfied is True
+ assert short.satisfied is False
+ assert exact.satisfied is True
+
+
+# ---------------------------------------------------------------------------
+# Group 5: time-shifting flexible work
+# ---------------------------------------------------------------------------
+
+SHIFT = {
+ "forecast.shift.flexible_fraction": 0.25,
+ "econ.power.price_kwh": 0.10, # USD/kWh, unshifted blended
+ "forecast.shift.price_shifted": 0.04, # USD/kWh
+ "econ.carbon.intensity_kg_per_kwh": 0.40, # kg/kWh, unshifted
+ "forecast.shift.carbon_intensity_shifted": 0.10,
+ "econ.job.dc_power": 1.0e6, # W
+ "training.tokens_per_sec": 2.0e5, # tokens/s
+ "econ.cost.per_token": 1.0e-6, # USD/token
+ "econ.carbon.price_per_tonne": 50.0, # USD/t
+}
+
+
+def test_energy_weighted_price_and_carbon_paid():
+ phi = SHIFT["forecast.shift.flexible_fraction"]
+ price = (1 - phi) * SHIFT["econ.power.price_kwh"] + phi * SHIFT["forecast.shift.price_shifted"]
+ carbon = (
+ (1 - phi) * SHIFT["econ.carbon.intensity_kg_per_kwh"]
+ + phi * SHIFT["forecast.shift.carbon_intensity_shifted"]
+ )
+ assert value("forecast.shift.price_paid", **SHIFT) == pytest.approx(price)
+ assert value("forecast.shift.carbon_intensity_paid", **SHIFT) == pytest.approx(carbon)
+
+
+def test_facility_energy_cost_and_carbon_per_token():
+ energy = SHIFT["econ.job.dc_power"] / SHIFT["training.tokens_per_sec"] # J/token
+ phi = SHIFT["forecast.shift.flexible_fraction"]
+ price = (1 - phi) * 0.10 + phi * 0.04
+ carbon = (1 - phi) * 0.40 + phi * 0.10
+ assert energy == pytest.approx(5.0)
+ assert value("forecast.shift.facility_energy_per_token", **SHIFT) == pytest.approx(energy)
+ assert value("forecast.shift.energy_cost_per_token", **SHIFT) == pytest.approx(
+ energy * price / 3.6e6
+ )
+ assert value("forecast.shift.carbon_per_token", **SHIFT) == pytest.approx(
+ energy * carbon / 3.6e6
+ )
+ assert value("forecast.shift.carbon_per_token_unshifted", **SHIFT) == pytest.approx(
+ energy * 0.40 / 3.6e6
+ )
+
+
+def test_shifted_cost_per_token_swaps_in_shifted_electricity_and_carbon_cost():
+ phi = 0.25
+ energy = 5.0
+ price_paid = (1 - phi) * 0.10 + phi * 0.04
+ m0 = energy * 0.40 / 3.6e6
+ m1 = energy * ((1 - phi) * 0.40 + phi * 0.10) / 3.6e6
+ expected = (
+ SHIFT["econ.cost.per_token"]
+ + energy * (price_paid - 0.10) / 3.6e6
+ + 50.0 * (m1 - m0) / 1000.0
+ )
+ assert value("forecast.shift.cost_per_token", **SHIFT) == pytest.approx(expected)
+ assert value("forecast.shift.cost_saving_per_token", **SHIFT) == pytest.approx(
+ SHIFT["econ.cost.per_token"] - expected
+ )
+ assert value("forecast.shift.carbon_saving_per_token", **SHIFT) == pytest.approx(m0 - m1)
+
+
+def test_no_shifting_leaves_cost_and_carbon_per_token_unchanged():
+ none = {**SHIFT, "forecast.shift.flexible_fraction": 0.0}
+ assert value("forecast.shift.price_paid", **none) == pytest.approx(0.10)
+ assert value("forecast.shift.cost_per_token", **none) == pytest.approx(SHIFT["econ.cost.per_token"])
+ assert value("forecast.shift.cost_saving_per_token", **none) == pytest.approx(0.0, abs=1e-18)
+ assert value("forecast.shift.carbon_saving_per_token", **none) == pytest.approx(0.0, abs=1e-18)
+
+
+def test_shifting_everything_pays_the_shifted_hours_price_and_intensity():
+ all_in = {**SHIFT, "forecast.shift.flexible_fraction": 1.0}
+ assert value("forecast.shift.price_paid", **all_in) == pytest.approx(0.04)
+ assert value("forecast.shift.carbon_intensity_paid", **all_in) == pytest.approx(0.10)
+
+
+def test_savings_grow_with_the_flexible_fraction_when_shifted_hours_are_cheaper_and_cleaner():
+ savings = [
+ value(
+ "forecast.shift.cost_saving_per_token",
+ **{**SHIFT, "forecast.shift.flexible_fraction": phi},
+ )
+ for phi in (0.0, 0.25, 0.5, 1.0)
+ ]
+ carbon = [
+ value(
+ "forecast.shift.carbon_saving_per_token",
+ **{**SHIFT, "forecast.shift.flexible_fraction": phi},
+ )
+ for phi in (0.0, 0.25, 0.5, 1.0)
+ ]
+ assert savings == sorted(savings) and savings[-1] > savings[0]
+ assert carbon == sorted(carbon) and carbon[-1] > carbon[0]
+
+
+def test_shifting_to_dearer_hours_gives_negative_cost_saving():
+ dearer = {**SHIFT, "forecast.shift.price_shifted": 0.20, "econ.carbon.price_per_tonne": 0.0}
+ assert value("forecast.shift.cost_saving_per_token", **dearer) < 0
+
+
+def test_economics_run_power_cost_agrees_with_facility_energy_per_token():
+ # Run electricity cost per token in the economics scope is
+ # price_ws * job power * wall-clock / tokens, and wall-clock = tokens / TPS,
+ # so it equals the facility energy per token times the blended price.
+ tokens = 1.0e12
+ tps = SHIFT["training.tokens_per_sec"]
+ wallclock = tokens / tps
+ run_cost = value(
+ "econ.run.power_cost",
+ **{
+ "econ.power.price_kwh": SHIFT["econ.power.price_kwh"],
+ "econ.job.dc_power": SHIFT["econ.job.dc_power"],
+ "training.wallclock": wallclock,
+ },
+ )
+ energy = value("forecast.shift.facility_energy_per_token", **SHIFT)
+ assert run_cost / tokens == pytest.approx(energy * SHIFT["econ.power.price_kwh"] / 3.6e6)
diff --git a/tests/test_research_heatwave_cost.py b/tests/test_research_heatwave_cost.py
new file mode 100644
index 0000000..c2309a7
--- /dev/null
+++ b/tests/test_research_heatwave_cost.py
@@ -0,0 +1,430 @@
+"""Tests for the Heatwave cost-timing module (dollars per token against carbon per token).
+
+The scheduler tests use small synthetic price and carbon series (labelled synthetic here, used only
+to check the code); the result-file test reads the committed result when it exists.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import pytest
+from scipy.optimize import linprog
+
+from gpu_stack.research import heatwave_cost as hc
+
+HERE = Path(__file__).resolve().parent
+RESULT = HERE.parent / "experiments" / "e003-heatwave" / "results" / "cost-v1.json"
+SCENARIO = HERE.parent / "experiments" / "e003-heatwave" / "cost-scenario-v1.json"
+
+
+def synthetic_series(n_hours: int, seed: int = 0):
+ """Synthetic daily cycle plus noise; price and carbon are correlated through a shared term."""
+ rng = np.random.default_rng(seed)
+ t = np.arange(n_hours)
+ shared = np.sin(2 * np.pi * t / 24.0)
+ price = 40.0 + 15.0 * shared + 6.0 * rng.normal(size=n_hours)
+ carbon = 200.0 + 60.0 * shared + 25.0 * np.sin(2 * np.pi * t / 24.0 + 1.0) + 15.0 * rng.normal(size=n_hours)
+ return price, carbon
+
+
+NORM = {"price_mean": 40.0, "price_std": 15.0, "carbon_mean": 200.0, "carbon_std": 60.0}
+
+
+# ---------------------------------------------------------------------------------------------
+# Model and planning problem
+# ---------------------------------------------------------------------------------------------
+
+
+def test_quantile_grid_matches_heatwave():
+ forecast = pytest.importorskip("heatwave.forecast")
+ assert tuple(forecast.QUANTILE_LEVELS) == hc.QUANTILE_LEVELS
+
+
+def test_flex_model_validation_and_rates():
+ m = hc.FlexModel(0.7, 0.3, 24)
+ assert m.flexible_rate == pytest.approx(0.21)
+ assert m.inflexible_load == pytest.approx(0.49)
+ assert m.flexible_cap == pytest.approx(0.51)
+ with pytest.raises(ValueError):
+ hc.FlexModel(1.2, 0.3, 24, power_cap=1.0)
+ with pytest.raises(ValueError):
+ hc.FlexModel(0.7, 1.5, 24)
+ with pytest.raises(ValueError):
+ hc.FlexModel(0.7, 0.3, 0)
+
+
+def test_arrived_and_due_window_one_means_immediate():
+ m = hc.FlexModel(0.7, 0.3, 1)
+ t = np.arange(10)
+ assert np.allclose(m.arrived(t), m.due(t))
+
+
+def _transport_optimum(cost, model, t0, run_before, terminal):
+ """Same problem as a transportation LP over (arrival hour, run hour) pairs, to check the cumulative form.
+
+ Work that arrived before ``t0`` and has not run yet is split into arrival cohorts oldest first
+ (first-in first-out is earliest-deadline first for one window length).
+ """
+ h = len(cost)
+ a = model.flexible_rate
+ arrivals = np.arange(0, t0 + h)
+ # cohort remaining mass: cohorts before t0 are served oldest first up to run_before
+ mass = np.full(arrivals.size, a)
+ left = run_before
+ for i in range(t0):
+ take = min(mass[i], left)
+ mass[i] -= take
+ left -= take
+ cohorts = [(i, mass[i]) for i in range(arrivals.size) if mass[i] > 1e-12 or i >= t0]
+ n_c = len(cohorts)
+ var = n_c * (h + 1) # last column = carried past the horizon
+ c = np.zeros(var)
+ bounds = []
+ for ci, (i, m) in enumerate(cohorts):
+ for k in range(h + 1):
+ idx = ci * (h + 1) + k
+ if k < h:
+ c[idx] = cost[k] - terminal
+ ok = (t0 + k >= i) and (t0 + k <= i + model.window_hours - 1)
+ else:
+ ok = i + model.window_hours - 1 >= t0 + h # deadline beyond the horizon
+ bounds.append((0.0, None) if ok else (0.0, 0.0))
+ a_eq = np.zeros((n_c, var))
+ b_eq = np.zeros(n_c)
+ for ci, (i, m) in enumerate(cohorts):
+ a_eq[ci, ci * (h + 1):(ci + 1) * (h + 1)] = 1.0
+ b_eq[ci] = m
+ # arrivals after t0 are known inside the horizon, but a cohort arriving at hour i >= t0 is only available from hour i
+ a_ub = np.zeros((h, var))
+ for k in range(h):
+ for ci in range(n_c):
+ a_ub[k, ci * (h + 1) + k] = 1.0
+ res = linprog(c, A_ub=a_ub, b_ub=np.full(h, model.flexible_cap), A_eq=a_eq, b_eq=b_eq, bounds=bounds, method="highs")
+ assert res.status == 0
+ return res.fun
+
+
+@pytest.mark.parametrize("window", [3, 6])
+def test_cumulative_form_matches_transportation_form(window):
+ # Only cohorts whose deadline lies inside the horizon are forced; others may be carried (cost = terminal value).
+ rng = np.random.default_rng(1)
+ m = hc.FlexModel(0.6, 0.5, window)
+ h, t0 = 12, 20
+ for _ in range(5):
+ cost = rng.normal(size=h)
+ run_before = float(m.due(t0 - 1)) + rng.uniform(0.0, 0.2)
+ y, fell_back = hc.plan_flexible(cost, m, t0, run_before)
+ assert not fell_back
+ terminal = float(cost.mean())
+ obj = float(np.dot(cost - terminal, y))
+ # the transportation form drops cohorts' carried mass cost (terminal - terminal = 0), so objectives are comparable
+ assert obj == pytest.approx(_transport_optimum(cost, m, t0, run_before, terminal), abs=1e-7)
+
+
+def test_plan_respects_bounds():
+ m = hc.FlexModel(0.7, 0.3, 24)
+ rng = np.random.default_rng(2)
+ t0, run_before = 48, float(m.due(47)) + 0.3
+ y, fb = hc.plan_flexible(rng.normal(size=48), m, t0, run_before)
+ assert not fb
+ t = np.arange(t0, t0 + 48)
+ cum = run_before + np.cumsum(y)
+ assert np.all(y >= -1e-9) and np.all(y <= m.flexible_cap + 1e-9)
+ assert np.all(cum >= m.due(t) - 1e-7)
+ assert np.all(cum <= m.arrived(t) + 1e-7)
+
+
+def test_rolling_schedule_meets_deadlines_and_conserves_work():
+ n_days = 12
+ price, carbon = synthetic_series(n_days * 24 + 48, seed=3)
+ for window in (24, 72):
+ m = hc.FlexModel(0.7, 0.5, window)
+ cost_by_day = np.stack([hc.mixed_cost(price[24 * d:24 * d + 48], carbon[24 * d:24 * d + 48], 0.5, NORM) for d in range(n_days)])
+ y, diag = hc.schedule_forecast(cost_by_day, m, [(0, n_days)])
+ assert diag["plan_fallbacks"] == 0
+ t = np.arange(y.size)
+ cum = np.cumsum(y)
+ assert np.all(cum >= m.due(t) - 1e-6), "work finished after its deadline"
+ assert np.all(cum <= m.arrived(t) + 1e-6), "work run before it arrived"
+ backlog = m.arrived(t[-1]) - cum[-1]
+ assert 0.0 <= backlog <= m.flexible_rate * window + 1e-6
+ assert np.all(y <= m.flexible_cap + 1e-9)
+
+
+def test_window_one_and_zero_flexible_fraction_equal_no_shift():
+ n_days = 5
+ price, carbon = synthetic_series(n_days * 24 + 48, seed=4)
+ for m in (hc.FlexModel(0.7, 0.4, 1), hc.FlexModel(0.7, 0.0, 24)):
+ cost_by_day = np.stack([hc.mixed_cost(price[24 * d:24 * d + 48], carbon[24 * d:24 * d + 48], 0.3, NORM) for d in range(n_days)])
+ y, _ = hc.schedule_forecast(cost_by_day, m, [(0, n_days)])
+ assert np.allclose(y, hc.schedule_none(m, n_days * 24), atol=1e-7)
+
+
+def test_perfect_foresight_lowers_cost_and_carbon_at_the_ends_of_lambda():
+ n_days = 20
+ n = n_days * 24
+ price, carbon = synthetic_series(n + 48, seed=5)
+ m = hc.FlexModel(0.7, 0.5, 24)
+ pf_p = hc.padded_future(price, [24 * d for d in range(n_days)], 48)
+ pf_c = hc.padded_future(carbon, [24 * d for d in range(n_days)], 48)
+ sums, flex, fb = hc.sweep_forecast_policy(pf_p, pf_c, [0.0, 1.0], m, NORM, [(0, n_days)], price[:n], carbon[:n])
+ none_s, _ = hc.sweep_none(m, 2, price[:n], carbon[:n])
+ p0, c0 = hc.effective_prices(none_s[0])
+ p_cost, c_cost = hc.effective_prices(sums[1]) # lam = 1: price only
+ p_carb, c_carb = hc.effective_prices(sums[0]) # lam = 0: carbon only
+ assert fb == 0
+ assert p_cost < p0 and c_carb < c0
+ assert p_cost <= p_carb + 1e-9 # planning on price gives the lower price
+ assert c_carb <= c_cost + 1e-9 # planning on carbon gives the lower carbon
+
+
+def test_risk_quantile_changes_plan_when_spread_differs_by_hour():
+ # Two hours with equal medians; the second has a wide upper tail. A median planner is indifferent, a q=0.9 planner avoids it.
+ m = hc.FlexModel(0.5, 1.0, 4)
+ cost_mid = np.array([1.0, 1.0, 5.0, 5.0])
+ cost_high = np.array([1.0, 3.0, 5.0, 5.0])
+ y_mid, _ = hc.plan_flexible(cost_mid, m, 0, 0.0)
+ y_high, _ = hc.plan_flexible(cost_high, m, 0, 0.0)
+ assert y_high[0] >= y_mid[0] - 1e-9
+ assert y_high[1] <= y_mid[1] + 1e-9
+
+
+def test_reactive_obeys_deadline_and_cap_and_constant_cost_runs_on_arrival():
+ n_days = 10
+ price, carbon = synthetic_series(n_days * 24 + 200, seed=6)
+ m = hc.FlexModel(0.7, 0.5, 24)
+ cost = hc.mixed_cost(price, carbon, 0.5, NORM)
+ y = hc.schedule_reactive(cost, m, [(200, n_days)], lookback=168)
+ t = np.arange(y.size)
+ cum = np.cumsum(y)
+ assert np.all(cum >= m.due(t) - 1e-9) and np.all(cum <= m.arrived(t) + 1e-9)
+ assert np.all(y <= m.flexible_cap + 1e-9)
+ flat = np.zeros(n_days * 24 + 200)
+ y_flat = hc.schedule_reactive(flat, m, [(200, n_days)], lookback=168)
+ assert np.allclose(y_flat, hc.schedule_none(m, n_days * 24))
+
+
+def test_trailing_median_matches_numpy():
+ x = np.array([5.0, 1.0, 3.0, 2.0, 9.0, 4.0])
+ out = hc.trailing_median(x, 3)
+ assert np.isnan(out[:2]).all()
+ assert out[2:] == pytest.approx([3.0, 2.0, 3.0, 4.0])
+
+
+# ---------------------------------------------------------------------------------------------
+# Accounting
+# ---------------------------------------------------------------------------------------------
+
+
+def test_no_shift_effective_values_are_plain_means():
+ n = 10 * 24
+ price, carbon = synthetic_series(n, seed=7)
+ m = hc.FlexModel(0.7, 0.3, 24)
+ s = hc.daily_sums(hc.schedule_none(m, n), m, price, carbon)
+ p, c = hc.effective_prices(s)
+ assert p == pytest.approx(price.mean()) and c == pytest.approx(carbon.mean())
+
+
+def test_shift_identity_matches_the_power_forecast_scope():
+ # price paid = (1 - phi) * unshifted price + phi * shifted price, with phi the flexible energy share
+ n_days = 15
+ n = n_days * 24
+ price, carbon = synthetic_series(n + 48, seed=8)
+ m = hc.FlexModel(0.7, 0.4, 24)
+ pf_p = hc.padded_future(price, [24 * d for d in range(n_days)], 48)
+ pf_c = hc.padded_future(carbon, [24 * d for d in range(n_days)], 48)
+ sums, flex, _ = hc.sweep_forecast_policy(pf_p, pf_c, [0.5], m, NORM, [(0, n_days)], price[:n], carbon[:n])
+ p, c = hc.effective_prices(sums[0])
+ fl = flex[0].sum(axis=0)
+ phi = fl[0] / sums[0][:, 0].sum()
+ p_sh, c_sh = fl[1] / fl[0], fl[2] / fl[0]
+ assert p == pytest.approx((1 - phi) * price[:n].mean() + phi * p_sh)
+ assert c == pytest.approx((1 - phi) * carbon[:n].mean() + phi * c_sh)
+
+
+def test_bootstrap_indices_shape_range_and_blocks():
+ idx = hc.block_bootstrap_indices(50, 200, 7, seed=0)
+ assert idx.shape == (200, 50) and idx.min() >= 0 and idx.max() < 50
+ # inside a block the days are consecutive (circular)
+ assert np.all((idx[:, 1:7] - idx[:, :6]) % 50 == 1)
+ iid = hc.block_bootstrap_indices(50, 100, 1, seed=0)
+ assert iid.shape == (100, 50)
+ assert np.array_equal(idx, hc.block_bootstrap_indices(50, 200, 7, seed=0))
+
+
+def test_effective_prices_under_bootstrap_has_boot_shape():
+ sums = np.abs(np.random.default_rng(0).normal(size=(30, 3))) + 1.0
+ idx = hc.block_bootstrap_indices(30, 40, 3, seed=1)
+ p, c = hc.effective_prices(sums, idx)
+ assert p.shape == (40,) and c.shape == (40,)
+ lo, hi = hc.percentile_interval(p)
+ assert lo <= np.median(p) <= hi
+
+
+def test_fill_hourly_uses_same_hour_of_neighbouring_days():
+ x = np.tile(np.arange(24, dtype=float), 6)
+ y = x.copy()
+ y[50] = np.nan
+ filled, n = hc.fill_hourly(y)
+ assert n == 1 and filled[50] == pytest.approx(x[50])
+
+
+def test_padded_future_copies_previous_day_past_the_end():
+ v = np.arange(48, dtype=float)
+ out = hc.padded_future(v, [40], 16)
+ assert out[0, 0] == 40 and out[0, 7] == 47
+ assert out[0, 8] == 24 + 0 # hour 48 -> hour 24
+ assert out[0, 15] == 31
+
+
+def test_quantile_index_requires_grid_level():
+ assert hc.quantile_index(0.5) == 10
+ with pytest.raises(ValueError):
+ hc.quantile_index(0.123)
+
+
+# ---------------------------------------------------------------------------------------------
+# Calculator terms and the feedback path
+# ---------------------------------------------------------------------------------------------
+
+REF = {"params": 7e9, "tokens": 2e12, "gpu": "H100-SXM", "n_gpus": 1024, "overrides": {}}
+
+
+def test_reference_terms_agree_with_calculator():
+ from gpu_stack.calculator import estimate
+
+ terms = hc.reference_terms(REF)
+ est = estimate(7e9, 2e12, gpu="H100-SXM", n_gpus=1024)
+ assert terms["energy_mwh_per_mtoken"] * est.tokens / 1e6 == pytest.approx(est.energy_mwh)
+ assert terms["facility_power_mw"] * est.training_days * 24 == pytest.approx(est.energy_mwh)
+ assert terms["tokens_per_gpu_hour"] * est.gpu_hours == pytest.approx(est.tokens)
+ default_elec = terms["energy_mwh_per_mtoken"] * 1000 * terms["calculator_default_electricity_price_usd_per_kwh"] * est.tokens / 1e6
+ assert default_elec == pytest.approx(est.electricity_cost)
+ assert terms["hardware_usd_per_mtoken"] * est.tokens / 1e6 == pytest.approx(est.hardware_cost)
+
+
+def test_per_mtoken_units():
+ terms = hc.reference_terms(REF)
+ r = hc.per_mtoken(50.0, 200.0, terms) # 50 USD/MWh, 200 g/kWh
+ assert r["electricity_usd_per_mtoken"] == pytest.approx(terms["energy_mwh_per_mtoken"] * 50.0)
+ assert r["co2_g_per_mtoken"] == pytest.approx(terms["energy_mwh_per_mtoken"] * 1000.0 * 200.0)
+ assert r["total_usd_per_mtoken"] == pytest.approx(r["electricity_usd_per_mtoken"] + terms["hardware_usd_per_mtoken"])
+
+
+def _toy_curve():
+ return {
+ "flexible_fraction": 0.3, "window_hours": 24, "utilization": 0.7, "lambda": [0.0, 1.0],
+ "price_paid": [36.0, 30.0], "carbon_paid": [100.0, 110.0], "price_shifted": [20.0, 10.0],
+ "carbon_shifted": [60.0, 90.0], "flexible_energy_share": [0.3, 0.3],
+ "price_unshifted": 40.0, "carbon_unshifted": 120.0,
+ "electricity_usd_per_mtoken": [0.0, 0.0], "total_usd_per_mtoken": [0.0, 0.0], "co2_g_per_mtoken": [0.0, 0.0],
+ }
+
+
+def test_timing_inputs_units_and_interpolation():
+ ti = hc.timing_inputs(_toy_curve(), 0.5)
+ assert ti["electricity_price"] == pytest.approx(33.0 / 1000.0)
+ assert ti["carbon_intensity_paid"] == pytest.approx(105.0 / 1000.0)
+ assert ti["price_shifted"] == pytest.approx(15.0 / 1000.0)
+ ti_eur = hc.timing_inputs(_toy_curve(), 0.0, usd_per_price_unit=1.1)
+ assert ti_eur["electricity_price"] == pytest.approx(36.0 / 1000.0 * 1.1)
+
+
+def test_estimate_with_timing_changes_only_electricity():
+ from gpu_stack.calculator import estimate
+
+ base = estimate(7e9, 2e12, gpu="H100-SXM", n_gpus=1024, electricity_price=0.040)
+ est, carbon_kg, ti = hc.estimate_with_timing(REF, _toy_curve(), 0.0)
+ assert est.energy_mwh == pytest.approx(base.energy_mwh)
+ assert est.hardware_cost == pytest.approx(base.hardware_cost)
+ assert est.electricity_cost == pytest.approx(est.energy_mwh * 1000.0 * ti["electricity_price"])
+ assert carbon_kg == pytest.approx(est.energy_mwh * 1000.0 * ti["carbon_intensity_paid"])
+
+
+def test_lookup_curve_keys():
+ doc = {"result": {"scenario": {"workload": {"utilization": 0.7}}, "pairs": {"x": {"curves": {
+ hc.curve_key("timestral_28M", 0.3, 24, 0.7, 0.5): {"tag": "a"},
+ hc.curve_key("no_shift", 0.3, 24, 0.7, None): {"tag": "b"}}}}}}
+ assert hc.lookup_curve(doc, "x", "timestral_28M", 0.3, 24)["tag"] == "a"
+ assert hc.lookup_curve(doc, "x", "no_shift", 0.3, 24)["tag"] == "b"
+
+
+# ---------------------------------------------------------------------------------------------
+# Forecast use is causal
+# ---------------------------------------------------------------------------------------------
+
+
+def test_rolling_forecast_only_reads_history():
+ forecast = pytest.importorskip("heatwave.forecast")
+ values = synthetic_series(24 * 40)[0]
+ fc = forecast.SeasonalNaiveForecaster(24, max_context=168)
+ t0 = [24 * 20, 24 * 21]
+ a = hc.rolling_forecast(values, t0, fc, 168, 48, batch=2)
+ changed = values.copy()
+ changed[24 * 20:] += 1000.0 # change everything at and after the first decision time
+ b = hc.rolling_forecast(changed, t0[:1], fc, 168, 48, batch=1)
+ assert a.shape == (2, 48, 21)
+ assert np.allclose(a[0], b[0])
+
+
+# ---------------------------------------------------------------------------------------------
+# The committed result
+# ---------------------------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not RESULT.exists(), reason="result file not written yet")
+def test_committed_result_has_provenance_and_consistent_curves():
+ doc = json.loads(RESULT.read_text())
+ prov = doc["provenance"]
+ assert prov["data_kind"] == "public"
+ assert len(prov["script_sha256"]) == 64 and prov["git"]["commit"] != "unknown"
+ assert any("public_caiso_np15_lmp_dam" in k for k in prov["inputs"])
+ res = doc["result"]
+ assert res["gpu_stack_git"]["commit"]
+ scenario = json.loads(SCENARIO.read_text())
+ assert res["scenario"]["scenario_id"] == scenario["scenario_id"]
+ for pair, block in res["pairs"].items():
+ assert block["same_region"] is True and block["plan_fallbacks_total"] == 0
+ for key, curve in block["curves"].items():
+ n = len(curve["lambda"])
+ for field in ("price_paid", "carbon_paid", "electricity_usd_per_mtoken", "total_usd_per_mtoken", "co2_g_per_mtoken"):
+ assert len(curve[field]) == n and len(curve[field + "_ci"]) == n
+ # total cost = electricity + the calculator's hardware cost, the same for every policy
+ hw = res["reference_run_terms"]["hardware_usd_per_mtoken"]
+ tot = np.array(curve["total_usd_per_mtoken"]) - np.array(curve["electricity_usd_per_mtoken"])
+ assert np.allclose(tot, hw)
+ none = block["curves"][hc.curve_key("no_shift", 0.3, 24, 0.7, None)]
+ assert np.allclose(none["saving_pct_electricity"], 0.0)
+ assert np.allclose(none["price_paid"], block["mean_price_test"])
+
+
+def test_offline_optimum_is_feasible_and_no_worse_than_rolling_perfect_foresight():
+ n_days = 10
+ n = n_days * 24
+ price, carbon = synthetic_series(n + 48, seed=9)
+ m = hc.FlexModel(0.7, 0.5, 24)
+ cost = hc.mixed_cost(price[:n], carbon[:n], 0.5, NORM)
+ y_off = hc.offline_optimum(cost, m)
+ t = np.arange(n)
+ cum = np.cumsum(y_off)
+ assert np.all(cum >= m.due(t) - 1e-6) and np.all(cum <= m.arrived(t) + 1e-6)
+ assert np.all(y_off <= m.flexible_cap + 1e-9) and np.all(y_off >= -1e-9)
+ pf = hc.padded_future(price, [24 * d for d in range(n_days)], 48)
+ pc = hc.padded_future(carbon, [24 * d for d in range(n_days)], 48)
+ cost_rows = np.stack([hc.mixed_cost(pf[d], pc[d], 0.5, NORM) for d in range(n_days)])
+ y_roll, _ = hc.schedule_forecast(cost_rows, m, [(0, n_days)])
+ ref = cost.mean()
+ obj_off = float(np.dot(cost - ref, y_off))
+ obj_roll = float(np.dot(cost - ref, y_roll))
+ assert obj_off <= obj_roll + 1e-7
+
+
+def test_covariate_windows_pad_past_the_end_with_nan():
+ known = np.arange(2 * 100, dtype=float).reshape(2, 100)
+ w = hc.covariate_windows(known, [90], 20, 24)
+ assert w.shape == (1, 2, 44)
+ assert w[0, 0, 0] == known[0, 70] and w[0, 0, 29] == known[0, 99]
+ assert np.isnan(w[0, :, 30:]).all()
diff --git a/tests/test_research_heatwave_outage.py b/tests/test_research_heatwave_outage.py
new file mode 100644
index 0000000..af8e116
--- /dev/null
+++ b/tests/test_research_heatwave_outage.py
@@ -0,0 +1,806 @@
+"""Tests for the Heatwave outage-response module.
+
+Every run here is simulated. Horizons are short and job lists tiny; the full
+experiment lives in ``experiments/e003-heatwave/run_outage.py``. Expected
+values are computed from the configuration the simulator builds, never typed
+in as power or checkpoint numbers.
+"""
+
+from __future__ import annotations
+
+import copy
+from dataclasses import replace
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from gpu_stack.research.heatwave_events import GridCurtailment, HardwareFault, Precursor
+from gpu_stack.research.heatwave_outage import (
+ EpisodeTracker,
+ ForecastWarningController,
+ OracleAlarm,
+ OracleController,
+ OutageSettings,
+ PersistenceForecaster,
+ TraceRecorder,
+ aggregate_runs,
+ bootstrap_ci,
+ build_event_table,
+ episodes_from_pmax,
+ exceedance_probability,
+ job_costs,
+ load_outage_scenario,
+ oracle_alarms_from_reference,
+ oracle_scenario,
+ oracle_settings,
+ periodic_local_jobs,
+ ratio_ci,
+ run_policy,
+ score_run,
+ seed_scenario,
+)
+from gpu_stack.research.heatwave_sim import Job, build_config
+
+ROOT = Path(__file__).resolve().parents[1]
+SCENARIO = ROOT / "experiments" / "e003-heatwave" / "outage-scenario-v1.json"
+
+LEVELS = np.linspace(0.05, 0.95, 21)
+
+
+def outage_and_settings():
+ outage = load_outage_scenario(SCENARIO)
+ return outage, OutageSettings.from_scenario(outage)
+
+
+def make_job(job_id="j1", n=8, start=0.0, duration=7200.0, step=4.0, ckpt_every=3600.0, ckpt_dur=30.0, **kw) -> Job:
+ values = dict(
+ job_id=job_id, n_gpus=n, submit_s=0.0, start_s=start, duration_s=duration, step_time_s=step,
+ checkpoint_every_s=ckpt_every, checkpoint_duration_s=ckpt_dur, value_per_gpu_hour=2.0, flexible=False, site_id="site-a",
+ )
+ values.update(kw)
+ return Job(**values)
+
+
+def tiny_scenario(outage, horizon_s, jobs, events=()):
+ sc = copy.deepcopy(outage["sim"])
+ sc["horizon_s"] = horizon_s
+ sc["events"] = [e.to_dict() for e in events]
+ sc["jobs"] = {"list": [j.to_dict() for j in jobs]}
+ sc["controller_interval_s"] = 60.0
+ return sc
+
+
+def fast_settings(S: OutageSettings) -> OutageSettings:
+ return replace(S, decision_interval_s=60.0, context_steps=12, horizon_steps=6, warmup_s=0.0, recheckpoint_s=300.0)
+
+
+# ---------------------------------------------------------------------------
+# scenario and seeds
+# ---------------------------------------------------------------------------
+
+
+def test_scenario_is_labeled_simulated_and_loads():
+ outage, S = outage_and_settings()
+ assert outage["data_kind"] == "simulated"
+ assert outage["sim"]["data_kind"] == "simulated"
+ assert set(outage["seeds"]["train"]).isdisjoint(outage["seeds"]["test"])
+ assert S.horizon_s == S.horizon_steps * S.decision_interval_s
+ assert outage["sim"]["controller_interval_s"] == S.decision_interval_s
+
+
+def test_seed_scenario_is_deterministic_and_event_free_variant_is_empty():
+ outage, S = outage_and_settings()
+ a, b = seed_scenario(outage, 201), seed_scenario(outage, 201)
+ assert a == b
+ assert seed_scenario(outage, 202)["events"] != a["events"]
+ assert a["events"], "the generator draws events"
+ assert all(e["onset_s"] >= S.warmup_s for e in a["events"])
+ assert seed_scenario(outage, 201, with_events=False)["events"] == []
+ assert a["jobs"]["generate"]["seed"] != seed_scenario(outage, 202)["jobs"]["generate"]["seed"]
+
+
+def test_periodic_local_jobs_splits_even_single_site_jobs_only():
+ even, odd = make_job("e", n=8), make_job("o", n=7)
+ already = make_job("m", n=8, multisite=True, site_id=None, average_every_steps=4)
+ out = periodic_local_jobs([even, odd, already])
+ assert out[0].multisite and out[0].average_every_steps == 8 and out[0].site_id is None
+ assert out[1] == odd
+ assert out[2] == already
+
+
+# ---------------------------------------------------------------------------
+# probabilities and warning episodes
+# ---------------------------------------------------------------------------
+
+
+def test_exceedance_probability_from_quantiles():
+ values = np.broadcast_to(np.linspace(10.0, 20.0, 21), (3, 4, 21)).copy()
+ hi = exceedance_probability(values, LEVELS, np.full((3, 4), 25.0))
+ lo = exceedance_probability(values, LEVELS, np.full((3, 4), 5.0))
+ assert np.allclose(hi, 1.0 - LEVELS[-1])
+ assert np.allclose(lo, 1.0 - LEVELS[0])
+ mid = np.stack([exceedance_probability(values, LEVELS, np.full((3, 4), x))[0, 0] for x in (11.0, 14.0, 17.0, 19.0)])
+ assert np.all(np.diff(mid) < 0)
+ assert np.isclose(exceedance_probability(values, LEVELS, np.full((3, 4), 15.0))[0, 0], 0.5)
+
+
+def test_exceedance_probability_zero_spread_and_missing_limit():
+ flat = np.full((2, 3, 21), 7.0)
+ assert np.all(exceedance_probability(flat, LEVELS, np.full((2, 3), 6.0)) == 1.0)
+ assert np.all(exceedance_probability(flat, LEVELS, np.full((2, 3), 8.0)) == 0.0)
+ assert np.all(exceedance_probability(flat, LEVELS, np.full((2, 3), np.nan)) == 0.0)
+ with pytest.raises(ValueError):
+ exceedance_probability(flat, LEVELS[:-1], np.full((2, 3), 1.0))
+
+
+def test_episode_tracker_hysteresis_and_end_time():
+ tr = EpisodeTracker(2, threshold=0.5, clear_decisions=2)
+ series = [(0, 0.1), (1, 0.6), (2, 0.2), (3, 0.7), (4, 0.2), (5, 0.2), (6, 0.9)]
+ for t, p in series:
+ tr.update(float(t), np.array([p, 0.0]))
+ eps = tr.finish()
+ assert [(e["k"], e["start_s"], e["end_s"]) for e in eps] == [(0, 1.0, 3.0), (0, 6.0, 6.0)]
+ assert eps[0]["peak_p"] == pytest.approx(0.7)
+ assert eps[0]["n_decisions"] == 2
+
+
+def test_episodes_from_pmax_matches_tracker_and_labels_with_index():
+ pm = np.zeros((6, 2))
+ pm[2:4, 1] = 0.8
+ eps = episodes_from_pmax(np.arange(6) * 60.0, pm, 0.5, 1)
+ assert len(eps) == 1 and eps[0]["k"] == 1 and eps[0]["start_s"] == 120.0 and eps[0]["end_s"] == 180.0
+
+
+# ---------------------------------------------------------------------------
+# simulator interface
+# ---------------------------------------------------------------------------
+
+
+def test_hardware_precursor_is_visible_in_the_faulty_gpu_temperature():
+ outage, _ = outage_and_settings()
+ strength, lead, onset, slot = 6.0, 1200.0, 2400.0, 2
+
+ def event(with_precursor: bool):
+ pre = (Precursor("gpu_temp_c", lead, strength),) if with_precursor else ()
+ return HardwareFault("h1", "site-a", "a1", slot, onset, 120.0, precursors=pre)
+
+ class Watch:
+ def __init__(self):
+ self.rows = []
+
+ def decide(self, t_s, vs):
+ self.rows.append((t_s, np.array(vs.now["gpu_temp_c"], dtype=float)))
+ return ()
+
+ peaks = []
+ for flag in (True, False):
+ sc = tiny_scenario(outage, 3600, [make_job(duration=3600.0)], [event(flag)])
+ w = Watch()
+ _, cfg = run_policy(sc, w, 0)
+ racks = [r for lay in cfg.layouts for r in lay.racks]
+ first = sum(r.n_gpus for r in racks[: [r.rack_id for r in racks].index("a1")])
+ peaks.append(max(row[first + slot] for t, row in w.rows if t < onset))
+ assert peaks[0] - peaks[1] > 0.5 * strength
+
+
+def test_job_costs_follow_the_simulator_configuration():
+ outage, _ = outage_and_settings()
+ cfg = build_config(tiny_scenario(outage, 3600, [make_job(n=8, ckpt_dur=30.0)]), seed=0)
+ job = cfg.jobs[0]
+ c = job_costs(job, cfg.transfer)
+ assert c["checkpoint"] == pytest.approx(job.n_gpus * job.checkpoint_duration_s / 3600.0)
+ assert c["move"] > c["checkpoint"]
+ twice = job_costs(replace(job, checkpoint_duration_s=2 * job.checkpoint_duration_s), cfg.transfer)
+ assert twice["checkpoint"] == pytest.approx(2 * c["checkpoint"])
+
+
+# ---------------------------------------------------------------------------
+# scoring and policies on a tiny simulated window
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def faulted():
+ """One non-flexible job, a hardware fault without precursor, and a long checkpoint interval."""
+ outage, S = outage_and_settings()
+ S = fast_settings(S)
+ sc = tiny_scenario(outage, 7200, [make_job(duration=7200.0, ckpt_every=3000.0)],
+ [HardwareFault("h1", "site-a", "a1", 2, 2900.0, 600.0)])
+ rec = TraceRecorder(S)
+ res, cfg = run_policy(sc, rec, 0)
+ table = build_event_table(cfg, res.events, S)
+ ex = rec.export()
+ ex["episodes"] = []
+ ref = float(run_policy(tiny_scenario(outage, 7200, [make_job(duration=7200.0, ckpt_every=3000.0)]), None, 0)[0].summary()["useful_gpu_hours"])
+ score = score_run(res, cfg, table, ex, S, ref)
+ return {"outage": outage, "S": S, "sc": sc, "res": res, "cfg": cfg, "table": table, "score": score, "ref": ref}
+
+
+def test_reactive_loses_work_and_attribution_conserves_it(faulted):
+ res, score = faulted["res"], faulted["score"]
+ lost = res.summary()["lost_gpu_hours"]
+ assert lost > 0.0
+ attributed = sum(e["lost_gpu_h"] for e in score["by_event"]) + score["unattributed"]["lost_gpu_h"]
+ assert attributed == pytest.approx(lost, rel=1e-6)
+ fault_rows = [e for e in score["by_event"] if e["kind"] == "hardware_fault"]
+ assert len(fault_rows) == 1 and fault_rows[0]["lost_gpu_h"] > 0.0 and not fault_rows[0]["caught"]
+ assert score["response"]["n_alarms"] == 0 and score["response"]["total_gpu_h"] == 0.0
+ assert score["data_kind"] == "simulated"
+
+
+def test_perfect_foresight_loses_less_than_reactive_and_catches_the_event(faulted):
+ S, sc, cfg, res, table = faulted["S"], faulted["sc"], faulted["cfg"], faulted["res"], faulted["table"]
+ So = oracle_settings(S, cfg.control_dt_s)
+ alarms = oracle_alarms_from_reference(cfg, res.events, [], So)
+ assert alarms and all(isinstance(a, OracleAlarm) for a in alarms)
+ ctrl = OracleController(So, alarms)
+ res_o, cfg_o = run_policy(oracle_scenario(sc), ctrl, 0)
+ score_o = score_run(res_o, cfg_o, table, ctrl.export(), So, faulted["ref"])
+ assert res_o.summary()["lost_gpu_hours"] < 0.5 * res.summary()["lost_gpu_hours"]
+ assert all(e["caught_in_time"] for e in score_o["by_event"] if e["kind"] == "hardware_fault")
+ assert score_o["response"]["total_gpu_h"] > 0.0
+
+
+def test_persistence_warning_on_announced_curtailment_checkpoints_then_resumes():
+ outage, S = outage_and_settings()
+ S = fast_settings(S)
+ flex = make_job("f", n=16, duration=7200.0, flexible=True, ckpt_every=3000.0)
+ sc = tiny_scenario(outage, 7200, [flex], [GridCurtailment("g1", "site-a", 1800.0, 900.0, 0.2, announce_lead_s=1500.0)])
+ ctrl = ForecastWarningController(S, PersistenceForecaster(), 0.5, None, "persistence")
+ res, cfg = run_policy(sc, ctrl, 0)
+ ex = ctrl.export()
+ site_eps = [e for e in ex["episodes"] if e["class"] == "curtail"]
+ assert site_eps and site_eps[0]["start_s"] < 1800.0
+ applied = [a for a in res.actions if a["status"] == "applied"]
+ assert applied, "the response ran at least one action"
+ table = build_event_table(cfg, res.events, S)
+ score = score_run(res, cfg, table, ex, S, 0.0)
+ row = [e for e in score["by_event"] if e["kind"] == "grid_curtailment"][0]
+ assert row["caught"] and row["lead_s"] is not None and row["lead_s"] > 0.0
+
+
+# ---------------------------------------------------------------------------
+# aggregation
+# ---------------------------------------------------------------------------
+
+
+def test_bootstrap_helpers_are_deterministic_and_pooled():
+ vals = [1.0, 2.0, 3.0, 4.0]
+ a, b = bootstrap_ci(vals, 200, 3), bootstrap_ci(vals, 200, 3)
+ assert a == b and a["mean"] == pytest.approx(2.5) and a["lo"] <= a["mean"] <= a["hi"]
+ r = ratio_ci([1, 2, 3], [2, 4, 6], 200, 1)
+ assert r["value"] == pytest.approx(0.5) and r["lo"] == pytest.approx(0.5) and r["hi"] == pytest.approx(0.5)
+ assert bootstrap_ci([], 50)["n"] == 0
+
+
+def test_aggregate_hides_detection_fields_without_a_warning_rule(faulted):
+ S = faulted["S"]
+ runs = [faulted["score"]] * 3
+ with_w = aggregate_runs(runs, S, has_warnings=True)
+ without = aggregate_runs(runs, S, has_warnings=False)
+ assert with_w["overall"]["recall"]["value"] == 0.0
+ assert without["overall"]["recall"] is None and without["has_warnings"] is False
+ assert without["no_precursor"]["n_events"] == 3
+ assert without["no_precursor"]["lost_gpu_h"]["mean"] == pytest.approx(with_w["no_precursor"]["lost_gpu_h"]["mean"])
+ assert with_w["data_kind"] == "simulated"
+
+
+# ---------------------------------------------------------------------------
+# per-GPU residual channel (outage scenario v2)
+# ---------------------------------------------------------------------------
+
+SCENARIO_V2 = ROOT / "experiments" / "e003-heatwave" / "outage-scenario-v2.json"
+
+
+def gpu_settings(S: OutageSettings) -> OutageSettings:
+ """Per-GPU channel on, 60 s decisions, short baseline so tiny windows can use it."""
+ return replace(
+ S, decision_interval_s=60.0, context_steps=12, horizon_steps=6, warmup_s=0.0, recheckpoint_s=300.0,
+ anomaly_baseline_window_s=3600.0, anomaly_baseline_gap_s=600.0, per_gpu_channel=True,
+ )
+
+
+def v2_and_settings():
+ outage = load_outage_scenario(SCENARIO_V2)
+ return outage, OutageSettings.from_scenario(outage)
+
+
+class Layout:
+ """Dry run without events: which GPU slots of one rack hold which job, and each GPU's gain."""
+
+ def __init__(self):
+ self.job_slots = {}
+ self.k_gpu = None
+ self.rack_start = None
+
+ def bind(self, run, cfg):
+ self.run = run
+ self.k_gpu = np.array(run.k_gpu)
+ self.rack_start = np.array(run.rack_start)
+
+ def decide(self, t_s, vs):
+ if not self.job_slots and t_s > 0:
+ run = self.run
+ for sh in run.shards:
+ if sh.gpus.size:
+ self.job_slots.setdefault(sh.job.job_id, []).extend(int(g) for g in sh.gpus)
+ return ()
+
+
+def layout_of(outage, jobs, horizon_s=3600):
+ lay = Layout()
+ sc = tiny_scenario(outage, horizon_s, jobs)
+ cfg = build_config(sc, seed=0)
+ from gpu_stack.research.heatwave_outage import OutageRun
+
+ run = OutageRun(cfg, lay, 0)
+ lay.bind(run, cfg)
+ run.run()
+ return lay, cfg
+
+
+def test_scenario_v2_differs_from_v1_only_by_the_per_gpu_channel():
+ v1, S1 = outage_and_settings()
+ v2, S2 = v2_and_settings()
+ assert not S1.per_gpu_channel and S2.per_gpu_channel
+ assert v1["seeds"] == v2["seeds"] and v1["event_generation"] == v2["event_generation"] and v1["seed_streams"] == v2["seed_streams"]
+ a = {k: v for k, v in v1["outage"].items()}
+ b = {k: v for k, v in v2["outage"].items() if k != "per_gpu_channel"}
+ assert a == b
+ assert v2["data_kind"] == "simulated" and v2["sim"]["data_kind"] == "simulated"
+ # the simulated inputs of a seed do not depend on the base scenario version (both seeds are set per run seed)
+ assert seed_scenario(v1, 201)["events"] == seed_scenario(v2, 201)["events"]
+ assert seed_scenario(v1, 201)["jobs"]["generate"] == seed_scenario(v2, 201)["jobs"]["generate"]
+
+
+def test_series_layout_appends_one_series_per_gpu_after_the_v1_series():
+ outage, S = v2_and_settings()
+ cfg = build_config(seed_scenario(outage, 201, with_events=False), seed=201)
+ from gpu_stack.research.heatwave_outage import SeriesIndex, gpu_unit_id
+
+ base, full = SeriesIndex.from_config(cfg), SeriesIndex.from_config(cfg, True)
+ n_gpus = sum(r.n_gpus for lay in cfg.layouts for r in lay.racks)
+ assert len(full) == len(base) + n_gpus and full.keys[: len(base)] == base.keys
+ assert len(set(full.gpu_units)) == n_gpus and full.gpu_units[0] == gpu_unit_id("a1", 0)
+ assert all(full.gpu_parts(u) == (u.rsplit("/", 1)[0], int(u.rsplit("/", 1)[1])) for u in full.gpu_units)
+ assert set(full.classes()[len(base):]) == {"hardware"}
+
+
+def test_residual_cancels_job_power_changes_that_a_plain_difference_keeps():
+ from gpu_stack.research.heatwave_outage import GpuResidualChannel
+
+ outage, S = v2_and_settings()
+ S = gpu_settings(S)
+ rng = np.random.default_rng(0)
+ gains = 0.0256 * (1.0 + 0.1 * rng.normal(size=12))
+ key = np.zeros(12, dtype=np.int64)
+ ch = GpuResidualChannel(12, S, floor_k=1.0)
+ power = np.where(np.arange(40) % 8 < 4, 700.0, 300.0) # the whole job changes power every 4 decisions
+ out, plain = [], []
+ for p in power:
+ excess = gains * p
+ out.append(ch.update(excess, key))
+ plain.append(excess - np.median(excess))
+ out, plain = np.array(out), np.array(plain)
+ late = out[20:]
+ assert np.isfinite(late).all()
+ assert np.abs(late).max() < 1e-9
+ assert np.abs(plain[20:]).max() > 0.5 # the plain difference to the peer median moves by several tenths of a kelvin or more
+
+
+def test_residual_sees_a_creep_on_one_gpu_and_marks_group_changes_missing():
+ from gpu_stack.research.heatwave_outage import GpuResidualChannel
+
+ outage, S = v2_and_settings()
+ S = gpu_settings(S)
+ gains = np.full(10, 0.025)
+ key = np.zeros(10, dtype=np.int64)
+ ch = GpuResidualChannel(10, S, floor_k=1.0)
+ res = []
+ for i in range(60):
+ excess = gains * 600.0
+ if i >= 40:
+ excess = excess.copy()
+ excess[3] += 0.25 * (i - 40) # 0.25 K per decision on GPU 3 only
+ if i == 50:
+ key = key.copy()
+ key[7] = 5 # GPU 7 changes job group: its sample is missing at that decision
+ res.append(ch.update(excess, key))
+ res = np.array(res)
+ assert res[59, 3] == pytest.approx(0.25 * 19, abs=1e-6)
+ others = np.delete(res[59], 3)
+ assert np.nanmax(np.abs(others)) < 1e-9
+ assert np.isnan(res[50, 7]) and np.isfinite(res[49, 7])
+
+
+def test_residual_channel_uses_no_future_samples():
+ from gpu_stack.research.heatwave_outage import GpuResidualChannel
+
+ outage, S = v2_and_settings()
+ S = gpu_settings(S)
+ rng = np.random.default_rng(3)
+ n, g = 50, 16
+ gains = 0.0256 * (1.0 + 0.1 * rng.normal(size=g))
+ inputs = [(gains * rng.uniform(300, 700) + rng.normal(0, 0.1, size=g), np.zeros(g, dtype=np.int64)) for _ in range(n)]
+ d = 30
+ a = GpuResidualChannel(g, S, 1.0)
+ out_a = [a.update(x, k) for x, k in inputs]
+ changed = list(inputs[: d + 1]) + [(rng.uniform(0, 50, size=g), rng.integers(0, 3, size=g)) for _ in range(n - d - 1)]
+ b = GpuResidualChannel(g, S, 1.0)
+ out_b = [b.update(x, k) for x, k in changed]
+ for i in range(d + 1):
+ assert np.array_equal(out_a[i], out_b[i], equal_nan=True)
+ assert not np.array_equal(out_a[-1], out_b[-1], equal_nan=True)
+
+
+def test_channel_rows_before_a_precursor_start_do_not_depend_on_the_event():
+ outage, S = v2_and_settings()
+ S = gpu_settings(S)
+ job = make_job(n=16, duration=7200.0)
+ onset, lead = 4800.0, 1200.0
+ pre = (Precursor("gpu_temp_c", lead, 6.0),)
+ rows = []
+ for ev in ([], [HardwareFault("h1", "site-a", "a1", 3, onset, 120.0, precursors=pre)]):
+ rec = TraceRecorder(S)
+ run_policy(tiny_scenario(outage, 7200, [job], ev), rec, 0)
+ rows.append(rec.trace())
+ t = rows[0]["times_s"]
+ before = t < onset - lead - 1e-9
+ assert before.sum() > 20
+ assert np.array_equal(rows[0]["samples"][:, before], rows[1]["samples"][:, before], equal_nan=True)
+ assert not np.array_equal(rows[0]["samples"][:, ~before], rows[1]["samples"][:, ~before], equal_nan=True)
+
+
+def test_v1_series_are_unchanged_when_the_per_gpu_channel_is_on():
+ outage, S = v2_and_settings()
+ S_on = gpu_settings(S)
+ S_off = replace(S_on, per_gpu_channel=False)
+ job = make_job(n=16, duration=3600.0)
+ traces = []
+ for st in (S_off, S_on):
+ rec = TraceRecorder(st)
+ run_policy(tiny_scenario(outage, 3600, [job]), rec, 0)
+ traces.append(rec.trace())
+ n_v1 = traces[0]["samples"].shape[0]
+ assert traces[1]["samples"].shape[0] > n_v1
+ assert np.array_equal(traces[0]["samples"], traces[1]["samples"][:n_v1], equal_nan=True)
+ assert "gpu_flex_move_gpu_h" not in traces[0] and traces[1]["gpu_nonflex_checkpoint_gpu_h"].shape[1] == traces[1]["samples"].shape[0] - n_v1
+
+
+def creep_scenario(outage, jobs, events=()):
+ """Six simulated hours with the scenario's own 600 s decision interval."""
+ sc = tiny_scenario(outage, 21600, jobs, events)
+ sc["controller_interval_s"] = 600.0
+ return sc
+
+
+@pytest.fixture(scope="module")
+def creep():
+ """Two 8-GPU jobs on rack a1; the lowest-gain GPU of job j1 creeps by 2.4 K over 3600 s, then faults.
+
+ 2.4 K is below the smallest gap between the rack's hottest GPU and this GPU in the same run without
+ the event during the ramp, so the rack maximum never shows the creep; the test checks that.
+ """
+ outage, S = v2_and_settings()
+ S = replace(S, warmup_s=0.0, context_steps=12)
+ jobs = [make_job("j1", n=8, duration=21600.0), make_job("j2", n=8, duration=21600.0)]
+ lay, cfg0 = layout_of(outage, jobs)
+ first_rack_gpus = set(range(int(cfg0.layouts[0].racks[0].n_gpus)))
+ j1 = [g for g in lay.job_slots["j1"] if g in first_rack_gpus]
+ j2 = [g for g in lay.job_slots["j2"] if g in first_rack_gpus]
+ assert len(j1) == 8 and len(j2) == 8
+ slot = min(j1, key=lambda g: lay.k_gpu[g])
+ strength, lead, onset = 2.4, 3600.0, 18000.0
+ ev = HardwareFault("h1", "site-a", "a1", int(slot), onset, 300.0, precursors=(Precursor("gpu_temp_c", lead, strength),))
+ return {"outage": outage, "S": S, "jobs": jobs, "slot": int(slot), "j2_gpus": j2, "strength": strength, "lead": lead,
+ "onset": onset, "event": ev, "lay": lay}
+
+
+class WatchRecorder(TraceRecorder):
+ """Trace recorder that also keeps one GPU's visible temperature and its rack's visible maximum."""
+
+ def __init__(self, S, g, rack_i):
+ super().__init__(S)
+ self.g, self.rack_i, self.seen = g, rack_i, []
+
+ def decide(self, t_s, vs):
+ self.seen.append((t_s, float(vs.now["gpu_temp_c"][self.g]), float(vs.now["max_gpu_temp_c"][self.rack_i])))
+ return super().decide(t_s, vs)
+
+
+def test_residual_sees_a_one_gpu_precursor_that_the_rack_maximum_hides(creep):
+ from gpu_stack.research.heatwave_outage import CH_GPU_RESID, CH_GPU_TEMP, SeriesIndex, gpu_unit_id
+
+ outage, S, onset, lead, strength = creep["outage"], creep["S"], creep["onset"], creep["lead"], creep["strength"]
+ traces, seen = {}, {}
+ for name, ev in (("with", [creep["event"]]), ("without", [])):
+ rec = WatchRecorder(S, creep["slot"], 0)
+ _, cfg = run_policy(creep_scenario(outage, creep["jobs"], ev), rec, 0)
+ traces[name], seen[name] = rec.trace(), np.array(rec.seen)
+ idx = SeriesIndex.from_config(cfg, True)
+ t = traces["with"]["times_s"]
+ k_res = idx.index(gpu_unit_id("a1", creep["slot"]), CH_GPU_RESID)
+ k_max = idx.index("a1", CH_GPU_TEMP)
+ pre = (t >= onset - lead) & (t < onset)
+ late_ramp = (t >= onset - 2 * S.decision_interval_s) & (t < onset)
+ # the creep is visible in the faulty GPU's own sensor ...
+ own = seen["with"][late_ramp, 1] - seen["without"][late_ramp, 1]
+ assert own.max() > 0.8 * strength
+ # ... but the rack's visible maximum is another GPU's temperature the whole time, and is unchanged by the creep
+ assert np.all(seen["with"][pre, 2] > seen["with"][pre, 1])
+ assert np.allclose(traces["with"]["samples"][k_max][pre], traces["without"]["samples"][k_max][pre])
+ # the residual follows the ramp (kelvin) and ends above the largest residual of any GPU in the run without the event
+ gk = [k for k, (u, c) in enumerate(idx.keys) if c == CH_GPU_RESID]
+ noise = np.nanmax(np.abs(traces["without"]["samples"][gk]))
+ ramp = traces["with"]["samples"][k_res]
+ assert np.nanmax(ramp[late_ramp]) > 0.5 * strength
+ assert np.nanmax(ramp[late_ramp]) > noise
+ # no other GPU of the cluster moves by more than the noise level of the run without the event (+ a small allowance)
+ others = [k for k in gk if k != k_res]
+ assert np.nanmax(np.abs(traces["with"]["samples"][others][:, t < onset])) < noise + 0.25
+
+
+def test_gpu_warning_protects_only_the_job_that_holds_the_gpu(creep):
+ from gpu_stack.research.heatwave_outage import Calibration, gpu_unit_id
+
+ outage, S, onset = creep["outage"], creep["S"], creep["onset"]
+ cal = Calibration({"supply_anom": 1.0e9, "gpu_anom": 1.0e9, "gpu_resid": 1.5}, 1.0, {}, "test")
+ ctrl = ForecastWarningController(S, PersistenceForecaster(), 0.5, cal, "persistence")
+ sc = creep_scenario(outage, creep["jobs"], [creep["event"]])
+ res, cfg = run_policy(sc, ctrl, 0)
+ unit = gpu_unit_id("a1", creep["slot"])
+ ex = ctrl.export()
+ gpu_eps = [e for e in ex["episodes"] if e["channel"] == "gpu_resid"]
+ assert gpu_eps and {e["unit"] for e in gpu_eps} == {unit} and {e["class"] for e in gpu_eps} == {"hardware"}
+ assert gpu_eps[0]["start_s"] < onset
+ applied = [a for a in res.actions if a["status"] == "applied"]
+ assert applied and {a["detail"]["job_id"] for a in applied} == {"j1"}
+ assert all(m["unit"] == unit for m in ex["meta_actions"])
+ table = build_event_table(cfg, res.events, S)
+ score = score_run(res, cfg, table, ex, S, 0.0)
+ row = [e for e in score["by_event"] if e["kind"] == "hardware_fault"][0]
+ assert row["caught"] and row["first_channel"] == "gpu_resid" and row["lead_s"] > 0.0
+ assert score["response"]["n_false_alarms"] == 0
+ alarm = [a for a in score["alarms"] if a["unit"] == unit][0]
+ assert alarm["true"] and alarm["channels"] == ["gpu_resid"] and alarm["cost_gpu_h"] > 0.0
+
+
+def test_a_warning_on_another_gpu_of_the_rack_does_not_count_as_catching_the_fault(faulted):
+ from gpu_stack.research.heatwave_outage import classify_alarms, match_events, merge_alarms
+
+ S, table = faulted["S"], faulted["table"]
+ ev = [e for e in table if e["kind"] == "hardware_fault"][0]
+ assert "a1/2" in ev["units"] and "a1" in ev["units"]
+
+ def episode(unit):
+ return {"unit": unit, "class": "hardware", "channel": "gpu_resid", "start_s": ev["onset_s"] - 1200.0,
+ "end_s": ev["onset_s"] - 120.0, "peak_p": 1.0, "n_decisions": 2, "k": 0}
+
+ right, wrong = [episode("a1/2")], [episode("a1/5")]
+ assert match_events(table, right, S)[ev["event_id"]]["caught"]
+ assert not match_events(table, wrong, S)[ev["event_id"]]["caught"]
+ cls = classify_alarms(merge_alarms(wrong, S.episode_merge_gap_s), table, S)
+ assert cls and not cls[0]["true"]
+ assert classify_alarms(merge_alarms(right, S.episode_merge_gap_s), table, S)[0]["true"]
+
+
+def test_replay_objective_prices_a_gpu_alarm_by_the_job_that_holds_the_gpu(creep):
+ from gpu_stack.research.heatwave_outage import SeriesIndex, gpu_unit_id, replay_objective
+
+ outage, S = creep["outage"], creep["S"]
+ rec = TraceRecorder(S)
+ res, cfg = run_policy(creep_scenario(outage, creep["jobs"]), rec, 0)
+ trace = rec.trace()
+ idx = SeriesIndex.from_config(cfg, True)
+ unit = gpu_unit_id("a1", creep["slot"])
+ ep = {"unit": unit, "class": "hardware", "channel": "gpu_resid", "start_s": 7200.0, "end_s": 7200.0, "peak_p": 1.0, "n_decisions": 1, "k": 0}
+ r = replay_objective([], [], [ep], trace, idx, S, 60.0)
+ expected = job_costs(cfg.jobs[0], cfg.transfer)["checkpoint"] # j1 is not flexible: one checkpoint of its 8 GPUs
+ assert r["alarm_gpu_h"] == pytest.approx(expected)
+ assert expected == pytest.approx(8 * cfg.jobs[0].checkpoint_duration_s / 3600.0)
+
+
+# ---------------------------------------------------------------------------
+# outage v3: per-class response design, same-decision checkpoint-and-act, accounting
+# ---------------------------------------------------------------------------
+
+SCENARIO_V3 = ROOT / "experiments" / "e003-heatwave" / "outage-scenario-v3.json"
+
+
+def v3_and_settings():
+ outage = load_outage_scenario(SCENARIO_V3)
+ return outage, OutageSettings.from_scenario(outage)
+
+
+def v3_fast(design=None, **changes):
+ """v3 settings with the control step as the decision interval; ``design`` replaces the scenario's response."""
+ from gpu_stack.research.heatwave_outage import ResponseDesign # noqa: F401
+
+ outage, S = v3_and_settings()
+ S = fast_settings(S)
+ S = replace(S, response=design if design is not None else S.response, response_tick_s=None, **changes)
+ return outage, S
+
+
+def pf_run(outage, S, jobs, events, ref_events=None):
+ """Perfect-foresight run of a tiny scenario with the response in ``S``; returns (result, config, controller)."""
+ sc = tiny_scenario(outage, 7200, jobs, events)
+ base, cfg0 = run_policy(sc, None, 0)
+ So = oracle_settings(S, cfg0.control_dt_s)
+ alarms = oracle_alarms_from_reference(cfg0, base.events, [], So)
+ ctrl = OracleController(So, alarms)
+ res, cfg = run_policy(oracle_scenario(sc), ctrl, 0)
+ return res, cfg, ctrl
+
+
+def actions_of(res, name, status="applied"):
+ return [a for a in res.actions if a["action"] == name and a["status"] == status]
+
+
+def test_scenario_v3_is_simulated_and_has_a_consistent_response_and_disjoint_seeds():
+ outage, S = v3_and_settings()
+ assert outage["data_kind"] == "simulated" and outage["sim"]["data_kind"] == "simulated"
+ assert S.response is not None and S.response.name
+ assert S.tick_s == S.response_tick_s == outage["sim"]["controller_interval_s"]
+ assert S.tick_s <= S.decision_interval_s and abs(S.decision_interval_s / S.tick_s - round(S.decision_interval_s / S.tick_s)) < 1e-9
+ seeds = outage["seeds"]
+ groups = [set(seeds["train"]), set(seeds["design_extra"]), set(seeds["test"])]
+ assert all(groups) and sum(len(g) for g in groups) == len(set().union(*groups))
+ v2, _ = v2_and_settings()
+ assert seeds["test"] == v2["seeds"]["test"] and seeds["train"] == v2["seeds"]["train"]
+ assert outage["event_generation"] == v2["event_generation"]
+
+
+def test_response_design_round_trip_changes_and_validation():
+ from gpu_stack.research.heatwave_outage import ClassRule, ResponseDesign
+
+ _, S = v3_and_settings()
+ d = S.response
+ assert ResponseDesign.from_dict(d.to_dict()) == d
+ e = d.with_changes(rule_hardware=("move", "pause"), protect_ratio=2.0, clear_decisions=3)
+ assert e.rule("hardware") == ClassRule("move", "pause") and e.protect_ratio == 2.0 and e.clear_decisions == 3
+ assert d.rule("hardware") != e.rule("hardware") # frozen: the original is unchanged
+ assert replace(S, response=e).clear_decisions_effective == 3
+ assert replace(S, response=d.with_changes(clear_decisions=None)).clear_decisions_effective == S.clear_decisions
+ with pytest.raises(ValueError):
+ ClassRule("explode", "none")
+ with pytest.raises(ValueError):
+ ClassRule("move", "move")
+ with pytest.raises(ValueError):
+ ResponseDesign("x", {"no-such-class": ClassRule()})
+ assert ResponseDesign("x", {}).rule("thermal").mode == "none"
+
+
+def test_checkpoint_and_pause_are_one_chain_with_no_unprotected_window():
+ from gpu_stack.research.heatwave_outage import ClassRule
+
+ outage, S0 = v3_and_settings()
+ design = S0.response.with_changes(rule_hardware=("pause", "checkpoint"), min_gap_s=0.0, protect_ratio=0.0)
+ outage, S = v3_fast(design)
+ job = make_job("f", n=16, duration=7200.0, flexible=True, ckpt_every=3000.0)
+ fault = HardwareFault("h1", "site-a", "a1", 2, 2900.0, 600.0)
+ res, cfg, ctrl = pf_run(outage, S, [job], [fault])
+ dt = cfg.control_dt_s
+ ck, pa = actions_of(res, "CheckpointNow"), actions_of(res, "PauseJob")
+ assert ck and pa
+ t_ck, t_pa = ck[0]["t_s"], pa[0]["t_s"]
+ assert t_ck < 2900.0 and 0 < t_pa - t_ck <= 2 * dt + 1e-9, "the pause follows the checkpoint within two control steps"
+ # the job is paused before the fault and the fault takes back at most what was run since the checkpoint
+ assert t_pa < 2900.0
+ ref, _ = run_policy(tiny_scenario(outage, 7200, [job], [fault]), None, 0)
+ assert res.summary()["lost_gpu_hours"] < 0.1 * ref.summary()["lost_gpu_hours"]
+ assert actions_of(res, "ResumeJob"), "the job is resumed once the warning ends"
+ assert ClassRule("pause", "checkpoint").mode == "pause"
+
+
+def test_checkpoint_only_design_never_pauses_or_moves_and_keeps_jobs_running():
+ outage, S0 = v3_and_settings()
+ design = S0.response.with_changes(**{f"rule_{c}": ("checkpoint", "checkpoint") for c in ("hardware", "power", "cooling", "thermal")})
+ outage, S = v3_fast(design)
+ job = make_job("f", n=16, duration=7200.0, flexible=True, ckpt_every=3000.0)
+ res, cfg, _ = pf_run(outage, S, [job], [HardwareFault("h1", "site-a", "a1", 2, 2900.0, 600.0)])
+ assert actions_of(res, "CheckpointNow")
+ assert not any(a["action"] in ("PauseJob", "MigrateJob", "ResumeJob") for a in res.actions)
+ assert all(not r["status"] == "paused" for r in res.jobs_table)
+
+
+def test_rule_none_takes_no_action():
+ outage, S0 = v3_and_settings()
+ design = S0.response.with_changes(rule_hardware=("none", "none"))
+ outage, S = v3_fast(design)
+ job = make_job("f", n=16, duration=7200.0, flexible=True, ckpt_every=3000.0)
+ res, _, _ = pf_run(outage, S, [job], [HardwareFault("h1", "site-a", "a1", 2, 2900.0, 600.0)])
+ assert not [a for a in res.actions if a["action"] in ("CheckpointNow", "PauseJob", "MigrateJob")]
+
+
+def test_move_goes_to_the_other_site_when_it_has_room_and_falls_back_when_it_is_full():
+ outage, S0 = v3_and_settings()
+ design = S0.response.with_changes(rule_hardware=("move", "checkpoint"), min_gap_s=0.0)
+ outage, S = v3_fast(design)
+ job = make_job("f", n=8, duration=7200.0, flexible=True, ckpt_every=3000.0)
+ fault = HardwareFault("h1", "site-a", "a1", 2, 2900.0, 600.0)
+ res, cfg, _ = pf_run(outage, S, [job], [fault])
+ moves = actions_of(res, "MigrateJob")
+ assert moves and moves[0]["detail"]["job_id"] == "f" and moves[0]["detail"].get("dest_site", "site-b") == "site-b"
+ n_b = sum(int(r.n_gpus) for r in cfg.layouts[1].racks)
+ filler = make_job("fill", n=n_b, duration=7200.0, site_id="site-b", flexible=False)
+ res2, _, _ = pf_run(outage, S, [job, filler], [fault])
+ assert not actions_of(res2, "MigrateJob")
+ assert [a for a in actions_of(res2, "CheckpointNow") if a["detail"]["job_id"] == "f"], "the fallback checkpoint is written"
+
+
+def test_protect_ratio_and_min_gap_gate_the_policy_checkpoint():
+ from types import SimpleNamespace
+
+ from gpu_stack.research.heatwave_outage import make_planner
+
+ outage, S0 = v3_and_settings()
+ job = make_job("j", n=8, duration=7200.0, ckpt_dur=30.0, step=4.0)
+ cfg = build_config(tiny_scenario(outage, 3600, [job]), seed=0)
+
+ def planner(**ch):
+ return make_planner(cfg, replace(S0, response=S0.response.with_changes(**ch)))
+
+ jv = lambda n: SimpleNamespace(status="running", steps_since_durable=n) # noqa: E731
+ assert planner(protect_ratio=0.0, min_gap_s=0.0)._checkpoint_due(100.0, job, jv(1))
+ assert not planner(protect_ratio=0.0, min_gap_s=0.0)._checkpoint_due(100.0, job, jv(0)), "nothing to protect"
+ p = planner(protect_ratio=2.0, min_gap_s=0.0) # needs 60 s of exposure = 15 steps of 4 s
+ assert not p._checkpoint_due(100.0, job, jv(14)) and p._checkpoint_due(100.0, job, jv(15))
+ assert p._checkpoint_due(100.0, job, jv(1), force=True)
+ g = planner(protect_ratio=0.0, min_gap_s=600.0)
+ assert g._checkpoint_due(100.0, job, jv(5))
+ g.last_checkpoint["j"] = 100.0
+ assert not g._checkpoint_due(600.0, job, jv(5)) and g._checkpoint_due(700.0, job, jv(5))
+ assert not g._checkpoint_due(700.0, job, SimpleNamespace(status="paused", steps_since_durable=5))
+
+
+def test_a_forecast_policy_response_does_not_depend_on_later_events():
+ """Moving a later event leaves every earlier action unchanged, while the perfect-foresight run does react to it."""
+ _, S0 = v3_and_settings()
+ outage, S = v3_fast(S0.response.with_changes(rule_curtail=("checkpoint", "checkpoint"), min_gap_s=0.0))
+ flex = make_job("f", n=16, duration=7200.0, flexible=True, ckpt_every=3000.0)
+ curtail = GridCurtailment("g1", "site-a", 1800.0, 900.0, 0.2, announce_lead_s=1500.0)
+ cut = 3500.0
+
+ def early_actions(fault_s, oracle):
+ sc = tiny_scenario(outage, 7200, [flex], [curtail, HardwareFault("h1", "site-a", "a1", 2, fault_s, 300.0)])
+ if oracle:
+ res, _, _ = pf_run(outage, S, [flex], [curtail, HardwareFault("h1", "site-a", "a1", 2, fault_s, 300.0)])
+ else:
+ ctrl = ForecastWarningController(S, PersistenceForecaster(), 0.5, None, "persistence")
+ res, _ = run_policy(sc, ctrl, 0)
+ return [(a["t_s"], a["action"], a["status"], str(a["detail"])) for a in res.actions if a["t_s"] < cut]
+
+ a, b = early_actions(3000.0, False), early_actions(3600.0, False)
+ assert a == b and a, "the warned run acts on the curtailment only, identically in both"
+ assert early_actions(3000.0, True) != early_actions(3600.0, True), "the oracle does read the event time (the test can see a leak)"
+
+
+def test_pause_records_and_deficit_decomposition_identity():
+ from gpu_stack.research.heatwave_outage import decompose_deficit, pause_records, response_costs
+
+ outage, S0 = v3_and_settings()
+ design = S0.response.with_changes(rule_hardware=("pause", "checkpoint"), min_gap_s=0.0)
+ outage, S = v3_fast(design)
+ job = make_job("f", n=16, duration=7200.0, flexible=True, ckpt_every=3000.0)
+ fault = HardwareFault("h1", "site-a", "a1", 2, 2900.0, 600.0)
+ sc = tiny_scenario(outage, 7200, [job], [fault])
+ base, cfg0 = run_policy(sc, None, 0)
+ table = build_event_table(cfg0, base.events, S)
+ So = oracle_settings(S, cfg0.control_dt_s)
+ ctrl = OracleController(So, oracle_alarms_from_reference(cfg0, base.events, [], So))
+ res, cfg = run_policy(oracle_scenario(sc), ctrl, 0)
+ free, _ = run_policy(tiny_scenario(outage, 7200, [job]), None, 0)
+ score = score_run(res, cfg, table, ctrl.export(), So, float(free.summary()["useful_gpu_hours"]))
+ pauses = score["pauses"]
+ assert pauses and pauses[0]["job_id"] == "f" and pauses[0]["gpu_h"] > 0.0 and pauses[0]["event_id"] == "h1"
+ assert 0.0 <= pauses[0]["after_event_gpu_h"] <= pauses[0]["gpu_h"]
+ dec = decompose_deficit([score], [free._run.ledger.to_gpu_hours()], 50)
+ assert dec["identity_max_abs_error_gpu_h"] < 1e-9 and dec["data_kind"] == "simulated"
+ comp = dec["components_gpu_h"]
+ assert comp["paused_increase"]["mean"] > 0.0
+ assert sum(comp[k]["mean"] for k in comp if k != "residual") + comp["residual"]["mean"] == pytest.approx(dec["deficit_gpu_h"]["mean"])
+ assert dec["by_event_type"]["hardware_fault"]["response_gpu_h_by_action"]["PauseJob"]["mean"] >= 0.0
+ costs = response_costs(res, {j.job_id: j for j in cfg.jobs}, cfg, ctrl.export()["meta_actions"])
+ assert [r["job_id"] for r in pause_records(costs, table, So)] == [p["job_id"] for p in pauses]
diff --git a/tests/test_research_heatwave_packing.py b/tests/test_research_heatwave_packing.py
new file mode 100644
index 0000000..1c484cb
--- /dev/null
+++ b/tests/test_research_heatwave_packing.py
@@ -0,0 +1,315 @@
+"""Tests for the e003 power-cap packing module (all simulated, short horizons)."""
+
+from __future__ import annotations
+
+import copy
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from gpu_stack.research import heatwave_packing as hp
+from gpu_stack.research.heatwave_sim import load_scenario
+
+ROOT = Path(__file__).resolve().parents[1]
+E003 = ROOT / "experiments" / "e003-heatwave"
+
+
+class StubForecaster:
+ """Persistence of the last ratio with a symmetric spread that grows with the level distance from the median."""
+
+ name = "stub"
+
+ def __init__(self, spread: float = 0.1) -> None:
+ self.spread = spread
+
+ def forecast(self, context, horizon, known_covariates=None, past_covariates=None):
+ last = np.asarray(context)[:, 0, -1]
+ levels = np.asarray(hp.QUANTILE_LEVELS)
+ vals = last[:, None, None, None] * (1.0 + self.spread * (levels - 0.5))[None, None, None, :]
+ return SimpleNamespace(values=np.broadcast_to(vals, (last.shape[0], 1, horizon, levels.size)).copy())
+
+
+@pytest.fixture(scope="module")
+def tiny():
+ packing = hp.load_packing_scenario(E003 / "packing-scenario-v1.json")
+ packing = copy.deepcopy(packing)
+ packing["episode"].update({"horizon_s": 10800, "warmup_s": 3600, "min_context_bins": 6, "warm_ratio_bins": 6})
+ packing["cap"]["calibration"] = {"idle_hours": 1, "full_hours": 1}
+ packing["jobs"]["generate"].update({"duration_range_s": [1800, 7200], "utilization_goal": 1.5})
+ sim = load_scenario(E003 / "sim-scenario-v3.json")
+ cal = hp.calibrate_sites(sim, packing)
+ return packing, sim, cal
+
+
+def make_controller(tiny, ep, mode="forecast", q=0.5, **kw):
+ packing, sim, cal = tiny
+ settings = hp.PackingSettings.from_scenario(packing)
+ site_ids = list(cal["models"])
+ if mode == "forecast":
+ kw.setdefault("forecaster", StubForecaster())
+ kw.setdefault("policy", hp.FixedQ(q))
+ return hp.PackingController(ep.jobs, cal["models"], site_ids, settings, ep.hour0_local, mode=mode, **kw), settings
+
+
+def test_levels_and_feature_names():
+ assert len(hp.QUANTILE_LEVELS) == hp.N_LEVELS == 21
+ assert hp.QUANTILE_LEVELS[hp.MEDIAN_INDEX] == 0.5
+ assert hp.level_index(0.9) == 18
+ with pytest.raises(ValueError):
+ hp.level_index(0.42)
+ assert len(hp.FEATURE_NAMES) == hp.N_FEATURES
+ assert "log_penalty" in hp.FEATURE_NAMES and "value_per_gpu_hour" in hp.FEATURE_NAMES
+
+
+def test_scenario_file_loads_and_rejects_bad_schema(tmp_path):
+ sc = hp.load_packing_scenario(E003 / "packing-scenario-v1.json")
+ assert sc["data_kind"] == "simulated"
+ bad = dict(sc, schema="nope")
+ p = tmp_path / "bad.json"
+ p.write_text(json.JSONEncoder().encode(bad))
+ with pytest.raises(ValueError):
+ hp.load_packing_scenario(p)
+
+
+def test_calibration_is_affine_and_caps_sit_between_idle_and_full(tiny):
+ _, _, cal = tiny
+ for sid, m in cal["models"].items():
+ assert m.full_w > m.idle_w > 0
+ assert m.idle_w < m.cap_w < m.full_w
+ assert m.beta_w_per_gpu == pytest.approx((m.full_w - m.idle_w) / m.n_gpus)
+ for row in cal["half_load_check"]:
+ assert abs(row["relative_error"]) < 0.01
+
+
+def test_build_episode_is_seeded_and_runs(tiny):
+ packing, sim, cal = tiny
+ a = hp.build_episode(packing, sim, cal["models"], 3)
+ b = hp.build_episode(packing, sim, cal["models"], 3)
+ c = hp.build_episode(packing, sim, cal["models"], 4)
+ assert [j.to_dict() for j in a.jobs] == [j.to_dict() for j in b.jobs]
+ assert [j.to_dict() for j in a.jobs] != [j.to_dict() for j in c.jobs]
+ assert all(j.flexible and j.site_id is None and j.start_s == j.submit_s for j in a.jobs)
+ for item in a.scenario["sites"]:
+ assert item["site"]["power_cap_w"] == pytest.approx(cal["models"][item["site"]["site_id"]].cap_w)
+ assert a.scenario["data_kind"] == "simulated"
+
+
+def test_slack_decreases_with_level_and_with_extra_jobs(tiny):
+ packing, sim, cal = tiny
+ ep = hp.build_episode(packing, sim, cal["models"], 1)
+ ctrl, st = make_controller(tiny, ep)
+ H = st.horizon_bins
+ levels = np.asarray(hp.QUANTILE_LEVELS)
+ paths = np.broadcast_to((1.0 + 0.1 * (levels - 0.5))[None, None, :], (2, H, 21)).copy()
+ bplan = np.zeros((2, H))
+ active = np.ones(H)
+ s_small = ctrl._slack((0,), (8.0,), active, bplan, paths)
+ s_big = ctrl._slack((0,), (64.0,), active, bplan, paths)
+ assert np.all(np.diff(s_small) <= 1e-12)
+ assert np.all(s_big <= s_small + 1e-12)
+ # a job that ends after the first bin only counts in that bin
+ short = np.zeros(H)
+ short[0] = 1.0
+ assert np.all(ctrl._slack((0,), (64.0,), short, bplan, paths) >= s_big - 1e-12)
+ # two sites: slack is the minimum over the sites
+ both = ctrl._slack((0, 1), (32.0, 32.0), active, bplan, paths)
+ one = [ctrl._slack((j,), (32.0,), active, bplan, paths) for j in (0, 1)]
+ assert np.allclose(both, np.minimum(one[0], one[1]))
+
+
+def test_jobs_wait_for_admission_and_higher_q_admits_no_more(tiny):
+ packing, sim, cal = tiny
+ ep = hp.build_episode(packing, sim, cal["models"], 5)
+ out = {}
+ for q in (0.5, 0.99):
+ ctrl, st = make_controller(tiny, ep, q=q)
+ res = hp.run_episode(ep, ctrl)
+ starts = {e["job_id"]: e["t_s"] for e in res.events if e["kind"] == "job_start"}
+ admitted_at = {a["job_id"]: a["t_s"] for a in ctrl.admissions}
+ # no job starts unless the controller admitted it, and never before the admission time
+ assert set(starts) <= set(admitted_at)
+ assert all(starts[j] >= admitted_at[j] - 1e-9 for j in starts)
+ out[q] = (len(ctrl.admissions), hp.episode_metrics(res, ctrl, ep, st))
+ # before the warm-up ends both runs use the same rule
+ assert out[0.5][1]["data_kind"] == "simulated"
+ assert out[0.5][0] >= out[0.99][0]
+ assert out[0.5][1]["gpu_hours"] >= out[0.99][1]["gpu_hours"] - 1e-6
+
+
+def test_warmup_decisions_do_not_depend_on_the_policy(tiny):
+ packing, sim, cal = tiny
+ ep = hp.build_episode(packing, sim, cal["models"], 6)
+ logs = []
+ for q in (0.1, 0.9):
+ ctrl, st = make_controller(tiny, ep, q=q)
+ hp.run_episode(ep, ctrl)
+ logs.append([(a["t_s"], a["job_id"]) for a in ctrl.admissions if a["t_s"] < st.warmup_s])
+ assert logs[0] == logs[1] and len(logs[0]) > 0
+
+
+def test_metrics_window_and_ranges(tiny):
+ packing, sim, cal = tiny
+ ep = hp.build_episode(packing, sim, cal["models"], 7)
+ ctrl, st = make_controller(tiny, ep, q=0.9)
+ res = hp.run_episode(ep, ctrl)
+ m = hp.episode_metrics(res, ctrl, ep, st)
+ assert 0.0 <= m["violation_rate"] <= 1.0
+ assert m["energy_above_cap_kwh"] >= 0.0 and m["seconds_above_cap"] >= 0.0
+ assert m["stranded_kwh"] >= 0.0
+ assert m["window_s"] == [st.warmup_s, st.horizon_s]
+ assert hp.net_value(m, 0.0) == pytest.approx(m["value_usd"])
+ assert hp.net_value(m, 5.0) == pytest.approx(m["value_usd"] - 5.0 * m["energy_above_cap_kwh"])
+
+
+def test_reactive_controller_throttles_when_over_cap(tiny):
+ packing, sim, cal = tiny
+ low = copy.deepcopy(packing)
+ low["cap"]["load_fraction"] = 0.05
+ models = {
+ sid: hp.SiteModel(m.site_id, m.n_gpus, m.idle_w, m.full_w, m.idle_w + 0.05 * (m.full_w - m.idle_w), m.utc_offset_h)
+ for sid, m in cal["models"].items()
+ }
+ ep = hp.build_episode(low, sim, models, 8)
+ ctrl, st = make_controller(tiny, ep, mode="reactive")
+ ctrl.cap = np.array([m.cap_w for m in models.values()])
+ res = hp.run_episode(ep, ctrl)
+ names = [a["action"] for a in res.actions]
+ assert "SetRackPowerCap" in names
+ assert "AdmitJob" in names
+
+
+def test_oracle_follows_the_reference_ratio(tiny):
+ packing, sim, cal = tiny
+ ep = hp.build_episode(packing, sim, cal["models"], 9)
+ ref, st = make_controller(tiny, ep, q=0.9)
+ hp.run_episode(ep, ref)
+ trace = ref.ratio_bins()
+ assert trace.shape[0] == 2 and trace.shape[1] > 0
+ ctrl, _ = make_controller(tiny, ep, mode="oracle", oracle_ratio=trace)
+ res = hp.run_episode(ep, ctrl)
+ m = hp.episode_metrics(res, ctrl, ep, st)
+ assert m["forecast_calls"] == 0
+ assert len(ctrl.admissions) > 0
+
+
+def test_ratio_bins_match_the_definition(tiny):
+ packing, sim, cal = tiny
+ ep = hp.build_episode(packing, sim, cal["models"], 10)
+ ctrl, st = make_controller(tiny, ep, q=0.9)
+ res = hp.run_episode(ep, ctrl)
+ p, b = ctrl.minute_records()
+ bins = ctrl.ratio_bins()
+ k = int(st.bin_s / st.call_s)
+ i = 3
+ expect = p[i * k:(i + 1) * k, 0].sum() / (ctrl.idle[0] + ctrl.beta[0] * b[i * k:(i + 1) * k, 0]).sum()
+ assert bins[0, i] == pytest.approx(expect)
+ # the recorded call-interval means equal the simulator's control-step samples averaged over the interval
+ m0 = res.site_power_w[: int(st.call_s / res.dt), 0].mean()
+ assert p[0, 0] == pytest.approx(m0)
+
+
+def test_hindsight_costs_hand_example():
+ models = [hp.SiteModel("s", 10, 1000.0, 2000.0, 1500.0, 0.0)] # beta = 100 W per GPU
+ M = 70
+ p_min = np.full((M, 1), 1450.0)
+ b_min = np.full((M, 1), 5.0) # idle + beta*b = 1500, s_real = 1450/1500
+ t = np.array([0.0])
+ # a candidate of 10 GPUs that was NOT admitted: admitting adds 10 * 100 * (1450/1500) = 966.67 W
+ sit = {
+ "t": t, "sites": np.array([[0]]), "shares": np.array([[10.0]]), "x": np.zeros((1, hp.N_FEATURES)),
+ "slack": np.array([np.linspace(0.1, -0.1, hp.N_LEVELS)]), "k": np.array([0]), "admit": np.array([False]),
+ "value": np.array([2.0]), "n_gpus": np.array([10.0]), "life_s": np.array([7200.0]),
+ }
+ costs, value = hp.hindsight_costs(sit, p_min, b_min, models, [0.0, 10.0], 60.0, 3600.0)
+ assert value[0] == pytest.approx(10 * 2.0 * 1.0)
+ fits = sit["slack"][0] >= 0
+ # no penalty: admitting is best, so refusing levels cost the value
+ assert np.allclose(costs[0][0][fits], 0.0)
+ assert np.allclose(costs[0][0][~fits], 20.0)
+ # excess energy if admitted: (1450 + 966.67 - 1500) W for one hour = 0.91667 kWh; penalty 10 => 9.1667 USD
+ excess_kwh = (1450.0 + 100 * 10 * 1450.0 / 1500.0 - 1500.0) / 1000.0
+ assert costs[1][0][fits][0] == pytest.approx(0.0 + 0.0) # admit reward 20 - 9.17 > refuse reward 0
+ assert costs[1][0][~fits][0] == pytest.approx(20.0 - 10.0 * excess_kwh)
+ # a very large penalty flips the best action to refusing
+ c3, _ = hp.hindsight_costs(sit, p_min, b_min, models, [100.0], 60.0, 3600.0)
+ assert np.allclose(c3[0][0][~fits], 0.0)
+ assert np.allclose(c3[0][0][fits], 100.0 * excess_kwh - 20.0)
+
+
+def test_cost_sensitive_model_learns_a_threshold_rule():
+ rng = np.random.default_rng(0)
+ n = 3000
+ X = rng.normal(size=(n, hp.N_FEATURES))
+ best = (X[:, 0] > 0).astype(int) * 10 + 5 # level 15 when the first feature is positive, else level 5
+ costs = np.abs(np.arange(hp.N_LEVELS)[None, :] - best[:, None]).astype(float)
+ value = np.ones(n)
+ model = hp.CostSensitiveModel(n_rff=128, bandwidth=1.0, ridge=1.0).fit(X[:2000], costs[:2000], value[:2000])
+ pick = np.argmin(model.predict_costs(X[2000:]), axis=1)
+ assert np.mean(np.abs(pick - best[2000:]) <= 2) > 0.9
+ again = hp.CostSensitiveModel.from_dict(model.to_dict())
+ assert np.allclose(again.predict_costs(X[2000:2010]), model.predict_costs(X[2000:2010]))
+ assert model.regret(X[2000:], costs[2000:], value[2000:]) < 1.0
+
+
+def test_bootstrap_ci():
+ v = [1.0, 2.0, 3.0, 4.0]
+ a = hp.bootstrap_ci(v, n_boot=500, seed=1)
+ b = hp.bootstrap_ci(v, n_boot=500, seed=1)
+ assert a == b
+ assert a["mean"] == pytest.approx(2.5) and a["lo"] <= 2.5 <= a["hi"] and a["n"] == 4
+ one = hp.bootstrap_ci([3.0], n_boot=50)
+ assert one["lo"] == one["hi"] == 3.0
+ with pytest.raises(ValueError):
+ hp.bootstrap_ci([])
+
+
+def test_forecast_refresh_reuses_paths(tiny):
+ packing, sim, cal = tiny
+ ep = hp.build_episode(packing, sim, cal["models"], 11)
+
+ class Counting(StubForecaster):
+ calls = 0
+ horizons = []
+
+ def forecast(self, context, horizon, known_covariates=None, past_covariates=None):
+ Counting.calls += 1
+ Counting.horizons.append(horizon)
+ return super().forecast(context, horizon)
+
+ runs = {}
+ for k in (1, 4):
+ Counting.calls, Counting.horizons = 0, []
+ ctrl, st = make_controller(tiny, ep, q=0.9, forecaster=Counting(), forecast_refresh_bins=k)
+ res = hp.run_episode(ep, ctrl)
+ m = hp.episode_metrics(res, ctrl, ep, st)
+ runs[k] = (Counting.calls, list(Counting.horizons), m)
+ assert runs[4][0] < runs[1][0]
+ assert set(runs[1][1]) == {st.horizon_bins} and set(runs[4][1]) == {st.horizon_bins + 3}
+ assert runs[4][2]["forecast_refresh_bins"] == 4 and runs[1][2]["forecast_refresh_bins"] == 1
+ with pytest.raises(ValueError):
+ make_controller(tiny, ep, q=0.9, forecast_refresh_bins=0)
+
+
+def test_planned_busy_counts_overlapping_bins_and_stretches_walltime(tiny):
+ from dataclasses import replace
+
+ packing, sim, cal = tiny
+ ep = hp.build_episode(packing, sim, cal["models"], 12)
+ ctrl, st = make_controller(tiny, ep)
+ job = ep.jobs[0]
+ jv = SimpleNamespace(job_id=job.job_id, status="running", site_ids=("site-a",), progress_steps=0.0, total_steps=1.0)
+ vs = SimpleNamespace(jobs=(jv,))
+ # remaining walltime of 1.2 bins: the job overlaps bins 0 and 1 under 'overlap', and bin 0 and 1 are
+ # active under the midpoint rule only when the bin midpoint is before the end
+ short = replace(job, duration_s=1.2 * st.bin_s / (1.0 + st.walltime_overhead))
+ ctrl.job_by_id[job.job_id] = short
+ n = short.n_gpus
+ busy = ctrl._planned_busy(0.0, vs)
+ assert busy[0, 0] == n and busy[0, 1] == n and busy[0, 2] == 0.0
+ ctrl.cfg = replace(st, end_rule="midpoint")
+ busy_mid = ctrl._planned_busy(0.0, vs)
+ assert busy_mid[0, 0] == n and busy_mid[0, 1] == 0.0
+ assert st.walltime_overhead == pytest.approx(hp.declared_checkpoint_overhead(packing))
diff --git a/tests/test_research_heatwave_sim.py b/tests/test_research_heatwave_sim.py
new file mode 100644
index 0000000..f8ae6b4
--- /dev/null
+++ b/tests/test_research_heatwave_sim.py
@@ -0,0 +1,715 @@
+"""Tests for the Heatwave cluster simulator and fault events.
+
+Every run here is simulated. Short horizons and small job lists keep the
+suite quick; the full seven-day scenario runs in the smoke script instead.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from gpu_stack.research.heatwave_events import (
+ BreakerTrip,
+ CoolingFailure,
+ EventSchedule,
+ GridCurtailment,
+ HardwareFault,
+ InverseTimeBreaker,
+ Precursor,
+ random_schedule,
+)
+from gpu_stack.research.heatwave_sim import (
+ AdmitJob,
+ CheckpointNow,
+ DeferJob,
+ Job,
+ MigrateJob,
+ PauseJob,
+ ReactiveSiteCapController,
+ ResumeJob,
+ SetRackPowerCap,
+ SiteEnvironment,
+ build_config,
+ generate_jobs,
+ measure_e002_levels,
+ plan_schedule,
+ planned_covariates,
+ run_scenario,
+)
+
+ROOT = Path(__file__).resolve().parents[1]
+SCENARIO_PATH = ROOT / "experiments" / "e003-heatwave" / "sim-scenario-v1.json"
+E001 = ROOT / "experiments" / "e001-beyond-one-datacenter"
+
+
+def base_scenario(horizon_s: int = 3600, jobs=None, events=None, **power_overrides) -> dict:
+ sc = json.loads(SCENARIO_PATH.read_text())
+ sc["horizon_s"] = horizon_s
+ sc["events"] = events or []
+ sc["outputs"] = {"gpu_windows": []}
+ sc["jobs"] = {"list": [j.to_dict() for j in (jobs or [])]}
+ sc["power_levels"].update(power_overrides)
+ return sc
+
+
+def job(job_id="j1", n=8, start=0.0, duration=3600.0, step=4.0, ckpt_every=600.0, ckpt_dur=30.0, **kw) -> Job:
+ values = dict(
+ job_id=job_id, n_gpus=n, submit_s=0.0, start_s=start, duration_s=duration, step_time_s=step,
+ checkpoint_every_s=ckpt_every, checkpoint_duration_s=ckpt_dur, value_per_gpu_hour=2.0, flexible=False,
+ )
+ values.update(kw)
+ return Job(**values)
+
+
+# ---------------------------------------------------------------------------
+# inputs and calibration
+# ---------------------------------------------------------------------------
+
+
+def test_scenario_is_labeled_simulated_and_builds():
+ sc = json.loads(SCENARIO_PATH.read_text())
+ assert sc["data_kind"] == "simulated"
+ cfg = build_config(sc)
+ assert [lay.site.site_id for lay in cfg.layouts] == ["site-a", "site-b"]
+ assert all(len(lay.racks) == 4 for lay in cfg.layouts)
+ assert all(8 <= r.n_gpus <= 32 for lay in cfg.layouts for r in lay.racks)
+ assert cfg.horizon_s == 7 * 86400
+ kinds = {e.kind for e in cfg.events.events}
+ assert kinds == {"cooling_failure", "breaker_trip", "grid_curtailment", "hardware_fault"}
+ truth = cfg.events.ground_truth()
+ for kind in kinds:
+ flags = {row["has_precursor"] for row in truth if row["kind"] == kind}
+ assert flags == {True, False}, kind
+
+
+def test_transfer_parameters_match_e001_files():
+ sc = json.loads(SCENARIO_PATH.read_text())
+ rec = json.loads((E001 / "recovery-scenario-v2.json").read_text())
+ scr = json.loads((E001 / "screening-scenario-v1.json").read_text())
+ tr = sc["transfer"]
+ assert tr["restore_bandwidth_bytes_per_second"] == rec["restore_bandwidth_bytes_per_second"]
+ assert tr["fixed_restart_latency_s"] == pytest.approx(rec["fixed_restart_latency_ns"] / 1e9)
+ assert tr["membership_reconfiguration_s"] == pytest.approx(rec["membership_reconfiguration_ns"] / 1e9)
+ link = scr["links"][0]
+ assert tr["wan_bandwidth_bytes_per_second"] == link["bandwidth_bytes_per_second"]
+ assert tr["wan_latency_s"] == pytest.approx(link["latency_ns"] / 1e9)
+ west = scr["sites"][0]
+ assert tr["checkpoint_bytes_per_gpu"] == pytest.approx(scr["checkpoint_bytes"] / west["accelerator_count"])
+ for item in sc["sites"]:
+ for key in ("accelerator_flops_per_second", "collective_bandwidth_bytes_per_second",
+ "state_transfer_bandwidth_bytes_per_second", "checkpoint_bandwidth_bytes_per_second",
+ "accelerator_power_w"):
+ assert item["site"][key] == west[key]
+
+
+def test_measured_levels_come_from_the_e002_file():
+ m = measure_e002_levels()
+ assert m["data_kind"] == "measured"
+ r = m["ratio_to_compute"]
+ assert 0.0 < r["idle-outside-active-window"] < 1.0
+ assert r["checkpoint-snapshot"] == pytest.approx(m["median_power_w_laptop"]["checkpoint-snapshot"] / m["median_power_w_laptop"]["canonical-healthy-compute"])
+ assert m["compute_relative_std"] > 0
+ # The telemetry cannot resolve a checkpoint dip: its update period exceeds the window length.
+ assert m["telemetry_update_period_s"] > m["median_window_duration_s"]["checkpoint-snapshot"]
+
+
+def test_levels_are_scaled_to_the_preset_tdp():
+ cfg = build_config(base_scenario())
+ lv = cfg.levels
+ assert lv.compute_w == pytest.approx(700.0 * 0.8)
+ ratio = cfg.measured_levels["ratio_to_compute"]["idle-outside-active-window"]
+ assert lv.idle_w == pytest.approx(lv.compute_w * ratio)
+ assert lv.host_w_per_gpu == pytest.approx(0.25 * lv.compute_w)
+ assert "measured" in lv.provenance["idle_w"]
+
+
+def test_job_validation_and_generator_determinism():
+ with pytest.raises(ValueError):
+ job(n=0)
+ with pytest.raises(ValueError):
+ job(start=-1.0)
+ with pytest.raises(ValueError):
+ Job(**{**job().to_dict(), "multisite": True, "n_gpus": 7})
+ kw = dict(
+ horizon_s=86_400.0, total_gpus=100, utilization_goal=0.6, gpu_choices=[8, 16], gpu_weights=[1, 1],
+ duration_range_s=(600.0, 7200.0), step_time_choices_s=[2.0, 4.0], checkpoint_every_range_s=(300.0, 900.0),
+ checkpoint_duration_range_s=(10.0, 30.0), value_range=(1.0, 2.0), flexible_fraction=0.5,
+ multisite_fraction=0.2, max_start_delay_s=600.0,
+ )
+ a = generate_jobs(5, **kw)
+ b = generate_jobs(5, **kw)
+ c = generate_jobs(6, **kw)
+ assert [j.to_dict() for j in a] == [j.to_dict() for j in b]
+ assert [j.to_dict() for j in a] != [j.to_dict() for j in c]
+ assert sum(j.n_gpus * min(j.duration_s, 86_400.0) for j in a) >= 0.6 * 100 * 86_400.0
+
+
+# ---------------------------------------------------------------------------
+# events
+# ---------------------------------------------------------------------------
+
+
+def test_breaker_curve_matches_iec_conventional_values():
+ b = InverseTimeBreaker()
+ assert b.trip_time_s(1.45) == pytest.approx(3600.0)
+ assert b.trip_time_s(1.13) == float("inf")
+ assert b.trip_time_s(2.0) < b.trip_time_s(1.45)
+ theta = np.array(0.0)
+ t = 0.0
+ while t < 5000:
+ theta, tripped = b.step(theta, 1.45, 10.0)
+ t += 10.0
+ if bool(tripped):
+ break
+ assert t == pytest.approx(3600.0, abs=10.0)
+ cooled, _ = b.step(np.array(1.2), 0.0, 3600.0)
+ assert float(cooled) < 1.2
+
+
+def test_event_schedule_roundtrip_and_validation():
+ events = EventSchedule(
+ (
+ CoolingFailure("c", "site-a", 100.0, 50.0, 0.3, rack_id="a1", precursors=(Precursor("supply_temp_c", 60.0, 2.0),)),
+ BreakerTrip("b", "site-a", "a1", 200.0, 100.0, "derate", 0.7, 30.0),
+ GridCurtailment("g", "site-a", 300.0, 60.0, 0.5, announce_lead_s=120.0),
+ HardwareFault("h", "site-a", "a1", 2, 400.0, 60.0, precursors=(Precursor("gpu_temp_c", 60.0, 5.0),)),
+ )
+ )
+ again = EventSchedule.from_dicts(events.to_dict()["events"])
+ assert again == events
+ with pytest.raises(ValueError):
+ EventSchedule((HardwareFault("x", "s", "r", 0, 1.0, 1.0), HardwareFault("x", "s", "r", 1, 1.0, 1.0)))
+ with pytest.raises(ValueError):
+ BreakerTrip("b", "s", "r", 1.0, 1.0, "forced", precursor_lead_s=5.0)
+ with pytest.raises(ValueError):
+ CoolingFailure("c", "s", 1.0, 1.0, 1.5)
+
+
+def test_precursor_ramp_and_announcement_visibility():
+ p = Precursor("supply_temp_c", 100.0, 4.0)
+ assert float(p.value(0.0, 1000.0)) == 0.0
+ assert float(p.value(950.0, 1000.0)) == pytest.approx(2.0)
+ assert float(p.value(1000.0, 1000.0)) == pytest.approx(4.0)
+ ev = EventSchedule((GridCurtailment("g", "s", 1000.0, 100.0, 0.5, announce_lead_s=300.0),
+ GridCurtailment("u", "s", 2000.0, 100.0, 0.5)))
+ assert ev.announced_curtailments(600.0) == []
+ assert [c["event_id"] for c in ev.announced_curtailments(800.0)] == ["g"]
+ assert ev.announced_curtailments(1200.0) == []
+ assert [c["event_id"] for c in ev.announced_curtailments(2050.0)] == [] # unannounced stays hidden
+
+
+def test_random_schedule_is_deterministic_and_mixes_precursors():
+ kw = dict(
+ horizon_s=7 * 86_400.0, site_racks={"site-a": [("a1", 32), ("a2", 16)], "site-b": [("b1", 24)]},
+ rates_per_day={"cooling_failure": 2.0, "breaker_trip": 2.0, "grid_curtailment": 2.0, "hardware_fault": 2.0},
+ precursor_probability=0.5, lead_range_s=(600.0, 3600.0),
+ strength_ranges={"supply_temp_c": (1.0, 3.0), "gpu_temp_c": (3.0, 8.0)},
+ duration_range_s=(600.0, 3600.0), remaining_capacity_range=(0.2, 0.6), derate_range=(0.6, 0.8),
+ limit_fraction_range=(0.5, 0.8), repair_range_s=(3600.0, 7200.0),
+ )
+ a = random_schedule(3, **kw)
+ b = random_schedule(3, **kw)
+ assert a == b and len(a.events) > 20
+ flags = {row["has_precursor"] for row in a.ground_truth()}
+ assert flags == {True, False}
+
+
+# ---------------------------------------------------------------------------
+# simulator
+# ---------------------------------------------------------------------------
+
+
+def test_hierarchy_sums_add_up_exactly():
+ sc = base_scenario(
+ horizon_s=3600, jobs=[job(n=16), job("j2", n=8, start=300.0)],
+ )
+ sc["outputs"] = {"gpu_windows": [[0, 3600]]}
+ res = run_scenario(sc, seed=1)
+ arr = res.to_arrays()
+ run = res._run
+ gpu = arr["gpu_power_w_window_0"]
+ assert gpu.shape == (run.G, 3600)
+ by_rack = np.add.reduceat(gpu, run.rack_start, axis=0).T
+ assert np.array_equal(by_rack, arr["rack_gpu_power_w"])
+ assert np.allclose(arr["rack_it_power_w"], arr["rack_gpu_power_w"] + np.repeat(run.rack_host_b, run.K, axis=0), rtol=1e-12, atol=0)
+ assert np.allclose(arr["rack_power_w"], arr["rack_it_power_w"] + arr["rack_cooling_power_w"], rtol=1e-12, atol=0)
+ site = np.stack([arr["rack_power_w"][:, run.rack_site == s].sum(axis=1) for s in range(run.S)], axis=1)
+ assert np.allclose(site, arr["site_power_w"], rtol=1e-12, atol=0)
+ hourly = arr["site_power_w"].reshape(1, 3600, run.S).sum(axis=1) / 3.6e6
+ assert np.allclose(hourly, arr["site_energy_kwh_1h"], rtol=1e-12)
+ assert np.allclose(arr["site_energy_kwh_1h"] * run.price_h, arr["site_cost_usd_1h"], rtol=1e-12)
+ assert np.allclose(arr["site_energy_kwh_1h"] * run.carbon_h / 1000.0, arr["site_carbon_kg_1h"], rtol=1e-12)
+ # Metadata labels everything as simulated.
+ meta = arr["meta"]
+ assert meta["data_kind"] == "simulated"
+ assert "measured e002" in meta["quantity_provenance"]["gpu_power_levels"]
+ assert meta["measured_e002"]["data_kind"] == "measured"
+
+
+def test_rack_temperature_output_is_one_minute():
+ res = run_scenario(base_scenario(3600, jobs=[job(n=16)]), seed=0)
+ arr = res.to_arrays()
+ assert arr["rack_temp_c_1min"].shape == (60, 8)
+ assert arr["rack_power_w"].shape == (3600, 8)
+
+
+def test_determinism_for_a_seed():
+ sc = base_scenario(3600, jobs=[job(n=16), job("j2", n=8)])
+ a = run_scenario(sc, seed=4).to_arrays()
+ b = run_scenario(sc, seed=4).to_arrays()
+ c = run_scenario(sc, seed=5).to_arrays()
+ assert np.array_equal(a["site_power_w"], b["site_power_w"])
+ assert np.array_equal(a["rack_temp_c_1min"], b["rack_temp_c_1min"])
+ assert not np.array_equal(a["site_power_w"], c["site_power_w"])
+
+
+def test_no_noise_waveform_levels_and_planned_covariates_agree():
+ sc = base_scenario(
+ 3600, jobs=[job(n=8, duration=1500.0, ckpt_every=300.0, ckpt_dur=40.0)],
+ noise_scale=0.0, comm_ratio=0.5, checkpoint_ratio=0.3,
+ )
+ sc["outputs"] = {"gpu_windows": [[0, 3600]]}
+ res = run_scenario(sc, seed=0)
+ run = res._run
+ lv = run.cfg.levels
+ gpu = res.gpu_windows[0]["gpu_power_w"][:, :1800]
+ job_rows = gpu[:8]
+ assert np.allclose(job_rows, job_rows[0]) # lockstep, no noise
+ cov = res.covariates("site", 1.0, 0.0, 1800.0)
+ # Planned scheduled-checkpoint indicator matches the simulated checkpoint stalls.
+ ckpt_planned = cov["scheduled_checkpoint_gpus"][:, 0] / 8.0
+ p = job_rows[0]
+ ckpt_sim = np.isclose(p, lv.checkpoint_w)
+ assert np.allclose(ckpt_planned, ckpt_sim.astype(float), atol=1e-9)
+ # Planned compute fraction matches the simulated power level while running.
+ running = ~ckpt_sim & (cov["planned_gpus_running"][:, 0] > 0)
+ fc_sim = (p - lv.comm_w) / (lv.compute_w - lv.comm_w)
+ assert np.allclose(cov["step_phase_compute_fraction"][running, 0], fc_sim[running], atol=1e-9)
+ # Idle GPUs sit at idle power.
+ assert np.allclose(gpu[8:, :], lv.idle_w)
+ # Levels: checkpoint dip is visible in the waveform.
+ assert p.min() < lv.comm_w < p.max()
+
+
+def test_event_free_run_matches_the_plan():
+ sc = base_scenario(
+ 6 * 3600,
+ jobs=[
+ job("a", n=64, start=0.0, duration=7200.0), job("b", n=32, start=100.0, duration=5400.0),
+ job("c", n=32, start=200.0, duration=3600.0, multisite=True),
+ job("d", n=16, start=300.0, duration=4000.0, site_id="site-b"),
+ job("e", n=64, start=400.0, duration=5000.0),
+ ],
+ )
+ res = run_scenario(sc, seed=0)
+ starts = {}
+ ends = {}
+ for e in res.events:
+ if e["kind"] == "job_start":
+ starts.setdefault(e["job_id"], e["t_s"])
+ if e["kind"] == "job_done":
+ ends[e["job_id"]] = max(ends.get(e["job_id"], 0.0), e["t_s"])
+ assert not any(e["kind"] == "job_fault" for e in res.events)
+ for pj in res.plan:
+ assert starts[pj.job_id] == pytest.approx(pj.start_s), pj.job_id
+ assert ends[pj.job_id] == pytest.approx(pj.end_s, abs=1e-6), pj.job_id
+ assert all(r["status"] == "done" for r in res.jobs_table)
+ assert all(r["lost_steps"] == 0 for r in res.jobs_table)
+
+
+def test_gpu_fault_loses_work_since_last_checkpoint():
+ j = job(n=8, duration=7200.0, step=4.0, ckpt_every=600.0, ckpt_dur=30.0)
+ ev = [{"kind": "hardware_fault", "event_id": "h", "site_id": "site-a", "rack_id": "a1", "gpu_slot": 3, "onset_s": 1500.0, "repair_s": 7200.0}]
+ res = run_scenario(base_scenario(3 * 3600, jobs=[j], events=ev), seed=0)
+ row = res.jobs_table[0]
+ # Checkpoints complete at wall 630 and 1260; the fault is at 1500, so 240 s of running time is lost.
+ assert row["lost_steps"] == pytest.approx(240.0 / 4.0)
+ assert row["lost_gpu_hours"] == pytest.approx(60.0 * 4.0 * 8 / 3600.0)
+ assert row["faults"] == 1 and row["restarts"] == 1
+ assert row["useful_steps"] <= row["total_steps"] + 1e-6
+
+
+def test_multisite_fault_loses_only_work_since_last_average():
+ j = job("m", n=16, duration=3200.0, step=4.0, multisite=True, flexible=False)
+ # A filler job leaves no spare GPU at site b, so the failed shard waits for the repair.
+ fill = job("fill", n=72, duration=6000.0, ckpt_every=6000.0, site_id="site-b")
+ ev = [{"kind": "hardware_fault", "event_id": "h", "site_id": "site-b", "rack_id": "b4", "gpu_slot": 0, "onset_s": 100.0, "repair_s": 600.0}]
+ res = run_scenario(base_scenario(3 * 3600, jobs=[j, fill], events=ev), seed=0)
+ row = {r["job_id"]: r for r in res.jobs_table}["m"]
+ stall = res._run.cfg.transfer.merge_stall_s(8)
+ cycle = 8 * 4.0 + stall
+ # Fault block starts at t=100: wall time into the cycle is 100 mod cycle (shard runs 32 s then stalls).
+ into = 100.0 % cycle
+ expected_lost = min(into, 32.0) / 4.0
+ assert row["lost_steps"] == pytest.approx(expected_lost, abs=0.6)
+ assert row["lost_steps"] <= 8.0
+ assert row["skipped_steps"] > 0
+ assert row["skipped_steps"] % 8 == pytest.approx(0.0, abs=1e-6) or row["skipped_steps"] % 8 == pytest.approx(8.0, abs=1e-6)
+ assert row["status"] == "done"
+ assert row["faults"] == 1
+
+
+def test_breaker_trip_forced_and_inverse_time():
+ forced = [{"kind": "breaker_trip", "event_id": "f", "site_id": "site-a", "rack_id": "a1", "onset_s": 600.0, "duration_s": 900.0, "mode": "forced"}]
+ res = run_scenario(base_scenario(3600, jobs=[job(n=32, duration=3000.0)], events=forced), seed=0)
+ trips = [e for e in res.events if e["kind"] == "breaker_trip"]
+ assert len(trips) == 1 and trips[0]["t_s"] == 600.0 and trips[0]["reason"] == "forced"
+ p = res.rack_it_power_w
+ k = res._run.K
+ rack_a1 = p[:, 0]
+ assert rack_a1[620 * 1] == 0.0 and rack_a1[1400] == 0.0
+ assert rack_a1[1600] > 0.0 # back on after the 900 s off window plus rounding
+ assert [e["t_s"] for e in res.events if e["kind"] == "rack_reset"] == [1500.0]
+ assert k == 10
+
+ derate = [{"kind": "breaker_trip", "event_id": "d", "site_id": "site-a", "rack_id": "a1", "onset_s": 300.0, "duration_s": 20000.0, "mode": "derate", "derate_fraction": 0.6}]
+ res = run_scenario(base_scenario(3 * 3600, jobs=[job(n=32, duration=10_000.0)], events=derate), seed=0)
+ trips = [e for e in res.events if e["kind"] == "breaker_trip"]
+ assert trips and trips[0]["reason"] == "overload"
+ run = res._run
+ brk = run.cfg.breaker
+ b0 = int(300 / run.cdt)
+ b1 = int(round(trips[0]["t_s"] / run.cdt))
+ ratios = run.rack_ratio_b[b0:b1, 0]
+ assert ratios.min() > brk.nontrip_ratio
+ waited = trips[0]["t_s"] - 300.0
+ # The element opens no earlier than a cold start at the highest ratio and no later than at the lowest.
+ assert waited >= brk.trip_time_s(float(ratios.max())) - run.cdt
+ assert waited <= brk.trip_time_s(float(ratios.min())) + 2 * run.cdt
+ # No trip below the non-tripping ratio: same load, no derate.
+ res0 = run_scenario(base_scenario(3 * 3600, jobs=[job(n=32, duration=10_000.0)]), seed=0)
+ assert not [e for e in res0.events if e["kind"] == "breaker_trip"]
+
+
+def test_cooling_failure_causes_throttling_and_less_power_and_work():
+ ev = [{"kind": "cooling_failure", "event_id": "c", "site_id": "site-a", "rack_id": "a1", "onset_s": 600.0, "duration_s": 3000.0, "remaining_capacity_fraction": 0.25}]
+ jobs = [job(n=32, duration=6000.0, ckpt_every=3000.0)]
+ healthy = run_scenario(base_scenario(7200, jobs=jobs), seed=0)
+ failed = run_scenario(base_scenario(7200, jobs=jobs, events=ev), seed=0)
+ assert not [e for e in healthy.events if e["kind"] == "thermal_throttle"]
+ assert [e for e in failed.events if e["kind"] == "thermal_throttle"]
+ assert failed._run.rack_air_b[:, 0].max() > healthy._run.rack_air_b[:, 0].max() + 10.0
+ # Throttle (or protective shutdown) lowers rack power and delivered work.
+ window = slice(1500, 3500)
+ assert failed.rack_it_power_w[window, 0].mean() < healthy.rack_it_power_w[window, 0].mean()
+ assert failed.summary()["useful_gpu_hours"] < healthy.summary()["useful_gpu_hours"]
+
+
+def test_throttle_threshold_reduces_gpu_power():
+ jobs = [job(n=32, duration=3000.0, ckpt_every=3000.0)]
+ hot = base_scenario(3600, jobs=jobs)
+ hot["throttle"]["throttle_threshold_c"] = 40.0 # below normal GPU temperature
+ a = run_scenario(base_scenario(3600, jobs=jobs), seed=0)
+ b = run_scenario(hot, seed=0)
+ assert b.rack_gpu_power_w[600:, 0].mean() < a.rack_gpu_power_w[600:, 0].mean()
+ assert b.jobs_table[0]["useful_steps"] < a.jobs_table[0]["useful_steps"]
+
+
+def test_cap_and_import_limit_violation_accounting():
+ sc = base_scenario(3600, jobs=[job("a", n=64), job("b", n=32, start=60.0)])
+ sc["sites"][0]["site"]["power_cap_w"] = 40_000.0 # valid: below cooling and import limits
+ ev = [{"kind": "grid_curtailment", "event_id": "g", "site_id": "site-a", "onset_s": 1200.0, "duration_s": 600.0, "limit_fraction": 0.3}]
+ sc["events"] = ev
+ res = run_scenario(sc, seed=0)
+ run = res._run
+ p = res.site_power_w[:, 0]
+ over = np.maximum(p - 40_000.0, 0.0)
+ assert res.site_cap_violation["seconds_above"][0] == (over > 0).sum() * run.dt
+ assert res.site_cap_violation["energy_above_kwh"][0] == pytest.approx(over.sum() * run.dt / 3.6e6)
+ assert res.site_cap_violation["max_excess_w"][0] == pytest.approx(over.max())
+ limit = 110_000.0 * np.where((np.arange(3600) >= 1200) & (np.arange(3600) < 1800), 0.3, 1.0)
+ over_i = np.maximum(p - limit, 0.0)
+ assert res.import_violation["seconds_above"][0] == (over_i > 0).sum()
+ assert res.import_violation["energy_above_kwh"][0] == pytest.approx(over_i.sum() / 3.6e6)
+ assert res.import_violation["seconds_above"][0] > 0
+ # Site b has no violation.
+ assert res.site_cap_violation["seconds_above"][1] == 0
+
+
+def test_policy_hook_can_reduce_cap_violation():
+ sc = base_scenario(3600, jobs=[job("a", n=64), job("b", n=32, start=60.0)])
+ sc["sites"][0]["site"]["power_cap_w"] = 80_000.0
+ plain = run_scenario(sc, seed=0)
+ capped = run_scenario(sc, controller=ReactiveSiteCapController(margin=0.97), seed=0)
+ assert plain.site_cap_violation["seconds_above"][0] > 0
+ assert capped.site_cap_violation["energy_above_kwh"][0] < plain.site_cap_violation["energy_above_kwh"][0]
+ assert any(a["action"] == "SetRackPowerCap" and a["status"] == "applied" for a in capped.actions)
+
+
+class Recorder:
+ """Controller that records the visible state at every decision."""
+
+ def __init__(self, actions=None):
+ self.states = {}
+ self.actions = actions or {}
+
+ def decide(self, t_s, visible_state):
+ self.states[t_s] = visible_state
+ return self.actions.get(t_s, ())
+
+
+def _snapshot(state):
+ return {
+ "now": {k: np.array(v, copy=True) for k, v in state.now.items()},
+ "history": {c: state.history(c, 7200.0) for c in ("rack_air_c", "rack_it_power_w", "site_power_w", "supply_c")},
+ "jobs": state.jobs,
+ "planned": state.known_future["planned_jobs"],
+ "announced": state.known_future["announced_curtailments"],
+ "forecast": state.known_future["outdoor_forecast_c"],
+ }
+
+
+def test_visible_state_has_no_future_leakage():
+ jobs = [job("a", n=32, duration=5000.0), job("b", n=16, start=200.0, duration=3000.0)]
+ future = [
+ {"kind": "hardware_fault", "event_id": "h", "site_id": "site-a", "rack_id": "a1", "gpu_slot": 1, "onset_s": 2500.0, "repair_s": 600.0},
+ {"kind": "grid_curtailment", "event_id": "g", "site_id": "site-b", "onset_s": 2700.0, "duration_s": 300.0, "limit_fraction": 0.4},
+ {"kind": "cooling_failure", "event_id": "c", "site_id": "site-a", "rack_id": "a2", "onset_s": 2800.0, "duration_s": 600.0, "remaining_capacity_fraction": 0.2},
+ ]
+ announced = {"kind": "grid_curtailment", "event_id": "ga", "site_id": "site-a", "onset_s": 3400.0, "duration_s": 300.0, "limit_fraction": 0.5, "announce_lead_s": 600.0}
+ cut = 2400.0
+ ra, rb = Recorder(), Recorder()
+ sc_a = base_scenario(3600, jobs=jobs, events=[announced])
+ sc_b = base_scenario(3600, jobs=jobs, events=future + [announced])
+ # Different realized outdoor temperature after the cut, identical forecast series.
+ cfg = build_config(sc_a)
+ env_a = dict(cfg.environment)
+ env_b = {}
+ for sid, e in env_a.items():
+ actual = e.outdoor_c.copy()
+ actual[1:] += 15.0 # changes the realized series after t = 0 hour... applied below only after the cut hour
+ env_b[sid] = SiteEnvironment(outdoor_c=np.where(np.arange(actual.size) * 3600.0 > cut + 1200.0, actual, e.outdoor_c),
+ price_usd_per_kwh=e.price_usd_per_kwh, carbon_g_per_kwh=e.carbon_g_per_kwh,
+ outdoor_forecast_c=e.forecast_c)
+ run_scenario(sc_a, controller=ra, seed=2)
+ run_scenario(sc_b, controller=rb, seed=2, environment=env_b)
+ times = sorted(t for t in ra.states if t <= cut)
+ assert times and times[-1] == cut
+ for t in times:
+ sa, sb = _snapshot(ra.states[t]), _snapshot(rb.states[t])
+ for key in sa["now"]:
+ assert np.array_equal(sa["now"][key], sb["now"][key], equal_nan=True), (t, key)
+ for key in sa["history"]:
+ assert np.array_equal(sa["history"][key][0], sb["history"][key][0]), (t, key)
+ assert np.array_equal(sa["history"][key][1], sb["history"][key][1], equal_nan=True), (t, key)
+ assert sa["history"][key][0].size == 0 or sa["history"][key][0].max() <= t
+ assert sa["jobs"] == sb["jobs"] and sa["planned"] == sb["planned"] and sa["announced"] == sb["announced"]
+ assert np.array_equal(sa["forecast"], sb["forecast"])
+ # After the future events happen the states do differ, so the test can tell.
+ late = sorted(t for t in ra.states if t > 2900.0)[0]
+ assert not np.array_equal(ra.states[late].now["rack_air_c"], rb.states[late].now["rack_air_c"]) or \
+ not np.array_equal(ra.states[late].now["gpu_temp_c"], rb.states[late].now["gpu_temp_c"], equal_nan=True)
+ # The announced curtailment appears 600 s ahead, not before.
+ before = ra.states[2700.0].known_future["announced_curtailments"]
+ after = ra.states[2880.0].known_future["announced_curtailments"]
+ assert before == () and [c["event_id"] for c in after] == ["ga"]
+ # The unannounced events never appear in the known-future schedule.
+ assert all(c["event_id"] == "ga" for s in ra.states.values() for c in s.known_future["announced_curtailments"])
+ assert all(c["event_id"] == "ga" for s in rb.states.values() for c in s.known_future["announced_curtailments"])
+ with pytest.raises(KeyError):
+ ra.states[cut].history("ground_truth_events", 60.0)
+
+
+def test_controller_actions_defer_pause_checkpoint_and_rejections():
+ jobs = [
+ job("flex", n=16, start=0.0, duration=3000.0, flexible=True),
+ job("rigid", n=8, start=0.0, duration=3000.0, flexible=False),
+ job("later", n=8, start=600.0, duration=1200.0, flexible=True),
+ ]
+ actions = {
+ 0.0: [DeferJob("later", 1200.0), PauseJob("rigid"), CheckpointNow("rigid"), SetRackPowerCap("nope", 1.0)],
+ 120.0: [PauseJob("flex")],
+ 600.0: [ResumeJob("flex")],
+ 660.0: [CheckpointNow("flex"), DeferJob("flex", 9999.0)],
+ }
+ res = run_scenario(base_scenario(3600, jobs=jobs), controller=Recorder(actions), seed=0)
+ log = {(a["t_s"], a["action"], a["detail"].get("job_id")): a["status"] for a in res.actions}
+ assert log[(0.0, "DeferJob", "later")] == "applied"
+ assert log[(0.0, "PauseJob", "rigid")].startswith("rejected")
+ assert log[(0.0, "CheckpointNow", "rigid")].startswith("rejected") # not yet running at t=0
+ assert log[(120.0, "PauseJob", "flex")] == "applied"
+ assert log[(660.0, "DeferJob", "flex")].startswith("rejected")
+ starts = {e["job_id"]: e["t_s"] for e in res.events if e["kind"] == "job_start"}
+ assert starts["later"] == 1200.0
+ rows = {r["job_id"]: r for r in res.jobs_table}
+ # The flexible job lost 480 s of progress to the pause (paused 120 to 600).
+ assert rows["flex"]["useful_steps"] < rows["rigid"]["useful_steps"]
+ assert any(a["status"].startswith("rejected") and a["action"] == "SetRackPowerCap" for a in res.actions)
+
+
+def test_rack_power_cap_limits_power_and_progress():
+ jobs = [job("a", n=32, duration=3000.0, ckpt_every=3000.0)]
+ cap = {0.0: [SetRackPowerCap("a1", 15_000.0)]}
+ free = run_scenario(base_scenario(3600, jobs=jobs), seed=0)
+ capped = run_scenario(base_scenario(3600, jobs=jobs), controller=Recorder(cap), seed=0)
+ assert capped.rack_it_power_w[300:, 0].max() < free.rack_it_power_w[300:, 0].mean()
+ assert capped.jobs_table[0]["useful_steps"] < free.jobs_table[0]["useful_steps"]
+
+
+def test_migration_moves_a_flexible_job_and_costs_transfer_time():
+ j = job("mig", n=16, start=0.0, duration=3000.0, flexible=True, site_id="site-a", ckpt_every=10_000.0)
+ res = run_scenario(
+ base_scenario(2 * 3600, jobs=[j]), controller=Recorder({300.0: [MigrateJob("mig", "site-b")]}), seed=0
+ )
+ row = res.jobs_table[0]
+ assert row["migrations"] == 1 and row["status"] == "done" and row["lost_steps"] == 0
+ tr = res._run.cfg.transfer
+ assert tr.transfer_s(16) == pytest.approx(16 * 4.375e9 / 3.125e9 + 0.02 + 0.012)
+ start = [e for e in res.events if e["kind"] == "migration_start"][0]
+ done = [e for e in res.events if e["kind"] == "job_done"][0]
+ # The job needs 3000 s of running time; the migration stalls it for a checkpoint, the transfer and a restore.
+ stall = j.checkpoint_duration_s + tr.transfer_s(16) + tr.restore_s(16)
+ assert done["t_s"] == pytest.approx(3000.0 + stall + 10_000.0 * 0, abs=3 * 10.0 + 1.0)
+ assert start["to_site_id"] == "site-b"
+ # Rack power moves from site a racks to site b racks.
+ assert res.rack_gpu_power_w[100:200, :4].sum() > res.rack_gpu_power_w[100:200, 4:].sum() * 0 + 1
+ late = slice(1500, 2000)
+ assert res.rack_gpu_power_w[late, 4:].mean(axis=0)[0] > res.rack_gpu_power_w[late, :4].mean(axis=0)[0]
+
+
+def test_admit_action_starts_a_queued_job_early():
+ j = job("early", n=8, start=2400.0, duration=600.0, flexible=True)
+ res = run_scenario(base_scenario(3600, jobs=[j]), controller=Recorder({120.0: [AdmitJob("early")]}), seed=0)
+ start = [e for e in res.events if e["kind"] == "job_start"][0]
+ assert start["t_s"] == 120.0
+
+
+def test_known_future_covariates_shapes_and_alignment():
+ sc = base_scenario(7200, jobs=[job("a", n=40, start=600.0, duration=2000.0)])
+ res = run_scenario(sc, seed=0)
+ cfg = res._run.cfg
+ cov = planned_covariates(cfg, plan_schedule(cfg), 60.0, "rack")
+ assert cov["planned_gpus_running"].shape == (120, 8)
+ assert cov["bin_start_s"][0] == 0.0 and cov["bin_start_s"][1] == 60.0
+ assert cov["planned_gpus_running"][5].sum() == 0
+ assert cov["planned_gpus_running"][30].sum() == pytest.approx(40.0)
+ site = planned_covariates(cfg, plan_schedule(cfg), 3600.0, "site")
+ assert site["planned_gpus_running"].shape == (2, 2)
+ assert site["outdoor_forecast_c"].shape == (2, 2)
+ assert np.all(np.isfinite(site["price_usd_per_kwh"]))
+ arrays = res.to_arrays()
+ assert arrays["cov_rack_1min_planned_gpus_running"].shape == (120, 8)
+ assert arrays["cov_site_1h_carbon_g_per_kwh"].shape == (2, 2)
+
+
+def test_environment_override_flows_into_cost_and_carbon():
+ sc = base_scenario(3600, jobs=[job(n=16)])
+ cfg = build_config(sc)
+ n = int(sc["horizon_s"] / 3600) + 1
+ env = {
+ sid: SiteEnvironment(outdoor_c=np.full(n, 20.0), price_usd_per_kwh=np.full(n, 0.5), carbon_g_per_kwh=np.full(n, 100.0))
+ for sid in cfg.environment
+ }
+ res = run_scenario(sc, seed=0, environment=env)
+ assert np.allclose(res.site_cost_usd_1h, res.site_energy_kwh_1h * 0.5)
+ assert np.allclose(res.site_carbon_kg_1h, res.site_energy_kwh_1h * 0.1)
+
+
+def test_job_override_accepts_an_external_schedule():
+ sc = base_scenario(3600, jobs=[job("x", n=8)])
+ ext = [job("trace-1", n=8, duration=600.0), job("trace-2", n=16, start=100.0, duration=600.0)]
+ res = run_scenario(sc, seed=0, jobs=ext)
+ assert [r["job_id"] for r in res.jobs_table] == ["trace-1", "trace-2"]
+
+
+def test_default_scenario_runs_for_a_few_hours():
+ sc = json.loads(SCENARIO_PATH.read_text())
+ sc["horizon_s"] = 4 * 3600
+ sc["outputs"] = {"gpu_windows": [[0, 600]]}
+ sc["events"] = [e for e in sc["events"] if e["onset_s"] < 4 * 3600 and e["event_id"] != "hw-1-precursor"]
+ res = run_scenario(copy.deepcopy(sc), seed=0)
+ s = res.summary()
+ assert s["data_kind"] == "simulated"
+ assert s["jobs"] > 0 and s["useful_gpu_hours"] > 0
+ assert res.to_arrays()["meta"]["data_kind"] == "simulated"
+
+
+def test_scenario_v2_calibration_notes_label_levels():
+ """v2 levels carry their data labels, and the run metadata says so."""
+ import copy
+ import json as _json
+ from pathlib import Path as _Path
+
+ from gpu_stack.research import heatwave_sim as _hs
+
+ path = _Path(__file__).resolve().parents[1] / "experiments/e003-heatwave/sim-scenario-v2.json"
+ sc = _hs.load_scenario(path)
+ assert sc["calibration"]["sha256"]
+ small = copy.deepcopy(sc)
+ small["horizon_s"] = 3 * 3600
+ small["outputs"] = {"gpu_windows": [[0, 60]]}
+ meta = _hs.run_scenario(small, seed=0).to_arrays()["meta"]
+ prov = meta["power_levels"]["provenance"]
+ assert prov["idle_w"].startswith("public data")
+ assert prov["checkpoint_w"].startswith("measured laptop data")
+ assert "calibration notes applied" in meta["quantity_provenance"]["gpu_power_levels"]
+ bad = copy.deepcopy(small)
+ bad["power_levels"]["calibration_notes"] = {"not_a_level": "x"}
+ import pytest as _pytest
+
+ with _pytest.raises(ValueError):
+ _hs.run_scenario(bad, seed=0)
+ _json.dumps(meta, default=str)
+
+
+def test_hardware_fault_precursor_is_visible_and_jobs_report_racks():
+ """A GPU-temperature precursor shows in visible now/history before onset; JobView lists racks."""
+ jobs = [job("a", n=32, duration=5000.0)]
+ pre = {"kind": "hardware_fault", "event_id": "hp", "site_id": "site-a", "rack_id": "a1", "gpu_slot": 1,
+ "onset_s": 3000.0, "repair_s": 600.0,
+ "precursors": [{"channel": "gpu_temp_c", "lead_s": 1200.0, "strength": 6.0}]}
+ ra, rb = Recorder(), Recorder()
+ run_scenario(base_scenario(3600, jobs=jobs, events=[]), controller=ra, seed=3)
+ run_scenario(base_scenario(3600, jobs=jobs, events=[pre]), controller=rb, seed=3)
+ t_early = max(t for t in ra.states if t <= 1700.0)
+ t_late = max(t for t in ra.states if t <= 2900.0)
+ sa, sb = ra.states[t_early], rb.states[t_early]
+ assert np.array_equal(sa.now["gpu_temp_c"], sb.now["gpu_temp_c"], equal_nan=True) # before the ramp
+ sa, sb = ra.states[t_late], rb.states[t_late]
+ rack = list(sb.rack_ids).index("a1")
+ diff = np.nan_to_num(sb.now["gpu_temp_c"]) - np.nan_to_num(sa.now["gpu_temp_c"])
+ assert diff.max() > 3.0 # ramp is about 6 K * (2900 - 1800) / 1200 by then
+ _, ha = sa.history("max_gpu_temp_c", 600.0)
+ _, hb = sb.history("max_gpu_temp_c", 600.0)
+ assert hb[-1, rack] > ha[-1, rack]
+ views = [v for v in sb.jobs if v.job_id == "a"]
+ assert views and views[0].rack_ids and all(r in sb.rack_ids for r in views[0].rack_ids)
+
+
+def test_per_gpu_temperature_history_shows_single_gpu_precursor():
+ """The per-GPU history channel carries a one-GPU precursor that a rack max can hide."""
+ jobs = [job("a", n=32, duration=5000.0)]
+ pre = {"kind": "hardware_fault", "event_id": "hp", "site_id": "site-a", "rack_id": "a1", "gpu_slot": 1,
+ "onset_s": 3000.0, "repair_s": 600.0,
+ "precursors": [{"channel": "gpu_temp_c", "lead_s": 1200.0, "strength": 6.0}]}
+ ra, rb = Recorder(), Recorder()
+ run_scenario(base_scenario(3600, jobs=jobs, events=[]), controller=ra, seed=3)
+ res_b = run_scenario(base_scenario(3600, jobs=jobs, events=[pre]), controller=rb, seed=3)
+ t_late = max(t for t in ra.states if t <= 2900.0)
+ _, ha = ra.states[t_late].history("gpu_temp_c", 300.0)
+ _, hb = rb.states[t_late].history("gpu_temp_c", 300.0)
+ gidx = rb.states[t_late].now["gpu_rack_index"]
+ rack = list(rb.states[t_late].rack_ids).index("a1")
+ g = int(np.flatnonzero(gidx == rack)[0] + 1)
+ assert hb[-1, g] - ha[-1, g] > 3.0
+ others = np.delete(np.arange(hb.shape[1]), g)
+ assert np.allclose(np.nan_to_num(hb[-1, others]), np.nan_to_num(ha[-1, others]))
+ arr = res_b.to_arrays()
+ assert arr["gpu_temp_c_1min"].shape[1] == gidx.size
+
+
+def test_scenario_v3_schedule_follows_run_seed():
+ from pathlib import Path as _Path
+
+ from gpu_stack.research import heatwave_sim as _hs
+
+ sc = _hs.load_scenario(_Path(__file__).resolve().parents[1] / "experiments/e003-heatwave/sim-scenario-v3.json")
+ a, b = _hs.build_config(sc, seed=100), _hs.build_config(sc, seed=101)
+ assert [j.job_id for j in a.jobs] != [j.job_id for j in b.jobs] or \
+ [j.duration_s for j in a.jobs] != [j.duration_s for j in b.jobs]
+ assert [j.duration_s for j in _hs.build_config(sc, seed=100).jobs] == [j.duration_s for j in a.jobs]
diff --git a/tests/test_research_heatwave_thermal.py b/tests/test_research_heatwave_thermal.py
new file mode 100644
index 0000000..debb070
--- /dev/null
+++ b/tests/test_research_heatwave_thermal.py
@@ -0,0 +1,243 @@
+"""Tests for the Heatwave RC thermal, throttle, and cooling power models."""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import pytest
+
+from gpu_stack.research.heatwave_thermal import (
+ CoolingPowerModel,
+ GPUThrottleModel,
+ RCThermalParams,
+ fit,
+ rc_step,
+ simulate,
+ simulate_full,
+ supply_setpoint,
+)
+
+
+def make_params(**over: float) -> RCThermalParams:
+ values = dict(
+ heat_capacity_j_per_k=2.0e5,
+ resistance_k_per_w=1.5e-3,
+ cooling_lag_s=120.0,
+ cooling_capacity_w=math.inf,
+ supply_setpoint_c=22.0,
+ outdoor_coupling_k_per_k=0.3,
+ outdoor_reference_c=25.0,
+ )
+ values.update(over)
+ return RCThermalParams(**values)
+
+
+def test_params_validation():
+ with pytest.raises(ValueError):
+ make_params(heat_capacity_j_per_k=0.0)
+ with pytest.raises(ValueError):
+ make_params(resistance_k_per_w=-1.0)
+ with pytest.raises(ValueError):
+ make_params(cooling_lag_s=-1.0)
+ with pytest.raises(ValueError):
+ make_params(cooling_capacity_w=-5.0)
+ p = make_params()
+ assert p.time_constant_s == pytest.approx(300.0)
+ assert p.to_dict()["cooling_capacity_w"] is None
+
+
+def test_steady_state_matches_formula():
+ p = make_params(outdoor_coupling_k_per_k=0.0, cooling_lag_s=0.0)
+ n = 4000
+ temps = simulate(p, np.full(n, 10_000.0), np.full(n, 20.0), 10.0, T0=22.0)
+ assert temps[-1] == pytest.approx(22.0 + 1.5e-3 * 10_000.0, abs=1e-6)
+
+
+def test_exact_in_dt_for_piecewise_constant_inputs():
+ p = make_params()
+ rng = np.random.default_rng(3)
+ coarse_p = rng.uniform(5_000, 15_000, size=40)
+ coarse_a = rng.uniform(15, 40, size=40)
+ dt = 60.0
+ coarse = simulate(p, coarse_p, coarse_a, dt, T0=30.0)
+ sub = 120
+ fine = simulate(p, np.repeat(coarse_p, sub), np.repeat(coarse_a, sub), dt / sub, T0=30.0)
+ # Same state at the coarse sample times, up to rounding.
+ assert np.allclose(coarse, fine[::sub], atol=1e-6)
+
+
+def test_lag_equals_time_constant_branch_is_continuous():
+ p_eq = make_params(cooling_lag_s=300.0)
+ p_near = make_params(cooling_lag_s=300.0 * (1 + 1e-6))
+ pw = np.full(200, 9_000.0)
+ am = np.linspace(20, 40, 200)
+ a = simulate(p_eq, pw, am, 30.0, 28.0)
+ b = simulate(p_near, pw, am, 30.0, 28.0)
+ assert np.allclose(a, b, atol=1e-3)
+
+
+def test_scalar_and_vector_steps_agree():
+ p = make_params(cooling_capacity_w=8_000.0)
+ pw = np.concatenate([np.full(50, 6_000.0), np.full(100, 14_000.0), np.full(100, 3_000.0)])
+ am = np.full(pw.size, 30.0)
+ scalar = simulate_full(p, pw, am, 20.0, T0=30.0)
+ T = np.array(30.0)
+ Ts = np.array(float(supply_setpoint(p, 30.0)))
+ out = []
+ for i in range(pw.size):
+ out.append(float(T))
+ T, Ts, _ = rc_step(
+ p.heat_capacity_j_per_k, p.resistance_k_per_w, p.cooling_lag_s,
+ p.cooling_capacity_w, supply_setpoint(p, am[i]), T, Ts, pw[i], 20.0,
+ )
+ assert np.allclose(out, scalar["temp_c"], atol=1e-9)
+
+
+def test_capacity_limit_and_energy_balance():
+ cap = 8_000.0
+ p = make_params(cooling_capacity_w=cap, outdoor_coupling_k_per_k=0.0)
+ pw = np.full(600, 14_000.0)
+ am = np.full(600, 25.0)
+ out = simulate_full(p, pw, am, 10.0, T0=30.0)
+ assert np.all(out["heat_removed_w"] <= cap + 1e-6)
+ # Heat above capacity warms the node at (P - cap) / C once saturated.
+ late = out["temp_c"][-200:]
+ slope = (late[-1] - late[0]) / (199 * 10.0)
+ assert slope == pytest.approx((14_000.0 - cap) / p.heat_capacity_j_per_k, rel=1e-6)
+ # Energy balance: sum of heat removed + stored equals energy in.
+ stored = p.heat_capacity_j_per_k * (out["temp_c"][-1] - out["temp_c"][0])
+ total_in = float(np.sum(pw)) * 10.0
+ # temps has the state at the start of each step, so add the last step.
+ last_step_removed = out["heat_removed_w"][-1] * 10.0
+ last_step_dt = 10.0 * (pw[-1] - out["heat_removed_w"][-1]) / p.heat_capacity_j_per_k
+ stored_total = stored + p.heat_capacity_j_per_k * last_step_dt
+ removed_total = float(np.sum(out["heat_removed_w"])) * 10.0
+ assert removed_total + stored_total == pytest.approx(total_in, rel=1e-9)
+ assert last_step_removed > 0
+
+
+def test_capacity_series_override_cooling_failure_raises_temperature():
+ p = make_params(outdoor_coupling_k_per_k=0.0)
+ n = 720
+ pw = np.full(n, 10_000.0)
+ am = np.full(n, 25.0)
+ cap = np.full(n, 20_000.0)
+ cap[300:] = 2_000.0
+ base = simulate(p, pw, am, 10.0, T0=22.0 + 15.0)
+ failed = simulate(p, pw, am, 10.0, T0=22.0 + 15.0, capacity_w=cap)
+ assert failed[-1] > base[-1] + 10.0
+ assert np.allclose(failed[:300], base[:300])
+
+
+def test_saturation_crossing_is_close_to_fine_integration():
+ p = make_params(cooling_capacity_w=9_000.0, outdoor_coupling_k_per_k=0.0)
+ n = 90
+ pw = np.concatenate([np.full(30, 6_000.0), np.full(60, 12_000.0)])
+ am = np.full(n, 25.0)
+ coarse = simulate(p, pw, am, 60.0, T0=22.0 + 9.0)
+ sub = 60
+ fine = simulate(p, np.repeat(pw, sub), np.repeat(am, sub), 1.0, T0=22.0 + 9.0)
+ assert np.max(np.abs(coarse - fine[::sub])) < 0.05
+
+
+def test_fit_recovers_known_parameters():
+ truth = make_params(cooling_lag_s=90.0)
+ rng = np.random.default_rng(11)
+ n = 3000
+ dt = 20.0
+ # Piecewise-constant power with random levels and a diurnal-like outdoor swing.
+ pw = np.repeat(rng.uniform(4_000, 16_000, size=n // 30), 30)
+ am = 28.0 + 8.0 * np.sin(2 * np.pi * np.arange(n) * dt / 86_400.0 * 6) + 4.0 * np.sin(
+ 2 * np.pi * np.arange(n) * dt / 5_000.0
+ )
+ temps = simulate(truth, pw, am, dt, T0=35.0)
+ temps_noisy = temps + rng.normal(0.0, 0.02, size=n)
+ fitted, report = fit(
+ pw, am, temps_noisy, dt,
+ free=("heat_capacity_j_per_k", "resistance_k_per_w", "cooling_lag_s", "supply_setpoint_c", "outdoor_coupling_k_per_k"),
+ fixed={"outdoor_reference_c": truth.outdoor_reference_c},
+ )
+ assert fitted.heat_capacity_j_per_k == pytest.approx(truth.heat_capacity_j_per_k, rel=0.05)
+ assert fitted.resistance_k_per_w == pytest.approx(truth.resistance_k_per_w, rel=0.05)
+ assert fitted.time_constant_s == pytest.approx(truth.time_constant_s, rel=0.05)
+ assert fitted.cooling_lag_s == pytest.approx(truth.cooling_lag_s, rel=0.05)
+ assert fitted.supply_setpoint_c == pytest.approx(truth.supply_setpoint_c, abs=0.5)
+ assert report["rmse_c"] < 0.1
+ assert report["r2"] > 0.99
+ assert report["time_constant_s"] == pytest.approx(fitted.time_constant_s)
+
+
+def test_fit_recovers_rc_without_lag():
+ truth = make_params(cooling_lag_s=0.0, outdoor_coupling_k_per_k=0.0)
+ rng = np.random.default_rng(5)
+ n = 1500
+ pw = np.repeat(rng.uniform(3_000, 18_000, size=n // 25), 25)
+ am = np.full(n, 20.0)
+ temps = simulate(truth, pw, am, 15.0, T0=25.0)
+ fitted, report = fit(pw, am, temps, 15.0)
+ assert fitted.heat_capacity_j_per_k == pytest.approx(truth.heat_capacity_j_per_k, rel=0.05)
+ assert fitted.resistance_k_per_w == pytest.approx(truth.resistance_k_per_w, rel=0.05)
+ assert report["success"]
+
+
+def test_fit_rejects_bad_inputs():
+ with pytest.raises(ValueError):
+ fit(np.ones(4), np.ones(4), np.ones(4), 1.0)
+ with pytest.raises(ValueError):
+ fit(np.ones(20), np.ones(20), np.ones(20), 1.0, free=("supply_setpoint_c",))
+
+
+def throttle() -> GPUThrottleModel:
+ return GPUThrottleModel(
+ throttle_threshold_c=80.0, throttle_span_c=10.0, min_clock_fraction=0.5, power_exponent=2.0
+ )
+
+
+def test_throttle_model_has_no_defaults():
+ with pytest.raises(TypeError):
+ GPUThrottleModel() # type: ignore[call-arg]
+
+
+def test_throttle_reduces_clock_and_power_monotonically():
+ m = throttle()
+ temps = np.array([60.0, 80.0, 85.0, 90.0, 120.0])
+ clk = m.clock_fraction(temps)
+ assert clk[0] == 1.0 and clk[1] == 1.0
+ assert np.all(np.diff(clk) <= 0)
+ assert clk[-1] == pytest.approx(0.5)
+ pw = m.power_w(560.0, 100.0, clk)
+ assert np.all(np.diff(pw) <= 1e-12)
+ assert pw[-1] < pw[0]
+ # Idle power does not scale.
+ assert m.power_w(100.0, 100.0, 0.5) == pytest.approx(100.0)
+
+
+def test_clock_for_power_inverts_power_w():
+ m = throttle()
+ clk = m.clock_for_power(300.0, 560.0, 100.0)
+ assert float(m.power_w(560.0, 100.0, clk)) == pytest.approx(300.0)
+ assert float(m.clock_for_power(1_000.0, 560.0, 100.0)) == 1.0
+ assert float(m.clock_for_power(50.0, 560.0, 100.0)) == pytest.approx(0.5)
+
+
+def test_cooling_power_pue_mode():
+ m = CoolingPowerModel(kind="pue", pue=1.2)
+ assert float(m.power_w(10_000.0, 30.0, 24.0)) == pytest.approx(2_000.0)
+
+
+def test_carnot_calibration_reproduces_pue_and_rises_with_outdoor_temperature():
+ m = CoolingPowerModel.carnot_from_pue(
+ pue=1.2, reference_outdoor_c=20.0, reference_supply_c=24.0,
+ evaporator_approach_k=10.0, condenser_approach_k=10.0, cop_min=1.5, cop_max=40.0,
+ )
+ assert float(m.power_w(10_000.0, 20.0, 24.0)) == pytest.approx(2_000.0)
+ cool = float(m.power_w(10_000.0, 10.0, 24.0))
+ hot = float(m.power_w(10_000.0, 40.0, 24.0))
+ assert cool < 2_000.0 < hot
+ assert m.to_dict()["kind"] == "carnot"
+ with pytest.raises(ValueError):
+ CoolingPowerModel.carnot_from_pue(
+ pue=1.2, reference_outdoor_c=20.0, reference_supply_c=24.0,
+ evaporator_approach_k=10.0, condenser_approach_k=10.0, cop_min=6.0, cop_max=40.0,
+ )
diff --git a/tests/test_scenario_audit_text_family_index.py b/tests/test_scenario_audit_text_family_index.py
index 307492a..fca6539 100644
--- a/tests/test_scenario_audit_text_family_index.py
+++ b/tests/test_scenario_audit_text_family_index.py
@@ -36,11 +36,11 @@ def test_scenario_audit_missing_families_text_indexes_missing_roots_by_family():
assert "Scenario audit:" in out
assert (
"pythia_70m_dgx_h100_us_2024_industrial_power: issues "
- "targets=1 issues=33 sourced=True"
+ "targets=1 issues=34 sourced=True"
) in out
assert (
"cost_per_token: issues target=econ.cost.per_token "
- "missing=33 violated_constraints=0"
+ "missing=34 violated_constraints=0"
) in out
assert "unresolved inputs:" not in out
assert " missing families:" in out
@@ -50,7 +50,8 @@ def test_scenario_audit_missing_families_text_indexes_missing_roots_by_family():
) in out
assert (
" family=econ.cluster boundary_category=symbolic-boundary "
- "primitive_boundary=False count=1 names=econ.cluster.facility_capex"
+ "primitive_boundary=False count=2 "
+ "names=econ.cluster.facility_capex, econ.cluster.facility_capex_rate"
) in out
assert (
" family=thermal.water boundary_category=primitive-root "