apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: {{ include "nccl-perftest.name" . }}
  labels:
    {{- include "nccl-perftest.labels" . | nindent 4 }}
spec:
  serviceName: {{ include "nccl-perftest.name" . }}
  replicas: {{ .Values.replicaCount }}
  selector:
    matchLabels:
      {{- include "nccl-perftest.selectorLabels" . | nindent 6 }}
  template:
    metadata:
      labels:
        {{- include "nccl-perftest.selectorLabels" . | nindent 8 }}
      annotations:
        k8s.v1.cni.cncf.io/networks: {{ .Values.network.nadName }}
    spec:
      containers:
        - name: worker
          image: {{ .Values.image.repository }}:{{ .Values.image.tag }}
          imagePullPolicy: {{ .Values.image.pullPolicy }}
          securityContext:
            privileged: true
            capabilities:
              add:
                - IPC_LOCK
          command:
            - bash
            - -lc
            - |
              cp /config_scripts/* /workspace/ && \
              chmod +x /workspace/run.sh && \
              cd /workspace && ./run.sh; \
              sleep infinity
          env:
            - name: NNODES
              value: {{ .Values.replicaCount | quote }}
            - name: NPROC_PER_NODE
              value: {{ .Values.training.nprocPerNode | quote }}
            - name: MASTER_ADDR
              value: "{{ include "nccl-perftest.name" . }}-0.{{ include "nccl-perftest.name" . }}"
            - name: MASTER_PORT
              value: {{ .Values.training.masterPort | quote }}
          resources:
            limits:
              nvidia.com/gpu: {{ .Values.resources.gpu | quote }}
              nvidia.com/hostdev: {{ .Values.resources.hostdev }}
            requests:
              nvidia.com/gpu: {{ .Values.resources.gpu | quote }}
              nvidia.com/hostdev: {{ .Values.resources.hostdev }}
          volumeMounts:
            - name: scripts
              mountPath: /config_scripts
            - name: shared-memory
              mountPath: /dev/shm
            - name: infiniband
              mountPath: /dev/infiniband
            - name: cifar-data
              mountPath: /workspace/data/cifar-10-batches-py
      volumes:
        - name: scripts
          configMap:
            name: {{ include "nccl-perftest.name" . }}-scripts
            defaultMode: 0755
        - name: shared-memory
          emptyDir:
            medium: Memory
            sizeLimit: {{ .Values.shmSize }}
        - name: infiniband
          hostPath:
            path: {{ .Values.volumes.infiniband.hostPath }}
            type: Directory
        - name: cifar-data
          hostPath:
            path: {{ .Values.volumes.dataset.hostPath }}
            type: Directory