Skip to content

Commit 006bfc9

Browse files
michielkallenbergmichielkallenberg
committed
slurm
1 parent df3186e commit 006bfc9

3 files changed

Lines changed: 350 additions & 0 deletions

File tree

‎cybench/runs/analysis/global_insights_lib.py‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -879,6 +879,129 @@ def build_horizon_skill_curves_payload(
879879
}
880880

881881

882+
def report_model_horizon_pairs(
883+
df: pd.DataFrame,
884+
model: str,
885+
*,
886+
crop: str | None = None,
887+
min_sample_ratio: float | None = None,
888+
) -> dict[str, Any]:
889+
"""Paired mid/qtr/eos report for one model (fair country intersection).
890+
891+
Inner-joins on (country, model) across all horizons present in ``df``.
892+
Optional ``min_sample_ratio`` drops pairs where min(n_samples)/max(n_samples)
893+
across horizons falls below the threshold (e.g. 0.95).
894+
"""
895+
horizons = horizons_in_data(df)
896+
if len(horizons) < 2:
897+
return {"model": model, "error": "Need at least two horizons in summaries."}
898+
899+
work = _filter_summary_work(df, crop=crop)
900+
work = work[work["model"] == model]
901+
if work.empty:
902+
return {"model": model, "crop": crop, "error": f"No rows for model {model!r}."}
903+
904+
group_keys = ["country", "model"]
905+
wide: pd.DataFrame | None = None
906+
samples_wide: pd.DataFrame | None = None
907+
for hz in horizons:
908+
hz_df = work[work["batch_horizon"] == hz]
909+
if hz_df.empty:
910+
return {"model": model, "crop": crop, "error": f"No {hz} rows for {model}."}
911+
nrmse_part = hz_df.groupby(group_keys, as_index=False)["nrmse"].median().rename(
912+
columns={"nrmse": f"nrmse_{hz}"}
913+
)
914+
wide = nrmse_part if wide is None else wide.merge(nrmse_part, on=group_keys, how="inner")
915+
if "n_samples" in hz_df.columns:
916+
samp_part = hz_df.groupby(group_keys, as_index=False)["n_samples"].sum().rename(
917+
columns={"n_samples": f"n_samples_{hz}"}
918+
)
919+
samples_wide = (
920+
samp_part if samples_wide is None else samples_wide.merge(samp_part, on=group_keys, how="inner")
921+
)
922+
923+
if wide is None or wide.empty:
924+
return {"model": model, "crop": crop, "error": "No country intersection across horizons."}
925+
926+
if samples_wide is not None:
927+
for _, row in samples_wide.iterrows():
928+
vals = [row.get(f"n_samples_{hz}") for hz in horizons]
929+
vals = [float(v) for v in vals if pd.notna(v) and float(v) > 0]
930+
if len(vals) >= 2:
931+
ratio = min(vals) / max(vals)
932+
wide.loc[wide["country"] == row["country"], "_sample_ratio"] = ratio
933+
934+
if min_sample_ratio is not None and "_sample_ratio" in wide.columns:
935+
wide = wide[wide["_sample_ratio"].fillna(0) >= min_sample_ratio].copy()
936+
937+
crop_work = _filter_summary_work(df, crop=crop)
938+
any_countries = (
939+
set(crop_work[crop_work["model"] == model]["country"].astype(str).unique())
940+
if "country" in crop_work.columns
941+
else set()
942+
)
943+
paired_countries = sorted(wide["country"].astype(str).unique())
944+
excluded = sorted(any_countries - set(paired_countries))
945+
946+
detail_rows: list[dict[str, Any]] = []
947+
for _, row in wide.sort_values("country").iterrows():
948+
country = str(row["country"])
949+
nrmse_by_hz = {hz: float(row[f"nrmse_{hz}"]) for hz in horizons}
950+
best_hz = min(nrmse_by_hz, key=nrmse_by_hz.get) # type: ignore[arg-type]
951+
entry: dict[str, Any] = {
952+
"country": country,
953+
**{f"nrmse_{hz}": round(nrmse_by_hz[hz], 4) for hz in horizons},
954+
"best_horizon": best_hz,
955+
}
956+
if samples_wide is not None:
957+
sw = samples_wide[samples_wide["country"] == country]
958+
if not sw.empty:
959+
for hz in horizons:
960+
col = f"n_samples_{hz}"
961+
if col in sw.columns:
962+
val = sw.iloc[0][col]
963+
entry[col] = int(val) if pd.notna(val) else None
964+
if "_sample_ratio" in row.index and pd.notna(row["_sample_ratio"]):
965+
entry["sample_ratio"] = round(float(row["_sample_ratio"]), 4)
966+
detail_rows.append(entry)
967+
968+
summary: dict[str, Any] = {"n_paired_countries": len(paired_countries)}
969+
for hz in horizons:
970+
col = f"nrmse_{hz}"
971+
vals = wide[col].astype(float)
972+
summary[f"median_{col}"] = round(float(vals.median()), 4)
973+
summary[f"mean_{col}"] = round(float(vals.mean()), 4)
974+
975+
if "eos" in horizons:
976+
for hz in horizons:
977+
if hz == "eos":
978+
continue
979+
col = f"nrmse_{hz}"
980+
delta = wide[col] - wide["nrmse_eos"]
981+
wins = delta > 0 # lower NRMSE is better; positive delta => eos better
982+
summary[f"eos_better_than_{hz}_rate"] = round(float(wins.mean()), 4)
983+
summary[f"mean_delta_{hz}_minus_eos"] = round(float(delta.mean()), 4)
984+
summary[f"median_delta_{hz}_minus_eos"] = round(float(delta.median()), 4)
985+
986+
return {
987+
"model": model,
988+
"crop": crop or "all",
989+
"horizons": list(horizons),
990+
"horizon_labels": {hz: HORIZON_DISPLAY_LABELS.get(hz, hz) for hz in horizons},
991+
"n_paired_countries": len(paired_countries),
992+
"paired_countries": paired_countries,
993+
"excluded_countries": excluded,
994+
"summary": summary,
995+
"detail": detail_rows,
996+
"interpretation": (
997+
f"Paired comparison for {model}: only countries with walk-forward summaries "
998+
f"at all horizons ({', '.join(horizons)}). "
999+
"delta = horizon_nrmse − eos_nrmse; positive ⇒ end-of-season has lower NRMSE. "
1000+
"best_horizon = lowest NRMSE within country."
1001+
),
1002+
}
1003+
1004+
8821005
def build_crop_comparison_payload(df: pd.DataFrame) -> dict[str, dict[str, Any]]:
8831006
"""Per horizon, pairwise crop NRMSE on shared countries only."""
8841007
if df.empty:

