Executar uma carga de trabalho interativa com os programas

Pré-lançamento

Este produto ou recurso está sujeito aos "Termos das Soluções Pré-GA" na seção Termos Gerais de Serviço dos Termos Específicos de Serviço. Os produtos e recursos pré-GA estão disponíveis no estado em que se encontram e podem ter assistência limitada. Para mais informações, consulte as descrições das fases de lançamento.

As cargas de trabalho interativas dos programas são cargas de trabalho JAX remotas executadas em uma VM que não faz parte do cluster do GKE que hospeda o cluster dos programas. 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 Jupyter como exemplo para demonstrar cargas de trabalho interativas.

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

Antes de começar

Você precisa ter:

• Criar um cluster do GKE usando o XPK
• XPK instalado
• Ferramentas do Kubernetes instaladas
• Instale a CLI gcloud.
• Ativou a API TPU
• Ativou a API Google Kubernetes Engine
• Verifique se os programas de aprendizado estão ativados no seu projeto do Google Cloud

Executar programas no modo interativo

É possível executar os programas de aprendizado no modo interativo usando xpk ou kubectl.

XPK

1. Configure as variáveis de ambiente a seguir:

   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 da Cloud TPU que você quer criar. Para mais informações sobre os tipos de TPU disponíveis em cada versão, consulte Versões de TPU.
   • PROJECT: o ID do projeto do Google Cloud
   • 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 programa no cluster. Para executar uma carga de trabalho sem interface gráfica, 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}

Neste ponto, sua carga de trabalho do JAX pode se conectar ao servidor proxy do IFRT.

kubectl

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

1. Substitua os marcadores de posição, copie o seguinte YAML e cole 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 da TPU que você quer usar. Exemplos 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 workers de programas de aprendizado) especificado 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. Também é necessário reiniciar os notebooks conectados para estabelecer a conexão com o novo cluster do programa.
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 uma fatia 2 v5p 2x2x2, 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 programa de aprendizado no modo interativo

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

Conectar usando o encaminhamento de portas

Neste ponto, é possível usar o encaminhamento de porta (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 do terminal. Defina as variáveis de ambiente JAX_PLATFORMS e JAX_BACKEND_TARGET e execute um script Python que importa 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 de hosts na VPC sem usar encaminhamento de portas

Se não quiser usar o encaminhamento de porta, conecte-se ao cluster do programa usando o Cloud DNS ou um balanceador de carga interno.

Conectar usando o Cloud DNS

Ao ativar o Cloud DNS no cluster, o provedor de DNS muda de kube-dns para Cloud DNS. Quando ativada, uma zona privada do Cloud DNS é criada na sua 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 são do GKE na sua nuvem privada virtual. Os nomes têm o formato <service_name>.<namespace>.svc.<custom_dns_domain>. O pod principal do programa 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 principal do Cloud DNS 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 sua VPC que aponta para a implantação de programas, 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 subconfiguração 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. 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 é escalonado além de 250 nós. Com o subagrupamento do ILB ativado, o GKE cria grupos de endpoint de rede em vez de grupos de instâncias, e apenas os nós que executam um dos pods de exibição do serviço são incluídos. A ativação da subdivisão do ILB tem uma latência de configuração única (~15 minutos). O comando a seguir mostra como ativar a subdefinição de ILB:

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

Depois que o subsetting do ILB é ativado, é possível criar um serviço do Kubernetes do tipo LoadBalancer usando o seguinte YAML. 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 seu ID de usuário do Google Cloud e salve o arquivo como pathways-headless-ilb.yaml.

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 de programas sem encaminhamento de porta em hosts na mesma VPC do 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 no comando e nos exemplos de saída anteriores.

Notebooks do Jupyter

É possível criar um notebook do Jupyter usando a Vertex AI ou um notebook do Jupyter autohospedado.

Criar uma instância do Vertex AI Workbench

Depois de configurar e verificar o cluster do programa, é possível acessar as VMs de TPU do GKE em um notebook Jupyter da Vertex AI. As instruções de configuração a seguir pressupõem que o cluster do GKE Pathways esteja 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 instância do Workbench na guia Instâncias usando o botão Criar nova. Verifique se a rede é a mesma do cluster do GKE. É possível usar a linha de comando para criar uma 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ê omitir, a rede padrão será usada.

Quando a instância for criada, navegue até ela e clique em Abrir JupyterLab.

Criar uma instância de notebook Jupyter autohospedada

O comando a seguir mostra como criar uma instância de notebook do Jupyter autohospedada usando 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 Jupyter autohospedada usando kubectl. Aplique o YAML a seguir depois que um cluster do Pathways sem interface gráfica for criado. Para mais informações, consulte Executar programas de aprendizado no modo interativo com 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

Navegue até a instância do notebook

Conecte-se ao notebook da 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

No navegador local, acesse http://localhost:8888?token=<var>your-token</var>. Substitua <your-token> pelo token dos registros do contêiner do notebook Jupyter.

kubectl logs ${MAIN_POD}

que vai 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 programa

1. No JupyterLab, crie um notebook Python 3.
2. Conectar-se ao servidor proxy do programa de aprendizado

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 ao servidor proxy do programa de aprendizado.

!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 "hello world" e imprima os dispositivos no cluster Pathways.

import pathwaysutils
import jax

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

Se tudo estiver funcionando bem, você vai receber uma mensagem indicando que o back-end do Pathways on Cloud foi detectado.

O número de dispositivos JAX listados precisa corresponder ao número de chips de TPU e ao número de frações especificadas ao criar o cluster do programa de aprendizado.

Adicionar seu código a um notebook

Adicione seu próprio código JAX e execute de forma interativa nas TPUs do cluster Pathways. O código a seguir mostra como realizar cálculos em duas frações 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 programa 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 usando o Pathways
• Cargas de trabalho em lote com os programas
• Modo interativo dos programas
• Como migrar cargas de trabalho do JAX para o Pathways
• Treinamento resiliente com programas de aprendizado
• Programas de solução de problemas