透過 Saxml 在 GKE 上使用 TPU 提供 Gemma 開放原始碼模型

標準

---

注意： 自 2025 年 4 月 24 日起，系統將降低 GKE 上 Saxml 的優先順序。這表示專案不會再收到更新。現有的 Saxml 部署作業會繼續運作，不受影響。強烈建議您遷移至 JetStream，這是 Google 最新的開放原始碼推論架構，可在 TPU 和 GPU 上提供高效能的 LLM 服務。JetStream 提供連續批次處理和量化功能，可提升處理量和記憶體效率。如需遷移範例，請參閱透過 JetStream 在 GKE 上使用 TPU 提供 Gemma。

本教學課程說明如何使用 Saxml 服務架構，透過 GKE 上的 TPU 部署及提供 Gemma 2 大型語言模型 (LLM)。本教學課程提供基礎知識，協助您瞭解及探索如何在代管型 Kubernetes 環境中，實際部署 LLM 以進行推論。您可以使用 Saxml 將預先建構的容器部署至 GKE。您也會設定 GKE，在執行階段從 Cloud Storage 載入 Gemma 2B 和 7B 的權重。

本教學課程的適用對象為機器學習 (ML) 工程師、平台管理員和操作員，以及有興趣使用 Kubernetes 容器自動化調度管理功能提供 LLM 服務的資料和 AI 專家。如要進一步瞭解Google Cloud 內容中提及的常見角色和範例工作，請參閱「常見的 GKE Enterprise 使用者角色和工作」。

閱讀本頁面之前，請先熟悉下列概念：

• 目前可用的 TPU 版本 (Cloud TPU 系統架構)
• GKE 中的 TPU

如果您需要統一管理的 AI 平台，以符合成本效益的方式快速建構及提供機器學習模型，建議您試用 Vertex AI 部署解決方案。

背景

本節說明本教學課程中使用的主要技術。

Gemma

Gemma 是一組開放授權的輕量型生成式 AI 模型，您可以在應用程式、硬體、行動裝置或代管服務中執行這些 AI 模型。您可以使用 Gemma 模型生成文字，也可以調整這些模型來執行專業工作。

詳情請參閱 Gemma 說明文件。

TPU

TPU 是 Google 開發的客製化特殊應用積體電路 (ASIC)，用於加速 TensorFlow、PyTorch 和 JAX 等資料處理架構。

本教學課程會提供 Gemma 2B 和 Gemma 7B 模型。 GKE 會在下列單一主機 TPU v5e 節點集區中代管這些模型：

• Gemma 2B：在 TPU v5e 節點集區中代管的指令調整模型，拓撲為 1x1，代表一個 TPU 晶片。節點的機器類型為 ct5lp-hightpu-1t。
• Gemma 7B：在 TPU v5e 節點集區中代管的指令調整模型，拓撲為 2x2，代表四個 TPU 晶片。節點的機器類型為 ct5lp-hightpu-4t。

Saxml

Saxml 是實驗性系統，可為推論提供 Paxml、JAX 和 PyTorch 模型。Saxml 系統包含下列元件：

• Saxml 儲存格或 Sax 叢集：管理員伺服器和一組模型伺服器。管理伺服器會追蹤模型伺服器、將發布的模型指派給模型伺服器以提供服務，並協助用戶端找出提供特定發布模型的模型伺服器。
• Saxml 用戶端：Saxml 系統的使用者導向程式設計介面。Saxml 用戶端包含指令列工具 (saxutil) 和一組 Python、C++ 和 Go 用戶端程式庫。

在本教學課程中，您也會使用 Saxml HTTP 伺服器。 Saxml HTTP 伺服器是自訂 HTTP 伺服器，可封裝 Saxml Python 用戶端程式庫，並公開 REST API，與 Saxml 系統互動。REST API 包含發布、列出、取消發布模型，以及產生預測結果的端點。

目標

1. 根據模型特性，準備具有建議 TPU 拓撲的 GKE Standard 叢集。
2. 在 GKE 上部署 Saxml 元件。
3. 取得並發布 Gemma 2B 或 Gemma 7B 參數模型。
4. 提供已發布的模型並與之互動。

架構

本節說明本教學課程使用的 GKE 架構。此架構包含 GKE Standard 叢集，可佈建 TPU 並代管 Saxml 元件，以部署及提供 Gemma 2B 或 7B 模型。下圖顯示這個架構的元件：

這個架構包含下列元件：

