#!/bin/bash
# shellcheck disable=SC2090,SC2086,SC2089,SC1091
#SBATCH -c 20
#SBATCH -w anon
#SBATCH --gres=gpu:2
#SBATCH --job-name=mt_dao_example
#SBATCH --tasks-per-node=1
#SBATCH --output=%x-%j.out
#SBATCH --time=01:00:00
#SBATCH --dependency=afterany:4871

#! Default project path
PROJECT_PATH="$HOME/projects/mt_dao"
#! Get the current date and time
DATETIME=$(date '+%Y%m%d_%H%M%S')

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

eval set -- "$OPTIONS"

while true; do
	case "$1" in
	-p | --project_path)
		PROJECT_PATH="$2"
		shift 2
		;;
	--)
		shift
		break
		;;
	*)
		break
		;;
	esac
done
echo "mt_dao_example.sh: PROJECT_PATH=$PROJECT_PATH"

RAY_TEMP_DIR="$PROJECT_PATH/ray"
# RAY_TEMP_DIR="$HOME/ray"

# rm -rf /dev/shm/*
# rm -rf /tmp/tmp*
rm -rf "${RAY_TEMP_DIR:?}"/*

#! Setting general ML parameters
COMM_BATCHES=1
USE_RAY=false
USE_SHM=true
USE_S3=false
LOCAL_BATCH_SIZE=32
LOCAL_STEPS=10
N_CLIENTS=4
TOTAL_STEPS=5120
WARMUP_STEPS=$(uv run python -c "print(int(0.05*$TOTAL_STEPS))")  # 5% of the total steps
COOLDOWN_STEPS=$(uv run python -c "print(int(0.2*$TOTAL_STEPS))") # 20% of the total steps
COOLDOWN_STEPS=0
N_ROUNDS=$((TOTAL_STEPS / (LOCAL_STEPS)))
FL_EVAL_PERIOD=null
LOCAL_EVAL_INTERVAL=$LOCAL_STEPS

#! Set and export the run UUID and the save path
export RUN_UUID="test-$DATETIME"
# export RUN_UUID="fed-125M-example-20250501_134201b"
# export SAVE_PATH="s3://anon/$RUN_UUID"
export SAVE_PATH="$PROJECT_PATH/runs/$RUN_UUID/$DATETIME"

#! Initialize the external configs
EXTERNAL_CONFIGS=""
EXTERNAL_CONFIGS="llm_config=deptoe" # This is the smallest model we have, so it is used for testing purposes
# EXTERNAL_CONFIGS="llm_config=mpt-16m" # This is the smallest model we have, so it is used for testing purposes
# EXTERNAL_CONFIGS="llm_config=mpt-125m"

#! Dataset configs
DATASET_NAME="smollm-corpus-shared"
export DATASET_CACHE_DIR="/nfs-share/datasets/mt_dao/dataset_cache"
# export DATASET_CACHE_DIR="$PROJECT_PATH/dataset_cache"
mkdir -p $DATASET_CACHE_DIR
EXTERNAL_CONFIGS+=" dataset=$DATASET_NAME"                                                      # Dataset name
EXTERNAL_CONFIGS+=" dataset.train.root_local=$DATASET_CACHE_DIR/$DATASET_NAME"                  # Path of the local cache for the training dataset
EXTERNAL_CONFIGS+=" dataset.val.root_local=$DATASET_CACHE_DIR/$DATASET_NAME"                    # Path of the local cache for the evaluation dataset
EXTERNAL_CONFIGS+=" dataset/streams@dataset.train.streams=smollm_corpus_4_clients_proportional" # Stream configuration for the training dataset
EXTERNAL_CONFIGS+=" dataset/streams@dataset.val.streams=smollm_corpus_4_clients_proportional"   # Stream configuration for the training dataset
EXTERNAL_CONFIGS+=" centralized.stream_id=null"
EXTERNAL_CONFIGS+=" llm_config.eval_loader=null"
EXTERNAL_CONFIGS+=" llm_config.tokenizer.name='HuggingFaceTB/SmolLM-1.7B'"
# EXTERNAL_CONFIGS+=" llm_config.model.vocab_size=50368"

#! Local optimizer configs
EXTERNAL_CONFIGS+=" llm_config.max_duration=${TOTAL_STEPS}ba"
# EXTERNAL_CONFIGS+=" llm_config.optimizer.name=adopt"
# EXTERNAL_CONFIGS+=" llm_config.optimizer.lr=1.0e-4"
EXTERNAL_CONFIGS+=" llm_config.scheduler.schedulers.lr.name=constant_with_sqrt_cooldown_with_warmup"
EXTERNAL_CONFIGS+=" llm_config.scheduler.schedulers.lr.t_max=${TOTAL_STEPS}ba"
EXTERNAL_CONFIGS+=" llm_config.scheduler.schedulers.lr.t_warmup=${WARMUP_STEPS}ba"
EXTERNAL_CONFIGS+=" ++llm_config.scheduler.schedulers.lr.t_cooldown=${COOLDOWN_STEPS}ba"
# EXTERNAL_CONFIGS+=" ~llm_config.scheduler.schedulers.lr.alpha_f"

#! Batch size configs (train and eval)
EXTERNAL_CONFIGS+=" llm_config.global_train_batch_size=$LOCAL_BATCH_SIZE"
EXTERNAL_CONFIGS+=" llm_config.device_train_microbatch_size=auto"
# EXTERNAL_CONFIGS+=" llm_config.device_train_microbatch_size=8"
EXTERNAL_CONFIGS+=" llm_config.device_eval_batch_size=$LOCAL_BATCH_SIZE"
EXTERNAL_CONFIGS+=" ++llm_config.device_eval_microbatch_size=auto"
EXTERNAL_CONFIGS+=" ++llm_config.device_eval_microbatch_size=1"

#! Precision and attention implementation configs
EXTERNAL_CONFIGS+=" llm_config.precision=amp_bf16"
# EXTERNAL_CONFIGS+=" llm_config.precision=amp_fp16"
# EXTERNAL_CONFIGS+=" llm_config.model.attn_config.attn_impl=flash"
# EXTERNAL_CONFIGS+=" llm_config.model.attn_config.attn_impl=torch"
# EXTERNAL_CONFIGS+=" ~llm_config.fsdp_config" # DDP
# EXTERNAL_CONFIGS+=" ++llm_config.fsdp_config.state_dict_type=sharded"

#! Federated learning configs
EXTERNAL_CONFIGS+=" fl.n_rounds=$N_ROUNDS"
EXTERNAL_CONFIGS+=" fl.n_total_clients=$N_CLIENTS"
EXTERNAL_CONFIGS+=" fl.n_clients_per_round=$N_CLIENTS"
EXTERNAL_CONFIGS+=" fl.n_local_steps=$LOCAL_STEPS"
EXTERNAL_CONFIGS+=" fl.eval_period=$FL_EVAL_PERIOD"
EXTERNAL_CONFIGS+=" fl.strategy_name=FEDAVG"
EXTERNAL_CONFIGS+=" fl.strategy_kwargs={}"
# EXTERNAL_CONFIGS+=" fl.strategy_kwargs.server_learning_rate=1.0"
# EXTERNAL_CONFIGS+=" fl.strategy_kwargs.server_momentum=0.0"
EXTERNAL_CONFIGS+=" fl.reset_optimizer=false"
EXTERNAL_CONFIGS+=" fl.dropout_ratio=0.0"
EXTERNAL_CONFIGS+=" fl.dropout_function_name=random"
# EXTERNAL_CONFIGS+=" fl.experts_assignment_policy.experts_overlapping_factor=$N_CLIENTS" # Full overlap
EXTERNAL_CONFIGS+=" fl.experts_assignment_policy.experts_overlapping_factor=1" # each expert gets to only one client

#! DEPT configs
EXTERNAL_CONFIGS+=" fl.use_unigram_metrics=true"
EXTERNAL_CONFIGS+=" fl.resize_vocab=null"
EXTERNAL_CONFIGS+=" fl.random_layers=[]"
EXTERNAL_CONFIGS+=" fl.random_init_freq=0"
EXTERNAL_CONFIGS+=" fl.truly_random_init=true"
EXTERNAL_CONFIGS+=" fl.personalized_layers=[]"
EXTERNAL_CONFIGS+=" fl.frozen_layers=null"
EXTERNAL_CONFIGS+=" fl.unfrozen_layers=null"
EXTERNAL_CONFIGS+=" fl.split_eval=false"
export COMPOSER_FAIL_ON_VOCAB_MISMATCH=0
export ALLOW_EMBEDDING_RESIZING=0
EXTERNAL_CONFIGS+=" fl.parameter_scheduler_kwargs.PARAMETERS=1"
EXTERNAL_CONFIGS+=" +fl.parameter_scheduler_kwargs.EXP_AVG=1"
EXTERNAL_CONFIGS+=" +fl.parameter_scheduler_kwargs.EXP_AVG_SQ=1"

#! mt_dao communication stack configs
EXTERNAL_CONFIGS+=" mt_dao.n_nodes=1"
if [ "$USE_RAY" = true ]; then
	EXTERNAL_CONFIGS+=" mt_dao.comm_stack.ray=true"
fi
if [ "$USE_SHM" = true ]; then
	EXTERNAL_CONFIGS+=" mt_dao.comm_stack.shm=true"
fi
if [ "$USE_S3" = true ]; then
	EXTERNAL_CONFIGS+=" mt_dao.comm_stack.s3=true"
fi
EXTERNAL_CONFIGS+=" mt_dao.comm_stack.n_batches=$COMM_BATCHES"

EXTERNAL_CONFIGS+=" use_wandb=true" # Disable wandb logging on the ServerApp

#! Additional configs
EXTERNAL_CONFIGS+=" llm_config.eval_subset_num_batches=1"
EXTERNAL_CONFIGS+=" ++llm_config.tp_config=null"
EXTERNAL_CONFIGS+=" llm_config.save_folder=null"
EXTERNAL_CONFIGS+=" llm_config.save_overwrite=true"
EXTERNAL_CONFIGS+=" llm_config.save_interval=${TOTAL_STEPS}ba"
EXTERNAL_CONFIGS+=" llm_config.eval_interval=${LOCAL_EVAL_INTERVAL}ba"
EXTERNAL_CONFIGS+=" llm_config.eval_first=false"
EXTERNAL_CONFIGS+=" llm_config.console_log_interval=50ba"
# EXTERNAL_CONFIGS+=" ~llm_config.loggers.wandb"
EXTERNAL_CONFIGS+=" ~llm_config.loggers.tensorboard"
EXTERNAL_CONFIGS+=" ~llm_config.callbacks.activation_monitor_full_model"
# EXTERNAL_CONFIGS+=" ++llm_config.callbacks.noise_scale_monitor={}"
# EXTERNAL_CONFIGS+=" llm_config.callbacks.activation_monitor_full_model.interval=1ba"
EXTERNAL_CONFIGS+=" llm_config.callbacks.optimizer_monitor.interval=1ba"

#! Evaluation gauntlet
EXTERNAL_CONFIGS+=" icl_tasks_config=empty eval_gauntlet_config=empty" # Empty gauntlet

#! mt_dao checkpointing
# EXTERNAL_CONFIGS+=" mt_dao.checkpoint=true"
#
# EXTERNAL_CONFIGS+=" mt_dao.cleanup_checkpoints=false"
# EXTERNAL_CONFIGS+=" mt_dao.cleanup_checkpoints_per_round=false"

#! Resume options
# EXTERNAL_CONFIGS+=" mt_dao.resume_round=-1"
# EXTERNAL_CONFIGS+=" mt_dao.copy_client_checkpoints=false"
# EXTERNAL_CONFIGS+=" mt_dao.restore_run_uuid=fed-125M-example-20250501_134201"
# EXTERNAL_CONFIGS+=" mt_dao.restore_cent.run_uuid=null"
# EXTERNAL_CONFIGS+=" mt_dao.restore_cent.run_batches=null"
# EXTERNAL_CONFIGS+=" pretrained_model_path=/nfs-share/anon/projects/mt_dao/runs/fed-125M-example-20250501_134201/server/21/current_server_parameters.npz"

#! External configs will be appended at the end of the command, thus rule over those
#! in the `mt_dao_base.sh` script
export EXTERNAL_CONFIGS

#! Get the master IP address
MASTER_IP=$(hostname -I | awk '{print $1}')
#! Get the Flower ports
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}']))")
CLIENTAPP_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}']))")
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}', '${CLIENTAPP_API_PORT}']))")
#! Set Flower addresses
SUPERLINK_ADDRESS=${SUPERLINK_ADDRESS:-"${MASTER_IP}:${SUPERLINK_PORT}"}
EXEC_API_ADDRESS=${EXEC_API_ADDRESS:-"${MASTER_IP}:${EXEC_API_PORT}"}
SERVERAPP_API_ADDRESS=${SERVERAPP_API_ADDRESS:-"${MASTER_IP}:${SERVERAPP_API_PORT}"}
#! Set Ray port
RAY_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}', '${CLIENTAPP_API_PORT}', '${SERVERAPP_API_PORT}']))")
#! Set Ray environmental variables
RAY_NODE_IP=$MASTER_IP
RAY_ADDRESS="$RAY_NODE_IP:$RAY_PORT"

#! Exporting the variables for the `mt_dao_base.sh` script
export SUPERLINK_ADDRESS
export EXEC_API_ADDRESS
export SERVERAPP_API_ADDRESS
export RAY_PORT
export RAY_ADDRESS
export RAY_NODE_IP

#! The following command will start the Ray head node. When running in a single node
#! setup, it is sufficient as both client and server will automatically attach to the
#! current Ray session.
if [ "$USE_RAY" = true ]; then
	uv run ray start --head --port=$RAY_PORT --temp-dir $RAY_TEMP_DIR &
fi

export CUDA_LAUNCH_BLOCKING=1

#! NOTE: Setting the environmental variable mt_dao_LAUNCHER_NO_SUPERLINK=1 will skip the
#! launch of the Flower Superlink. Setting mt_dao_LAUNCHER_NO_SERVERAPP=1 will skip the
#! launch of the Flower ServerApp. Setting mt_dao_LAUNCHER_NO_SUPERNODE=1 will skip the
#! launch of the Flower Supernode. Setting CLIENTAPP_LAUNCHER_INDEPENDENT=1 will launch
#! the Flower ClientApp in independent mode (1 per GPU). Setting CLIENTAPP_LAUNCHER_BI_INDEPENDENT=1
#! will launch the Flower ClientApp in bi-independent mode (1 per 2xGPUs). Setting
#! CLIENTAPP_NO_RELOAD=1 will pass the --no-reload option to the Flower ClientApp.

export CLIENTAPP_NO_RELOAD=1
export CLIENTAPP_LAUNCHER_INDEPENDENT=1
# export FLWR_LOG_LEVEL=DEBUG
# export RAY_BACKEND_LOG_LEVEL=debug

bash $HOME/projects/mt_dao/scripts/base_launchers/mt_dao_base.sh -p "$PROJECT_PATH"

# NOTE: For executing in a multinode setup, the Ray head node will need to be started
# only on the first node (e.g., `ray start --head --port=$RAY_PORT`), which can run the
# federated learning server, the Flower Superlink, and, potentially, a single client
#(e.g., `bash $HOME/projects/mt_dao/scripts/mt_dao_base.sh`). The other nodes
# that run, let's say, other clients, will need to attach to the Ray head node (e.g.,
# `ray start --address=$RAY_ADDRESS`, where `$RAY_ADDRESS` is the IP address and port of
# the head node) and run the `mt_dao_base.sh` script without launching Superlink, ServerApp.
