Exécuter une charge de travail interactive avec Pathways

Aperçu

Ce produit ou cette fonctionnalité est soumis aux "Conditions des offres pré-DG" de la section "Conditions d'utilisation générales du service" des Conditions spécifiques du service. Les produits et fonctionnalités pré-DG sont disponibles "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 distantes qui s'exécutent dans une VM ne faisant pas partie du cluster GKE hébergeant le cluster Pathways. Contrairement aux charges de travail par lot, la fin de l'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'IFRT 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 transparente 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
• Installer XPK
• Installer les outils Kubernetes
• Installer gcloud CLI
• Activer l'API TPU
• Activer l'API Google Kubernetes Engine
• Assurez-vous que Pathways est activé pour votre Google Cloud projet.

Exécuter Pathways 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 cette valeur sur 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 Google Cloud projet.
   • 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 au code YAML 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 du PathwaysJob.
   • TPU_MACHINE_TYPE : type de machine TPU que vous souhaitez utiliser. Exemples de valeurs compatibles : "ct6e-standard-8t", "ct5p-hightpu-4t".
   • TOPOLOGY : topologie 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 code YAML précédent, vous devez supprimer ce PathwaysJob et en créer un autre PathwaysJob avec le nombre de pools de nœuds mis à jour. Vous devez également redémarrer tous 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 une tranche v5p 2x2x2, où USER correspond à l'ID de l'utilisateur exécutant 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 (à partir de 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 appropriée pour 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 remplace le fournisseur Cloud DNS de kube-dns par Cloud DNS. Une fois activée, une zone Cloud DNS privée est créée dans votre cloud privé virtuel pour les noms Cloud DNS. Pour en savoir plus, consultez la page Utiliser Cloud DNS pour GKE.

Si vous activez Cloud DNS avec le champ d'application du cluster, du VPC additif ou du VPC, les noms Cloud DNS Kubernetes peuvent être résolus à partir de VM non-GKE dans votre cloud privé virtuel. Les noms ont le format <service_name>.<namespace>.svc.<custom_dns_domain>. Le pod principal 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 Pathways, créez un service soutenu par un équilibreur de charge interne. Votre cluster n'a pas besoin d'avoir Cloud DNS activé.

Pour les clusters comportant de nombreuses VM, nous vous recommandons d'activer le sous-paramètre ILB si vous créez des équilibreurs de charge internes. Pour en savoir plus, consultez la page 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 backend pour tous les équilibreurs de charge internes. Cette solution ne peut pas être mise à l'échelle au-delà de 250 nœuds. Lorsque le sous-paramètre 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 exécutant l'un des pods de diffusion du service sont inclus. L'activation du sous-paramètre ILB entraîne une latence de configuration unique (environ 15 minutes). La commande suivante montre comment activer le sous-paramètre ILB :

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

Une fois le sous-paramètre ILB activé, vous pouvez créer un service Kubernetes de type LoadBalancer à l'aide du code YAML suivant. GKE crée 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

Remplacez le USER par votre Google Cloud ID utilisateur 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 Pathways sans transfert 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 affichées dans les exemples de commande et de sortie précédents.

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

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. 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é autrement). Accédez à la console Vertex AI Workbench.

Créez une instance Workbench (à partir de l'onglet Instances) à l'aide du bouton Create new (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 Open Jupyterlab (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 code YAML suivant montre comment créer une instance de notebook Jupyter auto-hébergée à l'aide de kubectl. Appliquez le code YAML suivant après avoir créé un cluster Pathways sans interface graphique. Pour en savoir plus, consultez la section 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 ordinateur local à 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 des journaux du conteneur de notebook Jupyter.

kubectl logs ${MAIN_POD}

La sortie devrait être la suivante :

...
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é 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 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é s'affiche.

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

Étape suivante

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