Servir modelos abiertos de Gemma mediante TPUs en GKE con Saxml

Estándar

---

Precaución: A partir del 24 de abril del 2025, se reducirá la prioridad de Saxml en GKE. Esto significa que el proyecto no recibirá más actualizaciones. Las implementaciones de Saxml seguirán funcionando como hasta ahora sin interrupciones. Te recomendamos que migres a JetStream, el framework de inferencia de código abierto actualizado de Google para ofrecer LLMs de alto rendimiento en TPUs y GPUs. JetStream ofrece la creación de lotes y la cuantización continuas para mejorar el rendimiento y la eficiencia de la memoria. Para ver un ejemplo de migración, consulta Sirve Gemma mediante TPUs en GKE con JetStream.

En este tutorial se muestra cómo desplegar y servir un modelo de lenguaje extenso (LLM) Gemma 2 mediante TPUs en GKE con el framework de servicio Saxml. Este tutorial proporciona una base para comprender y explorar la implementación práctica de LLMs para la inferencia en un entorno de Kubernetes gestionado. Despliega un contenedor precompilado con Saxml en GKE. También configurarás GKE para que cargue los pesos de Gemma 2B y 7B desde Cloud Storage en el tiempo de ejecución.

Este tutorial está dirigido a ingenieros de aprendizaje automático, administradores y operadores de plataformas, y especialistas en datos e IA que quieran usar las funciones de orquestación de contenedores de Kubernetes para servir LLMs. Para obtener más información sobre los roles habituales y las tareas de ejemplo a las que hacemos referencia en el Google Cloud contenido, consulta Roles y tareas habituales de los usuarios de GKE.

Antes de leer esta página, asegúrese de que conoce los siguientes conceptos:

• Disponibilidad de la versión actual de TPU con la arquitectura del sistema de TPU de Cloud
• TPUs en GKE

Si necesitas una plataforma de IA gestionada unificada para crear y ofrecer modelos de aprendizaje automático rápidamente y de forma rentable, te recomendamos que pruebes nuestra solución de implementación Vertex AI.

Fondo

En esta sección se describen las tecnologías clave que se usan en este tutorial.

Gemma

Gemma es un conjunto de modelos de IA generativa ligeros y disponibles públicamente que se han lanzado con una licencia abierta. Estos modelos de IA se pueden ejecutar en tus aplicaciones, hardware, dispositivos móviles o servicios alojados. Puedes usar los modelos de Gemma para generar texto, además de ajustarlos para tareas especializadas.

Para obtener más información, consulta la documentación de Gemma.

TPUs

Las TPUs son circuitos integrados para aplicaciones específicas (ASIC) desarrollados a medida por Google que se utilizan para acelerar frameworks de procesamiento de datos como TensorFlow, PyTorch y JAX.

En este tutorial se usan los modelos Gemma 2B y Gemma 7B. GKE aloja estos modelos en los siguientes grupos de nodos de TPU v5e de un solo host:

• Gemma 2B: modelo ajustado para instrucciones alojado en un grupo de nodos de TPU v5e con una topología 1x1 que representa un chip de TPU. El tipo de máquina de los nodos es ct5lp-hightpu-1t.
• Gemma 7B: modelo ajustado para instrucciones alojado en un pool de nodos de TPU v5e con una topología 2x2 que representa cuatro chips de TPU. El tipo de máquina de los nodos es ct5lp-hightpu-4t.

Saxml

Saxml es un sistema experimental que ofrece modelos de Paxml, JAX y PyTorch para la inferencia. El sistema Saxml incluye los siguientes componentes:

• Celda Saxml o clúster Sax: un servidor de administración y un grupo de servidores de modelos. El servidor de administración monitoriza los servidores de modelos, asigna modelos publicados a los servidores de modelos para que los sirvan y ayuda a los clientes a localizar los servidores de modelos que sirven modelos publicados específicos.
• Cliente de Saxml: interfaz de programación orientada al usuario del sistema Saxml. El cliente de Saxml incluye una herramienta de línea de comandos (saxutil) y un conjunto de bibliotecas de cliente en Python, C++ y Go.

