Massimizza la larghezza di banda della rete della GPU con GPUDirect-TCPX e networking multiplo

Questa pagina mostra come massimizzare la larghezza di banda e la velocità effettiva della rete per carichi di lavoro GPU ad alte prestazioni nei cluster Google Kubernetes Engine (GKE) in modalità Standard. Questa pagina è rivolta agli ingegneri di machine learning (ML) e agli amministratori di piattaforma che facilitano i carichi di lavoro ML. Dovresti già avere familiarità con le tecnologie di networking come le schede di interfaccia di rete (NIC) e TCP, nonché con tecnologie di acceleratore come la NVIDIA Collective Communications Library (NCCL).

Le applicazioni di intelligenza artificiale (AI), ML e computing ad alte prestazioni (HPC) richiedono una potente accelerazione per ottimizzare le prestazioni riducendo i tempi di completamento dei job. Ad esempio, i modelli ML incentrati sull'AI conversazionale e sulla generazione di immagini richiedono un'elevata scalabilità e potenza di calcolo.

Informazioni sui supercomputer GPU Google Cloud

Google Cloud dispone di supercomputer ottimizzati per l'acceleratore, creati per modelli scalabili e di grandi dimensioni. Queste macchine offrono i seguenti vantaggi:

  • Otto GPU NVIDIA H100 per macchina.
  • Larghezza di banda fino a 200 Gbit/s sul NIC principale.
  • Fino a quattro NIC secondarie, ciascuna delle quali supporta una larghezza di banda fino a 200 Gbps per il trasferimento di dati GPU.

Per un elenco completo dei vantaggi, consulta Serie di macchine A3 nella documentazione di Compute Engine.

Il carico di lavoro GKE deve utilizzare tutte le GPU disponibili e tutte le NIC secondarie disponibili su un singolo nodo e utilizzare una parte significativa della larghezza di banda disponibile. La soluzione descritta in questo documento è ideale per carichi di lavoro che richiedono prestazioni elevate, velocità effettiva elevata e bassa latenza.

Funzionalità e caratteristiche richieste per massimizzare la larghezza di banda

Per massimizzare la larghezza di banda della rete nei nodi dei supercomputer GPU, utilizza tutte le seguenti funzionalità:

  • GPUDirect-TCPX: riduci l'overhead necessario per trasferire i payload dei pacchetti da e verso le GPU, il che migliora notevolmente la velocità effettiva su larga scala rispetto alle GPU che non utilizzano GPUDirect-TCPX.
  • gVNIC: attiva le funzionalità GPUDirect-TCPX, come la suddivisione dell'intestazione del pacchetto, il flusso di reindirizzamento e la gestione del buffer. Per utilizzare GPUDirect-TCPX, è richiesto gVNIC. Per maggiori dettagli su gVNIC, consulta Aumentare la velocità del traffico di rete per i nodi GPU.
  • Networking multiplo: aggiungi NIC secondarie alla macchina ottimizzata per l'acceleratore. Per i computer A3, vengono aggiunte altre quattro NIC. Ogni NIC è associato a una subnet separata nel proprio VPC per evitare conflitti. Per maggiori dettagli sul supporto di più reti, consulta Configurare il supporto di più reti per i pod.
  • Criteri di posizionamento: utilizza un criterio di posizionamento delle risorse per posizionare tutti i nodi GPU per un carico di lavoro specifico su server fisicamente chiusi per ridurre al minimo la latenza. Per maggiori dettagli, consulta Definire il posizionamento compatto per i nodi GKE.

Struttura della procedura

