경로를 사용하여 대화형 워크로드 실행

Pathways 대화형 워크로드는 클러스터의 Pathways 컨트롤러 구성요소 외부(예: VM)에서 실행되는 원격 JAX 워크로드입니다. 일괄 워크로드와 달리 대화형 워크로드가 완료되면 Pathways 클러스터 구성요소가 종료되지 않으며 다른 JAX 클라이언트의 연결에 계속 사용할 수 있습니다. 이 문서에서는 Jupyter 노트북을 예시로 사용하여 대화형 워크로드를 보여줍니다. Google Kubernetes Engine(GKE) 클러스터 또는 Slurm이 있는 Cluster Director 클러스터에서 Pathways를 사용하여 대화형 워크로드를 실행할 수 있습니다.

이 페이지에는 두 플랫폼에 대한 안내가 모두 포함되어 있습니다. 각 섹션에서 GKE 또는 Cluster Director 탭을 선택합니다.

JAX 사용자는 IFRT 인터페이스를 사용하여 Pathways 클러스터에 명령어를 전송합니다. 터미널, 노트북 또는 Python 호환 환경에서 실행되는 JAX 코드는 Pathways 리소스와 상호작용할 수 있습니다.

시작하기 전에

시작하기 전에

GKE

다음 기본 요건을 완료합니다.

Cluster Director

다음 기본 요건을 완료합니다.

대화형 모드로 Pathways 실행

GKE

  1. 다음 환경 변수를 설정합니다.

    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 버킷입니다.

  2. Pathways 아티팩트용 Cloud Storage 버킷을 아직 만들지 않았으면 버킷을 만듭니다.

    gcloud storage buckets create gs://${BUCKET_NAME} \
    --project=${PROJECT_ID} \
    --location=${REGION}

  3. 클러스터에 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-headless

    gcluster 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 프록시 서버에 연결할 수 있습니다.

Cluster Director

  1. 클러스터의 로그인 노드에 연결합니다.

    gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap

    다음을 바꿉니다.

    • CLUSTER_ID: Cluster Director 클러스터의 이름입니다.
    • PROJECT_ID:Google Cloud 프로젝트 ID입니다.
    • ZONE: 클러스터가 있는 영역입니다.

  2. 다음 환경 변수를 설정합니다.

    export SLICE_SHAPE=TOPOLOGY
export MAX_SLICES=NUM_SLICES
export MIN_SLICES=NUM_SLICES
export GCS_BUCKET=gs://BUCKET_NAME

    다음을 바꿉니다.

    • TOPOLOGY: 각 슬라이스의 TPU 토폴로지입니다(예: 2x4). TPU 노드 집합의 토폴로지와 일치해야 합니다.
    • NUM_SLICES: TPU 슬라이스 수입니다. 모든 슬라이스가 실행 중일 때만 Pathways 클러스터가 준비되도록 MIN_SLICES 및 MAX_SLICES를 동일한 값으로 설정합니다.
    • BUCKET_NAME: Pathways 아티팩트의 기존 Cloud Storage 버킷입니다.

    TPU 파티션이 클러스터의 기본 파티션이 아닌 경우 PARTITION을 해당 이름으로 설정합니다.

  3. 헤드리스 모드에서 Pathways 클러스터를 시작합니다.

    cd ~/pathways
sbatch start_pathways.sh

  4. 서비스 로그 pathways-service-JOB_ID.out(여기서 JOB_ID는 Slurm 작업 ID)에 다음과 유사한 출력이 표시될 때까지 기다립니다.

    Pathways Runtime Service is READY!
 Discovery file    : /home/USER/pathways/cluster.env
 IFRT Proxy Target : grpc://RM_IP:29000

    RM_IP는 프록시 서버의 내부 IP 주소입니다. ~/pathways/cluster.env 파일에 동일한 주소가 포함되어 있습니다.

이제 JAX 워크로드가 IFRT 프록시 서버에 연결할 수 있습니다.

대화형 모드에서 Pathways 클러스터에 연결

포트 전달을 사용하거나 사용하지 않고 Pathways 클러스터에 연결할 수 있습니다. 다음 섹션 중 하나를 사용하여 Pathways 클러스터에 연결합니다.

Cluster Director에서 JAX 클라이언트는 Pathways 서버와 동일한 JAX 버전인 0.11.1를 사용해야 합니다.

포트 전달을 사용하여 연결

GKE

클러스터의 컨트롤 플레인에 액세스할 수 있는 모든 호스트에서 로컬 포트 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

Cluster Director

워크스테이션 또는 클러스터 서브넷 외부의 VM과 같이 클러스터의 로그인 노드에 SSH로 연결할 수 있는 모든 호스트에서 로그인 노드를 통해 프록시 포트를 전달합니다.

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 29000:RM_IP:29000

동일한 호스트에서 새 터미널 창을 엽니다. JAX_PLATFORMS 및 JAX_BACKEND_TARGET 환경 변수를 설정하고 pathwaysutils 및 jax를 가져오는 Python 스크립트를 실행합니다. 호스트에 RM_IP에 대한 경로가 없으므로 클라이언트가 터널의 로컬 끝인 127.0.0.1:29000에 연결합니다.

