Benchmark directives on all rules + Slurm resource measurement tooling

workflow/collect_benchmarks.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+# Per-rule resource report for tuning Slurm requests.
+#
+#   bash collect_benchmarks.sh [SINCE]        # SINCE defaults to now-6hours
+#
+# sacct splits the data we need across two rows per job: the allocation row
+# carries JobName/ReqMem, the .batch row carries MaxRSS. This joins them on the
+# base job id and aggregates per rule, so `sacct -X` (which shows MaxRSS blank)
+# is never the thing you look at.
+set -uo pipefail
+
+SINCE="${1:-now-6hours}"
+
+sacct -S "$SINCE" -P --format=JobID,JobName,State,Elapsed,MaxRSS,ReqMem 2>/dev/null |
+awk -F'|' '
+  NR == 1 { next }
+  {
+    split($1, a, ".")
+    id = a[1]
+    if ($1 !~ /\./) { name[id] = $2; state[id] = $3; elapsed[id] = $4; req[id] = $6 }
+    else if ($1 ~ /\.batch$/) { gsub(/K$/, "", $5); rss[id] = $5 + 0 }
+  }
+  END {
+    for (id in name) {
+      n = name[id]
+      if (n == "batch" || n == "extern" || n == "" ) continue
+      if (n ~ /^(bash|uv|python|snakemake|ca2040_probe)$/) continue
+      cnt[n]++
+      sum_rss[n] += rss[id]
+      if (rss[id] > max_rss[n]) max_rss[n] = rss[id]
+      split(elapsed[id], t, ":")
+      secs = t[1] * 3600 + t[2] * 60 + t[3]
+      if (secs > max_secs[n]) max_secs[n] = secs
+      reqm[n] = req[id]
+      if (state[id] !~ /COMPLETED/) bad[n] = bad[n] " " state[id]
+    }
+    printf "%-34s %5s %11s %11s %10s %8s  %s\n", "RULE", "N", "MAXRSS_MB", "PEAK_ELAPSED", "REQ_MEM", "RATIO", "NONCOMPLETE"
+    for (n in cnt) {
+      mb = max_rss[n] / 1024
+      r = reqm[n]; gsub(/[MGn]$/, "", r)
+      if (reqm[n] ~ /G/) r = r * 1024
+      ratio = (mb > 0 && r > 0) ? sprintf("%.1fx", r / mb) : "-"
+      printf "%-34s %5d %11.0f %11s %10s %8s  %s\n", \
+        n, cnt[n], mb, \
+        sprintf("%d:%02d:%02d", max_secs[n]/3600, (max_secs[n]%3600)/60, max_secs[n]%60), \
+        reqm[n], ratio, bad[n]
+    }
+  }
+' | { read -r hdr; echo "$hdr"; sort -k3 -nr; }

workflow/probe_2019.sbatch

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Probe run: weather year 2019, both resolutions, purely to collect per-rule
+# resource data before launching the remaining 45 overlays.
+#
+#   sbatch probe_2019.sbatch
+#
+# This is the DRIVER job. It holds two `snakemake --cluster` processes, which
+# submit each rule as its own Slurm job. Its own CPU/memory needs are tiny; its
+# walltime must cover the whole build span including queue wait for rule jobs.
+#SBATCH -p serc
+#SBATCH -J ca2040_probe
+#SBATCH --time=06:00:00
+#SBATCH --cpus-per-task=2
+#SBATCH --mem=8G
+#SBATCH -o logs/drivers/probe-%j.out
+#SBATCH -e logs/drivers/probe-%j.err
+#SBATCH --mail-type=FAIL
+#SBATCH --mail-user=ctehran@stanford.edu
+
+set -uo pipefail
+cd /oak/stanford/groups/iazevedo/kamran/CH3/pypsa-usa/workflow
+
+export PARTITION=serc
+export EMAIL=ctehran@stanford.edu
+export JOBS=20          # concurrent Slurm rule-jobs per overlay
+export CONCURRENCY=2    # both overlays at once
+export RESTART_TIMES=0  # a retry would re-run with attempt*mem and skew MaxRSS
+
+echo "=== probe start $(date -Is) on $(hostname) ==="
+bash run_slurm.sh ca2040_wy2019_z4.yaml ca2040_wy2019_county.yaml
+rc=$?
+echo "=== probe end $(date -Is) rc=$rc ==="
+exit $rc

