#!/bin/bash
# shellcheck disable=SC2086,SC2004,SC2124,SC2206,SC2046,SC2181,SC2207,SC2155

#===============================================================================
# UNIFIED FEDERATED HYPERPARAMETER SWEEP SCRIPT (v6.2)
#===============================================================================
#
# This script launches a comprehensive hyperparameter grid sweep for federated
# learning experiments. It combines the detailed sweep control of a centralized
# training script with the federated execution logic for multi-node setups.
#
# Key Features:
#   - Federated Learning: Designed to run FL experiments with strategies like FedAvg/FedAdamW.
#   - Advanced Sweep Control: Sweeps over seeds, batch sizes, custom
#     optimizer parameter triplets, and federated sync frequencies.
#   - SLURM & Local Execution: Supports both SLURM-based job submission
#     with dependency chaining and sequential local execution.
#   - Range-Based Execution: Use '--range START-END' to split large sweeps.
#   - Dry Run Mode: Preview job configurations without execution via '--dry-run'.
#   - Sweep Hashing: Generates a unique ID for the sweep for clear identification.
#   - Safe Port Allocation: Uses a file lock to prevent port conflicts during
#     concurrent launches on the same node.
#
#===============================================================================

# --- 1. SWEEP CONTROL FLAGS ---
SWEEP_SEEDS=false
SWEEP_BSS=false
SWEEP_QH_PARAMS=false
SWEEP_SYNC_FREQS=false

# --- 2. EXECUTION AND SYSTEM CONFIGURATION ---
USE_SLURM=true
MAX_CONCURRENT_JOBS=1
TARGET_NODE="anon"
GPUS_PER_JOB=4
CPUS_PER_JOB=8
DRY_RUN=false

# Default server and port values (fallback)
DEFAULT_SERVER_IP="127.0.0.1"
DEFAULT_SUPERLINK_PORT=48151
DEFAULT_EXEC_API_PORT=48152
DEFAULT_SERVERAPP_API_PORT=48153

# Paths
PROJECT_PATH="$HOME/projects/mt_dao"
DATASET_CACHE_DIR="/nfs-share/datasets/mt_dao/dataset_cache"
PRE_TRAINED_MODEL_PATH=null

if [[ $USE_SLURM == false ]]; then
	DATASET_CACHE_DIR="/home/ubuntu/datasets/mt_dao/dataset_cache"
fi

# Default run range (0 to infinity, inclusive)
START_INDEX=0
END_INDEX=-1

# --- 3. HYPERPARAMETER GRID DEFINITION ---
declare -A LLM_CONFIGS
LLM_CONFIGS=(
	["qh_fl_long"]="smollm-3m"
)

# Base values
BASE_SEED=42
BASE_BS=16
BASE_LR="0.00390625"

# Sweep arrays
SEEDS=(42 1337)
BSS=(16 32)

# Sweep triplets: "new_v new_beta switch_scale"
# These parameters are applied by a callback after the warmup period.
BASE_QH_TRIPLET="0.98 0.999 2.8284271258"
QH_TRIPLETS=(
	"0.95 0.99 2.0"
	"0.95 0.999 2.0"
)

# Sweep quadruplets for sync frequencies: "BASE PARAMETERS EXP_AVG EXP_AVG_SQ"
BASE_SYNC_QUAD="256 1 1 1"
SYNC_QUADS=(
	"32 1 1 1"
	"32 2 1 1"
	"32 4 1 1"
	"32 8 1 1"
	"32 16 1 1"
	"32 1 2 1"
	"32 1 4 1"
	"32 1 8 1"
	"32 1 16 1"
	"8 1 1 1"
	"16 1 1 1"
)

# --- 4. FIXED & SCALING FEDERATED PARAMETERS ---
D_MODEL=256
BASE_MAX_DURATION=40960 # Total client steps at BASE_BS
BASE_T_WARMUP=2049
BASE_T_COOLDOWN=2048
FIXED_INIT_STD="0.02"
FIXED_WEIGHT_DECAY="0.0"

