Disponibilizar modelos abertos do Gemma usando TPUs no GKE com Saxml

Padrão

---

Neste guia, mostramos como exibir um modelo de linguagem grande (LLM) de modelos abertos do Gemma usando Unidades de Processamento de Tensor (TPUs) no Google Kubernetes Engine (GKE) com o Saxml. Neste guia, faça o download dos modelos Gemma ajustados por instruções de parâmetros 2B e 7B para o Cloud Storage e os implante em um cluster do GKE Standard usando contêineres que executam o Saxml.

Este guia é um bom ponto de partida se você precisar da escalonabilidade, resiliência e economia oferecidos pelos recursos do Kubernetes ao implantar o modelo no Saxml.

Se você precisa de uma plataforma de IA gerenciada unificada para criar e disponibilizar modelos de ML rapidamente e de maneira econômica, recomendamos testar nossa solução de implantação da Vertex AI.

Contexto

Ao disponibilizar o Gemma usando TPUs no GKE com o Saxml, é possível implementar uma solução de veiculação de inferência robusta e pronta para produção com todos os benefícios do Kubernetes, incluindo escalonabilidade eficiente e maior disponibilidade. Esta seção descreve as principais tecnologias usadas neste tutorial.

Gemma

O Gemma é um conjunto de modelos de IA generativa leve e abertamente lançados com licença aberta. Esses modelos de IA estão disponíveis para execução em aplicativos, hardware, dispositivos móveis ou serviços hospedados. Use o Modelos Gemma para geração de texto. É possível ajustá-los para tarefas especializadas.

Para saber mais, consulte a documentação do Gemma.

TPUs

