Executar uma carga de trabalho interativa com os programas

As cargas de trabalho interativas do Pathways são cargas de trabalho remotas do JAX que são executadas fora dos componentes do controlador do Pathways do cluster, por exemplo, em uma VM. Ao contrário das cargas de trabalho em lote, quando uma carga de trabalho interativa é concluída, ela não desliga os componentes do cluster do Pathways, que permanecem disponíveis para conexões de outros clientes JAX. Neste documento, usamos um notebook do Jupyter como exemplo para demonstrar cargas de trabalho interativas. É possível executar cargas de trabalho interativas com o Pathways em um cluster do Google Kubernetes Engine (GKE) ou em um cluster do Cluster Director com o Slurm.

Esta página tem instruções para as duas plataformas. Em cada seção, selecione a guia GKE ou Cluster Director.

Usando a interface IFRT, os usuários do JAX enviam comandos para um cluster do Pathways. O código JAX, seja executado em um terminal, notebook ou qualquer ambiente compatível com Python, pode interagir com os recursos do Pathways.

Antes de começar

• Selecione a guia para como planeja usar as amostras nesta página:

  gcloud

  Instale a Google Cloud CLI e faça login na CLI gcloud com sua identidade federada. Depois de fazer login, inicialize a Google Cloud CLI executando o seguinte comando:

  gcloud init

  REST

  Para usar as amostras da API REST desta página em um ambiente de desenvolvimento local, use as credenciais fornecidas para a CLI gcloud.

  Instale a Google Cloud CLI e faça login na CLI gcloud com sua identidade federada.

  Saiba mais em Autenticar para usar REST na documentação de autenticação do Google Cloud .

Antes de começar

GKE

Siga os seguintes pré-requisitos:

• Crie um cluster do GKE com o Pathways.
• Instale o Cluster Toolkit.
• Instale as ferramentas do Kubernetes.
• Instale a CLI do Google Cloud.
• Ative a API Kubernetes Engine.
• Verifique se o Pathways está ativado para o projeto do Google Cloud .

Cluster Director

Siga os seguintes pré-requisitos:

• Crie um cluster do Cluster Director e configure o ambiente do Pathways.
• Instale a CLI gcloud.
• Verifique se você tem acesso SSH ao nó de login do cluster.
• Instale o Python 3 com o módulo venv no host que executa o cliente JAX. No Debian e no Ubuntu, instale o pacote python3-venv.

Executar o Pathways no modo interativo

GKE

1. Configure as variáveis de ambiente a seguir:

   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

   Substitua:

   • WORKLOAD: um nome exclusivo para identificar sua carga de trabalho (22 caracteres ou menos devido aos limites de rótulo de DNS).
   • WORKLOAD_NODEPOOL_COUNT: o número de pools de nós usados por uma carga de trabalho do Pathways.
   • COMPUTE_TYPE: o tipo de máquina da TPU, por exemplo, ct6e-standard-4t. Para mais informações sobre os tipos de TPU compatíveis com cada versão de TPU, consulte Versões de TPU.
   • TOPOLOGY: a topologia de posicionamento da TPU, por exemplo, 2x4.
   • PROJECT_ID: o ID do projeto doGoogle Cloud .
   • REGION: a região em que o bucket e o cluster estão localizados.
   • ZONE: a zona em que você planeja executar a carga de trabalho.
   • CLUSTER: o nome do cluster do GKE.
   • BUCKET_NAME: o bucket do Cloud Storage para artefatos do Pathways.

2. Se você ainda não criou um bucket do Cloud Storage para artefatos do Pathways, crie um:

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

