Dimensione as cargas de trabalho de AA com o Ray

Este documento fornece detalhes sobre como executar cargas de trabalho de aprendizagem automática (ML) com o Ray e o JAX nas TPUs. Existem dois modos diferentes para usar TPUs com o Ray: Modo centrado no dispositivo (PyTorch/XLA) e Modo centrado no anfitrião (JAX).

Este documento pressupõe que já tem um ambiente de TPU configurado. Para mais informações, consulte os seguintes recursos:

• Cloud TPU: configure o ambiente do Cloud TPU e faça a gestão dos recursos da TPU
• Google Kubernetes Engine (GKE): implemente cargas de trabalho de TPU no GKE Autopilot ou implemente cargas de trabalho de TPU no GKE Standard

Modo centrado no dispositivo (PyTorch/XLA)

O modo centrado no dispositivo retém grande parte do estilo de programação do PyTorch clássico. Neste modo, adiciona um novo tipo de dispositivo XLA, que funciona como qualquer outro dispositivo PyTorch. Cada processo individual interage com um dispositivo XLA.

Este modo é ideal se já estiver familiarizado com o PyTorch com GPUs e quiser usar abstrações de programação semelhantes.

As secções seguintes descrevem como executar uma carga de trabalho do PyTorch/XLA num ou mais dispositivos sem usar o Ray e, em seguida, como executar a mesma carga de trabalho em vários anfitriões usando o Ray.

Crie uma TPU

1. Crie variáveis de ambiente para parâmetros de criação de TPUs.

   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

   Descrições das variáveis de ambiente

   Variável	Descrição
   PROJECT_ID	O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
   TPU_NAME	O nome da TPU.
   ZONE	A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
   ACCELERATOR_TYPE	O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
   RUNTIME_VERSION	A versão do software do Cloud TPU.

2. Use o seguinte comando para criar uma VM de TPU v5p com 8 núcleos:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE  \
      --version=$RUNTIME_VERSION

3. Estabeleça ligação à VM de TPU através do seguinte comando:

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

Se estiver a usar o GKE, consulte o guia KubeRay no GKE para obter informações de configuração.

Requisitos de instalação

Execute os seguintes comandos na VM de TPU para instalar as dependências necessárias:

1. Guarde o seguinte num ficheiro. Por exemplo, 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
   

2. Para instalar as dependências necessárias, execute:

   pip install -r requirements.txt
   

Se estiver a executar a sua carga de trabalho no GKE, recomendamos que crie um Dockerfile que instale as dependências necessárias. Por exemplo, consulte o artigo Execute a sua carga de trabalho em nós de fatia de TPU na documentação do GKE.

Execute uma carga de trabalho do PyTorch/XLA num único dispositivo

O exemplo seguinte demonstra como criar um tensor XLA num único dispositivo, que é um chip TPU. Isto é semelhante à forma como o PyTorch processa outros tipos de dispositivos.

1. Guarde o seguinte fragmento do código num ficheiro. Por exemplo, 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)
   

   A declaração import torch_xla import inicializa o PyTorch/XLA e a função xm.xla_device() devolve o dispositivo XLA atual, um chip de TPU.

2. Defina a variável de ambiente PJRT_DEVICE como TPU.

   export PJRT_DEVICE=TPU
   

3. Execute o script.

   python workload.py
   

   O resultado tem um aspeto semelhante ao seguinte. Certifique-se de que o resultado indica que o dispositivo XLA foi encontrado.

   xla:0
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   

Execute o PyTorch/XLA em vários dispositivos

1. Atualize o fragmento do código da secção anterior para ser executado em vários 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=())
   

2. Execute o script.

   python workload.py
   

   Se executar o fragmento de código numa TPU v5p-8, o resultado é semelhante ao seguinte:

   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() usa dois argumentos: uma função e uma lista de parâmetros. Cria um processo para cada dispositivo XLA disponível e chama a função especificada nos argumentos. Neste exemplo, existem 4 dispositivos TPU disponíveis, pelo que torch_xla.launch() cria 4 processos e chama _mp_fn() em cada dispositivo. Cada processo só tem acesso a um dispositivo, pelo que cada dispositivo tem o índice 0 e xla:0 é impresso para todos os processos.