TPUs são circuitos integrados de aplicação específica (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. Eles são usados para acelerar frameworks de processamento de dados, como TensorFlow, PyTorch e JAX.

Antes de usar TPUs no GKE, recomendamos que você conclua o seguinte programa de aprendizado:

1. Saiba mais sobre a disponibilidade atual da versão da TPU com a arquitetura do sistema do Cloud TPU.
2. Saiba mais sobre TPUs no GKE.

Este tutorial serve os modelos Gemma 2B e Gemma 7B. O GKE hospeda esses modelos nos seguintes pools de nós da TPU v5e de host único:

• Gemma 2B: modelo ajustado por instruções hospedado em um pool de nós da TPU v5e com topologia 1x1 que representa um chip de TPU. O tipo de máquina dos nós é ct5lp-hightpu-1t.
• Gemma 7B: modelo ajustado por instruções hospedado em um pool de nós da TPU v5e com topologia 2x2 que representa quatro chips de TPU. O tipo de máquina dos nós é ct5lp-hightpu-4t.

Saxml

O Saxml é um sistema experimental que exibe modelos Paxml, JAX e PyTorch para inferência. O sistema Saxml inclui os seguintes componentes:

• Célula Saxml ou cluster Sax: um servidor administrador e um grupo de servidores de modelo. O servidor de administrador monitora os servidores de modelos, atribui modelos publicados a servidores de modelo para exibição e ajuda os clientes a localizar servidores de modelos que atendem a modelos publicados específicos.
• Cliente Saxml: a interface de programação voltada ao usuário para o sistema Saxml. O cliente Saxml inclui uma ferramenta de linha de comando (saxutil) e um conjunto de bibliotecas de cliente em Python, C++ e Go.

Neste tutorial, você também vai usar o servidor HTTP Saxml. O servidor HTTP Saxml é um servidor HTTP personalizado que encapsula a biblioteca de cliente em Python Saxmldo e expõe as APIs REST para interagir com o sistema Saxml. As APIs REST incluem endpoints para publicar, listar, cancelar a publicação de modelos e gerar previsões.

Objetivos

Este tutorial é destinado a clientes de IA generativa que usam o JAX, além de usuários novos ou existentes do GKE que querem usar os recursos de orquestração de contêineres do Kubernetes para disponibilizar o Gemma, como engenheiros de ML, engenheiros de MLOps (DevOps) e administradores de plataformas.

Este tutorial inclui as etapas a seguir:

1. Prepare um cluster do GKE Standard com a topologia de TPU recomendada com base nas características do modelo.
2. Implantar componentes Saxml no GKE.
3. Conseguir e publicar o modelo de parâmetro Gemma 2B ou Gemma 7B.
4. Disponibilizar e interagir com os modelos publicados.

Arquitetura

Nesta seção, descrevemos a arquitetura do GKE usada neste tutorial. A arquitetura compreende um cluster do GKE Standard que provisiona TPUs e hospeda componentes Saxml para implantar e disponibilizar modelos Gemma 2B ou 7B. O diagrama a seguir mostra os componentes dessa arquitetura:

Essa arquitetura inclui os seguintes componentes:

• Um cluster zonal do GKE Standard.
• Um pool de nós de fração de TPU de host único que depende do modelo Gemma que você quer disponibilizar:
  • Gemma 2B: configurado com uma TPU v5e com uma topologia 1x1. Uma instância do servidor do modelo Saxml está configurada para usar esse pool de nós.
  • Gemma 7B: configurado com uma TPU v5e com uma topologia 2x2. Uma instância do servidor do modelo Saxml está configurada para usar esse pool de nós.
• Um pool de nós de CPU padrão em que o servidor de administração Saxml e o servidor HTTP Saxml são implantados.
• Dois buckets do Cloud Storage:
  • Um bucket do Cloud Storage armazena o estado gerenciado por um servidor administrador.
  • Um bucket do Cloud Storage armazena os checkpoints do modelo Gemma.

Essa arquitetura tem as seguintes características:

• Um Artifact Registry público gerencia as imagens de contêineres para os componentes do Saxml.
• O cluster do GKE usa a federação de identidade da carga de trabalho do GKE. Todos os componentes do Saxml usam uma federação de identidade da carga de trabalho que integra uma conta de serviço do IAM para acessar serviços externos, como buckets do Cloud Storage.
• Os registros gerados pelos componentes do Saxml são integrados ao Cloud Logging.
• É possível usar o Cloud Monitoring para analisar as métricas de desempenho dos pools de nós do GKE, como as criadas neste tutorial.

Antes de começar

• 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

  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 colunn 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.

     Acessar o IAM
  2. Selecionar um projeto.
  3. Clique em CONCEDER ACESSO.

  4. No campo Novos principais, insira seu identificador de usuário. Normalmente, é o endereço de e-mail de uma Conta do Google.

  5. Na lista Selecionar um papel, escolha um.
  6. Para conceder outros papéis, clique em Adicionar outro papel e adicione cada papel adicional.
  7. Clique em Salvar.

• Verifique se você tem cotas suficientes para cinco chips TPU v5e. Neste tutorial, você usará as instâncias sob demanda.
• Crie uma conta da Kaggle, caso ainda não tenha uma.

Preparar o ambiente para o Gemma

Iniciar o Cloud Shell

Neste tutorial, você usará o Cloud Shell para gerenciar recursos hospedados no Google Cloud. O Cloud Shell vem pré-instalado com o software necessário para este tutorial, incluindo kubectl e gcloud CLI.

1. No console do Google Cloud, inicie uma instância do Cloud Shell:
   Abrir o Cloud Shell

2. Defina as variáveis de ambiente padrão:

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

   Substitua os seguintes valores:

   • PROJECT_ID: o ID do projeto do Google Cloud.
   • LOCATION: o nome da zona do Compute Engine em que os tipos de máquina TPU v5e estão disponíveis.

Criar um cluster do GKE Standard

Nesta seção, você vai criar o cluster do GKE e o pool de nós.

Gemma 2B-it

Use o Cloud Shell para fazer o seguinte:

1. Crie um cluster padrão que use a federação de identidade da carga de trabalho do 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}
   

   A criação do cluster pode levar vários minutos.

2. Crie um pool de nós de TPU v5e com uma topologia 1x1 e um nó:

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

   Você exibe o modelo Gemma 2B neste pool de nós.

Gemma 7B-it

Use o Cloud Shell para fazer o seguinte:

1. Crie um cluster padrão que use a federação de identidade da carga de trabalho do 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}
   

   A criação do cluster pode levar vários minutos.

2. Crie um pool de nós de TPU v5e com uma topologia 2x2 e um nó:

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

   Você exibe o modelo Gemma 7B neste pool de nós.

Criar os buckets do Cloud Storage

Crie dois buckets do Cloud Storage para gerenciar o estado do servidor Saxml Admin e os checkpoints do modelo.

