Skip to content

Commit 27e2817

Browse files
michielkallenbergmichielkallenberg
committed
slurm multiple seeds
1 parent d40dc1d commit 27e2817

5 files changed

Lines changed: 210 additions & 7 deletions

File tree

‎cybench/runs/slurm/benchmark_completion_lib.py‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,88 @@ def write_manifest(path: Path, jobs: list[JobRow], *, header: str | None = None)
109109
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
110110

111111

112+
def target_walk_forward_seeds(*, base_seed: int, total_repetitions: int) -> list[int]:
113+
if total_repetitions < 1:
114+
raise ValueError(f"total_repetitions must be >= 1, got {total_repetitions}")
115+
return list(range(base_seed, base_seed + total_repetitions))
116+
117+
118+
def walk_forward_seeds_for_job(
119+
job: JobRow,
120+
*,
121+
baselines_dir: Path,
122+
horizon_tag_value: str,
123+
repo_root: Path,
124+
base_seed: int,
125+
total_repetitions: int,
126+
resume: bool,
127+
) -> list[int]:
128+
"""Seeds to run for one job (full target list, or missing only when resume)."""
129+
targets = target_walk_forward_seeds(
130+
base_seed=base_seed, total_repetitions=total_repetitions
131+
)
132+
if not resume:
133+
return targets
134+
run_dir = _latest_run(
135+
baselines_dir,
136+
crop=job.crop,
137+
country=job.country,
138+
model_slug=job.model,
139+
phase="walk_forward",
140+
horizon_tag_value=horizon_tag_value,
141+
repo_root=repo_root,
142+
)
143+
if run_dir is None:
144+
return targets
145+
existing = set(discover_run_seeds(run_dir))
146+
return [seed for seed in targets if seed not in existing]
147+
148+
149+
def expand_walk_forward_manifest_lines(
150+
jobs: list[JobRow],
151+
*,
152+
baselines_dir: Path,
153+
horizon: str,
154+
repo_root: Path,
155+
base_seed: int = 42,
156+
total_repetitions: int = 1,
157+
resume: bool = False,
158+
per_seed_for_gpu: bool = True,
159+
) -> list[str]:
160+
"""Expand manifest lines. GPU rows become one SLURM task per seed; CPU/naive stay bundled."""
161+
hz_tag = horizon_tag(horizon)
162+
lines: list[str] = []
163+
for job in jobs:
164+
if per_seed_for_gpu and job.needs_gpu == "yes":
165+
seeds = walk_forward_seeds_for_job(
166+
job,
167+
baselines_dir=baselines_dir,
168+
horizon_tag_value=hz_tag,
169+
repo_root=repo_root,
170+
base_seed=base_seed,
171+
total_repetitions=total_repetitions,
172+
resume=resume,
173+
)
174+
for seed in seeds:
175+
lines.append(f"{job.to_line()} {seed}")
176+
continue
177+
178+
if resume:
179+
seeds = walk_forward_seeds_for_job(
180+
job,
181+
baselines_dir=baselines_dir,
182+
horizon_tag_value=hz_tag,
183+
repo_root=repo_root,
184+
base_seed=base_seed,
185+
total_repetitions=total_repetitions,
186+
resume=True,
187+
)
188+
if not seeds:
189+
continue
190+
lines.append(job.to_line())
191+
return lines
192+
193+
112194
def horizon_tag(end_of_sequence: str) -> str:
113195
return prediction_horizon_tag(normalize_horizon(end_of_sequence))
114196

