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, 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. Puedes ejecutar cargas de trabajo interactivas con Pathways en un clúster de Google Kubernetes Engine (GKE) o en un clúster de Cluster Director con Slurm.

Esta página contiene instrucciones para ambas plataformas. En cada sección, selecciona la pestaña GKE o Cluster Director.

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

• Selecciona la pestaña según cómo planeas usar las muestras de esta página:

  gcloud

  Instala Google Cloud CLI y, luego, accede a gcloud CLI con tu identidad federada. Después de acceder, inicializa Google Cloud CLI con el siguiente comando:

  gcloud init

  REST

  Para usar las muestras de la API de REST incluidas en esta página en un entorno de desarrollo local, debes usar las credenciales que proporciones a la gcloud CLI.

  Instala Google Cloud CLI y, luego, accede a gcloud CLI con tu identidad federada.

  Para obtener más información, consulta Autentícate para usar REST en la documentación de autenticación de Google Cloud .

Antes de comenzar

GKE

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 .

Cluster Director

Completa los siguientes requisitos previos:

• Crea un clúster de Cluster Director y configura el entorno de Pathways.
• Instala gcloud CLI.
• Verifica que tengas acceso SSH al nodo de acceso del clúster.
• Instala Python 3 con el módulo venv en el host que ejecuta el cliente de JAX. En Debian y Ubuntu, instala el paquete python3-venv.

Ejecuta Pathways en modo interactivo

GKE

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.

Cluster Director

1. Conéctate al nodo de acceso del clúster:

   gcloud compute ssh CLUSTER_ID-login-001 \
       --project=PROJECT_ID \
       --zone=ZONE \
       --tunnel-through-iap

   Reemplaza lo siguiente:

   • CLUSTER_ID: Es el nombre de tu clúster de Cluster Director.
   • PROJECT_ID: Es el ID del proyecto deGoogle Cloud .
   • ZONE: Es la zona en la que se encuentra tu clúster.

2. Configura las siguientes variables de entorno:

   export SLICE_SHAPE=TOPOLOGY
   export MAX_SLICES=NUM_SLICES
   export MIN_SLICES=NUM_SLICES
   export GCS_BUCKET=gs://BUCKET_NAME

   Reemplaza lo siguiente:

   • TOPOLOGY: Es la topología de TPU de cada porción, por ejemplo, 2x4. Debe coincidir con la topología de tu conjunto de nodos de TPU.
   • NUM_SLICES: Es la cantidad de porciones de TPU. Establece MIN_SLICES y MAX_SLICES en el mismo valor para que el clúster de Pathways esté listo solo cuando se estén ejecutando todas las porciones.
   • BUCKET_NAME: Es un bucket de Cloud Storage existente para los artefactos de Pathways.

   Si tu partición de TPU no es la predeterminada del clúster, también establece PARTITION en su nombre.

3. Inicia el clúster de Pathways en modo sin interfaz gráfica:

   cd ~/pathways
   sbatch start_pathways.sh
   

4. Espera hasta que el registro del servicio, pathways-service-JOB_ID.out (en el que JOB_ID es el ID del trabajo de Slurm), muestre un resultado similar al siguiente:

   Pathways Runtime Service is READY!
    Discovery file    : /home/USER/pathways/cluster.env
    IFRT Proxy Target : grpc://RM_IP:29000
   

   RM_IP es la dirección IP interna del servidor proxy. El archivo ~/pathways/cluster.env contiene la misma dirección.

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.

En Cluster Director, el cliente de JAX debe usar la misma versión de JAX que el servidor de Pathways, 0.11.1.

Conéctate con la redirección de puertos

GKE

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

Cluster Director

Desde cualquier host que pueda establecer una conexión SSH con el nodo de acceso del clúster, como tu estación de trabajo o una VM fuera de la subred del clúster, reenvía el puerto del proxy a través del nodo de acceso:

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 29000:RM_IP:29000

