编译 triton server 的 PyTorch backend 插件

创建日期: 2025-01-17 16:57 | 作者: 风波 | 浏览次数: 14 | 分类: TritonServer

以打包 21.11 版本为例

1. 启动容器

使用的编译镜像是 nvcr.io/nvidia/tritonserver:21.11-py3

2. 下载代码

git clone https://github.com/triton-inference-server/pytorch_backend.git

3. 切换分支到 r21.11

git checkout r.21.11

4. 安装依赖

  1. 安装依赖
apt update
apt-get install patchelf rapidjson-dev python3-dev
  1. 安装加密库 cert
  2. 下载代码 http://server:10368/cert.git
  3. 编译并安装 ./build.sh install

5. 复制依赖的 include 文件和 so 文件

如下步骤在文件 CMakeLists.txt 可以找到。只需要将这部分命令复制出来修改一下即可。

  1. 需要从 nvcr.io/nvidia/pytorch:21.11-py3 中复制依赖的 libtorch include 文件。 执行命令 docker --rm -it -v $(pwd):$(pwd) nvcr.io/nvidia/pytorch:21.11-py3 bash,切换到目录后执行如下命令
mkdir -p "include/torchvision"
cp -a /opt/conda/lib/python3.8/site-packages/torch/include include/torch
cp -a /opt/pytorch/pytorch/torch/csrc/jit/codegen include/torch/torch/csrc/jit/codegen
cp -a /opt/pytorch/vision/torchvision/csrc include/torchvision/torchvision
  1. 复制 so 文件 因为使用的编译镜像就是 nvcr.io/nvidia/tritonserver:21.11-py3,这个镜像就是 triton 官方已经完整的镜像,所以不需要再复制 so 文件,直接使用镜像中的 /opt/tritonserver/backends/pytorch 目录就可以了。

6. 修改 CMakeLists.txt

  1. 分别修改如下三个 git repo tag 为 r21.11。在第 50 行左右。
  2. TRITON_BACKEND_REPO_TAG
  3. TRITON_CORE_REPO_TAG
  4. TRITON_COMMON_REPO_TAG

  5. 增加 INCLUDE_DIRECTORIES

INCLUDE_DIRECTORIES (
    ${CMAKE_CURRENT_SOURCE_DIR}/include/torch
    ${CMAKE_CURRENT_SOURCE_DIR}/include/torchvision
)
  1. 增加加密库 crypt (如果需要的话) 在 target_link_libraries 项中增加 crypt
target_link_libraries(
  triton-pytorch-backend
  PRIVATE
    triton-core-serverapi  # from repo-core
    triton-core-backendapi # from repo-core
    triton-core-serverstub # from repo-core
    triton-backend-utils   # from repo-backend
    ${TRITON_PYTORCH_LDFLAGS}
    -ltorch
    -ltorchvision
    crypt
)

7. 修改代码,增加加密库的使用

  1. 修改文件 src/libtorch.cc,增加头文件 #include <cert/crypto.h>

  2. 修改读取模型文件的函数 ModelState::LoadModel,增加如下代码,可能需要修改变量

  std::string model_content;
  {
    std::string model_file_path = *model_path;
    try{ 
        LOG_MESSAGE(
                TRITONSERVER_LOG_INFO,
                crypt::Version::buildtime().c_str());
        if(!crypt::Cert::license_verify({}, nullptr, nullptr)) {
            return TRITONSERVER_ErrorNew(
                    TRITONSERVER_ERROR_INTERNAL,
                    ("crypt license verify failed: " + model_file_path).c_str());
        }
        if(!crypt::Cert::decrypt(model_file_path, &model_content, {})) {
            return TRITONSERVER_ErrorNew(
                    TRITONSERVER_ERROR_INTERNAL,
                    ("crypt read file failed: " + model_file_path).c_str());
        }
        LOG_MESSAGE(
                TRITONSERVER_LOG_INFO,
                (std::string("crypt read file success: ") + model_file_path).c_str());
    } catch (const std::exception& e) { 
        LOG_MESSAGE(
                TRITONSERVER_LOG_ERROR,
                (std::string("crypt read file failed: ") + model_file_path).c_str());
        return TRITONSERVER_ErrorNew(
                TRITONSERVER_ERROR_INTERNAL,
                ("crypt license verify failed: " + model_file_path).c_str());
    }    
  }

8. 编译

build.sh 编译命令如下

mkdir -p build && \
cd build && \
cmake -DTRITON_BACKEND_REPO_TAG=r21.11 \
    -DTRITON_CORE_REPO_TAG=r21.11 \
    -DTRITON_COMMON_REPO_TAG=r21.11 \
    -DCMAKE_INSTALL_PREFIX:PATH=`pwd`/install \
    -DTRITON_PYTORCH_LIB_PATHS=/opt/tritonserver/backends/pytorch \
    -DTRITON_PYTORCH_DOCKER_IMAGE="nvcr.io/nvidia/pytorch:21.11-py3" .. && \
make install VERBOSE=on

编译完成后生成文件 build/libtriton_pytorch.so

14 浏览
0 评论