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

프리뷰

이 제품이나 기능에는 서비스별 약관이 적용됩니다. GA 이전 제품과 기능은 '있는 그대로' 제공되며 지원이 제한될 수 있습니다. 자세한 내용은 출시 단계 설명을 참조하세요.

경로 대화형 워크로드는 경로 클러스터를 호스팅하는 GKE 클러스터에 속하지 않는 VM 내에서 실행되는 원격 JAX 워크로드입니다. 배치 워크로드와 달리 대화형 워크로드 작업이 완료되어도 Pathways 클러스터 구성요소가 종료되지 않으며, 다른 JAX 클라이언트가 연결할 수 있는 상태로 유지됩니다. 이 문서에서는 Jupyter 노트북을 예로 사용하여 대화형 워크로드를 보여줍니다.

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

시작하기 전에

다음 사항이 필요합니다.

• XPK를 사용하여 GKE 클러스터를 만들었습니다.
• 설치된 XPK
• 설치된 Kubernetes 도구
• gcloud CLI 설치
• TPU API 사용 설정
• Google Kubernetes Engine API 사용 설정
• Google Cloud 프로젝트에 개발자 과정이 사용 설정되어 있는지 확인합니다.

대화형 모드로 Pathways 실행

xpk 또는 kubectl을 사용하여 대화형 모드에서 Pathways를 실행할 수 있습니다.

XPK

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

   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

   다음을 바꿉니다.

   • WORKLOAD: 워크로드를 식별하는 고유한 이름으로 설정합니다.
   • WORKLOAD_NODEPOOL_COUNT: Pathways 워크로드에서 사용하는 노드 풀 수
   • TPU_TYPE: TPU 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. TPU 버전마다 지원되는 TPU 유형에 대한 자세한 내용은 TPU 버전을 참고하세요.
   • PROJECT: Google Cloud 프로젝트 ID
   • ZONE: 워크로드를 실행할 영역
   • CLUSTER: GKE 클러스터의 이름

2. 클러스터에 Pathways 컨테이너를 만듭니다. 헤드리스 워크로드를 실행하려면 다음 명령어를 실행합니다.

   xpk workload create-pathways \
   --headless \
   --workload=${WORKLOAD} \
   --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
   --tpu-type=${TPU_TYPE} \
   --project=${PROJECT} \
   --zone=${ZONE} \
   --cluster=${CLUSTER}

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

kubectl

다음 YAML은 main 컨테이너를 지정하지 않는다는 점을 제외하고는 배치 워크로드 YAML과 동일합니다.

1. 자리표시자를 바꾸고 다음 YAML을 복사하여 pathways-headless-workload.yaml이라는 파일에 붙여넣습니다.

   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
       

   다음을 바꿉니다.
   • USERNAME : 사용자 이름
   • MAX_RESTARTS : PathwaysJob을 다시 시작할 수 있는 최대 횟수입니다.
   • TPU_MACHINE_TYPE: 사용할 TPU 머신 유형입니다. 지원되는 값의 예는 'ct6e-standard-8t', 'ct5p-hightpu-4t'입니다.
   • TOPOLOGY : TPU 토폴로지
   • WORKLOAD_NODEPOOL_COUNT : Pathways 워크로드에서 사용되는 노드 풀 수
   • BUCKET_NAME : 임시 파일을 저장하는 데 사용되는 Cloud Storage 버킷
   이전 YAML에서 WORKLOAD_NODEPOOL_COUNT로 지정된 노드 풀 수(경로-작업자 복제본)를 변경하려면 이 PathwaysJob를 삭제하고 업데이트된 노드 풀 수로 새 PathwaysJob를 만들어야 합니다. 새 Pathways 클러스터와의 연결을 설정하려면 연결된 노트북도 다시 시작해야 합니다.
2. pathways-headless-workload.yaml 파일을 적용합니다.

     kubectl apply -f pathways-headless-workload.yaml
     

3. kubectl get pods를 실행하여 포드의 모든 컨테이너가 실행 중인지 확인합니다. 다음 출력은 2 슬라이스 v5p 2x2x2에 대한 출력입니다. 여기서 USER은 명령어를 실행하는 사용자의 ID입니다.

       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
       

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

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

포트 전달을 사용하여 연결