Nota: Si ejecutas gcloud CLI en una VM de Compute Engine con su cuenta de servicio predeterminada, la VM necesita el permiso de acceso cloud-platform; de lo contrario, gcloud compute ssh falla con insufficient authentication scopes. Como alternativa, ejecuta gcloud auth login en la VM para usar tu propia cuenta.

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. El host no tiene una ruta a RM_IP, por lo que el cliente se conecta al extremo local del túnel, 127.0.0.1:29000:

python3 -m venv .venv
source .venv/bin/activate
pip install pathwaysutils jax==0.11.1

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(288,TPU_DEVICE,coords=[0,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 device(289,TPU_DEVICE,coords=[1,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 ...
 device(313,TPU_DEVICE,coords=[1,3,0,0],logical_task=1,slice=1,default_mem=device,mem_spaces=3)]
W1029 22:52:22.897483    2361 grpc_client.cc:91] IFRT proxy server disconnected: CANCELLED: Stream removed (CANCELLED)

La advertencia IFRT proxy server disconnected aparece cuando la secuencia de comandos sale y cierra su conexión. Es lo esperado.

Nota: Cada llamada de JAX pasa por el túnel SSH. Usa la redirección de puertos para verificar la conectividad. Para los notebooks, ejecuta el cliente de JAX en la subred del clúster.

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

GKE

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.

Cluster Director

Los hosts de la subred del clúster, como una VM de Compute Engine, pueden conectarse directamente al servidor proxy. El firewall del director del clúster permite el tráfico de la subred del clúster, pero no de otras subredes de la misma red de VPC:

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

Reemplaza RM_IP por la dirección del proxy de IFRT del registro del servicio.

Nota: El servidor proxy de IFRT no autentica a los clientes. Cualquier host que pueda comunicarse con RM_IP:29000 puede usar el clúster de Pathways.

Notebooks de Jupyter

GKE

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.

Cluster Director

Puedes ejecutar un notebook de Jupyter en una VM de Compute Engine existente en la subred del clúster, o puedes crear un notebook de Jupyter alojado de forma independiente en el nodo de acceso del clúster. En ambos casos, el notebook se conecta al servidor proxy sin redirección de puertos.

Usa una VM existente en la subred del clúster

En la VM, instala JupyterLab en un entorno virtual y, luego, inícialo:

python3 -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
jupyter lab --no-browser --port=8888

Crea una instancia de notebook de Jupyter autoalojada

En el nodo de acceso del clúster, instala JupyterLab y, luego, inícialo. Crea el entorno virtual en tu directorio principal con el Python del clúster. El Python del sistema del nodo de acceso es demasiado antiguo para las versiones actuales de JAX, por lo que pip no puede encontrar jax==0.11.1 allí:

cd ~
/home/pathways/venv/bin/python -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
source ~/pathways/cluster.env
jupyter lab --no-browser --port=8888

Conéctate a la instancia de notebook

GKE

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

Cluster Director

Para un notebook en una VM o en el nodo de acceso, conéctate desde tu máquina local con la redirección de puertos a ese host. Por ejemplo, para el nodo de acceso:

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 8888:localhost:8888

En tu navegador local, ve a http://localhost:8888?token=TOKEN. Reemplaza TOKEN por el token que JupyterLab imprimió cuando se inició.

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. En Cluster Director, instala jax==0.11.1 en la misma celda. En un notebook alojado de forma automática que iniciaste después del source ~/pathways/cluster.env, las variables ya están configuradas.

!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). En Cluster Director, usa 127.0.0.1 si usas la redirección de puertos o la dirección del proxy de IFRT del registro del servicio. 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. En Cluster Director, establece MIN_SLICES y MAX_SLICES 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

GKE

Cluster Toolkit

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

kubectl

kubectl delete job WORKLOAD

Cluster Director

En el nodo de acceso del clúster, cancela el trabajo del servicio de Pathways:

scancel -u ${USER} -n pathways-service

Este comando también cancela los trabajos de los trabajadores de Pathways. Las VMs de TPU son nodos estáticos, por lo que siguen ejecutándose y están disponibles para otros trabajos de Slurm.

Si tienes problemas, consulta Soluciona problemas de Pathways en Cloud.

¿Qué sigue?

• Crea un clúster de GKE con Pathways
• Crea un clúster de Cluster Director 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