在 GKE 上部署 Snowflake 的 Arctic Embed 和 TEI DLC

Snowflake 的 Arctic Embed 是一套文本嵌入模型，专注于创建高质量的检索模型，这些模型经过优化以提高性能，在其每个尺寸变体上都在 MTEB/BEIR 排行榜上实现了最先进 (SOTA) 的性能。文本嵌入推理 (TEI) 是 Hugging Face 开发的工具包，用于部署和提供开源文本嵌入和序列分类模型；为最流行的模型（包括 FlagEmbedding、Ember、GTE 和 E5）实现高性能提取。Google Kubernetes Engine (GKE) 是 Google Cloud 中的一项完全托管的 Kubernetes 服务，可用于使用 GCP 的基础设施大规模部署和操作容器化应用程序。

本示例展示了如何在 GKE 集群上部署 Hugging Face Hub 中的文本嵌入模型，该集群运行一个专用容器，用于在安全且受管理的环境中使用 Hugging Face TEI DLC 部署文本嵌入模型。

设置/配置

首先，您需要在本地机器上安装 gcloud 和 kubectl，它们分别是 Google Cloud 和 Kubernetes 的命令行工具，用于与 GCP 和 GKE 集群交互。

要安装 gcloud，请按照 Cloud SDK 文档 - 安装 gcloud CLI 中的说明进行操作。
要安装 kubectl，请按照 Kubernetes 文档 - 安装工具中的说明进行操作。

或者，为了简化本教程中命令的使用，您需要为 GCP 设置以下环境变量

export PROJECT_ID=your-project-id
export LOCATION=your-location
export CLUSTER_NAME=your-cluster-name

然后您需要登录到您的 GCP 账户，并将项目 ID 设置为您要用于部署 GKE 集群的项目。

gcloud auth login
gcloud auth application-default login  # For local development
gcloud config set project $PROJECT_ID

登录后，您需要启用 GCP 中必要的服务 API，例如 Google Kubernetes Engine API、Google Container Registry API 和 Google Container File System API，这些都是部署 GKE 集群和 Hugging Face TEI DLC 所必需的。

gcloud services enable container.googleapis.com
gcloud services enable containerregistry.googleapis.com
gcloud services enable containerfilesystem.googleapis.com

此外，要将 kubectl 与 GKE 集群凭据一起使用，您还需要安装 gke-gcloud-auth-plugin，可以如下使用 gcloud 进行安装

gcloud components install gke-gcloud-auth-plugin

无需通过 `gcloud` 特别安装 `gke-gcloud-auth-plugin`，如需了解更多替代安装方法，请访问 https://cloud.google.com/kubernetes-engine/docs/how-to/cluster-access-for-kubectl#install_plugin。

创建 GKE 集群

一切设置好后，您可以继续创建 GKE 集群和节点池。本例中将使用单个 CPU 节点，因为对于大多数工作负载，CPU 推理足以满足大多数文本嵌入模型的需求，尽管它可以通过 GPU 服务获得很大好处。

CPU 用于在文本嵌入模型之上运行推理，以展示 TEI 的当前功能。但切换到 GPU 就像将 `spec.containers[0].image` 替换为 `us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-embeddings-inference-cu122.1-4.ubuntu2204`，然后更新所请求的资源以及 `deployment.yaml` 文件中的 `nodeSelector` 要求一样简单。有关更多信息，请参阅 `gpu-config` 目录，其中包含使用 NVIDIA Tesla T4 GPU（计算能力为 7.5，即 TEI 原生支持）在 TEI 中提供 GPU 服务的预定义配置。

部署 GKE 集群时，将使用“Autopilot”模式，因为它是大多数工作负载的推荐模式，因为底层基础设施由 Google 管理。或者，您也可以使用“Standard”模式。

在创建 GKE Autopilot 集群之前，请务必查看 GKE 文档 - 通过选择机器系列优化 Autopilot Pod 性能，因为并非所有集群版本都支持所有 CPU。GPU 支持也适用同样的情况，例如，GKE 集群版本 1.28.3 或更低版本不支持 `nvidia-l4`。

gcloud container clusters create-auto $CLUSTER_NAME \
    --project=$PROJECT_ID \
    --location=$LOCATION \
    --release-channel=stable \
    --cluster-version=1.28 \
    --no-autoprovisioning-enable-insecure-kubelet-readonly-port

要选择您所在位置的特定 GKE 集群版本，您可以运行以下命令

gcloud container get-server-config \
    --flatten="channels" \
    --filter="channels.channel=STABLE" \
    --format="yaml(channels.channel,channels.defaultVersion)" \
    --location=$LOCATION

欲了解更多信息，请访问 https://cloud.google.com/kubernetes-engine/versioning#specifying_cluster_version。

GKE Cluster in the GCP Console

GKE 集群创建完成后，您可以使用以下命令通过 kubectl 获取访问它的凭据

