Escala las cargas de trabajo de AA con Ray
En este documento, se proporcionan detalles sobre cómo ejecutar cargas de trabajo de aprendizaje automático (AA) con Ray y JAX en TPU. Hay dos modos diferentes para usar TPU con Ray: modo centrado en el dispositivo (PyTorch/XLA) y modo centrado en el host (JAX).
En este documento, se supone que ya tienes configurado un entorno de TPU. Para obtener más información, consulta los siguientes recursos:
- Cloud TPU: Configura el entorno de Cloud TPU y administra los recursos de TPU
- Google Kubernetes Engine (GKE): Implementa cargas de trabajo de TPU en GKE Autopilot o Implementa cargas de trabajo de TPU en GKE Standard
Modo centrado en el dispositivo (PyTorch/XLA)
El modo centrado en el dispositivo conserva gran parte del estilo programático de PyTorch clásico. En este modo, agregas un nuevo tipo de dispositivo XLA, que funciona como cualquier otro dispositivo de PyTorch. Cada proceso individual interactúa con un dispositivo XLA.
Este modo es ideal si ya conoces PyTorch con GPUs y quieres usar abstracciones de programación similares.
En las siguientes secciones, se describe cómo ejecutar una carga de trabajo de PyTorch/XLA en uno o más dispositivos sin usar Ray y, luego, cómo ejecutar la misma carga de trabajo en varios hosts con Ray.
Crear una TPU
Crea variables de entorno para los parámetros de creación de TPU.
export PROJECT_ID=your-project-id export TPU_NAME=your-tpu-name export ZONE=europe-west4-b export ACCELERATOR_TYPE=v5p-8 export RUNTIME_VERSION=v2-alpha-tpuv5
Descripciones de las variables de entorno
Variable Descripción PROJECT_ID
El Google Cloud ID de tu proyecto. Usa un proyecto existente o crea uno nuevo. TPU_NAME
El nombre de la TPU. ZONE
Es la zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas admitidas, consulta Regiones y zonas de TPU. ACCELERATOR_TYPE
El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU. RUNTIME_VERSION
La versión de software de Cloud TPU Usa el siguiente comando para crear una VM de TPU v5p con 8 núcleos:
gcloud compute tpus tpu-vm create $TPU_NAME \ --zone=$ZONE \ --accelerator-type=$ACCELERATOR_TYPE \ --version=$RUNTIME_VERSION
Conéctate a la VM de TPU con el siguiente comando:
gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
Si usas GKE, consulta la guía KubeRay en GKE para obtener información sobre la configuración.
Requisitos de instalación
Ejecuta los siguientes comandos en tu VM de TPU para instalar las dependencias requeridas:
Guarda lo siguiente en un archivo. Por ejemplo,
requirements.txt
.--find-links https://storage.googleapis.com/libtpu-releases/index.html --find-links https://storage.googleapis.com/libtpu-wheels/index.html torch~=2.6.0 torch_xla[tpu]~=2.6.0 ray[default]==2.40.0
Para instalar las dependencias requeridas, ejecuta lo siguiente:
pip install -r requirements.txt
Si ejecutas tu carga de trabajo en GKE, te recomendamos que crees un Dockerfile que instale las dependencias requeridas. Para ver un ejemplo, consulta Ejecuta tu carga de trabajo en nodos de porción TPU en la documentación de GKE.
Ejecuta una carga de trabajo de PyTorch/XLA en un solo dispositivo
En el siguiente ejemplo, se muestra cómo crear un tensor de XLA en un solo dispositivo, que es un chip de TPU. Esto es similar a la forma en que PyTorch controla otros tipos de dispositivos.
Guarda el siguiente fragmento de código en un archivo. Por ejemplo,
workload.py
.import torch import torch_xla import torch_xla.core.xla_model as xm t = torch.randn(2, 2, device=xm.xla_device()) print(t.device) print(t)
La sentencia de importación
import torch_xla
inicializa PyTorch/XLA, y la funciónxm.xla_device()
muestra el dispositivo XLA actual, un chip TPU.Establece la variable de entorno
PJRT_DEVICE
como TPU.export PJRT_DEVICE=TPU
Ejecuta la secuencia de comandos.
python workload.py
El resultado es similar al siguiente. Asegúrate de que el resultado indique que se encontró el dispositivo XLA.
xla:0 tensor([[ 0.6220, -1.4707], [-1.2112, 0.7024]], device='xla:0')
Ejecuta PyTorch/XLA en varios dispositivos
Actualiza el fragmento de código de la sección anterior para que se ejecute en varios dispositivos.
import torch import torch_xla import torch_xla.core.xla_model as xm def _mp_fn(index): t = torch.randn(2, 2, device=xm.xla_device()) print(t.device) print(t) if __name__ == '__main__': torch_xla.launch(_mp_fn, args=())
Ejecuta la secuencia de comandos.
python workload.py
Si ejecutas el fragmento de código en una TPU v5p-8, el resultado será similar al siguiente:
xla:0 xla:0 xla:0 tensor([[ 1.2309, 0.9896], [ 0.5820, -1.2950]], device='xla:0') xla:0 tensor([[ 1.2309, 0.9896], [ 0.5820, -1.2950]], device='xla:0') tensor([[ 1.2309, 0.9896], [ 0.5820, -1.2950]], device='xla:0') tensor([[ 1.2309, 0.9896], [ 0.5820, -1.2950]], device='xla:0')
torch_xla.launch()
toma dos argumentos: una función y una lista de parámetros. Crea un proceso para cada dispositivo XLA disponible y llama a la función especificada en los argumentos. En este ejemplo, hay 4 dispositivos TPU disponibles, por lo que torch_xla.launch()
crea 4 procesos y llama a _mp_fn()
en cada dispositivo. Cada proceso solo tiene acceso a un dispositivo, por lo que cada dispositivo tiene el índice 0 y se imprime xla:0
para todos los procesos.
Ejecuta PyTorch/XLA en varios hosts con Ray
En las siguientes secciones, se muestra cómo ejecutar el mismo fragmento de código en una porción de TPU de varios hosts más grande. Para obtener más información sobre la arquitectura de TPU de varios hosts, consulta Arquitectura del sistema.
En este ejemplo, configurarás Ray de forma manual. Si ya conoces la configuración de Ray, puedes pasar a la última sección, Ejecuta una carga de trabajo de Ray. Para obtener más información sobre cómo configurar Ray para un entorno de producción, consulta los siguientes recursos:
Crea una VM de TPU de varios hosts
Crea variables de entorno para los parámetros de creación de TPU.
export PROJECT_ID=your-project-id export TPU_NAME=your-tpu-name export ZONE=europe-west4-b export ACCELERATOR_TYPE=v5p-16 export RUNTIME_VERSION=v2-alpha-tpuv5
Descripciones de las variables de entorno
Variable Descripción PROJECT_ID
El Google Cloud ID de tu proyecto. Usa un proyecto existente o crea uno nuevo. TPU_NAME
El nombre de la TPU. ZONE
Es la zona en la que se creará la VM de TPU. Si deseas obtener más información sobre las zonas admitidas, consulta Regiones y zonas de TPU. ACCELERATOR_TYPE
El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU. RUNTIME_VERSION
La versión de software de Cloud TPU Crea una TPU v5p de varios hosts con 2 hosts (un v5p-16, con 4 chips de TPU en cada host) con el siguiente comando:
gcloud compute tpus tpu-vm create $TPU_NAME \ --zone=$ZONE \ --accelerator-type=$ACCELERATOR_TYPE \ --version=$RUNTIME_VERSION
Configura Ray
Una TPU v5p-16 tiene 2 hosts de TPU, cada uno con 4 chips TPU. En este ejemplo, iniciarás el nodo principal de Ray en un host y agregarás el segundo host como nodo trabajador al clúster de Ray.
Conéctate al primer host con SSH.
gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=0
Instala las dependencias con el mismo archivo de requisitos que en la sección Instala los requisitos.
pip install -r requirements.txt
Inicia el proceso de Ray.
ray start --head --port=6379
El resultado es similar al siguiente:
Enable usage stats collection? This prompt will auto-proceed in 10 seconds to avoid blocking cluster startup. Confirm [Y/n]: y Usage stats collection is enabled. To disable this, add `--disable-usage-stats` to the command that starts the cluster, or run the following command: `ray disable-usage-stats` before starting the cluster. See https://docs.ray.io/en/master/cluster/usage-stats.html for more details. Local node IP: 10.130.0.76 -------------------- Ray runtime started. -------------------- Next steps To add another node to this Ray cluster, run ray start --address='10.130.0.76:6379' To connect to this Ray cluster: import ray ray.init() To terminate the Ray runtime, run ray stop To view the status of the cluster, use ray status
Este host de TPU ahora es el nodo principal de Ray. Anota las líneas que muestran cómo agregar otro nodo al clúster de Ray, de manera similar a la siguiente:
To add another node to this Ray cluster, run ray start --address='10.130.0.76:6379'
Usarás este comando en un paso posterior.
Verifica el estado del clúster de Ray:
ray status
El resultado es similar al siguiente:
======== Autoscaler status: 2025-01-14 22:03:39.385610 ======== Node status --------------------------------------------------------------- Active: 1 node_bc0c62819ddc0507462352b76cc06b462f0e7f4898a77e5133c16f79 Pending: (no pending nodes) Recent failures: (no failures) Resources --------------------------------------------------------------- Usage: 0.0/208.0 CPU 0.0/4.0 TPU 0.0/1.0 TPU-v5p-16-head 0B/268.44GiB memory 0B/119.04GiB object_store_memory 0.0/1.0 your-tpu-name Demands: (no resource demands)
El clúster solo contiene 4 TPU (
0.0/4.0 TPU
) porque solo agregaste el nodo principal hasta el momento.Ahora que el nodo principal se está ejecutando, puedes agregar el segundo host al clúster.
Conéctate al segundo host con SSH.
gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=1
Instala las dependencias con el mismo archivo de requisitos que en la sección Requisitos de instalación.
pip install -r requirements.txt
Inicia el proceso de Ray. Para agregar este nodo al clúster de Ray existente, usa el comando del resultado del comando
ray start
. Asegúrate de reemplazar la dirección IP y el puerto en el siguiente comando:ray start --address='10.130.0.76:6379'
El resultado es similar al siguiente:
Local node IP: 10.130.0.80 [2025-01-14 22:30:07,397 W 75572 75572] global_state_accessor.cc:463: Retrying to get node with node ID 35f9ac0675c91429805cdc1b97c3713422d97eee783ccb0c0304f5c1 -------------------- Ray runtime started. -------------------- To terminate the Ray runtime, run ray stop
Vuelve a verificar el estado de Ray:
ray status
El resultado es similar al siguiente:
======== Autoscaler status: 2025-01-14 22:45:21.485617 ======== Node status --------------------------------------------------------------- Active: 1 node_bc0c62819ddc0507462352b76cc06b462f0e7f4898a77e5133c16f79 1 node_35f9ac0675c91429805cdc1b97c3713422d97eee783ccb0c0304f5c1 Pending: (no pending nodes) Recent failures: (no failures) Resources --------------------------------------------------------------- Usage: 0.0/416.0 CPU 0.0/8.0 TPU 0.0/1.0 TPU-v5p-16-head 0B/546.83GiB memory 0B/238.35GiB object_store_memory 0.0/2.0 your-tpu-name Demands: (no resource demands)
El segundo host de TPU ahora es un nodo en el clúster. La lista de recursos disponibles ahora muestra 8 TPU (
0.0/8.0 TPU
).
Ejecuta una carga de trabajo de Ray
Actualiza el fragmento de código para que se ejecute en el clúster de Ray:
import os import torch import torch_xla import torch_xla.core.xla_model as xm import ray import torch.distributed as dist import torch_xla.runtime as xr from torch_xla._internal import pjrt # Defines the local PJRT world size, the number of processes per host. LOCAL_WORLD_SIZE = 4 # Defines the number of hosts in the Ray cluster. NUM_OF_HOSTS = 4 GLOBAL_WORLD_SIZE = LOCAL_WORLD_SIZE * NUM_OF_HOSTS def init_env(): local_rank = int(os.environ['TPU_VISIBLE_CHIPS']) pjrt.initialize_multiprocess(local_rank, LOCAL_WORLD_SIZE) xr._init_world_size_ordinal() # This decorator signals to Ray that the `print_tensor()` function should be run on a single TPU chip. @ray.remote(resources={"TPU": 1}) def print_tensor(): # Initializes the runtime environment on each Ray worker. Equivalent to # the `torch_xla.launch call` in the Run PyTorch/XLA on multiple devices section. init_env() t = torch.randn(2, 2, device=xm.xla_device()) print(t.device) print(t) ray.init() # Uses Ray to dispatch the function call across available nodes in the cluster. tasks = [print_tensor.remote() for _ in range(GLOBAL_WORLD_SIZE)] ray.get(tasks) ray.shutdown()
Ejecuta la secuencia de comandos en el nodo principal de Ray. Reemplaza ray-workload.py por la ruta de acceso a tu secuencia de comandos.
python ray-workload.py
El resultado es similar al siguiente:
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU. xla:0 xla:0 xla:0 xla:0 xla:0 tensor([[ 0.6220, -1.4707], [-1.2112, 0.7024]], device='xla:0') tensor([[ 0.6220, -1.4707], [-1.2112, 0.7024]], device='xla:0') xla:0 xla:0 tensor([[ 0.6220, -1.4707], [-1.2112, 0.7024]], device='xla:0') tensor([[ 0.6220, -1.4707], [-1.2112, 0.7024]], device='xla:0') tensor([[ 0.6220, -1.4707], [-1.2112, 0.7024]], device='xla:0') tensor([[ 0.6220, -1.4707], [-1.2112, 0.7024]], device='xla:0') tensor([[ 0.6220, -1.4707], [-1.2112, 0.7024]], device='xla:0') xla:0 tensor([[ 0.6220, -1.4707], [-1.2112, 0.7024]], device='xla:0')
El resultado indica que se llamó correctamente a la función en cada dispositivo XLA (8 dispositivos en este ejemplo) en la porción de TPU de varios hosts.
Modo centrado en el host (JAX)
En las siguientes secciones, se describe el modo centrado en el host con JAX. JAX usa un paradigma de programación funcional y admite semánticas de un solo programa y varios datos (SPMD) de nivel superior. En lugar de que cada proceso interactúe con un solo dispositivo XLA, el código JAX está diseñado para operar en varios dispositivos en un solo host de forma simultánea.
JAX está diseñado para la computación de alto rendimiento y puede usar TPU de manera eficiente para el entrenamiento y la inferencia a gran escala. Este modo es ideal si conoces los conceptos de programación funcional para que puedas aprovechar todo el potencial de JAX.
En estas instrucciones, se supone que ya tienes configurado un entorno de Ray y TPU, incluido un entorno de software que incluye JAX y otros paquetes relacionados. Para crear un clúster de TPU de Ray, sigue las instrucciones que se indican en Cómo iniciar Google Cloud un clúster de GKE con TPU para KubeRay. Para obtener más información sobre el uso de TPU con KubeRay, consulta Cómo usar TPU con KubeRay.
Ejecuta una carga de trabajo de JAX en una TPU de host único
En la siguiente secuencia de comandos de ejemplo, se muestra cómo ejecutar una función de JAX en un clúster de Ray con una TPU de un solo host, como una v6e-4. Si tienes una TPU de varios hosts, esta secuencia de comandos deja de responder debido al modelo de ejecución de varios controladores de JAX. Para obtener más información sobre cómo ejecutar Ray en una TPU de varios hosts, consulta Cómo ejecutar una carga de trabajo de JAX en una TPU de varios hosts.
Crea variables de entorno para los parámetros de creación de TPU.
export PROJECT_ID=your-project-id export TPU_NAME=your-tpu-name export ZONE=europe-west4-a export ACCELERATOR_TYPE=v6e-4 export RUNTIME_VERSION=v2-alpha-tpuv6e
Descripciones de las variables de entorno
Variable Descripción PROJECT_ID
El Google Cloud ID de tu proyecto. Usa un proyecto existente o crea uno nuevo. TPU_NAME
El nombre de la TPU. ZONE
Es la zona en la que se creará la VM de TPU. Si deseas obtener más información sobre las zonas admitidas, consulta Regiones y zonas de TPU. ACCELERATOR_TYPE
El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU. RUNTIME_VERSION
La versión de software de Cloud TPU Usa el siguiente comando para crear una VM de TPU v6e con 4 núcleos:
gcloud compute tpus tpu-vm create $TPU_NAME \ --zone=$ZONE \ --accelerator-type=$ACCELERATOR_TYPE \ --version=$RUNTIME_VERSION
Conéctate a la VM de TPU con el siguiente comando:
gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
Instala JAX y Ray en tu TPU.
pip install ray jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Guarda el siguiente código en un archivo. Por ejemplo,
ray-jax-single-host.py
.import ray import jax @ray.remote(resources={"TPU": 4}) def my_function() -> int: return jax.device_count() h = my_function.remote() print(ray.get(h)) # => 4
Si estás acostumbrado a ejecutar Ray con GPUs, hay algunas diferencias clave cuando se usan TPU:
- En lugar de configurar
num_gpus
, especificaTPU
como un recurso personalizado y establece la cantidad de chips TPU. - Especifica la TPU con la cantidad de chips por nodo trabajador de Ray. Por ejemplo,
si usas una v6e-4, ejecutar una función remota con
TPU
establecido en 4 consume todo el host de TPU. - Esto es diferente de la forma en que se ejecutan las GPU, con un proceso por host.
No se recomienda establecer
TPU
en un número que no sea 4.- Excepción: Si tienes un
v6e-8
ov5litepod-8
de un solo host, debes establecer este valor en 8.
- Excepción: Si tienes un
- En lugar de configurar
Ejecuta la secuencia de comandos.
python ray-jax-single-host.py
Ejecuta una carga de trabajo de JAX en una TPU de varios hosts
En la siguiente secuencia de comandos de ejemplo, se muestra cómo ejecutar una función de JAX en un clúster de Ray con una TPU de varios hosts. La secuencia de comandos de ejemplo usa un v6e-16.
Crea variables de entorno para los parámetros de creación de TPU.
export PROJECT_ID=your-project-id export TPU_NAME=your-tpu-name export ZONE=europe-west4-a export ACCELERATOR_TYPE=v6e-16 export RUNTIME_VERSION=v2-alpha-tpuv6e
Descripciones de las variables de entorno
Variable Descripción PROJECT_ID
El Google Cloud ID de tu proyecto. Usa un proyecto existente o crea uno nuevo. TPU_NAME
El nombre de la TPU. ZONE
Es la zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas admitidas, consulta Regiones y zonas de TPU. ACCELERATOR_TYPE
El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU. RUNTIME_VERSION
La versión de software de Cloud TPU Usa el siguiente comando para crear una VM de TPU v6e con 16 núcleos:
gcloud compute tpus tpu-vm create $TPU_NAME \ --zone=$ZONE \ --accelerator-type=$ACCELERATOR_TYPE \ --version=$RUNTIME_VERSION
Instala JAX y Ray en todos los trabajadores de TPU.
gcloud compute tpus tpu-vm ssh $TPU_NAME \ --zone=$ZONE \ --worker=all \ --command="pip install ray jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
Guarda el siguiente código en un archivo. Por ejemplo,
ray-jax-multi-host.py
.import ray import jax @ray.remote(resources={"TPU": 4}) def my_function() -> int: return jax.device_count() ray.init() num_tpus = ray.available_resources()["TPU"] num_hosts = int(num_tpus) # 4 h = [my_function.remote() for _ in range(num_hosts)] print(ray.get(h)) # [16, 16, 16, 16]
Si estás acostumbrado a ejecutar Ray con GPUs, hay algunas diferencias clave cuando se usan TPU:
- Similar a las cargas de trabajo de PyTorch en GPUs:
- Las cargas de trabajo de JAX en TPUs se ejecutan de forma multicontrolador, un solo programa, múltiples datos (SPMD).
- El framework de aprendizaje automático controla los colectivos entre dispositivos.
- A diferencia de las cargas de trabajo de PyTorch en GPUs, JAX tiene una vista global de los dispositivos disponibles en el clúster.
- Similar a las cargas de trabajo de PyTorch en GPUs:
Copia la secuencia de comandos en todos los trabajadores de TPU.
gcloud compute tpus tpu-vm scp ray-jax-multi-host.py $TPU_NAME: --zone=$ZONE --worker=all
Ejecuta la secuencia de comandos.
gcloud compute tpus tpu-vm ssh $TPU_NAME \ --zone=$ZONE \ --worker=all \ --command="python ray-jax-multi-host.py"
Ejecuta una carga de trabajo de JAX de Multislice
Multislice te permite ejecutar cargas de trabajo que abarcan varias porciones de TPU dentro de un solo pod de TPU o en varios pods a través de la red del centro de datos.
Puedes usar el paquete ray-tpu
para simplificar las interacciones de Ray con las porciones de TPU.
Instala ray-tpu
usando pip
.
pip install ray-tpu
Para obtener más información sobre el uso del paquete ray-tpu
, consulta Cómo comenzar en el repositorio de GitHub. Para ver un ejemplo de uso de Multislice, consulta Cómo ejecutar en Multislice.
Organiza cargas de trabajo con Ray y MaxText
Para obtener más información sobre el uso de Ray con MaxText, consulta Cómo ejecutar un trabajo de entrenamiento con MaxText.
Recursos de TPU y Ray
Ray trata las TPU de manera diferente a las GPUs para adaptarse a la diferencia en el uso. En el siguiente ejemplo, hay nueve nodos de Ray en total:
- El nodo principal de Ray se ejecuta en una VM
n1-standard-16
. - Los nodos de trabajo de Ray se ejecutan en dos TPU
v6e-16
. Cada TPU constituye cuatro trabajadores.
$ ray status
======== Autoscaler status: 2024-10-17 09:30:00.854415 ========
Node status
---------------------------------------------------------------
Active:
1 node_e54a65b81456cee40fcab16ce7b96f85406637eeb314517d9572dab2
1 node_9a8931136f8d2ab905b07d23375768f41f27cc42f348e9f228dcb1a2
1 node_c865cf8c0f7d03d4d6cae12781c68a840e113c6c9b8e26daeac23d63
1 node_435b1f8f1fbcd6a4649c09690915b692a5bac468598e9049a2fac9f1
1 node_3ed19176e9ecc2ac240c818eeb3bd4888fbc0812afebabd2d32f0a91
1 node_6a88fe1b74f252a332b08da229781c3c62d8bf00a5ec2b90c0d9b867
1 node_5ead13d0d60befd3a7081ef8b03ca0920834e5c25c376822b6307393
1 node_b93cb79c06943c1beb155d421bbd895e161ba13bccf32128a9be901a
1 node_9072795b8604ead901c5268ffcc8cc8602c662116ac0a0272a7c4e04
Pending:
(no pending nodes)
Recent failures:
(no failures)
Resources
---------------------------------------------------------------
Usage:
0.0/727.0 CPU
0.0/32.0 TPU
0.0/2.0 TPU-v6e-16-head
0B/5.13TiB memory
0B/1.47TiB object_store_memory
0.0/4.0 tpu-group-0
0.0/4.0 tpu-group-1
Demands:
(no resource demands)
Descripciones de los campos de uso de recursos:
CPU
: Es la cantidad total de CPUs disponibles en el clúster.TPU
: Es la cantidad de chips TPU en el clúster.TPU-v6e-16-head
: Es un identificador especial para el recurso que corresponde al trabajador 0 de una porción de TPU. Esto es importante para acceder a las porciones individuales de TPU.memory
: Es la memoria del montón de trabajo que usa tu aplicación.object_store_memory
: Es la memoria que se usa cuando tu aplicación crea objetos en el almacén de objetos conray.put
y cuando muestra valores de funciones remotas.tpu-group-0
ytpu-group-1
: Son identificadores únicos para las porciones individuales de TPU. Esto es importante para ejecutar trabajos en rebanadas. Estos campos se establecen en 4 porque hay 4 hosts por porción de TPU en una v6e-16.