python3 -m venv .venv
source .venv/bin/activate
pip install pathwaysutils jax==0.11.1

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(288,TPU_DEVICE,coords=[0,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 device(289,TPU_DEVICE,coords=[1,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 ...
 device(313,TPU_DEVICE,coords=[1,3,0,0],logical_task=1,slice=1,default_mem=device,mem_spaces=3)]
W1029 22:52:22.897483    2361 grpc_client.cc:91] IFRT proxy server disconnected: CANCELLED: Stream removed (CANCELLED)

스크립트가 종료되고 연결이 닫히면 IFRT proxy server disconnected 경고가 표시됩니다. 이는 정상적인 동작입니다.

포트 전달을 사용하지 않고 VPC의 호스트에서 연결

GKE

포트 전달을 사용하지 않으려면 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 클러스터에 연결하는 방법을 보여줍니다.

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

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

Cluster Director

Compute Engine VM과 같은 클러스터 서브넷의 호스트는 프록시 서버에 직접 연결할 수 있습니다. 클러스터 디렉터의 방화벽은 클러스터의 서브넷에서 트래픽을 허용하지만 동일한 VPC 네트워크의 다른 서브넷에서는 허용하지 않습니다.

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://RM_IP:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

RM_IP을 서비스 로그의 IFRT 프록시 주소로 바꿉니다.

Jupyter 노트북

GKE

Vertex AI를 사용하여 Jupyter 노트북을 만들거나 자체 호스팅 Jupyter 노트북을 만들 수 있습니다.

Vertex AI Workbench 인스턴스 생성

Pathways 클러스터를 설정하고 확인한 후 Vertex AI Workbench Jupyter 노트북에서 GKE TPU VM에 액세스할 수 있습니다. 다음 설정 안내에서는 GKE Pathways 클러스터가 동일한 VPC 네트워크에 있다고 가정합니다. 이 네트워크는 달리 구성하지 않는 한 기본 네트워크입니다.

  1. Google Cloud 콘솔에서 Vertex AI Workbench 페이지로 이동합니다.

  2. 인스턴스 탭에서 새로 만들기를 클릭하고 네트워크가 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 네트워크를 사용합니다.

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

클러스터의 서브넷에 있는 기존 Compute Engine VM에서 Jupyter 노트북을 실행하거나 클러스터의 로그인 노드에서 자체 호스팅 Jupyter 노트북을 만들 수 있습니다. 두 경우 모두 노트북은 포트 전달 없이 프록시 서버에 연결됩니다.

클러스터의 서브넷에서 기존 VM 사용

VM에서 가상 환경에 JupyterLab을 설치하고 시작합니다.

python3 -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
jupyter lab --no-browser --port=8888

자체 호스팅 Jupyter 노트북 인스턴스 만들기

클러스터의 로그인 노드에서 JupyterLab을 설치하고 시작합니다. 클러스터의 Python을 사용하여 홈 디렉터리에 가상 환경을 만듭니다. 로그인 노드의 시스템 Python이 현재 JAX 출시 버전에 비해 너무 오래되었으므로 pip에서 jax==0.11.1를 찾을 수 없습니다.

cd ~
/home/pathways/venv/bin/python -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
source ~/pathways/cluster.env
jupyter lab --no-browser --port=8888

노트북 인스턴스에 연결

GKE

포트 전달을 사용하여 로컬 머신에서 노트북에 연결합니다.

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

Cluster Director

VM 또는 로그인 노드의 노트북의 경우 해당 호스트에 대한 포트 전달을 사용하여 로컬 머신에서 연결합니다. 예를 들어 로그인 노드의 경우 다음과 같습니다.

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 8888:localhost:8888

로컬 브라우저에서 http://localhost:8888?token=TOKEN로 이동합니다. TOKEN을 JupyterLab이 시작될 때 출력한 토큰으로 바꿉니다.

Pathways 클러스터에 대한 노트북 연결

  1. JupyterLab 내에서 새 Python 3 노트북을 만듭니다.
  2. 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 주소)로 바꿉니다. Cluster Director에서 포트 전달을 사용하는 경우 127.0.0.1를 사용하거나 서비스 로그의 IFRT 프록시 주소를 사용합니다. 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 워크로드를 만듭니다. Cluster Director에서 MIN_SLICES 및 MAX_SLICES를 2 이상으로 설정합니다. 슬라이스가 하나인 경우 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 대화형 워크로드 삭제

GKE

Cluster Toolkit

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

kubectl

kubectl delete job WORKLOAD

Cluster Director

클러스터의 로그인 노드에서 Pathways 서비스 작업을 취소합니다.

scancel -u ${USER} -n pathways-service

이 명령어는 Pathways 워커 작업도 취소합니다. TPU VM은 정적 노드이므로 계속 실행되며 다른 Slurm 작업에 사용할 수 있습니다.

문제가 발생하면 Cloud에서 Pathways 문제 해결을 참조하세요.

다음 단계