gcloud container clusters get-credentials $CLUSTER_NAME --location=$LOCATION

部署 TEI

现在，您可以继续 Kubernetes 部署 Hugging Face TEI DLC，提供来自 Hugging Face Hub 的 `Snowflake/snowflake-arctic-embed-m` 模型。

最近，Hugging Face Hub 团队在 Hub 中包含了 `text-embeddings-inference` 标签，因此请随时在 https://huggingface.co/models?other=text-embeddings-inference 探索所有可通过 TEI 提供的嵌入模型。

Hugging Face TEI 的 DLC 将通过 `kubectl` 部署，配置文件位于 `cpu-config/` 或 `gpu-config/` 目录中，具体取决于您是想使用 CPU 还是 GPU 加速器。

`deployment.yaml`：包含 Pod 的部署详情，包括对 Hugging Face TEI DLC 的引用，将 `MODEL_ID` 设置为 `Snowflake/snowflake-arctic-embed-m`。
`service.yaml`：包含 Pod 的服务详情，暴露 TEI 服务的 8080 端口。
（可选）`ingress.yaml`：包含 Pod 的 ingress 详情，将服务暴露给外部世界，以便可以通过 ingress IP 访问。

git clone https://github.com/huggingface/Google-Cloud-Containers
kubectl apply -f Google-Cloud-Containers/examples/gke/tei-deployment/cpu-config

如前所述，对于此示例，您将把容器部署在 CPU 节点上，但部署 TEI 在 GPU 节点上的配置也可在 `gpu-config` 目录中找到，因此如果您想在 GPU 节点上部署 TEI，请运行 `kubectl apply -f Google-Cloud-Containers/examples/gke/tei-deployment/gpu-config` 而不是 `kubectl apply -f Google-Cloud-Containers/examples/gke/tei-deployment/cpu-config`。

GKE Deployment in the GCP Console

Kubernetes 部署可能需要几分钟才能准备就绪，您可以使用以下命令检查部署状态：

kubectl get pods

或者，您也可以使用以下命令等待部署准备就绪：

kubectl wait --for=condition=Available --timeout=700s deployment/tei-deployment

使用 TEI 进行推理

要在部署的 TEI 服务上运行推理，您可以选择以下方法：

将已部署的 TEI 服务端口转发到 8080 端口，以便通过 `localhost` 访问，命令如下：
```
kubectl port-forward service/tei-service 8080:8080
```
通过入口的外部 IP 访问 TEI 服务，这是此处的默认情况，因为您已经在 `cpu-config/ingress.yaml` 或 `gpu-config/ingress.yaml` 文件中定义了入口配置（但可以跳过而选择端口转发），可以通过以下命令检索：
```
kubectl get ingress tei-ingress -o jsonpath='{.status.loadBalancer.ingress[0].ip}'
```

TEI 根据模型提供的任务公开不同的推理端点

文本嵌入：文本嵌入模型公开端点 `/embed`，期望有效载荷中包含 `inputs` 键，该键可以是字符串或字符串列表，用于嵌入。
重排序：重排序模型公开端点 `/rerank`，期望有效载荷中包含 `query` 和 `texts` 键，其中 `query` 是用于对 `texts` 中每个文本的相似性进行排序的参考。
序列分类：经典序列分类模型公开端点 `/predict`，期望有效载荷中包含 `inputs` 键，该键可以是字符串或字符串列表，用于分类。更多信息请参见 https://huggingface.co/docs/text-embeddings-inference/quick_tour。

通过 cURL

要使用 `cURL` 向 TEI 服务发送 POST 请求，您可以运行以下命令：

curl https://:8080/embed \
    -X POST \
    -d '{"inputs":"What is Deep Learning?"}' \
    -H 'Content-Type: application/json'

或者向 ingress IP 发送 POST 请求：

curl http://$(kubectl get ingress tei-ingress -o jsonpath='{.status.loadBalancer.ingress[0].ip}'):8080/embed \
    -X POST \
    -d '{"inputs":"What is Deep Learning?"}' \
    -H 'Content-Type: application/json'

生成以下输出（为简洁起见已截断，但原始张量长度为 768，这是您正在服务的 `Snowflake/snowflake-arctic-embed-m` 模型的嵌入维度）

[[-0.01483098,0.010846359,-0.024679236,0.012507628,0.034231555,...]]

删除 GKE 集群

最后，一旦您在 GKE 集群上使用 TEI 完成，您可以安全地删除 GKE 集群，以避免不必要的成本。

gcloud container clusters delete $CLUSTER_NAME --location=$LOCATION

或者，如果您想保留集群，也可以将已部署的 Pod 的副本数缩减到 0，因为默认情况下部署的 GKE 集群在 GKE Autopilot 模式下只运行一个 e2-small 实例。

kubectl scale --replicas=0 deployment/tei-deployment

📍 在 GitHub 这里找到完整示例！

< > 在 GitHub 上更新