Exécuter une charge de travail interactive avec Pathways

Preview

Ce produit ou cette fonctionnalité sont soumis aux "Conditions des offres de pré-DG" de la section "Conditions générales du service" des Conditions spécifiques du service. Les produits et fonctionnalités de pré-DG sont proposés "en l'état" et peuvent avoir une compatibilité limitée. Pour en savoir plus, consultez les descriptions des étapes de lancement.

Les charges de travail interactives Pathways sont des charges de travail JAX à distance qui s'exécutent dans une VM qui ne fait pas partie du cluster GKE hébergeant le cluster Pathways. Contrairement aux charges de travail par lot, l'exécution d'une opération de charge de travail interactive n'arrête pas les composants du cluster Pathways, qui restent disponibles pour la connexion par d'autres clients JAX. Ce document utilise un notebook Jupyter comme exemple pour illustrer les charges de travail interactives.

À l'aide de l'interface IFRT, les utilisateurs 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 de manière fluide avec les ressources Pathways.

Avant de commencer

Vérifiez que vous disposez bien des éléments suivants :

• Créer un cluster GKE à l'aide de XPK
• XPK installé
• Outils Kubernetes installés
• gcloud CLI installé
• Activer l'API TPU
• Activer l'API Google Kubernetes Engine
• Assurez-vous que Pathways est activé pour votre projet Google Cloud .

Exécuter des parcours en mode interactif

Vous pouvez exécuter Pathways en mode interactif à l'aide de xpk ou kubectl.

XPK

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

   export WORKLOAD=WORKLOAD
   export WORKLOAD_NODEPOOL_COUNT=WORKLOAD_NODEPOOL_COUNT
   export TPU_TYPE=TPU_TYPE
   export PROJECT_ID=PROJECT
   export ZONE=ZONE \
   export CLUSTER=CLUSTER

   Remplacez les éléments suivants :

   • WORKLOAD : définissez un nom unique pour identifier votre charge de travail.
   • WORKLOAD_NODEPOOL_COUNT : nombre de pools de nœuds utilisés par une charge de travail Pathways
   • TPU_TYPE : le type de TPU spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types de TPU compatibles avec chaque version de TPU, consultez Versions de TPU.
   • PROJECT : ID de votre projet Google Cloud
   • ZONE : zone dans laquelle vous prévoyez d'exécuter votre charge de travail.
   • CLUSTER : nom de votre cluster GKE

2. Créez les conteneurs Pathways sur le cluster. Pour exécuter une charge de travail sans interface graphique, exécutez la commande suivante :

   xpk workload create-pathways \
   --headless \
   --workload=${WORKLOAD} \
   --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
   --tpu-type=${TPU_TYPE} \
   --project=${PROJECT} \
   --zone=${ZONE} \
   --cluster=${CLUSTER}

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

kubectl

Le code YAML suivant est identique à celui de la charge de travail par lot, sauf qu'il ne spécifie pas le conteneur main.

1. Remplacez les espaces réservés, copiez le code YAML suivant et collez-le dans un fichier nommé pathways-headless-workload.yaml.

   apiVersion: pathways-job.pathways.domain/v1
   kind: PathwaysJob
   metadata:
     name: pathways-USERNAME
   spec:
     maxRestarts: MAX_RESTARTS
     workers:
       - type: TPU_MACHINE_TYPE
         topology: TOPOLOGY
         numSlices: WORKLOAD_NODEPOOL_COUNT
     pathwaysDir: gs://BUCKET_NAME
     controller:
       deploymentMode: default
       

   Remplacez les éléments suivants :
   • USERNAME : votre nom d'utilisateur
   • MAX_RESTARTS : nombre maximal de redémarrages de PathwaysJob
   • TPU_MACHINE_TYPE : type de machine TPU que vous souhaitez utiliser. Exemples de valeurs acceptées : "ct6e-standard-8t", "ct5p-hightpu-4t"
   • TOPOLOGY : topologie de TPU
   • WORKLOAD_NODEPOOL_COUNT : nombre de pools de nœuds utilisés par une charge de travail Pathways
   • BUCKET_NAME : bucket Cloud Storage utilisé pour stocker les fichiers temporaires
   Pour modifier le nombre de pools de nœuds (répliques de pathways-worker) spécifié par WORKLOAD_NODEPOOL_COUNT dans le fichier YAML précédent, vous devez supprimer ce PathwaysJob et créer un PathwaysJob avec le nombre de pools de nœuds mis à jour. Vous devez également redémarrer les notebooks connectés pour établir la connexion avec le nouveau cluster Pathways.