‎cybench/runs/slurm/README.md‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,25 @@ manifest. Use `--cpu` on large GPU batches if the queue is backed up.
159159
`baselines_DE_eos_v1` (orchestrated). Completion resolves output/manifest paths
160160
case-insensitively and submits using the on-disk folder name.
161161

162+
**Restore full manifests after `--submit`:** `orchestrate_benchmark_complete.sh`
163+
overwrites `manifests/<batch>/benchmark_jobs.txt` with the incomplete retry list.
164+
Regenerate the full crop×model list before the next `--list` audit:
165+
166+
```bash
167+
export CYBENCH_OUTPUT_ROOT=/lustre/backup/SHARED/AIN/agml/output
168+
169+
# All countries with baselines_*_qtr_v2 on disk
170+
cybench/runs/slurm/refresh_batch_manifests.sh --horizon qtr --version 2 --backup
171+
172+
# Or specific countries (e.g. after partial qtr GPU reruns)
173+
cybench/runs/slurm/refresh_batch_manifests.sh \
174+
--countries BR BG CN FR --horizon qtr --version 2 --backup
175+
176+
# Then re-audit
177+
cybench/runs/slurm/orchestrate_benchmark_complete.sh \
178+
--all-countries --horizons qtr --version 2 --list
179+
```
180+
162181
Manual per-manifest submits (fine-grained control) are below.
163182

164183
## 2. Submit screening

