Menjalankan NCCL di cluster GKE yang menggunakan A3 Edge

Halaman ini menjelaskan cara menjalankan pengujian NVIDIA Collective Communications Library (NCCL) di cluster GKE kustom yang menggunakan protokol jaringan A3 Edge dengan GPUDirect-TCPX. Cluster GKE kustom adalah cluster yang Anda buat dengan menggunakan perintah gcloud.

Anda dapat menggunakan pengujian yang dijelaskan di halaman ini untuk skenario berikut:

Sebelum memulai

Pengujian di halaman ini menggunakan JobSet dan Kueue dengan Topology Aware Scheduling (TAS). Sebelum menjalankan pengujian apa pun, Anda harus menyiapkan cluster dan melakukan hal berikut:

  1. Instal JobSet.

  2. Instal Kueue.

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

Menyiapkan cluster dengan Jobset dan Kueue

Setelah Anda menginstal JobSet dan Kueue, lakukan langkah-langkah berikut:

  1. Simpan manifes berikut sebagai file kueue-config.yaml:

    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-edge-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-edge-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-edge-flavor
          resources:
          - name: "cpu"
            nominalQuota: 1000
          - name: "memory"
            nominalQuota: 1000Ti
          - name: "nvidia.com/gpu"
            nominalQuota: 1000
        - name: a3-edge-dws-flavor
          resources:
          - name: "cpu"
            nominalQuota: 1000
          - name: "memory"
            nominalQuota: 1000Ti
          - name: "nvidia.com/gpu"
            nominalQuota: 1000
      admissionChecksStrategy:
        admissionChecks:
        - name: "dws-prov"
          onFlavors: [a3-edge-dws-flavor]
    ---
    apiVersion: kueue.x-k8s.io/v1beta2
    kind: LocalQueue
    metadata:
      namespace: default
      name: lq-tas
    spec:
      clusterQueue: cq-tas
    

  2. Terapkan manifes:

    kubectl apply -f kueue-config.yaml
    

Saat menjalankan workload dengan Topology Aware Scheduling (TAS) diaktifkan, Anda dapat menentukan seberapa ketat batasan topologi diterapkan dengan menggunakan salah satu anotasi berikut dalam manifes workload:

  • kueue.x-k8s.io/podset-required-topology: Jika Anda menggunakan anotasi ini, Kueue akan memblokir penjadwalan hingga beban kerja dapat dijadwalkan dalam batasan topologi yang diminta. Gunakan anotasi ini untuk memastikan pod ditempatkan bersama untuk performa yang optimal.

  • kueue.x-k8s.io/podset-preferred-topology: Jika Anda menggunakan anotasi ini, Kueue akan mencoba menjadwalkan pod dalam batasan topologi yang diminta, tetapi jika hal itu tidak memungkinkan, Kueue akan mengizinkan workload tanpa memenuhi batasan topologi.

Untuk setiap anotasi, tentukan salah satu nilai berikut sebagai batasan topologi:

  • cloud.google.com/gce-topology-block: Menjadwalkan pod dalam blok jaringan yang sama.
  • cloud.google.com/gce-topology-subblock: Menjadwalkan pod dalam rak yang sama.
  • cloud.google.com/gce-topology-host: Menjadwalkan pod di host fisik yang sama.

Menguji pada dua node flex-start

Untuk menjalankan pengujian NCCL di cluster GKE yang menggunakan VM mulai fleksibel A3 Edge, gunakan prosedur berikut. Prosedur ini menggunakan manifes JobSet untuk menjalankan pengujian NCCL pada dua node.

  1. Simpan manifes berikut sebagai file nccl-tas-jobset.yaml:

    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
    

  2. Terapkan manifes ke cluster Anda:

    kubectl apply -f nccl-tas-jobset.yaml
    
  3. Pastikan JobSet diterima dan berjalan:

    kubectl get jobset nccl-tas-test
    

    Tunggu hingga JobSet dilanjutkan dan Pod mencapai status Running.

  4. Picu pengujian NCCL dengan menjalankan skrip allgather.sh dari Pod pekerja pertama:

    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
    

    Output untuk pengujian dua node mirip dengan berikut ini:

      #                                                              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
    

Men-deploy workload pengujian NCCL dengan TAS

Jika Anda memiliki lebih dari dua node, sebaiknya gunakan pengujian berikut, yang menggunakan (TAS). Untuk menjalankan pengujian NCCL dengan TAS di cluster GKE yang menggunakan VM dengan mulai fleksibel A3 Edge, gunakan prosedur berikut.

  1. Simpan manifes berikut sebagai file nccl-jobset-test.yaml. Ganti NUM_NODES dengan jumlah node di node pool:

    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
    

  2. Terapkan manifes:

    kubectl apply -f nccl-jobset-test.yaml
    
  3. Pastikan beban kerja diterima dan mencapai status Completed.

  4. Ambil log untuk Pod yang cocok dengan nccl-ag-worker-0-0-.* untuk melihat hasilnya:

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

Langkah berikutnya