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

Snowflake 的 Arctic Embed 是一套文本嵌入模型，专注于创建针对性能优化的高质量检索模型，在其每个尺寸变体上都实现了最先进 (SOTA) 的 MTEB/BEIR 排行榜性能。文本嵌入推理 (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 集群的项目 ID。

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，这些 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

创建 GKE 集群

一切设置完成后，您可以继续创建 GKE 集群和节点池，在本例中将是一个单一的 CPU 节点，因为对于大多数工作负载来说，CPU 推理足以服务大多数文本嵌入模型，虽然它可以从 GPU 服务中获益良多。

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

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

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

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

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

部署 TEI

现在您可以继续进行 Hugging Face DLC for TEI 的 Kubernetes 部署，该部署服务于来自 Hugging Face Hub 的 Snowflake/snowflake-arctic-embed-m 模型。

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

• deployment.yaml：包含 Pod 的部署详细信息，包括对 Hugging Face DLC for TEI 的引用，并将 MODEL_ID 设置为 Snowflake/snowflake-arctic-embed-m。
• service.yaml：包含 Pod 的服务详细信息，将端口 8080 公开给 TEI 服务。
• (可选) ingress.yaml：包含 Pod 的入口详细信息，将服务公开给外部世界，以便可以通过入口 IP 访问。

kubectl apply -f cpu-config/

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}'

通过 cURL

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

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

或向入口 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 集群仅运行一个 e2-small 实例。

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