‎cybench/runs/slurm/refresh_batch_manifests.sh‎

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
#!/usr/bin/env bash
2+
#
3+
# Regenerate full per-country benchmark_jobs.txt from generate_job_manifest.py.
4+
#
5+
# Use after orchestrate_benchmark_complete.sh --submit overwrote a working manifest
6+
# with an "incomplete jobs only" retry list.
7+
#
8+
# Usage (from repo root on anunna):
9+
# cybench/runs/slurm/refresh_batch_manifests.sh
10+
# cybench/runs/slurm/refresh_batch_manifests.sh --horizon qtr --version 2
11+
# cybench/runs/slurm/refresh_batch_manifests.sh --countries BR FR --horizon qtr --version 2
12+
# cybench/runs/slurm/refresh_batch_manifests.sh --dry-run
13+
#
14+
set -euo pipefail
15+
16+
usage() {
17+
cat <<'EOF'
18+
Usage: refresh_batch_manifests.sh [options]
19+
20+
Regenerate manifests/baselines_<CC>_<hz>_vN/benchmark_jobs.txt for batches that
21+
exist under the output root (or for explicit --countries).
22+
23+
Options:
24+
--horizon H eos | mid | qtr (default: qtr)
25+
--version N Batch version suffix (default: 2)
26+
--countries CC.. Only these countries (default: all matching batch dirs)
27+
--output-root DIR Parent of baselines_* (default: $CYBENCH_OUTPUT_ROOT or lustre)
28+
--data-dir DIR Passed to generate_job_manifest.py
29+
--backup Keep timestamped copy of existing benchmark_jobs.txt
30+
--dry-run Print commands only
31+
-h, --help This help
32+
33+
Examples:
34+
refresh_batch_manifests.sh --horizon qtr --version 2
35+
refresh_batch_manifests.sh --countries BR BG CN FR --horizon qtr --version 2 --backup
36+
EOF
37+
}
38+
39+
if [[ -f "${SLURM_SUBMIT_DIR:-}/cybench/runs/slurm/slurm_common.sh" ]]; then
40+
SLURM_DIR="${SLURM_SUBMIT_DIR}/cybench/runs/slurm"
41+
elif [[ -n "${REPO_ROOT:-}" && -f "${REPO_ROOT}/cybench/runs/slurm/slurm_common.sh" ]]; then
42+
SLURM_DIR="${REPO_ROOT}/cybench/runs/slurm"
43+
else
44+
SLURM_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
45+
fi
46+
# shellcheck source=cybench/runs/slurm/slurm_common.sh
47+
source "${SLURM_DIR}/slurm_common.sh"
48+
49+
REPO_ROOT="${REPO_ROOT:-$(cd "${SLURM_DIR}/../../.." && pwd)}"
50+
GENERATE_PY="${SLURM_DIR}/generate_job_manifest.py"
51+
52+
HORIZON_KEY="qtr"
53+
VERSION="2"
54+
COUNTRIES=()
55+
OUTPUT_ROOT="${CYBENCH_OUTPUT_ROOT:-}"
56+
DATA_DIR=""
57+
BACKUP=false
58+
DRY_RUN=false
59+
60+
while [[ $# -gt 0 ]]; do
61+
case "$1" in
62+
--horizon)
63+
HORIZON_KEY=$2
64+
shift 2
65+
;;
66+
--version)
67+
VERSION=$2
68+
shift 2
69+
;;
70+
--countries)
71+
shift
72+
while [[ $# -gt 0 && ! "$1" =~ ^-- ]]; do
73+
COUNTRIES+=("$(printf '%s' "$1" | tr '[:lower:]' '[:upper:]')")
74+
shift
75+
done
76+
;;
77+
--output-root)
78+
OUTPUT_ROOT=$2
79+
shift 2
80+
;;
81+
--data-dir)
82+
DATA_DIR=$2
83+
shift 2
84+
;;
85+
--backup)
86+
BACKUP=true
87+
shift
88+
;;
89+
--dry-run)
90+
DRY_RUN=true
91+
shift
92+
;;
93+
-h | --help)
94+
usage
95+
exit 0
96+
;;
97+
*)
98+
echo "Unknown option: $1" >&2
99+
usage >&2
100+
exit 1
101+
;;
102+
esac
103+
done
104+
105+
case "${HORIZON_KEY}" in
106+
eos) BATCH_HZ="eos"; SLURM_HORIZON="eos" ;;
107+
mid | middle-of-season | middle_of_season | mid-season | mid_season)
108+
BATCH_HZ="mid"
109+
SLURM_HORIZON="middle-of-season"
110+
;;
111+
qtr | quarter-of-season | quarter_of_season | quarter-season | quarter_season)
112+
BATCH_HZ="qtr"
113+
SLURM_HORIZON="quarter-of-season"
114+
;;
115+
*)
116+
echo "Unknown horizon: ${HORIZON_KEY}" >&2
117+
exit 1
118+
;;
119+
esac
120+
121+
if [[ -z "${OUTPUT_ROOT}" ]]; then
122+
if [[ -d /lustre/backup/SHARED/AIN/agml/output ]]; then
123+
OUTPUT_ROOT=/lustre/backup/SHARED/AIN/agml/output
124+
else
125+
OUTPUT_ROOT="${REPO_ROOT}/../output"
126+
fi
127+
fi
128+
129+
if [[ ! -d "${OUTPUT_ROOT}" ]]; then
130+
echo "Output root not found: ${OUTPUT_ROOT}" >&2
131+
exit 1
132+
fi
133+
134+
parse_country_from_batch_dir() {
135+
local name=$1
136+
if [[ "${name}" =~ ^baselines_([A-Za-z]{2})_${BATCH_HZ}_v${VERSION}$ ]]; then
137+
printf '%s' "${BASH_REMATCH[1]}" | tr '[:lower:]' '[:upper:]'
138+
return 0
139+
fi
140+
return 1
141+
}
142+
143+
declare -a TARGET_COUNTRIES=()
144+
if [[ ${#COUNTRIES[@]} -gt 0 ]]; then
145+
TARGET_COUNTRIES=("${COUNTRIES[@]}")
146+
else
147+
shopt -s nullglob
148+
declare -A seen=()
149+
for d in "${OUTPUT_ROOT}"/baselines_*_"${BATCH_HZ}"_v"${VERSION}"; do
150+
[[ -d "${d}" ]] || continue
151+
cc=$(parse_country_from_batch_dir "$(basename "${d}")") || continue
152+
if [[ -z "${seen[${cc}]:-}" ]]; then
153+
seen["${cc}"]=1
154+
TARGET_COUNTRIES+=("${cc}")
155+
fi
156+
done
157+
shopt -u nullglob
158+
if [[ ${#TARGET_COUNTRIES[@]} -eq 0 ]]; then
159+
echo "No baselines_*_${BATCH_HZ}_v${VERSION} directories under ${OUTPUT_ROOT}" >&2
160+
exit 1
161+
fi
162+
IFS=$'\n' TARGET_COUNTRIES=($(printf '%s\n' "${TARGET_COUNTRIES[@]}" | sort))
163+
unset IFS
164+
fi
165+
166+
run_generate() {
167+
local cc=$1
168+
local batch="baselines_${cc}_${BATCH_HZ}_v${VERSION}"
169+
local manifest_dir
170+
manifest_dir="$(manifest_batch_dir "${SLURM_DIR}" "${batch}")"
171+
local out="${manifest_dir}/benchmark_jobs.txt"
172+
mkdir -p "${manifest_dir}"
173+
174+
if [[ -f "${out}" && "${BACKUP}" == true ]]; then
175+
local stamp
176+
stamp=$(date -u +%Y%m%dT%H%M%SZ)
177+
local backup="${out}.bak.${stamp}"
178+
if [[ "${DRY_RUN}" == true ]]; then
179+
echo "[DRY-RUN] cp ${out} ${backup}"
180+
else
181+
cp "${out}" "${backup}"
182+
echo "[backup] ${backup}"
183+
fi
184+
fi
185+
186+
local cmd=(
187+
poetry run python "${GENERATE_PY}"
188+
--countries "${cc}"
189+
--horizon "${SLURM_HORIZON}"
190+
-o "${out}"
191+
)
192+
if [[ -n "${DATA_DIR}" ]]; then
193+
cmd+=(--data-dir "${DATA_DIR}")
194+
fi
195+
196+
echo "Generating ${cc} -> ${out}"
197+
if [[ "${DRY_RUN}" == true ]]; then
198+
printf ' '; printf '%q ' "${cmd[@]}"; printf '\n'
199+
else
200+
(cd "${REPO_ROOT}" && "${cmd[@]}")
201+
fi
202+
}
203+
204+
echo "Refresh manifests | horizon=${SLURM_HORIZON} (${BATCH_HZ}) | version=${VERSION} | output=${OUTPUT_ROOT}"
205+
for cc in "${TARGET_COUNTRIES[@]}"; do
206+
run_generate "${cc}"
207+
done
208+
echo "[DONE] Refreshed ${#TARGET_COUNTRIES[@]} manifest(s)"

0 commit comments

Comments
 (0)