Interaktive Arbeitslast mit Pathways ausführen

Interaktive Pathways-Arbeitslasten sind Remote-JAX-Arbeitslasten, die in einer VM ausgeführt werden, die nicht Teil des GKE-Cluster ist, in dem der Pathways-Cluster gehostet wird. Im Gegensatz zu Batch-Arbeitslasten werden die Pathways-Clusterkomponenten nach Abschluss des Vorgangs für interaktive Arbeitslasten nicht heruntergefahren, sondern bleiben für die Verbindung durch andere JAX-Clients verfügbar. In diesem Dokument wird ein Jupyter-Notebook als Beispiel verwendet, um interaktive Arbeitslasten zu veranschaulichen.

Ü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 nahtlos mit Pathways-Ressourcen interagieren.

Hinweis

Sie benötigen Folgendes:

• GKE-Cluster mit XPK erstellt
• XPK installiert
• Kubernetes-Tools installiert
• gcloud-CLI installiert
• TPU API aktiviert
• Google Kubernetes Engine API aktiviert
• Pathways für Ihr Google Cloud Projekt aktiviert

Pathways im interaktiven Modus ausführen

Sie können Pathways im interaktiven Modus mit xpk oder kubectl ausführen.

XPK

1. Legen Sie die folgenden Umgebungsvariablen fest:

   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

   Ersetzen Sie Folgendes:

   • WORKLOAD: Legen Sie einen eindeutigen Namen fest, um Ihre Arbeitslast zu identifizieren.
   • WORKLOAD_NODEPOOL_COUNT: Die Anzahl der Knotenpools, die von einer Pathways-Arbeitslast verwendet werden.
   • TPU_TYPE: Der TPU-Typ gibt die Version und Größe der zu erstellenden Cloud TPU an. Weitere Informationen zu den unterstützten TPU-Typen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen
   • PROJECT: Ihre Google Cloud Projekt-ID
   • ZONE: Die Zone, in der Sie Ihre Arbeitslast ausführen möchten
   • CLUSTER: Der Name Ihres GKE-Cluster

2. Erstellen Sie die Pathways-Container im Cluster. Führen Sie den folgenden Befehl aus, um eine Headless-Arbeitslast auszuführen:

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

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

kubectl

Die folgende YAML-Datei ist mit der YAML-Datei für Batch-Arbeitslasten identisch, mit der Ausnahme, dass der main-Container nicht angegeben ist.

1. Ersetzen Sie die Platzhalter, kopieren Sie die folgende YAML-Datei und fügen Sie sie in eine Datei mit dem Namen pathways-headless-workload.yaml ein.

   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
       

   Ersetzen Sie Folgendes:
   • USERNAME : Ihr Nutzername
   • MAX_RESTARTS : Die maximale Anzahl der Neustarts des PathwaysJob
   • TPU_MACHINE_TYPE : Der TPU-Maschinentyp, den Sie verwenden möchten. Beispiel für unterstützte Werte: „ct6e-standard-8t“, „ct5p-hightpu-4t“
   • TOPOLOGY : Die TPU-Topologie
   • WORKLOAD_NODEPOOL_COUNT : Die Anzahl der Knotenpools, die von einer Pathways-Arbeitslast verwendet werden.
   • BUCKET_NAME : Ein Cloud Storage-Bucket, der zum Speichern temporärer Dateien verwendet wird
   Wenn Sie die Anzahl der Knotenpools (Pathways-Worker-Replikate) ändern möchten, die durch WORKLOAD_NODEPOOL_COUNT in der vorherigen YAML-Datei angegeben wurden, müssen Sie diesen PathwaysJob löschen und einen neuen PathwaysJob mit der aktualisierten Anzahl von Knotenpools erstellen. Außerdem müssen Sie alle verbundenen Notebooks neu starten, um die Verbindung zum neuen Pathways-Cluster herzustellen.