‎cybench/runs/slurm/expand_walk_forward_manifest.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
#!/usr/bin/env python3
2+
"""Expand walk-forward GPU manifest rows to one SLURM task per seed."""
3+
4+
from __future__ import annotations
5+
6+
import argparse
7+
from pathlib import Path
8+
9+
from cybench.runs.slurm.benchmark_completion_lib import (
10+
default_output_root,
11+
expand_walk_forward_manifest_lines,
12+
read_manifest,
13+
resolve_batch_dir,
14+
)
15+
16+
_REPO_ROOT = Path(__file__).resolve().parents[3]
17+
18+
19+
def main(argv: list[str] | None = None) -> int:
20+
parser = argparse.ArgumentParser(description=__doc__)
21+
parser.add_argument("--input", type=Path, required=True)
22+
parser.add_argument("--output", type=Path, required=True)
23+
parser.add_argument("--batch", required=True, help="Hydra experiment.name / baselines folder")
24+
parser.add_argument("--horizon", default="eos")
25+
parser.add_argument("--repetitions", type=int, default=1)
26+
parser.add_argument("--base-seed", type=int, default=42)
27+
parser.add_argument("--resume", action="store_true")
28+
parser.add_argument(
29+
"--per-seed",
30+
action="store_true",
31+
help="Expand needs_gpu=yes rows to one line per seed (default for gpu manifests)",
32+
)
33+
parser.add_argument("--output-root", type=Path)
34+
parser.add_argument("--repo-root", type=Path, default=_REPO_ROOT)
35+
args = parser.parse_args(argv)
36+
37+
repo_root = args.repo_root.resolve()
38+
output_root = args.output_root or default_output_root(repo_root)
39+
baselines_dir, _ = resolve_batch_dir(output_root, args.batch)
40+
jobs = read_manifest(args.input.resolve())
41+
lines = expand_walk_forward_manifest_lines(
42+
jobs,
43+
baselines_dir=baselines_dir,
44+
horizon=args.horizon,
45+
repo_root=repo_root,
46+
base_seed=args.base_seed,
47+
total_repetitions=args.repetitions,
48+
resume=args.resume,
49+
per_seed_for_gpu=args.per_seed,
50+
)
51+
header = (
52+
"# crop country model framework hp_search feature_design needs_gpu [seed]\n"
53+
"# GPU rows: one SLURM task per seed. CPU/naive: bundled (no seed column)."
54+
)
55+
args.output.write_text(
56+
header + ("\n" + "\n".join(lines) if lines else "") + "\n",
57+
encoding="utf-8",
58+
)
59+
print(f"[INFO] Expanded {len(jobs)} job(s) -> {len(lines)} SLURM task(s) -> {args.output}")
60+
return 0
61+
62+
63+
if __name__ == "__main__":
64+
raise SystemExit(main())

‎cybench/runs/slurm/orchestrate_benchmark_complete.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -520,10 +520,14 @@ def main(argv: list[str] | None = None) -> int:
520520

521521
args._targets = targets # type: ignore[attr-defined]
522522
exit_code = 0
523-
for batch, horizon in targets:
523+
n_targets = len(targets)
524+
for idx, (batch, horizon) in enumerate(targets, start=1):
525+
print(f"\n[INFO] ({idx}/{n_targets}) {batch} | horizon={horizon}")
524526
code = _process_batch(batch=batch, horizon=horizon, args=args)
525527
if code != 0:
526528
exit_code = code
529+
if n_targets:
530+
print(f"\n[DONE] Processed {n_targets} batch×horizon target(s)")
527531
return exit_code
528532

529533

‎cybench/runs/slurm/slurm_common.sh‎

Lines changed: 42 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,11 @@ slurm_task_job_name() {
141141
*) p=$(echo "${phase}" | cut -c1-2) ;;
142142
esac
143143
crop_s=$(slurm_crop_short "${CROP}")
144-
echo "cb_${p}_${MODEL}_${crop_s}${COUNTRY}"
144+
if [[ -n "${WF_SEED:-}" ]]; then
145+
echo "cb_${p}_${MODEL}_${crop_s}${COUNTRY}_s${WF_SEED}"
146+
else
147+
echo "cb_${p}_${MODEL}_${crop_s}${COUNTRY}"
148+
fi
145149
}
146150

147151
# JobId for scontrol update: must be ArrayJobId_ArrayTaskId, not ArrayJobId alone
@@ -301,7 +305,12 @@ read_benchmark_job() {
301305
echo "No job for array task ${SLURM_ARRAY_TASK_ID} in ${JOB_MANIFEST}" >&2
302306
exit 1
303307
fi
304-
read -r CROP COUNTRY MODEL FRAMEWORK HP_SEARCH FEATURE_DESIGN NEEDS_GPU <<< "${line}"
308+
read -r CROP COUNTRY MODEL FRAMEWORK HP_SEARCH FEATURE_DESIGN NEEDS_GPU _extra <<< "${line}"
309+
WF_SEED=""
310+
if [[ $(awk '{print NF}' <<< "${line}") -ge 8 ]]; then
311+
WF_SEED=$(awk '{print $NF}' <<< "${line}")
312+
fi
313+
export WF_SEED
305314
}
306315

