Exécuter une charge de travail interactive avec Pathways

Les charges de travail interactives Pathways sont des charges de travail JAX distantes qui s'exécutent en dehors des composants du contrôleur Pathways du cluster, par exemple sur une VM. Contrairement aux charges de travail par lot, lorsqu'une charge de travail interactive se termine, elle n'arrête pas les composants du cluster Pathways, qui restent disponibles pour les connexions d'autres clients JAX. Ce document utilise un notebook Jupyter comme exemple pour illustrer les charges de travail interactives. Vous pouvez exécuter des charges de travail interactives avec Pathways sur un cluster Google Kubernetes Engine (GKE) ou sur un cluster Cluster Director avec Slurm.

Cette page contient des instructions pour les deux plates-formes. Dans chaque section, sélectionnez l'onglet GKE ou Cluster Director.

À l'aide de l'interface IFRT, les utilisateurs de JAX envoient des commandes à un cluster Pathways. Le code JAX, qu'il soit exécuté à partir d'un terminal, d'un notebook ou de tout environnement compatible avec Python, peut interagir avec les ressources Pathways.

Avant de commencer

• Sélectionnez l'onglet correspondant à la façon dont vous prévoyez d'utiliser les exemples de cette page :

  gcloud

  Installez la Google Cloud CLI, puis connectez-vous à la gcloud CLI avec votre identité fédérée. Après vous être connecté, initialisez la Google Cloud CLI en exécutant la commande suivante :

  gcloud init

  REST

  Pour utiliser les exemples API REST de cette page dans un environnement de développement local, vous devez utiliser les identifiants que vous fournissez à la gcloud CLI.

  Installez la Google Cloud CLI, puis connectez-vous à la gcloud CLI avec votre identité fédérée.

  Pour en savoir plus, consultez la section S'authentifier pour utiliser REST dans la documentation sur l'authentification Google Cloud .

Avant de commencer

GKE

Remplissez les conditions préalables suivantes :

• Créer un cluster GKE avec Pathways
• Installez Cluster Toolkit.
• Installez les outils Kubernetes.
• Installez la Google Cloud CLI.
• Activez l'API Kubernetes Engine.
• Assurez-vous que Pathways est activé pour votre projet Google Cloud .

Cluster Director

Remplissez les conditions préalables suivantes :

• Créez un cluster Cluster Director et configurez l'environnement Pathways.
• Installez la gcloud CLI.
• Vérifiez que vous disposez d'un accès SSH au nœud de connexion du cluster.
• Installez Python 3 avec le module venv sur l'hôte qui exécute le client JAX. Sous Debian et Ubuntu, installez le package python3-venv.

Exécuter Pathways en mode interactif

GKE

1. Définissez les variables d'environnement suivantes :

   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

   Remplacez les éléments suivants :

   • WORKLOAD : nom unique permettant d'identifier votre charge de travail (22 caractères ou moins en raison des limites d'étiquettes DNS).
   • WORKLOAD_NODEPOOL_COUNT : nombre de pools de nœuds utilisés par une charge de travail Pathways.
   • COMPUTE_TYPE : type de machine TPU, par exemple ct6e-standard-4t. Pour en savoir plus sur les types de TPU compatibles avec chaque version de TPU, consultez Versions de TPU.
   • TOPOLOGY : topologie de l'emplacement du TPU, par exemple 2x4.
   • PROJECT_ID : ID de votre projetGoogle Cloud .
   • REGION : région où se trouvent votre bucket et votre cluster.
   • ZONE : zone dans laquelle vous prévoyez d'exécuter votre charge de travail.
   • CLUSTER : nom de votre cluster GKE.
   • BUCKET_NAME : bucket Cloud Storage pour les artefacts Pathways.

2. Si vous n'avez pas encore créé de bucket Cloud Storage pour les artefacts Pathways, créez-en un :

   gcloud storage buckets create gs://${BUCKET_NAME} \
       --project=${PROJECT_ID} \
       --location=${REGION}
   

