Interaktive Arbeitslast mit Pathways ausführen

Interaktive Pathways-Arbeitslasten sind Remote-JAX-Arbeitslasten, die außerhalb der Pathways-Controller-Komponenten des Clusters ausgeführt werden, z. B. auf einer VM. Im Gegensatz zu Batch-Arbeitslasten werden die Pathways-Clusterkomponenten nach Abschluss einer interaktiven Arbeitslast nicht heruntergefahren. Sie bleiben für Verbindungen von anderen JAX-Clients verfügbar. In diesem Dokument wird ein Jupyter-Notebook als Beispiel für interaktive Arbeitslasten verwendet. Sie können interaktive Arbeitslasten mit Pathways in einem Google Kubernetes Engine-Cluster (GKE) oder in einem Cluster Director-Cluster mit Slurm ausführen.

Diese Seite enthält Anleitungen für beide Plattformen. Wählen Sie in jedem Abschnitt den Tab GKE oder Cluster Director aus.

Über die IFRT-Schnittstelle senden JAX-Nutzer Befehle an einen Pathways-Cluster. JAX-Code, der über ein Terminal, ein Notebook oder eine beliebige Python-kompatible Umgebung ausgeführt wird, kann mit Pathways-Ressourcen interagieren.

Hinweis

• Wählen Sie den Tab für die gewünschte Verwendung der Beispiele auf dieser Seite aus:

  gcloud

  Installieren Sie die Google Cloud CLI und melden Sie sich dann mit Ihrer föderierten Identität in der gcloud CLI an. Nach der Anmeldung initialisieren Sie die Google Cloud CLI mit dem folgenden Befehl:

  gcloud init

  REST

  Wenn Sie die REST API-Beispiele auf dieser Seite in einer lokalen Entwicklungsumgebung verwenden möchten, verwenden Sie die Anmeldedaten, die Sie der gcloud CLI bereitstellen.

  Installieren Sie die Google Cloud CLI und melden Sie sich dann mit Ihrer föderierten Identität in der gcloud CLI an.

  Weitere Informationen finden Sie in der Dokumentation zur Google Cloud -Authentifizierung unter Für die Verwendung von REST authentifizieren.

Hinweis

GKE

Erfüllen Sie die folgenden Voraussetzungen:

• GKE-Cluster mit Pathways erstellen
• Cluster Toolkit installieren
• Kubernetes-Tools installieren
• Installieren Sie die Google Cloud CLI.
• Aktivieren Sie die Kubernetes Engine API.
• Achten Sie darauf, dass Pathways für Ihr Google Cloud -Projekt aktiviert ist.

Cluster Director

Erfüllen Sie die folgenden Voraussetzungen:

• Cluster Director-Cluster erstellen und Pathways-Umgebung einrichten
• Installieren Sie die gcloud CLI.
• Prüfen Sie, ob Sie SSH-Zugriff auf den Login-Knoten des Clusters haben.
• Installieren Sie Python 3 mit dem Modul venv auf dem Host, auf dem der JAX-Client ausgeführt wird. Installieren Sie unter Debian und Ubuntu das Paket python3-venv.

Pathways im interaktiven Modus ausführen

GKE

1. Legen Sie die folgenden Umgebungsvariablen fest:

   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

   Ersetzen Sie Folgendes:

   • WORKLOAD: Ein eindeutiger Name zur Identifizierung Ihrer Arbeitslast (maximal 22 Zeichen aufgrund von DNS-Label-Limits).
   • WORKLOAD_NODEPOOL_COUNT: Die Anzahl der Knotenpools, die von einer Pathways-Arbeitslast verwendet werden.
   • COMPUTE_TYPE: der TPU-Maschinentyp, z. B. ct6e-standard-4t. Weitere Informationen zu den unterstützten TPU-Typen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.
   • TOPOLOGY: die TPU-Platzierungstopologie, z. B. 2x4.
   • PROJECT_ID: IhreGoogle Cloud -Projekt-ID.
   • REGION: die Region, in der sich Ihr Bucket und Ihr Cluster befinden.
   • ZONE: die Zone, in der Sie Ihre Arbeitslast ausführen möchten.
   • CLUSTER: der Name Ihres GKE-Cluster.
   • BUCKET_NAME: der Cloud Storage-Bucket für Pathways-Artefakte.