307316
# CPU tabular: one Optuna trial at a time, sklearn uses all SLURM CPUs.
@@ -470,3 +479,34 @@ plan_walk_forward_seeds() {
470479
WF_RUN_REPS="${#missing[@]}"
471480
return 0
472481
}
482+
483+
# Single-seed walk-forward task (GPU manifest with 8th column). Sets WF_RUN_DIR, WF_START_SEED, WF_RUN_REPS=1.
484+
# Returns 0 run, 1 skip (seed present), 2 error (seed>base without run dir).
485+
plan_walk_forward_single_seed() {
486+
local crop=$1 country=$2 model_slug=$3 seed=$4
487+
local base=${WF_BASE_SEED:-42}
488+
local run_dir=""
489+
490+
run_dir=$(find_latest_walk_forward_run_dir "${crop}" "${country}" "${model_slug}")
491+
if [[ -n "${run_dir}" && -d "${run_dir}" ]]; then
492+
local -a existing=()
493+
mapfile -t existing < <(discover_run_seeds_py "${run_dir}")
494+
local e
495+
for e in "${existing[@]}"; do
496+
if [[ "${e}" == "${seed}" ]]; then
497+
echo "[SKIP] Walk-forward seed ${seed} already in ${run_dir}" >&2
498+
return 1
499+
fi
500+
done
501+
WF_RUN_DIR="$(cd "${run_dir}" && pwd)"
502+
elif [[ "${seed}" != "${base}" ]]; then
503+
echo "[ERROR] Seed ${seed} requires existing walk-forward run (run seed ${base} first)" >&2
504+
return 2
505+
else
506+
WF_RUN_DIR=""
507+
fi
508+
509+
WF_START_SEED="${seed}"
510+
WF_RUN_REPS=1
511+
return 0
512+
}

‎cybench/runs/slurm/walk_forward.sh‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
#!/bin/bash
22
#
3-
# Walk-forward: one array task = one (crop, country, model), using frozen screening artifacts.
3+
# Walk-forward: one array task = one (crop, country, model) [× seed for GPU manifest].
44
# Auto-discovers the latest screening run under CYBENCH_EXPERIMENT_NAME (default: baselines):
55
# ../output/<batch>/<crop>_<country>_<model>_screening_<horizon>_<timestamp>/<test_years>/optimal_model.yaml
66
#
7+
# GPU manifest rows with an 8th column (seed) run one seed per SLURM task.
8+
# CPU/naive manifests bundle seeds in one task (experiment.n_repetitions).
9+
#
710
# Submit after screening jobs finished:
811
# sbatch cybench/runs/slurm/walk_forward.sh
912
#
@@ -48,15 +51,25 @@ WF_RUN_DIR=""
4851
WF_START_SEED="${WF_BASE_SEED}"
4952
WF_RUN_REPS="${WF_REPETITIONS}"
5053

51-
if ! plan_walk_forward_seeds "${CROP}" "${COUNTRY}" "${MODEL}"; then
54+
if [[ -n "${WF_SEED:-}" ]]; then
55+
if ! plan_walk_forward_single_seed "${CROP}" "${COUNTRY}" "${MODEL}" "${WF_SEED}"; then
56+
exit 0
57+
fi
58+
elif ! plan_walk_forward_seeds "${CROP}" "${COUNTRY}" "${MODEL}"; then
5259
exit 0
5360
fi
5461

5562
resume_note=""
5663
if [[ -n "${WF_RUN_DIR}" ]]; then
57-
resume_note=" | resume=${WF_RUN_DIR} | seeds=${WF_START_SEED}..$((WF_START_SEED + WF_RUN_REPS - 1))"
64+
resume_note=" | resume=${WF_RUN_DIR}"
65+
fi
66+
seed_note=" | seed=${WF_START_SEED}"
67+
if [[ -n "${WF_SEED:-}" ]]; then
68+
seed_note=" | task_seed=${WF_SEED}"
69+
elif [[ "${WF_RUN_REPS}" != "1" ]]; then
70+
seed_note=" | seeds=${WF_START_SEED}..$((WF_START_SEED + WF_RUN_REPS - 1))"
5871
fi
59-
echo "Walk-forward | ${CROP}/${COUNTRY} | model=${MODEL} | device=$(device_mode_label) | horizon=${PREDICTION_HORIZON} | batch=${CYBENCH_EXPERIMENT_NAME} | target_repetitions=${WF_REPETITIONS}${resume_note} | frozen=${FROZEN_DIR}"
72+
echo "Walk-forward | ${CROP}/${COUNTRY} | model=${MODEL} | device=$(device_mode_label) | horizon=${PREDICTION_HORIZON} | batch=${CYBENCH_EXPERIMENT_NAME} | target_repetitions=${WF_REPETITIONS}${seed_note}${resume_note} | frozen=${FROZEN_DIR}"
6073

6174
COMMON=(
6275
"dataset/crop=${CROP}"

0 commit comments

Comments
 (0)