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 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 que se conecten 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 sin problemas con los recursos de Pathways.

Antes de comenzar

Asegúrate de tener lo siguiente:

• Crear un clúster de GKE con Pathways
• Instalaste Cluster Toolkit
• Herramientas de Kubernetes instaladas
• Instalado gcloud CLI
• Habilitar la API de Google 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: Establece un nombre único para identificar tu carga de trabajo (debe tener 22 caracteres o menos debido a los límites de las etiquetas de DNS).
   • WORKLOAD_NODEPOOL_COUNT: 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 de Google Cloud .
   • REGION: 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 de tu clúster de GKE.
   • BUCKET_NAME: 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} --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
   

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

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:

Cluster Toolkit

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

Cuando habilitas Cloud DNS en tu clúster, el proveedor de DNS en la nube cambia 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 de clúster, VPC adicional o 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 Service 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 de Cloud DNS del líder se pueda resolver desde un host que no sea de GKE:

   Cluster Toolkit

   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:

   Cluster Toolkit

   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 Cloud DNS habilitado.

En el caso de los clústeres con muchas VMs, te recomendamos que habilites la subdivisión de ILB si crearás 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 escala más allá de los 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 entrega del servicio. La habilitación de la creación de subconjuntos de ILB tiene una latencia de configuración única (aproximadamente 15 minutos). En el siguiente comando, se muestra cómo habilitar el parámetro secundario 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 Service 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 ID de usuario de Google Cloud 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 más tarde), 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 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 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 VM 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 la hayas configurado de otra manera). 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 de tu clúster de GKE. Puedes usar la línea de comandos para crear una instancia de Workbench nueva.

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 hasta 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 YAML, se muestra cómo crear una instancia de notebook de Jupyter alojada de forma automática con kubectl. Aplica el siguiente 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 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:

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 | grep jupyter-notebook-USERNAME | awk '{print $1}')
  kubectl port-forward pod/${MAIN_POD} 8888:8888

En tu navegador local, ve 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ías obtener el siguiente resultado:

...
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 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 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 "hello world" 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 indican 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 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.

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 carga de trabajo interactiva de Pathways

Cluster Toolkit

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

kubectl

kubectl delete -f pathways-headless-workload.yaml

¿Qué sigue?

• Crea un clúster de GKE con Pathways
• Ejecuta una carga de trabajo por lotes con Pathways
• Porting JAX workloads to Pathways
• Entrenamiento resiliente con Pathways
• Rutas de aprendizaje para solucionar problemas en la nube