2. Wenden Sie die Datei pathways-headless-workload.yaml an:

     kubectl apply -f pathways-headless-workload.yaml
     

3. Führen Sie kubectl get pods aus, um zu prüfen, ob alle Container im Pod ausgeführt werden. Die folgende Ausgabe gilt für ein v5p-Slice mit 2 × 2 × 2. Dabei ist USER die ID des Nutzers, der den Befehl ausführt:

       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
       

Verbindung zum Pathways-Cluster im interaktiven Modus herstellen

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

Verbindung über Portweiterleitung herstellen

An diesem Punkt können Sie die Portweiterleitung (von jedem Host mit Zugriff auf die Steuerungsebene Ihres Clusters) verwenden, um auf den Proxyserver zuzugreifen:

Verwenden Sie den für Ihre Arbeitslast geeigneten Befehl:

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}

Die Ausgabe sollte etwa so aussehen:

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}

Die Ausgabe sollte etwa so aussehen:

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

Ö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:

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 sollte etwa so aussehen:

[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

Verbindung von Hosts in der VPC ohne Portweiterleitung herstellen

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

Verbindung über Cloud DNS herstellen

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

Wenn Sie Cloud DNS mit dem Cluster, dem additiven VPC-Bereich oder dem VPC-Bereich aktivieren, können die Kubernetes Cloud DNS-Namen von Nicht-GKE-VMs in Ihrer Virtual Private Cloud aufgelöst werden. Die Namen haben das Format <service_name>.<namespace>.svc.<custom_dns_domain>. Der Head Pod von Pathways hat einen Dienst mit dem Namen <jobset_name>-pathways-head-0-0.<jobset_name>.<namespace>.svc.<custom_dns_domain>.

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

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

   XPK

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

   Die Ausgabe sollte etwa so aussehen:

   <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

   Die Ausgabe sollte etwa so aussehen:

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

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

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

   Die Ausgabe sollte etwa so aussehen:

   [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 Ihrer VPC, die auf Ihre Pathways-Bereitstellung verweist, einen Dienst, der von einem internen Load-Balancer unterstützt wird. Dazu muss Cloud DNS für Ihren Cluster nicht aktiviert sein.

Für Cluster mit vielen VMs empfehlen wir, die Teilmengeneinstellung für interne Load-Balancer zu aktivieren, wenn Sie interne Load-Balancer erstellen. Weitere Informationen finden Sie unter GKE-Teilmengeneinstellung in einem vorhandenen Cluster aktivieren. Wenn die Teilmengeneinstellung für interne Load-Balancer nicht aktiviert ist, sind alle Knoten im Cluster Teil der Backend-Instanzgruppe für alle internen Load-Balancer. Dies ist nicht über 250 Knoten hinaus skalierbar. Wenn die Teilmengeneinstellung für interne Load-Balancer aktiviert ist, erstellt GKE anstelle von Instanzgruppen Netzwerknutzendpunktgruppen. Es werden nur Knoten berücksichtigt, auf denen einer der bereitstellenden Pods des Dienstes ausgeführt wird. Die Aktivierung der Teilmengeneinstellung für interne Load-Balancer hat eine einmalige Einrichtungslatenz von etwa 15 Minuten. Der folgende Befehl zeigt, wie Sie die Teilmengeneinstellung für interne Load-Balancer aktivieren:

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

Nachdem die Teilmengeneinstellung für interne Load-Balancer aktiviert wurde, können Sie mit der folgenden YAML-Datei einen Kubernetes-Dienst vom Typ LoadBalancer erstellen. Dadurch erstellt GKE einen internen Load-Balancer in der VPC Ihres Clusters:

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

Ersetzen Sie USER durch Ihre Google Cloud Nutzer-ID und speichern Sie die Datei als pathways-headless-ilb.yaml.

Wenden Sie das Manifest an:

kubectl apply -f pathways-headless-ilb.yaml

Nachdem der Load-Balancer erstellt wurde (etwa eine Minute später), hat die Spalte EXTERNAL-IP einen Wert:

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

Sie können ohne Portweiterleitung auf Hosts in derselben VPC wie Ihr Cluster auf die Pathways-Bereitstellung zugreifen:

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

Hinweis: Ersetzen Sie die IP-Adresse des JAX-Back-Ends (10.0.0.22) im vorherigen Befehl durch den Wert, der von kubectl get services angezeigt wird.

Die Ausgabe sollte etwa so aussehen:

[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 in den vorherigen Befehls- und Ausgabebeispielen gezeigten.

Jupyter-Notebooks

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 geprüft haben, können Sie über ein Vertex AI Jupyter-Notebook auf die GKE-TPU-VMs zugreifen. Bei der folgenden Einrichtung wird davon ausgegangen, dass sich Ihr GKE-Pathways-Cluster im selben Virtual Private Cloud-Netzwerk befindet (das Standardnetzwerk, sofern Sie nichts anderes konfiguriert haben). Rufen Sie die Vertex AI Workbench-Konsole auf.

Erstellen Sie auf dem Tab Instanzen mit der Schaltfläche Neu erstellen eine neue Workbench-Instanz. Achten Sie darauf, dass das Netzwerk mit dem Netzwerk Ihres GKE-Cluster identisch ist. Sie können über die Befehlszeile eine neue Workbench-Instanz erstellen.

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

Hinweis: das Flag --network ist optional. Wenn Sie es weglassen, wird das Standardnetzwerk verwendet.

Nachdem die Instanz erstellt wurde, rufen Sie sie auf und klicken Sie auf JupyterLab öffnen.

Selbst gehostete Jupyter-Notebook-Instanz erstellen

Der folgende Befehl zeigt, wie Sie mit XPK eine selbst gehostete Jupyter-Notebook-Instanz erstellen:

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"

Die folgende YAML-Datei zeigt, wie Sie mit kubectl eine selbst gehostete Jupyter-Notebook-Instanz erstellen. Wenden Sie die folgende YAML-Datei an, nachdem ein monitorloser Pathways-Cluster erstellt wurde. Weitere Informationen finden Sie unter Pathways im interaktiven Modus mit kubectl ausführen.

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

Zur Notebook-Instanz wechseln

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

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

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

kubectl logs ${MAIN_POD}

Die Ausgabe sollte so aussehen:

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

Notebook-Verbindung zum Pathways-Cluster

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

Fügen Sie im Notebook eine Zelle hinzu, um pathwaysutils zu installieren, legen Sie JAX_PLATFORMS auf proxy und JAX_BACKEND_TARGET auf PROXY_ADDRESS fest.

Hinweis: Der Wert für PROXY_ADDRESS hängt von der Verbindungsoption ab, die Sie unter Verbindung zum Pathways-Proxyserver herstellen ausgewählt haben.

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

Fügen Sie eine zweite Zelle als „Hallo Welt“-Prüfung hinzu und geben Sie die Geräte im Pathways-Cluster aus.

import pathwaysutils
import jax

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

Wenn alles funktioniert, wird eine Meldung angezeigt, dass das Pathways-on-Cloud-Back-End erkannt wurde.

Die Anzahl der aufgeführten JAX-Geräte muss mit der Anzahl der TPU-Chips und der Anzahl der Slices übereinstimmen, 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 Sie Berechnungen auf zwei Slices über ein einzelnes Notebook ausführen.

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)

Interaktiven Pathways-Cluster löschen

XPK

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

kubectl

kubectl delete -f pathways-headless-workload.yaml

Nächste Schritte

• GKE-Cluster mit Pathways erstellen
• Inferenz auf mehreren Hosts mit Pathways
• Batch-Arbeitslasten mit Pathways
• Interaktiver Modus von Pathways
• JAX-Arbeitslasten zu Pathways portieren
• Resilientes Training mit Pathways
• Fehlerbehebung bei Pathways