Pathways 대화형 워크로드는 Google Kubernetes Engine(GKE) 클러스터의 Pathways 컨트롤러 구성요소 외부(예: VM)에서 실행되는 원격 JAX 워크로드입니다. 일괄 워크로드와 달리 대화형 워크로드가 완료되면 Pathways 클러스터 구성요소가 종료되지 않으며 다른 JAX 클라이언트의 연결에 계속 사용할 수 있습니다. 이 문서에서는 Jupyter 노트북을 예시로 사용하여 대화형 워크로드를 보여줍니다.
JAX 사용자는 IFRT 인터페이스를 사용하여 Pathways 클러스터에 명령어를 전송합니다. 터미널, 노트북 또는 Python 호환 환경에서 실행되는 JAX 코드는 Pathways 리소스와 상호작용할 수 있습니다.
시작하기 전에
다음 기본 요건을 완료합니다.
- Pathways로 GKE 클러스터를 만듭니다.
- Cluster Toolkit을 설치합니다.
- Kubernetes 도구 설치
- Google Cloud CLI를 설치합니다.
- Kubernetes Engine API를 사용 설정합니다.
- Google Cloud 프로젝트에 Pathways가 사용 설정되어 있는지 확인합니다.
대화형 모드로 Pathways 실행
다음 환경 변수를 설정합니다.
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
다음을 바꿉니다.
WORKLOAD: 워크로드를 식별하는 고유한 이름입니다(DNS 라벨 제한으로 인해 22자(영문 기준) 이하).WORKLOAD_NODEPOOL_COUNT: Pathways 워크로드에서 사용하는 노드 풀 수입니다.COMPUTE_TYPE: TPU 머신 유형(예:ct6e-standard-4t) 각 TPU 버전에서 지원되는 TPU 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.TOPOLOGY: TPU 배치 토폴로지(예:2x4)PROJECT_ID:Google Cloud 프로젝트 ID입니다.REGION: 버킷과 클러스터가 있는 리전입니다.ZONE: 워크로드를 실행할 영역입니다.CLUSTER: GKE 클러스터의 이름입니다.BUCKET_NAME: Pathways 아티팩트의 Cloud Storage 버킷입니다.
Pathways 아티팩트용 Cloud Storage 버킷을 아직 만들지 않았으면 버킷을 만듭니다.
gcloud storage buckets create gs://${BUCKET_NAME} \ --project=${PROJECT_ID} \ --location=${REGION}클러스터에 Pathways 컨테이너를 만듭니다. 헤드리스 워크로드를 실행하려면
gcluster job submit명령어를 사용합니다. 자세한 내용은 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-headlessgcluster job submit명령어는 Pathways 헤드 포드를 실행하기 위해cpu-np또는pathways-np라는 CPU 노드 풀을 찾습니다. CPU 노드 풀의 이름이 다른 경우(예: 클러스터 생성 가이드의cpu-pathways-np), 이 페이지의gcluster job submit명령어에--pathways-head-np=NODE_POOL_NAME플래그를 추가합니다. 그렇지 않으면 명령어가failed to auto-detect a suitable CPU node pool오류와 함께 실패합니다.
이제 JAX 워크로드가 IFRT 프록시 서버에 연결할 수 있습니다.
대화형 모드에서 Pathways 클러스터에 연결
포트 전달을 사용하거나 사용하지 않고 Pathways 클러스터에 연결할 수 있습니다. 다음 섹션 중 하나를 사용하여 Pathways 클러스터에 연결합니다.
포트 전달을 사용하여 연결
클러스터의 컨트롤 플레인에 액세스할 수 있는 모든 호스트에서 로컬 포트 29000를 Pathways 헤드 포드로 전달하여 프록시 서버에 액세스합니다.
PROXY_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
PROXY_PORT=29000
kubectl port-forward ${PROXY_POD} ${PROXY_PORT}:${PROXY_PORT}
출력은 다음과 비슷합니다.
Forwarding from 127.0.0.1:29000 -> 29000
Forwarding from [::1]:29000 -> 29000
동일한 호스트에서 새 터미널 창을 엽니다. Python 3.11 이상을 사용하여 가상 환경을 만들고, pathwaysutils 및 jax 패키지를 설치하고, JAX_PLATFORMS 및 JAX_BACKEND_TARGET 환경 변수를 설정하고, pathwaysutils 및 jax를 가져오는 Python 스크립트를 실행합니다.
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())'
출력은 다음과 비슷합니다.
[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
포트 전달을 사용하지 않고 VPC의 호스트에서 연결
포트 전달을 사용하지 않으려면 Cloud DNS 또는 내부 부하 분산기를 사용하여 Pathways 클러스터에 연결하면 됩니다.
Cloud DNS를 사용하여 연결
클러스터에서 Cloud DNS를 사용 설정하면 DNS 제공업체가 kube-dns에서 Cloud DNS로 전환됩니다. 사용 설정하면 GKE가 DNS 이름에 대해 가상 프라이빗 클라우드(VPC) 네트워크에 비공개 Cloud DNS 영역을 만듭니다. 자세한 내용은 GKE용 Cloud DNS 사용을 참조하세요.
VPC 범위 또는 추가 VPC 범위로 Cloud DNS를 사용 설정하면 VPC 네트워크 내의 GKE가 아닌 VM에서 Kubernetes DNS 이름을 변환할 수 있습니다. 클러스터 범위에서는 클러스터의 노드만 이러한 이름을 확인할 수 있습니다. 이름의 형식은 SERVICE_NAME.NAMESPACE.svc.CUSTOM_DNS_DOMAIN입니다.
Pathways 헤드 포드에는 WORKLOAD-pathways-head-0-0.WORKLOAD.NAMESPACE.svc.CUSTOM_DNS_DOMAIN라는 서비스가 있습니다.
다음 명령어는 Cloud DNS를 사용하여 Pathways 클러스터에 연결하는 방법을 보여줍니다.
GKE가 아닌 호스트에서 리더 DNS 항목을 확인할 수 있는지 확인합니다.
host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN다음을 바꿉니다.
WORKLOAD: Pathways 워크로드의 이름입니다.CUSTOM_DNS_DOMAIN: 클러스터에 구성된 커스텀 DNS 도메인입니다.
출력은 다음과 비슷합니다.
WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN has address 10.0.2.75
DNS 이름을 사용하여 Pathways 클러스터에 연결합니다.
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())'
출력은 다음과 비슷합니다.
[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
내부 부하 분산기를 사용하여 연결
Pathways 배포를 가리키는 VPC 네트워크의 비공개 IP 주소의 경우 내부 부하 분산기(ILB)에서 지원하는 서비스를 만듭니다. 이렇게 하면 클러스터에 Cloud DNS를 사용 설정하지 않아도 됩니다.
VM이 많은 클러스터의 경우 내부 부하 분산기를 만들 때 ILB 하위 설정을 사용 설정합니다. 자세한 내용은 기존 클러스터에서 GKE 하위 설정 사용 설정을 참조하세요. ILB 하위 설정이 사용 설정되지 않은 경우 클러스터의 모든 노드가 모든 내부 부하 분산기의 백엔드 인스턴스 그룹에 속하며, 이 그룹은 250개 노드를 초과하여 확장되지 않습니다. ILB 하위 설정을 사용 설정하면 GKE는 인스턴스 그룹 대신 네트워크 엔드포인트 그룹을 만들고 서비스의 제공 포드 중 하나를 실행하는 노드만 포함합니다. ILB 하위 설정을 사용 설정하면 일회성 설정 지연 시간이 약 15분 발생합니다. 다음 명령어는 ILB 하위 설정을 사용 설정합니다.
gcloud container clusters update ${CLUSTER} \ --project=${PROJECT_ID} \ --location=CLUSTER_LOCATION \ --enable-l4-ilb-subsetting
CLUSTER_LOCATION을 클러스터가 있는 리전 또는 영역으로 바꿉니다.
ILB 하위 설정을 사용 설정한 후 다음 YAML 매니페스트를 사용하여 LoadBalancer 유형의 Kubernetes 서비스를 만듭니다. 이렇게 하면 GKE가 클러스터의 VPC 네트워크 내에 내부 부하 분산기를 만듭니다.
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
WORKLOAD을 Pathways 워크로드 이름으로 바꾸고 파일을 pathways-headless-ilb.yaml로 저장합니다.
매니페스트를 적용합니다.
kubectl apply -f pathways-headless-ilb.yaml
GKE가 부하 분산기를 만든 후(약 1분 후) EXTERNAL-IP 열에 IP 주소가 표시되는지 확인합니다.
kubectl get services
출력은 다음과 비슷합니다.
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
클러스터와 동일한 VPC 네트워크의 호스트에서 포트 전달 없이 Pathways 배포에 액세스할 수 있습니다.
JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://EXTERNAL_IP:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'
EXTERNAL_IP를 kubectl get services에 표시된 EXTERNAL-IP 값(예: 10.0.0.22)으로 바꿉니다.
출력은 다음과 비슷합니다.
[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 노트북
Vertex AI를 사용하여 Jupyter 노트북을 만들거나 자체 호스팅 Jupyter 노트북을 만들 수 있습니다.
Vertex AI Workbench 인스턴스 생성
Pathways 클러스터를 설정하고 확인한 후 Vertex AI Workbench Jupyter 노트북에서 GKE TPU VM에 액세스할 수 있습니다. 다음 설정 안내에서는 GKE Pathways 클러스터가 동일한 VPC 네트워크에 있다고 가정합니다. 이 네트워크는 달리 구성하지 않는 한 기본 네트워크입니다.
- Google Cloud 콘솔에서 Vertex AI Workbench 페이지로 이동합니다.
인스턴스 탭에서 새로 만들기를 클릭하고 네트워크가 GKE 클러스터의 네트워크와 일치하는지 확인합니다.
또는 명령줄을 사용하여 새 Vertex AI Workbench 인스턴스를 만듭니다.
gcloud workbench instances create INSTANCE_NAME \ --project=PROJECT_ID \ --machine-type=e2-standard-4 \ --data-disk-size=100 \ --location=ZONE \ --network=NETWORK
다음을 바꿉니다.
INSTANCE_NAME: Vertex AI Workbench 인스턴스 이름PROJECT_ID:Google Cloud 프로젝트 ID입니다.ZONE: 인스턴스를 만들 영역입니다.NETWORK: GKE 클러스터의 VPC 네트워크 이름입니다.--network플래그는 선택사항입니다. 생략하면 명령어가default네트워크를 사용합니다.
Vertex AI가 인스턴스를 만든 후 인스턴스로 이동하여 JupyterLab 열기를 클릭합니다.
자체 호스팅 Jupyter 노트북 인스턴스 만들기
다음 명령어는 gcluster job submit 명령어를 사용하여 자체 호스팅 Jupyter 노트북 인스턴스를 만드는 방법을 보여줍니다.
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"
다음 YAML 매니페스트는 kubectl를 사용하여 자체 호스팅 Jupyter 노트북 인스턴스를 만드는 방법을 보여줍니다. 헤드리스 Pathways 클러스터를 만든 후 다음 YAML 매니페스트를 적용합니다. 자세한 내용은 대화형 모드에서 Pathways 실행을 참조하세요.
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
hostNetwork: true 필드를 사용하면 노트북의 JAX 클라이언트가 Pathways 프록시 서버와 애플리케이션 레이어 전송 보안(ALTS) 핸드셰이크를 완료할 수 있습니다. 이 필드가 없으면 GKE용 워크로드 아이덴티티 제휴를 사용하는 노드 풀에서 연결이 실패합니다. dnsPolicy: ClusterFirstWithHostNet 필드를 사용하면 노트북에서 Pathways 헤드 포드 이름과 같은 클러스터 내 DNS 이름을 확인할 수 있습니다. 자세한 내용은 Cloud의 Pathways 문제 해결을 참조하세요.
노트북 인스턴스에 연결
포트 전달을 사용하여 로컬 머신에서 노트북에 연결합니다.
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
로컬 브라우저에서 http://localhost:8888?token=TOKEN로 이동합니다.
thought
TOKEN을 Jupyter 노트북 컨테이너 로그의 토큰으로 바꿉니다.
kubectl logs ${MAIN_POD}
출력은 다음과 비슷합니다.
... 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은 노트북 포드를 실행하는 GKE 노드의 호스트 이름입니다(예: gke-my-cluster-cpu-np-1a2b3c4d-ef56).
Pathways 클러스터에 대한 노트북 연결
- JupyterLab 내에서 새 Python 3 노트북을 만듭니다.
- Pathways 프록시 서버에 연결합니다.
노트북에서 셀을 추가하여 pathwaysutils를 설치하고 JAX_PLATFORMS를 proxy로 설정하고 JAX_BACKEND_TARGET를 grpc://PROXY_ADDRESS:29000로 설정합니다.
!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
PROXY_ADDRESS을 대화형 모드에서 Pathways 클러스터에 연결의 Pathways 프록시 서버의 호스트 또는 IP 주소(예: 127.0.0.1, Cloud DNS 호스트 이름 또는 내부 부하 분산기 IP 주소)로 바꿉니다. kubectl를 사용하여 자체 호스팅 노트북을 배포한 경우 PROXY_ADDRESS를 Pathways 헤드 포드의 클러스터 내 DNS 이름으로 바꿉니다. 그러면 JAX_BACKEND_TARGET 값이 grpc://WORKLOAD-pathways-head-0-0.WORKLOAD:29000가 됩니다.
두 번째 셀을 추가하여 연결을 확인하고 Pathways 클러스터의 기기를 출력합니다.
import pathwaysutils
import jax
pathwaysutils.initialize()
print(jax.devices())
연결에 성공하면 출력에 JAX가 Pathways-on-Cloud 백엔드를 감지했음을 나타내는 메시지가 포함됩니다.
나열된 JAX 기기 수는 Pathways 클러스터를 만들 때 지정한 TPU 칩 수 및 슬라이스 수와 일치합니다.
노트북에 코드 추가
자체 JAX 코드를 추가하고 Pathways 클러스터의 TPU에서 대화형으로 실행합니다. 다음 코드는 단일 노트북에서 두 슬라이스에 걸쳐 계산을 수행하는 방법을 보여줍니다. 전체 예시를 실행하려면 WORKLOAD_NODEPOOL_COUNT를 2 이상으로 설정하여 슬라이스가 두 개 이상인 Pathways 워크로드를 만듭니다. 슬라이스가 하나인 경우 slice1_devices는 비어 있고 예시의 마지막 부분은 ValueError로 인해 실패합니다.
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)
Pathways 대화형 워크로드 삭제
Cluster Toolkit
gcluster job cancel WORKLOAD \ --cluster=CLUSTER \ --project=PROJECT_ID \ --location=ZONE
kubectl
kubectl delete job WORKLOAD