Interaktive Arbeitslast mit Pathways ausführen

Vorschau

Dieses Produkt oder diese Funktion unterliegt den „Pre-GA-Angebotsbedingungen“ im Abschnitt „Allgemeine Dienstbedingungen“ der dienstspezifischen Nutzungsbedingungen. Pre-GA-Produkte und ‑Funktionen sind „wie gesehen“ verfügbar und bieten möglicherweise nur eingeschränkten Support. Weitere Informationen finden Sie in den Beschreibungen der Produkteinführungsphasen.

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

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

Hinweise

Sie benötigen Folgendes:

• GKE-Cluster mit XPK erstellt
• Installierte XPKs
• Installierte Kubernetes-Tools
• gcloud CLI installiert
• TPU API aktiviert
• Google Kubernetes Engine API aktiviert
• Achten Sie darauf, dass Pathways für Ihr Google Cloud -Projekt aktiviert ist.

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 hier einen eindeutigen Namen fest, um Ihren Workload 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 Cloud TPU an, die Sie erstellen möchten. 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 dieselbe wie die YAML-Datei für Batcharbeitslasten, mit der Ausnahme, dass der main-Container nicht angegeben ist.

1. Ersetzen Sie die Platzhalter, kopieren Sie das folgende YAML-Manifest und fügen Sie es 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 von 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, in dem temporäre Dateien gespeichert werden.
   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 diese PathwaysJob löschen und eine neue 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 bezieht sich auf einen 2-Slice-v5p-2x2x2-Job, wobei USER die ID des Nutzers ist, 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 einem beliebigen 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, in dem pathwaysutils und jax importiert werden:

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 herstellen, ohne Portweiterleitung zu verwenden

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 Cloud DNS-Anbieter von kube-dns auf Cloud DNS umgestellt. 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- 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 Pathways-Head-Pod 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 über 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 mit dem 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 einen Dienst, der von einem internen Load Balancer unterstützt wird, um eine private IP-Adresse in Ihrer VPC zu erhalten, die auf Ihre Pathways-Bereitstellung verweist. Dazu muss Cloud DNS nicht für Ihren Cluster aktiviert sein.

Bei Clustern mit vielen VMs empfehlen wir, die Teilmengeneinstellung des internen Load Balancers zu aktivieren, wenn Sie interne Load Balancer erstellen. Weitere Informationen finden Sie unter GKE-Teilmengeneinstellung 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. Diese Methode ist nicht für mehr als 250 Knoten geeignet. Wenn die ILB-Teileinstellung aktiviert ist, erstellt GKE Netzwerk-Endpunktgruppen anstelle von Instanzgruppen. Es werden nur Knoten berücksichtigt, auf denen einer der Bereitstellungs-Pods des Dienstes ausgeführt wird. Die Aktivierung der ILB-Teilmengeneinstellung hat eine einmalige Einrichtungsverzögerung von etwa 15 Minuten. Der folgende Befehl zeigt, wie das ILB-Subsetting aktiviert wird:

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

Nachdem ILB-Subsetting aktiviert wurde, können Sie mit dem folgenden YAML-Code einen Kubernetes-Dienst vom Typ „LoadBalancer“ erstellen. Dadurch wird in der VPC Ihres Clusters ein interner Load-Balancer erstellt:

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 (nach etwa einer Minute), enthält 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 auf die Bereitstellung von Pathways zugreifen, ohne Portweiterleitung auf Hosts in derselben VPC wie Ihr Cluster zu verwenden:

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-Backends (10.0.0.22) im vorherigen Befehl durch den von kubectl get services angezeigten Wert.

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 IP-Adressen, die in den Beispielen für den vorherigen Befehl und die Ausgabe gezeigt werden.

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 bestätigt haben, können Sie über ein Vertex AI Jupyter-Notebook auf die GKE-TPU-VMs zugreifen. Bei der folgenden Einrichtungsanleitung wird davon ausgegangen, dass sich Ihr GKE Pathways-Cluster im selben VPC-Netzwerk (Virtual Private Cloud) befindet. Das ist 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 übereinstimmt. Sie können eine neue Workbench-Instanz über die Befehlszeile 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.

Wenn 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"

Das folgende YAML zeigt, wie Sie mit kubectl eine selbst gehostete Jupyter-Notebookinstanz erstellen. Wenden Sie die folgende YAML-Datei an, nachdem ein Headless-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 Portweiterleitung eine Verbindung zum Notebook von Ihrem lokalen Computer aus 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. Verbindung zum Pathways-Proxyserver herstellen

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

Hinweis: Der Wert für PROXY_ADDRESS hängt von der Verbindungsoption ab, die Sie unter Mit dem Pathways-Proxyserver verbinden 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 „Hello World“-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, sollten Sie eine Meldung sehen, dass das Pathways-on-Cloud-Backend erkannt wurde.

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

Code 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 für zwei Slices aus einem einzelnen Notebook ausgeführt werden.

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 Cluster für Pathways 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 für Lernpfade
• JAX-Arbeitslasten zu Pathways migrieren
• Belastbares Training mit Pathways
• Pfade zur Fehlerbehebung