2. Wenn Sie noch keinen Cloud Storage-Bucket für Pathways-Artefakte erstellt haben, erstellen Sie einen:

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

3. Erstellen Sie die Pathways-Container im Cluster. Verwenden Sie den Befehl gcluster job submit, um eine Headless-Arbeitslast auszuführen. Weitere Informationen finden Sie im Leitfaden zum Einreichen von Jobs im 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
   

   Der Befehl gcluster job submit sucht nach einem CPU-Knotenpool mit dem Namen cpu-np oder pathways-np, um den Pathways-Head-Pod auszuführen. Wenn Ihr CPU-Knotenpool einen anderen Namen hat (z. B. cpu-pathways-np aus der Anleitung zum Erstellen von Clustern), fügen Sie das Flag --pathways-head-np=NODE_POOL_NAME zu den gcluster job submit-Befehlen auf dieser Seite hinzu. Andernfalls schlägt der Befehl mit dem Fehler failed to auto-detect a suitable CPU node pool fehl.

An diesem Punkt kann Ihre JAX-Arbeitslast eine Verbindung zum IFRT-Proxyserver herstellen.

Cluster Director

1. Stellen Sie eine Verbindung zum Login-Knoten des Clusters her:

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

   Ersetzen Sie Folgendes:

   • CLUSTER_ID: der Name Ihres Cluster Director-Clusters.
   • PROJECT_ID: IhreGoogle Cloud -Projekt-ID.
   • ZONE: Die Zone, in der sich der Cluster befindet.

2. Legen Sie die folgenden Umgebungsvariablen fest:

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

   Ersetzen Sie Folgendes:

   • TOPOLOGY: die TPU-Topologie jedes Slices, z. B. 2x4. Sie muss mit der Topologie Ihres TPU-Knotensatzes übereinstimmen.
   • NUM_SLICES: die Anzahl der TPU-Slices. Legen Sie für MIN_SLICES und MAX_SLICES denselben Wert fest, damit der Pathways-Cluster erst bereit ist, wenn alle Slices ausgeführt werden.
   • BUCKET_NAME: ein vorhandener Cloud Storage-Bucket für Pathways-Artefakte.

   Wenn Ihre TPU-Partition nicht die Standardpartition des Clusters ist, legen Sie für PARTITION auch den Namen fest.

3. Starten Sie den Pathways-Cluster im Headless-Modus:

   cd ~/pathways
   sbatch start_pathways.sh
   

4. Warten Sie, bis das Dienstprotokoll pathways-service-JOB_ID.out (wobei JOB_ID die Slurm-Job-ID ist) eine Ausgabe wie die folgende anzeigt:

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

   RM_IP ist die interne IP-Adresse des Proxyservers. Die Datei ~/pathways/cluster.env enthält dieselbe Adresse.

An diesem Punkt kann Ihre JAX-Arbeitslast eine Verbindung zum IFRT-Proxyserver herstellen.

Im interaktiven Modus eine Verbindung zum Pathways-Cluster herstellen

Sie können mit oder ohne Portweiterleitung eine Verbindung zum Pathways-Cluster herstellen. Verwenden Sie einen der folgenden Abschnitte, um eine Verbindung zum Pathways-Cluster herzustellen.

Auf Cluster Director muss der JAX-Client dieselbe JAX-Version wie der Pathways-Server verwenden, 0.11.1.

Verbindung über Portweiterleitung herstellen

GKE

Leiten Sie von einem beliebigen Host mit Zugriff auf die Steuerungsebene Ihres Clusters den lokalen Port 29000 an den Pathways-Head-Pod weiter, um auf den Proxyserver zuzugreifen:

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

Die Ausgabe sieht etwa so aus:

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

Öffnen Sie auf demselben Host ein neues Terminalfenster. Erstellen Sie mit Python 3.11 oder höher eine virtuelle Umgebung, installieren Sie die Pakete pathwaysutils und jax, legen Sie die Umgebungsvariablen JAX_PLATFORMS und JAX_BACKEND_TARGET fest und führen Sie ein Python-Skript aus, das pathwaysutils und jax importiert:

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

