Executar o NCCL em clusters do GKE que usam A3 Mega ou A3 High

Esta página descreve como executar testes da Biblioteca de Comunicação Coletiva da NVIDIA (NCCL, na sigla em inglês) em clusters personalizados do GKE que usam A3 Mega e A3 High com protocolos de rede GPUDirect-TCPXO e GPUDirect-TCPX. Um cluster do GKE personalizado é um cluster que você cria usando comandos gcloud.

É possível usar os testes descritos nesta página para os seguintes cenários:

Antes de começar

Os testes nesta página usam JobSet e Kueue com programação consciente da topologia (TAS, na sigla em inglês). Antes de executar qualquer teste, configure o cluster e faça o seguinte:

  1. Instale o JobSet.

  2. Instale o Kueue.

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/v0.16.5/manifests.yaml
    

Configurar o cluster com o Jobset e o Kueue

Depois de instalar o JobSet e o Kueue, siga estas etapas:

  1. Salve o seguinte manifesto como um arquivo kueue-config.yaml:

    A3 Alto

      apiVersion: kueue.x-k8s.io/v1beta2
      kind: Topology
      metadata:
        name: "gke-default"
      spec:
        levels:
        - nodeLabel: "cloud.google.com/gce-topology-block"
        - nodeLabel: "cloud.google.com/gce-topology-subblock"
        - nodeLabel: "cloud.google.com/gce-topology-host"
        - nodeLabel: "kubernetes.io/hostname"
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: ResourceFlavor
      metadata:
        name: a3-high-flavor
      spec:
        nodeLabels:
          cloud.google.com/gke-accelerator: nvidia-h100-80gb
        topologyName: "gke-default"
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: ResourceFlavor
      metadata:
        name: a3-high-dws-flavor
      spec:
        nodeLabels:
          cloud.google.com/gke-accelerator: nvidia-h100-80gb
        topologyName: "gke-default"
        tolerations:
        - key: "cloud.google.com/gke-queued"
          operator: "Exists"
          effect: NoSchedule
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: AdmissionCheck
      metadata:
        name: dws-prov
      spec:
        controllerName: kueue.x-k8s.io/provisioning-request
        parameters:
          apiGroup: kueue.x-k8s.io
          kind: ProvisioningRequestConfig
          name: dws-config
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: ProvisioningRequestConfig
      metadata:
        name: dws-config
      spec:
        provisioningClassName: queued-provisioning.gke.io
        podSetUpdates:
        - key: autoscaling.gke.io/provisioning-request
        valueFromProvisioningClassDetail: ResizeRequestName
        managedResources:
        - nvidia.com/gpu
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: ClusterQueue
      metadata:
        name: cq-tas
      spec:
        namespaceSelector: {}
        clusterQueueingStrategy: BestEffortFIFO
        resourceGroups:
        - flavors:
          - name: a3-high-flavor
            resources:
            - name: "cpu"
              nominalQuota: 1000
            - name: "memory"
              nominalQuota: 1000Ti
            - name: "nvidia.com/gpu"
              nominalQuota: 1000
          - name: a3-high-dws-flavor
            resources:
            - name: "cpu"
              nominalQuota: 1000
            - name: "memory"
              nominalQuota: 1000Ti
            - name: "nvidia.com/gpu"
              nominalQuota: 1000
        admissionChecksStrategy:
          admissionChecks:
          - name: "dws-prov"
            onFlavors: [a3-high-dws-flavor]
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: LocalQueue
      metadata:
        namespace: default
        name: lq-tas
      spec:
        clusterQueue: cq-tas
    

    A3 Mega

      apiVersion: kueue.x-k8s.io/v1beta2
      kind: Topology
      metadata:
        name: "gke-default"
      spec:
        levels:
        - nodeLabel: "cloud.google.com/gce-topology-block"
        - nodeLabel: "cloud.google.com/gce-topology-subblock"
        - nodeLabel: "cloud.google.com/gce-topology-host"
        - nodeLabel: "kubernetes.io/hostname"
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: ResourceFlavor
      metadata:
        name: a3-mega-flavor
      spec:
        nodeLabels:
          cloud.google.com/gke-accelerator: nvidia-h100-mega-80gb
        topologyName: "gke-default"
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: ResourceFlavor
      metadata:
        name: a3-mega-dws-flavor
      spec:
        nodeLabels:
          cloud.google.com/gke-accelerator: nvidia-h100-mega-80gb
        topologyName: "gke-default"
        tolerations:
        - key: "cloud.google.com/gke-queued"
          operator: "Exists"
          effect: NoSchedule
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: AdmissionCheck
      metadata:
        name: dws-prov
      spec:
        controllerName: kueue.x-k8s.io/provisioning-request
        parameters:
          apiGroup: kueue.x-k8s.io
          kind: ProvisioningRequestConfig
          name: dws-config
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: ProvisioningRequestConfig
      metadata:
        name: dws-config
      spec:
        provisioningClassName: queued-provisioning.gke.io
        podSetUpdates:
        - key: autoscaling.gke.io/provisioning-request
        valueFromProvisioningClassDetail: ResizeRequestName
        managedResources:
        - nvidia.com/gpu
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: ClusterQueue
      metadata:
        name: cq-tas
      spec:
        namespaceSelector: {}
        clusterQueueingStrategy: BestEffortFIFO
        resourceGroups:
        - flavors:
          - name: a3-mega-flavor
            resources:
            - name: "cpu"
              nominalQuota: 1000
            - name: "memory"
              nominalQuota: 1000Ti
            - name: "nvidia.com/gpu"
              nominalQuota: 1000
          - name: a3-mega-dws-flavor
            resources:
            - name: "cpu"
              nominalQuota: 1000
            - name: "memory"
              nominalQuota: 1000Ti
            - name: "nvidia.com/gpu"
              nominalQuota: 1000
        admissionChecksStrategy:
          admissionChecks:
          - name: "dws-prov"
            onFlavors: [a3-mega-dws-flavor]
      ---
      apiVersion: kueue.x-k8s.io/v1beta2
      kind: LocalQueue
      metadata:
        namespace: default
        name: lq-tas
      spec:
        clusterQueue: cq-tas
    

  2. Aplique o manifesto:

    kubectl apply -f kueue-config.yaml
    

