Executar uma carga de trabalho interativa com os programas

As cargas de trabalho interativas do Pathways são cargas de trabalho JAX remotas executadas em uma VM que não faz parte do cluster do GKE que hospeda o cluster do Pathways. Ao contrário das cargas de trabalho em lote, a conclusão da operação de carga de trabalho interativa não desativa os componentes do cluster do Pathways, que permanecem disponíveis para conexão por outros clientes JAX. Este documento usa um notebook do Jupyter como exemplo para demonstrar cargas de trabalho interativas.

Usando a IFRT interface, os usuários do JAX enviam comandos para um cluster do Pathways. O código JAX, executado em um terminal, notebook ou qualquer ambiente compatível com Python, pode interagir perfeitamente com os recursos do Pathways.

Antes de começar

Você precisa ter:

• Criado um cluster do GKE usando o XPK
• Instalado o XPK
• Instalado as ferramentas do Kubernetes
• Instalado a CLI gcloud
• Ativado a API TPU
• API do Google Kubernetes Engine ativada
• Verifique se o Pathways está ativado para o seu Google Cloud projeto

Executar o Pathways no modo interativo

É possível executar o Pathways no modo interativo usando xpk ou kubectl.

XPK

1. Defina as seguintes variáveis de ambiente:

   export WORKLOAD=WORKLOAD
   export WORKLOAD_NODEPOOL_COUNT=WORKLOAD_NODEPOOL_COUNT
   export TPU_TYPE=TPU_TYPE
   export PROJECT_ID=PROJECT
   export ZONE=ZONE \
   export CLUSTER=CLUSTER

   Substitua:

   • WORKLOAD: defina um nome exclusivo para identificar sua carga de trabalho
   • WORKLOAD_NODEPOOL_COUNT: o número de pools de nós usados por uma carga de trabalho do Pathways
   • TPU_TYPE: o tipo de TPU especifica a versão e o tamanho do Cloud TPU que você quer criar. Para mais informações sobre os tipos de TPU compatíveis com cada versão, consulte Versões de TPU.
   • PROJECT: ID do seu Google Cloud projeto
   • ZONE: a zona em que você planeja executar sua carga de trabalho
   • CLUSTER: o nome do cluster do GKE

2. Crie os contêineres do Pathways no cluster. Para executar uma carga de trabalho sem periféricos, execute o seguinte comando:

   xpk workload create-pathways \
   --headless \
   --workload=${WORKLOAD} \
   --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
   --tpu-type=${TPU_TYPE} \
   --project=${PROJECT} \
   --zone=${ZONE} \
   --cluster=${CLUSTER}

Nesse momento, sua carga de trabalho JAX pode se conectar ao servidor proxy IFRT.

kubectl

O YAML a seguir é o mesmo que o YAML de carga de trabalho em lote, exceto que ele não especifica o contêiner main.

1. Substitua os marcadores de posição, copie o YAML a seguir e cole-o em um arquivo chamado pathways-headless-workload.yaml.

   apiVersion: pathways-job.pathways.domain/v1
   kind: PathwaysJob
   metadata:
     name: pathways-USERNAME
   spec:
     maxRestarts: MAX_RESTARTS
     workers:
       - type: TPU_MACHINE_TYPE
         topology: TOPOLOGY
         numSlices: WORKLOAD_NODEPOOL_COUNT
     pathwaysDir: gs://BUCKET_NAME
     controller:
       deploymentMode: default
       

   Substitua:
   • USERNAME : seu nome de usuário
   • MAX_RESTARTS : o número máximo de vezes que o PathwaysJob pode ser reiniciado
   • TPU_MACHINE_TYPE : o tipo de máquina TPU que você quer usar. Exemplo de valores aceitos: "ct6e-standard-8t", "ct5p-hightpu-4t"
   • TOPOLOGY : a topologia da TPU
   • WORKLOAD_NODEPOOL_COUNT : o número de pools de nós usados por uma carga de trabalho do Pathways
   • BUCKET_NAME : um bucket do Cloud Storage usado para armazenar arquivos temporários
   Para mudar o número de pools de nós (réplicas de worker do Pathways) especificados por WORKLOAD_NODEPOOL_COUNT no YAML anterior, exclua este PathwaysJob e crie um novo PathwaysJob com o número atualizado de pools de nós. Você também precisa reiniciar todos os notebooks conectados para estabelecer a conexão com o novo cluster do Pathways.
2. Aplique o arquivo pathways-headless-workload.yaml:

     kubectl apply -f pathways-headless-workload.yaml
     

