Administra recursos de TPU

En esta página, se describe cómo administrar los recursos de Cloud TPU con lo siguiente:

Requisitos previos

Antes de ejecutar estos procedimientos, debes instalar Google Cloud CLI, crear un proyecto de Google Cloud y habilitar la API de Cloud TPU. Para obtener instrucciones, consulta Configura el entorno de Cloud TPU.

Si usas Google Cloud CLI, puedes ejecutar comandos con Cloud Shell, una VM de Compute Engine o tu máquina local. Cloud Shell te permite interactuar con Cloud TPU sin tener que instalar ningún software. Cloud Shell se desconecta después de un período de inactividad. Si ejecutas comandos de larga duración, te recomendamos que instales Google Cloud CLI en tu máquina local. Para obtener más información sobre Google Cloud CLI, consulta la Referencia de gcloud.

Aprovisiona Cloud TPU

Puedes aprovisionar una Cloud TPU con gcloud, la consola de Google Cloud o la API de Cloud TPU.

Hay dos métodos para aprovisionar TPU con gcloud:

  1. Usa recursos en cola: gcloud compute tpus queued-resources create
  2. Con la API de Create Node: gcloud compute tpus tpu-vm create

La práctica recomendada es aprovisionar TPU con recursos en cola. Cuando solicitas recursos en cola, la solicitud se agrega a una cola que mantiene el servicio de Cloud TPU. Cuando el recurso solicitado está disponible, se asigna a tu proyecto de Google Cloud para tu uso exclusivo inmediato. Para obtener más información, consulta Recursos en cola administrados.

Cuando usas Multislice, debes usar recursos en cola y especificar los siguientes parámetros adicionales:

export NODE_COUNT=node_count
export NODE_PREFIX=your_tpu_prefix # Optional

Donde:

  • ${NODE_COUNT} es la cantidad de porciones que se crearán.
  • ${NODE_PREFIX} es el prefijo que especificas a fin de generar nombres para cada porción. Se agrega un número al prefijo para cada porción. Por ejemplo, si configuras ${NODE_PREFIX} como mySlice, las secciones se denominan mySlice-0, mySlice-1, y continúa numéricamente para cada porción.

Para obtener más información sobre Multislice, consulta la Introducción a Multislice.

Crea una Cloud TPU con la API de Crear nodo

Cuando creas una Cloud TPU, debes especificar la imagen de VM de TPU (también llamada versión de software de TPU). Para determinar qué imagen de VM debes usar, consulta Imágenes de VM de TPU.

También debes especificar la configuración de TPU en términos de TensorCores o chips de TPU. Para obtener más información, consulta la sección de la versión de TPU que usas en Arquitectura del sistema.

gcloud

Para crear una TPU con la API de Create Node, usa el comando gcloud compute tpus tpu-vm create.

El siguiente comando usa una configuración basada en TensorCore:

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=us-central2-b \
  --accelerator-type=v4-8 \
  --version=tpu-software-version

Descripciones de las marcas de comandos

zone
La zona en la que deseas crear la Cloud TPU.
accelerator-type
El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Si deseas obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta las versiones de TPU.
version
La versión del software de la TPU.
shielded-secure-boot (opcional)
Especifica que las instancias de TPU se crean con el arranque seguro habilitado. Esto las convierte de manera implícita en instancias de VM protegida. Consulta ¿Qué es una VM protegida? para obtener más información.

Con el siguiente comando, se crea una TPU con una topología específica:

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=us-central2-b \
  --type=v4 \
  --topology=2x2x1 \
  --version=tpu-software-version

Marcas obligatorias

tpu-name
Es el nombre de la VM de TPU que estás creando.
zone
La zona en la que crearás la Cloud TPU.
type
La versión de TPU que quieres usar. Para obtener más información, consulta Versiones de TPU.
topology
La disposición física de los chips TPU, que especifica la cantidad de chips en cada dimensión. Si deseas obtener más información sobre las topologías compatibles con cada versión de TPU, consulta las versiones de TPU.
version
La versión de software de TPU que quieres usar. Para obtener más información, consulta las versiones de software de TPU.

Consola

  1. En la consola de Google Cloud, ve a la página TPU:

    Ir a TPU

  2. Haz clic en Crear TPU.

  3. En el campo Nombre, ingresa un nombre para la TPU.

  4. En el cuadro Zona, selecciona la zona en la que se creará la TPU.

  5. En el cuadro Tipo de TPU, selecciona un tipo de acelerador. El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Si deseas obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta las versiones de TPU.

  6. En el cuadro Versión de software de TPU, selecciona una versión de software. Cuando se crea una VM de Cloud TPU, la versión de software de TPU especifica la versión del entorno de ejecución de TPU que se instalará. Para obtener más información, consulta Imágenes de VM de TPU.

  7. Haz clic en Crear para crear tus recursos.

curl