En este tutorial, también se usa el servidor HTTP de Saxml. El servidor HTTP de Saxml es un servidor HTTP personalizado que encapsula la biblioteca de cliente de Python de Saxml y expone APIs REST para interactuar con el sistema Saxml. Las APIs REST incluyen endpoints para publicar, listar y anular la publicación de modelos, así como para generar predicciones.

Objetivos

1. Prepara un clúster de GKE Standard con la topología de TPU recomendada en función de las características del modelo.
2. Despliega los componentes de Saxml en GKE.
3. Obtén y publica el modelo de parámetros Gemma 2B o Gemma 7B.
4. Ofrecer e interactuar con los modelos publicados.

Arquitectura

En esta sección se describe la arquitectura de GKE que se usa en este tutorial. La arquitectura incluye un clúster estándar de GKE que aprovisiona TPUs y aloja componentes de Saxml para desplegar y servir modelos de Gemma 2B o 7B. En el siguiente diagrama se muestran los componentes de esta arquitectura:

Esta arquitectura incluye los siguientes componentes:

• Un clúster zonal de GKE Standard.
• Un grupo de nodos de segmento de TPU de un solo host que depende del modelo de Gemma que quieras servir:
  • Gemma 2B: configurada con una TPU v5e con una topología 1x1. Se ha configurado una instancia del servidor de modelos de Saxml para usar este grupo de nodos.
  • Gemma 7B: configurada con una TPU v5e con una topología 2x2. Se ha configurado una instancia del servidor de modelos de Saxml para usar este grupo de nodos.
• Un grupo de nodos de CPU predeterminado en el que se implementan el servidor de administrador de Saxml y el servidor HTTP de Saxml.
• Dos cubos de Cloud Storage:
  • Un segmento de Cloud Storage almacena el estado gestionado por un servidor de administrador.
  • Un segmento de Cloud Storage almacena los puntos de control del modelo Gemma.

Esta arquitectura tiene las siguientes características:

• Un Artifact Registry público gestiona las imágenes de contenedor de los componentes de Saxml.
• El clúster de GKE usa Workload Identity Federation for GKE. Todos los componentes de Saxml usan una federación de identidades de carga de trabajo que integra una cuenta de servicio de gestión de identidades y accesos para acceder a servicios externos, como los contenedores de Cloud Storage.
• Los registros generados por los componentes de Saxml se integran en Cloud Logging.
• Puedes usar Cloud Monitoring para analizar las métricas de rendimiento de los grupos de nodos de GKE, como los que se crean en este tutorial.