Ao executar cargas de trabalho com a programação consciente da topologia (TAS) ativada, é possível especificar a rigidez com que as restrições de topologia são aplicadas usando uma das seguintes anotações no manifesto da carga de trabalho:

  • kueue.x-k8s.io/podset-required-topology: se você usar essa anotação, o Kueue vai bloquear o agendamento até que a carga de trabalho possa ser agendada dentro da restrição de topologia solicitada. Use essa anotação para garantir que os pods sejam colocados juntos para otimizar a performance.

  • kueue.x-k8s.io/podset-preferred-topology: se você usar essa anotação, o Kueue vai tentar agendar pods dentro da restrição de topologia solicitada, mas, se isso não for possível, ele vai admitir a carga de trabalho sem atender às restrições de topologia.

Para qualquer anotação, especifique um dos seguintes valores como a restrição de topologia:

  • cloud.google.com/gce-topology-block: agenda pods no mesmo bloco de rede.
  • cloud.google.com/gce-topology-subblock: agenda pods no mesmo rack.
  • cloud.google.com/gce-topology-host: agenda pods no mesmo host físico.

Teste em dois nós de início flexível

Para executar testes de NCCL em um cluster do GKE que usa VMs de início flexível A3 Mega ou A3 High, use o procedimento a seguir. Esse procedimento usa um JobSet para executar um teste de NCCL em dois nós.

  1. Salve o seguinte manifesto como um arquivo nccl-tas-jobset.yaml:

    A3 Alto

      apiVersion: v1
      kind: ConfigMap
      metadata:
        name: nccl-config
            data:
              allgather.sh: |
                #!/bin/bash
                for script in /configs/*; do
                  name=$(basename $script)
                  cp $script "/scripts/$name"
                  chmod +x "/scripts/$name"
                done
                /scripts/init_ssh.sh ${@};
                pushd /scripts;
                /scripts/gen_hostfiles.sh ${@};
                popd;
                /scripts/run-allgather.sh 8 eth1,eth2,eth3,eth4 1M 512M ${#};
      ---
      apiVersion: jobset.x-k8s.io/v1alpha2
      kind: JobSet
      metadata:
        name: nccl-tas-test
        labels:
          kueue.x-k8s.io/queue-name: lq-tas
      spec:
        suspend: true
        network:
          enableDNSHostnames: true
        replicatedJobs:
        - name: worker
          replicas: 2
          template:
            spec:
              parallelism: 1
              completions: 1
              template:
                metadata:
                  annotations:
                    kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"
                    networking.gke.io/default-interface: 'eth0'
                    networking.gke.io/interfaces: |
                            [
                              {"interfaceName":"eth0","network":"default"},
                              {"interfaceName":"eth1","network":"vpc0"},
                              {"interfaceName":"eth2","network":"vpc1"},
                              {"interfaceName":"eth3","network":"vpc2"},
                              {"interfaceName":"eth4","network":"vpc3"}
                            ]
                spec:
                  terminationGracePeriodSeconds: 0
                  nodeSelector:
                    cloud.google.com/gke-accelerator: nvidia-h100-80gb
                  tolerations:
                  - key: cloud.google.com/gke-queued
                    effect: NoSchedule
                    value: "true"
                  - key: "nvidia.com/gpu"
                    operator: "Exists"
                    effect: "NoSchedule"
                  setHostnameAsFQDN: true
                  containers:
                  - name: tcpx-daemon
                    image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd-dev:v2.0.11
                    command:
                      - /tcpgpudmarxd/build/app/tcpgpudmarxd
                      - --gpu_nic_preset
                      - a3vm
                      - --gpu_shmem_type
                      - fd
                      - --uds_path
                      - /run/tcpx
                      - --setup_param
                      - "--verbose 128 2 0 "
                    securityContext:
                      privileged: true
                      capabilities:
                        add:
                          - NET_ADMIN
                    volumeMounts:
                      - name: libraries
                        mountPath: /usr/local/nvidia/lib64
                      - name: tcpx-socket
                        mountPath: /run/tcpx
                      - name: sys
                        mountPath: /hostsysfs
                      - name: proc-sys
                        mountPath: /hostprocsysfs
                    env:
                      - name: LD_LIBRARY_PATH
                        value: /usr/local/nvidia/lib64
                  - name: nccl-test
                    image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx-dev:v3.1.8
                    command:
                      - bash
                      - -c
                      - |
                        /scripts/container_entry.sh daemon;
                        sleep infinity;
                    securityContext:
                      privileged: true
                    volumeMounts:
                      - name: tcpx-socket
                        mountPath: /tmp
                      - name: libraries
                        mountPath: /usr/local/nvidia/lib64
                      - name: nccl-config
                        mountPath: /configs
                      - name: shared-memory
                        mountPath: /dev/shm
                    resources:
                      limits:
                        cpu: "200"
                        memory: "1800Gi"
                        nvidia.com/gpu: 8
                      requests:
                        cpu: "200"
                        memory: "1800Gi"
                        nvidia.com/gpu: 8
                  volumes:
                  - name: libraries
                    hostPath:
                      path: /home/kubernetes/bin/nvidia/lib64
                  - name: tcpx-socket
                    emptyDir: {}
                  - name: sys
                    hostPath:
                      path: /sys
                  - name: proc-sys
                    hostPath:
                      path: /proc/sys
                  - name: shared-memory
                    emptyDir:
                      medium: Memory
                      sizeLimit: 250Gi
                  - name: nccl-config
                    configMap:
                      name: nccl-config
                      defaultMode: 0777
    

    A3 Mega

      apiVersion: v1
      kind: ConfigMap
      metadata:
        name: nccl-configmap
            data:
              allgather.sh: |
                #!/bin/bash
                service ssh restart;
                /scripts/init_ssh.sh ${@};
                pushd /scripts;
                /scripts/gen_hostfiles.sh ${@};
                popd;
                # Set up environment variables for GPUDirect-TCPXO
                export LD_LIBRARY_PATH=/usr/local/nvidia/lib64
                export NCCL_FASTRAK_CTRL_DEV=eth0
                export NCCL_FASTRAK_IFNAME=eth1,eth2,eth3,eth4,eth5,eth6,eth7,eth8
                export NCCL_SOCKET_IFNAME=eth0
                export NCCL_CROSS_NIC=0
                export NCCL_ALGO=Ring,Tree
                export NCCL_PROTO=Simple
                export NCCL_NET_GDR_LEVEL=PIX
                # Run the benchmark
                /scripts/demo-run-nccl-test-tcpxo-via-mpi.sh
      ---
      apiVersion: jobset.x-k8s.io/v1alpha2
      kind: JobSet
      metadata:
        name: nccl-tas-test
        labels:
          kueue.x-k8s.io/queue-name: lq-tas
      spec:
        ttlSecondsAfterFinished: 1200
        suspend: true
        network:
          enableDNSHostnames: true
        replicatedJobs:
          - name: worker
            replicas: 2
            template:
              spec:
                parallelism: 1
                completions: 1
                template:
                  metadata:
                    annotations:
                      kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"
                      networking.gke.io/default-interface: 'eth0'
                      networking.gke.io/interfaces: |
                        [
                          {"interfaceName":"eth0","network":"default"},
                          {"interfaceName":"eth1","network":"vpc0"},
                          {"interfaceName":"eth2","network":"vpc1"},
                          {"interfaceName":"eth3","network":"vpc2"},
                          {"interfaceName":"eth4","network":"vpc3"},
                          {"interfaceName":"eth5","network":"vpc4"},
                          {"interfaceName":"eth6","network":"vpc5"},
                          {"interfaceName":"eth7","network":"vpc6"},
                          {"interfaceName":"eth8","network":"vpc7"}
                        ]
                  spec:
                    activeDeadlineSeconds: 3600
                    restartPolicy: Never
                    nodeSelector:
                      cloud.google.com/gke-accelerator: nvidia-h100-mega-80gb
                    tolerations:
                    - key: cloud.google.com/gke-queued
                      effect: NoSchedule
                      value: "true"
                    - key: "nvidia.com/gpu"
                      operator: "Exists"
                      effect: "NoSchedule"
                    setHostnameAsFQDN: true
                    volumes:
                    - name: nvidia
                      hostPath:
                        path: /home/kubernetes/bin/nvidia
                    - name: lib64
                      hostPath:
                        path: /lib64
                    - name: proc
                      hostPath:
                        path: /proc
                    - name: shared-memory
                      emptyDir:
                        medium: "Memory"
                        sizeLimit: 250Gi
                    - name: nccl-config
                      configMap:
                        name: nccl-configmap
                        defaultMode: 0755
                    containers:
                    - name: nccl-test
                      image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/nccl-plugin-gpudirecttcpx-dev:v1.0.15
                      stdin: true
                      tty: true
                      securityContext:
                        privileged: true
                      env:
                      - name: LD_LIBRARY_PATH
                        value: /usr/local/nvidia/lib64
                      volumeMounts:
                      - name: nvidia
                        mountPath: /usr/local/nvidia
                      - name: shared-memory
                        mountPath: /dev/shm
                      - name: nccl-config
                        mountPath: /configs
                      resources:
                        limits:
                          cpu: "200"
                          memory: "3700Gi"
                          nvidia.com/gpu: 8
                        requests:
                          cpu: "200"
                          memory: "3700Gi"
                          nvidia.com/gpu: 8
                    - name: tcpxo-daemon
                      image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/tcpgpudmarxd-dev:v1.0.21
                      imagePullPolicy: Always
                      command: ["/bin/sh", "-c"]
                      args:
                        - |
                          set -ex
                          chmod 755 /fts/entrypoint_rxdm_container.sh
                          /fts/entrypoint_rxdm_container.sh --num_hops=2 --num_nics=8 --uid= --alsologtostderr
                      securityContext:
                        privileged: true
                        capabilities:
                          add:
                            - NET_ADMIN
                            - NET_BIND_SERVICE
                      volumeMounts:
                      - name: nvidia
                        mountPath: /usr/local/nvidia/lib64
                      - name: proc
                        mountPath: /proc
                      env:
                      - name: LD_LIBRARY_PATH
                        value: /usr/local/nvidia/lib64
    

  2. Aplique o manifesto ao cluster:

    kubectl apply -f nccl-tas-jobset.yaml
    
  3. Verifique se o JobSet foi admitido e está em execução:

    kubectl get jobset nccl-tas-test
    

    Aguarde até que o JobSet seja retomado e os pods atinjam o status Running.

  4. Acione o teste de NCCL executando o script allgather.sh do primeiro pod de trabalho:

    kubectl exec --stdin --tty --container=nccl-test nccl-tas-test-worker-0-0 -- /configs/allgather.sh nccl-tas-test-worker-0-0 nccl-tas-test-worker-1-0
    

    A saída de um teste de dois nós é semelhante a esta:

    A3 Alto

      #                                                              out-of-place                       in-place
      #       size         count      type   redop    root     time   algbw   busbw #wrong     time   algbw   busbw #wrong
      #        (B)    (elements)                               (us)  (GB/s)  (GB/s)            (us)  (GB/s)  (GB/s)
          1048576         16384     float    none      -1    696.8    1.50    1.41      0    729.0    1.44    1.35      0
          ...
          536870912       8388608     float    none      -1   7101.7   75.60   70.87      0   7060.9   76.03   71.28      0
      # Out of bounds values : 0 OK
      # Avg bus bandwidth    : 29.8293
    

    A3 Mega

      #                                                              out-of-place                       in-place
      #        size         count      type   redop    root     time   algbw   busbw #wrong     time   algbw   busbw #wrong
      #         (B)    (elements)                               (us)  (GB/s)  (GB/s)            (us)  (GB/s)  (GB/s)
          0                 0         float    none      -1     0.24    0.00    0.00      0     0.18    0.00    0.00      0
          ...
          8589934592     134217728    float    none      -1    42603  201.63  189.03      0    42670  201.31  188.73      0
      # Out of bounds values : 0 OK
      # Avg bus bandwidth    : 45.7587
    

Implantar uma carga de trabalho de teste de NCCL com TAS

Se você tiver mais de dois nós, recomendamos usar o teste a seguir, que usa a programação consciente da topologia (TAS). Para executar testes de NCCL com TAS em um cluster do GKE que usa VMs de início flexível A3 Mega ou A3 High, use o procedimento a seguir.

  1. Salve o seguinte manifesto como um arquivo nccl-jobset-test.yaml. Substitua NUM_NODES pelo número de nós no pool de nós:

    A3 Alto

        apiVersion: jobset.x-k8s.io/v1alpha2
        kind: JobSet
        metadata:
          name: nccl-ag
          labels:
            kueue.x-k8s.io/queue-name: lq-tas
        spec:
          ttlSecondsAfterFinished: 1200
          suspend: true
          network:
            enableDNSHostnames: true
          replicatedJobs:
            - name: worker
              template:
                spec:
                  parallelism: NUM_NODES
                  completions: NUM_NODES
                  template:
                    metadata:
                      annotations:
                        kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-subblock"
                        networking.gke.io/default-interface: 'eth0'
                        networking.gke.io/interfaces: |
                            [
                              {"interfaceName":"eth0","network":"default"},
                              {"interfaceName":"eth1","network":"vpc0"},
                              {"interfaceName":"eth2","network":"vpc1"},
                              {"interfaceName":"eth3","network":"vpc2"},
                              {"interfaceName":"eth4","network":"vpc3"}
                            ]
                    spec:
                      activeDeadlineSeconds: 3600
                      restartPolicy: Never
                      nodeSelector:
                        cloud.google.com/gke-accelerator: nvidia-h100-80gb
                      tolerations:
                      - key: cloud.google.com/gke-queued
                        operator: "Exists"
                        effect: NoSchedule
                      - key: "nvidia.com/gpu"
                        operator: "Exists"
                        effect: "NoSchedule"
                      setHostnameAsFQDN: true
                      volumes:
                      - name: proc
                        hostPath:
                          path: /proc
                      - name: nvidia
                        hostPath:
                          path: /home/kubernetes/bin/nvidia
                      - name: libraries
                        hostPath:
                          path: /home/kubernetes/bin/nvidia/lib64
                      - name: tcpx-socket
                        emptyDir: {}
                      - name: shared-memory
                        emptyDir:
                          medium: "Memory"
                          sizeLimit: 250Gi
                      containers:
                      - name: tcpx-daemon
                        image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd-dev:v2.0.11
                        command:
                          - /tcpgpudmarxd/build/app/tcpgpudmarxd
                          - --gpu_nic_preset
                          - a3vm
                          - --gpu_shmem_type
                          - fd
                          - --uds_path
                          - /run/tcpx
                          - --setup_param
                          - "--verbose 128 2 0 "
                        securityContext:
                          privileged: true
                        volumeMounts:
                          - name: tcpx-socket
                            mountPath: /run/tcpx
                          - name: libraries
                            mountPath: /usr/local/nvidia/lib64
                      - name: nccl-test
                        stdin: true
                        tty: true
                        image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx-dev:v3.1.8
                        securityContext:
                          privileged: true
                        env:
                        - name: MY_NODE_NAME
                          valueFrom:
                            fieldRef:
                              fieldPath: spec.nodeName
                        - name: OMPI_ALLOW_RUN_AS_ROOT
                          value: "1"
                        - name: OMPI_ALLOW_RUN_AS_ROOT_CONFIRM
                          value: "1"
                        - name: N_NODES
                          value: "NUM_NODES"
                        - name: LD_LIBRARY_PATH
                          value: /usr/local/nvidia/lib64
                        command:
                        - bash
                        - -c
                        - |
                          /scripts/container_entry.sh daemon &
                          export POSTFIX=$(hostname | cut -d . -f 2-)
                          export WORKERS_BASENAME=$(hostname | cut -d . -f 1 | rev | cut -d - -f 2- | rev )
                          export NODE_RANK=$JOB_COMPLETION_INDEX
                          for i in `seq 0 $(($N_NODES-1))`; do
                            OTHER=${WORKERS_BASENAME}-${i}.${POSTFIX}
                            until ssh -p 222 -o StrictHostKeyChecking=no $OTHER hostname; do
                              sleep 10
                            done
                            echo ${OTHER} port=222 slots=8 | tee -a /tmp/hostfile;
                          done
                          if [[ "${NODE_RANK}" -eq "0" ]]; then
                              /scripts/run-allgather.sh 8 eth1,eth2,eth3,eth4 1M 512M ${N_NODES}
                          else
                              while ping -c 1 ${WORKERS_BASENAME}-0.${POSTFIX}; do
                              sleep 5
                          done
                          fi
                          exit 0
                        volumeMounts:
                        - name: nvidia
                          mountPath: /usr/local/nvidia
                        - name: tcpx-socket
                          mountPath: /tmp
                        - name: libraries
                          mountPath: /usr/local/nvidia/lib64
                        - name: shared-memory
                          mountPath: /dev/shm
                        resources:
                          limits:
                            cpu: "200"
                            memory: "1800Gi"
                            nvidia.com/gpu: 8
                          requests:
                            cpu: "200"
                            memory: "1800Gi"
                            nvidia.com/gpu: 8
    

    A3 Mega

        apiVersion: jobset.x-k8s.io/v1alpha2
        kind: JobSet
        metadata:
          name: nccl-ag
          labels:
            kueue.x-k8s.io/queue-name: lq-tas
        spec:
          ttlSecondsAfterFinished: 1200
          suspend: true
          network:
            enableDNSHostnames: true
          replicatedJobs:
            - name: worker
              template:
                spec:
                  parallelism: NUM_NODES
                  completions: NUM_NODES
                  template:
                    metadata:
                      annotations:
                        kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-subblock"
                        networking.gke.io/default-interface: 'eth0'
                        networking.gke.io/interfaces: |
                          [
                            {"interfaceName":"eth0","network":"default"},
                            {"interfaceName":"eth1","network":"vpc0"},
                            {"interfaceName":"eth2","network":"vpc1"},
                            {"interfaceName":"eth3","network":"vpc2"},
                            {"interfaceName":"eth4","network":"vpc3"},
                            {"interfaceName":"eth5","network":"vpc4"},
                            {"interfaceName":"eth6","network":"vpc5"},
                            {"interfaceName":"eth7","network":"vpc6"},
                            {"interfaceName":"eth8","network":"vpc7"}
                          ]
                    spec:
                      activeDeadlineSeconds: 3600
                      restartPolicy: Never
                      nodeSelector:
                        cloud.google.com/gke-accelerator: nvidia-h100-mega-80gb
                      tolerations:
                      - key: cloud.google.com/gke-queued
                        operator: "Exists"
                        effect: NoSchedule
                      - key: "nvidia.com/gpu"
                        operator: "Exists"
                        effect: "NoSchedule"
                      setHostnameAsFQDN: true
                      volumes:
                      - name: proc
                        hostPath:
                          path: /proc
                      - name: nvidia
                        hostPath:
                          path: /home/kubernetes/bin/nvidia
                      - name: lib64
                        hostPath:
                          path: /lib64
                      - name: shared-memory
                        emptyDir:
                          medium: "Memory"
                          sizeLimit: 250Gi
                      containers:
                      - name: nccl-test
                        stdin: true
                        tty: true
                        image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/nccl-plugin-tcpxo-diagnostic:v1.0.6
                        securityContext:
                          privileged: true
                        env:
                        - name: MY_NODE_NAME
                          valueFrom:
                            fieldRef:
                              fieldPath: spec.nodeName
                        - name: OMPI_ALLOW_RUN_AS_ROOT
                          value: "1"
                        - name: OMPI_ALLOW_RUN_AS_ROOT_CONFIRM
                          value: "1"
                        - name: N_NODES
                          value: "NUM_NODES"
                        - name: NCCL_SOCKET_IFNAME
                          value: eth0
                        - name: NCCL_FASTRAK_CTRL_DEV
                          value: eth0
                        - name: NCCL_FASTRAK_IFNAME
                          value: eth1,eth2,eth3,eth4,eth5,eth6,eth7,eth8
                        - name: NCCL_CROSS_NIC
                          value: "0"
                        - name: NCCL_ALGO
                          value: Ring,Tree
                        - name: NCCL_PROTO
                          value: Simple
                        - name: NCCL_NET_GDR_LEVEL
                          value: PIX
                        - name: LD_LIBRARY_PATH
                          value: /usr/local/nvidia/lib64
                        command:
                        - bash
                        - -c
                        - |
                          set -x
                          /scripts/container_entry.sh daemon &
                          export POSTFIX=$(hostname | cut -d . -f 2-)
                          export WORKERS_BASENAME=$(hostname | cut -d . -f 1 | rev | cut -d - -f 2- | rev )
                          export NODE_RANK=$JOB_COMPLETION_INDEX
                          for i in `seq 0 $(($N_NODES-1))`; do
                            OTHER=${WORKERS_BASENAME}-${i}.${POSTFIX}
                            until ssh -p 222 -o StrictHostKeyChecking=no $OTHER hostname; do
                              sleep 10
                            done
                            echo ${OTHER} port=222 slots=8 | tee -a /tmp/hostfile;
                          done
                          if [[ "${NODE_RANK}" -eq "0" ]]; then
                              export NCCL_TESTS_SPLIT_MASK="0x0";
                              ENV_VARS=$(echo ${!NCCL*} ${!OMPI*} LD_LIBRARY_PATH PATH | sed 's/ / -x /g')
                              mpirun --hostfile /tmp/hostfile \
                                -x $ENV_VARS  \
                                -mca plm_rsh_no_tree_spawn 1 \
                                --mca orte_keep_fqdn_hostnames 1 \
                                --mca btl self,tcp \
                                --mca btl_tcp_if_include eth0 \
                                --bind-to none \
                                --mca plm_rsh_agent "ssh -q -o LogLevel=ERROR -o StrictHostKeyChecking=no -p 222" \
                                /third_party/nccl-tests/build/all_gather_perf -b 1K -e 8G -f 2 -g 1 -w 5 --iters 100 -c 1
                          else
                              while ping -c 1 ${WORKERS_BASENAME}-0.${POSTFIX}; do
                              sleep 5
                          done
                          fi
                          exit 0
                        volumeMounts:
                        - name: nvidia
                          mountPath: /usr/local/nvidia
                        - name: lib64
                          mountPath: /lib64
                        - name: shared-memory
                          mountPath: /dev/shm
                        resources:
                          limits:
                            cpu: "200"
                            memory: "3700Gi"
                            nvidia.com/gpu: 8
                          requests:
                            cpu: "200"
                            memory: "3700Gi"
                            nvidia.com/gpu: 8
                      - name: tcpxo-daemon
                        image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/tcpxo-daemon:v1.0.1
                        imagePullPolicy: Always
                        command:
                        - bash
                        - -c
                        - |
                          /usr/bin/tcpxo_daemon
                        securityContext:
                          privileged: true
                        volumeMounts:
                        - name: nvidia
                          mountPath: /usr/local/nvidia
                        - name: proc
                          mountPath: /proc
                        env:
                        - name: LD_LIBRARY_PATH
                          value: /usr/local/nvidia/lib64
    

  2. Aplique o manifesto:

    kubectl apply -f nccl-jobset-test.yaml
    
  3. Verifique se a carga de trabalho foi admitida e atingiu o estado Completed.

  4. Busque registros do pod correspondente a nccl-ag-worker-0-0-.* para conferir os resultados:

    kubectl logs $(kubectl get pods -o go-template='{{range .items}}{{.metadata.name}}{{"\n"}}{{end}}' | grep nccl-ag-worker-0-0)
    

A seguir