3. Créez les conteneurs Pathways sur le cluster. Pour exécuter une charge de travail sans interface graphique, utilisez la commande gcluster job submit (pour en savoir plus, consultez le guide d'envoi de jobs 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
   

   La commande gcluster job submit recherche un pool de nœuds de processeur nommé cpu-np ou pathways-np pour exécuter le pod principal Pathways. Si votre pool de nœuds de processeur porte un nom différent (tel que cpu-pathways-np dans le guide de création de cluster), ajoutez l'option --pathways-head-np=NODE_POOL_NAME aux commandes gcluster job submit de cette page. Sinon, la commande échoue avec l'erreur failed to auto-detect a suitable CPU node pool.

À ce stade, votre charge de travail JAX peut se connecter au serveur proxy IFRT.

Cluster Director

1. Connectez-vous au nœud de connexion du cluster :

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

   Remplacez les éléments suivants :

   • CLUSTER_ID : nom de votre cluster Cluster Director.
   • PROJECT_ID : ID de votre projetGoogle Cloud .
   • ZONE : zone dans laquelle se trouve votre cluster.

2. Définissez les variables d'environnement suivantes :

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

   Remplacez les éléments suivants :

   • TOPOLOGY : topologie TPU de chaque tranche, par exemple 2x4. Il doit correspondre à la topologie de votre ensemble de nœuds TPU.
   • NUM_SLICES : nombre de tranches de TPU. Définissez MIN_SLICES et MAX_SLICES sur la même valeur afin que le cluster Pathways ne soit prêt que lorsque toutes les tranches sont en cours d'exécution.
   • BUCKET_NAME : bucket Cloud Storage existant pour les artefacts Pathways.

   Si votre partition TPU n'est pas la partition par défaut du cluster, définissez également PARTITION sur son nom.

3. Démarrez le cluster Pathways en mode sans interface graphique :

   cd ~/pathways
   sbatch start_pathways.sh
   

4. Attendez que le journal de service, pathways-service-JOB_ID.out (où JOB_ID est l'ID de job Slurm), affiche un résultat semblable à celui-ci :

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

   RM_IP est l'adresse IP interne du serveur proxy. Le fichier ~/pathways/cluster.env contient la même adresse.

À ce stade, votre charge de travail JAX peut se connecter au serveur proxy IFRT.

Se connecter au cluster Pathways en mode interactif

Vous pouvez vous connecter au cluster Pathways avec ou sans transfert de port. Utilisez l'une des sections suivantes pour vous connecter au cluster Pathways.

Sur Cluster Director, le client JAX doit utiliser la même version de JAX que le serveur Pathways, 0.11.1.

Se connecter à l'aide du transfert de port

GKE

Depuis n'importe quel hôte ayant accès au plan de contrôle de votre cluster, transférez le port local 29000 vers le pod principal Pathways pour accéder au serveur 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}

Le résultat ressemble à ce qui suit :

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

Sur le même hôte, ouvrez une nouvelle fenêtre de terminal. À l'aide de Python 3.11 ou version ultérieure, créez un environnement virtuel, installez les packages pathwaysutils et jax, définissez les variables d'environnement JAX_PLATFORMS et JAX_BACKEND_TARGET, puis exécutez un script Python qui importe pathwaysutils et 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())'

Le résultat ressemble à ce qui suit :

[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

Depuis n'importe quel hôte pouvant se connecter en SSH au nœud de connexion du cluster, tel que votre poste de travail ou une VM en dehors du sous-réseau du cluster, transférez le port proxy via le nœud de connexion :

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

Remarque : Si vous exécutez gcloud CLI sur une VM Compute Engine avec son compte de service par défaut, la VM doit disposer du niveau d'accès cloud-platform, sinon gcloud compute ssh échoue avec insufficient authentication scopes. Vous pouvez également exécuter gcloud auth login sur la VM pour utiliser votre propre compte.

Sur le même hôte, ouvrez une nouvelle fenêtre de terminal. Définissez les variables d'environnement JAX_PLATFORMS et JAX_BACKEND_TARGET, puis exécutez un script Python qui importe pathwaysutils et jax. L'hôte n'a pas de route vers RM_IP. Le client se connecte donc à l'extrémité locale du tunnel, 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())'

Le résultat ressemble à ce qui suit :

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

L'avertissement IFRT proxy server disconnected s'affiche lorsque le script se ferme et met fin à sa connexion. C'est normal.

Remarque : Chaque appel JAX passe par le tunnel SSH. Utilisez le transfert de port pour vérifier la connectivité. Pour les notebooks, exécutez le client JAX dans le sous-réseau du cluster.

Se connecter à partir d'hôtes du VPC sans utiliser le transfert de port

GKE

Si vous ne souhaitez pas utiliser le transfert de port, vous pouvez vous connecter au cluster Pathways à l'aide de Cloud DNS ou d'un équilibreur de charge interne.

Se connecter à l'aide de Cloud DNS

L'activation de Cloud DNS dans votre cluster fait passer le fournisseur DNS de kube-dns à Cloud DNS. Lorsqu'elle est activée, GKE crée une zone Cloud DNS privée dans votre réseau cloud privé virtuel (VPC) pour les noms DNS. Pour en savoir plus, consultez Utiliser Cloud DNS pour GKE.

Si vous activez Cloud DNS avec un champ d'application à l'échelle du VPC ou un champ d'application additif à l'échelle du VPC, les noms DNS Kubernetes peuvent être résolus à partir de VM non GKE dans votre réseau VPC. Avec le champ d'application du cluster, seuls les nœuds du cluster peuvent résoudre ces noms. Les noms sont au format SERVICE_NAME.NAMESPACE.svc.CUSTOM_DNS_DOMAIN. Le pod principal Pathways dispose d'un service nommé WORKLOAD-pathways-head-0-0.WORKLOAD.NAMESPACE.svc.CUSTOM_DNS_DOMAIN.

Les commandes suivantes montrent comment se connecter à votre cluster Pathways à l'aide de Cloud DNS :

1. Vérifiez que l'entrée DNS du leader peut être résolue à partir d'un hôte non GKE :

   host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN

   Remplacez les éléments suivants :

   • WORKLOAD : nom de votre charge de travail Pathways.
   • CUSTOM_DNS_DOMAIN : domaine DNS personnalisé configuré pour votre cluster.

   Le résultat ressemble à ce qui suit :

   WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN has address 10.0.2.75

2. Connectez-vous au cluster Pathways à l'aide du nom 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())'

   Le résultat ressemble à ce qui suit :

   [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
   

Se connecter à l'aide d'un équilibreur de charge interne

Pour une adresse IP privée de votre réseau VPC qui pointe vers votre déploiement Pathways, créez un service soutenu par un équilibreur de charge interne (ILB). Pour cela, Cloud DNS ne doit pas être activé sur votre cluster.

Pour les clusters comportant de nombreuses VM, activez le sous-paramètre d'équilibreur de charge interne si vous créez des équilibreurs de charge internes. Pour en savoir plus, consultez Activer le sous-paramétrage GKE dans un cluster existant. Lorsque le sous-paramètre ILB n'est pas activé, tous les nœuds du cluster font partie du groupe d'instances backend pour tous les équilibreurs de charge internes, qui ne peuvent pas évoluer au-delà de 250 nœuds. Lorsque le sous-paramètre ILB est activé, GKE crée des groupes de points de terminaison du réseau au lieu de groupes d'instances, et n'inclut que les nœuds qui exécutent l'un des pods de diffusion du service. L'activation du sous-paramètre ILB entraîne une latence de configuration unique d'environ 15 minutes. La commande suivante active le sous-paramétrage ILB :

gcloud container clusters update ${CLUSTER} \
    --project=${PROJECT_ID} \
    --location=CLUSTER_LOCATION \
    --enable-l4-ilb-subsetting

Remplacez CLUSTER_LOCATION par la région ou la zone où se trouve votre cluster.

Après avoir activé le sous-paramètre ILB, créez un service Kubernetes de type LoadBalancer à l'aide du fichier manifeste YAML suivant. GKE crée alors un équilibreur de charge interne dans le réseau VPC de votre cluster :

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

Remplacez WORKLOAD par le nom de votre charge de travail Pathways, puis enregistrez le fichier sous le nom pathways-headless-ilb.yaml.

Appliquez le fichier manifeste :

kubectl apply -f pathways-headless-ilb.yaml

Une fois que GKE a créé l'équilibreur de charge (environ une minute plus tard), vérifiez que la colonne EXTERNAL-IP affiche une adresse IP :

kubectl get services

Le résultat ressemble à ce qui suit :

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

Vous pouvez accéder au déploiement Pathways sans transfert de port depuis des hôtes du même réseau VPC que votre cluster :

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

Remplacez EXTERNAL_IP par la valeur EXTERNAL-IP affichée par kubectl get services (par exemple, 10.0.0.22).

Le résultat ressemble à ce qui suit :

[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

Remarque : Vos adresses IP diffèrent des adresses affichées dans les exemples de commandes et de résultats précédents.

Cluster Director

Les hôtes du sous-réseau du cluster, tels qu'une VM Compute Engine, peuvent se connecter directement au serveur proxy. Le pare-feu du Cluster Director autorise le trafic provenant du sous-réseau du cluster, mais pas des autres sous-réseaux du même réseau 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())'

Remplacez RM_IP par l'adresse du proxy IFRT du journal de service.

Remarque : Le serveur proxy IFRT n'authentifie pas les clients. Tout hôte pouvant atteindre RM_IP:29000 peut utiliser le cluster Pathways.

Notebooks Jupyter

GKE

Vous pouvez créer un notebook Jupyter à l'aide de Vertex AI ou créer un notebook Jupyter auto-hébergé.

Créer une instance Vertex AI Workbench

Après avoir configuré et vérifié votre cluster Pathways, vous pouvez accéder aux VM TPU GKE à partir d'un notebook Jupyter Vertex AI Workbench. Les instructions de configuration suivantes supposent que votre cluster GKE Pathways réside dans le même réseau VPC (qui est le réseau par défaut, sauf si vous l'avez configuré autrement).

1. Accédez à la page Vertex AI Workbench dans la console Google Cloud .

2. Dans l'onglet Instances, cliquez sur Créer et assurez-vous que le réseau correspond à celui de votre cluster GKE.

   Vous pouvez également utiliser la ligne de commande pour créer une instance 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

   Remplacez les éléments suivants :

   • INSTANCE_NAME : nom de votre instance Vertex AI Workbench.
   • PROJECT_ID : ID de votre projetGoogle Cloud .
   • ZONE : zone dans laquelle vous souhaitez créer l'instance.
   • NETWORK : nom du réseau VPC de votre cluster GKE. L'option --network est facultative. Si vous l'omettez, la commande utilise le réseau default.

3. Une fois que Vertex AI a créé l'instance, accédez-y et cliquez sur Ouvrir JupyterLab.

Créer une instance de notebook Jupyter auto-hébergée

La commande suivante montre comment créer une instance de notebook Jupyter auto-hébergée à l'aide de la commande 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"

Le fichier manifeste YAML suivant montre comment créer une instance de notebook Jupyter auto-hébergée à l'aide de kubectl. Appliquez le fichier manifeste YAML suivant après avoir créé un cluster Pathways sans adresse IP. Pour en savoir plus, consultez Exécuter Pathways en mode interactif :

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

Le champ hostNetwork: true permet au client JAX du notebook d'effectuer le handshake ALTS (Application Layer Transport Security) avec le serveur proxy Pathways. Sans ce champ, la connexion échoue sur les pools de nœuds qui utilisent Workload Identity Federation pour GKE. Le champ dnsPolicy: ClusterFirstWithHostNet permet au notebook de résoudre les noms DNS du cluster, tels que le nom du pod principal Pathways. Pour en savoir plus, consultez Résoudre les problèmes liés à Pathways sur Cloud.

Cluster Director

Vous pouvez exécuter un notebook Jupyter sur une VM Compute Engine existante dans le sous-réseau du cluster, ou créer un notebook Jupyter auto-hébergé sur le nœud de connexion du cluster. Dans les deux cas, le notebook se connecte au serveur proxy sans transfert de port.

Utiliser une VM existante dans le sous-réseau du cluster

Sur la VM, installez JupyterLab dans un environnement virtuel et démarrez-le :

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

Créer une instance de notebook Jupyter auto-hébergée

Sur le nœud de connexion du cluster, installez JupyterLab et démarrez-le. Créez l'environnement virtuel dans votre répertoire d'accueil avec la version Python du cluster. La version de Python du système du nœud de connexion est trop ancienne pour les versions actuelles de JAX. Par conséquent, pip ne peut pas y trouver jax==0.11.1 :

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

Se connecter à l'instance de notebook

GKE

Connectez-vous au notebook depuis votre ordinateur local à l'aide du transfert de port :

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

Dans votre navigateur local, accédez à http://localhost:8888?token=TOKEN. Remplacez TOKEN par le jeton des journaux du conteneur du notebook Jupyter :

kubectl logs ${MAIN_POD}

Le résultat ressemble à ce qui suit :

...
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 est le nom d'hôte du nœud GKE qui exécute le pod du notebook (par exemple, gke-my-cluster-cpu-np-1a2b3c4d-ef56).

Cluster Director

Pour un notebook sur une VM ou sur le nœud de connexion, connectez-vous depuis votre ordinateur local en utilisant le transfert de port vers cet hôte. Par exemple, pour le nœud de connexion :

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

Dans votre navigateur local, accédez à http://localhost:8888?token=TOKEN. Remplacez TOKEN par le jeton affiché par JupyterLab au démarrage.

Connectivité du notebook au cluster Pathways

1. Dans JupyterLab, créez un notebook Python 3.
2. Connectez-vous au serveur proxy Pathways.

Dans le notebook, ajoutez une cellule pour installer pathwaysutils, définissez JAX_PLATFORMS sur proxy et définissez JAX_BACKEND_TARGET sur grpc://PROXY_ADDRESS:29000 :

Remarque : La valeur de PROXY_ADDRESS dépend de l'option de connexion que vous avez choisie dans Se connecter au cluster Pathways en mode interactif. Sur Cluster Director, installez jax==0.11.1 dans la même cellule. Sur un notebook autohébergé que vous avez démarré après le source ~/pathways/cluster.env, les variables sont déjà définies.

!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

Remplacez PROXY_ADDRESS par l'hôte ou l'adresse IP de votre serveur proxy Pathways à partir de Se connecter au cluster Pathways en mode interactif (par exemple, 127.0.0.1, le nom d'hôte Cloud DNS ou l'adresse IP de l'équilibreur de charge interne). Sur Cluster Director, utilisez 127.0.0.1 si vous utilisez le transfert de port, ou l'adresse du proxy IFRT du journal. Si vous avez déployé un notebook auto-hébergé à l'aide de kubectl, remplacez PROXY_ADDRESS par le nom DNS du pod principal Pathways dans le cluster, de sorte que la valeur JAX_BACKEND_TARGET soit grpc://WORKLOAD-pathways-head-0-0.WORKLOAD:29000.

Ajoutez une deuxième cellule pour vérifier la connectivité et afficher les appareils du cluster Pathways :

import pathwaysutils
import jax

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

Si la connexion réussit, la sortie inclut un message indiquant que JAX a détecté le backend Pathways-on-Cloud.

Le nombre d'appareils JAX listés correspond au nombre de puces TPU et au nombre de tranches que vous avez spécifiés lors de la création du cluster Pathways.

Ajouter votre code à un notebook

Ajoutez votre propre code JAX et exécutez-le de manière interactive sur les TPU du cluster Pathways. Le code suivant montre comment effectuer des calculs sur deux tranches à partir d'un seul notebook. Pour exécuter l'exemple complet, créez votre charge de travail Pathways avec au moins deux tranches en définissant WORKLOAD_NODEPOOL_COUNT sur 2 ou plus. Sur Cluster Director, définissez MIN_SLICES et MAX_SLICES sur 2 ou plus. Avec une tranche, slice1_devices est vide et la dernière partie de l'exemple échoue avec une erreur 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)

Supprimer votre charge de travail interactive Pathways

GKE

Cluster Toolkit

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

kubectl

kubectl delete job WORKLOAD

Cluster Director

Sur le nœud de connexion du cluster, annulez le job du service Pathways :

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

Cette commande annule également les jobs de nœud de calcul Pathways. Les VM TPU sont des nœuds statiques. Elles continuent donc de s'exécuter et deviennent disponibles pour d'autres jobs Slurm.

Si vous rencontrez des problèmes, consultez Résoudre les problèmes liés à Pathways sur Cloud.

Étapes suivantes

• Créer un cluster GKE avec Pathways
• Créer un cluster Cluster Director avec Pathways
• Exécuter une charge de travail par lot avec Pathways
• Transférer des charges de travail JAX vers Pathways
• Entraînement résilient avec Pathways
• Résoudre les problèmes liés à Pathways sur Cloud