workflow/repo_data/config/config.california.yaml

@@ -69,6 +69,17 @@ clustering:
   aggregation_strategies:
     generators:
       p_min_pu: 'capacity_weighted_average'   # keep the UC feasibility invariant through clustering
+      # CPUC out-of-state units (REMOTE_PREFIX "R ") are attached to CA buses with
+      # p_nom_extendable=False, while in-state units of the same carrier are extendable
+      # because their carrier is in extendable_carriers.Generator. Clustering groups
+      # one-ports by (bus, carrier), so those two land in one group and PyPSA's default
+      # `consense` raises. Same situation `committable: any` already handles in the base.
+      p_nom_extendable: any
+      # Hydro units disagree on up_time_before ({0, 1}). `max` keeps the integer
+      # dtype (capacity_weighted_average would yield a float) and is the
+      # UC-feasibility-safe direction: the cluster inherits the longest online
+      # history, so it is never spuriously forced to stay committed.
+      up_time_before: max
 
 solving:
   mem: 30000

workflow/report_benchmarks.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+"""Turn snakemake benchmark TSVs into per-rule Slurm resource recommendations.
+
+    ../.venv/bin/python report_benchmarks.py [benchmarks_dir]
+
+Snakemake writes one TSV per job under benchmarks/, with columns
+s / h:m:s / max_rss / max_vms / max_uss / max_pss / io_in / io_out /
+mean_load / cpu_time.  Memory is in MB, s in seconds.
+
+The rule name is recovered from the path, which mirrors each rule's `log:`
+layout, so it is stripped of run name, interconnect, and wildcard-bearing
+prefixes/suffixes.
+
+Recommendations use headroom on the observed peak:
+  mem   = max(2 GB, ceil(1.5 x peak max_rss / 500 MB) * 500 MB)
+  time  = max(10 min, ceil(3 x peak runtime / 5 min) * 5 min)
+The multipliers differ because memory overrun is a hard kill while a generous
+walltime only costs queue priority; 3x also absorbs a slower/contended node.
+"""
+
+import re
+import sys
+from pathlib import Path
+
+import pandas as pd
+
+root = Path(sys.argv[1] if len(sys.argv) > 1 else "benchmarks")
+
+# Wildcard debris to strip from the file stem to recover a rule name. The simpl
+# wildcard is a number or a word like "county"/"all", so anchor on the known
+# forms -- a bare `_s\w+$` would eat `_shapes` and `_substations`.
+SIMPL = r"(?:\d+|county|all)"
+STEM_SUBS = [
+    (rf"_?elec_s{SIMPL}(?:_c\d+[a-z]?)?", ""),  # elec_s75_c4 / elec_scounty
+    (rf"_s{SIMPL}$", ""),  # trailing _s75 / _scounty
+    (r"_c?\d{2,}[a-z]?$", ""),  # trailing cluster counts / horizons
+    (r"^elec_", ""),
+]
+
+
+# Benchmark paths are not uniform: most are benchmarks/<run>/<interconnect>/<stem>,
+# but some rules (cluster_network) write benchmarks/<rule>/<interconnect>/<stem>.
+# So resolve against the real rule list rather than inferring from position.
+KNOWN_RULES = set()
+for smk in Path("rules").glob("*.smk"):
+    KNOWN_RULES.update(re.findall(r"^rule\s+(\w+)\s*:", smk.read_text(), re.M))
+
+
+def rule_from(path: Path) -> str:
+    # A rule name appearing as a directory component wins outright.
+    for part in path.parts:
+        if part in KNOWN_RULES:
+            return part
+    stem = path.stem
+    # Longest known rule that prefixes the stem, so build_renewable_profiles_onwind
+    # keeps its tech suffix but build_shapes is not truncated.
+    hits = [r for r in KNOWN_RULES if stem.startswith(r)]
+    if hits:
+        best = max(hits, key=len)
+        return stem if stem.startswith(f"{best}_") and best != stem else best
+    for pat, rep in STEM_SUBS:
+        stem = re.sub(pat, rep, stem)
+    stem = stem.strip("_")
+    return stem or path.parent.name
+
+
+rows = []
+for f in sorted(root.rglob("*")):
+    if not f.is_file():
+        continue
+    try:
+        df = pd.read_csv(f, sep="\t")
+    except Exception:
+        continue
+    if df.empty or "s" not in df.columns:
+        continue
+    parts = f.relative_to(root).parts
+    run = parts[0] if len(parts) > 1 else "-"
+    r = df.iloc[0]
+    rows.append(
+        {
+            "rule": rule_from(f),
+            "run": run,
+            "secs": float(r["s"]),
+            "max_rss": float(r.get("max_rss", 0) or 0),
+            "max_pss": float(r.get("max_pss", 0) or 0),
+            "cpu_time": float(r.get("cpu_time", 0) or 0),
+            "mean_load": float(r.get("mean_load", 0) or 0),
+        },
+    )
+
+if not rows:
+    sys.exit(f"no benchmark files under {root}")
+
+d = pd.DataFrame(rows)
+
+
+def fmt_hms(secs: float) -> str:
+    s = int(round(secs))
+    return f"{s // 3600:d}:{(s % 3600) // 60:02d}:{s % 60:02d}"
+
+
+def rec_mem(peak_mb: float) -> int:
+    import math
+
+    return max(2000, int(math.ceil(1.5 * peak_mb / 500.0) * 500))
+
+
+def rec_time(peak_s: float) -> str:
+    import math
+
+    mins = max(10, int(math.ceil(3 * peak_s / 60.0 / 5.0) * 5))
+    return f"{mins // 60:02d}:{mins % 60:02d}:00"
+
+
+agg = (
+    d.groupby("rule")
+    .agg(
+        n=("secs", "size"),
+        peak_s=("secs", "max"),
+        peak_rss=("max_rss", "max"),
+        mean_load=("mean_load", "max"),
+    )
+    .sort_values("peak_rss", ascending=False)
+)
+
+agg["rec_mem_mb"] = agg.peak_rss.map(rec_mem)
+agg["rec_walltime"] = agg.peak_s.map(rec_time)
+agg["peak_time"] = agg.peak_s.map(fmt_hms)
+
+print(
+    f"{'RULE':<44} {'N':>2} {'PEAK_RSS_MB':>11} {'PEAK_TIME':>10} {'LOAD%':>6}   {'REC_MEM_MB':>10} {'REC_WALLTIME':>12}"
+)
+print("-" * 108)
+for name, r in agg.iterrows():
+    print(
+        f"{name:<44} {int(r.n):>2} {r.peak_rss:>11.0f} {r.peak_time:>10} "
+        f"{r.mean_load:>6.0f}   {int(r.rec_mem_mb):>10} {r.rec_walltime:>12}",
+    )
+
+print()
+print(f"total jobs benchmarked: {len(d)}   distinct rules: {len(agg)}")
+print(f"sum of peak runtimes (serial lower bound): {fmt_hms(agg.peak_s.sum())}")