En el siguiente comando, se usa curl para crear una TPU.

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d "{accelerator_type: 'v4-8', \
runtime_version:'tpu-vm-tf-2.16.1-pjrt', \
network_config: {enable_external_ips: true}, \
shielded_instance_config: { enable_secure_boot: true }}" \
https://tpu.googleapis.com/v2/projects/project-id/locations/us-central2-b/nodes?node_id=node_name

Campos obligatorios

runtime_version
La versión del entorno de ejecución de Cloud TPU que quieres usar.
project
El nombre del proyecto de Google Cloud inscrito
zone
La zona en la que crearás tu Cloud TPU.
node_name
Es el nombre de la VM de TPU que estás creando.

Ejecuta una secuencia de comandos de inicio

Puedes ejecutar una secuencia de comandos de inicio en cada VM de TPU si especificas la marca --metadata startup-script cuando crees la VM de TPU. Con el siguiente comando, se crea una VM de TPU con una secuencia de comandos de inicio.

$ gcloud compute tpus tpu-vm create tpu-name \
    --zone=us-central2-b \
    --accelerator-type=tpu-type \
    --version=tpu-vm-tf-2.16.1-pjrt \
    --metadata startup-script='#! /bin/bash
      pip3 install numpy
      EOF'

Conéctate a una Cloud TPU

gcloud

Conéctate a tu Cloud TPU con SSH:

$ gcloud compute tpus tpu-vm ssh tpu-name --zone=zone

Cuando solicitas una porción más grande que un solo host, Cloud TPU crea una VM de TPU para cada host. La cantidad de chips TPU por host depende de la versión de TPU.

Para instalar objetos binarios o ejecutar código, conéctate a cada VM de TPU con el tpu-vm ssh command.

$ gcloud compute tpus tpu-vm ssh tpu-name

Para conectarte a una VM de TPU específica mediante SSH, usa la marca --worker, que sigue un índice basado en 0:

$ gcloud compute tpus tpu-vm ssh tpu-name --worker=1

Para ejecutar un comando en todas las VM de TPU con un solo comando, usa las marcas --worker=all y --command:

$ gcloud compute tpus tpu-vm ssh tpu-name \
  --project=your_project_ID \
  --zone=zone \
  --worker=all \
  --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

En el caso de Multislice, puedes ejecutar un comando en una VM única con el nombre de TPU enumerado, con cada prefijo de porción y el número adjunto. Para ejecutar un comando en todas las VM de TPU en todas las porciones, usa las marcas --node=all, --worker=all y --command, con una marca --batch-size opcional.

$ gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
  --project=project_ID \
  --zone=zone \
  --node=all \
  --worker=all \
  --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html' \
  --batch-size=4

Consola

Para conectarte a tus TPU en la consola de Google Cloud, usa SSH en el navegador:

  1. En la consola de Google Cloud, ve a la página TPU:

    Ir a TPU

  2. En la lista de VM de TPU, haz clic en SSH en la fila de la VM de TPU a la que deseas conectarte.

Haz una lista de tus recursos de Cloud TPU

Puedes enumerar todas las Cloud TPU en una zona especificada.

gcloud

$ gcloud compute tpus tpu-vm list --zone=zone

Consola

En la consola de Google Cloud, ve a la página TPU:

Ir a TPU

Recupera información sobre tu Cloud TPU

Puedes recuperar información sobre una Cloud TPU especificada.

gcloud

$ gcloud compute tpus tpu-vm describe tpu-name \
  --zone=zone

Consola

  1. En la consola de Google Cloud, ve a la página TPU:

    Ir a TPU

  2. Haz clic en el nombre de tu Cloud TPU. En la consola, se muestra la página de detalles de Cloud TPU.

Detén tus recursos de Cloud TPU

Puedes detener una sola Cloud TPU para que deje de generar cargos sin perder la configuración ni el software de la VM.

gcloud

$ gcloud compute tpus tpu-vm stop tpu-name \
  --zone=zone

Consola

  1. En la consola de Google Cloud, ve a la página TPU:

    Ir a TPU

  2. Selecciona la casilla de verificación junto a tu Cloud TPU.

  3. Haz clic en Detener.

Inicia tus recursos de Cloud TPU

Puedes iniciar una Cloud TPU cuando está detenida.

gcloud

$ gcloud compute tpus tpu-vm start tpu-name \
  --zone=zone

Consola

  1. En la consola de Google Cloud, ve a la página TPU:

    Ir a TPU

  2. Selecciona la casilla de verificación junto a tu Cloud TPU.

  3. Haz clic en Start.

Borra una Cloud TPU

Borra las porciones de VM de TPU al final de la sesión.

gcloud

$ gcloud compute tpus tpu-vm delete tpu-name \
  --project=project-id \
  --zone=zone \
  --quiet

Descripciones de las marcas de comandos

zone
La zona en la que planeas borrar la Cloud TPU.

Consola

  1. En la consola de Google Cloud, ve a la página TPU:

    Ir a TPU

  2. Selecciona la casilla de verificación junto a tu Cloud TPU.

  3. Haz clic en Borrar.

