TorchServe 配置记录

准备镜像

在生产环境中，最佳选择是通过容器部署和管理模型。为此，我们按照TorchServe官方提供的Docker镜像指南进行配置。首先，执行以下步骤克隆官方项目：

git clone https://github.com/pytorch/serve.git
cd serve/docker

接下来，根据实际需求使用build_image.sh脚本来创建镜像，脚本的关键参数包括：

使用基于GPU的基础镜像：参数-g, --gpu，指定是否使用基于GPU的基础镜像
基础镜像：参数-bi, --baseimage，指定 NVIDIA CUDA Runtime 以启用对 ONNX 的支持，比如：nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu20.04
CUDA版本：参数-cv, --cudaversion，指定 CUDA 版本，比如11.8.0对应cu118
Python版本：参数-py, --pythonversion，指定 Python 版本，比如3.9
创建类型：参数-bt, --buildtype，指定创建何种类型的镜像，可以从production、dev、ci中选择
镜像标签：参数-t, --tag，指定镜像的标签，比如torchserve:0.1、torchserve:0.1rc1等

示例：

./build_image.sh -g -cv cu118 -py 3.9 -t torchserve:dev-gpu
./build_image.sh -cv cu118 -py 3.9 -bi breakstring/gpt-sovits:latest -t torchserve:dev-base &> build.log

把构建镜像的日志导出到build.log是很好的实践，可以方便追踪镜像构建过程中的错误。

使用过程中需要注意几点，-cv参数如果在-bi参数之后会覆盖掉-bi参数指定的基础镜像，此外，Dockerfile是根据BASE_IMAGE名称中是否包含cuda:来判断是否支持CUDA的，如果名称中不包含cuda:但又想安装支持CUDA的PyTorch，需要在Dockerfile中进行调整：

if echo "$BASE_IMAGE" | grep -q "cuda:"; then \
# 改为
if echo "$BASE_IMAGE" | grep -q "cuda:" || [ -n "$USE_CUDA_VERSION" ]; then \

在镜像制作完成后，如需额外安装其他依赖，可以新建一个Dockerfile文件，并运行docker build . -t torchserve:dev-gpu来构建：

FROM torchserve:dev-base

COPY requirements.txt /tmp/
RUN pip install --no-cache-dir -r /tmp/requirements.txt
CMD ""

制作.mar文件

制作.mar文件需要先安装torch-model-archiver，并使用命令torch-model-archiver来创建.mar文件。对于非常复杂的项目来说，可以参考该issue指出的waveglow mar creation script，将所需的代码和模型都打包到压缩包，再在handler中进行解压操作。

示例如下：

#!/bin/bash
set -euxo pipefail

model_name="tts-gpt-sovits"
version="0.0.1"
model_filename="${model_name}-v${version}"

cd /tmp
rm -rf main.zip
wget https://github.com/RVC-Boss/GPT-SoVITS/archive/main.zip
rm -rf GPT-SoVITS-main
unzip main.zip
cd -
rm -rf gpt-sovits.zip
rm -rf GPT-SoVITS
mv /tmp/GPT-SoVITS-main/ GPT-SoVITS/
# copy additional files to GPT-SoVITS
cp inference.py GPT-SoVITS/GPT_SoVITS/
cp -r models/pretrained_models GPT-SoVITS/GPT_SoVITS
cp -r models/GPT_weights GPT-SoVITS/
cp -r models/SoVITS_weights GPT-SoVITS/
zip -r gpt-sovits.zip GPT-SoVITS
torch-model-archiver --model-name ${model_filename} --version ${version} --handler handler.py --extra-files gpt-sovits.zip
mv "${model_filename}.mar" ../model_store/
rm -rf gpt-sovits.zip
rm -rf GPT-SoVITS

启动容器

启动容器所用的docker-compose.yml示例如下：

version: "3"
services:
  torchserve:
    image: torchserve:dev-gpu
    container_name: torchserve-dev
    volumes:
      - ./model_store:/home/model-server/model-store/:ro
      - ./config.properties:/home/model-server/config.properties/:ro
    ports:
      - ${PORT_INFERENCE}:8080
      - ${PORT_MANAGEMENT}:8081
    shm_size: 8g
    tty: true
    stdin_open: true
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['${GPU_ID}']
              capabilities: [gpu]
    command: >
      torchserve --start --ncs
      --model-store model-store
      --models
        tts-gpt-sovits-v0.0.1.mar

建议额外写一个简单的脚本来设置环境变量，并启动容器:

#!/bin/bash

set -eu

if [ "$#" -ne 1 ]; then
    echo "Error: Invalid number of arguments."
    echo "Usage: $0 [GPU_ID]"
    exit 1
fi

export PORT_INFERENCE=23380
export PORT_MANAGEMENT=23381
export GPU_ID=$1
echo "GPU: $GPU_ID"

docker compose -f docker-compose.yml up --remove-orphans

_(°:з」∠)_/TorchServe 配置记录

准备镜像

制作.mar文件

启动容器