workflow/rules/build_electricity.smk

@@ -35,6 +35,8 @@ rule build_shapes:
         county_shapes=GEOSPATIAL + "{interconnect}/county_shapes.geojson",
     log:
         "logs/build_shapes/{interconnect}.log",
+    benchmark:
+        BENCHMARKS + "{interconnect}/build_shapes"
     threads: 1
     resources:
         walltime=config_provider("walltime", "build_shapes", default="00:30:00"),
@@ -83,6 +85,8 @@ rule build_base_network:
         network=NETWORKS + "{interconnect}/elec_base_network.nc",
     log:
         "logs/create_network/{interconnect}.log",
+    benchmark:
+        BENCHMARKS + "{interconnect}/build_base_network"
     threads: 1
     resources:
         mem_mb=5000,
@@ -114,6 +118,8 @@ rule build_bus_regions:
         regions_offshore=GEOSPATIAL + "{interconnect}/regions_offshore.geojson",
     log:
         "logs/build_bus_regions/{interconnect}.log",
+    benchmark:
+        BENCHMARKS + "{interconnect}/build_bus_regions"
     threads: 1
     resources:
         mem_mb=3000,
@@ -137,6 +143,8 @@ rule build_cost_data:
         sector_costs=COSTS + "sector_costs_{year}.csv",
     log:
         LOGS + "costs_{year}.log",