3. Crie os contêineres do Pathways no cluster. Para executar uma carga de trabalho headless, use o comando gcluster job submit (para mais informações, consulte o Guia de envio de jobs do 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
   

   O comando gcluster job submit procura um pool de nós de CPU chamado cpu-np ou pathways-np para executar o pod principal do Pathways. Se o pool de nós de CPU tiver um nome diferente (como cpu-pathways-np do guia de criação de cluster), adicione a flag --pathways-head-np=NODE_POOL_NAME aos comandos gcluster job submit nesta página. Caso contrário, o comando falhará com o erro failed to auto-detect a suitable CPU node pool.

Neste ponto, sua carga de trabalho do JAX pode se conectar ao servidor proxy do IFRT.

Cluster Director

1. Conecte-se ao nó de login do cluster:

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

   Substitua:

   • CLUSTER_ID: o nome do cluster do Cluster Director.
   • PROJECT_ID: o ID do projeto doGoogle Cloud .
   • ZONE: a zona em que o cluster está localizado.

2. Configure as variáveis de ambiente a seguir:

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

   Substitua:

   • TOPOLOGY: a topologia de TPU de cada fração, por exemplo, 2x4. Ela precisa corresponder à topologia do nodeset de TPU.
   • NUM_SLICES: o número de frações de TPU. Defina MIN_SLICES e MAX_SLICES com o mesmo valor para que o cluster do Pathways esteja pronto apenas quando todas as frações estiverem em execução.
   • BUCKET_NAME: um bucket do Cloud Storage para artefatos do Pathways.

   Se a partição de TPU não for a padrão do cluster, defina PARTITION como o nome dela.

3. Inicie o cluster do Pathways no modo headless:

   cd ~/pathways
   sbatch start_pathways.sh
   

4. Aguarde até que o registro do serviço, pathways-service-JOB_ID.out (em que JOB_ID é o ID da tarefa do Slurm), mostre uma saída semelhante a esta:

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

   RM_IP é o endereço IP interno do servidor proxy. O arquivo ~/pathways/cluster.env contém o mesmo endereço.

Neste ponto, sua carga de trabalho do JAX pode se conectar ao servidor proxy do IFRT.

Conectar-se ao cluster do Pathways no modo interativo

É possível se conectar ao cluster do Pathways com ou sem encaminhamento de portas. Use uma das seções a seguir para se conectar ao cluster do Pathways.

No Cluster Director, o cliente JAX precisa usar a mesma versão do JAX que o servidor Pathways, 0.11.1.

Conectar usando o encaminhamento de portas

GKE

Em qualquer host com acesso ao plano de controle do cluster, encaminhe a porta local 29000 para o pod principal do Pathways para acessar o servidor proxy:

PROXY_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
PROXY_PORT=29000
kubectl port-forward ${PROXY_POD} ${PROXY_PORT}:${PROXY_PORT}

O resultado será o seguinte:

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

No mesmo host, abra uma nova janela do terminal. Usando o Python 3.11 ou mais recente, crie um ambiente virtual, instale os pacotes pathwaysutils e jax, defina as variáveis de ambiente JAX_PLATFORMS e JAX_BACKEND_TARGET e execute um script em Python que importe pathwaysutils e jax:

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

O resultado será o seguinte:

[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

Em qualquer host que possa se conectar por SSH ao nó de login do cluster, como sua estação de trabalho ou uma VM fora da sub-rede do cluster, encaminhe a porta do proxy pelo nó de login:

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

Observação: se você executar a CLI gcloud em uma VM do Compute Engine com a conta de serviço padrão, a VM precisará do escopo de acesso cloud-platform. Caso contrário, gcloud compute ssh falhará com insufficient authentication scopes. Como alternativa, execute gcloud auth login na VM para usar sua própria conta.

No mesmo host, abra uma nova janela do terminal. Defina as variáveis de ambiente JAX_PLATFORMS e JAX_BACKEND_TARGET e execute um script Python que importe pathwaysutils e jax. O host não tem rota para RM_IP, então o cliente se conecta à extremidade local do túnel, 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())'

O resultado será o seguinte:

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

O aviso IFRT proxy server disconnected aparece quando o script é encerrado e fecha a conexão. Isso é esperado.

Observação: cada chamada do JAX passa pelo túnel SSH. Use o encaminhamento de portas para verificar a conectividade. Para notebooks, execute o cliente JAX na sub-rede do cluster.

Conectar-se de hosts na VPC sem usar o encaminhamento de portas

GKE

Se você não quiser usar o encaminhamento de portas, conecte-se ao cluster do Pathways usando o Cloud DNS ou um balanceador de carga interno.

Conectar usando o Cloud DNS

Ativar o Cloud DNS no cluster muda o provedor de DNS de kube-dns para o Cloud DNS. Quando ativado, o GKE cria uma zona de DNS particular do Cloud DNS na rede de nuvem privada virtual (VPC) para os nomes DNS. Para mais informações, consulte Como usar o Cloud DNS para GKE.

Se você ativar o Cloud DNS com escopo de VPC ou escopo de VPC aditivo, os nomes DNS do Kubernetes poderão ser resolvidos em VMs que não são do GKE na sua rede VPC. Com o escopo do cluster, apenas os nós no cluster podem resolver esses nomes. Os nomes têm o formato SERVICE_NAME.NAMESPACE.svc.CUSTOM_DNS_DOMAIN. O pod principal do Pathways tem um serviço chamado WORKLOAD-pathways-head-0-0.WORKLOAD.NAMESPACE.svc.CUSTOM_DNS_DOMAIN.

Os comandos a seguir mostram como se conectar ao cluster do Pathways usando o Cloud DNS:

1. Confirme se a entrada DNS do líder pode ser resolvida em um host que não seja do GKE:

   host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN

   Substitua:

   • WORKLOAD: o nome da carga de trabalho do Pathways.
   • CUSTOM_DNS_DOMAIN: o domínio DNS personalizado configurado para o cluster.

   O resultado será o seguinte:

   WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN has address 10.0.2.75

2. Conecte-se ao cluster do Pathways usando o nome DNS:

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

   O resultado será o seguinte:

   [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
   

Conectar usando um balanceador de carga interno

Para um endereço IP particular na rede VPC que aponta para a implantação do Pathways, crie um serviço apoiado por um balanceador de carga interno (ILB, na sigla em inglês). Isso não exige que o cluster tenha o Cloud DNS ativado.

Para clusters com muitas VMs, ative o subconjunto de ILB se você criar balanceadores de carga internos. Para mais informações, consulte Ativar a criação de subconjuntos do GKE em um cluster atual. Quando a subconfiguração de ILB não está ativada, todos os nós no cluster fazem parte do grupo de instâncias de back-end para todos os balanceadores de carga internos, que não escalonam além de 250 nós. Com o subagrupamento de ILB ativado, o GKE cria grupos de endpoints de rede em vez de grupos de instâncias e inclui apenas nós que executam um dos pods de exibição do serviço. A ativação do subconjunto de ILB tem uma latência de configuração única de cerca de 15 minutos. O comando a seguir ativa o subconjunto de ILB:

gcloud container clusters update ${CLUSTER} \
    --project=${PROJECT_ID} \
    --location=CLUSTER_LOCATION \
    --enable-l4-ilb-subsetting

Substitua CLUSTER_LOCATION pela região ou zona em que seu cluster está localizado.

Depois de ativar o subconjunto de ILB, crie um serviço do Kubernetes do tipo LoadBalancer usando o manifesto YAML a seguir. Isso faz com que o GKE crie um balanceador de carga interno na rede VPC do cluster:

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

Substitua WORKLOAD pelo nome da carga de trabalho do Pathways e salve o arquivo como pathways-headless-ilb.yaml.

Aplique o manifesto:

kubectl apply -f pathways-headless-ilb.yaml

Depois que o GKE criar o balanceador de carga (cerca de um minuto depois), verifique se a coluna EXTERNAL-IP mostra um endereço IP:

kubectl get services

O resultado será o seguinte:

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

É possível acessar a implantação do Pathways sem encaminhamento de portas de hosts na mesma rede VPC do cluster:

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

Substitua EXTERNAL_IP pelo valor de EXTERNAL-IP exibido por kubectl get services, por exemplo, 10.0.0.22.

O resultado será o seguinte:

[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

Observação: seus endereços IP são diferentes dos endereços mostrados nos exemplos de comando e saída anteriores.

Cluster Director

Os hosts na sub-rede do cluster, como uma VM do Compute Engine, podem se conectar diretamente ao servidor proxy. O firewall do Cluster Director permite o tráfego da sub-rede do cluster, mas não de outras sub-redes da mesma rede 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())'

Substitua RM_IP pelo endereço de proxy do IFRT do registro do serviço.

Observação: o servidor proxy do IFRT não autentica clientes. Qualquer host que possa alcançar RM_IP:29000 pode usar o cluster do Pathways.

Notebooks do Jupyter

GKE

É possível criar um notebook do Jupyter usando a Vertex AI ou criar um notebook do Jupyter auto-hospedado.

Criar uma instância do Vertex AI Workbench

Depois de configurar e verificar o cluster do Pathways, é possível acessar as VMs de TPU do GKE em um notebook do Jupyter do Vertex AI Workbench. As instruções de configuração a seguir pressupõem que o cluster do GKE Pathways esteja na mesma rede VPC (que é a rede padrão, a menos que você tenha configurado de outra forma).

1. Acesse a página do Vertex AI Workbench no console do Google Cloud .

2. Na guia Instâncias, clique em Criar nova e verifique se a rede corresponde à rede do cluster do GKE.

   Como alternativa, use a linha de comando para criar uma nova instância do 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

   Substitua:

   • INSTANCE_NAME: o nome da sua instância do Vertex AI Workbench.
   • PROJECT_ID: o ID do projeto doGoogle Cloud .
   • ZONE: a zona em que você quer criar a instância.
   • NETWORK: o nome da rede VPC do cluster do GKE. A flag --network é opcional. Se você omiti-la, o comando usará a rede default.

3. Depois que a Vertex AI criar a instância, acesse-a e clique em Abrir o JupyterLab.

Criar uma instância de notebook do Jupyter auto-hospedada

O comando a seguir mostra como criar uma instância de notebook do Jupyter auto-hospedada usando o comando gcluster job submit:

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"

O manifesto YAML a seguir mostra como criar uma instância de notebook do Jupyter auto-hospedada usando kubectl. Aplique o seguinte manifesto YAML depois de criar um cluster do Pathways headless. Para mais informações, consulte Executar o Pathways no modo interativo:

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

O campo hostNetwork: true permite que o cliente JAX no notebook conclua o handshake do Application Layer Transport Security (ALTS) com o servidor proxy do Pathways. Sem esse campo, a conexão falha em pools de nós que usam a Federação de Identidade da Carga de Trabalho para GKE. O campo dnsPolicy: ClusterFirstWithHostNet permite que o notebook resolva nomes de DNS no cluster, como o nome do pod principal do Pathways. Para mais informações, consulte Resolver problemas de Pathways no Cloud.

Cluster Director

É possível executar um notebook do Jupyter em uma VM do Compute Engine na sub-rede do cluster ou criar um notebook do Jupyter auto-hospedado no nó de login do cluster. Em ambos os casos, o notebook se conecta ao servidor proxy sem encaminhamento de portas.

Usar uma VM atual na sub-rede do cluster

Na VM, instale o JupyterLab em um ambiente virtual e inicie-o:

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

Criar uma instância de notebook do Jupyter auto-hospedada

No nó de login do cluster, instale o JupyterLab e inicie-o. Crie o ambiente virtual no diretório inicial com o Python do cluster. O Python do sistema do nó de login é muito antigo para as versões atuais do JAX. Portanto, o pip não consegue encontrar o 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

Conectar-se à instância de notebook

GKE

Conecte-se ao notebook pela máquina local usando o encaminhamento de portas:

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

No navegador local, acesse http://localhost:8888?token=TOKEN. Substitua TOKEN pelo token dos registros do contêiner do notebook do Jupyter:

kubectl logs ${MAIN_POD}

O resultado será o seguinte:

...
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 é o nome do host do nó do GKE que executa o pod do notebook (por exemplo, gke-my-cluster-cpu-np-1a2b3c4d-ef56).

Cluster Director

Para um notebook em uma VM ou no nó de login, conecte-se pela máquina local usando o encaminhamento de portas para esse host. Por exemplo, para o nó de login:

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

No navegador local, acesse http://localhost:8888?token=TOKEN. Substitua TOKEN pelo token que o JupyterLab exibiu quando foi iniciado.

Conectividade do notebook com o cluster do Pathways

1. No JupyterLab, crie um notebook Python 3.
2. Conecte-se ao servidor proxy do Pathways.

No notebook, adicione uma célula para instalar pathwaysutils, defina JAX_PLATFORMS como proxy e defina JAX_BACKEND_TARGET como grpc://PROXY_ADDRESS:29000:

Observação: o valor de PROXY_ADDRESS depende da opção de conexão escolhida em Conectar-se ao cluster do Pathways no modo interativo. No Cluster Director, instale jax==0.11.1 na mesma célula. Em um notebook auto-hospedado que você iniciou após source ~/pathways/cluster.env, as variáveis já estão definidas.

!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

Substitua PROXY_ADDRESS pelo host ou endereço IP do servidor proxy do Pathways de Conectar-se ao cluster do Pathways no modo interativo (por exemplo, 127.0.0.1, o nome do host do Cloud DNS ou o endereço IP do balanceador de carga interno). No Cluster Director, use 127.0.0.1 se você usar o encaminhamento de portas ou o endereço de proxy do IFRT do registro do serviço. Se você implantou um notebook auto-hospedado usando kubectl, substitua PROXY_ADDRESS pelo nome DNS no cluster do pod principal do Pathways, para que o valor de JAX_BACKEND_TARGET seja grpc://WORKLOAD-pathways-head-0-0.WORKLOAD:29000.

Adicione uma segunda célula para verificar a conectividade e mostrar os dispositivos no cluster do Pathways:

import pathwaysutils
import jax

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

Se a conexão for bem-sucedida, a saída vai incluir uma mensagem indicando que o JAX detectou o back-end Pathways-on-Cloud.

O número de dispositivos JAX listados corresponde ao número de chips de TPU e ao número de frações que você especificou ao criar o cluster do Pathways.

Adicionar seu código a um notebook

Adicione seu próprio código JAX e execute-o de forma interativa nas TPUs no cluster do Pathways. O código a seguir mostra como realizar cálculos em duas frações de um único notebook. Para executar o exemplo completo, crie sua carga de trabalho do Pathways com pelo menos duas frações definindo WORKLOAD_NODEPOOL_COUNT como 2 ou mais. No Cluster Director, defina MIN_SLICES e MAX_SLICES como 2 ou mais. Com uma fração, slice1_devices fica vazio, e a última parte do exemplo falha com um 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)

Excluir a carga de trabalho interativa do Pathways

GKE

Cluster Toolkit

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

kubectl

kubectl delete job WORKLOAD

Cluster Director

No nó de login do cluster, cancele o job de serviço do Pathways:

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

Esse comando também cancela os jobs de worker do Pathways. As VMs de TPU são nós estáticos, então continuam em execução e ficam disponíveis para outros jobs do Slurm.

Se você tiver problemas, consulte Resolver problemas do Pathways no Cloud.

A seguir

• Criar um cluster do GKE com o Pathways
• Criar um cluster do Cluster Director com o Pathways
• Executar uma carga de trabalho em lote com o Pathways
• Portar cargas de trabalho do JAX para o Pathways
• Treinamento resiliente com o Pathways
• Resolver problemas de Pathways no Cloud