No Cloud Shell, execute o seguinte:

1. Crie um bucket do Cloud Storage para armazenar as configurações do servidor administrador do Saxml.

   gcloud storage buckets create gs://ADMIN_BUCKET_NAME
   

   Substitua ADMIN_BUCKET_NAME pelo nome do bucket do Cloud Storage que armazena o servidor Administrador do Saxml.

2. Crie um bucket do Cloud Storage para armazenar checkpoints do modelo:

   gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME
   

   Substitua CHECKPOINTS_BUCKET_NAME pelo nome do bucket do Cloud Storage que armazena os checkpoints do modelo.

Configurar o acesso das cargas de trabalho usando a federação de identidade da carga de trabalho do GKE

Atribua uma Kubernetes ServiceAccount ao aplicativo e configure-a para atuar como uma conta de serviço do IAM.

1. Configure kubectl para se comunicar com o cluster:

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

2. Crie uma conta de serviço do Kubernetes para o aplicativo usar:

   gcloud iam service-accounts create wi-sax
   

3. Adicione uma vinculação de política do IAM para que sua conta de serviço do IAM leia e grave no 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. Permita que a conta de serviço do Kubernetes atue como a conta de serviço do IAM adicionando uma vinculação de política do IAM entre as duas. Essa vinculação permite que a Kubernetes ServiceAccount atue como a conta de serviço do Google.

   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. Anote a conta de serviço do Kubernetes com o endereço de e-mail da conta de serviço do IAM

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

Receber acesso ao modelo

Para ter acesso aos modelos Gemma para implantação no GKE, faça login na plataforma Kaggle, assine o contrato de consentimento de licença e receba um token de API da Kaggle. Neste tutorial, você usa um Secret do Kubernetes para as credenciais do Kaggle.

Assinar o contrato de consentimento de licença

É necessário assinar o contrato de consentimento para usar o Gemma. Siga estas instruções:

1. Acesse a página de consentimento do modelo em Kaggle.com.
2. Faça login no Kaggle, caso ainda não o tenha feito.
3. Clique em Solicitar acesso.
4. Na seção Escolher conta para consentimento, selecione Verificar via conta do Kaggle para usar sua conta do Kaggle para conceder permissão.
5. Aceite os Termos e Condições do modelo.

Gerar um token de acesso

Para acessar o modelo pelo Kaggle, você precisa de um token da API Kaggle.

Siga as etapas abaixo para gerar um novo token, caso ainda não tenha um:

1. No navegador, acesse as configurações da Kaggle.
2. Na seção API, clique em Criar novo token.

O Kaggle faz o download de um arquivo chamado kaggle.json.

Fazer upload do token de acesso no Cloud Shell

No Cloud Shell, é possível fazer upload do token da API Kaggle para seu projeto do Google Cloud:

1. No Cloud Shell, clique em Mais > Fazer upload.
2. Selecione "Arquivo" e clique em Escolher arquivos.
3. Abra o arquivo kaggle.json.
4. Clique em Fazer upload.

Criar um secret do Kubernetes para credenciais do Kaggle

No Cloud Shell, realize as seguintes etapas:

1. Configure kubectl para se comunicar com o cluster:

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

2. Crie um Secret para armazenar as credenciais do Kaggle:

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

Implantar o Saxml

Nesta seção, você vai implantar o servidor de administrador Saxml, os servidores de modelo e o servidor HTTP.

Implantar o servidor de administrador Saxml

1. Crie o seguinte manifesto 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

   Substitua ADMIN_BUCKET_NAME pelo nome do bucket criado na seção Criar buckets do Cloud Storage. Não inclua o prefixo gs://.

2. Aplique o manifesto:

   kubectl apply -f saxml-admin-server.yaml
   

3. Verifique a implantação do servidor de administrador:

   kubectl get deployment
   

   A saída será assim:

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

Implantar o servidor do modelo Saxml

Siga estas instruções para implantar o servidor para os modelos Gemma 2B ou Gemma 7B.

Gemma 2B-it

1. Crie o seguinte manifesto 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

   Substitua ADMIN_BUCKET_NAME pelo nome do bucket criado na seção Criar buckets do Cloud Storage. Não inclua o prefixo gs://.

