# coding=utf-8
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export VLLM_ENGINE_ITERATION_TIMEOUT_S=1000
timestamp=$(date +"%Y%m%d_%H%M%S")
export LOCAL_CKPT_DIR=${LOAD_CKPT_DIR:=../}
export LOCAL_CHAT_TEMPLATE_PATH=${CHAT_TEMPLATE_PATH:=$SCRIPT_DIR/chat_template_74B_1124.jinja}

if [[ ! -d ${OUTPUT_TEXT_DIR} ]]; then
    export OUTPUT_TEXT_DIR=$PROJECT_DIR/outputs
    mkdir -p ${OUTPUT_TEXT_DIR}
fi
 

apply_patch_to_package() {
    local PACKAGE_NAME=$1
    local PATCH_FILE=$2
    local PACKAGE_DIR
    echo "apply patch ${PATCH_FILE} to ${PACKAGE_NAME}..."

    PACKAGE_DIR=$(python -m pip show "$PACKAGE_NAME" | grep 'Editable project location' | awk -F': ' '{print $2}')
    (
        cd "$PACKAGE_DIR" || exit 1  

        if git apply --reverse --check "$PATCH_FILE" >/dev/null 2>&1; then
            echo "Patch for ${PACKAGE_NAME} already applied. Skipping."
            return 0
        fi

        if git apply --check "$PATCH_FILE" >/dev/null 2>&1; then	
            git apply "$PATCH_FILE"
            if [ $? -ne 0 ]; then
                echo "Failed to apply patch to ${PACKAGE_NAME} due to git apply error."
        		exit 1
            fi
        else
            echo "Failed to apply patch to ${PACKAGE_NAME}. Patch file is either incorrect or conflicts exist."
        	exit 1
        fi
    ) || exit 1
}

apply_patch_to_package "vllm-ascend" "$SCRIPT_DIR/vllm-ascend_v0.11.0rc0.patch"
apply_patch_to_package "vllm" "$SCRIPT_DIR/vllm_v0.11.0rc3.patch"

source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/omni_custom_ops/bin/set_env.bash

# 指定 HOST=127.0.0.1（本地主机）表示服务器只能从主设备访问。
# 指定 HOST=0.0.0.0 允许从同一网络上的其他设备甚至从互联网访问 vLLM 服务器，前提是网络配置正确（例如，防火墙规则、端口转发）。
HOST=xxx.xxx.xxx.xxx

# 拉起模型服务
python $SCRIPT_DIR/vllm_register.py \
	--model $LOCAL_CKPT_DIR \
	--served-model-name ${SERVED_MODEL_NAME:=pangu_7b} \
	--tensor-parallel-size ${TENSOR_PARALLEL_SIZE:=8} \
	--trust-remote-code \
	--host ${HOST} \
	--port ${PORT:=8000} \
	--max-num-seqs ${MAX_NUM_SEQS:=256} \
	--max-model-len ${MAX_MODEL_LEN:=131072} \
	--tokenizer-mode "slow" \
	--dtype bfloat16 \
	--enable-log-requests \
	--distributed-executor-backend mp \
	--gpu-memory-utilization 0.9 \
  	--max-num-batched-tokens ${MAX_NUM_BATCHED_TOKENS:=4096} \
	--no-enable-prefix-caching \
	--reasoning-parser pangu \
	--chat_template ${LOCAL_CHAT_TEMPLATE_PATH} \
	--enable-auto-tool-choice \
	--tool-call-parser pangu  \
	--trust-request-chat-template \
	--enable_chunked_prefill \
	--additional_config '{"ascend_scheduler_config":{"enabled":true,"enable_chunked_prefill":true,"chunked_prefill_enabled":true}}'   | tee $OUTPUT_TEXT_DIR/inference.log