DeepHealth/train_key_models_multiseed_linux.sh

#!/usr/bin/env bash
#
# Re-run the key DeepHealth experiments with additional random seeds.
#
# The selected configurations form a compact evidence chain:
#   1. FFN + next_token + absolute + exponential
#      Delphi2M reproduction baseline.
#   2. FFN + all_future + absolute + exponential
#      Isolates the target change from next_token to all_future.
#   3. TrajMixer + all_future + absolute + exponential
#      Isolates the architecture change from FFN to TrajMixer.
#   4. TrajMixer + all_future + relative + exponential
#      Isolates the time-mode change and is the current disease-best candidate.
#   5. TrajMixer + all_future + relative + Weibull
#      Unified disease/death candidate.
#   6. FFN + all_future + relative + Weibull
#      Current mortality-best control and matched architecture comparison for (5).
#
# Each task uses one GPU. Tasks assigned to the same GPU run sequentially,
# while different GPUs run in parallel.
#
# Examples:
#   bash train_key_models_multiseed_linux.sh --gpus 0,1,2,3
#   bash train_key_models_multiseed_linux.sh --gpus 0,1 --seeds 43,44
#

set -uo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

GPU_CSV=""
SEED_CSV="43,44,45"
NUM_WORKERS=4
PYTHON_BIN="${PYTHON_BIN:-python}"
CAMPAIGN_NAME="key_models_multiseed_smoking_alcohol_bmi"
DRY_RUN=0

BATCH_SIZE=256
EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_smoking_alcohol_bmi.txt"

usage() {
    cat <<'EOF'
Usage:
  bash train_key_models_multiseed_linux.sh --gpus GPU_LIST [options]

Required:
  --gpus LIST            Comma-separated GPU ids, for example 0,1,2,3.

Options:
  --seeds LIST           Additional seeds (default: 43,44,45).
                         Use 43,44 for two additional seeds.
  --num-workers N        DataLoader workers per task (default: 4).
  --python PATH          Python executable (default: $PYTHON_BIN or python).
  --campaign NAME        Output campaign name.
  --dry-run              Print all commands without running them.
  -h, --help             Show this help message.

Fixed experiment settings:
  batch_size             256
  extra_info_types       extra_info_types_smoking_alcohol_bmi.txt
  continuous scaling     required train-split RobustScale
  model size             Defaults from the training entrypoints
  tasks per seed         6

Outputs:
  runs/<campaign>/seed_<seed>/<architecture>/...
  batch_logs/<campaign>/seed_<seed>/<job>.log
EOF
}

while (($# > 0)); do
    case "$1" in
        --gpus)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --gpus requires a value." >&2
                exit 2
            }
            GPU_CSV="$2"
            shift 2
            ;;
        --seeds)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --seeds requires a value." >&2
                exit 2
            }
            SEED_CSV="$2"
            shift 2
            ;;
        --num-workers)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --num-workers requires a value." >&2
                exit 2
            }
            NUM_WORKERS="$2"
            shift 2
            ;;
        --python)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --python requires a value." >&2
                exit 2
            }
            PYTHON_BIN="$2"
            shift 2
            ;;
        --campaign)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --campaign requires a value." >&2
                exit 2
            }
            CAMPAIGN_NAME="$2"
            shift 2
            ;;
        --dry-run)
            DRY_RUN=1
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "ERROR: unknown argument: $1" >&2
            usage >&2
            exit 2
            ;;
    esac
done

[[ -n "$GPU_CSV" ]] || {
    echo "ERROR: --gpus is required." >&2
    usage >&2
    exit 2
}
[[ -n "$SEED_CSV" ]] || {
    echo "ERROR: --seeds must not be empty." >&2
    exit 2
}
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
    echo "ERROR: --num-workers must be a non-negative integer." >&2
    exit 2
}
[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || {
    echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2
    exit 2
}
[[ -f "$EXTRA_INFO_TYPES_FILE" ]] || {
    echo "ERROR: missing extra-info file: $EXTRA_INFO_TYPES_FILE" >&2
    exit 2
}
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
    echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
    exit 2
}

IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
declare -A SEEN_GPUS=()
for gpu in "${GPU_IDS[@]}"; do
    [[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
        echo "ERROR: invalid GPU id: $gpu" >&2
        exit 2
    }
    [[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
        echo "ERROR: duplicate GPU id: $gpu" >&2
        exit 2
    }
    SEEN_GPUS["$gpu"]=1
done

IFS=',' read -r -a SEEDS <<< "$SEED_CSV"
declare -A SEEN_SEEDS=()
for seed in "${SEEDS[@]}"; do
    [[ "$seed" =~ ^[0-9]+$ ]] || {
        echo "ERROR: invalid seed: $seed" >&2
        exit 2
    }
    [[ -z "${SEEN_SEEDS[$seed]+x}" ]] || {
        echo "ERROR: duplicate seed: $seed" >&2
        exit 2
    }
    SEEN_SEEDS["$seed"]=1
done

RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME"
LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME"
if ((!DRY_RUN)); then
    mkdir -p "$RUNS_ROOT" "$LOG_ROOT"
fi

declare -a JOB_NAMES=()
declare -a JOB_SEEDS=()
declare -a JOB_ENTRYPOINTS=()
declare -a JOB_ARCHITECTURES=()
declare -a JOB_TIME_MODES=()
declare -a JOB_DIST_MODES=()

add_job() {
    JOB_NAMES+=("$1")
    JOB_SEEDS+=("$2")
    JOB_ENTRYPOINTS+=("$3")
    JOB_ARCHITECTURES+=("$4")
    JOB_TIME_MODES+=("$5")
    JOB_DIST_MODES+=("$6")
}

for seed in "${SEEDS[@]}"; do
    add_job \
        "ffn_next_token_absolute_exponential" \
        "$seed" \
        "train_next_step.py" \
        "transformer_ffn_v1" \
        "absolute" \
        "exponential"

    add_job \
        "ffn_all_future_absolute_exponential" \
        "$seed" \
        "train_all_future.py" \
        "transformer_ffn_v1" \
        "absolute" \
        "exponential"

    add_job \
        "traj_mixer_all_future_absolute_exponential" \
        "$seed" \
        "train_all_future.py" \
        "traj_mixer_v5" \
        "absolute" \
        "exponential"

    add_job \
        "traj_mixer_all_future_relative_exponential" \
        "$seed" \
        "train_all_future.py" \
        "traj_mixer_v5" \
        "relative" \
        "exponential"

    add_job \
        "traj_mixer_all_future_relative_weibull" \
        "$seed" \
        "train_all_future.py" \
        "traj_mixer_v5" \
        "relative" \
        "weibull"

    add_job \
        "ffn_all_future_relative_weibull" \
        "$seed" \
        "train_all_future.py" \
        "transformer_ffn_v1" \
        "relative" \
        "weibull"
done

print_command() {
    printf '%q ' "$@"
    printf '\n'
}

run_job() {
    local job_index="$1"
    local gpu="$2"
    local job_name="${JOB_NAMES[$job_index]}"
    local seed="${JOB_SEEDS[$job_index]}"
    local entrypoint="${JOB_ENTRYPOINTS[$job_index]}"
    local architecture="${JOB_ARCHITECTURES[$job_index]}"
    local time_mode="${JOB_TIME_MODES[$job_index]}"
    local dist_mode="${JOB_DIST_MODES[$job_index]}"
    local seed_runs_root="$RUNS_ROOT/seed_$seed"
    local seed_log_root="$LOG_ROOT/seed_$seed"
    local log_file="$seed_log_root/$job_name.log"
    local -a command=(
        "$PYTHON_BIN"
        -u
        "$SCRIPT_DIR/$entrypoint"
        --runs_root "$seed_runs_root"
        --seed "$seed"
        --batch_size "$BATCH_SIZE"
        --extra_info_types_file "$EXTRA_INFO_TYPES_FILE"
        --model_architecture "$architecture"
        --num_workers "$NUM_WORKERS"
        --device cuda
    )

    if ((!DRY_RUN)); then
        mkdir -p "$seed_runs_root" "$seed_log_root"
    fi

    if [[ "$entrypoint" == "train_all_future.py" ]]; then
        command+=(
            --time_mode "$time_mode"
            --dist_mode "$dist_mode"
        )
    fi

    echo "[$(date '+%F %T')] START seed=$seed job=$job_name gpu=$gpu"
    echo "  log=$log_file"
    if ((DRY_RUN)); then
        printf '  CUDA_VISIBLE_DEVICES=%q ' "$gpu"
        print_command "${command[@]}"
        return 0
    fi

    if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
        "${command[@]}" >"$log_file" 2>&1; then
        echo "[$(date '+%F %T')] DONE  seed=$seed job=$job_name gpu=$gpu"
        return 0
    else
        local exit_code=$?
        echo "[$(date '+%F %T')] FAIL  seed=$seed job=$job_name gpu=$gpu exit=$exit_code" >&2
        echo "  See: $log_file" >&2
        return "$exit_code"
    fi
}

worker() {
    local slot="$1"
    local gpu="${GPU_IDS[$slot]}"
    local job_index
    local failed=0

    for ((job_index = slot; job_index < ${#JOB_NAMES[@]}; job_index += ${#GPU_IDS[@]})); do
        run_job "$job_index" "$gpu" || failed=1
    done
    return "$failed"
}

echo "Campaign: $CAMPAIGN_NAME"
echo "Additional seeds: ${SEEDS[*]}"
echo "GPUs: ${GPU_IDS[*]}"
echo "Configurations per seed: 6"
echo "Total tasks: ${#JOB_NAMES[@]}"
echo "Runs root: $RUNS_ROOT"
echo "Log root: $LOG_ROOT"
echo

declare -a WORKER_PIDS=()
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
    worker "$slot" &
    WORKER_PIDS+=("$!")
done

overall_status=0
for pid in "${WORKER_PIDS[@]}"; do
    wait "$pid" || overall_status=1
done

if ((overall_status != 0)); then
    echo "One or more training tasks failed. Inspect logs under: $LOG_ROOT" >&2
    exit 1
fi

if ((DRY_RUN)); then
    echo "Dry run completed successfully."
else
    echo "All training tasks completed successfully."
fi