2. Aplique o manifesto:

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

3. Verifique o status da implantação do servidor de modelo:

   kubectl get deployment
   

   A saída será assim:

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

Sucesso: você implantou um servidor de modelo em um pool de nós da TPU v5e com uma fatia 1 x 1 para exibir o modelo Gemma 2B.

Gemma 7B-it

1. Crie o seguinte manifesto 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

   Substitua ADMIN_BUCKET_NAME pelo nome do bucket criado na seção Criar buckets do Cloud Storage. Não inclua o prefixo gs://.

2. Aplique o manifesto:

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

3. Verifique o status da implantação do servidor de modelo:

   kubectl get deployment
   

   A saída será assim:

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

Sucesso: você implantou um servidor de modelo em um pool de nós da TPU v5e com uma fatia 2x2 para exibir o modelo Gemma 7B.

Implantar o servidor HTTP Saxml

Nesta seção, você implanta o servidor HTTP Saxml e cria um serviço IP de cluster que você usa para acessar o servidor.

1. Crie o seguinte manifesto 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

   Substitua ADMIN_BUCKET_NAME pelo nome do bucket do Cloud Storage que armazena o servidor Administrador do Saxml.

2. Aplique o manifesto:

   kubectl apply -f saxml-http.yaml
   

3. Verifique o status da implantação do servidor HTTP Saxml:

   kubectl get deployment
   

   Gemma 2B-it

   A saída será assim:

     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

   A saída será assim:

     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
   

Fazer o download do checkpoint do modelo

Nesta seção, você executará um job do Kubernetes que busca, faz o download e armazena o checkpoint do modelo. Siga as etapas do modelo Gemma que você quer usar:

Gemma 2B-it

1. Crie o seguinte manifesto 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

   Substitua CHECKPOINTS_BUCKET_NAME pelo nome do bucket criado na seção Criar buckets do Cloud Storage. Não inclua o prefixo gs://.

2. Aplique o manifesto:

   kubectl apply -f job-2b.yaml
   

3. Aguarde a conclusão do job:

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

   A saída será assim:

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

4. Verifique se o job foi concluído com sucesso:

   kubectl get job/data-loader-2b
   

   A saída será assim:

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

5. Visualize os registros do job:

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

O job faz upload do checkpoint para gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000.

Gemma 7B-it

1. Crie o seguinte manifesto 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

   Substitua CHECKPOINTS_BUCKET_NAME pelo nome do bucket criado na seção Criar buckets do Cloud Storage. Inclua o prefixo gs://.

2. Aplique o manifesto:

   kubectl apply -f job-7b.yaml
   

3. Aguarde a conclusão do job:

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

   A saída será assim:

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

4. Verifique se o job foi concluído com sucesso:

   kubectl get job/data-loader-7b
   

   A saída será assim:

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

5. Visualize os registros do job:

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

O job faz upload do checkpoint para gs://CHECKPOINTS_BUCKET_NAME/gemma_7b_it/checkpoint_00000000.

Expor o servidor HTTP Saxml

É possível acessar o servidor HTTP Saxml pelo Serviço ClusterIP que você criou durante a implantação do servidor HTTP Saxml. Os serviços ClusterIP só podem ser acessados dentro do cluster. Portanto, para acessar o serviço de fora do cluster, conclua as seguintes etapas:

1. Estabeleça uma sessão de encaminhamento de portas:

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

2. Para verificar se é possível acessar o servidor HTTP Saxml, abra um novo terminal e execute o seguinte comando:

   curl -s localhost:8888
   

   A saída será assim:

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

No servidor HTTP Saxml, a interface do cliente é encapsulada no sistema Saxml e ela é exposta por um conjunto de APIs REST. Você usará essas APIs para publicar, gerenciar e interagir com os modelos Gemma 2B e Gemma 7B.

Publicar o modelo Gemma

Em seguida, é possível publicar o modelo Gemma em um servidor do modelo que é executado em um pool de nós de frações de TPU. Use a API publish do servidor HTTP Saxml para publicar um modelo. Siga estas etapas para publicar o modelo de parâmetro Gemma 2B ou 7B.

Para saber mais sobre a API do servidor HTTP Saxml, consulte APIs HTTP do Saxml.

Gemma 2B-it