Execute o PyTorch/XLA em vários anfitriões com o Ray

As secções seguintes mostram como executar o mesmo fragmento de código numa fatia de TPU de vários anfitriões maior. Para mais informações sobre a arquitetura de TPU de vários anfitriões, consulte o artigo Arquitetura do sistema.

Neste exemplo, configura manualmente o Ray. Se já estiver familiarizado com a configuração do Ray, pode avançar para a última secção, Execute uma carga de trabalho do Ray. Para mais informações sobre a configuração do Ray para um ambiente de produção, consulte os seguintes recursos:

• Lançamento de clusters Ray
• Inicie Google Cloud o cluster do GKE com TPUs para o KubeRay

Crie uma VM de TPU com vários anfitriões

1. Crie variáveis de ambiente para parâmetros de criação de TPUs.

   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

   Descrições das variáveis de ambiente

   Variável	Descrição
   PROJECT_ID	O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
   TPU_NAME	O nome da TPU.
   ZONE	A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
   ACCELERATOR_TYPE	O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
   RUNTIME_VERSION	A versão do software do Cloud TPU.

2. Crie uma TPU v5p com vários anfitriões e 2 anfitriões (uma v5p-16, com 4 chips de TPU em cada anfitrião) através do seguinte comando:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE \
      --version=$RUNTIME_VERSION

Configure o Ray

Uma TPU v5p-16 tem 2 anfitriões de TPU, cada um com 4 chips de TPU. Neste exemplo, vai iniciar o nó principal do Ray num anfitrião e adicionar o segundo anfitrião como um nó de trabalho ao cluster do Ray.

1. Ligue-se ao primeiro anfitrião através de SSH.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=0

2. Instale as dependências com o mesmo ficheiro de requisitos que na secção de requisitos de instalação.

   pip install -r requirements.txt
   

3. Inicie o processo do Ray.

   ray start --head --port=6379
   

   Nota: se receber um erro como Command 'ray' not found, pode ter de adicionar o diretório que contém o executável ray à variável de ambiente PATH do seu sistema. Também pode especificar o caminho completo quando executar comandos ray (por exemplo, ~/.local/bin/ray).
   Nota: o comando ray start pede-lhe que ative ou desative a recolha de estatísticas de utilização. Para mais informações, consulte a secção Recolha de estatísticas de utilização na documentação do Ray.

   O resultado tem um aspeto semelhante ao seguinte:

   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 anfitrião de TPU é agora o nó principal do Ray. Tome nota das linhas que mostram como adicionar outro nó ao cluster do Ray, semelhante ao seguinte:

   To add another node to this Ray cluster, run
       ray start --address='10.130.0.76:6379'
   

   Vai usar este comando num passo posterior.

4. Verifique o estado do cluster do Ray:

   ray status
   

   O resultado tem um aspeto semelhante ao seguinte:

   ======== 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)
   

   O cluster só contém 4 TPUs (0.0/4.0 TPU) porque só adicionou o nó principal até agora.

   Agora que o nó principal está em execução, pode adicionar o segundo anfitrião ao cluster.

5. Ligue-se ao segundo anfitrião através de SSH.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=1

6. Instale as dependências com o mesmo ficheiro de requisitos que na secção Requisitos de instalação.

   pip install -r requirements.txt
   

7. Inicie o processo do Ray. Para adicionar este nó ao cluster do Ray existente, use o comando do resultado do comando ray start. Certifique-se de que substitui o endereço IP e a porta no seguinte comando:

   ray start --address='10.130.0.76:6379'

   O resultado tem um aspeto semelhante ao seguinte:

   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
   

