Ejecuta una carga de trabajo interactiva con Pathways

Las cargas de trabajo interactivas de Pathways son cargas de trabajo de JAX remotas que se ejecutan fuera de los componentes del controlador de Pathways del clúster de Google Kubernetes Engine (GKE), por ejemplo, en una VM. A diferencia de las cargas de trabajo por lotes, cuando finaliza una carga de trabajo interactiva, no se cierran los componentes del clúster de Pathways, que permanecen disponibles para las conexiones de otros clientes de JAX. En este documento, se usa un notebook de Jupyter como ejemplo para demostrar las cargas de trabajo interactivas.

Con la interfaz de IFRT, los usuarios de JAX envían comandos a un clúster de Pathways. El código de JAX, ya sea que se ejecute desde una terminal, un notebook o cualquier entorno compatible con Python, puede interactuar con los recursos de Pathways.

Antes de comenzar

Completa los siguientes requisitos previos:

• Crea un clúster de GKE con Pathways.
• Instala Cluster Toolkit.
• Instala las herramientas de Kubernetes.
• Instala Google Cloud CLI.
• Habilita la API de Kubernetes Engine.
• Asegúrate de que Pathways esté habilitado para tu proyecto de Google Cloud .

Ejecuta Pathways en modo interactivo

1. Configura las siguientes variables de entorno:

   export WORKLOAD=WORKLOAD
   export WORKLOAD_NODEPOOL_COUNT=WORKLOAD_NODEPOOL_COUNT
   export COMPUTE_TYPE="COMPUTE_TYPE"
   export TOPOLOGY="TOPOLOGY"
   export PROJECT_ID=PROJECT_ID
   export REGION=REGION
   export ZONE=ZONE
   export CLUSTER=CLUSTER
   export BUCKET_NAME=BUCKET_NAME

   Reemplaza lo siguiente:

   • WORKLOAD: Es un nombre único para identificar tu carga de trabajo (22 caracteres o menos debido a los límites de las etiquetas de DNS).
   • WORKLOAD_NODEPOOL_COUNT: es la cantidad de grupos de nodos que usa una carga de trabajo de Pathways.
   • COMPUTE_TYPE: Es el tipo de máquina de TPU, por ejemplo, ct6e-standard-4t. Para obtener más información sobre los tipos de TPU compatibles con cada versión de TPU, consulta Versiones de TPU.
   • TOPOLOGY: la topología de ubicación de la TPU, por ejemplo, 2x4.
   • PROJECT_ID: Es el ID del proyecto deGoogle Cloud .
   • REGION: Es la región en la que se encuentran tu bucket y tu clúster.
   • ZONE: Es la zona en la que planeas ejecutar tu carga de trabajo.
   • CLUSTER: Es el nombre del clúster de GKE.
   • BUCKET_NAME: Es el bucket de Cloud Storage para los artefactos de Pathways.

2. Si aún no creaste un bucket de Cloud Storage para los artefactos de Pathways, crea uno:

   gcloud storage buckets create gs://${BUCKET_NAME} \
       --project=${PROJECT_ID} \
       --location=${REGION}
   

3. Crea los contenedores de Pathways en el clúster. Para ejecutar una carga de trabajo sin interfaz gráfica, usa el comando gcluster job submit (para obtener más información, consulta la Guía de envío de trabajos del Cluster Toolkit):

   gcluster job submit \
       --pathways \
       --pathways-gcs-location="gs://${BUCKET_NAME}/pathways-artifacts" \
       --name=${WORKLOAD} \
       --cluster=${CLUSTER} \
       --project=${PROJECT_ID} \
       --location=${ZONE} \
       --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
       --compute-type=${COMPUTE_TYPE} \
       --topology=${TOPOLOGY} \
       --pathways-headless
   

   El comando gcluster job submit busca un grupo de nodos de CPU llamado cpu-np o pathways-np para ejecutar el Pod principal de Pathways. Si tu grupo de nodos de CPU tiene un nombre diferente (como cpu-pathways-np de la guía de creación de clústeres), agrega la marca --pathways-head-np=NODE_POOL_NAME a los comandos gcluster job submit de esta página. De lo contrario, el comando falla con el error failed to auto-detect a suitable CPU node pool.