Per utilizzare tutte queste funzionalità insieme, procedi nel seguente modo:

  1. Crea Virtual Private Cloud (VPC) e subnet
  2. Crea l'ambiente GKE:
    1. Crea un cluster con più reti abilitate
    2. Crea un pool di nodi con le seguenti caratteristiche:
      1. gVNIC attivato
      2. Subnet di networking multiplo specificate per ogni NIC secondario
      3. Serie di macchine A3 con GPU H100 (quattro NIC secondarie e otto GPU) a supporto dei nodi
      4. I driver NVIDIA più recenti installati
  3. Installa il programma binario GPUDirect-TCPX e il plug-in NCCL
  4. Esegui il deployment di un carico di lavoro di test per verificare la configurazione di GPUDirect-TCPX

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti attività:

  • Abilita l'API Google Kubernetes Engine.
    • Abilita l'API Google Kubernetes Engine
  • Se vuoi utilizzare Google Cloud CLI per questa attività, installa e initialize gcloud CLI. Se hai già installato gcloud CLI, scarica la versione più recente eseguendo gcloud components update.
  • Assicurati di disporre di una quota sufficiente per le GPU H100. Per richiedere un aumento della quota, consulta la pagina Quote delle GPU.

Requisiti

  • GPUDirect-TCPX è supportato su GKE versione 1.27 o successive e richiede:
    • Per GKE versione 1.27, utilizza la patch GKE 1.27.7-gke.1121000 o successive.
    • Per GKE versione 1.28, utilizza la patch GKE 1.28.8-gke.1095000 o successive.
    • Per GKE versione 1.29, utilizza la patch GKE 1.29.3-gke.1093000 o successive.
  • I nodi GPU devono utilizzare il driver NVIDIA versione 535 o successiva.
  • Devi utilizzare GKE Dataplane V2.

Limitazioni

Si applicano le seguenti limitazioni:

  • Non puoi utilizzare GPUDirect-TCPX nei cluster Autopilot
  • Puoi utilizzare GPUDirect-TCPX solo su GKE versione 1.27 o successive e con le seguenti versioni patch:
    • Per GKE versione 1.27, utilizza la patch GKE 1.27.7-gke.1121000 o successive.
    • Per GKE versione 1.28, utilizza la patch GKE 1.28.8-gke.1095000 o successive.
    • Per GKE versione 1.29, utilizza la patch GKE 1.29.3-gke.1093000 o successive.
  • Non puoi utilizzare GPUDirect-TCPX con GPU multi-istanza o condivisione del tempo GPU.
  • Non puoi utilizzare NCCL FastSocket
  • Il tuo ambiente deve supportare l'impostazione hostNetwork: true nella specifica del pod

  • Per utilizzare gli SSD locali per l'archiviazione dei pod, devi specificare in modo esplicito il numero esatto di SSD locali da collegare alla VM A3 sottostante utilizzando il flag --ephemeral-storage-local-ssd=count=SSD_COUNT per l'archiviazione temporanea o il flag --local-nvme-ssd-block=count=SSD_COUNT per il blocco dell'accesso. Se ometti questo flag, non potrai utilizzare gli SSD locali nei tuoi pod. Questi flag sono obbligatori solo se vuoi utilizzare SSD locale per l'accesso ai dati.

    Le dimensioni della macchina supportate in GKE sono a3-highgpu-8g e il numero di SSD locali corrispondenti è 16.

Crea VPC e subnet

Crea reti VPC separate nel progetto per ogni NIC virtuale che aggiungerai ai nodi. Ogni VPC deve avere una subnet e una regola firewall che consenta il traffico di rete interno. Per massimizzare la larghezza di banda, consigliamo di creare quattro nuove reti.

  1. Aggiorna la subnet VPC predefinita nel tuo progetto per aggiungere intervalli di indirizzi IP secondari per i pod e i servizi:

    gcloud compute networks subnets update DEFAULT_NETWORK \
    --region=REGION \
    --add-secondary-ranges="CLUSTER_NAME-pods=POD_IP_ADDRESS_RANGE,CLUSTER_NAME-services=SERVICE_IP_ADDRESS_RANGE"

    Sostituisci quanto segue:

    • DEFAULT_NETWORK: il nome della subnet predefinita nel tuo progetto.
    • REGION: la regione della subnet predefinita.
    • CLUSTER_NAME: il nome del tuo cluster GKE.
    • POD_IP_ADDRESS_RANGE: l'intervallo di indirizzi IP in notazione CIDR per i pod nel cluster. Ad esempio, 10.64.0.0/19.
    • SERVICE_IP_ADDRESS_RANGE: l'intervallo di indirizzi IP da utilizzare per i servizi nel cluster, in notazione CIDR. Deve essere diverso dall'intervallo del pod. Ad esempio, 10.65.0.0/19.

  2. Crea le reti VPC per GPUDirect-TCPX nel tuo progetto, ciascuna con una subnet e una regola firewall:

    for N in $(seq 1 4); do