1. Verifique se a sessão de encaminhamento de portas ainda está ativa:

   curl -s localhost:8888
   

2. Publique o 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"
   }'
   

   A saída será assim:

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

   Confira a próxima etapa para monitorar o progresso da implantação.

3. Monitore o progresso observando os registros em um pod de servidor de modelo da implantação sax-model-server-v5e-1x1.

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

   Essa implantação pode levar até cinco minutos para ser concluída. Aguarde até uma mensagem semelhante a esta aparecer:

   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. Verifique se você pode acessar o modelo exibindo as informações dele:

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

   A saída será assim:

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

Sucesso: você publicou um modelo Gemma 2B. Agora é possível interagir com o modelo.

Gemma 7B-it

1. Verifique se a sessão de encaminhamento de portas ainda está ativa:

   curl -s localhost:8888
   

2. Publique o 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"
   }'
   

   A saída será assim:

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

   Confira a próxima etapa para monitorar o progresso da implantação.

3. Monitore o progresso observando os registros em um pod de servidor de modelo da implantação sax-model-server-v5e-2x2.

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

   Aguarde até uma mensagem semelhante a esta aparecer:

   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. Verifique se o modelo foi publicado exibindo as informações dele:

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

   O resultado será assim:

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

Sucesso: você publicou um modelo Gemma 7B. Agora é possível interagir com o modelo.

use o modelo

É possível interagir com os modelos Gemma 2B ou 7B. Use a API generate do servidor HTTP Saxml para enviar um prompt ao modelo.

Gemma 2B-it

Disponibilize uma solicitação de prompt usando o endpoint generate do servidor HTTP 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?"
}'

Confira a seguir um exemplo de resposta do modelo. A saída real varia de acordo com o comando que você disponibiliza:

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

É possível executar o comando com diferentes parâmetros query. Também é possível modificar parâmetros extras, como temperature, top_k e topc_p, usando a API generate. Para saber mais sobre a API do servidor HTTP Saxml, consulte APIs HTTP do Saxml.

Gemma 7B-it

Disponibilize uma solicitação de prompt usando o endpoint generate do servidor HTTP 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?"
}'

Confira a seguir um exemplo de resposta do modelo. A saída pode variar em cada comando exibido:

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

É possível executar o comando com diferentes parâmetros query. Também é possível modificar parâmetros extras, como temperature, top_k e topc_p, usando a API generate. Para saber mais sobre a API do servidor HTTP Saxml, consulte APIs HTTP do Saxml.

Cancelar a publicação do modelo

Siga estas etapas para cancelar a publicação do modelo:

Gemma 2B-it

Para cancelar a publicação do modelo Gemma 2B-it, execute o seguinte comando:

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

A saída será assim:

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

É possível executar o comando com diferentes solicitações transmitidas no parâmetro query.

Gemma 7B-it

Para cancelar a publicação do modelo Gemma 7B-it, execute o seguinte comando:

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

A saída será assim:

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

É possível executar o comando com diferentes solicitações transmitidas no parâmetro query.

Resolver problemas

• Se você receber a mensagem Empty reply from server, é possível que o contêiner não tenha concluído o download dos dados do modelo. Verifique os registros do pod novamente para ver a mensagem Connected, que indica que o modelo está pronto para ser disponibilizado.
• Se aparecer Connection refused, verifique se o encaminhamento de portas está ativo.

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

Excluir os recursos implantados

Para evitar cobranças na sua conta do Google Cloud pelos recursos criados neste guia, execute o seguinte 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

Substitua:

• ADMIN_BUCKET_NAME: o nome do bucket do Cloud Storage que armazena o servidor Administrador do Saxml.
• CHECKPOINTS_BUCKET_NAME: o nome do bucket do Cloud Storage que armazena os checkpoints do modelo.

A seguir

• Saiba mais sobre TPUs no GKE.
• Conheça o repositório do Saxml no GitHub, incluindo as APIs HTTP do Saxml.
• Conhecer o Model Garden da Vertex AI.
• Descubra como executar cargas de trabalho de IA/ML otimizadas com os recursos de orquestração da plataforma GKE.
• Confira arquiteturas de referência, diagramas, tutoriais e práticas recomendadas do Google Cloud. Confira o Centro de arquitetura do Cloud.