# Core FL Parameters
N_CLIENTS=4
N_CLIENTS_PER_ROUND=4

# --- 5. STATIC PARAMETERS FOR ALL RUNS ---
STATIC_PARAMETERS=""
# General & Precision
STATIC_PARAMETERS+=" use_wandb=true is_test=false"
STATIC_PARAMETERS+=" llm_config.precision=amp_bf16"
# Core Model
STATIC_PARAMETERS+=" llm_config.model.d_model=${D_MODEL}"
STATIC_PARAMETERS+=" llm_config.model.n_heads=4"
STATIC_PARAMETERS+=" llm_config.model.n_layers=4"
STATIC_PARAMETERS+=" llm_config.model.max_seq_len=2048"
STATIC_PARAMETERS+=" ++llm_config.model.no_bias=True"
STATIC_PARAMETERS+=" ++llm_config.model.vocab_size=50368"
STATIC_PARAMETERS+=" +llm_config.model.use_cache=true"
# Attention & RoPE
STATIC_PARAMETERS+=" llm_config.model.attn_config.attn_impl=flash"
STATIC_PARAMETERS+=" llm_config.model.attn_config.rope=true"
STATIC_PARAMETERS+=" +llm_config.model.attn_config.rope_impl=dail"
STATIC_PARAMETERS+=" llm_config.model.attn_config.rope_theta=10000"
# FFN Base Config
STATIC_PARAMETERS+=" llm_config.model.ffn_config.ffn_act_fn.name=silu"
STATIC_PARAMETERS+=" llm_config.model.ffn_hidden_size=1024" # 4 * d_model
# Init, muP, and complete-P
STATIC_PARAMETERS+=" llm_config.model.init_config.name=baseline_"
STATIC_PARAMETERS+=" llm_config.model.init_config.init_std=${FIXED_INIT_STD}"
STATIC_PARAMETERS+=" llm_config.model.init_config.emb_init_std=${FIXED_INIT_STD}"
STATIC_PARAMETERS+=" ++llm_config.model.use_peri_norm=True"
STATIC_PARAMETERS+=" ++llm_config.model.use_embedding_norm=True"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_enabled=true"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_disable_attention_scaling=true"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_disable_hidden_lr_scaling=false"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_width_multiplier=1.0"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_input_alpha=1.0"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_output_alpha=1.0"
STATIC_PARAMETERS+=" llm_config.model.mup_config.completep_depth_alpha_enabled=true"
STATIC_PARAMETERS+=" llm_config.model.mup_config.completep_depth_multiplier=1.0"
STATIC_PARAMETERS+=" llm_config.model.mup_config.completep_depth_alpha_exp=1.0"
STATIC_PARAMETERS+=" llm_config.model.mup_config.completep_eps_scaling_enabled=true"
STATIC_PARAMETERS+=" llm_config.model.mup_config.mup_scale_expert_sel=true"
# Client-side Optimizer & Scheduler
STATIC_PARAMETERS+=" llm_config.optimizer.name=qhadopt"
STATIC_PARAMETERS+=" llm_config.optimizer.weight_decay=${FIXED_WEIGHT_DECAY}"
STATIC_PARAMETERS+=" llm_config.optimizer.betas=[0.95,0.9999]"
STATIC_PARAMETERS+=" llm_config.optimizer.eps=1e-6"
STATIC_PARAMETERS+=" ++llm_config.optimizer.v1=1.0"
STATIC_PARAMETERS+=" llm_config.scheduler.schedulers.lr.name=constant_with_sqrt_cooldown_with_warmup_and_switch"
# Dataset & Tokenizer
STATIC_PARAMETERS+=" dataset=smollm-corpus-shared-with-val"
STATIC_PARAMETERS+=" dataset.train.root_local=${DATASET_CACHE_DIR}/smollm-corpus-shared"
STATIC_PARAMETERS+=" dataset.val.root_local=${DATASET_CACHE_DIR}/smollm-corpus-shared-val"
STATIC_PARAMETERS+=" dataset/streams@dataset.train.streams=smollm_corpus_4_clients_proportional"
STATIC_PARAMETERS+=" dataset/streams@dataset.val.streams=smollm_corpus_4_clients_proportional"
STATIC_PARAMETERS+=" centralized.stream_id=null"
STATIC_PARAMETERS+=" llm_config.tokenizer.name='HuggingFaceTB/SmolLM-1.7B'"
# Training & Device
STATIC_PARAMETERS+=" llm_config.device_train_microbatch_size=auto"
STATIC_PARAMETERS+=" +llm_config.device_eval_microbatch_size=auto"
STATIC_PARAMETERS+=" ++llm_config.compile_config={}"
STATIC_PARAMETERS+=" ~llm_config.fsdp_config"
# Evaluation
STATIC_PARAMETERS+=" eval_gauntlet_config=empty"
STATIC_PARAMETERS+=" eval_gauntlet_config.destination_dir=${DATASET_CACHE_DIR}/eval"
STATIC_PARAMETERS+=" icl_tasks_config=empty"
STATIC_PARAMETERS+=" icl_tasks_config.root_dir=${DATASET_CACHE_DIR}"
STATIC_PARAMETERS+=" llm_config.eval_subset_num_batches=128"
# Callbacks & Logging
STATIC_PARAMETERS+=" ~llm_config.loggers.tensorboard"
STATIC_PARAMETERS+=" llm_config.console_log_interval=64ba"
STATIC_PARAMETERS+=" ++llm_config.callbacks.quasi_hyper_switch_callback={}"
STATIC_PARAMETERS+=" ++llm_config.callbacks.activation_monitor_full_model={}"
STATIC_PARAMETERS+=" llm_config.callbacks.activation_monitor_full_model.interval=16ba"
STATIC_PARAMETERS+=" ++llm_config.callbacks.optimizer_monitor={}"
STATIC_PARAMETERS+=" llm_config.callbacks.optimizer_monitor.interval=16ba"
STATIC_PARAMETERS+=" ++llm_config.callbacks.betas_monitor={}"
STATIC_PARAMETERS+=" ++llm_config.callbacks.quasi_hyperbolic_parameter_monitor={}"
# Federated Learning Strategy
STATIC_PARAMETERS+=" fl.strategy_name=FEDAVG"
STATIC_PARAMETERS+=" fl.eval_period=null fl.reset_optimizer=false llm_config.eval_loader=null"
# STATIC_PARAMETERS+=" fl.delta_modes=[all_params]"
STATIC_PARAMETERS+=" fl.n_total_clients=${N_CLIENTS}"
STATIC_PARAMETERS+=" fl.n_clients_per_round=${N_CLIENTS_PER_ROUND}"