gcloud compute networks create PROJECT_ID-net-$N \
    --subnet-mode=custom \
    --mtu=8244

gcloud compute networks subnets create PROJECT_ID-sub-$N \
    --network=PROJECT_ID-net-$N \
    --region=REGION \
    --range=SUBNET_RANGE

gcloud compute firewall-rules create PROJECT_ID-internal-$N \
  --network=PROJECT_ID-net-$N \
  --action=ALLOW \
  --rules=tcp:0-65535,udp:0-65535,icmp \
  --source-ranges=SOURCE_RANGE
done

    Sostituisci quanto segue:

    • PROJECT_ID: il tuo ID progetto Google Cloud.
    • REGION: la regione di Compute Engine per ogni subnet.
    • SUBNET_RANGE: l'intervallo di indirizzi IP di ogni subnet in notazione CIDR. Questo comando di esempio viene ripetuto per quattro subnet, quindi utilizza una variabile per modificare l'indirizzo IP di ogni subnet. Ad esempio, specifica 192.168.$N.0/24 in modo che la prima subnet utilizzi 192.168.1.0/24, la seconda subnet utilizzi 192.168.2.0/24 e così via.
    • SOURCE_RANGE: l'intervallo di indirizzi IP di origine in notazione CIDR per consentire alla regola firewall di consentire il traffico in entrata. Ad esempio, 192.168.0.0/16.

  3. Verifica che le reti siano state create:

    gcloud compute networks list

Crea l'ambiente GKE

Crea un nuovo cluster GKE che utilizza il networking multiplo (anteprima) e crea un pool di nodi GPU che utilizzi macchine A3 con GPU H100 collegate e quattro NIC aggiuntive. Non puoi aggiornare un cluster esistente per utilizzare il networking multiplo.

  1. Crea un cluster:

    gcloud container clusters create CLUSTER_NAME \
    --location=LOCATION \
    --cluster-version=VERSION \
    --enable-dataplane-v2 --enable-ip-alias \
    --enable-multi-networking \
    --no-enable-autoupgrade \
    --cluster-secondary-range-name=CLUSTER_NAME-pods \
    --services-secondary-range-name=CLUSTER_NAME-services

    Sostituisci quanto segue:

    • CLUSTER_NAME: il nome del nuovo cluster
    • LOCATION: la regione di Compute Engine per il cluster
    • VERSION: la versione di GKE per il cluster. Deve essere una versione supportata come descritto nella sezione Requisiti.

    Questo comando specifica in modo esplicito anche l'indirizzo IP secondario per pod e servizi per il cluster che hai creato nella sezione precedente.

  2. Crea nel cluster risorse di rete e GKENetworkParamSet corrispondenti alle reti e alle subnet VPC che hai creato:

    kubectl apply -f - <<EOF
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: vpc1
spec:
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: vpc1
  type: Device
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: vpc2
spec:
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: vpc2
  type: Device
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: vpc3
spec:
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: vpc3
  type: Device
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: vpc4
spec:
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: vpc4
  type: Device
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: vpc1
spec:
  vpc: PROJECT_ID-net-1
  vpcSubnet: PROJECT_ID-sub-1
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: vpc2
spec:
  vpc: PROJECT_ID-net-2
  vpcSubnet: PROJECT_ID-sub-2
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: vpc3
spec:
  vpc: PROJECT_ID-net-3
  vpcSubnet: PROJECT_ID-sub-3
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: vpc4
spec:
  vpc: PROJECT_ID-net-4
  vpcSubnet: PROJECT_ID-sub-4
  deviceMode: NetDevice
