使用 SAX 在 TPU v5e 上部署和提供 Llama 2 模型


本教程将指导您完成以下操作:

  • 创建 Cloud TPU VM 以部署Llama 2 系列大语言模型 (LLM),可提供不同大小(700 亿、1300 亿或 7000 亿)
  • 为模型准备检查点并将其部署在 SAX 上
  • 通过 HTTP 端点与模型交互

Serving for AGI Experiments (SAX) 是一个实验性系统,应用 PaxmlJAXPyTorch 模型以便进行推理。SAX 的代码和文档位于 Saxml Git 代码库中。支持 TPU v5e 的当前稳定版本是 v1.1.0。

SAX 单元格简介

SAX 单元(或集群)是提供模型的核心单元。它由两个主要部分组成:

  • 管理服务器:此服务器会跟踪您的模型服务器、将模型分配给这些模型服务器,并帮助客户端找到要与之交互的合适模型服务器。
  • 模型服务器:这些服务器用于运行您的模型。它们负责处理传入请求并生成响应。

下图显示了 SAX 单元的示意图:

包含管理服务器和模型服务器的 SAX 单元

图 1. 包含管理服务器和模型服务器的 SAX 单元。

您可以使用用 Python、C++ 或 Go 编写的客户端与 SAX 单元格进行交互,也可以直接通过 HTTP 服务器进行交互。下图显示了外部客户端如何与 SAX 单元格进行交互:

与 SAX 单元格交互的外部客户端和 HTTP 端点

图 2. 与 SAX 交互的外部客户端的运行时架构 单元格。

目标

  • 设置 TPU 资源以提供服务
  • 创建 SAX 集群
  • 发布 Llama 2 模型
  • 与模型交互

费用

在本文档中,您将使用 Google Cloud 的以下收费组件:

  • Cloud TPU
  • Compute Engine
  • Cloud Storage

您可使用价格计算器根据您的预计使用情况来估算费用。 Google Cloud 新用户可能有资格申请免费试用

准备工作

按照设置 Cloud TPU 环境中的说明设置 Google Cloud 项目、激活 Cloud TPU API 并创建服务账号。

创建 TPU

以下步骤展示了如何创建用于您的模型的 TPU 虚拟机。

  1. 创建环境变量:

    export PROJECT_ID=PROJECT_ID
    export ACCELERATOR_TYPE=ACCELERATOR_TYPE
    export ZONE=ZONE
    export RUNTIME_VERSION=v2-alpha-tpuv5-lite
    export SERVICE_ACCOUNT=SERVICE_ACCOUNT
    export TPU_NAME=TPU_NAME
    export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID

    环境变量说明

    PROJECT_ID
    您的 Google Cloud 项目的 ID。
    ACCELERATOR_TYPE
    加速器类型指定 您要创建的 Cloud TPU。不同的 Llama 2 模型大小具有不同的 TPU 大小要求:
    • 7Bv5litepod-4 或更大
    • 13Bv5litepod-8 或更大
    • 700 亿v5litepod-16 或更大
    ZONE
    您要在哪个可用区中创建 Cloud TPU。
    SERVICE_ACCOUNT
    您要关联到 Cloud TPU 的服务账号。
    TPU_NAME
    Cloud TPU 的名称。
    QUEUED_RESOURCE_ID
    已加入队列的资源请求的标识符。
  2. 在有效的 Google Cloud CLI 配置中设置项目 ID 和区域:

    gcloud config set project $PROJECT_ID  && gcloud config set compute/zone $ZONE
    
  3. 创建 TPU 虚拟机:

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id ${TPU_NAME} \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --accelerator-type ${ACCELERATOR_TYPE} \
      --runtime-version ${RUNTIME_VERSION} \
      --service-account ${SERVICE_ACCOUNT}
    
  4. 检查 TPU 是否处于活跃状态:

    gcloud compute tpus queued-resources list --project $PROJECT_ID --zone $ZONE
    

设置检查点转化节点

如需在 SAX 集群上运行 LLama 模型,您需要将原始 Llama 检查点设置为 SAX 兼容格式。

转换需要大量内存资源,具体取决于模型大小:

模型 机器类型
70 亿 50-60 GB 内存
130 亿 120 GB 内存
700 亿 500-600 GB 内存(N2 或 M1 机器类型)

