Optimize calibration evaluation

evaluate_calibration_all_runs_linux.sh

@@ -29,6 +29,7 @@ LOG_ROOT="$SCRIPT_DIR/batch_logs/evaluate_calibration_all_runs"
 GPU_CSV="0"
 PYTHON_BIN="${PYTHON_BIN:-python}"
 NUM_WORKERS=4
+NUM_WORKERS_CALIBRATION=0
 BATCH_SIZE=128
 DISEASE_CHUNK_SIZE=64
 HORIZONS=""
@@ -52,6 +53,8 @@ Options:
   --python PATH               Python executable
                               (default: $PYTHON_BIN or python).
   --num-workers N             DataLoader workers per job (default: 4).
+  --num-workers-calibration N CPU calibration workers per job. Default: 0,
+                              which divides all logical CPUs across GPUs.
   --batch-size N              Evaluation batch size (default: 128).
   --disease-chunk-size N      Disease projection chunk size (default: 64).
   --horizons LIST             Optional comma-separated horizons in years.
@@ -107,6 +110,14 @@ while (($# > 0)); do
             NUM_WORKERS="$2"
             shift 2
             ;;
+        --num-workers-calibration)
+            [[ $# -ge 2 ]] || {
+                echo "ERROR: --num-workers-calibration requires a value." >&2
+                exit 2
+            }
+            NUM_WORKERS_CALIBRATION="$2"
+            shift 2
+            ;;
         --batch-size)
             [[ $# -ge 2 ]] || {
                 echo "ERROR: --batch-size requires a value." >&2
@@ -163,6 +174,10 @@ done
     echo "ERROR: --num-workers must be a non-negative integer." >&2
     exit 2
 }
+[[ "$NUM_WORKERS_CALIBRATION" =~ ^[0-9]+$ ]] || {
+    echo "ERROR: --num-workers-calibration must be a non-negative integer." >&2
+    exit 2
+}
 [[ "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || {
     echo "ERROR: --batch-size must be a positive integer." >&2
     exit 2
@@ -203,6 +218,14 @@ for gpu in "${GPU_IDS[@]}"; do
     SEEN_GPUS["$gpu"]=1
 done
 
+if ((NUM_WORKERS_CALIBRATION == 0)); then
+    TOTAL_CPUS="$(nproc)"
+    NUM_WORKERS_CALIBRATION=$(( (TOTAL_CPUS + ${#GPU_IDS[@]} - 1) / ${#GPU_IDS[@]} ))
+    if ((NUM_WORKERS_CALIBRATION < 1)); then
+        NUM_WORKERS_CALIBRATION=1
+    fi
+fi
+
 declare -a JOB_RUN_DIRS=()
 declare -a JOB_LOG_FILES=()
 
@@ -281,6 +304,7 @@ run_job() {
         --eval_split test
         --device cuda
         --num_workers "$NUM_WORKERS"
+        --num_workers_calibration "$NUM_WORKERS_CALIBRATION"
         --batch_size "$BATCH_SIZE"
         --disease_chunk_size "$DISEASE_CHUNK_SIZE"
     )
@@ -341,6 +365,7 @@ worker() {
 
 echo "Runs root: $RUNS_ROOT"
 echo "GPUs: ${GPU_IDS[*]}"
+echo "Calibration CPU workers per GPU job: $NUM_WORKERS_CALIBRATION"
 echo "Runs discovered: $run_count"
 echo "Incomplete runs skipped: $incomplete_count"
 echo "next_token runs skipped: $next_token_count"