HP_trials=20

GPU threshold at 105 regions

Author: michielkallenberg

cybench/runs/slurm/benchmark_submit_lib.py

@@ -12,7 +12,7 @@
 from cybench.config import DATASETS, PATH_DATA_DIR
 
 # Countries with at least this many regions use the gpu SLURM partition for torch/TabPFN.
-DEFAULT_GPU_REGION_THRESHOLD = 50
+DEFAULT_GPU_REGION_THRESHOLD = 105
 
 _BATCH_RE = re.compile(
     r"^baselines_(?P<country>[A-Za-z]{2})_(?P<batch_hz>eos|mid|qtr|early)_v(?P<version>\d+)$"

cybench/runs/slurm/orchestrate_benchmark_complete.sh

@@ -39,7 +39,7 @@ Options:
   --dry-run           With --submit: no sbatch
   --cpu               Force torch/TabPFN group to CPU partition
   --force-gpu         Use gpu partition even when regions < threshold
-  --region-threshold N  gpu when country has >= N regions (default: 50)
+  --region-threshold N  gpu when country has >= N regions (default: 105)
 
 Examples:
   orchestrate_benchmark_complete.sh --country DE --horizons eos mid --list

cybench/runs/slurm/orchestrate_benchmark_submit.sh

@@ -12,6 +12,8 @@
 #
 set -euo pipefail
 
+export HP_TRIALS="${HP_TRIALS:-20}"
+
 usage() {
   cat <<'EOF'
 Usage: orchestrate_benchmark_submit.sh [options]
@@ -28,7 +30,7 @@ Options:
   --dry-run           Print submit_benchmark.sh commands without sbatch
   --countries CC ...  Limit to these countries (default: all with data on disk)
   --horizon H ...     eos | mid | qtr | early (alias early-season; repeatable)
-  --region-threshold N  gpu partition when country has >= N regions (default: 50)
+  --region-threshold N  gpu partition when country has >= N regions (default: 105)
   --version N         Batch version suffix (default: 3)
   --phase MODE        screening | walk_forward | all (default: all)
   --repetitions N     Walk-forward seeds 42..42+N-1 (passed to submit_benchmark.sh)
@@ -41,7 +43,7 @@ Options:
   --data-dir DIR      Override cybench/data
 
 Environment:
-  HP_TRIALS           Optuna trials in screening (default: 100)
+  HP_TRIALS           Optuna trials in screening (default: 20)
   WF_REPETITIONS      Default for --repetitions when omitted (default: 1)
 
 Examples:
@@ -71,7 +73,7 @@ FORCE=false
 ALL_COUNTRIES=false
 PHASE_MODE="all"
 VERSION=3
-REGION_THRESHOLD=50
+REGION_THRESHOLD=105
 MAX_BATCHES=0
 MANIFEST_ROOT="${SLURM_DIR}/manifests"
 DATA_DIR=""

cybench/runs/slurm/slurm_common.sh

@@ -354,7 +354,7 @@ configure_parallelism() {
 configure_hpo_extras() {
   local -n _extra=$1
   if [[ "${HP_SEARCH}" == "yes" ]]; then
-    _extra+=(+hp_search=bayesian hp_search.n_trials="${HP_TRIALS:-100}")
+    _extra+=(+hp_search=bayesian hp_search.n_trials="${HP_TRIALS:-20}")
     _extra+=(
       "hp_search.storage.url=sqlite:///${TMPDIR:-/tmp}/optuna_${SLURM_JOB_ID}_${SLURM_ARRAY_TASK_ID}.db"
     )
@@ -482,13 +482,35 @@ plan_walk_forward_seeds() {
 }
 
 # Single-seed walk-forward task (GPU manifest with 8th column). Sets WF_RUN_DIR, WF_START_SEED, WF_RUN_REPS=1.
-# Returns 0 run, 1 skip (seed present), 2 error (seed>base without run dir).
+# Returns 0 run, 1 skip (seed present), 2 error (timed out waiting for base-seed run dir).
 plan_walk_forward_single_seed() {
   local crop=$1 country=$2 model_slug=$3 seed=$4
   local base=${WF_BASE_SEED:-42}
   local run_dir=""
+  local wait_secs=${WF_SEED_WAIT_SECS:-7200}
+  local interval=${WF_SEED_WAIT_INTERVAL:-30}
+  local elapsed=0
+
+  # Parallel per-seed arrays may start seed 43+ before seed ${base} creates the run dir.
+  while true; do
+    run_dir=$(find_latest_walk_forward_run_dir "${crop}" "${country}" "${model_slug}")
+    if [[ -n "${run_dir}" && -d "${run_dir}" ]]; then
+      break
+    fi
+    if [[ "${seed}" == "${base}" ]]; then
+      break
+    fi
+    if [[ "${elapsed}" -eq 0 ]]; then
+      echo "[WAIT] Seed ${seed} waiting for walk-forward run dir (seed ${base})..." >&2
+    fi
+    if [[ "${elapsed}" -ge "${wait_secs}" ]]; then
+      echo "[ERROR] Seed ${seed} timed out after ${wait_secs}s waiting for walk-forward run (seed ${base})" >&2
+      return 2
+    fi
+    sleep "${interval}"
+    elapsed=$((elapsed + interval))
+  done
 
-  run_dir=$(find_latest_walk_forward_run_dir "${crop}" "${country}" "${model_slug}")
   if [[ -n "${run_dir}" && -d "${run_dir}" ]]; then
     local -a existing=()
     mapfile -t existing < <(discover_run_seeds_py "${run_dir}")
@@ -500,9 +522,6 @@ plan_walk_forward_single_seed() {
       fi
     done
     WF_RUN_DIR="$(cd "${run_dir}" && pwd)"
-  elif [[ "${seed}" != "${base}" ]]; then
-    echo "[ERROR] Seed ${seed} requires existing walk-forward run (run seed ${base} first)" >&2
-    return 2
   else
     WF_RUN_DIR=""
   fi

cybench/runs/slurm/walk_forward.sh

@@ -52,9 +52,13 @@ WF_START_SEED="${WF_BASE_SEED}"
 WF_RUN_REPS="${WF_REPETITIONS}"
 
 if [[ -n "${WF_SEED:-}" ]]; then
-  if ! plan_walk_forward_single_seed "${CROP}" "${COUNTRY}" "${MODEL}" "${WF_SEED}"; then
-    exit 0
-  fi
+  plan_walk_forward_single_seed "${CROP}" "${COUNTRY}" "${MODEL}" "${WF_SEED}"
+  wf_plan_rc=$?
+  case "${wf_plan_rc}" in
+    0) ;;
+    1) exit 0 ;; # seed already present
+    *) exit 1 ;; # timeout / planning error — do not mask as success
+  esac
 elif ! plan_walk_forward_seeds "${CROP}" "${COUNTRY}" "${MODEL}"; then
   exit 0
 fi

tests/runs/test_benchmark_submit_lib.py

@@ -82,6 +82,20 @@ def test_resolve_batch_dir_and_parse_batch_name(tmp_path: Path):
     assert note is not None
 
 
+def test_gpu_region_threshold_default_routes_fr_to_cpu(tmp_path: Path, monkeypatch):
+    from cybench.runs.slurm import benchmark_submit_lib as lib
+
+    assert lib.DEFAULT_GPU_REGION_THRESHOLD == 105
+
+    def _fake_count(country: str, data_dir=None) -> int:
+        return {"FR": 100, "IT": 105, "DE": 200}[country.upper()]
+
+    monkeypatch.setattr(lib, "count_regions", _fake_count)
+    assert lib.gpu_partition_for_country("FR")[0] is False
+    assert lib.gpu_partition_for_country("IT")[0] is True
+    assert lib.gpu_partition_for_country("DE")[0] is True
+
+
 def test_expand_all_country_targets():
     from cybench.runs.slurm.benchmark_completion_lib import expand_all_country_targets