#!/bin/bash
# llama.cpp 服务启动 + 推理任务脚本
# 用法: ./llama_inference.sh

set -e

# ==================== 配置区域 ====================
LLAMA_CMD="./llama-server -m /path/to/model.gguf --port 8080"  # 启动命令
LLAMA_PORT="8080"                                                # 服务端口
HEALTH_URL="http://localhost:${LLAMA_PORT}/health"               # 健康检查地址
MAX_WAIT=300                                                     # 最长等待秒数

# 推理任务（按顺序执行）
TASKS=(
    'curl -s http://localhost:8080/completion -d "{\"prompt\": \"你好\", \"n_predict\": 50}"'
    'curl -s http://localhost:8080/completion -d "{\"prompt\": \"测试\", \"n_predict\": 50}"'
)

# ==================== 启动服务 ====================
echo "[$(date '+%H:%M:%S')] 启动 llama.cpp 服务..."
$LLAMA_CMD &
LLAMA_PID=$!
echo "[$(date '+%H:%M:%S')] 服务 PID: $LLAMA_PID"

# ==================== 等待服务就绪 ====================
echo "[$(date '+%H:%M:%S')] 等待模型加载..."
WAITED=0
while ! curl -s "$HEALTH_URL" > /dev/null 2>&1; do
    if [ $WAITED -ge $MAX_WAIT ]; then
        echo "[$(date '+%H:%M:%S')] ❌ 等待超时，服务未就绪"
        kill $LLAMA_PID 2>/dev/null || true
        exit 1
    fi
    sleep 2
    WAITED=$((WAITED + 2))
    echo "[$(date '+%H:%M:%S')] 已等待 ${WAITED}s..."
done
echo "[$(date '+%H:%M:%S')] ✅ 服务已就绪"

# ==================== 执行推理任务 ====================
echo "[$(date '+%H:%M:%S')] 开始执行推理任务..."
TASK_NUM=1
for TASK in "${TASKS[@]}"; do
    echo "[$(date '+%H:%M:%S')] --- 任务 $TASK_NUM ---"
    eval "$TASK"
    echo ""
    TASK_NUM=$((TASK_NUM + 1))
done

# ==================== 清理（可选）====================
# 如果任务执行完要关闭服务，取消下面注释
# echo "[$(date '+%H:%M:%S')] 关闭服务..."
# kill $LLAMA_PID

echo "[$(date '+%H:%M:%S')] ✅ 全部完成"