2. Appliquez le fichier pathways-headless-workload.yaml :

     kubectl apply -f pathways-headless-workload.yaml
     

3. Exécutez kubectl get pods pour vérifier que tous les conteneurs du pod sont en cours d'exécution. La sortie suivante concerne un v5p 2x2x2 à deux tranches, où USER est l'ID de l'utilisateur qui exécute la commande :

       NAME                                         READY   STATUS    RESTARTS   AGE
       pathways-USER-pathways-head-0-0-n848j      2/2     Running   0          49s
       pathways-USER-pathways-workers-0-0-jxt2z   1/1     Running   0          71s
       pathways-USER-pathways-workers-0-1-cxmhc   1/1     Running   0          70s
       pathways-USER-pathways-workers-1-0-5kmz9   1/1     Running   0          71s
       pathways-USER-pathways-workers-1-1-vg5n4   1/1     Running   0          71s
       

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.

Se connecter à l'aide du transfert de port

À ce stade, vous pouvez utiliser le transfert de port (depuis n'importe quel hôte ayant accès au plan de contrôle de votre cluster) pour accéder au serveur proxy :

Utilisez la commande adaptée à votre charge de travail :

XPK

PROXY_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
PROXY_PORT=29000
kubectl port-forward ${PROXY_POD} ${PROXY_PORT}:${PROXY_PORT}

Un résultat semblable à celui-ci s'affiche :

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}

Un résultat semblable à celui-ci s'affiche :

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

Vous devriez voir une sortie semblable à 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

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

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 permet de remplacer le fournisseur kube-dns par Cloud DNS. Lorsqu'elle est activée, une zone Cloud DNS privée est créée dans votre Virtual Private Cloud pour les noms Cloud DNS. Pour en savoir plus, consultez Utiliser Cloud DNS pour GKE.

Si vous activez Cloud DNS avec le champ d'application du cluster, additif à l'échelle du VPC ou à l'échelle du VPC, les noms Kubernetes Cloud DNS peuvent être résolus à partir de VM non GKE dans votre cloud privé virtuel. Les noms sont au format <service_name>.<namespace>.svc.<custom_dns_domain>. Le pod de tête Pathways possède un service nommé <jobset_name>-pathways-head-0-0.<jobset_name>.<namespace>.svc.<custom_dns_domain>.

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

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

   XPK

   host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.USERNAME-test

   Un résultat semblable à celui-ci s'affiche :

   <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

   Un résultat semblable à celui-ci s'affiche :

   pathways-<user>-pathways-head-0-0.pathways-<user>.default.svc.<user>-test has address 10.0.2.75

2. Connectez-vous au cluster Pathways à l'aide du nom Cloud DNS :

   XPK

   JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.USERNAME-test:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

   kubectl

   JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://pathways-USERNAME-pathways-head-0-0.pathways-USERNAME.default.svc.USERNAME-test:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

   Un résultat semblable à celui-ci s'affiche :

   [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 dans votre VPC qui pointe vers votre déploiement de parcours, créez un service soutenu par un équilibreur de charge interne. Cloud DNS ne doit pas nécessairement être activé sur votre cluster.

Pour les clusters comportant de nombreuses VM, nous vous recommandons d'activer 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ètre 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 de backend pour tous les équilibreurs de charge internes. Cette méthode ne s'adapte pas au-delà de 250 nœuds. Lorsque le sous-ensemble ILB est activé, GKE crée des groupes de points de terminaison réseau au lieu de groupes d'instances, et seuls les nœuds qui exécutent l'un des pods de diffusion du service sont inclus. L'activation du sous-ensemble ILB entraîne une latence de configuration unique (environ 15 minutes). La commande suivante montre comment activer le sous-ensemble d'ILB :

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

Une fois le sous-ensemble ILB activé, vous pouvez créer un service Kubernetes de type LoadBalancer à l'aide du fichier YAML suivant. GKE créera alors un équilibreur de charge interne dans le VPC de votre cluster :

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

Mettez à jour USER avec votre ID utilisateur Google Cloud et enregistrez le fichier sous le nom pathways-headless-ilb.yaml.

Appliquez le fichier manifeste :

kubectl apply -f pathways-headless-ilb.yaml

Une fois l'équilibreur de charge créé (environ une minute plus tard), la colonne EXTERNAL-IP aura une valeur :

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

Vous pouvez accéder au déploiement des pathways sans redirection de port sur les hôtes du même VPC que votre cluster :

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())'

