Ejecuta una carga de trabajo interactiva con Pathways

Vista previa

Este producto o esta función se encuentra sujeto a las "Condiciones de la oferta de la fase previa a la DG" en la sección Condiciones del servicio generales de las Condiciones específicas del servicio. Los productos y las funciones de la fase previa a la DG están disponibles “tal como están” y pueden estar sujetos a asistencia limitada. Para obtener más información, consulta las descripciones de la etapa de lanzamiento.

Las cargas de trabajo interactivas de Pathways son cargas de trabajo JAX remotas que se ejecutan dentro de una VM que no forma parte del clúster de GKE que aloja el clúster de Pathways. A diferencia de las cargas de trabajo por lotes, la finalización de la operación de carga de trabajo interactiva no cierra los componentes del clúster de Pathways, que permanecen disponibles para la conexión de otros clientes JAX. En este documento, se usa un notebook de Jupyter como ejemplo para demostrar las cargas de trabajo interactivas.

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

Antes de comenzar

Asegúrate de tener lo siguiente:

• Creaste un clúster de GKE con XPK.
• Instalaste XPK.
• Instalaste herramientas de Kubernetes.
• Instalaste gcloud CLI.
• Habilitaste la API de TPU.
• Habilitaste la API de Google Kubernetes Engine.
• Asegúrate de que Pathways esté habilitado para tu Google Cloud proyecto.

Ejecuta Pathways en modo interactivo

Puedes ejecutar Pathways en modo interactivo con xpk o kubectl.

XPK

1. Configura las siguientes variables de entorno:

   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

   Reemplaza lo siguiente:

   • WORKLOAD: establécelo en un nombre único para identificar tu carga de trabajo.
   • WORKLOAD_NODEPOOL_COUNT: es la cantidad de grupos de nodos que usa una carga de trabajo de Pathways.
   • TPU_TYPE: el tipo de TPU 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 TPU compatibles con cada versión de TPU, consulta Versiones de TPU
   • PROJECT: es el ID de tu Google Cloud proyecto.
   • ZONE: es la zona en la que planeas ejecutar tu carga de trabajo.
   • CLUSTER: es el nombre de tu clúster de GKE.

2. Crea los contenedores de Pathways en el clúster. Para ejecutar una carga de trabajo sin encabezado, ejecuta el siguiente comando:

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

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

kubectl

El siguiente YAML es el mismo que el YAML de carga de trabajo por lotes, excepto que no especifica el contenedor main.

1. Reemplaza los marcadores de posición, copia el siguiente YAML y pégalo en un archivo llamado 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
       

   Reemplaza lo siguiente:
   • USERNAME : tu nombre de usuario
   • MAX_RESTARTS : la cantidad máxima de veces que se puede reiniciar PathwaysJob
   • TPU_MACHINE_TYPE : el tipo de máquina de TPU que deseas usar, por ejemplo, valores compatibles: "ct6e-standard-8t", "ct5p-hightpu-4t"
   • TOPOLOGY : la topología de TPU
   • WORKLOAD_NODEPOOL_COUNT : la cantidad de grupos de nodos que usa una carga de trabajo de Pathways
   • BUCKET_NAME : un bucket de Cloud Storage que se usa para almacenar archivos temporales
   Para cambiar la cantidad de grupos de nodos (réplicas de trabajadores de Pathways) especificados por WORKLOAD_NODEPOOL_COUNT en el YAML anterior, debes borrar este PathwaysJob y crear un PathwaysJob nuevo con la cantidad actualizada de grupos de nodos. También debes reiniciar los notebooks conectados para establecer la conexión con el nuevo clúster de Pathways.
2. Aplica el archivo pathways-headless-workload.yaml:

     kubectl apply -f pathways-headless-workload.yaml
     

3. Ejecuta kubectl get pods para verificar que todos los contenedores del Pod se estén ejecutando. El siguiente resultado es para una porción 2 de v5p 2x2x2, en la que USER es el ID del usuario que ejecuta el 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
       

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

En este punto, puedes usar la redirección de puertos (desde cualquier host con acceso al plano de control de tu clúster) para acceder al servidor proxy:

Usa el comando adecuado para tu carga de trabajo:

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}

Deberías ver un resultado similar a este:

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}

Deberías ver un resultado similar a este:

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

En el mismo host, abre una nueva ventana de terminal. Configura 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())'

Deberías ver un resultado como el 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 deseas 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

Si habilitas Cloud DNS en tu clúster, el proveedor de Cloud DNS cambia de kube-dns a Cloud DNS. Cuando está habilitada, se crea una zona privada de Cloud DNS en tu nube privada virtual para los nombres de Cloud DNS. Para obtener más información, consulta Usa Cloud DNS para GKE.