Die Ausgabe sieht etwa so aus:

[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

Leiten Sie von einem beliebigen Host, der eine SSH-Verbindung zum Login-Knoten des Clusters herstellen kann, z. B. von Ihrer Workstation oder einer VM außerhalb des Subnetzes des Clusters, den Proxy-Port über den Login-Knoten weiter:

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

Hinweis: Wenn Sie die gcloud CLI auf einer Compute Engine-VM mit ihrem Standarddienstkonto ausführen, benötigt die VM den Zugriffsbereich cloud-platform. Andernfalls schlägt gcloud compute ssh mit insufficient authentication scopes fehl. Alternativ können Sie gcloud auth login auf der VM ausführen, um Ihr eigenes Konto zu verwenden.

Öffnen Sie auf demselben Host ein neues Terminalfenster. Legen Sie die Umgebungsvariablen JAX_PLATFORMS und JAX_BACKEND_TARGET fest und führen Sie ein Python-Skript aus, das pathwaysutils und jax importiert. Der Host hat keine Route zu RM_IP, daher stellt der Client eine Verbindung zum lokalen Ende des Tunnels (127.0.0.1:29000) her:

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

Die Ausgabe sieht etwa so aus:

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

Die Warnung IFRT proxy server disconnected wird angezeigt, wenn das Script beendet und die Verbindung geschlossen wird. Das ist zu erwarten.

Hinweis: Jeder JAX-Aufruf wird über den SSH-Tunnel geleitet. Verwenden Sie die Portweiterleitung, um die Verbindung zu prüfen. Führen Sie für Notebooks den JAX-Client im Subnetz des Clusters aus.

Verbindung von Hosts in der VPC ohne Portweiterleitung herstellen

GKE

Wenn Sie keine Portweiterleitung verwenden möchten, können Sie über Cloud DNS oder einen internen Load Balancer eine Verbindung zum Pathways-Cluster herstellen.

Verbindung über Cloud DNS herstellen

Wenn Sie Cloud DNS in Ihrem Cluster aktivieren, wird der DNS-Anbieter von kube-dns zu Cloud DNS gewechselt. Wenn diese Option aktiviert ist, erstellt die GKE eine private Cloud DNS-Zone in Ihrem VPC-Netzwerk (Virtual Private Cloud) für die DNS-Namen. Weitere Informationen finden Sie unter Cloud DNS für GKE verwenden.

Wenn Sie Cloud DNS mit VPC-Bereich oder additivem VPC-Bereich aktivieren, können die Kubernetes-DNS-Namen von Nicht-GKE-VMs in Ihrem VPC-Netzwerk aufgelöst werden. Mit dem Clusterbereich können nur Knoten im Cluster diese Namen auflösen. Die Namen haben das Format SERVICE_NAME.NAMESPACE.svc.CUSTOM_DNS_DOMAIN. Der Pathways-Head-Pod hat einen Dienst namens WORKLOAD-pathways-head-0-0.WORKLOAD.NAMESPACE.svc.CUSTOM_DNS_DOMAIN.

Die folgenden Befehle zeigen, wie Sie über Cloud DNS eine Verbindung zu Ihrem Pathways-Cluster herstellen:

1. Prüfen Sie, ob der DNS-Eintrag des Leaders von einem Nicht-GKE-Host aufgelöst werden kann:

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

   Ersetzen Sie Folgendes:

   • WORKLOAD: der Name Ihrer Pathways-Arbeitslast.
   • CUSTOM_DNS_DOMAIN: die benutzerdefinierte DNS-Domain, die für Ihren Cluster konfiguriert ist.

   Die Ausgabe sieht etwa so aus:

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

2. Stellen Sie über den DNS-Namen eine Verbindung zum Pathways-Cluster her:

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

   Die Ausgabe sieht etwa so aus:

   [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
   

Verbindung über einen internen Load Balancer herstellen

Erstellen Sie für eine private IP-Adresse in Ihrem VPC-Netzwerk, die auf Ihre Pathways-Bereitstellung verweist, einen Dienst, der von einem internen Load Balancer (ILB) unterstützt wird. Dafür muss Cloud DNS nicht für Ihren Cluster aktiviert sein.

Aktivieren Sie für Cluster mit vielen VMs die ILB-Teilmengeneinstellung, wenn Sie interne Load Balancer erstellen. Weitere Informationen finden Sie unter GKE-Teilmengen in einem vorhandenen Cluster aktivieren. Wenn die ILB-Teilmengeneinstellung nicht aktiviert ist, sind alle Knoten im Cluster Teil der Backend-Instanzgruppe für alle internen Load Balancer, die nicht über 250 Knoten hinaus skaliert werden. Wenn die ILB-Teileinstellung aktiviert ist, erstellt GKE Netzwerk-Endpunktgruppen anstelle von Instanzgruppen und schließt nur Knoten ein, auf denen einer der Bereitstellungs-Pods des Dienstes ausgeführt wird. Die Aktivierung der ILB-Teilmengeneinstellung hat eine einmalige Einrichtungslatenz von etwa 15 Minuten. Mit dem folgenden Befehl wird die ILB-Teilmengeneinstellung aktiviert:

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

Ersetzen Sie CLUSTER_LOCATION durch die Region oder Zone, in der sich der Cluster befindet.

Nachdem Sie die ILB-Teilmengeneinstellung aktiviert haben, erstellen Sie mit dem folgenden YAML-Manifest einen Kubernetes-Dienst vom Typ LoadBalancer. Dadurch erstellt GKE einen internen Load Balancer im VPC-Netzwerk Ihres Clusters:

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

Ersetzen Sie WORKLOAD durch den Namen Ihrer Pathways-Arbeitslast und speichern Sie die Datei als pathways-headless-ilb.yaml.

Wenden Sie das Manifest an:

kubectl apply -f pathways-headless-ilb.yaml

Nachdem GKE den Load Balancer erstellt hat (etwa eine Minute später), prüfen Sie, ob in der Spalte EXTERNAL-IP eine IP-Adresse angezeigt wird:

kubectl get services

Die Ausgabe sieht etwa so aus:

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

Sie können von Hosts im selben VPC-Netzwerk wie Ihr Cluster ohne Portweiterleitung auf die Pathways-Bereitstellung zugreifen:

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

Ersetzen Sie EXTERNAL_IP durch den EXTERNAL-IP-Wert, der von kubectl get services angezeigt wird, z. B. 10.0.0.22.

Die Ausgabe sieht etwa so aus:

[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

Hinweis: Ihre IP-Adressen unterscheiden sich von den Adressen, die in den vorherigen Befehls- und Ausgabebeispielen gezeigt wurden.

Cluster Director

Hosts im Subnetz des Clusters, z. B. eine Compute Engine-VM, können direkt eine Verbindung zum Proxyserver herstellen. Die Firewall des Cluster Directors lässt Traffic aus dem Subnetz des Clusters zu, aber nicht aus anderen Subnetzen desselben VPC-Netzwerk:

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

Ersetzen Sie RM_IP durch die IFRT-Proxy-Adresse aus dem Dienstprotokoll.

Hinweis: Der IFRT-Proxyserver authentifiziert keine Clients. Jeder Host, der RM_IP:29000 erreichen kann, kann den Pathways-Cluster verwenden.

Jupyter-Notebooks

GKE

Sie können ein Jupyter-Notebook mit Vertex AI oder ein selbst gehostetes Jupyter-Notebook erstellen.

Vertex AI Workbench-Instanz erstellen

Nachdem Sie Ihren Pathways-Cluster eingerichtet und überprüft haben, können Sie über ein Vertex AI Workbench Jupyter-Notebook auf die GKE-TPU-VMs zugreifen. In der folgenden Einrichtungsanleitung wird davon ausgegangen, dass sich Ihr GKE-Pathways-Cluster im selben VPC-Netzwerk befindet (das ist das Standardnetzwerk, sofern Sie nichts anderes konfiguriert haben).

1. Rufen Sie in der Google Cloud Console die Seite Vertex AI Workbench auf.

2. Klicken Sie auf dem Tab Instanzen auf Neu erstellen und achten Sie darauf, dass das Netzwerk mit dem Netzwerk Ihres GKE-Cluster übereinstimmt.

   Alternativ können Sie die Befehlszeile verwenden, um eine neue Vertex AI Workbench-Instanz zu erstellen:

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

   Ersetzen Sie Folgendes:

   • INSTANCE_NAME: Der Name Ihrer Vertex AI Workbench-Instanz.
   • PROJECT_ID: IhreGoogle Cloud -Projekt-ID.
   • ZONE: die Zone, in der Sie die Instanz erstellen möchten.
   • NETWORK: der Name des VPC-Netzwerk Ihres GKE-Cluster. Das Flag --network ist optional. Wenn Sie es weglassen, verwendet der Befehl das Netzwerk default.

3. Nachdem Vertex AI die Instanz erstellt hat, rufen Sie die Instanz auf und klicken Sie auf JupyterLab öffnen.

Selbst gehostete Jupyter-Notebook-Instanz erstellen

Der folgende Befehl zeigt, wie Sie mit dem Befehl gcluster job submit eine selbst gehostete Jupyter-Notebook-Instanz erstellen:

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"

Das folgende YAML-Manifest zeigt, wie Sie mit kubectl eine selbst gehostete Jupyter-Notebook-Instanz erstellen. Wenden Sie das folgende YAML-Manifest an, nachdem Sie einen Pathways-Cluster ohne Head erstellt haben. Weitere Informationen finden Sie unter Pathways im interaktiven Modus ausführen:

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

Über das Feld hostNetwork: true kann der JAX-Client im Notebook den ALTS-Handshake (Application Layer Transport Security) mit dem Pathways-Proxyserver abschließen. Ohne dieses Feld schlägt die Verbindung in Knotenpools fehl, die Workload Identity Federation for GKE verwenden. Mit dem Feld dnsPolicy: ClusterFirstWithHostNet kann das Notebook clusterinterne DNS-Namen auflösen, z. B. den Namen des Pathways-Head-Pods. Weitere Informationen finden Sie unter Fehlerbehebung für Pathways in der Cloud.

Cluster Director

Sie können ein Jupyter-Notebook auf einer vorhandenen Compute Engine-VM im Subnetz des Clusters ausführen oder ein selbst gehostetes Jupyter-Notebook auf dem Login-Knoten des Clusters erstellen. In beiden Fällen stellt das Notebook ohne Portweiterleitung eine Verbindung zum Proxyserver her.

Vorhandene VM im Subnetz des Clusters verwenden

Installieren Sie JupyterLab auf der VM in einer virtuellen Umgebung und starten Sie es:

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

Selbst gehostete Jupyter-Notebook-Instanz erstellen

Installieren Sie JupyterLab auf dem Login-Knoten des Clusters und starten Sie es. Erstellen Sie die virtuelle Umgebung in Ihrem Basisverzeichnis mit dem Python des Clusters. Die Python-Version des Login-Knotens ist zu alt für aktuelle JAX-Releases, daher kann pip dort jax==0.11.1 nicht finden:

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

Verbindung zur Notebook-Instanz herstellen

GKE

Stellen Sie über die Portweiterleitung eine Verbindung vom lokalen Computer zum Notebook her:

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

Rufen Sie in Ihrem lokalen Browser http://localhost:8888?token=TOKEN auf. Ersetzen Sie TOKEN durch das Token aus den Logs des Jupyter-Notebook-Containers:

kubectl logs ${MAIN_POD}

Die Ausgabe sieht etwa so aus:

...
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 ist der Hostname des GKE-Knotens, auf dem der Notebook-Pod ausgeführt wird (z. B. gke-my-cluster-cpu-np-1a2b3c4d-ef56).

Cluster Director

Wenn Sie ein Notebook auf einer VM oder auf dem Login-Knoten verwenden, stellen Sie von Ihrem lokalen Computer aus eine Verbindung her, indem Sie die Portweiterleitung zu diesem Host verwenden. Beispiel für den Login-Knoten:

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

Rufen Sie in Ihrem lokalen Browser http://localhost:8888?token=TOKEN auf. Ersetzen Sie TOKEN durch das Token, das JupyterLab beim Start ausgegeben hat.

Notebook-Verbindung zum Pathways-Cluster

1. Erstellen Sie in JupyterLab ein neues Python 3-Notebook.
2. Verbindung zum Pathways-Proxyserver herstellen

Fügen Sie im Notebook eine Zelle hinzu, um pathwaysutils zu installieren, JAX_PLATFORMS auf proxy und JAX_BACKEND_TARGET auf grpc://PROXY_ADDRESS:29000 festzulegen:

Hinweis: Der Wert für PROXY_ADDRESS hängt von der Verbindungsoption ab, die Sie unter Im interaktiven Modus mit dem Pathways-Cluster verbinden ausgewählt haben. Installieren Sie jax==0.11.1 in Cluster Director in derselben Zelle. In einem selbst gehosteten Notebook, das Sie nach dem source ~/pathways/cluster.env gestartet haben, sind die Variablen bereits festgelegt.

!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

Ersetzen Sie PROXY_ADDRESS durch den Host oder die IP-Adresse Ihres Pathways-Proxyservers aus Im interaktiven Modus eine Verbindung zum Pathways-Cluster herstellen (z. B. 127.0.0.1, der Cloud DNS-Hostname oder die IP-Adresse des internen Load Balancers). Verwenden Sie auf Cluster Director 127.0.0.1, wenn Sie die Portweiterleitung verwenden, oder die IFRT-Proxy-Adresse aus dem Dienstprotokoll. Wenn Sie ein selbst gehostetes Notebook mit kubectl bereitgestellt haben, ersetzen Sie PROXY_ADDRESS durch den clusterinternen DNS-Namen des Pathways-Head-Pods, sodass der Wert für JAX_BACKEND_TARGET grpc://WORKLOAD-pathways-head-0-0.WORKLOAD:29000 ist.

Fügen Sie eine zweite Zelle hinzu, um die Verbindung zu prüfen und die Geräte im Pathways-Cluster auszugeben:

import pathwaysutils
import jax

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

Wenn die Verbindung erfolgreich ist, enthält die Ausgabe eine Meldung, dass JAX das Pathways-on-Cloud-Backend erkannt hat.

Die Anzahl der aufgeführten JAX-Geräte entspricht der Anzahl der TPU-Chips und der Anzahl der Slices, die Sie beim Erstellen des Pathways-Clusters angegeben haben.

Code zu einem Notebook hinzufügen

Fügen Sie Ihren eigenen JAX-Code hinzu und führen Sie ihn interaktiv auf den TPUs im Pathways-Cluster aus. Der folgende Code zeigt, wie Berechnungen über zwei Slices aus einem einzelnen Notebook ausgeführt werden. Wenn Sie das vollständige Beispiel ausführen möchten, erstellen Sie Ihre Pathways-Arbeitslast mit mindestens zwei Slices, indem Sie WORKLOAD_NODEPOOL_COUNT auf 2 oder mehr festlegen. Legen Sie in Cluster Director für MIN_SLICES und MAX_SLICES mindestens 2 fest. Bei einem Slice ist slice1_devices leer und der letzte Teil des Beispiels schlägt mit einem ValueError fehl:

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)

Interaktive Pathways-Arbeitslast löschen

GKE

Cluster Toolkit

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

kubectl

kubectl delete job WORKLOAD

Cluster Director

Brechen Sie auf dem Login-Knoten des Clusters den Pathways-Dienstjob ab:

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

Mit diesem Befehl werden auch die Pathways-Worker-Jobs abgebrochen. Die TPU-VMs sind statische Knoten, sodass sie weiter ausgeführt werden und für andere Slurm-Jobs verfügbar sind.

Wenn Probleme auftreten, lesen Sie den Abschnitt Fehlerbehebung für Pathways in der Cloud.

Nächste Schritte

• GKE-Cluster mit Pathways erstellen
• Cluster Director-Cluster mit Pathways erstellen
• Batcharbeitslast mit Pathways ausführen
• JAX-Arbeitslasten zu Pathways portieren
• Ausfallsicheres Training mit Pathways
• Fehlerbehebung für Pathways in der Cloud