Configuración avanzada

Especifica recursos de red personalizados

Cuando creas la TPU, puedes elegir especificar una red o subred.

gcloud

Para especificar la red o subred con la CLI de gcloud, usa las siguientes marcas de comando:

--network NETWORK --subnetwork SUBNETWORK

curl

Para especificar la red o subred en una llamada a curl, agrega los siguientes parámetros al cuerpo de la solicitud:

network_config: {network: 'NETWORK', subnet: 'SUBNETWORK', enable_external_ips: true}

Red

De manera opcional, puedes especificar la red que se usará para la TPU. Si no se especifica, se usa la red default.

Formatos de red válidos:

https://www.googleapis.com/compute/{version}/projects/{proj-id}/global/networks/{network}
compute/{version}/projects/{proj-id}/global/networks/{network}
compute/{version}/projects/{proj-##}/global/networks/{network}
projects/{proj-id}/global/networks/{network}
projects/{proj-##}/global/networks/{network}
global/networks/{network}
{network}

Subred

Puedes especificar una subred específica para usar en la TPU. La subred especificada debe estar en la misma región que la zona en la que se ejecuta la TPU.

Formatos válidos:

https://www.googleapis.com/compute/{version}/projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
compute/{version}/projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
compute/{version}/projects/{proj-##}/regions/{region}/subnetworks/{subnetwork}
projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
projects/{proj-##}/regions/{region}/subnetworks/{subnetwork}
regions/{region}/subnetworks/{subnetwork}
{subnetwork}

Habilita el Acceso privado a Google

Para conectarte a las VM de TPU con SSH, debes agregar configuraciones de acceso para las VM de TPU o activar el Acceso privado a Google en la subred a la que están conectadas las VM de TPU.

Para agregar configuraciones de acceso, debes establecer enable_external_ips. Cuando creas una TPU, enable_external_ips se establece de forma predeterminada.

Si deseas inhabilitar esta opción, habilita las IP internas:

gcloud

Usa la marca --internal-ips cuando crees una TPU:

--internal-ips

curl

Agrega los siguientes parámetros al cuerpo de la solicitud:

network_config: {enable_external_ips: false}

Después de configurar el Acceso privado a Google, conéctate a la VM mediante SSH.

Adjunta una cuenta de servicio personalizada

Cada VM de TPU tiene una cuenta de servicio asociada que usa para realizar solicitudes a la API en tu nombre. Las VM de TPU usan esta cuenta de servicio para llamar a las APIs de Cloud TPU y acceder a Cloud Storage y otros servicios. De forma predeterminada, la VM de TPU usa la cuenta de servicio predeterminada de Compute Engine.

La cuenta de servicio debe definirse en el mismo proyecto de Google Cloud en el que creas la VM de TPU. Las cuentas de servicio personalizadas que se usan en las VM de TPU deben tener la función de Visualizador de TPU para llamar a la API de Cloud TPU. Si el código que se ejecuta en tu VM de TPU llama a otros servicios de Google Cloud, debe tener las funciones necesarias para acceder a esos servicios.

Cuando creas una TPU, puedes especificar una cuenta de servicio personalizada con la marca --service-account. Para obtener más información sobre las cuentas de servicio, consulta Cuentas de servicio.

Usa los siguientes comandos para especificar una cuenta de servicio personalizada.

gcloud

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=us-central2-b \
  --accelerator-type=tpu-type \
  --version=tpu-vm-tf-2.16.1-pjrt \
  --service-account=your-service-account

curl

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d "{accelerator_type: 'v4-8', \
runtime_version:'tpu-vm-tf-2.16.1-pjrt', \
network_config: {enable_external_ips: true}, \
shielded_instance_config: { enable_secure_boot: true }}" \
service_account: {email: 'your-service-account'} \
https://tpu.googleapis.com/v2/projects/project-id/locations/us-central2-b/nodes?node_id=node_name

Habilita métodos SSH personalizados

  1. Configurar un firewall para SSH

    La red predeterminada está preconfigurada para permitir el acceso SSH a todas las VMs. Si no usas la red predeterminada o si cambiaste la configuración de red predeterminada, es posible que debas habilitar explícitamente el acceso SSH mediante la adición de una regla de firewall:

    $ gcloud compute tpus tpu-vm compute firewall-rules create \
      --network=network allow-ssh \
      --allow=tcp:22
    
  2. Conéctate a las VMs de TPU con SSH.

    $ gcloud compute tpus tpu-vm ssh tpu-name \
      --zone=us-central2-b \
      --project=project-id
    

    Campos obligatorios

    • tpu-name: Es el nombre de la VM de TPU.
    • zone: Es la zona en la que creaste la VM de TPU.
    • project-id: Es el nombre de tu proyecto de Google Cloud.

    Para obtener una lista de los campos opcionales, consulta la documentación de la API de gcloud.