Antes de empezar

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.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

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

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the API

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

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

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

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

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.

     Ir a IAM
  2. Selecciona el proyecto.
  3. Haz clic en Conceder acceso.

  4. En el campo Nuevos principales, introduce tu identificador de usuario. Normalmente, se trata de la dirección de correo de una cuenta de Google.

  5. En la lista Selecciona un rol, elige un rol.
  6. Para conceder más roles, haz clic en Añadir otro rol y añade cada rol adicional.
  7. Haz clic en Guardar.
  • Asegúrate de que tienes suficientes cuotas para 5 chips TPU v5e. En este tutorial, usarás instancias bajo demanda.
  • Crea una cuenta de Kaggle si aún no tienes una.

  Preparar el entorno para Gemma

  Abrir Cloud Shell

  En este tutorial, usarás Cloud Shell para gestionar los recursos alojados en Google Cloud. Cloud Shell tiene preinstalado el software que necesitas para este tutorial, como kubectl y la CLI de gcloud.

  1. En la Google Cloud consola, inicia una instancia de Cloud Shell:
     Abrir Cloud Shell

  2. Define las variables de entorno predeterminadas:

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

     Sustituye los siguientes valores:

     • PROJECT_ID: tu Google Cloud ID de proyecto.
     • LOCATION: nombre de la zona de Compute Engine en la que están disponibles los tipos de máquina TPU v5e.

  Crear un clúster estándar de GKE

  En esta sección, creará el clúster de GKE y el grupo de nodos.

  Gemma 2B-it

  Usa Cloud Shell para hacer lo siguiente:

  1. Crea un clúster estándar que use Workload Identity Federation para GKE:

     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}
     

     La creación del clúster puede tardar varios minutos.

  2. Crea un grupo de nodos de TPU v5e con una topología 1x1 y un nodo:

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

     Ofreces el modelo Gemma 2B en este grupo de nodos.

  Gemma 7B-it

  Usa Cloud Shell para hacer lo siguiente:

  1. Crea un clúster estándar que use Workload Identity Federation para GKE:

     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}
     

     La creación del clúster puede tardar varios minutos.

  2. Crea un grupo de nodos de TPU v5e con una topología 2x2 y un nodo:

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

     Ofreces el modelo Gemma 7B en este grupo de nodos.

  Crear los segmentos de Cloud Storage

  Crea dos segmentos de Cloud Storage para gestionar el estado del servidor de administrador de Saxml y los puntos de control del modelo.

  En Cloud Shell, ejecuta lo siguiente:

  1. Crea un segmento de Cloud Storage para almacenar las configuraciones del servidor de administrador de Saxml.

     gcloud storage buckets create gs://ADMIN_BUCKET_NAME
     

     Sustituye ADMIN_BUCKET_NAME por el nombre del segmento de Cloud Storage que almacena el servidor de administrador de Saxml.

  2. Crea un segmento de Cloud Storage para almacenar los puntos de control del modelo:

     gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME
     

     Sustituye CHECKPOINTS_BUCKET_NAME por el nombre del segmento de Cloud Storage que almacena los puntos de control del modelo.

  Configurar el acceso de las cargas de trabajo mediante Workload Identity Federation para GKE

  Asigna una cuenta de servicio de Kubernetes a la aplicación y configura esa cuenta de servicio de Kubernetes para que actúe como cuenta de servicio de gestión de identidades y accesos.

  1. Configura kubectl para que se comunique con tu clúster:

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

  2. Crea una cuenta de servicio de Kubernetes para que la use tu aplicación:

     gcloud iam service-accounts create wi-sax
     

  3. Añade una vinculación de política de gestión de identidades y accesos a tu cuenta de servicio de gestión de identidades y accesos para leer y escribir en 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. Permite que la cuenta de servicio de Kubernetes suplante la identidad de la cuenta de servicio de gestión de identidades y accesos añadiendo un enlace de política de gestión de identidades y accesos entre las dos cuentas de servicio. Este enlace permite que la cuenta de servicio de Kubernetes actúe como cuenta de servicio de gestión de identidades y accesos:

     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. Anota la cuenta de servicio de Kubernetes con la dirección de correo de la cuenta de servicio de gestión de identidades y accesos:

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

  Acceder al modelo

  Para acceder a los modelos de Gemma e implementarlos en GKE, debes iniciar sesión en la plataforma Kaggle, firmar el contrato de licencia y obtener un token de API de Kaggle. En este tutorial, usarás un secreto de Kubernetes para las credenciales de Kaggle.

  Firmar el contrato de consentimiento de licencia

  Debes firmar el acuerdo de consentimiento para usar Gemma. Te las indicamos a continuación:

  1. Accede a la página de consentimiento del modelo en Kaggle.com.
  2. Inicia sesión en Kaggle si aún no lo has hecho.
  3. Haz clic en Solicitar acceso.
  4. En la sección Choose Account for Consent (Elegir cuenta para el consentimiento), selecciona Verify via Kaggle Account (Verificar con la cuenta de Kaggle) para usar tu cuenta de Kaggle y dar el consentimiento.
  5. Acepta los Términos y Condiciones del modelo.

  Generar un token de acceso

  Para acceder al modelo a través de Kaggle, necesitas un token de la API de Kaggle.

  Si aún no tienes un token, sigue estos pasos para generarlo:

  1. En tu navegador, ve a Configuración de Kaggle.
  2. En la sección API, haz clic en Create New Token (Crear token).

  Kaggle descarga un archivo llamado kaggle.json.

  Subir el token de acceso a Cloud Shell

  En Cloud Shell, puedes subir el token de la API de Kaggle a tu proyecto: Google Cloud

  1. En Cloud Shell, haz clic en Más > Subir.
  2. Selecciona Archivo y haz clic en Elegir archivos.
  3. Abre el archivo kaggle.json.
  4. Haz clic en Subir.

  Crear un secreto de Kubernetes para las credenciales de Kaggle

  En Cloud Shell, sigue estos pasos:

  1. Configura kubectl para que se comunique con tu clúster:

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

  2. Crea un secreto para almacenar las credenciales de Kaggle:

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

  Desplegar Saxml

  En esta sección, implementará el servidor de administración de Saxml, los servidores de modelos y el servidor HTTP. En este tutorial se usan manifiestos de despliegue de Kubernetes. Un Deployment es un objeto de la API de Kubernetes que te permite ejecutar varias réplicas de pods distribuidas entre los nodos de un clúster.

  Implementar el servidor de administración de Saxml

  En esta sección, implementará el servidor de administración de Saxml.

  1. Crea el siguiente archivo de manifiesto 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

     Sustituye ADMIN_BUCKET_NAME por el nombre del segmento que has creado en la sección Crear segmentos de Cloud Storage. No incluya el prefijo gs://.

  2. Aplica el archivo de manifiesto:

     kubectl apply -f saxml-admin-server.yaml
     

  3. Verifica la implementación del servidor de administración:

     kubectl get deployment
     

     El resultado es similar al siguiente:

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

  Desplegar el servidor de modelos de Saxml

  Sigue estas instrucciones para implementar el servidor de modelos de Gemma 2B o Gemma 7B.

  Gemma 2B-it

  1. Crea el siguiente archivo de manifiesto 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

     Sustituye ADMIN_BUCKET_NAME por el nombre del segmento que has creado en la sección Crear segmentos de Cloud Storage. No incluya el prefijo gs://.

  2. Aplica el archivo de manifiesto:

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

  3. Verifica el estado del despliegue del servidor de modelos:

     kubectl get deployment
     

     El resultado es similar al siguiente:

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

  Correcto: has desplegado un servidor de modelos en un grupo de nodos de TPU v5e con un slice de 1x1 para servir el modelo Gemma 2B.

  Gemma 7B-it

  1. Crea el siguiente archivo de manifiesto 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

     Sustituye ADMIN_BUCKET_NAME por el nombre del segmento que has creado en la sección Crear segmentos de Cloud Storage. No incluya el prefijo gs://.

  2. Aplica el archivo de manifiesto:

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

  3. Verifica el estado del despliegue del servidor de modelos:

     kubectl get deployment
     

     El resultado es similar al siguiente:

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

  Correcto: has desplegado un servidor de modelos en un grupo de nodos de TPU v5e con un 2x2 slice para servir el modelo Gemma 7B.

  Desplegar el servidor HTTP de Saxml

  En esta sección, implementarás el servidor HTTP de Saxml y crearás un servicio de IP de clúster que usarás para acceder al servidor.

  1. Crea el siguiente archivo de manifiesto 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

     Sustituye ADMIN_BUCKET_NAME por el nombre del segmento de Cloud Storage que almacena el servidor de administrador de Saxml.

  2. Aplica el archivo de manifiesto:

     kubectl apply -f saxml-http.yaml
     

  3. Verifica el estado del despliegue del servidor HTTP de Saxml:

     kubectl get deployment
     

     Gemma 2B-it

     El resultado es similar al siguiente:

     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

     El resultado es similar al siguiente:

     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
     

  Descargar el punto de control del modelo

  En esta sección, ejecutarás un trabajo de Kubernetes que obtendrá, descargará y almacenará el punto de control del modelo. Un controlador de trabajo de Kubernetes crea uno o varios pods y se asegura de que ejecuten correctamente una tarea específica.

  Sigue los pasos correspondientes al modelo de Gemma que quieras usar:

  Gemma 2B-it

  1. Crea el siguiente archivo de manifiesto 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

     Sustituye CHECKPOINTS_BUCKET_NAME por el nombre del segmento que has creado en la sección Crear segmentos de Cloud Storage. No incluya el prefijo gs://.

  2. Aplica el archivo de manifiesto:

     kubectl apply -f job-2b.yaml
     

  3. Espera a que se complete el trabajo:

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

     El resultado es similar al siguiente:

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

  4. Comprueba que el trabajo se haya completado correctamente:

     kubectl get job/data-loader-2b
     

     El resultado es similar al siguiente:

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

  5. Para ver los registros de la tarea, haz lo siguiente:

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

  El trabajo sube el punto de control a gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000.

  Gemma 7B-it

  1. Crea el siguiente archivo de manifiesto 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

     Sustituye CHECKPOINTS_BUCKET_NAME por el nombre del segmento que has creado en la sección Crear segmentos de Cloud Storage. Incluye el prefijo gs://.

  2. Aplica el archivo de manifiesto:

     kubectl apply -f job-7b.yaml
     

  3. Espera a que se complete el trabajo:

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

     El resultado es similar al siguiente:

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

  4. Comprueba que el trabajo se haya completado correctamente:

     kubectl get job/data-loader-7b
     

     El resultado es similar al siguiente:

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

  5. Para ver los registros de la tarea, haz lo siguiente:

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

  El trabajo sube el punto de control a gs://CHECKPOINTS_BUCKET_NAME/gemma_7b_it/checkpoint_00000000.

  Exponer el servidor HTTP de Saxml

  Puedes acceder al servidor HTTP de Saxml a través del servicio ClusterIP que creaste al desplegar el servidor HTTP de Saxml. Solo se puede acceder a los servicios ClusterIP desde el clúster. Por lo tanto, para acceder al servicio desde fuera del clúster, sigue estos pasos:

  1. Establece una sesión de redirección de puertos:

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

  2. Comprueba que puedes acceder al servidor HTTP de Saxml abriendo un nuevo terminal y ejecutando el siguiente comando:

     curl -s localhost:8888
     

     El resultado es similar al siguiente:

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

  El servidor HTTP de Saxml encapsula la interfaz de cliente del sistema Saxml y la expone a través de un conjunto de APIs REST. Estas APIs se usan para publicar, gestionar e interactuar con los modelos Gemma 2B y Gemma 7B.

  Publicar el modelo de Gemma

  A continuación, puedes publicar el modelo Gemma en un servidor de modelos que se ejecute en un pool de nodos de segmento de TPU. Utilizas la API publish del servidor HTTP de Saxml para publicar un modelo. Sigue estos pasos para publicar el modelo de parámetros Gemma 2B o 7B.

  Para obtener más información sobre la API del servidor HTTP de Saxml, consulta las APIs HTTP de Saxml.

  Gemma 2B-it

  1. Comprueba que tu sesión de reenvío de puertos siga activa:

     curl -s localhost:8888
     

  2. Publica el parámetro 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"
     }'
     

     El resultado es similar al siguiente:

     {
         "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
     }
     

     Consulta el siguiente paso para monitorizar el progreso de la implementación.

  3. Para monitorizar el progreso, consulta los registros de un pod del servidor de modelos de la implementación de sax-model-server-v5e-1x1.

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

     Este despliegue puede tardar hasta cinco minutos en completarse. Espera hasta que veas un mensaje similar al siguiente:

     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. Verifica que puedes acceder al modelo mostrando su información:

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

     El resultado es similar al siguiente:

     {
         "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
     }
     

  Correcto: has publicado un modelo Gemma 2B. Ahora puedes interactuar con el modelo.

  Gemma 7B-it

  1. Comprueba que tu sesión de reenvío de puertos siga activa:

     curl -s localhost:8888
     

  2. Publica el parámetro 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"
     }'
     

     El resultado es similar al siguiente:

     {
         "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
     }
     

     Consulta el siguiente paso para monitorizar el progreso de la implementación.

  3. Para monitorizar el progreso, consulta los registros de un pod del servidor de modelos de la implementación de sax-model-server-v5e-2x2.

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

     Espera hasta que veas un mensaje similar al siguiente:

     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. Verifica que el modelo se ha publicado mostrando su información:

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

     El resultado debería ser similar al siguiente:

     {
         "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
     }
     

  Correcto: has publicado un modelo Gemma 7B. Ahora puedes interactuar con el modelo.

  Usar el modelo

  Puedes interactuar con los modelos Gemma 2B o 7B. Usa la API generate del servidor HTTP de Saxml para enviar una petición al modelo.

  Gemma 2B-it

  Sirve una solicitud de prompt mediante el endpoint generate del servidor HTTP de Saxml:

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

  A continuación, se muestra un ejemplo de la respuesta del modelo. El resultado real varía en función de la petición que envíes:

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

  Puedes ejecutar el comando con diferentes parámetros de query. También puedes modificar parámetros adicionales, como temperature, top_k y topc_p, mediante la API generate. Para obtener más información sobre la API del servidor HTTP de Saxml, consulta APIs HTTP de Saxml.

  Gemma 7B-it

  Sirve una solicitud de prompt mediante el endpoint generate del servidor HTTP de Saxml:

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

  A continuación, se muestra un ejemplo de la respuesta del modelo. El resultado puede variar en cada petición que sirvas:

  [
      [
          "\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
      ]
  ]
  

  Puedes ejecutar el comando con diferentes parámetros de query. También puedes modificar parámetros adicionales, como temperature, top_k y topc_p, mediante la API generate. Para obtener más información sobre la API del servidor HTTP de Saxml, consulta APIs HTTP de Saxml.

  Anular la publicación del modelo

  Sigue estos pasos para dejar de publicar tu modelo:

  Gemma 2B-it

  Para dejar de publicar el modelo Gemma 2B-it, ejecuta el siguiente comando:

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

  El resultado es similar al siguiente:

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

  Puedes ejecutar el comando con diferentes peticiones que se pasan en el parámetro query.

  Gemma 7B-it

  Para dejar de publicar el modelo Gemma 7B-it, ejecuta el siguiente comando:

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

  El resultado es similar al siguiente:

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

  Puedes ejecutar el comando con diferentes peticiones que se pasan en el parámetro query.

  Solucionar problemas

  • Si recibes el mensaje Empty reply from server, es posible que el contenedor no haya terminado de descargar los datos del modelo. Vuelve a consultar los registros del pod para ver el mensaje Connected, que indica que el modelo está listo para ofrecerse.
  • Si ves Connection refused, verifica que el reenvío de puertos esté activo.

  Limpieza

  Para evitar que los recursos utilizados en este tutorial se cobren en tu cuenta de Google Cloud, elimina el proyecto que contiene los recursos o conserva el proyecto y elimina los recursos.

  Eliminar los recursos desplegados

  Para evitar que se apliquen cargos en tu cuenta de Google Cloud por los recursos que has creado en esta guía, ejecuta el siguiente comando:

  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
  

  Haz los cambios siguientes:

  • ADMIN_BUCKET_NAME: nombre del segmento de Cloud Storage que almacena el servidor de administrador de Saxml.
  • CHECKPOINTS_BUCKET_NAME: nombre del segmento de Cloud Storage que almacena los puntos de control del modelo.

  Siguientes pasos

  • Consulta más información sobre las TPUs en GKE.
  • Consulta el repositorio de GitHub de Saxml, incluidas las APIs HTTP de Saxml.
  • Consulta Vertex AI Model Garden.
  • Descubre cómo ejecutar cargas de trabajo de IA y aprendizaje automático optimizadas con las funciones de orquestación de la plataforma GKE.
  • Consulta arquitecturas de referencia, diagramas y prácticas recomendadas sobre Google Cloud. Consulta nuestro Centro de arquitectura de Cloud.