AIBrix

AIBrix 是放在 vLLM 前面的推理网关,把 OpenAI 兼容请求转到后面的引擎。本文在一张昇腾卡上启动一个 vLLM 引擎,再用 local mode 把一次 chat completion 转到这个引擎。

前置条件

硬件

Atlas 800T / 900 A2 训练系列,Ascend 910B。本文示例为单卡。

软件

类别

要求

CANN

toolkit 与驱动固件已安装,并可 source set_env.sh

Python

满足当前 CANN 镜像和 vLLM-Ascend 安装说明

Go

1.22.6,与上游 go.mod 的 toolchain 一致,见下文安装

Envoy

官方 linux-aarch64 发行版。下文安装当前版本

vLLM-Ascend

vllm 与 vllm-ascend 均为 0.23.0,见下文安装。版本说明见 vLLM-Ascend 安装说明

模型

Qwen/Qwen2.5-0.5B-Instruct

本文验证环境

项目

内容

镜像

swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:9.1.0-910b-ubuntu22.04-py3.12

设备

单卡 Ascend 910B

Python

3.12,来自上面的镜像

阅读本文前,请先按 快速安装昇腾环境 准备好 CANN 与驱动。

1. 加载 CANN 环境

加载 CANN 与 ATB,并把 /usr/local/sbin 加入 PATH。

source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/latest/atb/set_env.sh
export PATH=/usr/local/sbin:/usr/sbin:$PATH

2. 检查环境是否就绪

2.1 确认 NPU 在线

npu-smi info

如果 npu-smi 找不到,回到 快速安装昇腾环境 检查驱动与设备挂载。

2.2 确认工具可用

确认 CANN 已加载,并且 npu-smi 与 python 都在 PATH 里。

set -eu
test -n "$ASCEND_HOME_PATH"
echo "ASCEND_HOME_PATH set"
command -v npu-smi
python --version

输出中应包含:

ASCEND_HOME_PATH set
...npu-smi
Python ...

3. 安装依赖

3.1 安装 iproute2

安装 iproute2,其中包含 ss。

set -eu
apt-get update
apt-get install -y iproute2
ss --version

完整输出较长,其中应包含:

...ss utility, iproute2-...

3.2 安装 Go

安装上游 go.mod 所写的 Go 1.22.6 linux-arm64 包,解压到 .aibrix-quick-start/toolchain/go。

set -eu
mkdir -p .aibrix-quick-start/toolchain
curl -fL --retry 3 --retry-delay 5 --connect-timeout 30 \
  -o .aibrix-quick-start/go.tar.gz \
  https://dl.google.com/go/go1.22.6.linux-arm64.tar.gz
rm -rf .aibrix-quick-start/toolchain/go
tar -C .aibrix-quick-start/toolchain -xzf .aibrix-quick-start/go.tar.gz
.aibrix-quick-start/toolchain/go/bin/go version

输出结果如下:

go version go1.22.6 linux/arm64

3.3 安装 Envoy

从 Envoy 当前发行版匹配 linux-aarch64 包,放到 .aibrix-quick-start/bin/envoy。

