在 Vertex AI 上部署#

使用本文档了解如何在 Google Cloud Vertex AI 上部署 Nemo Retriever 文本嵌入 NIM。

注意

目前，Vertex AI 仅支持 NV-EmbedQA-E5-v5 模型。

前提条件#

在 GCP Vertex AI 上部署 Nemo Retriever 文本嵌入 NIM 之前，您需要以下条件

环境设置#

在 Artifact Registry 中创建一个仓库，用于存储 NIM Docker 镜像。

在您的终端中设置以下 Google Cloud 环境变量

export REGION={REGION}
export PROJECT_ID={PROJECT_ID}
export ARTIFACT_NAME_ON_GCP_ARTIFACT_REGISTRY={ARTIFACT_REGISTRY}

下载并推送最新的嵌入 NIM 镜像#

  1. 设置您的 NGC API KEY 以进行 NGC 身份验证。

    export NGC_API_KEY= <your_ngc_api_key>

  2. 从 NGC 下载镜像到本地机器。

    docker pull nvcr.io/nvstaging/nim/nv-embedqa-e5-v5:1.3.0

  3. 使用以下命令标记镜像，以便 Google Cloud 可以识别存储镜像的位置和仓库。

    docker tag nvcr.io/nvstaging/nim/nv-embedqa-e5-v5:1.3.0 ${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_NAME_ON_GCP_ARTIFACT_REGISTRY}/nv-embedqa-e5-v5:1.3.0

  4. 使用您的 Google 帐户验证 Google Cloud SDK。

    gcloud auth login

  5. 配置 Docker 以使用 gcloud 身份验证，从而拉取和推送镜像到仓库。

    gcloud auth configure-docker ${REGION}-docker.pkg.dev

  6. 将镜像推送到您上面创建的 GCP Artifact Registry。

    docker push ${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_NAME_ON_GCP_ARTIFACT_REGISTRY}/nv-embedqa-e5-v5:1.3.0

上传 NIM 到 Vertex AI#

通过运行以下代码，将镜像作为模型资源上传到 Vertex AI。

gcloud ai models upload \
  --region=${REGION} \
  --display-name=nv-embedqa-e5-v5:1.3.0 \
  --container-image-uri=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_NAME_ON_GCP_ARTIFACT_REGISTRY}/nv-embedqa-e5-v5:1.3.0 \
  --container-ports=8080 \
  --container-predict-route="/v1/embeddings" \
  --container-health-route="/v1/health/ready" \
  --container-shared-memory-size-mb=16000 \
  --container-env-vars="NGC_API_KEY=$NGC_API_KEY"

创建 Vertex AI 端点#

通过运行以下代码，创建一个 Triton 客户端用于发送请求的端点。

gcloud ai endpoints create \
  --region=${REGION} \
  --display-name="nv-embedqa-endpoint"

通过运行以下代码，提取 MODEL_ID 和 ENDPOINT_ID。

export MODEL_ID=$(gcloud ai models list --format="value(MODEL_ID)")
export ENDPOINT_ID=$(gcloud ai endpoints list --format="value(ENDPOINT_ID)")

将模型部署到端点#

通过运行以下代码，将模型部署到端点。

gcloud ai endpoints deploy-model ${ENDPOINT_ID} \
--region=${REGION}   \
--model=${MODEL_ID}   \
--display-name=nv-embedqa-e5-v5:1.3.0   \
--machine-type=a2-ultragpu-1g   \
--accelerator=type=nvidia-a100-80gb,count=1   \
--traffic-split=0=100

测试端点#

通过运行以下代码，确认端点处于活动状态并准备好接收请求。

curl -X POST "https://${REGION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${REGION}/endpoints/${ENDPOINT_ID}:rawPredict" \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "accept: application/json" \
  -H "Content-Type: application/json" \
  -d '{
    "input": ["Hello world"],
    "model": "nvidia/nv-embedqa-e5-v5",
    "input_type": "query"
  }'
在本页