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 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 Pathways
• Cluster Toolkit instalado
• Herramientas de Kubernetes instaladas
• Instalaste la CLI de gcloud
• 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

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 este parámetro en un nombre único para identificar tu carga de trabajo (debe tener 22 caracteres o menos debido a los límites de 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: Es la topología de ubicación de la TPU (por ejemplo, 2x4).
   • PROJECT_ID: ID del proyecto de Google 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 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 encabezado, usa el comando gcluster job submit (para obtener más información, consulta la Guía de envío de trabajos de 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 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:

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 el reenvío de puertos

Si no quieres usar el reenvío 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 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 de clúster, de VPC adicional o 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 la cabeza 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:

   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 rutas, crea un servicio respaldado por un balanceador de cargas interno. Esto no requiere que tu clúster tenga habilitado Cloud DNS.

En el caso de 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 la subdivisión del 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. Esta opción 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 división en subconjuntos del ILB tiene una latencia de configuración única (alrededor de 15 minutos). El siguiente comando 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 subdivisión 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 de 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 el comando anterior y en los ejemplos de resultado.

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 Instances) con el botón Create new. 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 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"

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 interfaz gráfica. Para obtener más información, consulta Cómo ejecutar 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 el reenvío 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, 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 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 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
• Inferencia multihost con Pathways
• Ejecuta una carga de trabajo por lotes con Pathways
• Cómo portar cargas de trabajo de JAX a Pathways
• Entrenamiento resiliente con Pathways
• Soluciona problemas de Pathways on Cloud