# Parameter/optimizer state synchronization schedule
STATIC_PARAMETERS+=" fl.parameter_scheduler_name=FREQ"

STATIC_PARAMETERS+=" llm_config.save_num_checkpoints_to_keep=-1"

STATIC_PARAMETERS+=" mt_dao.checkpointing.local_checkpoint=true"
STATIC_PARAMETERS+=" mt_dao.checkpointing.remote_checkpoint=true"

STATIC_PARAMETERS+=" mt_dao.restore_cent.run_uuid=qhadopt_long_8a8d1799_0-20250911_122145"

STATIC_PARAMETERS+=" mt_dao.restore_cent.run_batches=2048"
STATIC_PARAMETERS+=" mt_dao.restore_cent.checkpoint_dir=s3://anon/qhadopt_long_8a8d1799_0-20250911_122145"

# STATIC_PARAMETERS+=" mt_dao.comm_stack.shm=true"

# --- 6. SCRIPT LOGIC ---

print_run_summary() {
	local job_name=$1
	local run_index=$2
	local total_jobs=$3
	local seed=$4
	local bs=$5
	local new_v=$6
	local new_beta=$7
	local switch_scale=$8
	local base_freq=$9
	local param_freq=${10}
	local exp_avg_freq=${11}
	local exp_avg_sq_freq=${12}
	local max_dur=${13}
	local n_rounds=${14}
	local warmup=${15}
	local cooldown=${16}

	echo "----------------------------------------------------"
	echo "Preparing Job: ${job_name}"
	echo "Sweep Index:   ${run_index} / $((total_jobs - 1))"
	echo "===================================================="
	echo "  >> FEDERATED RUN CONFIGURATION << "
	echo "----------------------------------------------------"
	echo "  Core Hyperparameters:"
	echo "    Seed:               ${seed}"
	echo "    Global Batch Size:  ${bs}"
	echo "    Base Client LR:     ${BASE_LR}"
	echo "  Swept QH Parameters (post-warmup):"
	echo "    New v:              ${new_v}"
	echo "    New beta_2:         ${new_beta}"
	echo "    LR Switch Scale:    ${switch_scale}"
	echo "  Sync Frequencies (Client -> Server):"
	echo "    BASE:               ${base_freq}"
	echo "    PARAMETERS:         ${param_freq}"
	echo "    EXP_AVG:            ${exp_avg_freq}"
	echo "    EXP_AVG_SQ:         ${exp_avg_sq_freq}"
	echo "  Duration (scaled):"
	echo "    Total Client Steps: ${max_dur}ba"
	echo "    Warmup:             ${warmup}ba"
	echo "    Cooldown:           ${cooldown}ba"
	echo "  Federated Learning:"
	echo "    Local Steps:        ${FL_LOCAL_STEPS}"
	echo "    Num Rounds:         ${n_rounds}"
	echo "    Clients per Round:  ${N_CLIENTS_PER_ROUND}"
	echo "----------------------------------------------------"
}