+    benchmark:
+        BENCHMARKS + "build_cost_data_{year}"
     threads: 1
     resources:
         mem_mb=5000,
@@ -754,6 +762,8 @@ rule build_powerplants:
         powerplants="resources/powerplants/powerplants.csv",
     log:
         "logs/build_powerplants.log",
+    benchmark:
+        BENCHMARKS + "build_powerplants"
     resources:
         mem_mb=30000,
         walltime=config_provider("walltime", "build_powerplants", default="00:30:00"),
@@ -889,6 +899,8 @@ rule aggregate_to_substations:
         busmap=BUSMAPS + "{interconnect}/busmap_b.csv",
     log:
         "logs/aggregate_to_substations/{interconnect}.log",
+    benchmark:
+        BENCHMARKS + "{interconnect}/aggregate_to_substations"
     threads: 1
     resources:
         mem_mb=lambda wildcards, input, attempt: (input.size // 150000) * attempt * 1.5,
@@ -938,6 +950,8 @@ rule cluster_resources:
         busmap=BUSMAPS + "{interconnect}/busmap_s{simpl}.csv",
     log:
         "logs/cluster_resources/{interconnect}/elec_s{simpl}.log",
+    benchmark:
+        BENCHMARKS + "{interconnect}/cluster_resources_elec_s{simpl}"
     threads: 1
     resources:
         mem_mb=lambda wildcards, input, attempt: (input.size // 150000) * attempt * 1.5,
@@ -956,6 +970,8 @@ rule build_servm_load_weights:
         weights=DEMAND + "{interconnect}/servm_load_weights_s{simpl}.csv",
     log:
         LOGS + "{interconnect}/build_servm_load_weights_s{simpl}.log",
+    benchmark:
+        BENCHMARKS + "{interconnect}/build_servm_load_weights_s{simpl}"
     threads: 1
     resources:
         mem_mb=lambda wildcards, input, attempt: (input.size // 200000) * attempt * 2,
@@ -1071,6 +1087,8 @@ rule add_extra_components:
         NETWORKS + "{interconnect}/elec_s{simpl}_c{clusters}_ec.nc",
     log:
         "logs/add_extra_components/{interconnect}/elec_s{simpl}_c{clusters}_ec.log",
+    benchmark:
+        BENCHMARKS + "{interconnect}/add_extra_components_elec_s{simpl}_c{clusters}_ec"
     threads: 1
     resources:
         mem_mb=lambda wildcards, input, attempt: (input.size // 100000) * attempt * 2,
@@ -1109,6 +1127,9 @@ rule prepare_network:
         NETWORKS + "{interconnect}/elec_s{simpl}_c{clusters}_ec_l{ll}_{opts}.nc",
     log:
         "logs/prepare_network/{interconnect}/elec_s{simpl}_c{clusters}_ec_l{ll}_{opts}.log",
+    benchmark:
+        BENCHMARKS
+        +"{interconnect}/prepare_network_elec_s{simpl}_c{clusters}_ec_l{ll}_{opts}"
     threads: 1
     resources:
         walltime=config_provider("walltime", "prepare_network", default="00:30:00"),

workflow/rules/build_sector.smk

@@ -57,6 +57,9 @@ rule add_sectors:
         + "{interconnect}/elec_s{simpl}_c{clusters}_ec_l{ll}_{opts}_{sector}.nc",
     log:
         "logs/add_sectors/{interconnect}/elec_s{simpl}_c{clusters}_ec_l{ll}_{opts}_{sector}.log",
+    benchmark:
+        BENCHMARKS
+        +"{interconnect}/add_sectors_elec_s{simpl}_c{clusters}_ec_l{ll}_{opts}_{sector}"
     group:
         "prepare"
     threads: 1