이 시점에서 클러스터의 컨트롤 플레인에 액세스할 수 있는 호스트에서 포트 포워딩을 사용하여 프록시 서버에 액세스할 수 있습니다.

워크로드에 적합한 명령어를 사용합니다.

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}

다음과 비슷한 출력이 표시됩니다.

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}

다음과 비슷한 출력이 표시됩니다.

Forwarding from 127.0.0.1:29000 -> 29000
Forwarding from [::1]:29000 -> 29000

동일한 호스트에서 새 터미널 창을 엽니다. 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를 사용 설정하면 Cloud DNS 제공업체가 kube-dns에서 Cloud DNS로 전환됩니다. 사용 설정하면 Cloud DNS 이름에 대해 가상 프라이빗 클라우드에 비공개 Cloud DNS 영역이 생성됩니다. 자세한 내용은 GKE에 Cloud DNS 사용을 참고하세요.

클러스터, 추가 VPC 또는 VPC 범위로 Cloud DNS를 사용 설정하면 Virtual Private Cloud 내의 GKE 이외의 VM에서 Kubernetes Cloud DNS 이름을 확인할 수 있습니다. 이름은 <service_name>.<namespace>.svc.<custom_dns_domain> 형식입니다. 경로 헤드 포드에는 <jobset_name>-pathways-head-0-0.<jobset_name>.<namespace>.svc.<custom_dns_domain>라는 서비스가 있습니다.

다음 명령어는 Cloud DNS를 사용하여 Pathways 클러스터에 연결하는 방법을 보여줍니다.

1. 리더 Cloud DNS 항목이 GKE가 아닌 호스트에서 확인할 수 있는지 확인합니다.

   XPK

   host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.USERNAME-test

   다음과 비슷한 출력이 표시됩니다.

   <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

   다음과 비슷한 출력이 표시됩니다.

   pathways-<user>-pathways-head-0-0.pathways-<user>.default.svc.<user>-test has address 10.0.2.75

2. Cloud DNS 이름을 사용하여 Pathways 클러스터에 연결합니다.

   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())'

   다음과 비슷한 출력이 표시됩니다.

   [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

내부 부하 분산기를 사용하여 연결

VPC에서 경로 배포를 가리키는 비공개 IP 주소의 경우 내부 부하 분산기로 지원되는 서비스를 만듭니다. 클러스터에서 Cloud DNS를 사용 설정하지 않아도 됩니다.

VM이 많은 클러스터의 경우 내부 부하 분산기를 만드는 경우 ILB 하위 설정을 사용 설정하는 것이 좋습니다. 자세한 내용은 기존 클러스터에서 GKE 하위 설정 사용 설정을 참고하세요. ILB 하위 설정이 사용 설정되지 않은 경우 클러스터의 모든 노드가 모든 내부 부하 분산기의 백엔드 인스턴스 그룹에 속하게 됩니다. 이 방법은 250개 노드를 초과하여 확장되지 않습니다. ILB 하위 설정이 사용 설정된 경우 GKE는 인스턴스 그룹 대신 네트워크 엔드포인트 그룹을 만들고 서비스의 서비스 제공 포드 중 하나를 실행하는 노드만 포함됩니다. ILB 하위 설정을 사용 설정하면 일회성 설정 지연 시간 (~15분)이 발생합니다. 다음 명령어는 ILB 서브세팅을 사용 설정하는 방법을 보여줍니다.

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

ILB 서브세팅이 사용 설정되면 다음 yaml을 사용하여 LoadBalancer 유형의 Kubernetes 서비스를 만들 수 있습니다. 이렇게 하면 GKE가 클러스터의 VPC 내부에 내부 부하 분산기를 만듭니다.

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

USER를 Google Cloud 사용자 ID로 업데이트하고 파일을 pathways-headless-ilb.yaml로 저장합니다.

매니페스트를 적용합니다.

kubectl apply -f pathways-headless-ilb.yaml

부하 분산기가 생성된 후 (약 1분 후) EXTERNAL-IP 열에 값이 표시됩니다.

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

클러스터와 동일한 VPC의 호스트에서 포트 포워딩 없이 경로 배포에 액세스할 수 있습니다.

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())'

참고: 이전 명령어에서 JAX 백엔드 (10.0.0.22)의 IP 주소를 kubectl get services에 표시된 값으로 바꿉니다.

