Servir modelos abiertos de Gemma mediante TPUs en GKE con Saxml

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:

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.

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:

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 un enlace de política de gestión de identidades y accesos a tu cuenta de servicio de gestión de identidades y accesos para que pueda 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.

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 de la implementación 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
    

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 de la implementación 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
    

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 de la implementación 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
    

Descarga 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. Para publicar un modelo, se usa la API publish del servidor HTTP de Saxml. 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
    }
    

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
    }
    

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