Ejecutar una carga de trabajo de lote pequeña con TPUs y VMs de inicio flexible

Autopilot Standard

Vista previa — flex-start

Esta función está sujeta a los "Términos de las Ofertas de Acceso Previo a la Disponibilidad General" de la sección Términos Generales de los Servicios de los Términos Específicos de los Servicios. Las funciones previas a la disponibilidad general están disponibles tal cual y pueden tener una compatibilidad y asistencia limitadas. Para obtener más información, consulta las descripciones de las fases de lanzamiento.

En esta guía se explica cómo optimizar el aprovisionamiento de TPUs para cargas de trabajo de entrenamiento a mediana y pequeña escala mediante el uso de máquinas virtuales de inicio flexible. Las VMs de inicio flexible se crean mediante la opción de consumo inicio flexible. En esta guía, usarás VMs de inicio flexible para desplegar una carga de trabajo que consta de un grupo de nodos de segmento de TPU.

Esta guía está dirigida a ingenieros de aprendizaje automático, administradores y operadores de plataformas, y especialistas en datos e IA que estén interesados en usar las funciones de orquestación de contenedores de Kubernetes para ejecutar cargas de trabajo por lotes. Para obtener más información sobre los roles habituales y las tareas de ejemplo a las que hacemos referencia en el contenido, consulta Roles y tareas habituales de los usuarios de GKE. Google Cloud

Precios de inicio flexible

Se recomienda el inicio flexible si tu carga de trabajo requiere recursos aprovisionados dinámicamente según sea necesario, durante un máximo de siete días con reservas a corto plazo, sin una gestión de cuotas compleja y con un acceso rentable. La función de inicio flexible se basa en Dynamic Workload Scheduler y se factura según los precios de Dynamic Workload Scheduler:

• Descuento (hasta el 53%) en vCPUs, GPUs y TPUs.
• Pagas a medida que avanzas.

Antes de empezar

Antes de empezar, asegúrate de que has realizado las siguientes tareas:

• Habilita la API de Google Kubernetes Engine.
Habilitar la API de Google Kubernetes Engine
• Si quieres usar Google Cloud CLI para esta tarea, instálala y, a continuación, inicialízala. Si ya has instalado la gcloud CLI, obtén la versión más reciente ejecutando gcloud components update.
  Nota: Si ya tienes instalada la CLI de gcloud, asegúrate de definir la compute/region propiedad. Si usas principalmente clústeres zonales, define compute/zone en su lugar. Si defines una ubicación predeterminada, puedes evitar errores en gcloud CLI como el siguiente: One of [--zone, --region] must be supplied: Please specify location. Es posible que tengas que especificar la ubicación en determinados comandos si la ubicación de tu clúster es diferente de la predeterminada que hayas definido.

• Comprueba que tienes un clúster de Autopilot o un clúster estándar que ejecute la versión 1.33.0-gke.1712000 o una posterior.
• Comprueba que conoces las limitaciones de la hora de inicio flexible.
• Cuando uses un clúster estándar, comprueba que mantienes al menos un grupo de nodos sin la función de inicio flexible habilitada para que el clúster funcione correctamente.
• Comprueba que tengas cuota para TPUs no garantizadas en las ubicaciones de tus nodos.

Crear un grupo de nodos con inicio flexible

Si usas un clúster en el modo Autopilot, sáltate esta sección y ve a la sección Ejecutar una carga de trabajo por lotes.

Para crear un grupo de nodos con flex-start habilitado en un clúster Standard, usa la CLI de gcloud.

Puedes crear un grupo de nodos de un slice de TPU de un solo host o de varios hosts con inicio flexible:

1. Crea un grupo de nodos con inicio flexible:

   Un solo host

   gcloud container node-pools create NODE_POOL_NAME \
       --cluster=CLUSTER_NAME \
       --location=LOCATION_NAME \
       --node-locations=NODE_ZONES \
       --machine-type=MACHINE_TYPE \
       --reservation-affinity=none \
       --enable-autoscaling \
       --flex-start \
       --num-nodes 0 \
       --min-nodes=0 \
       --max-nodes=1
   

   Haz los cambios siguientes:

   • NODE_POOL_NAME: el nombre que elijas para tu grupo de nodos.
   • CLUSTER_NAME: el nombre del clúster.
   • LOCATION_NAME: la región de computación del plano de control del clúster.
   • NODE_ZONES: lista separada por comas de una o varias zonas en las que GKE crea el grupo de nodos.
   • MACHINE_TYPE: el tipo de máquina que se va a usar en los nodos. Para obtener más información sobre los tipos de máquinas compatibles con TPU, consulta la tabla de la sección Elegir la versión de TPU.

   Varios hosts

   gcloud container node-pools create NODE_POOL_NAME \
       --cluster=CLUSTER_NAME \
       --location=LOCATION_NAME \
       --node-locations=NODE_ZONES \
       --machine-type=MACHINE_TYPE \
       --tpu-topology=TPU_TOPOLOGY \
       --flex-start \
       --enable-autoscaling \
       --num-nodes=0 \
       --max-nodes=2 \
       --reservation-affinity=none \
       --no-enable-autorepair
   

   Haz los cambios siguientes:

   • NODE_POOL_NAME: el nombre que elijas para tu grupo de nodos.
   • CLUSTER_NAME: el nombre del clúster.
   • LOCATION_NAME: la región de computación del plano de control del clúster.
   • NODE_ZONES: lista separada por comas de una o varias zonas en las que GKE crea el grupo de nodos.
   • MACHINE_TYPE: el tipo de máquina que se va a usar en los nodos. Por ejemplo, puedes usar ct6e-standard-4t para TPU Trillium. Para obtener más información sobre los tipos de máquinas disponibles, consulta Elegir la versión de TPU.
   • TPU_TOPOLOGY: la topología física del segmento de TPU. El formato de la topología depende de la versión de la TPU. Para obtener más información sobre las topologías de TPU, consulta la tabla de la sección Elegir una topología.

   El comando anterior usa las siguientes marcas obligatorias al crear un grupo de nodos con VMs de inicio flexible:

   • --enable-autoscaling: flex-start solo proporciona los recursos de computación necesarios cuando se ejecuta la carga de trabajo. Debes definir los siguientes parámetros:

     • --num-nodes=0
     • --min-nodes=0

     • --max-nodes se establece en el número de máquinas virtuales que requiere tu segmento de TPU.

       Por ejemplo, el comando de creación del grupo de nodos puede incluir los siguientes parámetros:

       ...
       --machine-type=ct6e-standard-4t \
       --tpu-topology=4x4 \
       --enable-autoscaling \
       --num-nodes=0 \
       --max-nodes=4 \
       

       Este comando asigna el valor --max-nodes al campo 4 porque una topología 4x4 consta de 16 chips y cada VM ct6e-standard-4t tiene 4 chips.

     El autoescalador de clústeres aumenta el número de nodos hasta el número que requiere tu carga de trabajo. Una vez que se haya completado la carga de trabajo, el autoescalador de clúster se reducirá a cero nodos.

   • --reservation-affinity=none: flex-start no usa ni requiere reservas.

2. Verifica el estado de inicio flexible en el grupo de nodos:

   gcloud container node-pools describe NODE_POOL_NAME \
       --cluster CLUSTER_NAME \
       --location LOCATION_NAME \
       --format="get(config.flexStart)"
   

   Si flex-start está habilitado en el grupo de nodos, el campo flexStart se define como True.

Ejecutar una carga de trabajo por lotes

En esta sección, crearás un trabajo que programará un nodo de TPU con VMs de inicio flexible. Un controlador de trabajo de Kubernetes crea uno o varios pods y se asegura de que ejecuten correctamente una tarea específica.

1. En la consola deGoogle Cloud , inicia una sesión de Cloud Shell haciendo clic en Activar Cloud Shell. Se abrirá una sesión en el panel inferior de la consola Google Cloud .