다음과 비슷한 출력이 표시됩니다.

[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

참고: IP 주소는 앞의 명령어와 출력 예에 표시된 주소와 다릅니다.

Jupyter 노트북

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

Vertex AI Workbench 인스턴스 만들기

경로 클러스터를 설정하고 확인한 후 Vertex AI Jupyter 노트북에서 GKE TPU VM에 액세스할 수 있습니다. 다음 설정 안내에서는 GKE Pathways 클러스터가 동일한 Virtual Private Cloud 네트워크에 있다고 가정합니다 (별도로 구성하지 않은 경우 기본 네트워크). Vertex AI Workbench 콘솔로 이동합니다.

인스턴스 탭에서 새로 만들기 버튼을 사용하여 새 Workbench 인스턴스를 만듭니다. 네트워크가 GKE 클러스터의 네트워크와 동일한지 확인합니다. 명령줄을 사용하여 새 워크벤치 인스턴스를 만들 수 있습니다.

gcloud workbench instances create INSTANCE_NAME \
--machine-type=e2-standard-4 \
--data-disk-size=100 \
--location=ZONE \
[--network=NETWORK]

참고: --network 플래그는 선택사항입니다. 생략하면 기본 네트워크가 사용됩니다.

인스턴스가 생성되면 인스턴스로 이동하여 JupyterLab 열기를 클릭합니다.

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

다음 명령어는 XPK를 사용하여 자체 호스팅 Jupyter 노트북 인스턴스를 만드는 방법을 보여줍니다.

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"

다음 YAML은 kubectl을 사용하여 자체 호스팅 Jupyter 노트북 인스턴스를 만드는 방법을 보여줍니다. 헤드리스 Pathways 클러스터가 생성된 후 다음 YAML을 적용합니다. 자세한 내용은 kubectl을 사용하여 대화형 모드로 Pathways 실행을 참고하세요.

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

노트북 인스턴스로 이동

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

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

로컬 브라우저에서 http://localhost:8888?token=<var>your-token</var>으로 이동합니다. <your-token>을 Jupyter 노트북 컨테이너의 로그에 있는 토큰으로 바꿉니다.

kubectl logs ${MAIN_POD}

다음과 같이 출력됩니다.

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

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

1. Jupyterlab에서 새 Python 3 노트북을 만듭니다.
2. Pathways 프록시 서버에 연결

노트북에서 pathwaysutils를 설치하고 JAX_PLATFORMS를 proxy로 설정하고 JAX_BACKEND_TARGET를 PROXY_ADDRESS로 설정하는 셀을 추가합니다.

참고: PROXY_ADDRESS 값은 학습 프로그램 프록시 서버에 연결에서 선택한 연결 옵션에 따라 다릅니다.

!pip install pathwaysutils
%env JAX_PLATFORMS=proxy
# Replace your proxy address below:
%env JAX_BACKEND_TARGET=PROXY_ADDRESS

두 번째 셀을 'hello world' 유형 검사로 추가하고 Pathways 클러스터의 기기를 출력합니다.

import pathwaysutils
import jax

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

모든 것이 제대로 작동하면 Pathways-on-Cloud 백엔드가 감지되었음을 나타내는 메시지가 표시됩니다.

나열된 JAX 기기 수는 Pathways 클러스터를 만들 때 지정한 TPU 칩 수 및 슬라이스 수와 일치해야 합니다.

노트북에 코드 추가

자체 JAX 코드를 추가하고 Pathways 클러스터의 TPU에서 대화형으로 실행합니다. 다음 코드는 단일 노트북에서 두 슬라이스에 걸쳐 계산을 실행하는 방법을 보여줍니다.

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)

학습 여정 대화형 클러스터 삭제

XPK

xpk workload delete --workload=WORKLOAD --cluster=CLUSTER --project=PROJECT --zone=ZONE

kubectl

kubectl delete -f pathways-headless-workload.yaml

다음 단계

• 경로를 사용하여 GKE 클러스터 만들기
• Pathways를 사용한 멀티 호스트 추론
• 경로가 있는 일괄 워크로드
• 학습 과정 대화형 모드
• JAX 워크로드를 Pathways로 포팅
• 학습 과정으로 복원력 있는 학습
• 문제 해결 경로