8. Verifique novamente o estado do Ray:

   ray status
   

   O resultado tem um aspeto semelhante ao seguinte:

   ======== 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)
   

   O segundo anfitrião de TPU é agora um nó no cluster. A lista de recursos disponíveis mostra agora 8 TPUs (0.0/8.0 TPU).

   Nota: tem de iniciar o processo do Ray em cada uma das suas VMs de TPU antes de executar a carga de trabalho.

Execute uma carga de trabalho do Ray

1. Atualize o fragmento do código para ser executado no cluster do 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()
   

2. Execute o guião no nó principal do Ray. Substitua ray-workload.py pelo caminho para o seu script.

   python ray-workload.py

   O resultado tem um aspeto semelhante ao seguinte:

   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')
   

   O resultado indica que a função foi chamada com êxito em cada dispositivo XLA (8 dispositivos neste exemplo) na fatia de TPU com vários anfitriões.

Modo centrado no anfitrião (JAX)

As secções seguintes descrevem o modo centrado no anfitrião com o JAX. O JAX usa um paradigma de programação funcional e suporta semântica de programa único de nível superior, vários dados (SPMD). Em vez de cada processo interagir com um único dispositivo XLA, o código JAX foi concebido para funcionar em vários dispositivos num único anfitrião em simultâneo.

O JAX foi concebido para computação de alto desempenho e pode usar as TPUs de forma eficiente para inferência e preparação em grande escala. Este modo é ideal se tiver conhecimentos dos conceitos de programação funcional para poder tirar partido de todo o potencial do JAX.

Estas instruções pressupõem que já tem um ambiente Ray e TPU configurado, incluindo um ambiente de software que inclui o JAX e outros pacotes relacionados. Para criar um cluster de TPU do Ray, siga as instruções em Iniciar Google Cloud o cluster do GKE com TPUs para o KubeRay. Para mais informações sobre a utilização de TPUs com o KubeRay, consulte o artigo Use TPUs with KubeRay (Use TPUs com o KubeRay).

Execute uma carga de trabalho JAX numa TPU de host único

O script de exemplo seguinte demonstra como executar uma função JAX num cluster Ray com uma TPU de anfitrião único, como uma v6e-4. Se tiver uma TPU com vários anfitriões, este script deixa de responder devido ao modelo de execução com vários controladores do JAX. Para mais informações sobre a execução do Ray numa TPU com vários anfitriões, consulte o artigo Execute uma carga de trabalho JAX numa TPU com vários anfitriões.

1. Crie variáveis de ambiente para parâmetros de criação de TPUs.

   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

   Descrições das variáveis de ambiente

   Variável	Descrição
   PROJECT_ID	O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
   TPU_NAME	O nome da TPU.
   ZONE	A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
   ACCELERATOR_TYPE	O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
   RUNTIME_VERSION	A versão do software do Cloud TPU.

2. Use o seguinte comando para criar uma VM de TPU v6e com 4 núcleos:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE  \
      --version=$RUNTIME_VERSION

3. Estabeleça ligação à VM de TPU através do seguinte comando:

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

4. Instale o JAX e o Ray na TPU.

   pip install ray jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   

5. Guarde o seguinte código num ficheiro. Por exemplo, 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
   

   Se estiver habituado a executar o Ray com GPUs, existem algumas diferenças importantes quando usa TPUs:

   • Em vez de definir num_gpus, especifique TPU como um recurso personalizado e defina o número de chips de TPU.
   • Especifique a TPU usando o número de chips por nó de trabalho do Ray. Por exemplo, se estiver a usar um v6e-4, a execução de uma função remota com TPU definida como 4 consome todo o anfitrião da TPU.
   • Isto é diferente da forma como as GPUs são normalmente executadas, com um processo por anfitrião. Não é recomendado definir TPU para um número diferente de 4.
     • Exceção: se tiver um único anfitrião v6e-8 ou v5litepod-8, deve definir este valor como 8.