对于 7B 和 13B 模型,您可以在 TPU VM 上运行转换。对于 70B 模型,则需要创建一个 Compute Engine 实例, 大约 1TB 磁盘空间:

gcloud compute instances create INSTANCE_NAME --project=$PROJECT_ID --zone=$ZONE \
  --machine-type=n2-highmem-128 \
  --network-interface=network-tier=PREMIUM,stack-type=IPV4_ONLY,subnet=default \
  --maintenance-policy=MIGRATE --provisioning-model=STANDARD \
  --service-account=$SERVICE_ACCOUNT \
  --scopes=https://www.googleapis.com/auth/cloud-platform \
  --tags=http-server,https-server \
  --create-disk=auto-delete=yes,boot=yes,device-name=bk-workday-dlvm,image=projects/ml-images/global/images/c0-deeplearning-common-cpu-v20230925-debian-10,mode=rw,size=500,type=projects/$PROJECT_ID/zones/$ZONE/diskTypes/pd-balanced \
  --no-shielded-secure-boot \
  --shielded-vtpm \
  --shielded-integrity-monitoring \
  --labels=goog-ec-src=vm_add-gcloud \
  --reservation-affinity=any

无论您使用 TPU 还是 Compute Engine 实例作为转化服务器, 设置您的服务器以转换 Llama 2 检查点:

  1. 对于 7B 和 13B 模型,将 将服务器名称环境变量设置为 TPU 的名称:

    export CONV_SERVER_NAME=$TPU_NAME
    

    对于 70B 模型,请将服务器名称环境变量设置为 Compute Engine 实例的名称:

    export CONV_SERVER_NAME=INSTANCE_NAME
  2. 使用 SSH 连接到转换节点。

    如果您的转化节点是 TPU,请连接到 TPU:

    gcloud compute tpus tpu-vm ssh $CONV_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE
    

    如果您的转化节点是 Compute Engine 实例,请连接到 Compute Engine 虚拟机:

    gcloud compute ssh $CONV_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE
    
  3. 在转化节点上安装所需的软件包:

    sudo apt update
    sudo apt-get install python3-pip
    sudo apt-get install git-all
    
    pip3 install paxml==1.1.0
    pip3 install torch
    pip3 install jaxlib==0.4.14
    
  4. 下载 Llama 检查点转换脚本:

    gcloud storage cp gs://cloud-tpu-inference-public/sax-tokenizers/llama/convert_llama_ckpt.py .
    

下载 Llama 2 权重

在转换模型之前,您需要下载 Llama 2 权重。在本教程中,您必须使用原始 Llama 2 权重(例如 meta-llama/Llama-2-7b),而不是已转换为 Hugging Face Transformer 格式的权重(例如 meta-llama/Llama-2-7b-hf)。

如果您已经有了 Llama 2 权重,请直接跳到将 权重

要从 Hugging Face Hub 下载权重,您需要将 设置一个用户访问令牌,并 请求访问 Llama 2 模型。如需申请访问权限,请按照说明操作 找到想要使用的模型,例如 meta-llama/Llama-2-7b.

  1. 为权重创建一个目录:

    sudo mkdir WEIGHTS_DIRECTORY
  2. 从 Hugging Face Hub 获取 Llama2 权重:

    1. 安装 Hugging Face Hub CLI:

      pip install -U "huggingface_hub[cli]"
      
    2. 切换到 weights 目录:

      cd WEIGHTS_DIRECTORY
    3. 下载 Llama 2 文件:

      python3
      from huggingface_hub import login
      login()
      from huggingface_hub import hf_hub_download, snapshot_download
      import os
      PATH=os.getcwd()
      snapshot_download(repo_id="meta-llama/LLAMA2_REPO", local_dir_use_symlinks=False, local_dir=PATH)

      LLAMA2_REPO 替换为 Hugging Face 代码库的名称 您要下载的文件:Llama-2-7bLlama-2-13bLlama-2-70b

转换权重

