#!/bin/bash
#
# mt_dao Centralised Training Script
# ==================================
#
# This script sets up and launches centralised training for machine learning models
# using the mt_dao framework. It handles environment setup, GPU detection, distributed
# training configuration, and process management.
#
# Usage:
#   ./base_centralised_training.sh [-p|--project_path PATH]
#
# Arguments:
#   -p, --project_path PATH  : Path to the mt_dao project directory
#                              (default: $HOME/projects/mt_dao)
#
# Environment Variables:
#   RUN_UUID                 : Unique identifier for this training run
#   MODEL_SIZE              : Size of the model being trained
#   SAVE_PATH               : Base path for saving outputs
#   mt_dao_SAVE_PATH        : mt_dao-specific save path
#   LLM_CONFIG              : LLM configuration parameters
#   EXTERNAL_CONFIGS        : External configuration parameters
#   PYTORCH_*               : PyTorch distributed training parameters
#
# shellcheck disable=SC2090,SC2086,SC2089,SC1091

set -euo pipefail # Exit on error, undefined variables, and pipe failures

# ============================================================================
# CONFIGURATION AND ARGUMENT PARSING
# ============================================================================

# Default configuration
readonly DEFAULT_PROJECT_PATH="$HOME/projects/mt_dao"
PROJECT_PATH="$DEFAULT_PROJECT_PATH"

# Parse command-line options
if ! OPTIONS=$(getopt -o p: --long project_path: -n 'parse-options' -- "$@"); then
	echo "base_centralised_training.sh: Error parsing options" >&2
	exit 1
fi

eval set -- "$OPTIONS"

# Process command-line arguments
while true; do
	case "$1" in
	-p | --project_path)
		PROJECT_PATH="$2"
		shift 2
		;;
	--)
		shift
		break
		;;
	*)
		break
		;;
	esac
done

echo "base_centralised_training.sh: PROJECT_PATH=$PROJECT_PATH"

# ============================================================================
# ENVIRONMENT SETUP
# ============================================================================

#! Moving to the project folder
cd "$PROJECT_PATH" || {
	echo "base_centralised_training.sh: Error: Cannot access project path: $PROJECT_PATH" >&2
	exit 1
}

#! Preparing environment based on execution context
if [[ $(hostname) == *'gpu-q'* ]]; then
	echo "base_centralised_training.sh: Detected CSD3 environment - using HPC setup"
	#! Executing the HPC environment preparation script
	#! NOTE: Must use "bash" to execute, "sh" doesn't work
	bash "$PROJECT_PATH/scripts/install_hpc_env.sh" -p "$PROJECT_PATH"
else
	echo "base_centralised_training.sh: Detected non-CSD3 environment - using standard setup"
	#! Executing the standard environment preparation script
	#! NOTE: Must use "bash" to execute, "sh" doesn't work
	bash "$PROJECT_PATH/scripts/install_env.sh" -p "$PROJECT_PATH"
fi

#! Export the endpoint of the S3 object store
export S3_ENDPOINT_URL='http://anon.anon:9000'

# ============================================================================
# RUN CONFIGURATION AND PATH SETUP
# ============================================================================

#! Set the date and time if not already set
if [ -z "${DATETIME:-}" ]; then
	DATETIME=$(date '+%Y%m%d_%H%M%S')
	export DATETIME
	echo "mt_dao_base.sh: Generated DATETIME: $DATETIME"
else
	echo "mt_dao_base.sh: Using existing DATETIME: $DATETIME"
fi

#! Set run UUID if not already provided
if [[ -z ${RUN_UUID:-} ]]; then
	export RUN_UUID="cen-${MODEL_SIZE:-unknown}-$DATETIME"
fi
echo "base_centralised_training.sh: RUN_UUID=$RUN_UUID"

#! Configure save paths
if [[ -z ${SAVE_PATH:-} ]]; then
	export SAVE_PATH="$PROJECT_PATH/$RUN_UUID"
	mkdir -p "$SAVE_PATH"
fi

if [[ -z ${mt_dao_SAVE_PATH:-} ]]; then
	export mt_dao_SAVE_PATH="$PROJECT_PATH/runs/$RUN_UUID"
fi

# Create necessary directories
mkdir -p "$mt_dao_SAVE_PATH"
mkdir -p "$mt_dao_SAVE_PATH/$DATETIME"

echo "base_centralised_training.sh: SAVE_PATH=$SAVE_PATH"
echo "base_centralised_training.sh: mt_dao_SAVE_PATH=$mt_dao_SAVE_PATH"

# ============================================================================
# GPU DETECTION AND CUDA CONFIGURATION
# ============================================================================
#! Detect available GPUs and configure CUDA devices
N_GPUS=$(uv run python -c 'import torch; print(torch.cuda.device_count())')
readonly N_GPUS

if [[ $N_GPUS -eq 0 ]]; then
	echo "base_centralised_training.sh: Warning: No GPUs detected. Training will run on CPU."
	CUDA_VISIBLE_DEVICES=""