Remarque : Remplacez l'adresse IP du backend JAX (10.0.0.22) dans la commande précédente par la valeur affichée par kubectl get services.

Un résultat semblable à celui-ci s'affiche :

[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 seront différentes de celles qui apparaissent dans la commande précédente et les exemples de résultats.

Notebooks Jupyter

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

Créer une instance Vertex AI Workbench

Une fois votre cluster Pathways configuré et validé, vous pouvez accéder aux VM TPU GKE depuis un notebook Jupyter Vertex AI. Les instructions de configuration suivantes supposent que votre cluster GKE Pathways réside dans le même réseau de cloud privé virtuel (qui est le réseau par défaut, sauf si vous l'avez configuré différemment). Accédez à la console Vertex AI Workbench.

Créez une instance Workbench (à partir de l'onglet Instances) à l'aide du bouton Créer. Assurez-vous que le réseau est le même que celui de votre cluster GKE. Vous pouvez utiliser la ligne de commande pour créer une instance Workbench.

gcloud workbench instances create INSTANCE_NAME \
--machine-type=e2-standard-4 \
--data-disk-size=100 \
--location=ZONE \
[--network=NETWORK]

Remarque : L'option --network est facultative. Si vous l'omettez, le réseau par défaut est utilisé.

Une fois l'instance créée, 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 XPK :

xpk workload create-pathways \
--workload=${WORKLOAD} \
--num-slices=${WORKLOAD_NODEPOOL_COUNT} \
--tpu-type=${TPU_TYPE} \
--project=${PROJECT} \
--zone=${ZONE} \
--cluster=${CLUSTER} \
--docker-image=jupyter/base-notebook \
--command "start-notebook.sh"

Le fichier YAML suivant montre comment créer une instance de notebook Jupyter auto-hébergée à l'aide de kubectl. Appliquez le fichier YAML suivant après la création d'un cluster Pathways sans interface graphique. Pour en savoir plus, consultez Exécuter Pathways en mode interactif avec 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

Accéder à l'instance de notebook

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

XPK

  MAIN_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
  kubectl port-forward pod/${MAIN_POD} 8888:8888

kubectl

  MAIN_POD=$(kubectl get pods | grep jupyter-notebook-USERNAME | awk '{print $1}')
  kubectl port-forward pod/${MAIN_POD} 8888:8888

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

kubectl logs ${MAIN_POD}

Le résultat devrait être le suivant :

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

Connectivité des notebooks au cluster Pathways

1. Dans JupyterLab, créez un notebook Python 3.
2. Se connecter au serveur proxy Pathways

Dans le notebook, ajoutez une cellule pour installer pathwaysutils, définissez JAX_PLATFORMS sur proxy et JAX_BACKEND_TARGET sur PROXY_ADDRESS.

Remarque : La valeur de PROXY_ADDRESS dépend de l'option de connexion que vous avez choisie dans Se connecter au serveur proxy Pathways.

!pip install pathwaysutils
%env JAX_PLATFORMS=proxy
# Replace your proxy address below:
%env JAX_BACKEND_TARGET=PROXY_ADDRESS

Ajoutez une deuxième cellule en tant que vérification de type "hello world" et imprimez les appareils du cluster Pathways.

import pathwaysutils
import jax

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

Si tout fonctionne correctement, un message indiquant que le backend Pathways-on-Cloud a été détecté devrait s'afficher.

Le nombre d'appareils JAX listés doit correspondre 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 même 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)

Supprimer votre cluster interactif Pathways

XPK

xpk workload delete --workload=WORKLOAD --cluster=CLUSTER --project=PROJECT --zone=ZONE

kubectl

kubectl delete -f pathways-headless-workload.yaml

Étapes suivantes

• Créer un cluster GKE avec Pathways
• Inférence multihôte avec Pathways
• Charges de travail par lot avec Pathways
• Mode interactif des parcours
• Transférer des charges de travail JAX vers Pathways
• Entraînement résilient avec Pathways
• Parcours de dépannage