修改转化脚本,然后运行转化脚本来转换模型 权重。

  1. 创建一个目录来存放转换后的权重:

    sudo mkdir CONVERTED_WEIGHTS
  2. 将 Saxml GitHub 代码库克隆到您已进行读取、写入、 和执行权限:

      git clone https://github.com/google/saxml.git -b r1.1.0
    
  3. 切换到 saxml 目录:

    cd saxml
    
  4. 打开 saxml/tools/convert_llama_ckpt.py 文件。

  5. saxml/tools/convert_llama_ckpt.py 文件中,将第 169 行从:

    'scale': pytorch_vars[0]['layers.%d.attention_norm.weight' % (layer_idx)].type(torch.float16).numpy()
    

    收件人:

    'scale': pytorch_vars[0]['norm.weight'].type(torch.float16).numpy()
    
  6. 运行 saxml/tools/init_cloud_vm.sh 脚本:

    saxml/tools/init_cloud_vm.sh
    
  7. 仅适用于 70B:关闭测试模式:

    1. 打开 saxml/server/pax/lm/params/lm_cloud.py 文件。
    2. saxml/server/pax/lm/params/lm_cloud.py 将第 344 行从:

      return True
      

      收件人:

      return False
      
  8. 转换权重:

    python3 saxml/tools/convert_llama_ckpt.py --base-model-path WEIGHTS_DIRECTORY \
      --pax-model-path CONVERTED_WEIGHTS \
      --model-size MODEL_SIZE

    替换以下内容:

    • WEIGHTS_DIRECTORY:原始权重的目录。
    • CONVERTED_WEIGHTS:转换后的权重的目标路径。
    • MODEL_SIZE7b13b70b

准备检查点目录

转换检查点后,检查点目录应包含 以下结构:

checkpoint_00000000
  metadata/
      metadata
    state/
        mdl_vars.params.lm*/
        ...
        ...
        step/

创建一个名为 commit_success.txt 的空文件,并将其副本放入 checkpoint_00000000metadatastate 目录,这会让 SAX 知道此检查点已完全转换并准备好加载:

  1. 切换到检查点目录:

    cd CONVERTED_WEIGHTS/checkpoint_00000000
  2. 创建一个名为 commit_success.txt 的空文件:

    touch commit_success.txt
    
  3. 切换到元数据目录,然后创建一个名为 commit_success.txt 的空文件:

    cd metadata && touch commit_success.txt
    
  4. 切换到 state 目录,然后创建一个名为 commit_success.txt 的空文件:

    cd .. && cd state && touch commit_success.txt
    

检查点目录现在应具有以下结构:

checkpoint_00000000
    commit_success.txt
metadata/
    commit_success.txt
    metadata
    state/
        commit_success.txt
        mdl_vars.params.lm*/
        ...
        ...
        step/

创建 Cloud Storage 存储桶

您需要将转换后的检查点存储在 Cloud Storage 存储桶中,以便在发布模型时使用。

  1. 为 Cloud Storage 存储桶的名称设置环境变量:

    export GSBUCKET=BUCKET_NAME
  2. 创建存储桶:

    gcloud storage buckets create gs://${GSBUCKET}
    
  3. 将转换后的检查点文件复制到您的存储桶:

    gcloud storage cp -r CONVERTED_WEIGHTS/checkpoint_00000000  gs://$GSBUCKET/sax_models/llama2/SAX_LLAMA2_DIR/

    SAX_LLAMA2_DIR 替换为适当的值:

    • 70 亿saxml_llama27b
    • 130 亿saxml_llama213b
    • 700 亿saxml_llama270b

创建 SAX 集群

如需创建 SAX 集群,您需要执行以下操作:

在典型的部署中,您将在 Compute Engine 实例上运行管理服务器,并在 TPU 或 GPU 上运行模型服务器。在本教程中,您将在同一 TPU v5e 实例上部署管理服务器和模型服务器。

创建管理服务器

创建管理服务器 Docker 容器:

  1. 在转化服务器上,安装 Docker:

    sudo apt-get update
    sudo apt-get install docker.io
    
  2. 启动管理服务器 Docker 容器:

    sudo docker run --name sax-admin-server \
      -it \
      -d \
      --rm \
      --network host \
      --env GSBUCKET=${GSBUCKET} us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.1.0
    

您可以运行不带 -d 选项的 docker run 命令,以查看日志并确保管理服务器正确启动。

创建模型服务器

以下部分介绍了如何创建模型服务器。

7b 模型

启动模型服务器 Docker 容器:

sudo docker run --privileged \
  -it \
  -d \
  --rm \
  --network host \
  --name "sax-model-server" \
  --env SAX_ROOT=gs://${GSBUCKET}/sax-root us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
  --sax_cell="/sax/test" \
  --port=10001 \
  --platform_chip=tpuv5e \
  --platform_topology='4'

13b 模型