set -eu
mkdir -p .aibrix-quick-start/bin
envoy_url=$(curl -fsSL -o /dev/null -w '%{url_effective}' https://github.com/envoyproxy/envoy/releases/latest)
envoy_ver=$(printf '%s\n' "$envoy_url" | grep -oE '[0-9]+\.[0-9]+\.[0-9]+$')
test -n "$envoy_ver"
curl -fL --retry 3 --retry-delay 5 --connect-timeout 30 \
  -o .aibrix-quick-start/bin/envoy \
  "https://github.com/envoyproxy/envoy/releases/download/v${envoy_ver}/envoy-${envoy_ver}-linux-aarch_64"
chmod +x .aibrix-quick-start/bin/envoy
.aibrix-quick-start/bin/envoy --version

输出中应包含:

...envoy...version...

3.4 安装 vLLM-Ascend

按 vLLM-Ascend 安装说明 安装当前带预编译包的 vllm-ascend,并安装同一版本号的 vllm。版本取下面这次安装能直接装到的预编译包。

set -eu
mkdir -p .aibrix-quick-start
python -m pip install --dry-run --no-deps --retries 3 --only-binary=vllm-ascend \
  --extra-index-url https://download.pytorch.org/whl/cpu/ \
  --extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi \
  --report .aibrix-quick-start/vllm-ascend-report.json \
  vllm-ascend
ascend_ver=$(python -c 'import json; print(next(item["metadata"]["version"] for item in json.load(open(".aibrix-quick-start/vllm-ascend-report.json"))["install"] if item["metadata"]["name"].replace("_", "-")=="vllm-ascend"))')
vllm_ver=${ascend_ver%%.post*}
test -n "$vllm_ver"
python -m pip install --retries 3 --only-binary=vllm \
  --extra-index-url https://download.pytorch.org/whl/cpu/ \
  "vllm==${vllm_ver}"
python -m pip install --retries 3 --only-binary=vllm-ascend \
  --extra-index-url https://download.pytorch.org/whl/cpu/ \
  --extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi \
  "vllm-ascend==${ascend_ver}"
python -m pip install --retries 3 --force-reinstall --no-deps --only-binary=triton-ascend \
  --extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi \
  triton-ascend
python -m pip show vllm vllm-ascend triton-ascend

完整输出较长,其中应包含:

...
Name: vllm
...
Name: vllm_ascend
...
Name: triton_ascend
...

4. 获取 AIBrix 源码并编译网关

克隆 release tag。

编译 gateway-plugins。GOPATH 与 GOCACHE 放在工作目录。

set -eu
rm -rf .aibrix-quick-start/aibrix
git clone --depth 1 --branch <ref> https://github.com/vllm-project/aibrix.git .aibrix-quick-start/aibrix
export GOPATH="$PWD/.aibrix-quick-start/gopath"
export GOCACHE="$PWD/.aibrix-quick-start/gocache"
mkdir -p "$GOPATH" "$GOCACHE"
cd .aibrix-quick-start/aibrix
CGO_ENABLED=0 "$PWD/../toolchain/go/bin/go" build -tags=nozmq -o bin/gateway-plugins cmd/plugins/main.go
"$PWD/../toolchain/go/bin/go" version -m bin/gateway-plugins

完整输出较长,其中应包含:

bin/gateway-plugins: go...
...

Note

请将 <ref> 替换为社区最新的 release 版本。

5. 启动 vLLM-Ascend 后端

在单卡上启动一个小规模 vLLM 服务,并等待 /health 就绪。

export PYTHONUNBUFFERED=1
mkdir -p .aibrix-quick-start
: > .aibrix-quick-start/vllm.log
setsid bash -c '
  echo $$ > .aibrix-quick-start/vllm.pid
  exec vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --served-model-name Qwen/Qwen2.5-0.5B-Instruct \
    --host 127.0.0.1 \
    --port 8000 \
    --max-model-len 2048 \
    --max-num-seqs 4 \
    --gpu-memory-utilization 0.2
' </dev/null >> .aibrix-quick-start/vllm.log 2>&1 &
for i in $(seq 1 180); do
  pid=$(cat .aibrix-quick-start/vllm.pid 2>/dev/null || true)
  if [ -n "${pid}" ] && ! kill -0 "${pid}" 2>/dev/null; then
    echo 'vLLM exited before /health succeeded. Full log:'
    cat .aibrix-quick-start/vllm.log
    exit 1
  fi
  if [ -n "${pid}" ] \
      && curl -sf --connect-timeout 2 -- 'http://127.0.0.1:8000/health' >/dev/null \
      && ss -ltnp 'sport = :8000' | grep -q "pid=${pid},"; then
    echo 'vLLM /health OK'
    exit 0
  fi
  sleep 2
done
echo 'timed out waiting for http://127.0.0.1:8000/health. Full log:'
cat .aibrix-quick-start/vllm.log
exit 1

输出结果如下:

vLLM /health OK

在 vLLM 日志中查找 backend=hccl。

grep -F 'backend=hccl' .aibrix-quick-start/vllm.log

完整输出较长,其中应包含:

...backend=hccl...

6. 配置网关并启动 local mode

把下面的内容写入 .aibrix-quick-start/endpoints.yaml。模型名写 Qwen/Qwen2.5-0.5B-Instruct,引擎写 vllm,后端地址写 127.0.0.1:8000。

models:
  - name: "Qwen/Qwen2.5-0.5B-Instruct"
    engine: "vllm"
    endpoints:
      - "127.0.0.1:8000"

启动 local mode。

export PATH="$PWD/.aibrix-quick-start/bin:$PATH"
bash .aibrix-quick-start/aibrix/deployment/local/run-local.sh \
  -e "$PWD/.aibrix-quick-start/endpoints.yaml"

完整输出较长,其中应包含:

...
AIBrix gateway is running!
...

7. 发一次推理

通过网关向已经启动的 vLLM 发一条对话,请模型用一句话打招呼。请求发到 127.0.0.1:10080,seed 是 42。输出与下面的结果相同。请用 python 执行:

import json
import urllib.request

payload = {
    'model': 'Qwen/Qwen2.5-0.5B-Instruct',
    'messages': [{'role': 'user', 'content': 'Say hi in one sentence.'}],
    'max_tokens': 32,
    'temperature': 0,
    'seed': 42,
}
request = urllib.request.Request(
    'http://127.0.0.1:10080/v1/chat/completions',
    data=json.dumps(payload).encode(),
    headers={'Content-Type': 'application/json'},
    method='POST',
)
with urllib.request.urlopen(request, timeout=120) as response:
    body = json.load(response)
content = (body['choices'][0]['message']['content'] or '').strip()
print('model', body['model'])
print(content)

输出结果如下:

model Qwen/Qwen2.5-0.5B-Instruct
Hello! How can I assist you today?

8. 更多用法

多引擎路由、Kubernetes 部署与其余模块的用法与社区文档相同。