else
	CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((N_GPUS - 1)))
	echo "base_centralised_training.sh: Detected $N_GPUS GPU(s)"
fi

echo "base_centralised_training.sh: CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"

# ============================================================================
# HYDRA CONFIGURATION AND RESOLVER
# ============================================================================

#! Configure LLM settings with run UUID
LLM_CONFIG="run_uuid=$RUN_UUID"

echo "base_centralised_training.sh: LLM_CONFIG=$LLM_CONFIG"
echo "base_centralised_training.sh: EXTERNAL_CONFIGS=${EXTERNAL_CONFIGS:-}"

#! Execute Hydra resolver to process configuration
echo "base_centralised_training.sh: Running Hydra resolver..."
HYDRA_FULL_ERROR=1 uv run python -m mt_dao.hydra_resolver \
	$LLM_CONFIG \
	${EXTERNAL_CONFIGS:-} \
	hydra/job_logging=none \
	hydra/hydra_logging=none \
	2>&1 | tee "$mt_dao_SAVE_PATH/$DATETIME/hydra_resolver.log"

# ============================================================================
# PYTORCH DISTRIBUTED TRAINING CONFIGURATION
# ============================================================================

#! Configure PyTorch distributed training parameters with sensible defaults
#! Set master address for distributed training (localhost for single-node)
if [[ -z ${PYTORCH_MASTER_ADDRESS:-} ]]; then
	export PYTORCH_MASTER_ADDRESS="127.0.0.1"
fi

#! Set node rank in multi-node setup (0 for single-node)
if [[ -z ${PYTORCH_NODE_RANK:-} ]]; then
	export PYTORCH_NODE_RANK="0"
fi

#! Set world size to match number of available GPUs
if [[ -z ${PYTORCH_WORLD_SIZE:-} ]]; then
	export PYTORCH_WORLD_SIZE=$N_GPUS
fi

#! Set base rank for this node (0 for single-node)
if [[ -z ${PYTORCH_BASE_RANK:-} ]]; then
	export PYTORCH_BASE_RANK=0
fi

#! Set master port, finding a free port if not specified
if [[ -z ${PYTORCH_MASTER_PORT:-} ]]; then
	PYTORCH_MASTER_PORT=$(uv run python -c "from mt_dao.port_utils import get_free_tcp_port; print(get_free_tcp_port())")
	export PYTORCH_MASTER_PORT
	echo "base_centralised_training.sh: Auto-assigned PYTORCH_MASTER_PORT=$PYTORCH_MASTER_PORT"
fi

# Display distributed training configuration
echo "base_centralised_training.sh: Distributed Training Configuration:"
echo "  PYTORCH_MASTER_ADDRESS=$PYTORCH_MASTER_ADDRESS"
echo "  PYTORCH_MASTER_PORT=$PYTORCH_MASTER_PORT"
echo "  PYTORCH_WORLD_SIZE=$PYTORCH_WORLD_SIZE"
echo "  PYTORCH_NODE_RANK=$PYTORCH_NODE_RANK"
echo "  PYTORCH_BASE_RANK=$PYTORCH_BASE_RANK"

# ============================================================================
# TRAINING EXECUTION
# ============================================================================

echo "base_centralised_training.sh: Starting centralised training..."

#! Launch centralised training script with optimized CUDA configuration
#! NOTE: Adding `NCCL_BLOCKING_WAIT=1` breaks the optimizer's checkpointing. We don't know why yet.
# TORCH_LOGS="+dynamo" TORCHDYNAMO_VERBOSE=1

# Configure environment variables for training
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export APPOINTED_CUDA_DEVICE=$CUDA_VISIBLE_DEVICES
export CUDA_LAUNCH_BLOCKING=1
export HYDRA_FULL_ERROR=1

# Execute training with composer
uv run composer \
	--world_size "$PYTORCH_WORLD_SIZE" \
	--base_rank "$PYTORCH_BASE_RANK" \
	--node_rank "$PYTORCH_NODE_RANK" \
	--master_addr "$PYTORCH_MASTER_ADDRESS" \
	--master_port "$PYTORCH_MASTER_PORT" \
	"$PROJECT_PATH/mt_dao/centralised_train.py" \
	2>&1 | tee "$mt_dao_SAVE_PATH/$DATETIME/centralised_train.log" &

#! Keep the process ID and wait for training to complete
readonly TRAINING_PID=$!
echo "base_centralised_training.sh: Training started with PID $TRAINING_PID"

# Wait for training process to complete
wait $TRAINING_PID
training_exit_code=$?

# Report training completion status
if [[ $training_exit_code -eq 0 ]]; then
	echo "base_centralised_training.sh: Training completed successfully"
else
	echo "base_centralised_training.sh: Training failed with exit code $training_exit_code" >&2
fi

echo "base_centralised_training.sh: Script completed at $(date)"
exit $training_exit_code
