#!/usr/bin/env bash

# Pull config from config.env
[ -f "config/config.env" ] && source config/config.env

# If the user has not provided VENV_PATH, we will assume $(pwd)/.venv
if [ -z "${VENV_PATH}" ]; then
    # what if we have VIRTUAL_ENV? use that instead
    if [ -n "${VIRTUAL_ENV}" ]; then
        export VENV_PATH="${VIRTUAL_ENV}"
    else
        export VENV_PATH="$(pwd)/.venv"
    fi
fi
if [ -z "${DISABLE_LD_OVERRIDE}" ]; then
    export NVJITLINK_PATH="$(find "${VENV_PATH}" -name nvjitlink -type d)/lib"
    # if it's not empty, we will add it to LD_LIBRARY_PATH at the front:
    if [ -n "${NVJITLINK_PATH}" ]; then
        export LD_LIBRARY_PATH="${NVJITLINK_PATH}:${LD_LIBRARY_PATH}"
    fi
fi

export TOKENIZERS_PARALLELISM=false
export PLATFORM
PLATFORM=$(uname -s)
if [[ "$PLATFORM" == "Darwin" ]]; then
    export MIXED_PRECISION="no"
fi

if [ -z "${ACCELERATE_EXTRA_ARGS}" ]; then
    ACCELERATE_EXTRA_ARGS=""
fi

if [ -z "${TRAINING_NUM_PROCESSES}" ]; then
    echo "Set custom env vars permanently in config/config.env:"
    printf "TRAINING_NUM_PROCESSES not set, defaulting to 1.\n"
    TRAINING_NUM_PROCESSES=1
fi

if [ -z "${TRAINING_NUM_MACHINES}" ]; then
    printf "TRAINING_NUM_MACHINES not set, defaulting to 1.\n"
    TRAINING_NUM_MACHINES=1
fi

if [ -z "${MIXED_PRECISION}" ]; then
    printf "MIXED_PRECISION not set, defaulting to bf16.\n"
    MIXED_PRECISION=bf16
fi

if [ -z "${TRAINING_DYNAMO_BACKEND}" ]; then
    printf "TRAINING_DYNAMO_BACKEND not set, defaulting to no.\n"
    TRAINING_DYNAMO_BACKEND="no"
fi


export ENV="default"

export ENV_PATH=""
if [[ "$ENV" != "default" ]]; then
    export ENV_PATH="${ENV}/"
fi

if [ -z "${CONFIG_BACKEND}" ]; then
    if [ -n "${CONFIG_TYPE}" ]; then
        export CONFIG_BACKEND="${CONFIG_TYPE}"
    fi
fi

if [ -z "${CONFIG_BACKEND}" ]; then
    export CONFIG_BACKEND="env"
    export CONFIG_PATH="config/${ENV_PATH}config"
    if [ -f "${CONFIG_PATH}.json" ]; then
        export CONFIG_BACKEND="json"
    elif [ -f "${CONFIG_PATH}.toml" ]; then
        export CONFIG_BACKEND="toml"
    elif [ -f "${CONFIG_PATH}.env" ]; then
        export CONFIG_BACKEND="env"
    fi
    echo "Using ${CONFIG_BACKEND} backend: ${CONFIG_PATH}.${CONFIG_BACKEND}"
fi

# Update dependencies
#if [ -z "${DISABLE_UPDATES}" ]; then
#    echo 'Updating dependencies. Set DISABLE_UPDATES to prevent this.'
#    if [ -f "pyproject.toml" ] && [ -f "poetry.lock" ]; then
#        nvidia-smi 2> /dev/null && poetry install
#        uname -s | grep -q Darwin && poetry install -C install/apple
#        rocm-smi 2> /dev/null && poetry install -C install/rocm
#    fi
#fi
# Run the training script.
#!/bin/bash

# 设置默认退出状态码
EXIT_STATUS=0

# 训练函数
run_training() {
    if [ -f "${ACCELERATE_CONFIG_PATH}" ]; then
        echo "Using Accelerate config file: ${ACCELERATE_CONFIG_PATH}"
        accelerate launch --config_file="${ACCELERATE_CONFIG_PATH}" train.py
    else
        echo "Accelerate config file not found: ${ACCELERATE_CONFIG_PATH}. Using values from config.env."
        accelerate launch ${ACCELERATE_EXTRA_ARGS} --mixed_precision="${MIXED_PRECISION}" --num_processes="${TRAINING_NUM_PROCESSES}" --num_machines="${TRAINING_NUM_MACHINES}" --dynamo_backend="${TRAINING_DYNAMO_BACKEND}" train.py
    fi
}

# 运行训练并捕获输出
TRAINING_OUTPUT=$(run_training 2>&1)
TRAINING_EXIT_CODE=$?

# 检查训练退出代码和输出
if [ $TRAINING_EXIT_CODE -eq 0 ]; then
    # 检查输出中是否包含成功信息
    if echo "$TRAINING_OUTPUT" | grep -q "Training has completed"; then
        echo "Training has completed successfully."
        EXIT_STATUS=0
    else
        echo "Training exited with code 0 but did not print completion message."
        echo "Training output:"
        echo "$TRAINING_OUTPUT"
        EXIT_STATUS=1
    fi
else
    echo "Training failed with exit code: $TRAINING_EXIT_CODE"
    echo "Training output:"
    echo "$TRAINING_OUTPUT"
    EXIT_STATUS=$TRAINING_EXIT_CODE
fi

exit $EXIT_STATUS