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, der den Pathways-Cluster hostet. Im Gegensatz zu Batch-Arbeitslasten werden die Pathways-Clusterkomponenten nach Abschluss des interaktiven Arbeitslastvorgangs nicht heruntergefahren. Sie bleiben für die Verbindung 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.

Hinweis

Sie benötigen Folgendes:

GKE-Cluster mit Pathways erstellt
Installiertes Cluster Toolkit
Installierte Kubernetes-Tools
gcloud CLI installiert
Google Kubernetes Engine API aktivieren
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 Ihrer Arbeitslast fest (aufgrund von DNS-Label-Limits maximal 22 Zeichen).
- 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, erstellen Sie einen:
```
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 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

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

Verbindung zum Pathways-Cluster im interaktiven Modus 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.

Verbindung über Portweiterleitung herstellen

An dieser Stelle 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-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 ü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-, additiven 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 über 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 über den 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 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. Dafür muss Cloud DNS nicht für Ihren Cluster aktiviert sein.

Für Cluster mit vielen VMs empfehlen wir, die ILB-Teilmengeneinstellung 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. Dies ist nicht über 250 Knoten hinaus skalierbar. Wenn die ILB-Teileinstellung aktiviert ist, erstellt GKE Netzwerk-Endpunktgruppen anstelle von Instanzgruppen. Es werden nur Knoten einbezogen, auf denen einer der Bereitstellungs-Pods des Dienstes ausgeführt wird. Die Aktivierung der ILB-Teilmengen führt zu einer einmaligen Einrichtungslatenz (ca. 15 Minuten). Der folgende Befehl zeigt, wie Sie die ILB-Teilmengeneinstellung aktivieren:

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 der folgenden YAML-Datei einen Kubernetes-Dienst vom Typ „LoadBalancer“ erstellen. Dadurch wird GKE veranlasst, einen internen Load Balancer in der VPC Ihres Clusters zu erstellen:

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

Aktualisieren Sie USER mit Ihrer 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), 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 ohne Portweiterleitung auf Hosts in derselben VPC wie Ihr Cluster 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())'

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

Erstellen Sie auf dem Tab Instanzen mit dem Button Neu erstellen eine neue Workbench-Instanz. Achten Sie darauf, dass das Netzwerk mit dem Netzwerk Ihres GKE-Cluster übereinstimmt. Sie können die Befehlszeile verwenden, um eine neue Workbench-Instanz zu 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"

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

Notebookinstanz aufrufen

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 | 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 im Notebook eine Zelle hinzu, um pathwaysutils zu installieren, setzen Sie JAX_PLATFORMS auf proxy und JAX_BACKEND_TARGET auf PROXY_ADDRESS.

!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“-Typprüfung hinzu und geben Sie die Geräte im Pathways-Cluster aus.

import pathwaysutils
import jax

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

Wenn alles korrekt 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 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.

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 ohne Portweiterleitung herstellen

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

Notebookinstanz aufrufen

Cluster Toolkit

kubectl

Notebook-Verbindung zum Pathways-Cluster

Code zu einem Notebook hinzufügen

Interaktive Pathways-Arbeitslast löschen

Cluster Toolkit

kubectl

Nächste Schritte

Interaktive Arbeitslast mit Pathways ausführen