Ejecuta una carga de trabajo interactiva con Pathways

Versión preliminar

Este producto o función está sujeto a las "Condiciones de la Oferta de la Fase Previa a la DG" en la sección Condiciones Generales del Servicio de las Condiciones Específicas del Servicio. Los productos y las funciones de la fase previa a la DG están disponibles sin modificaciones 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 de 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 de JAX. En este documento, se usa un notebook de Jupyter como ejemplo para demostrar 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 sin problemas con los recursos de Pathways.

Antes de comenzar

Asegúrate de tener lo siguiente:

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

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: Establece 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: El ID de tu proyecto de Google Cloud
   • ZONE: Es la zona en la que planeas ejecutar tu carga de trabajo.
   • CLUSTER: Es el nombre del 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 de JAX puede conectarse al servidor proxy de IFRT.

kubectl

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

1. Reemplaza los marcadores de posición, copia el siguiente código 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 : Es la cantidad máxima de veces que se puede reiniciar el PathwaysJob.
   • TPU_MACHINE_TYPE: El tipo de máquina de TPU que deseas usar. Algunos valores admitidos son "ct6e-standard-8t" y "ct5p-hightpu-4t".
   • TOPOLOGY : Es la topología de la TPU.
   • WORKLOAD_NODEPOOL_COUNT : Es la cantidad de grupos de nodos que usa una carga de trabajo de Pathways.
   • BUCKET_NAME : Es 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 rutas de acceso) especificada por WORKLOAD_NODEPOOL_COUNT en el YAML anterior, debes borrar este PathwaysJob y crear un nuevo PathwaysJob 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 un v5p de 2 cortes 2x2x2, en el 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 reenvío 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 el reenvío 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. 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())'

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 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 Cloud DNS de kube-dns a Cloud DNS. Cuando se habilita, 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 permiso del clúster, el permiso adicional de VPC 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 de encabezado 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 conectarte a tu clúster de Pathways con Cloud DNS:

1. Confirma que la entrada principal de Cloud DNS 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 rutas, 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 del 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 no se habilita la subdivisión del ILB, 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 puede escalar más allá de los 250 nodos. Con la subdivisión del 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. Habilitar la subdivisión de ILB tiene una latencia de configuración única (alrededor de 15 minutos). En el siguiente comando, se muestra cómo habilitar el subconjunto del ILB:

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

Una vez que se habilita la división en subconjuntos del 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 USER con tu ID de usuario Google Cloud y guarda el archivo comopathways-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 rutas sin reenvío de puertos en hosts de 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 a 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 crear un notebook de Jupyter alojado por tu cuenta.

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 GKE Pathways reside en la misma red de nube privada virtual (que es la red predeterminada, a menos que hayas configurado otra). Navega a la consola de Vertex AI Workbench.

Crea una instancia nueva de Workbench (en la pestaña Instancias) con el botón Crear nueva. Asegúrate de que la red sea la misma que la 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 usará 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 alojada por tu cuenta

El siguiente comando muestra cómo crear una instancia de notebook de Jupyter alojada por el usuario 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"

El siguiente código YAML muestra cómo crear una instancia de notebook de Jupyter alojada por el usuario con kubectl. Aplica el siguiente código 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 del 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}

El resultado debería ser el 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 del notebook al clúster de Pathways

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

En el notebook, agrega una celda para instalar pathwaysutils, establecer JAX_PLATFORMS en proxy y establecer 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" y, luego, 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 indica que se detectó el backend de Pathways on Cloud.

La cantidad de dispositivos JAX que se muestran debe coincidir con la cantidad de chips TPU y la cantidad de segmentos 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 segmentos 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 Rutas de aprendizaje

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 multihost con Pathways
• Cargas de trabajo por lotes con Pathways
• Modo interactivo de Rutas de aprendizaje
• Cómo portar cargas de trabajo de JAX a Pathways
• Capacitación resiliente con Pathways
• Rutas de solución de problemas