EOF

    Queste risorse indicano a GKE di configurare le NIC per il traffico della GPU in modalità passthrough. GKE non applica a questo traffico la programmazione di networking integrata tramite eBPF.

  3. Crea un pool di nodi per le GPU H100:

    gcloud container node-pools create NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --location=LOCATION \
    --machine-type=a3-highgpu-8g \
    --accelerator=type=nvidia-h100-80gb,count=8,gpu-driver-version=LATEST \
    --additional-node-network=network=PROJECT_ID-net-1,subnetwork=PROJECT_ID-sub-1 \
    --additional-node-network=network=PROJECT_ID-net-2,subnetwork=PROJECT_ID-sub-2 \
    --additional-node-network=network=PROJECT_ID-net-3,subnetwork=PROJECT_ID-sub-3 \
    --additional-node-network=network=PROJECT_ID-net-4,subnetwork=PROJECT_ID-sub-4 \
    --enable-gvnic \
    --no-enable-autoupgrade \
    [--ephemeral-storage-local-ssd=count=16]

    Sostituisci NODE_POOL_NAME con il nome del pool di nodi.

    Se questo comando non riesce, la quota GPU H100 potrebbe non essere sufficiente nel progetto. Assicurati di avere una quota sufficiente e riprova a eseguire il comando.

  4. Recupera un elenco di nodi nel cluster:

    kubectl get nodes

  5. Verifica che ogni nodo GPU abbia otto GPU:

    kubectl describe node NODE_NAME

    L'output è simile al seguente:

    Capacity:
  ...
  nvidia.com/gpu:             8
Allocatable:
  ...
  nvidia.com/gpu:             8

Installa GPUDirect-TCPX e configura NCCL

Questa sezione mostra come installare il programma binario GPUDirect-TCPX e un determinato NCCL utilizzando un DaemonSet.

  1. Esamina il manifest del DaemonSet:

    apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: nccl-tcpx-installer
  namespace: kube-system
  labels:
    k8s-app: nccl-tcpx-installer
spec:
  selector:
    matchLabels:
      k8s-app: nccl-tcpx-installer
  updateStrategy:
    type: RollingUpdate
  template:
    metadata:
      labels:
        name: nccl-tcpx-installer
        k8s-app: nccl-tcpx-installer
    spec:
      priorityClassName: system-node-critical
      affinity:
        nodeAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            nodeSelectorTerms:
              - matchExpressions:
                  - key: cloud.google.com/gke-accelerator
                    operator: In
                    values:
                      - nvidia-h100-80gb
      tolerations:
        - operator: "Exists"
      hostNetwork: true
      hostPID: true
      volumes:
        - name: var-lib
          hostPath:
            path: /var/lib
        - name: tcpx
          hostPath:
            path: /var/lib/tcpx
        - name: library-dir-host
          hostPath:
            path: /home/kubernetes/bin
      initContainers:
        - image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx-dev:v3.1.9
          name: nccl-tcpx-installer
          resources:
            requests:
              cpu: 150m
          securityContext:
            privileged: true
          volumeMounts:
            - name: var-lib
              mountPath: /var/lib
            - name: library-dir-host
              mountPath: /usr/local
          command: ["/bin/sh", "-c"]
          args:
            - |
              set -ex
              /scripts/container_entry.sh install --install-nccl
              mkdir -p /usr/local/nvidia/lib64
              cp -r /var/lib/tcpx/lib64/. /usr/local/nvidia/lib64
              echo "installation finishes"
      containers:
        - image: "gcr.io/google-containers/pause:2.0"
          name: pause

    Questo DaemonSet svolge le seguenti operazioni:

    1. Installa una libreria NCCL e il programma binario GPUDirect-TCPX sul nodo.
    2. Archivia la libreria e il programma binario nella directory /home/kubernetes/bin/nvidia/lib64 sulla VM. Per impostazione predefinita, GKE monta questa directory nel percorso /usr/local/nvidia/lib64 nei container GPU che devono utilizzare NCCL e GPUDirect-TCPX.

  2. Esegui il deployment del DaemonSet:

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-tcpx-installer.yaml

    L'esecuzione del plug-in NCCL richiede circa due minuti.

  3. Verifica lo stato dei pod DaemonSet:

    kubectl get pods -n=kube-system -l=name=nccl-tcpx-installer

    L'output è simile al seguente:

    nccl-tcpx-installer-6c2pv                    1/1     Running   0          2m11s