lm_cloud.py 中缺少 LLaMA13BFP16TPUv5e 的配置。通过 以下步骤显示了如何更新 lm_cloud.py 并提交新的 Docker 映像。

  1. 启动模型服务器:

    sudo docker run --privileged \
      -it \
      -d \
      --rm \
      --network host \
      --name "sax-model-server" \
      --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
      us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
      --sax_cell="/sax/test" \
      --port=10001 \
      --platform_chip=tpuv5e \
      --platform_topology='8'
    
  2. 使用 SSH 连接到 Docker 容器:

    sudo docker exec -it sax-model-server bash
    
  3. 在 Docker 映像中安装 Vim:

    $ apt update
    $ apt install vim
    
  4. 打开 saxml/server/pax/lm/params/lm_cloud.py 文件。搜索 LLaMA13B。您应该会看到以下代码:

    @servable_model_registry.register
    @quantization.for_transformer(quantize_on_the_fly=False)
    class LLaMA13B(BaseLLaMA):
      """13B model on a A100-40GB.
    
      April 12, 2023
      Latency = 5.06s with 128 decoded tokens. 38ms per output token.
      """
    
      NUM_LAYERS = 40
      VOCAB_SIZE = 32000
      DIMS_PER_HEAD = 128
      NUM_HEADS = 40
      MODEL_DIMS = 5120
      HIDDEN_DIMS = 13824
      ICI_MESH_SHAPE = [1, 1, 1]
    
      @property
      def test_mode(self) -> bool:
        return True
    
  5. 注释或删除以 @quantization 开头的行。在此之后 更改,文件内容应如下所示:

    @servable_model_registry.register
    class LLaMA13B(BaseLLaMA):
      """13B model on a A100-40GB.
    
      April 12, 2023
      Latency = 5.06s with 128 decoded tokens. 38ms per output token.
      """
    
      NUM_LAYERS = 40
      VOCAB_SIZE = 32000
      DIMS_PER_HEAD = 128
      NUM_HEADS = 40
      MODEL_DIMS = 5120
      HIDDEN_DIMS = 13824
      ICI_MESH_SHAPE = [1, 1, 1]
    
      @property
      def test_mode(self) -> bool:
        return True
    
  6. 添加以下代码以支持 TPU 配置。

    @servable_model_registry.register
    class LLaMA13BFP16TPUv5e(LLaMA13B):
    """13B model on TPU v5e-8.
    
    """
    
    BATCH_SIZE = [1]
    BUCKET_KEYS = [128]
    MAX_DECODE_STEPS = [32]
    ENABLE_GENERATE_STREAM = False
    
    ICI_MESH_SHAPE = [1, 1, 8]
    
    @property
    def test_mode(self) -> bool:
      return False
    
  7. 退出 Docker 容器 SSH 会话:

    exit
    
  8. 将更改提交到新的 Docker 映像:

    sudo docker commit sax-model-server sax-model-server:v1.1.0-mod
    
  9. 检查是否已创建新的 Docker 映像:

    sudo docker images
    

    您可以将 Docker 映像发布到项目的 Artifact Registry,但本教程将继续使用本地映像。

  10. 停止模型服务器。本教程的其余部分将使用更新后的模型服务器。

    sudo docker stop sax-model-server
    
  11. 使用更新后的 Docker 映像启动模型服务器。请务必指定 更新后的映像名称 sax-model-server:v1.1.0-mod

    sudo docker run --privileged \
      -it \
      -d \
      --rm \
      --network host \
      --name "sax-model-server" \
      --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
      sax-model-server:v1.1.0-mod \
      --sax_cell="/sax/test" \
      --port=10001 \
      --platform_chip=tpuv5e \
      --platform_topology='8'
    

70B 模型

使用 SSH 连接到 TPU 并启动模型服务器:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
 --project ${PROJECT_ID} \
 --zone ${ZONE} \
 --worker=all \
 --command="
   gcloud auth configure-docker \
     us-docker.pkg.dev
   # Pull SAX model server image
   sudo docker pull us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0
   # Run model server
   sudo docker run \
     --privileged  \
     -it \
     -d \
     --rm \
     --network host \
     --name "sax-model-server"  \
     --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
     us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
     --sax_cell="/sax/test" \
     --port=10001 \
     --platform_chip=tpuv5e \
     --platform_topology='16'
"

检查日志

检查模型服务器日志,确保模型服务器已正确启动:

docker logs -f sax-model-server

如果模型服务器未启动,请参阅问题排查部分了解详情。

对于 70B 模型,请对每个 TPU 虚拟机重复以下步骤:

  1. 使用 SSH 连接到 TPU:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --worker=WORKER_NUMBER

    WORKER_NUMBER 是一个从 0 开始的索引,表示您要运行哪个 TPU 虚拟机 连接。

  2. 查看日志:

    sudo docker logs -f sax-model-server
    

    三个 TPU 虚拟机应显示它们已连接到其他实例:

    I1117 00:16:07.196594 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.3:10001
    I1117 00:16:07.197484 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.87:10001
    I1117 00:16:07.199437 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.13:10001
    

    其中一个 TPU 虚拟机的日志应显示模型服务器正在启动:

    I1115 04:01:29.479170 139974275995200 model_service_base.py:867] Started joining SAX cell /sax/test
    ERROR: logging before flag.Parse: I1115 04:01:31.479794       1 location.go:141] Calling Join due to address update
    ERROR: logging before flag.Parse: I1115 04:01:31.814721       1 location.go:155] Joined 10.182.0.44:10000
    

发布模型

SAX 附带一个名为 saxutil 的命令行工具,可简化与 SAX 模型服务器的交互。在本教程中,您将使用 saxutil 发布模型。如需查看 saxutil 命令的完整列表,请参阅 Saxml 自述文件 文件

  1. 切换到克隆 Saxml GitHub 代码库的目录:

    cd  saxml
    
  2. 对于 70B 模型,请连接到转化服务器:

    gcloud compute ssh ${CONV_SERVER_NAME} \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
    
  3. 安装 Bazel:

    sudo apt-get install bazel
    
  4. 设置别名,以便通过 Cloud Storage 存储桶运行 saxutil

    alias saxutil='bazel run saxml/bin:saxutil -- --sax_root=gs://${GSBUCKET}/sax-root'
    
  5. 使用 saxutil 发布模型。在 TPU 上大约需要 10 分钟 v5litepod-8。

    saxutil --sax_root=gs://${GSBUCKET}/sax-root publish '/sax/test/MODEL' \
        saxml.server.pax.lm.params.lm_cloud.PARAMETERS \
        gs://${GSBUCKET}/sax_models/llama2/SAX_LLAMA2_DIR/checkpoint_00000000/ \
        1

    执行以下变量替换操作:

    模型大小
    70 亿 MODEL:llama27b
    PARAMETERS:saxml.server.pax.lm.params.lm_cloud.LLaMA7BFP16TPUv5e
    SAX_LLAMA2_DIR:saxml_llama27b
    130 亿 MODEL:llama213b
    PARAMETERS:saxml.server.pax.lm.params.lm_cloud.LLaMA13BFP16TPUv5e
    SAX_LLAMA2_DIR:saxml_llama213b
    700 亿 MODEL:llama270b
    PARAMETERS:saxml.server.pax.lm.params.lm_cloud.LLaMA70BFP16TPUv5e
    SAX_LLAMA2_DIR:saxml_llama270b

测试部署

如需检查部署是否成功,请使用 saxutil ls 命令:

saxutil ls /sax/test/MODEL

成功的部署的副本数量应该大于零, 与以下代码类似:

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root is /sax/test/1lama27b

+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+
| MODEL    | MODEL PATH                                            | CHECKPOINT PATH                                                       | # OF REPLICAS | (SELECTED) REPLICAADDRESS |
+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+
| llama27b | saxml.server.pax.lm.params.lm_cloud.LLaMA7BFP16TPUv5e | gs://${MODEL_BUCKET}/sax_models/llama2/7b/pax_7B/checkpoint_00000000/ | 1             | 10.182.0.28:10001         |
+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+

模型服务器的 Docker 日志将类似于以下内容:

I1114 17:31:03.586631 140003787142720 model_service_base.py:532] Successfully loaded model for key: /sax/test/llama27b

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root is /sax/test/1lama27b

问题排查

如果部署失败,请检查模型服务器日志:

sudo docker logs -f sax-model-server

如果部署成功,您应该会看到以下输出:

Successfully loaded model for key: /sax/test/llama27b

如果日志未显示模型已部署,请检查模型配置和模型检查点的路径。

生成回答

您可以使用 saxutil 工具生成对提示的回答。

生成问题回复:

saxutil lm.generate -extra="temperature:0.2"  /sax/test/MODEL "Q: Who is Harry Potter's mother? A:"

输出应类似如下所示:

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root' lm.generate /sax/test/llama27b 'Q: Who is Harry Potter's mother? A: `
+-------------------------------+------------+
| GENERATE                      | SCORE      |
+-------------------------------+------------+
| 1. Harry Potter's mother is   | -20.214787 |
| Lily Evans. 2. Harry Potter's |            |
| mother is Petunia Evans       |            |
| (Dursley).                    |            |
+-------------------------------+------------+

从客户端与模型交互

SAX 代码库包含可用于与 SAX 单元交互的客户端。 客户端支持 C++、Python 和 Go。以下示例展示了 如何构建 Python 客户端。

  1. 构建 Python 客户端:

    bazel build saxml/client/python:sax.cc --compile_one_dependency
    
  2. 将客户端添加到 PYTHONPATH。本示例假定您在主目录下有 saxml

    export PYTHONPATH=${PYTHONPATH}:$HOME/saxml/bazel-bin/saxml/client/python/
    
  3. 从 Python shell 与 SAX 交互:

    $ python3
    Python 3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0] on linux
    Type "help", "copyright", "credits" or "license" for more information.
    >>> import sax
    >>>
    

从 HTTP 端点与模型交互

如需从 HTTP 端点与模型进行交互,请创建一个 HTTP 客户端:

  1. 创建 Compute Engine 虚拟机:

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    export HTTP_SERVER_NAME=HTTP_SERVER_NAME
    export SERVICE_ACCOUNT=SERVICE_ACCOUNT
    export MACHINE_TYPE=e2-standard-8
    gcloud compute instances create $HTTP_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE \
      --machine-type=$MACHINE_TYPE \
      --network-interface=network-tier=PREMIUM,stack-type=IPV4_ONLY,subnet=default \
      --maintenance-policy=MIGRATE --provisioning-model=STANDARD \
      --service-account=$SERVICE_ACCOUNT \
      --scopes=https://www.googleapis.com/auth/cloud-platform \
      --tags=http-server,https-server \
      --create-disk=auto-delete=yes,boot=yes,device-name=$HTTP_SERVER_NAME,image=projects/ml-images/global/images/c0-deeplearning-common-cpu-v20230925-debian-10,mode=rw,size=500,type=projects/$PROJECT_ID/zones/$ZONE/diskTypes/pd-balanced \
      --no-shielded-secure-boot \
      --shielded-vtpm \
      --shielded-integrity-monitoring \
      --labels=goog-ec-src=vm_add-gcloud \
      --reservation-affinity=any
  2. 使用 SSH 连接到 Compute Engine 虚拟机:

    gcloud compute ssh $HTTP_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE
    
  3. 克隆 AI on GKE GitHub 代码库:

    git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git
    
  4. 切换到 HTTP 服务器目录:

    cd ai-on-gke/tools/saxml-on-gke/httpserver
    
  5. 构建 Docker 文件:

    docker build -f Dockerfile -t sax-http .
    
  6. 运行 HTTP 服务器:

    docker run -e SAX_ROOT=gs://${GSBUCKET}/sax-root -p 8888:8888 -it sax-http
    

使用以下命令从本地机器或其他有权访问端口 8888 的服务器测试端点:

  1. 导出服务器 IP 地址和端口的环境变量:

    export LB_IP=HTTP_SERVER_EXTERNAL_IP
    export PORT=8888
  2. 设置包含模型和查询的 JSON 载荷:

    json_payload=$(cat  << EOF
    {
      "model": "/sax/test/MODEL",
      "query": "Example query"
    }
    EOF
    )
  3. 发送请求:

    curl --request POST --header "Content-type: application/json" -s $LB_IP:$PORT/generate --data "$json_payload"
    

清理

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用,请删除包含这些资源的项目,或者保留项目但删除各个资源。

完成本教程后,请按照以下步骤清理 资源。

  1. 删除您的 Cloud TPU。

    $ gcloud compute tpus tpu-vm delete $TPU_NAME  --zone $ZONE
  2. 删除您的 Compute Engine 实例(如果您创建了)。

    gcloud compute instances delete INSTANCE_NAME
  3. 删除您的 Cloud Storage 存储桶及其内容。

    gcloud storage rm --recursive gs://BUCKET_NAME

后续步骤