Si habilitas Cloud DNS con el clúster, la VPC adicional o el permiso de VPC, los nombres de Cloud DNS de Kubernetes se pueden resolver desde VMs que no son de GKE dentro de tu nube privada virtual. Los nombres tienen el formato <service_name>.<namespace>.svc.<custom_dns_domain>. El Pod principal de Pathways tiene un servicio llamado <jobset_name>-pathways-head-0-0.<jobset_name>.<namespace>.svc.<custom_dns_domain>.

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

1. Confirma que la entrada de Cloud DNS principal se pueda resolver desde un host que no sea de GKE:

   XPK

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

   Deberías ver un resultado similar a este:

   <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

   Deberías ver un resultado similar a este:

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

2. Conéctate al clúster de Pathways con el nombre de 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())'

   Deberías ver un resultado similar a este:

   [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 VPC que apunte a tu implementación de Pathways, crea un servicio respaldado por un balanceador de cargas interno. Esto no requiere que tu clúster tenga habilitado Cloud DNS.

Para los clústeres con muchas VMs, te recomendamos que habilites la subdivisión de ILB si creas balanceadores de cargas internos. Para obtener más información, consulta Habilita la subdivisión de GKE en un clúster existente. Cuando la subdivisión de ILB no está habilitada, todos los nodos del clúster formarán parte del grupo de instancias de backend para todos los balanceadores de cargas internos. Esto no se 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 se incluyen los nodos que ejecutan uno de los Pods de servicio del servicio. La habilitación de la subdivisión de ILB tiene una latencia de configuración única (~15 minutos). En el siguiente comando, se muestra cómo habilitar la subdivisión de ILB:

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

Una vez que se habilita la subdivisión de ILB, puedes crear un servicio de Kubernetes de tipo LoadBalancer con el siguiente YAML. Esto hará que GKE cree un balanceador de cargas interno dentro de la VPC de tu clúster:

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

Actualiza el USER con tu Google Cloud ID de usuario y guarda el archivo como pathways-headless-ilb.yaml.

Aplica el manifiesto

kubectl apply -f pathways-headless-ilb.yaml

Después de crear el balanceador de cargas (aproximadamente un minuto después), la columna EXTERNAL-IP tendrá un 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

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

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

Nota: Reemplaza la dirección IP del backend de JAX (10.0.0.22) en el comando anterior por el valor que muestra kubectl get services.

Deberías ver un resultado similar a este:

[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 serán diferentes 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 un notebook de Jupyter autoalojado.

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. En las siguientes instrucciones de configuración, se supone que tu clúster de Pathways de GKE reside en la misma red de nube privada virtual (que es la red predeterminada, a menos que hayas configurado lo contrario). Navega a la consola de Vertex AI Workbench.

Crea una instancia nueva de Workbench (desde la pestaña Instancias) con el botón Crear nueva. Asegúrate de que la red sea la misma que la red de tu clúster de GKE. Puedes usar la línea de comandos para crear una instancia nueva de Workbench.

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

Nota: la marca --network es opcional. Si la omites, se usa la red predeterminada.

Una vez que se cree la instancia, navega 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 autoalojada con 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"

En el siguiente YAML, se muestra cómo crear una instancia de notebook de Jupyter autoalojada con kubectl. Aplica el siguiente YAML después de crear un clúster de Pathways sin encabezado. Para obtener más información, consulta Ejecuta Pathways en modo interactivo con 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

Navega a la instancia de notebook

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

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

En tu navegador local, navega a http://localhost:8888?token=<var>your-token</var>. Reemplaza <your-token> por el token de los registros del contenedor del notebook de Jupyter.

kubectl logs ${MAIN_POD}

Debería mostrar lo siguiente:

...
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>

Conectividad de notebook al clúster de Pathways

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

En el notebook, agrega una celda para instalar pathwaysutils, establece JAX_PLATFORMS en proxy y establece JAX_BACKEND_TARGET en PROXY_ADDRESS.

Nota: El valor de PROXY_ADDRESS depende de la opción de conexión que elegiste en Conéctate al servidor proxy de Pathways.

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

Agrega una segunda celda como una verificación de tipo "Hola mundo" e imprime los dispositivos en el clúster de Pathways.

import pathwaysutils
import jax

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

Si todo funciona bien, deberías ver un mensaje que indique que se detectó el backend de Pathways en Cloud.

La cantidad de dispositivos JAX que se muestran debe coincidir 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 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.

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)

Borra tu clúster interactivo de Pathways

XPK

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

kubectl

kubectl delete -f pathways-headless-workload.yaml

¿Qué sigue?

• Crea un clúster de GKE con Pathways.
• Inferencia de varios hosts con Pathways
• Cargas de trabajo por lotes con Pathways
• Modo interactivo de Pathways
• Cómo transferir cargas de trabajo JAX a Pathways
• Entrenamiento resistente con Pathways
• Solución de problemas de Pathways