En este punto, tu carga de trabajo de JAX puede conectarse al servidor proxy de IFRT.

Conéctate al clúster de Pathways en modo interactivo

Puedes conectarte al clúster de Pathways con o sin redirección de puertos. Usa una de las siguientes secciones para conectarte al clúster de Pathways.

Conéctate con la redirección de puertos

Desde cualquier host con acceso al plano de control de tu clúster, reenvía el puerto local 29000 al Pod principal de Pathways para acceder al servidor proxy:

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

El resultado es similar a lo siguiente:

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

En el mismo host, abre una nueva ventana de terminal. Con Python 3.11 o una versión posterior, crea un entorno virtual, instala los paquetes pathwaysutils y jax, establece las variables de entorno JAX_PLATFORMS y JAX_BACKEND_TARGET, y ejecuta una secuencia de comandos de Python que importe pathwaysutils y 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())'

El resultado es similar a lo siguiente:

[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

Conéctate desde hosts en la VPC sin usar la redirección de puertos

Si no quieres usar la redirección de puertos, puedes conectarte al clúster de Pathways con Cloud DNS o un balanceador de cargas interno.

Conéctate con Cloud DNS

Habilitar Cloud DNS en tu clúster cambia el proveedor de DNS de kube-dns a Cloud DNS. Cuando se habilita, GKE crea una zona privada de Cloud DNS en tu red de nube privada virtual (VPC) para los nombres de DNS. Para obtener más información, consulta Usa Cloud DNS para GKE.

Si habilitas Cloud DNS con el permiso de VPC o el permiso de VPC adicional, los nombres de DNS de Kubernetes se pueden resolver desde VMs que no son de GKE dentro de tu red de VPC. Con el permiso de clúster, solo los nodos del clúster pueden resolver estos nombres. Los nombres tienen el formato SERVICE_NAME.NAMESPACE.svc.CUSTOM_DNS_DOMAIN. El Pod principal de Pathways tiene un Service llamado WORKLOAD-pathways-head-0-0.WORKLOAD.NAMESPACE.svc.CUSTOM_DNS_DOMAIN.

En los siguientes comandos, se muestra cómo conectarte a tu clúster de Pathways con Cloud DNS:

1. Confirma que la entrada de DNS líder se pueda resolver desde un host que no sea de GKE:

   host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN

   Reemplaza lo siguiente:

   • WORKLOAD: Es el nombre de tu carga de trabajo de Pathways.
   • CUSTOM_DNS_DOMAIN: Es el dominio de DNS personalizado configurado para tu clúster.

   El resultado es similar a lo siguiente:

   WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN has address 10.0.2.75

2. Conéctate al clúster de Pathways con el nombre de DNS:

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

   El resultado es similar a lo siguiente:

   [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
   

Conéctate con un balanceador de cargas interno

Para una dirección IP privada en tu red de VPC que apunte a tu implementación de Pathways, crea un Service respaldado por un balanceador de cargas interno (ILB). Esto no requiere que tu clúster tenga Cloud DNS habilitado.

En el caso de los clústeres con muchas VMs, habilita la subdivisión de ILB si creas balanceadores de cargas internos. Para obtener más información, consulta Habilita la creación de subconjuntos de GKE en un clúster existente. Cuando la subdivisión de ILB no está habilitada, todos los nodos del clúster forman parte del grupo de instancias de backend para todos los balanceadores de cargas internos, que no escala más allá de 250 nodos. Con la subdivisión de ILB habilitada, GKE crea grupos de extremos de red en lugar de grupos de instancias y solo incluye nodos que ejecutan uno de los Pods de entrega del Service. La habilitación de la subdivisión de ILB tiene una latencia de configuración única de alrededor de 15 minutos. El siguiente comando habilita la creación de subconjuntos de ILB:

gcloud container clusters update ${CLUSTER} \
    --project=${PROJECT_ID} \
    --location=CLUSTER_LOCATION \
    --enable-l4-ilb-subsetting

Reemplaza CLUSTER_LOCATION por la región o zona en la que se encuentra el clúster.

Después de habilitar la subdivisión de ILB, crea un Service de Kubernetes de tipo LoadBalancer con el siguiente manifiesto de YAML. Esto hace que GKE cree un balanceador de cargas interno dentro de la red de VPC del clúster:

apiVersion: v1
kind: Service
metadata:
  name: WORKLOAD-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: WORKLOAD
    jobset.sigs.k8s.io/replicatedjob-name: pathways-head
  ports:
  - name: tcp-port
    protocol: TCP
    port: 29000
    targetPort: 29000

Reemplaza WORKLOAD por el nombre de tu carga de trabajo de Pathways y guarda el archivo como pathways-headless-ilb.yaml.

Aplica el manifiesto

kubectl apply -f pathways-headless-ilb.yaml

Después de que GKE cree el balanceador de cargas (aproximadamente un minuto después), verifica que la columna EXTERNAL-IP muestre una dirección IP:

kubectl get services

El resultado es similar a lo siguiente:

NAME             TYPE           CLUSTER-IP      EXTERNAL-IP   PORT(S)           AGE
WORKLOAD       ClusterIP      None            <none>        <none>            30m
WORKLOAD-ilb   LoadBalancer   34.118.232.46   10.0.0.22     29000:31246/TCP   2m41s

Puedes acceder a la implementación de Pathways sin redirección de puertos desde hosts en la misma red de VPC que tu clúster:

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

Reemplaza EXTERNAL_IP por el valor EXTERNAL-IP que muestra kubectl get services, por ejemplo, 10.0.0.22.

El resultado es similar a lo siguiente:

[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

Nota: Tus direcciones IP difieren de las que se muestran en los ejemplos de comandos y resultados anteriores.

Notebooks de Jupyter

Puedes crear un notebook de Jupyter con Vertex AI o puedes crear un notebook de Jupyter alojado de forma automática.

Crea una instancia de Vertex AI Workbench

Después de configurar y verificar tu clúster de Pathways, puedes acceder a las VMs de TPU de GKE desde un notebook de Jupyter de Vertex AI Workbench. En las siguientes instrucciones de configuración, se supone que tu clúster de GKE Pathways reside en la misma red de VPC (que es la red predeterminada, a menos que hayas configurado otra).

1. Ve a la página Vertex AI Workbench en la consola de Google Cloud .

2. En la pestaña Instancias, haz clic en Crear nueva y asegúrate de que la red coincida con la red de tu clúster de GKE.

   Como alternativa, usa la línea de comandos para crear una instancia nueva de Vertex AI Workbench:

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

   Reemplaza lo siguiente:

   • INSTANCE_NAME: Es el nombre de tu instancia de Vertex AI Workbench.
   • PROJECT_ID: Es el ID del proyecto deGoogle Cloud .
   • ZONE: Es la zona en la que deseas crear la instancia.
   • NETWORK: Es el nombre de la red de VPC de tu clúster de GKE. La marca --network es opcional. Si la omites, el comando usa la red default.

3. Después de que Vertex AI cree la instancia, ve a ella y haz clic en Abrir JupyterLab.

Crea una instancia de notebook de Jupyter autoalojada

En el siguiente comando, se muestra cómo crear una instancia de notebook de Jupyter autohospedada con el comando gcluster job submit:

gcluster job submit \
    --pathways \
    --name=${WORKLOAD} \
    --cluster=${CLUSTER} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
    --compute-type=${COMPUTE_TYPE} \
    --topology=${TOPOLOGY} \
    --pathways-gcs-location="gs://${BUCKET_NAME}/pathways-artifacts" \
    --image=jupyter/base-notebook \
    --command="start-notebook.sh"

En el siguiente manifiesto de YAML, se muestra cómo crear una instancia de notebook de Jupyter autohospedada con kubectl. Aplica el siguiente manifiesto de YAML después de crear un clúster de Pathways sin interfaz gráfica. Para obtener más información, consulta Ejecuta Pathways en modo interactivo:

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

El campo hostNetwork: true permite que el cliente de JAX en el notebook complete el protocolo de enlace de la seguridad de transporte de la capa de la aplicación (ALTS) con el servidor proxy de Pathways. Sin este campo, la conexión falla en los grupos de nodos que usan Workload Identity Federation for GKE. El campo dnsPolicy: ClusterFirstWithHostNet permite que el notebook resuelva los nombres de DNS en el clúster, como el nombre del Pod principal de Pathways. Para obtener más información, consulta Soluciona problemas de Pathways en Cloud.

Conéctate a la instancia de notebook

Conéctate al notebook desde tu máquina local con la redirección de puertos:

Cluster Toolkit

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 \
  --selector=batch.kubernetes.io/job-name=WORKLOAD \
  --output=jsonpath='{.items[0].metadata.name}')
kubectl port-forward pod/${MAIN_POD} 8888:8888

En tu navegador local, ve a http://localhost:8888?token=TOKEN. Reemplaza TOKEN por el token de los registros del contenedor del notebook de Jupyter:

kubectl logs ${MAIN_POD}

El resultado es similar a lo siguiente:

...
Or copy and paste one of these URLs:
  http://NODE_NAME:8888/lab?token=TOKEN
  http://127.0.0.1:8888/lab?token=TOKEN

NODE_NAME es el nombre de host del nodo de GKE que ejecuta el Pod del notebook (por ejemplo, gke-my-cluster-cpu-np-1a2b3c4d-ef56).

Conectividad del notebook al clúster de Pathways

1. En JupyterLab, crea un notebook de Python 3 nuevo.
2. Conéctate al servidor proxy de Pathways.

En el notebook, agrega una celda para instalar pathwaysutils, establece JAX_PLATFORMS en proxy y JAX_BACKEND_TARGET en grpc://PROXY_ADDRESS:29000:

Nota: El valor de PROXY_ADDRESS depende de la opción de conexión que elegiste en Conéctate al clúster de Pathways en modo interactivo.

!pip install pathwaysutils
# Set JAX_PLATFORMS and JAX_BACKEND_TARGET if you're not using `gcluster job submit` to deploy the notebook
%env JAX_PLATFORMS=proxy
  # Replace your proxy address in the following line:
%env JAX_BACKEND_TARGET=grpc://PROXY_ADDRESS:29000

Reemplaza PROXY_ADDRESS por el host o la dirección IP de tu servidor proxy de Pathways de Conéctate al clúster de Pathways en modo interactivo (por ejemplo, 127.0.0.1, el nombre de host de Cloud DNS o la dirección IP del balanceador de cargas interno). Si implementaste un notebook alojado de forma automática con kubectl, reemplaza PROXY_ADDRESS por el nombre de DNS en el clúster del Pod principal de Pathways, de modo que el valor de JAX_BACKEND_TARGET sea grpc://WORKLOAD-pathways-head-0-0.WORKLOAD:29000.

Agrega una segunda celda para verificar la conectividad y mostrar los dispositivos en el clúster de Pathways:

import pathwaysutils
import jax

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

Si la conexión se realiza correctamente, el resultado incluye un mensaje que indica que JAX detectó el backend de Pathways-on-Cloud.

La cantidad de dispositivos JAX que se indican coincide con la cantidad de chips TPU y la cantidad de porciones que especificaste cuando creaste el clúster de Pathways.

Agrega tu código a un notebook

Agrega tu propio código de JAX y ejecútalo de forma interactiva en las TPU del clúster de Pathways. En el siguiente código, se muestra cómo realizar cálculos en dos porciones desde un solo notebook. Para ejecutar el ejemplo completo, crea tu carga de trabajo de Pathways con al menos dos porciones configurando WORKLOAD_NODEPOOL_COUNT en 2 o más. Con una porción, slice1_devices está vacío y la última parte del ejemplo falla con un ValueError:

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.
x = jax.device_put(np.pi / 2., jax.devices()[-1])
jax.jit(jnp.sin)(x)

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

Borra tu carga de trabajo interactiva de Pathways

Cluster Toolkit

gcluster job cancel WORKLOAD \
  --cluster=CLUSTER \
  --project=PROJECT_ID \
  --location=ZONE

kubectl

kubectl delete job WORKLOAD

¿Qué sigue?

• Crea un clúster de GKE con Pathways
• Ejecuta una carga de trabajo por lotes con Pathways
• Porta cargas de trabajo de JAX a Pathways
• Entrenamiento resiliente con Pathways
• Soluciona problemas de Pathways en Cloud