6. Execute o script.

   python ray-jax-single-host.py

Execute uma carga de trabalho JAX numa TPU multi-anfitrião

O script de exemplo seguinte demonstra como executar uma função JAX num cluster do Ray com uma TPU multi-host. O script de exemplo usa um v6e-16.

1. Crie variáveis de ambiente para parâmetros de criação de TPUs.

   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

   Descrições das variáveis de ambiente

   Variável	Descrição
   PROJECT_ID	O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
   TPU_NAME	O nome da TPU.
   ZONE	A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
   ACCELERATOR_TYPE	O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
   RUNTIME_VERSION	A versão do software do Cloud TPU.

2. Use o seguinte comando para criar uma VM de TPU v6e com 16 núcleos:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE  \
      --version=$RUNTIME_VERSION

3. Instale o JAX e o Ray em todos os trabalhadores da 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"

4. Guarde o seguinte código num ficheiro. Por exemplo, 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]
   

   Se estiver habituado a executar o Ray com GPUs, existem algumas diferenças importantes quando usa TPUs:

   • Semelhante às cargas de trabalho do PyTorch em GPUs:
     • As cargas de trabalho JAX nas TPUs são executadas num formato de vários controladores, um único programa e vários dados (SPMD).
     • Os coletivos entre dispositivos são processados pela estrutura de aprendizagem automática.
   • Ao contrário das cargas de trabalho do PyTorch em GPUs, o JAX tem uma vista global dos dispositivos disponíveis no cluster.

5. Copie o script para todos os trabalhadores da TPU.

   gcloud compute tpus tpu-vm scp ray-jax-multi-host.py $TPU_NAME: --zone=$ZONE --worker=all

6. Execute o script.

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --zone=$ZONE \
      --worker=all \
      --command="python ray-jax-multi-host.py"

Execute uma carga de trabalho JAX de vários fragmentos

A multislice permite-lhe executar cargas de trabalho que abrangem várias fatias de TPU num único TPU Pod ou em vários pods na rede do centro de dados.

Pode usar o pacote ray-tpu para simplificar as interações do Ray com as fatias de TPU.

Instale o ray-tpu através do pip.

pip install ray-tpu

Para mais informações sobre a utilização do pacote ray-tpu, consulte o artigo Introdução no repositório do GitHub. Para ver um exemplo com o Multislice, consulte o artigo Execução no Multislice.

Orquestre cargas de trabalho com o Ray e o MaxText

Para mais informações sobre a utilização do Ray com o MaxText, consulte o artigo Execute uma tarefa de preparação com o MaxText.

Recursos da TPU e do Ray

O Ray trata as TPUs de forma diferente das GPUs para ter em conta a diferença de utilização. No exemplo seguinte, existem nove nós do Ray no total:

• O nó principal do Ray está a ser executado numa VM n1-standard-16.
• Os nós de trabalho do Ray estão a ser executados em duas v6e-16TPUs. Cada TPU constitui quatro trabalhadores.

$ 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)

Descrições dos campos de utilização de recursos:

• CPU: o número total de CPUs disponíveis no cluster.
• TPU: o número de chips de TPU no cluster.
• TPU-v6e-16-head: um identificador especial para o recurso que corresponde ao trabalhador 0 de uma fatia de TPU. Isto é importante para aceder a fatias individuais de TPUs.
• memory: memória heap do worker usada pela sua aplicação.
• object_store_memory: memória usada quando a sua aplicação cria objetos no armazenamento de objetos através de ray.put e quando devolve valores de funções remotas.
• tpu-group-0 e tpu-group-1: identificadores exclusivos das divisões individuais da TPU. Isto é importante para executar tarefas em fatias. Estes campos estão definidos como 4 porque existem 4 anfitriões por fatia de TPU num v6e-16.

Nota: todos estes recursos são detetados automaticamente quando o tempo de execução do Ray é iniciado (ray start), pelo que não precisa de os definir explicitamente.