• GKE Standard 可用區叢集。
• 單一主機 TPU 配量節點集區，取決於您想服務的 Gemma 模型：
  • Gemma 2B：使用 1x1 拓撲設定的 TPU v5e。Saxml 模型伺服器的一個例項已設定為使用這個節點集區。
  • Gemma 7B：使用 TPU v5e 設定，並採用 2x2 拓撲。Saxml 模型伺服器的一個例項已設定為使用這個節點集區。
• 預設 CPU 節點集區，用於部署 Saxml 管理伺服器和 Saxml HTTP 伺服器。
• 兩個 Cloud Storage bucket：
  • 一個 Cloud Storage 值區會儲存由管理伺服器管理的狀態。
  • 一個 Cloud Storage bucket 儲存 Gemma 模型檢查點。

這個架構具有下列特徵：

• 公開 Artifact Registry 會管理 Saxml 元件的容器映像檔。
• GKE 叢集使用 Workload Identity Federation for GKE。所有 Saxml 元件都使用工作負載身分聯盟，整合 IAM 服務帳戶來存取外部服務，例如 Cloud Storage 值區。
• Saxml 元件產生的記錄檔會整合至 Cloud Logging。
• 您可以使用 Cloud Monitoring 分析 GKE 節點集區的效能指標，例如本教學課程建立的節點集區。

事前準備

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the required API.

Enable the API

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the required API.

Enable the API

• Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin, roles/iam.policyAdmin, roles/iam.securityAdmin, roles/iam.roleAdmin

  Check for the roles

  1. In the Google Cloud console, go to the IAM page.

     Go to IAM
  2. Select the project.

  3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.

  4. For all rows that specify or include you, check the Role column to see whether the list of roles includes the required roles.

  Grant the roles

  1. In the Google Cloud console, go to the IAM page.

     前往「IAM」頁面
  2. 選取專案。
  3. 按一下「授予存取權」。

  4. 在「New principals」(新增主體) 欄位中，輸入您的使用者 ID。 這通常是 Google 帳戶的電子郵件地址。

  5. 在「Select a role」(選取角色) 清單中，選取角色。
  6. 如要授予其他角色，請按一下 「新增其他角色」，然後新增每個其他角色。
  7. 按一下 [Save]。
  • 確認您有足夠的配額，可使用 5 個 TPU v5e 晶片。在本教學課程中，您將使用隨選執行個體。
  • 如果沒有 Kaggle 帳戶，請先建立帳戶。

  準備 Gemma 環境

  啟動 Cloud Shell

  在本教學課程中，您將使用 Cloud Shell 管理託管在 Google Cloud上的資源。Cloud Shell 已預先安裝本教學課程所需的軟體，包括 kubectl 和 gcloud CLI。

  1. 在 Google Cloud 控制台中啟動 Cloud Shell 執行個體：
     開啟 Cloud Shell

  2. 設定預設環境變數：

     gcloud config set project PROJECT_ID
     export PROJECT_ID=$(gcloud config get project)
     export LOCATION=LOCATION
     export CLUSTER_NAME=saxml-tpu
     

     替換下列值：

     • PROJECT_ID：您的 Google Cloud 專案 ID。
     • LOCATION：Compute Engine 區域的名稱，TPU v5e 機器類型適用於該區域。

  建立 GKE Standard 叢集

  在本節中，您將建立 GKE 叢集和節點集區。

  Gemma 2B-it

  使用 Cloud Shell 執行下列操作：

  1. 建立使用 Workload Identity Federation for GKE 的 Standard 叢集：

     gcloud container clusters create ${CLUSTER_NAME} \
         --enable-ip-alias \
         --machine-type=e2-standard-4 \
         --num-nodes=2 \
         --release-channel=rapid \
         --workload-pool=${PROJECT_ID}.svc.id.goog \
         --location=${LOCATION}
     

     建立叢集可能需要幾分鐘的時間。

  2. 建立具有 1x1 拓撲和一個節點的 TPU v5e 節點集區：

     gcloud container node-pools create tpu-v5e-1x1 \
         --cluster=${CLUSTER_NAME} \
         --machine-type=ct5lp-hightpu-1t \
         --num-nodes=1 \
         --location=${LOCATION}
     

     您會在節點集區中提供 Gemma 2B 模型。

  Gemma 7B-it

  使用 Cloud Shell 執行下列操作：

  1. 建立使用 Workload Identity Federation for GKE 的 Standard 叢集：

     gcloud container clusters create ${CLUSTER_NAME} \
         --enable-ip-alias \
         --machine-type=e2-standard-4 \
         --num-nodes=2 \
         --release-channel=rapid \
         --workload-pool=${PROJECT_ID}.svc.id.goog \
         --location=${LOCATION}
     

     建立叢集可能需要幾分鐘的時間。

  2. 建立具有 2x2 拓撲和一個節點的 TPU v5e 節點集區：

     gcloud container node-pools create tpu-v5e-2x2 \
         --cluster=${CLUSTER_NAME} \
         --machine-type=ct5lp-hightpu-4t \
         --num-nodes=1 \
         --location=${LOCATION}
     

     您會在節點集區中提供 Gemma 7B 模型。

  建立 Cloud Storage bucket

  建立兩個 Cloud Storage bucket，管理 Saxml Admin 伺服器和模型檢查點的狀態。

  在 Cloud Shell 中執行下列指令：

  1. 建立 Cloud Storage bucket，用來儲存 Saxml Admin 伺服器設定。

     gcloud storage buckets create gs://ADMIN_BUCKET_NAME
     

     將 ADMIN_BUCKET_NAME 替換為儲存 Saxml 管理伺服器的 Cloud Storage 值區名稱。

  2. 建立 Cloud Storage bucket 來儲存模型檢查點：

     gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME
     

     將 CHECKPOINTS_BUCKET_NAME 替換為儲存模型檢查點的 Cloud Storage bucket 名稱。

  使用 Workload Identity Federation for GKE 設定工作負載存取權

  將 Kubernetes ServiceAccount 指派給應用程式，並將該 Kubernetes ServiceAccount 設定為 IAM 服務帳戶。

  1. 設定 kubectl 與叢集通訊：

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

  2. 為應用程式建立要使用的 Kubernetes ServiceAccount：

     gcloud iam service-accounts create wi-sax
     

  3. 為 IAM 服務帳戶新增 IAM 政策繫結，以便讀取及寫入 Cloud Storage：

     gcloud projects add-iam-policy-binding ${PROJECT_ID} \
         --member "serviceAccount:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
         --role roles/storage.objectUser
     
     gcloud projects add-iam-policy-binding ${PROJECT_ID} \
         --member "serviceAccount:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
         --role roles/storage.insightsCollectorService
     

  4. 在兩個服務帳戶之間新增 IAM 政策繫結，允許 Kubernetes ServiceAccount模擬 IAM 服務帳戶。這個繫結可讓 Kubernetes ServiceAccount 做為 IAM 服務帳戶：

     gcloud iam service-accounts add-iam-policy-binding wi-sax@${PROJECT_ID}.iam.gserviceaccount.com \
         --role roles/iam.workloadIdentityUser \
         --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]"
     

  5. 使用 IAM 服務帳戶的電子郵件地址註解 Kubernetes 服務帳戶：

     kubectl annotate serviceaccount default \
         iam.gke.io/gcp-service-account=wi-sax@${PROJECT_ID}.iam.gserviceaccount.com
     

  取得模型存取權

  如要部署至 GKE，您必須登入 Kaggle 平台、簽署授權同意協議，並取得 Kaggle API 權杖，才能存取 Gemma 模型。在本教學課程中，您會使用 Kubernetes Secret 儲存 Kaggle 憑證。

  簽署授權同意聲明協議

  您必須簽署同意聲明協議，才能使用 Gemma。請按照下列步驟操作：

  1. 前往 Kaggle.com 的模型同意聲明頁面。
  2. 如果尚未登入 Kaggle，請先登入。
  3. 按一下「要求存取權」。
  4. 在「選擇同意聲明適用的帳戶」部分，選取「透過 Kaggle 帳戶驗證」，使用 Kaggle 帳戶授予同意聲明。
  5. 接受模型條款及細則。

  產生存取權杖

  如要透過 Kaggle 存取模型，您需要 Kaggle API 權杖。

  如果沒有權杖，請按照下列步驟產生新權杖：

  1. 在瀏覽器中前往 Kaggle 設定。
  2. 在「API」部分，按一下「Create New Token」。

  Kaggle 會下載名為 kaggle.json 的檔案。

  將存取權杖上傳至 Cloud Shell

  在 Cloud Shell 中，您可以將 Kaggle API 權杖上傳至專案： Google Cloud

  1. 在 Cloud Shell 中，依序點選「更多」>「上傳」。
  2. 選取「檔案」，然後按一下「選擇檔案」。
  3. 開啟 kaggle.json 檔案。
  4. 按一下「上傳」。

  為 Kaggle 憑證建立 Kubernetes Secret

  在 Cloud Shell 中執行下列步驟：

  1. 設定 kubectl 與叢集通訊：

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

  2. 建立 Secret 來儲存 Kaggle 憑證：

     kubectl create secret generic kaggle-secret \
         --from-file=kaggle.json
     

  部署 Saxml

  在本節中，您將部署 Saxml 管理伺服器、模型伺服器和 HTTP 伺服器。本教學課程使用 Kubernetes Deployment 資訊清單。Deployment 是 Kubernetes API 物件，可讓您執行多個 Pod 副本，並將這些副本分配到叢集中的節點。

  部署 Saxml 管理伺服器

  在本節中，您將部署 Saxml 管理伺服器。

  1. 建立下列 saxml-admin-server.yaml 資訊清單：

     apiVersion: apps/v1
     kind: Deployment
     metadata:
       name: sax-admin-server
     spec:
       replicas: 1
       selector:
         matchLabels:
           app: sax-admin-server
       template:
         metadata:
           labels:
             app: sax-admin-server
         spec:
           hostNetwork: false
           containers:
           - name: sax-admin-server
             image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.2.0
             securityContext:
               privileged: true
             ports:
             - containerPort: 10000
             env:
             - name: GSBUCKET
               value: ADMIN_BUCKET_NAME

     將 ADMIN_BUCKET_NAME 替換為您在「建立 Cloud Storage 值區」一節中建立的值區名稱。請勿加入 gs:// 前置字元。

  2. 套用資訊清單：

     kubectl apply -f saxml-admin-server.yaml
     

  3. 驗證管理伺服器部署作業：

     kubectl get deployment
     

     輸出看起來類似以下內容：

     NAME                              READY   UP-TO-DATE   AVAILABLE   AGE
     sax-admin-server                  1/1     1            1           ##s
     

  部署 Saxml 模型伺服器

  請按照下列操作說明，為 Gemma 2B 或 Gemma 7B 模型部署模型伺服器。

  Gemma 2B-it

  1. 建立下列 saxml-model-server-1x1.yaml 資訊清單：

     apiVersion: apps/v1
     kind: Deployment
     metadata:
       name: sax-model-server-v5e-1x1
     spec:
       replicas: 1
       selector:
         matchLabels:
           app: gemma-server
       strategy:
         type: Recreate
       template:
         metadata:
           labels:
             app: gemma-server
             ai.gke.io/model: gemma-2b-it
             ai.gke.io/inference-server: saxml
             examples.ai.gke.io/source: user-guide
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-topology: 1x1
             cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
           hostNetwork: false
           restartPolicy: Always
           containers:
           - name: inference-server
             image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
             args:
             - "--jax_platforms=tpu"
             - "--platform_chip=tpuv5e"
             - "--platform_topology=1x1"
             - "--port=10001"
             - "--sax_cell=/sax/test"
             ports:
             - containerPort: 10001
             securityContext:
               privileged: true
             env:
             - name: SAX_ROOT
               value: "gs://ADMIN_BUCKET_NAME/sax-root"
             resources:
               requests:
                 google.com/tpu: 1
               limits:
                 google.com/tpu: 1

     將 ADMIN_BUCKET_NAME 替換為您在「建立 Cloud Storage 值區」一節中建立的值區名稱。請勿加入 gs:// 前置字元。

  2. 套用資訊清單：

     kubectl apply -f saxml-model-server-1x1.yaml
     

  3. 確認模型伺服器部署作業的狀態：

     kubectl get deployment
     

     輸出看起來類似以下內容：

     NAME                                              READY   STATUS    RESTARTS   AGE
     sax-admin-server                                  1/1     Running   0          ##m
     sax-model-server-v5e-1x1                          1/1     Running   0          ##s
     

  成功： 您已在 TPU v5e 節點集區上部署模型伺服器，並使用 1x1 切片提供 Gemma 2B 模型。

  Gemma 7B-it

  1. 建立下列 saxml-model-server-2x2.yaml 資訊清單：

     apiVersion: apps/v1
     kind: Deployment
     metadata:
       name: sax-model-server-v5e-2x2
     spec:
       replicas: 1
       selector:
         matchLabels:
           app: gemma-server
       strategy:
         type: Recreate
       template:
         metadata:
           labels:
             app: gemma-server
             ai.gke.io/model: gemma-7b-it
             ai.gke.io/inference-server: saxml
             examples.ai.gke.io/source: user-guide
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-topology: 2x2
             cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
           hostNetwork: false
           restartPolicy: Always
           containers:
           - name: inference-server
             image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
             args:
             - "--jax_platforms=tpu"
             - "--platform_chip=tpuv5e"
             - "--platform_topology=2x2"
             - "--port=10001"
             - "--sax_cell=/sax/test"
             ports:
             - containerPort: 10001
             securityContext:
               privileged: true
             env:
             - name: SAX_ROOT
               value: "gs://ADMIN_BUCKET_NAME/sax-root"
             resources:
               requests:
                 google.com/tpu: 4
               limits:
                 google.com/tpu: 4

     將 ADMIN_BUCKET_NAME 替換為您在「建立 Cloud Storage 值區」一節中建立的值區名稱。請勿加入 gs:// 前置字元。

  2. 套用資訊清單：

     kubectl apply -f saxml-model-server-2x2.yaml
     

  3. 確認模型伺服器部署作業的狀態：

     kubectl get deployment
     

     輸出看起來類似以下內容：

     NAME                                              READY   STATUS    RESTARTS   AGE
     sax-admin-server                                  1/1     Running   0          ##m
     sax-model-server-v5e-2x2                          1/1     Running   0          ##s
     

  成功： 您已在 TPU v5e 節點集區中部署模型伺服器，並使用 2x2 切片提供 Gemma 7B 模型。

  部署 Saxml HTTP 伺服器

  在本節中，您將部署 Saxml HTTP 伺服器，並建立用於存取伺服器的叢集 IP 服務。

  1. 建立下列 saxml-http.yaml 資訊清單：

     apiVersion: apps/v1
     kind: Deployment
     metadata:
       name: sax-http
     spec:
       replicas: 1
       selector:
         matchLabels:
           app: sax-http
       template:
         metadata:
           labels:
             app: sax-http
         spec:
           hostNetwork: false
           containers:
           - name: sax-http
             image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-http:v1.2.0
             imagePullPolicy: Always
             ports:
             - containerPort: 8888
             env:
             - name: SAX_ROOT
               value: "gs://ADMIN_BUCKET_NAME/sax-root"
     ---
     apiVersion: v1
     kind: Service
     metadata:
       name: sax-http-svc
     spec:
       selector:
         app: sax-http
       ports:
       - protocol: TCP
         port: 8888
         targetPort: 8888
       type: ClusterIP

     將 ADMIN_BUCKET_NAME 替換為儲存 Saxml 管理伺服器的 Cloud Storage 值區名稱。

  2. 套用資訊清單：

     kubectl apply -f saxml-http.yaml
     

  3. 驗證 Saxml HTTP 伺服器部署作業的狀態：

     kubectl get deployment
     

     Gemma 2B-it

     輸出看起來類似以下內容：

     NAME                                              READY   STATUS    RESTARTS   AGE
     sax-admin-server                                  1/1     Running   0          ##m
     sax-model-server-v5e-1x1                          1/1     Running   0          ##m
     sax-http                                          1/1     Running   0          ##s
     

     Gemma 7B-it

     輸出看起來類似以下內容：

     NAME                                              READY   STATUS    RESTARTS   AGE
     sax-admin-server                                  1/1     Running   0          ##m
     sax-model-server-v5e-2x2                          1/1     Running   0          ##m
     sax-http                                          1/1     Running   0          ##s
     

  下載模型檢查點

  在本節中，您將執行 Kubernetes Job，擷取、下載及儲存模型檢查點。Kubernetes 中的 Job 控制器會建立一或多個 Pod，並確保這些 Pod 成功執行特定工作。

  根據您要使用的 Gemma 模型，遵循步驟執行：

  Gemma 2B-it

  1. 建立下列 job-2b.yaml 資訊清單：

     apiVersion: v1
     kind: ConfigMap
     metadata:
       name: fetch-model-scripts
     data:
       fetch_model.sh: |-
         #!/usr/bin/bash -x
         pip install kaggle --break-system-packages && \
     
         MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
         VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \
     
         mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
         kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
         echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \
     
         gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
         echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
     ---
     apiVersion: batch/v1
     kind: Job
     metadata:
       name: data-loader-2b
       labels:
         app: data-loader-2b
     spec:
       ttlSecondsAfterFinished: 120
       template:
         metadata:
           labels:
             app: data-loader-2b
         spec:
           restartPolicy: OnFailure
           containers:
           - name: gcloud
             image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
             command:
             - /scripts/fetch_model.sh
             env:
             - name: BUCKET_NAME
               value: CHECKPOINTS_BUCKET_NAME
             - name: KAGGLE_CONFIG_DIR
               value: /kaggle
             - name: MODEL_PATH
               value: "google/gemma/pax/2b-it/2"
             volumeMounts:
             - mountPath: "/kaggle/"
               name: kaggle-credentials
               readOnly: true
             - mountPath: "/scripts/"
               name: scripts-volume
               readOnly: true
           volumes:
           - name: kaggle-credentials
             secret:
               defaultMode: 0400
               secretName: kaggle-secret
           - name: scripts-volume
             configMap:
               defaultMode: 0700
               name: fetch-model-scripts

     將 CHECKPOINTS_BUCKET_NAME 替換為您在「建立 Cloud Storage 值區」一節中建立的值區名稱。請勿加入 gs:// 前置字元。

  2. 套用資訊清單：

     kubectl apply -f job-2b.yaml
     

  3. 等待工作完成：

     kubectl wait --for=condition=complete --timeout=180s job/data-loader-2b
     

     輸出看起來類似以下內容：

     job.batch/data-loader-2b condition met
     

  4. 確認工作是否順利完成：

     kubectl get job/data-loader-2b
     

     輸出看起來類似以下內容：

     NAME             COMPLETIONS   DURATION   AGE
     data-loader-2b   1/1           ##s        #m##s
     

  5. 查看工作的記錄：

     kubectl logs --follow job/data-loader-2b
     

  這項工作會將檢查點上傳至 gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000。

  Gemma 7B-it

  1. 建立下列 job-7b.yaml 資訊清單：

     apiVersion: v1
     kind: ConfigMap
     metadata:
       name: fetch-model-scripts
     data:
       fetch_model.sh: |-
         #!/usr/bin/bash -x
         pip install kaggle --break-system-packages && \
     
         MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
         VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \
     
         mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
         kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
         echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \
     
         gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
         echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
     ---
     apiVersion: batch/v1
     kind: Job
     metadata:
       name: data-loader-7b
       labels:
         app: data-loader-7b
     spec:
       ttlSecondsAfterFinished: 120
       template:
         metadata:
           labels:
             app: data-loader-7b
         spec:
           restartPolicy: OnFailure
           containers:
           - name: gcloud
             image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
             command:
             - /scripts/fetch_model.sh
             env:
             - name: BUCKET_NAME
               value: CHECKPOINTS_BUCKET_NAME
             - name: KAGGLE_CONFIG_DIR
               value: /kaggle
             - name: MODEL_PATH
               value: "google/gemma/pax/7b-it/2"
             volumeMounts:
             - mountPath: "/kaggle/"
               name: kaggle-credentials
               readOnly: true
             - mountPath: "/scripts/"
               name: scripts-volume
               readOnly: true
           volumes:
           - name: kaggle-credentials
             secret:
               defaultMode: 0400
               secretName: kaggle-secret
           - name: scripts-volume
             configMap:
               defaultMode: 0700
               name: fetch-model-scripts

     將 CHECKPOINTS_BUCKET_NAME 替換為您在「建立 Cloud Storage 值區」一節中建立的值區名稱。請務必加入 gs:// 前置字元。

  2. 套用資訊清單：

     kubectl apply -f job-7b.yaml
     

  3. 等待工作完成：

     kubectl wait --for=condition=complete --timeout=360s job/data-loader-7b
     

     輸出看起來類似以下內容：

     job.batch/data-loader-7b condition met
     

  4. 確認工作是否順利完成：

     kubectl get job/data-loader-7b
     

     輸出看起來類似以下內容：

     NAME             COMPLETIONS   DURATION   AGE
     data-loader-7b   1/1           ##s        #m##s
     

  5. 查看工作的記錄：

     kubectl logs --follow job/data-loader-7b
     

  這項工作會將檢查點上傳至 gs://CHECKPOINTS_BUCKET_NAME/gemma_7b_it/checkpoint_00000000。

  公開 Saxml HTTP 伺服器

  您可以透過部署 Saxml HTTP 伺服器時建立的 ClusterIP 服務，存取 Saxml HTTP 伺服器。ClusterIP 服務只能從叢集內部存取。因此，如要從叢集外部存取 Service，請完成下列步驟：

  1. 建立通訊埠轉送工作階段：

     kubectl port-forward service/sax-http-svc 8888:8888
     

  2. 開啟新終端機並執行下列指令，確認您可以存取 Saxml HTTP 伺服器：

     curl -s localhost:8888
     

     輸出看起來類似以下內容：

     {
         "Message": "HTTP Server for SAX Client"
     }
     

  Saxml HTTP 伺服器會封裝 Saxml 系統的用戶端介面，並透過一組 REST API 公開。您可以使用這些 API 發布、管理 Gemma 2B 和 Gemma 7B 模型，以及與這些模型互動。

  發布 Gemma 模型

  接著，您可以將 Gemma 模型發布至在 TPU 節點集區中執行的模型伺服器。您可以使用 Saxml HTTP 伺服器的 publish API 發布模型。請按照下列步驟發布 Gemma 2B 或 7B 參數模型。

  如要進一步瞭解 Saxml HTTP 伺服器的 API，請參閱「Saxml HTTP API」。

  Gemma 2B-it

  1. 確認連接埠轉送工作階段仍處於有效狀態：

     curl -s localhost:8888
     

  2. 發布 Gemma 2B 參數：

     curl --request POST \
     --header "Content-type: application/json" \
     -s \
     localhost:8888/publish \
     --data \
     '{
         "model": "/sax/test/gemma2bfp16",
         "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
         "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
         "replicas": "1"
     }'
     

     輸出看起來類似以下內容：

     {
         "model": "/sax/test/gemma2bfp16",
         "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
         "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
         "replicas": 1
     }
     

     請參閱下一個步驟，瞭解如何監控部署作業的進度。

  3. 在 sax-model-server-v5e-1x1 部署作業的模型伺服器 Pod 中觀察記錄，監控進度。

     kubectl logs --follow deployment/sax-model-server-v5e-1x1
     

     部署作業最多可能需要五分鐘才能完成。等待畫面顯示類似以下的訊息：

     I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
     I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma2bfp16
     

  4. 顯示模型資訊，確認您可以存取模型：

     curl --request GET \
     --header "Content-type: application/json" \
     -s \
     localhost:8888/listcell \
     --data \
     '{
         "model": "/sax/test/gemma2bfp16"
     }'
     

     輸出看起來類似以下內容：

     {
         "model": "/sax/test/gemma2bfp16",
         "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
         "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
         "max_replicas": 1,
         "active_replicas": 1
     }
     

  成功： 您已發布 Gemma 2B 模型。現在可以與模型互動。

  Gemma 7B-it

  1. 確認連接埠轉送工作階段仍處於有效狀態：

     curl -s localhost:8888
     

  2. 發布 Gemma 7B 參數：

     curl --request POST \
     --header "Content-type: application/json" \
     -s \
     localhost:8888/publish \
     --data \
     '{
         "model": "/sax/test/gemma7bfp16",
         "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
         "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
         "replicas": "1"
     }'
     

     輸出看起來類似以下內容：

     {
         "model": "/sax/test/gemma7bfp16",
         "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
         "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
         "replicas": 1
     }
     

     請參閱下一個步驟，瞭解如何監控部署作業的進度。

  3. 在 sax-model-server-v5e-2x2 部署作業的模型伺服器 Pod 中觀察記錄，監控進度。

     kubectl logs --follow deployment/sax-model-server-v5e-2x2
     

     等待畫面顯示類似以下的訊息：

     I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
     I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma7bfp16
     

  4. 顯示模型資訊，確認模型已發布：

     curl --request GET \
     --header "Content-type: application/json" \
     -s \
     localhost:8888/listcell \
     --data \
     '{
         "model": "/sax/test/gemma7bfp16"
     }'
     

     輸出結果會與下列內容相似：

     {
         "model": "/sax/test/gemma7bfp16",
         "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
         "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
         "max_replicas": 1,
         "active_replicas": 1
     }
     

  成功： 您已發布 Gemma 7B 模型。現在可以與模型互動。

  使用模型

  您可以與 Gemma 2B 或 7B 模型互動。使用 Saxml HTTP 伺服器的 generate API 將提示傳送至模型。

  Gemma 2B-it

  使用 Saxml HTTP 伺服器的 generate 端點，提供提示要求：

  curl --request POST \
  --header "Content-type: application/json" \
  -s \
  localhost:8888/generate \
  --data \
  '{
    "model": "/sax/test/gemma2bfp16",
    "query": "What are the top 5 most popular programming languages?"
  }'
  

  以下是模型回覆的範例。 實際輸出內容會因您提供的提示而異：

  [
      [
          "\n\n1. **Python**\n2. **JavaScript**\n3. **Java**\n4. **C++**\n5. **Go**",
          -3.0704939365386963
      ]
  ]
  

  您可以搭配不同的 query 參數執行指令。您也可以使用 generate API 修改額外參數，例如 temperature、top_k、topc_p。如要進一步瞭解 Saxml HTTP 伺服器的 API，請參閱「Saxml HTTP API」。

  Gemma 7B-it

  使用 Saxml HTTP 伺服器的 generate 端點，提供提示要求：

  curl --request POST \
  --header "Content-type: application/json" \
  -s \
  localhost:8888/generate \
  --data \
  '{
    "model": "/sax/test/gemma7bfp16",
    "query": "What are the top 5 most popular programming languages?"
  }'
  

  以下是模型回覆的範例。您提供的每個提示輸出內容可能有所不同：

  [
      [
          "\n\n**1. JavaScript**\n\n* Most widely used language on the web.\n* Used for front-end development, such as websites and mobile apps.\n* Extensive libraries and frameworks available.\n\n**2. Python**\n\n* Known for its simplicity and readability.\n* Versatile, used for various tasks, including data science, machine learning, and web development.\n* Large and active community.\n\n**3. Java**\n\n* Object-oriented language widely used in enterprise applications.\n* Used for web applications, mobile apps, and enterprise software.\n* Strong ecosystem and support.\n\n**4. Go**\n\n",
          -16.806324005126953
      ]
  ]
  

  您可以搭配不同的 query 參數執行指令。您也可以使用 generate API 修改額外參數，例如 temperature、top_k、topc_p。如要進一步瞭解 Saxml HTTP 伺服器的 API，請參閱「Saxml HTTP API」。

  取消發布模型

  如要取消發布模型，請按照下列步驟操作：

  Gemma 2B-it

  如要取消發布 Gemma 2B-it 模型，請執行下列指令：

  curl --request POST \
  --header "Content-type: application/json" \
  -s \
  localhost:8888/unpublish \
  --data \
  '{
      "model": "/sax/test/gemma2bfp16"
  }'
  

  輸出看起來類似以下內容：

  {
      "model": "/sax/test/gemma2bfp16"
  }
  

  您可以搭配 query 參數傳遞的不同提示來執行指令。

  Gemma 7B-it

  如要取消發布 Gemma 7B-it 模型，請執行下列指令：

  curl --request POST \
  --header "Content-type: application/json" \
  -s \
  localhost:8888/unpublish \
  --data \
  '{
      "model": "/sax/test/gemma7bfp16"
  }'
  

  輸出看起來類似以下內容：

  {
      "model": "/sax/test/gemma7bfp16"
  }
  

  您可以搭配 query 參數傳遞的不同提示來執行指令。

  排解問題

  • 如果收到 Empty reply from server 訊息，可能是容器尚未完成下載模型資料。再次檢查 Pod 的記錄，確認是否有 Connected 訊息，這表示模型已準備好提供服務。
  • 如果看到 Connection refused，請確認連接埠轉送功能是否已啟用。

  清除所用資源

  如要避免系統向您的 Google Cloud 帳戶收取本教學課程中所用資源的相關費用，請刪除含有該項資源的專案，或者保留專案但刪除個別資源。

  刪除已部署的資源

  如要避免系統向您的 Google Cloud 帳戶收取本指南所建立資源的費用，請執行下列指令：

  gcloud container clusters delete ${CLUSTER_NAME} --location=${LOCATION}
  gcloud iam service-accounts delete --quiet wi-sax@${PROJECT_ID}.iam.gserviceaccount.com
  gcloud storage rm --recursive gs://ADMIN_BUCKET_NAME
  gcloud storage rm --recursive gs://CHECKPOINTS_BUCKET_NAME
  

  更改下列內容：

  • ADMIN_BUCKET_NAME：儲存 Saxml Admin 伺服器的 Cloud Storage 值區名稱。
  • CHECKPOINTS_BUCKET_NAME：儲存模型檢查點的 Cloud Storage 值區名稱。

  後續步驟

  • 進一步瞭解 GKE 中的 TPU。
  • 探索 Saxml GitHub 存放區，包括 Saxml HTTP API。
  • 探索 Vertex AI Model Garden。
  • 瞭解如何運用 GKE 平台的自動化調度管理功能，執行最佳化的 AI/機器學習工作負載。
  • 探索 Google Cloud 的參考架構、圖表和最佳做法。 歡迎瀏覽我們的雲端架構中心。