#!/bin/bash
# demo usage: source setup_env.sh
# then run your script

# Default values, you need to set these variables before running the script
DEFAULT_PROJECT_NAME=verl_train_genrm
DEFAULT_WANDB_API_KEY=""
DEFAULT_WANDB_OFFICIAL=1
DEFAULT_HDFS_DATA_PATH=""
DEFAULT_HDFS_MODEL_PATH=""
DEFAULT_HDFS_CHECKPOINT_PATH=""
DEFAULT_HDFS_LOG_PATH=""
DEFAULT_RUN_NAME=verl_train
DEFAULT_GENRMM_MODEL_PATH=""
DEFAULT_ARNOLD_WORKER_NUM=1 # default number with single node training


# Use external variables if set, otherwise use defaults
export PROJECT_NAME=${PROJECT_NAME:-$DEFAULT_PROJECT_NAME}
export WANDB_API_KEY=${WANDB_API_KEY:-$DEFAULT_WANDB_API_KEY}
export WANDB_OFFICIAL=${WANDB_OFFICIAL:-$DEFAULT_WANDB_OFFICIAL}
export HDFS_DATA_PATH=${HDFS_DATA_PATH:-$DEFAULT_HDFS_DATA_PATH}
export HDFS_MODEL_PATH=${HDFS_MODEL_PATH:-$DEFAULT_HDFS_MODEL_PATH}
export HDFS_CHECKPOINT_PATH=${HDFS_CHECKPOINT_PATH:-$DEFAULT_HDFS_CHECKPOINT_PATH}
export HDFS_LOG_PATH=${HDFS_LOG_PATH:-$DEFAULT_HDFS_LOG_PATH}
export RUN_NAME=${RUN_NAME:-$DEFAULT_RUN_NAME}
export GENRMM_MODEL_PATH=${GENRMM_MODEL_PATH:-$DEFAULT_GENRMM_MODEL_PATH}
export ARNOLD_WORKER_NUM=${ARNOLD_WORKER_NUM:-$DEFAULT_ARNOLD_WORKER_NUM} # the number of nodes
# export VLLM_USE_V1=1
export VLLM_ATTENTION_BACKEND=XFORMERS

# Optional: Print current configuration
echo "Current configuration:"
echo "PROJECT_NAME: $PROJECT_NAME"
echo "WANDB_API_KEY: $WANDB_API_KEY"
echo "WANDB_OFFICIAL: $WANDB_OFFICIAL"
echo "VLLM_ATTENTION_BACKEND: $VLLM_ATTENTION_BACKEND"
echo "HDFS_DATA_PATH: $HDFS_DATA_PATH"
echo "HDFS_MODEL_PATH: $HDFS_MODEL_PATH"
echo "HDFS_CHECKPOINT_PATH: $HDFS_CHECKPOINT_PATH"
echo "HDFS_LOG_PATH: $HDFS_LOG_PATH"
echo "GENRMM_MODEL_PATH: $GENRMM_MODEL_PATH"
echo "ARNOLD_WORKER_NUM: $ARNOLD_WORKER_NUM"