2. Crea un archivo llamado dws-flex-start.yaml:

   Un solo host

   Usa el siguiente manifiesto para el archivo dws-flex-start.yaml:

   apiVersion: batch/v1
   kind: Job
   metadata:
     name: job-1
   spec:
     template:
       spec:
         nodeSelector:
           cloud.google.com/gke-flex-start: "true"
           cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
           cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
         containers:
         - name: container-1
           image: gcr.io/k8s-staging-perf-tests/sleep:latest
           args: ["3600s"] # Sleep for 1 hour
           resources:
             requests:
                 google.com/tpu: NUM_CHIPS
             limits:
                 google.com/tpu: NUM_CHIPS
         restartPolicy: OnFailure
   

   Varios hosts

   Usa el siguiente manifiesto para el archivo dws-flex-start.yaml:

   apiVersion: v1
   kind: Service
   metadata:
     name: headless-svc
   spec:
     clusterIP: None
     selector:
       job-name: job-1
   ---
   apiVersion: batch/v1
   kind: Job
   metadata:
     name: job-1
   spec:
     backoffLimit: 0
     completions: 2
     parallelism: 2
     completionMode: Indexed
     template:
       spec:
         subdomain: headless-svc
         restartPolicy: Never
         nodeSelector:
             cloud.google.com/gke-flex-start: "true"
             cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
             cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
         containers:
         - name: tpu-job
           image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
           ports:
           - containerPort: 8471 # Default port using which TPU VMs communicate
           - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
           securityContext:
             privileged: true
           command:
           - bash
           - -c
           - |
             python -c 'import jax; print("TPU cores:", jax.device_count())'
           resources:
             requests:
               google.com/tpu: NUM_CHIPS
             limits:
               google.com/tpu: NUM_CHIPS
   

   Haz los cambios siguientes:

   • ACCELERATOR_TYPE: el tipo de acelerador de TPU que has usado al crear los grupos de nodos. Por ejemplo, tpu-v4-podslice o tpu-v5-lite-podslice.
   • TPU_TOPOLOGY: la topología física de la porción de TPU. Por ejemplo, el valor puede ser 4x4x4 o 2x2, según la versión de la TPU.
   • NUM_CHIPS: el número de chips de TPU de cada máquina virtual es uno, cuatro u ocho. Para obtener más información, consulta Versiones de TPU.

3. Aplica el manifiesto dws-flex-start.yaml:

   kubectl apply -f dws-flex-start.yaml
   

4. Verifica que los trabajos se estén ejecutando en el mismo nodo:

   kubectl get pods
   

   El resultado debería ser similar al siguiente:

   NAME    READY   STATUS      RESTARTS   AGE   IP       NODE               NOMINATED NODE   READINESS GATES
   job-1   0/1     Completed   0          19m   10.(...) gke-flex-zonal-a2  <none>           <none>
   

   Éxito: has desplegado correctamente una carga de trabajo por lotes mediante máquinas virtuales de inicio flexible.

Limpieza

Para evitar que se apliquen cargos en tu cuenta Google Cloud por los recursos que has usado en esta página, elimina el proyecto que contiene los recursos o conserva el proyecto y elimina los recursos.

Eliminar el proyecto

Precaución: Eliminar un proyecto tiene los siguientes efectos:
• Se elimina todo el contenido del proyecto. Si has usado un proyecto que ya existía para las tareas de este documento, cuando lo elimines, también se eliminará cualquier otro trabajo que hayas realizado en él.
• Se pierden los IDs de proyecto personalizados. Cuando creaste este proyecto, es posible que hayas creado un ID de proyecto personalizado que quieras usar en el futuro. Para conservar las URLs que usan el ID del proyecto, como una URL appspot.com, elimina los recursos seleccionados dentro del proyecto en lugar de eliminar todo el proyecto.
1. In the Google Cloud console, go to the Manage resources page.

   Go to Manage resources

2. In the project list, select the project that you want to delete, and then click Delete.
3. In the dialog, type the project ID, and then click Shut down to delete the project.

Eliminar el recurso concreto

1. Elimina los trabajos:

   kubectl delete job -l "job-name in (job-1,job-2)"
   

2. Elimina el grupo de nodos:

   gcloud container node-pools delete NODE_POOL_NAME \
         --location LOCATION_NAME
   

3. Elimina el clúster:

   gcloud container clusters delete CLUSTER_NAME
   

Siguientes pasos

• Consulta más información sobre las TPUs en GKE.
• Consulta más información sobre el aprovisionamiento automático de nodos.
• Consulta más información sobre las prácticas recomendadas para ejecutar cargas de trabajo por lotes en GKE.