3. Execute kubectl get pods para verificar se todos os contêineres no pod estão em execução. A saída a seguir é para um v5p 2x2x2 de 2 fatias, em que USER é o ID do usuário que executa o comando:

       NAME                                         READY   STATUS    RESTARTS   AGE
       pathways-USER-pathways-head-0-0-n848j      2/2     Running   0          49s
       pathways-USER-pathways-workers-0-0-jxt2z   1/1     Running   0          71s
       pathways-USER-pathways-workers-0-1-cxmhc   1/1     Running   0          70s
       pathways-USER-pathways-workers-1-0-5kmz9   1/1     Running   0          71s
       pathways-USER-pathways-workers-1-1-vg5n4   1/1     Running   0          71s
       

Como se conectar ao cluster do Pathways no modo interativo

É possível se conectar ao cluster do Pathways com ou sem encaminhamento de portas. Use uma das seções a seguir para se conectar ao cluster do Pathways.

Conectar usando o encaminhamento de portas

Nesse momento, é possível usar o encaminhamento de portas (de qualquer host com acesso ao plano de controle do cluster) para acessar o servidor proxy:

Use o comando adequado para sua carga de trabalho:

XPK

PROXY_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
PROXY_PORT=29000
kubectl port-forward ${PROXY_POD} ${PROXY_PORT}:${PROXY_PORT}

Você verá uma saída como:

Forwarding from 127.0.0.1:29000 -> 29000
Forwarding from [::1]:29000 -> 29000

kubectl

PROXY_POD=$(kubectl get pods | grep pathways-${USER}-pathways-head | awk '{print $1}')
PROXY_PORT=29000
kubectl port-forward ${PROXY_POD} ${PROXY_PORT}:${PROXY_PORT}

Você verá uma saída como:

Forwarding from 127.0.0.1:29000 -> 29000
Forwarding from [::1]:29000 -> 29000

No mesmo host, abra uma nova janela de terminal. Defina as variáveis de ambiente JAX_PLATFORMS e JAX_BACKEND_TARGET e execute um script Python que importe pathwaysutils e jax:

python3 -m venv .venv
source .venv/bin/activate
pip install pathwaysutils jax

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

O resultado será assim:

[device(144,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(145,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(146,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(147,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(148,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(149,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(150,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(151,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(162,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(163,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(164,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(165,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(166,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(167,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(168,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(169,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
Waiting up to 5 seconds.
Sent all pending logs.
2024-11-13 21:38:51.267523: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled

Conectar-se a hosts na VPC sem usar o encaminhamento de portas

Se você não quiser usar o encaminhamento de portas, poderá se conectar ao cluster do Pathways usando o Cloud DNS ou um balanceador de carga interno.

Conectar usando o Cloud DNS

A ativação do Cloud DNS no cluster muda o provedor de DNS do kube-dns para o Cloud DNS. Quando ativada, uma zona de DNS privada do Cloud DNS é criada na nuvem privada virtual para os nomes do Cloud DNS. Para mais informações, consulte Como usar o Cloud DNS para GKE.

Se você ativar o Cloud DNS com o cluster, a VPC aditiva ou o escopo da VPC, os nomes do Cloud DNS do Kubernetes poderão ser resolvidos em VMs que não sejam do GKE na nuvem privada virtual. Os nomes têm o formato <service_name>.<namespace>.svc.<custom_dns_domain>. O pod principal do Pathways tem um serviço chamado <jobset_name>-pathways-head-0-0.<jobset_name>.<namespace>.svc.<custom_dns_domain>.

Os comandos a seguir mostram como se conectar ao cluster do Pathways usando o Cloud DNS:

1. Confirme se a entrada do Cloud DNS líder pode ser resolvida em um host que não seja do GKE:

   XPK

   host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.USERNAME-test

   Você verá uma saída como:

   <WORKLOAD>-pathways-head-0-0.<WORKLOAD>.default.svc.<user>-test has address 10.0.2.75

   kubectl

   host pathways-USERNAME-pathways-head-0-0.pathways-USERNAME.default.svc.USERNAME-test

   Você verá uma saída como:

   pathways-<user>-pathways-head-0-0.pathways-<user>.default.svc.<user>-test has address 10.0.2.75

2. Conecte-se ao cluster do Pathways usando o nome do Cloud DNS:

   XPK

   JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.USERNAME-test:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

   kubectl

   JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://pathways-USERNAME-pathways-head-0-0.pathways-USERNAME.default.svc.USERNAME-test:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

   Você verá uma saída como:

   [device(216,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(217,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(218,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(219,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(220,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(221,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(222,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(223,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(234,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(235,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(236,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(237,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(238,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
   device(239,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
   device(240,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
   device(241,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
   Waiting up to 5 seconds.
   Sent all pending logs.
   2024-11-14 00:02:49.882044: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled

Conectar usando um balanceador de carga interno

Para um endereço IP particular na VPC que aponta para a implantação do Pathways, crie um serviço com suporte de um balanceador de carga interno. Isso não exige que o cluster tenha o Cloud DNS ativado.

Para clusters com muitas VMs, recomendamos que você ative a criação de subconjuntos do ILB se estiver criando balanceadores de carga internos. Para mais informações, consulte Ativar a criação de subconjuntos do GKE em um cluster atual. Quando a criação de subconjuntos do ILB não está ativada, todos os nós do cluster fazem parte do grupo de instâncias de back-end para todos os balanceadores de carga internos. Isso não é escalonável além de 250 nós. Com a criação de subconjuntos do ILB ativada, o GKE cria grupos de endpoints de rede em vez de grupos de instâncias, e apenas os nós que estão executando um dos pods de serviço do serviço são incluídos. A ativação da criação de subconjuntos do ILB tem uma latência de configuração única (~15 minutos). O comando a seguir mostra como ativar a criação de subconjuntos do ILB:

gcloud container clusters update ${CLUSTER} \
  --project=${PROJECT} \
  [--zone=${ZONE} | --region=${REGION}] \
  --enable-l4-ilb-subsetting

Depois que a criação de subconjuntos do ILB estiver ativada, você poderá criar um serviço do Kubernetes do tipo LoadBalancer usando o YAML a seguir. Isso fará com que o GKE crie um balanceador de carga interno na VPC do cluster:

apiVersion: v1
kind: Service
metadata:
  name: pathways-USERNAME-ilb
  annotations:
    networking.gke.io/load-balancer-type: "Internal"
    networking.gke.io/internal-load-balancer-allow-global-access: "true"
spec:
  type: LoadBalancer
  externalTrafficPolicy: Local
  selector:
    jobset.sigs.k8s.io/jobset-name: pathways-USER
    jobset.sigs.k8s.io/replicatedjob-name: pathways-head
  ports:
  - name: tcp-port
    protocol: TCP
    port: 29000
    targetPort: 29000

Atualize o USER com o ID do usuário e salve o arquivo como pathways-headless-ilb.yaml. Google Cloud

Aplique o manifesto:

kubectl apply -f pathways-headless-ilb.yaml

Depois que o balanceador de carga for criado (cerca de um minuto depois), a coluna EXTERNAL-IP terá um valor:

kubectl get services
NAME                  TYPE           CLUSTER-IP      EXTERNAL-IP   PORT(S)        AGE
pathways-$USER       ClusterIP      None            <none>        <none>         30m
pathways-$USER-ilb   LoadBalancer   34.118.232.46   10.0.0.22     80:31246/TCP   2m41s

É possível acessar a implantação do Pathways sem encaminhamento de portas em hosts na mesma VPC que o cluster:

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://10.0.0.22:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

Observação: substitua o endereço IP do back-end JAX (10.0.0.22) no comando anterior pelo valor exibido por kubectl get services.

Você verá uma saída como:

[device(288,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(289,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(290,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(291,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(292,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(293,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(294,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(295,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(306,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(307,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(308,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(309,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(310,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(311,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(312,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(313,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
Waiting up to 5 seconds.
Sent all pending logs.
2024-11-14 00:30:07.296917: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled

Observação: Seus endereços IP serão diferentes dos mostrados nos exemplos de comando e saída anteriores.

Notebooks do Jupyter

É possível criar um notebook do Jupyter usando o Vertex AI ou criar um notebook do Jupyter auto-hospedado.

Criar uma instância do Vertex AI Workbench

Depois de configurar e verificar o cluster do Pathways, é possível acessar as VMs da TPU do GKE em um notebook do Jupyter da Vertex AI. As instruções de configuração a seguir pressupõem que o cluster do GKE Pathways reside na mesma rede de nuvem privada virtual (que é a rede padrão, a menos que você tenha configurado de outra forma). Acesse o console do Vertex AI Workbench.

Crie uma nova instância do Workbench (na guia Instâncias ) com o botão Criar novo. Verifique se a rede é a mesma do cluster do GKE. É possível usar a linha de comando para criar uma nova instância do Workbench.

gcloud workbench instances create INSTANCE_NAME \
--machine-type=e2-standard-4 \
--data-disk-size=100 \
--location=ZONE \
[--network=NETWORK]

Observação: a flag --network é opcional. Se você a omitir, a rede padrão será usada.

Depois que a instância for criada, acesse-a e clique em Abrir Jupyterlab.

Criar uma instância de notebook do Jupyter auto-hospedada

O comando a seguir mostra como criar uma instância de notebook do Jupyter auto-hospedada usando o XPK:

xpk workload create-pathways \
--workload=${WORKLOAD} \
--num-slices=${WORKLOAD_NODEPOOL_COUNT} \
--tpu-type=${TPU_TYPE} \
--project=${PROJECT} \
--zone=${ZONE} \
--cluster=${CLUSTER} \
--docker-image=jupyter/base-notebook \
--command "start-notebook.sh"

O YAML a seguir mostra como criar uma instância de notebook do Jupyter auto-hospedada usando o kubectl. Aplique o YAML a seguir depois que um cluster do Pathways sem periféricos for criado. Para mais informações, consulte Executar o Pathways no modo interativo com o kubectl.

apiVersion: batch/v1
kind: Job
metadata:
  name: jupyter-notebook-USERNAME
spec:
  template:
    spec:
      restartPolicy: OnFailure
      containers:
      - name: jupyter-notebook
        image: jupyter/base-notebook  # Use the appropriate Jupyter image
        ports:
        - containerPort: 8888

Acessar a instância do notebook

Conecte-se ao notebook da sua máquina local usando o encaminhamento de portas:

XPK

  MAIN_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
  kubectl port-forward pod/${MAIN_POD} 8888:8888

kubectl

  MAIN_POD=$(kubectl get pods | grep jupyter-notebook-USERNAME | awk '{print $1}')
  kubectl port-forward pod/${MAIN_POD} 8888:8888

Navegue no navegador local até http://localhost:8888?token=<var>your-token</var>. Substitua <your-token> pelo token dos registros do contêiner do notebook do Jupyter.

kubectl logs ${MAIN_POD}

Que deve gerar:

...
Or copy and paste one of these URLs:
  http://jupyter-notebook-<user>-bbbdh:8888/lab?token=<token>
  http://127.0.0.1:8888/lab?token=<token>

Conectividade do notebook com o cluster do Pathways

1. No Jupyterlab, crie um novo notebook Python 3
2. Conecte-se ao servidor proxy do Pathways

No notebook, adicione uma célula para instalar pathwaysutils, defina JAX_PLATFORMS como proxy e defina JAX_BACKEND_TARGET como PROXY_ADDRESS.

**Observação**: O valor de PROXY_ADDRESS depende da opção de conexão escolhida em Conectar-se ao servidor proxy do Pathways.

!pip install pathwaysutils
%env JAX_PLATFORMS=proxy
# Replace your proxy address below:
%env JAX_BACKEND_TARGET=PROXY_ADDRESS

Adicione uma segunda célula como uma verificação do tipo "Olá, mundo" e imprima os dispositivos no cluster do Pathways.

import pathwaysutils
import jax

pathwaysutils.initialize()
print(jax.devices())

Se tudo estiver funcionando bem, você verá uma mensagem indicando que o back-end do Pathways no Cloud foi detectado.

O número de dispositivos JAX listados precisa corresponder ao número de chips de TPU e ao número de fatias especificados ao criar o cluster do Pathways.

Adicionar seu código a um notebook

Adicione seu próprio código JAX e execute interativamente nas TPUs no cluster do Pathways. O código a seguir mostra como realizar cálculos em duas fatias de um único notebook.

import jax
import jax.numpy as jnp
from jax import lax
import numpy as np

# You can use JAX APIs as usual across any of the devices.
jax.jit(jnp.sin, device=jax.devices()[-1])(np.pi / 2.)

# pmap can run across all devices on all slices
num_tpus = jax.device_count()
f = jax.pmap(lambda x: lax.psum(1, 'i'), 'i')
x = jnp.arange(num_tpus)
y = f(x)
print(y)

# You can also target devices from a specific slice
slice0_devices = [d for d in jax.devices() if d.slice_index == 0]
f = jax.pmap(lambda x: lax.psum(1, 'i'), 'i', devices=slice0_devices)
x = jnp.arange(len(slice0_devices))
y = f(x)
print(y)
print(y.global_shards)

# You can send data produced on one slice to another slice
slice1_devices = [d for d in jax.devices() if d.slice_index == 1]
g = jax.pmap(lambda x: x + lax.axis_index('i'), 'i', devices=slice1_devices)
z = g(y)
print(z)
print(z.global_shards)

Excluir o cluster interativo do Pathways

XPK

xpk workload delete --workload=WORKLOAD --cluster=CLUSTER --project=PROJECT --zone=ZONE

kubectl

kubectl delete -f pathways-headless-workload.yaml

A seguir

• Criar um cluster do GKE com o Pathways
• Inferência com vários hosts com o Pathways
• Cargas de trabalho em lote com o Pathways
• Modo interativo do Pathways
• Como migrar cargas de trabalho JAX para o Pathways
• Treinamento resiliente com o Pathways
• Solução de problemas do Pathways