nccl-tcpx-installer-qgg82                    1/1     Running   0          2m11s

Esegui il deployment di un carico di lavoro di test

In questa sezione eseguirai il deployment di un carico di lavoro di esempio per verificare che NCCL e GPUDirect-TCPX funzionino come previsto. Questo carico di lavoro include un container collaterale denominato tcpx-daemon, che esegue un servizio che consente al pod di utilizzare GPUDirect-TCPX. Devi aggiungere questo container collaterale a tutti i pod nel tuo ambiente che devono utilizzare GPUDirect-TCPX. Per uno snippet dei campi obbligatori da aggiungere ai manifest, consulta Aggiungere GPUDirect-TCPX al file manifest in questo documento.

  1. Esamina il manifest di ConfigMap di nccl-config-default.yaml in GitHub. Questo manifest esegue il deployment di scrip che inizializzano un test complessivo di NCCL e imposta le variabili di ambiente specifiche per NCCL.

  2. Esamina il manifest di nccl-test.yaml in GitHub. Il file manifest svolge le seguenti operazioni:

    1. Esegue il deployment di due pod, ognuno dei quali viene eseguito in un nodo con GPU H100.
    2. Esegue il deployment di un container collaterale denominato tcpx-daemon in ogni pod per consentire a questi pod di utilizzare GPUDirect-TCPX.

  3. Esegui il deployment dell'oggetto ConfigMap e del carico di lavoro di test:

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-config-default.yaml
kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-test.yaml

  4. Esegui questi comandi per attivare un test completo dell'NCCL per i nodi:

    head_pod=$(kubectl get pods --output='custom-columns=POD:.metadata.name' --no-headers | head -n1)

nodes=($(kubectl get pods --output='custom-columns=NODE:.spec.nodeName' --no-headers))

kubectl exec --stdin --tty --container=nccl-test ${head_pod} -- /configs/allgather.sh ${nodes[@]}

    L'output è simile al seguente:

    #                                                              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
     2097152         32768     float    none      -1    776.4    2.70    2.53      0    726.7    2.89    2.71      0
     4194304         65536     float    none      -1    774.3    5.42    5.08      0    805.1    5.21    4.88      0
     8388608        131072     float    none      -1    812.1   10.33    9.68      0    817.6   10.26    9.62      0
    16777216        262144     float    none      -1   1035.2   16.21   15.19      0   1067.8   15.71   14.73      0
    33554432        524288     float    none      -1   1183.3   28.36   26.59      0   1211.8   27.69   25.96      0
    67108864       1048576     float    none      -1   1593.4   42.12   39.49      0   1510.5   44.43   41.65      0
   134217728       2097152     float    none      -1   2127.8   63.08   59.13      0   2312.7   58.03   54.41      0
   268435456       4194304     float    none      -1   3603.0   74.50   69.85      0   3586.2   74.85   70.17      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

Utilizza le variabili di ambiente NCCL per migliorare le prestazioni

