Google verwendet KI-Technologie, um Inhalte in Ihre bevorzugte Sprache zu übersetzen. KI-Übersetzungen können Fehler enthalten.

Interaktive Arbeitslast mit Pathways ausführen

Interaktive Pathways-Arbeitslasten sind Remote-JAX-Arbeitslasten, die auf einer VM ausgeführt werden, die nicht Teil des GKE-Cluster ist, in 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 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 Pathways erstellt
Cluster Toolkit installiert
Installierte Kubernetes-Tools
gcloud CLI installiert
Google Kubernetes Engine API aktiviert
Achten Sie darauf, dass Pathways für Ihr Google Cloud Projekt aktiviert ist.

Pathways im interaktiven Modus ausführen

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: Legen Sie hier einen eindeutigen Namen zur Identifizierung Ihres Workloads fest. Aufgrund von DNS-Labelbeschränkungen darf der Name maximal 22 Zeichen lang sein.
- 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: Ihre Google 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
Wenn Sie noch keinen Cloud Storage-Bucket für Pathways-Artefakte erstellt haben, tun Sie dies jetzt:
```
gcloud storage buckets create gs://${BUCKET_NAME} --location=${REGION}
```

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 Cluster Toolkit-Jobs:

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

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

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:

Cluster Toolkit

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-Script 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 Cluster-, additivem VPC- oder 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 mit Cloud DNS eine Verbindung zu Ihrem Pathways-Cluster herstellen:

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

Cluster Toolkit

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

Stellen Sie mit dem Cloud DNS-Namen eine Verbindung zum Pathways-Cluster her:

Cluster Toolkit

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.

Für Cluster 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 Teilmengeneinstellung für den internen Load-Balancer nicht aktiviert ist, sind alle Knoten im Cluster Teil der Backend-Instanzgruppe für alle internen Load-Balancer. Dies ist nicht für mehr als 250 Knoten skalierbar. Wenn die ILB-Teileinstellung aktiviert ist, erstellt GKE Netzwerk-Endpunktgruppen anstelle von Instanzgruppen. Nur Knoten, auf denen einer der Bereitstellungs-Pods des Dienstes ausgeführt wird, werden berücksichtigt. Die Aktivierung von ILB-Subsetting 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 die ILB-Teilmengeneinstellung aktiviert wurde, können Sie mit dem folgenden YAML einen Kubernetes-Dienst vom Typ LoadBalancer erstellen. Dadurch wird in GKE ein interner Load Balancer im VPC Ihres Clusters 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 das Pathways-Deployment 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())'

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

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. Dies ist das Standardnetzwerk, sofern Sie nichts anderes konfiguriert haben. Rufen Sie die Vertex AI Workbench-Konsole auf.

Erstellen Sie eine neue Workbench-Instanz (auf dem Tab Instanzen) mit der Schaltfläche Neu erstellen. 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]

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 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 zeigt, wie Sie mit kubectl eine selbst gehostete Jupyter-Notebook-Instanz erstellen. Wenden Sie das folgende YAML 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 Portweiterleitung eine Verbindung zum Notebook von Ihrem lokalen Computer aus 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 | 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

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

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

!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, sollte eine Meldung angezeigt werden, 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)

Interaktive Pathways-Arbeitslast löschen

Cluster Toolkit

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

kubectl

kubectl delete -f pathways-headless-workload.yaml

Interaktive Arbeitslast mit Pathways ausführen Mit Sammlungen den Überblick behalten Sie können Inhalte basierend auf Ihren Einstellungen speichern und kategorisieren.

Hinweis

Pathways im interaktiven Modus ausführen

Verbindung zum Pathways-Cluster im interaktiven Modus herstellen

Verbindung über Portweiterleitung herstellen

Cluster Toolkit

kubectl

Verbindung von Hosts in der VPC herstellen, ohne Portweiterleitung zu verwenden

Verbindung über Cloud DNS herstellen

Cluster Toolkit

kubectl

Cluster Toolkit

kubectl

Verbindung über einen internen Load Balancer herstellen

Jupyter-Notebooks

Vertex AI Workbench-Instanz erstellen

Selbst gehostete Jupyter-Notebook-Instanz erstellen

Zur Notebook-Instanz wechseln

Cluster Toolkit

kubectl

Notebook-Verbindung zum Pathways-Cluster

Code einem Notebook hinzufügen

Interaktive Pathways-Arbeitslast löschen

Cluster Toolkit

kubectl

Nächste Schritte

Interaktive Arbeitslast mit Pathways ausführen