# Argument Parsing
while (("$#")); do
	case "$1" in
	--range)
		if [[ $2 =~ ^[0-9]+-[0-9]+$ ]]; then
			START_INDEX=$(echo "$2" | cut -d'-' -f1)
			END_INDEX=$(echo "$2" | cut -d'-' -f2)
		else
			echo "Error: --range requires format START-END (e.g., 0-99)" >&2
			exit 1
		fi
		shift 2
		;;
	--no-slurm)
		USE_SLURM=false
		shift
		;;
	--dry-run)
		DRY_RUN=true
		shift
		;;
	-h | --help)
		echo "Usage: $0 [OPTIONS]"
		echo "Options:"
		echo "  --range START-END   Execute jobs in this 0-indexed, inclusive range (e.g., 0-49)."
		echo "  --no-slurm          Run jobs sequentially on the local machine."
		echo "  --dry-run           Print job configurations without executing."
		echo "  -h, --help          Show this help message."
		exit 0
		;;
	*)
		echo "Error: Unsupported flag $1" >&2
		exit 1
		;;
	esac
done

if [[ $USE_SLURM == false ]]; then
	DATASET_CACHE_DIR="/home/ubuntu/datasets/mt_dao/dataset_cache"
fi


# Conditionally define arrays to iterate over
if [[ $SWEEP_SEEDS == true ]]; then SEEDS_TO_SWEEP=("${SEEDS[@]}"); else SEEDS_TO_SWEEP=(${BASE_SEED}); fi
if [[ $SWEEP_BSS == true ]]; then BSS_TO_SWEEP=("${BSS[@]}"); else BSS_TO_SWEEP=(${BASE_BS}); fi
if [[ $SWEEP_QH_PARAMS == true ]]; then QH_TRIPLETS_TO_SWEEP=("${QH_TRIPLETS[@]}"); else QH_TRIPLETS_TO_SWEEP=("${BASE_QH_TRIPLET}"); fi
if [[ $SWEEP_SYNC_FREQS == true ]]; then SYNC_QUADS_TO_SWEEP=("${SYNC_QUADS[@]}"); else SYNC_QUADS_TO_SWEEP=("${BASE_SYNC_QUAD}"); fi