Facoltativamente, puoi impostare variabili di ambiente specifiche per migliorare le prestazioni dei carichi di lavoro che utilizzano NCCL. Il ConfigMap nccl-config-default.yaml di cui esegui il deployment nella sezione Esegui il deployment di un carico di lavoro di test imposta alcune variabili NCCL per impostazione predefinita. La configurazione della variabile viene archiviata nello script run-nccl.sh in ConfigMap.

Per modificare le variabili di ambiente NCCL, esegui il deployment di un manifest ConfigMap aggiornato con le variabili modificate. Il manifest nccl-config-latest.yaml in GitHub contiene ogni variabile consigliata con uno script run-nccl.sh aggiornato.

Il seguente comando aggiorna il ConfigMap esistente che contiene le variabili predefinite con il ConfigMap aggiornato nccl-config-latest.yaml:

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-config-latest.yaml

Kubernetes richiede circa due minuti per aggiornare ConfigMap.

Per verificare le variabili di ambiente NCCL, esegui questo comando:

head_pod=$(kubectl get pods --output='custom-columns=POD:.metadata.name' --no-headers | head -n1)

kubectl exec --stdin --tty --container=nccl-test ${head_pod} -- cat /configs/run-nccl.sh

Aggiungere GPUDirect-TCPX ai tuoi manifest

Questa sezione fornisce i campi obbligatori che devi aggiungere ai file manifest Kubernetes affinché i pod possano utilizzare GPUDirect-TCPX.

  1. Aggiungi i seguenti campi alla specifica del pod:

    spec:
  hostNetwork: true
  dnsPolicy: ClusterFirstWithHostNet
  volumes:
  - name: libraries
    hostPath:
      path: /home/kubernetes/bin/nvidia/lib64
  - name: tcpx-socket
    hostPath:
      path: /run/tcpx

  2. Aggiungi il container seguente al manifest per eseguire il servizio tcpx-daemon:

    - name: tcpx-daemon
  image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd-dev:v2.0.9
  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: libraries
      mountPath: /usr/local/nvidia/lib64
    - name: tcpx-socket
      mountPath: /run/tcpx
  env:
    - name: LD_LIBRARY_PATH
      value: /usr/local/nvidia/lib64

  3. Aggiungi i seguenti montaggi di volumi a tutti i container che richiedono GPU:

    volumeMounts:
- name: tcpx-socket
  mountPath: /tmp
- name: libraries
  mountPath: /usr/local/nvidia/lib64

  4. Aggiungi la seguente variabile di ambiente a ogni container GPU:

    env:
- name: LD_LIBRARY_PATH
  value: /usr/local/nvidia/lib64

  5. Se vuoi, aggiungi variabili di ambiente per configurare le opzioni NCCL. Per maggiori dettagli, consulta la sezione Utilizzare le variabili di ambiente NCCL per migliorare le prestazioni in questo documento.

Una specifica di pod completa è simile alla seguente:

apiVersion: v1
kind: Pod
metadata:
  name: example-pod
  labels:
    name: example-pod
spec:
  hostNetwork: true
  dnsPolicy: ClusterFirstWithHostNet
  containers:
  - name: tcpx-daemon
    image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd-dev:v2.0.9
    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: libraries
        mountPath: /usr/local/nvidia/lib64
      - name: tcpx-socket
        mountPath: /run/tcpx
    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:v3.1.2
      imagePullPolicy: Always
      command:
        - /bin/sh
        - -c
        - "while true; do echo hello; sleep 1; done"
      env:
        - name: LD_LIBRARY_PATH
          value: /usr/local/nvidia/lib64
      volumeMounts:
        - name: tcpx-socket
          mountPath: /run/tcpx
        - name: libraries
          mountPath: /usr/local/nvidia/lib64
      resources:
        limits:
          nvidia.com/gpu: 8
  volumes:
    - name: libraries
      hostPath:
        path: /home/kubernetes/bin/nvidia/lib64
    - name: tcpx-socket
      hostPath:
        path: /run/tcpx