# --- Sweep Identification and Summary ---
TOTAL_JOBS=$((${#LLM_CONFIGS[@]} * ${#SEEDS_TO_SWEEP[@]} * ${#BSS_TO_SWEEP[@]} * ${#QH_TRIPLETS_TO_SWEEP[@]} * ${#SYNC_QUADS_TO_SWEEP[@]}))
SWEEP_CONFIG_STRING="SEEDS:${SWEEP_SEEDS}-(${SEEDS[*]});;BSS:${SWEEP_BSS}-(${BSS[*]});;QH_PARAMS:${SWEEP_QH_PARAMS}-(${QH_TRIPLETS[*]});;SYNC_FREQS:${SWEEP_SYNC_FREQS}-(${SYNC_QUADS[*]})"
SWEEP_HASH=$(echo -n "$SWEEP_CONFIG_STRING" | sha1sum | cut -c1-8)

echo "🚀 Starting Federated Hyperparameter Sweep..."
echo "Sweep Configuration ID (Hash): ${SWEEP_HASH}"
if ((TOTAL_JOBS > 0)); then
	echo "Total potential jobs in this sweep: ${TOTAL_JOBS} (indices 0 to $((TOTAL_JOBS - 1)))"
else
	echo "Total potential jobs in this sweep: 0"
fi
if ((END_INDEX != -1)); then
	echo "Executing jobs from index ${START_INDEX} to ${END_INDEX}"
else
	echo "Executing jobs from index ${START_INDEX} to end"
fi
if [[ $USE_SLURM == false ]]; then echo "Running in local mode (--no-slurm)."; fi
if [[ $DRY_RUN == true ]]; then echo "Running in dry-run mode."; fi
mkdir -p "$DATASET_CACHE_DIR"

# --- Main Submission Loop ---
RUN_INDEX=-1
SUBMITTED_JOB_COUNT=0
declare -a CHAIN_LAST_JOB_IDS
for ((i = 0; i < MAX_CONCURRENT_JOBS; i++)); do CHAIN_LAST_JOB_IDS[i]=""; done

for short_name in "${!LLM_CONFIGS[@]}"; do
	llm_config=${LLM_CONFIGS[$short_name]}
	for seed in "${SEEDS_TO_SWEEP[@]}"; do
		for bs in "${BSS_TO_SWEEP[@]}"; do
			# Scale durations based on batch size relative to the base batch size
			max_duration=$(printf "%.0f\n" $(echo "scale=4; ${BASE_MAX_DURATION} * ${BASE_BS} / ${bs}" | bc))
			t_warmup=$(printf "%.0f\n" $(echo "scale=4; ${BASE_T_WARMUP} * ${BASE_BS} / ${bs}" | bc))
			t_cooldown=$(printf "%.0f\n" $(echo "scale=4; ${BASE_T_COOLDOWN} * ${BASE_BS} / ${bs}" | bc))

			for qh_triplet in "${QH_TRIPLETS_TO_SWEEP[@]}"; do
				for sync_quad in "${SYNC_QUADS_TO_SWEEP[@]}"; do
					RUN_INDEX=$((RUN_INDEX + 1))
					if ((RUN_INDEX < START_INDEX || (END_INDEX != -1 && RUN_INDEX > END_INDEX))); then continue; fi

					read -r new_v new_beta switch_scale <<<"$qh_triplet"
					read -r base_freq param_freq exp_avg_freq exp_avg_sq_freq <<<"$sync_quad"

					SUBMITTED_JOB_COUNT=$((SUBMITTED_JOB_COUNT + 1))
					JOB_NAME=$(printf "%s_%s_%d" "${short_name}" "${SWEEP_HASH}" "${RUN_INDEX}")
					DATETIME=$(date '+%Y%m%d_%H%M%S')
					RUN_UUID="${JOB_NAME}-${DATETIME}"
					SAVE_PATH="${PROJECT_PATH}/outputs/${RUN_UUID}"

					# --- Construct the full parameter string for this run ---
					run_params="${STATIC_PARAMETERS}"
					run_params+=" mt_dao.n_nodes=${GPUS_PER_JOB}" # For FL, nodes correspond to GPUs
					run_params+=" llm_config=${llm_config}"
					run_params+=" run_uuid=${RUN_UUID}"
					run_params+=" pretrained_model_path=${PRE_TRAINED_MODEL_PATH}"
					# Append sweep parameters
					run_params+=" llm_config.global_seed=${seed}"
					run_params+=" llm_config.global_train_batch_size=${bs}"
					run_params+=" llm_config.optimizer.lr=${BASE_LR}"
					# Append scaled duration & FL round parameters
					FL_LOCAL_STEPS=${base_freq}
					fl_n_rounds=$((max_duration / FL_LOCAL_STEPS))
					[[ $fl_n_rounds -eq 0 ]] && fl_n_rounds=1

					run_params+=" fl.n_local_steps=${FL_LOCAL_STEPS}"
					run_params+=" fl.n_rounds=${fl_n_rounds}"
					run_params+=" llm_config.max_duration=${max_duration}ba"
					run_params+=" llm_config.scheduler.schedulers.lr.t_warmup=${t_warmup}ba"
					run_params+=" llm_config.scheduler.schedulers.lr.t_max=${max_duration}ba"
					run_params+=" llm_config.scheduler.schedulers.lr.t_cooldown=${t_cooldown}ba"
					# Append quasi-hyperbolic sweep parameters from the triplet
					run_params+=" ++llm_config.scheduler.schedulers.lr.switch_scale=${switch_scale}"
					run_params+=" ++llm_config.callbacks.quasi_hyper_switch_callback.t_switch=${t_warmup}ba"
					run_params+=" ++llm_config.callbacks.quasi_hyper_switch_callback.new_vs=[${new_v}]"
					run_params+=" ++llm_config.callbacks.quasi_hyper_switch_callback.new_betas=[${new_beta},0.9999]"
					# Append sync frequency parameters from the quadruplet

					run_params+=" ++fl.parameter_scheduler_kwargs.PARAMETERS=${param_freq}"
					run_params+=" ++fl.parameter_scheduler_kwargs.EXP_AVG=${exp_avg_freq}"
					run_params+=" ++fl.parameter_scheduler_kwargs.EXP_AVG_SQ=${exp_avg_sq_freq}"
					# Append save intervals
					run_params+=" llm_config.save_folder=${SAVE_PATH}"
					run_params+=" llm_config.save_interval=${max_duration}ba"
					run_params+=" llm_config.eval_interval=${max_duration}ba"

					print_run_summary "${JOB_NAME}" "${RUN_INDEX}" "${TOTAL_JOBS}" "${seed}" "${bs}" "${new_v}" "${new_beta}" "${switch_scale}" "${base_freq}" "${param_freq}" "${exp_avg_freq}" "${exp_avg_sq_freq}" "${max_duration}" "${fl_n_rounds}" "${t_warmup}" "${t_cooldown}"

					if [ "$DRY_RUN" = true ]; then
						echo "  [DRY RUN] Would execute with params: ${run_params}"
						echo "----------------------------------------------------"
						continue
					fi

					if [[ $USE_SLURM == true ]]; then
						chain_index=$(((SUBMITTED_JOB_COUNT - 1) % MAX_CONCURRENT_JOBS))
						last_job_in_chain=${CHAIN_LAST_JOB_IDS[$chain_index]}
						SBATCH_OPTS="--parsable"
						if [[ -n $last_job_in_chain ]]; then
							SBATCH_OPTS+=" --dependency=afterok:${last_job_in_chain}"
							echo "Submitting to SLURM (chain ${chain_index}, depends on Job ID: ${last_job_in_chain})"
						else
							echo "Submitting to SLURM (starting chain ${chain_index})"
						fi

						job_id=$(
							sbatch ${SBATCH_OPTS} <<EOF
#!/bin/bash
#SBATCH -c ${CPUS_PER_JOB}
#SBATCH -w ${TARGET_NODE}
#SBATCH --gres=gpu:${GPUS_PER_JOB}
#SBATCH --job-name=${JOB_NAME}
#SBATCH --output=%j-${JOB_NAME}.out
#SBATCH --error=%j-${JOB_NAME}.err
#SBATCH --tasks-per-node=1
echo "==================================================================="
echo "STARTING SLURM JOB: ${JOB_NAME} (ID: \$SLURM_JOB_ID)"
echo "Sweep Index: ${RUN_INDEX} / $((TOTAL_JOBS - 1))"
echo "Node: \$(hostname) at \$(date)"
echo "==================================================================="
export PROJECT_PATH="${PROJECT_PATH}"
export RUN_UUID="${RUN_UUID}"
export DATASET_CACHE_DIR="${DATASET_CACHE_DIR}"
LOCKFILE="/tmp/mt_dao_port_finder.lock"
PORT_ENV_FILE="/tmp/ports_\${SLURM_JOB_ID}.env"
(
    flock 200
    # Attempt to dynamically find server IP and free ports
    MASTER_IP=\$(hostname -I | awk '{print \$1}')
    SUPERLINK_PORT=\$(uv run python -c "from mt_dao.port_utils import get_free_tcp_port; print(get_free_tcp_port())")
    EXEC_API_PORT=\$(uv run python -c "from mt_dao.port_utils import get_free_tcp_port; print(get_free_tcp_port(['\${SUPERLINK_PORT}']))")
    SERVERAPP_API_PORT=\$(uv run python -c "from mt_dao.port_utils import get_free_tcp_port; print(get_free_tcp_port(['\${SUPERLINK_PORT}', '\${EXEC_API_PORT}']))")

    # Check for empty values and fall back to defaults if necessary
    if [[ -z "\$MASTER_IP" ]]; then
        echo "Warning: Could not determine MASTER_IP. Falling back to default: ${DEFAULT_SERVER_IP}"
        MASTER_IP="${DEFAULT_SERVER_IP}"
    fi
    if [[ -z "\$SUPERLINK_PORT" ]]; then
        echo "Warning: Could not find free Superlink port. Falling back to default: ${DEFAULT_SUPERLINK_PORT}"
        SUPERLINK_PORT="${DEFAULT_SUPERLINK_PORT}"
    fi
    if [[ -z "\$EXEC_API_PORT" ]]; then
        echo "Warning: Could not find free Exec API port. Falling back to default: ${DEFAULT_EXEC_API_PORT}"
        EXEC_API_PORT="${DEFAULT_EXEC_API_PORT}"
    fi
    if [[ -z "\$SERVERAPP_API_PORT" ]]; then
        echo "Warning: Could not find free ServerApp API port. Falling back to default: ${DEFAULT_SERVERAPP_API_PORT}"
        SERVERAPP_API_PORT="${DEFAULT_SERVERAPP_API_PORT}"
    fi

    # Export the final determined addresses to a temp file
    echo "export SUPERLINK_ADDRESS=\${MASTER_IP}:\${SUPERLINK_PORT}" > "\$PORT_ENV_FILE"
    echo "export EXEC_API_ADDRESS=\${MASTER_IP}:\${EXEC_API_PORT}" >> "\$PORT_ENV_FILE"
    echo "export SERVERAPP_API_ADDRESS=\${MASTER_IP}:\${SERVERAPP_API_PORT}" >> "\$PORT_ENV_FILE"
) 200>\$LOCKFILE
source "\$PORT_ENV_FILE" && rm "\$PORT_ENV_FILE"
echo "Superlink Address: \$SUPERLINK_ADDRESS"
export EXTERNAL_CONFIGS="${run_params}"
export CLIENTAPP_LAUNCHER_INDEPENDENT="true"
export CUDA_LAUNCH_BLOCKING=1
export CLIENTAPP_NO_RELOAD=1
bash "\${PROJECT_PATH}/scripts/base_launchers/mt_dao_base.sh" -p "\${PROJECT_PATH}"
echo "JOB FINISHED: \$(date)"
EOF
						)
						if [[ $? -eq 0 && -n $job_id ]]; then
							CHAIN_LAST_JOB_IDS[$chain_index]=$job_id
							echo "  Successfully submitted. Job ID: ${job_id}"
						else
							echo "❌ Error submitting job ${JOB_NAME}. The dependency chain is now broken. Aborting."
							exit 1
						fi
						sleep 30
					else # Local Execution
						echo "Executing locally..."
						pkill -9 -f flower &>/dev/null
						pkill -9 -f "mt_dao_base.sh" &>/dev/null

						export PROJECT_PATH="${PROJECT_PATH}"
						export RUN_UUID="${RUN_UUID}"
						export DATASET_CACHE_DIR="${DATASET_CACHE_DIR}"

						MASTER_IP=$(hostname -I | awk '{print $1}')
						SUPERLINK_PORT=$(uv run python -c "from mt_dao.port_utils import get_free_tcp_port; print(get_free_tcp_port())")

						# Check for empty values and fall back to defaults if necessary
						if [[ -z $MASTER_IP ]]; then
							echo "Warning: Could not determine MASTER_IP. Falling back to default: ${DEFAULT_SERVER_IP}"
							MASTER_IP="${DEFAULT_SERVER_IP}"
						fi
						if [[ -z $SUPERLINK_PORT ]]; then
							echo "Warning: Could not find free Superlink port. Falling back to default: ${DEFAULT_SUPERLINK_PORT}"
							SUPERLINK_PORT="${DEFAULT_SUPERLINK_PORT}"
						fi

						export SUPERLINK_ADDRESS="${MASTER_IP}:${SUPERLINK_PORT}"
						echo "Superlink Address: $SUPERLINK_ADDRESS"

						export EXTERNAL_CONFIGS="${run_params}"
						export CLIENTAPP_LAUNCHER_INDEPENDENT="true"
						export CUDA_LAUNCH_BLOCKING=1
						export CLIENTAPP_NO_RELOAD=1
						TRAIN_SCRIPT_PATH="${PROJECT_PATH}/scripts/base_launchers/mt_dao_base.sh"

						if [ ! -f "$TRAIN_SCRIPT_PATH" ]; then
							echo "❌ Training script not found at ${TRAIN_SCRIPT_PATH}. Aborting."
							exit 1
						fi
						bash "$TRAIN_SCRIPT_PATH" -p "${PROJECT_PATH}"
						if [[ $? -ne 0 ]]; then
							echo "❌ Local execution of job ${JOB_NAME} failed. Aborting."
							exit 1
						fi
						echo "  Local execution finished."
					fi
				done
			done
		done
	done
done

echo "===================================================="
if ((SUBMITTED_JOB_COUNT == 0)); then
	echo "No jobs were submitted. Please check your run range and sweep flags. 🤔"
else
	if [[ $DRY_RUN == true ]]; then
		echo "Dry run complete. Displayed configurations for ${SUBMITTED_JOB_COUNT} jobs. ✅"
	elif [[ $USE_SLURM == true ]]; then
		echo "All ${SUBMITTED_JOB_COUNT} jobs have been submitted to SLURM. ✅"
	else
		echo "All ${SUBMITTED_JOB_COUNT} jobs have been executed sequentially. ✅"
	fi
fi
echo "===================================================="
