apiVersion: v1
kind: ConfigMap
metadata:
  name: nccl-scripts
data:
  run.sh: |
    #!/usr/bin/env bash
    set -euo pipefail

    # ----- 기본 분산 설정 -----
    export MASTER_ADDR="${MASTER_ADDR:-nccl-0.nccl}"     # 첫 파드
    export MASTER_PORT="${MASTER_PORT:-23456}"
    export WORLD_SIZE="${WORLD_SIZE:-2}"                  # 파드 수
    export NUM_GPUS="${NUM_GPUS:-1}"                      # 파드당 프로세스(=GPU) 수

    # 파드 이름이 nccl-0, nccl-1 ... 이므로, 끝의 ordinal을 RANK로 사용
    ordinal="$(hostname | awk -F- '{print $NF}')"
    export RANK="${RANK:-$ordinal}"

    # ----- (선택) 특정 GPU만 사용하고 싶다면: 2번 GPU만 보이도록 고정 -----
    # 예: kubectl patch 로 아래 값을 "2"로 넣으면 그 노드의 물리 GPU #2만 노출됨
    if [[ -n "${CUDA_VISIBLE_DEVICES_OVERRIDE:-}" ]]; then
      export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES_OVERRIDE}"
    fi

    # ----- RDMA/IB 강제 -----
    export NCCL_DEBUG=INFO
    export NCCL_DEBUG_SUBSYS=INIT,NET
    export NCCL_NET=IB
    export NCCL_IB_DISABLE=0
    export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-net1}"   # HostDeviceNetwork로 붙인 NIC 이름
    export NCCL_IB_HCA="${NCCL_IB_HCA:-mlx5*}"
    # (RoCE 환경에 따라 필요 시) export NCCL_IB_GID_INDEX=3

    # 프로세스별 로그 파일 경로 (노드/프로세스마다 별도 파일)
    export NCCL_DEBUG_FILE="/tmp/nccl-%h-%p.log"

    # ----- 학습 실행 (각 파드가 1개 프로세스씩 구동) -----
    echo "[RUN] RANK=$RANK WORLD_SIZE=$WORLD_SIZE on $(hostname)"
    torchrun \
      --nnodes="${WORLD_SIZE}" \
      --nproc_per_node="${NUM_GPUS}" \
      --rdzv_backend=c10d \
      --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \
      /workspace/train_nccl.py 2>&1 | tee /tmp/nccl-aggregate.log

    # ----- RDMA 사용 여부 집계 -----
    total=0; ok=0
    for f in /tmp/nccl-*.log; do
      [[ -f "$f" ]] || continue
      total=$((total+1))
      if grep -q "NET/IB" "$f"; then
        ok=$((ok+1))
      fi
    done
    echo "[RDMA-CHECK] used IB on $ok / $total local ranks (processes) in $(hostname)"

    # 모든 랭크가 IB를 썼는지(파드 단위) 판단: 로컬 프로세스 수와 비교
    if [[ "$ok" -eq "$NUM_GPUS" ]]; then
      echo "[RDMA-CHECK] ✅ RDMA path OK for all local GPUs"
      exit 0
    else
      echo "[RDMA-CHECK] ❌ RDMA path NOT used by all local GPUs"
      echo "---- NCCL NET/IB related lines ----"
      grep -h "NET/" /tmp/nccl-*.log || true
      exit 1
    fi
  train_nccl.py: |
    import os, time, torch
    import torch.distributed as dist

    def main():
        # 분산 초기화
        rank = int(os.getenv("RANK", "0"))
        world_size = int(os.getenv("WORLD_SIZE", "1"))
        local_rank = int(os.getenv("LOCAL_RANK", "0"))  # torchrun이 설정

        # (선택) 특정 GPU 인덱스로 고정하고 싶을 때: LOCAL_RANK 대신 고정값 사용
        device_index = int(os.getenv("GPU_INDEX", str(local_rank)))
        torch.cuda.set_device(device_index)

        dist.init_process_group(
            backend="nccl",
            init_method="env://",
            world_size=world_size,
            rank=rank,
            timeout=torch.distributed.constants.default_pg_timeout
        )

        # 더미 학습 스텝: all-reduce로 통신 발생
        x = torch.ones(8 * 1024 * 1024, device=device_index, dtype=torch.float32)  # ~32MB
        for i in range(5):
            dist.all_reduce(x, op=dist.ReduceOp.SUM)
            torch.cuda.synchronize()
            if rank == 0:
                print(f"[STEP {i}] all_reduce done; tensor sum={x[0].item():.1f}")

        # 약간 대기해서 로그 flush
        time.sleep(2.0)
        dist.destroy_process_group()

    if __name__ == "__main__":
        main()
---
apiVersion: v1
kind: Service
metadata:
  name: nccl
  labels: { app: nccl }
spec:
  clusterIP: None  # Headless
  selector: { app: nccl }
  ports:
    - name: rdv
      port: 23456
---
apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: nccl
spec:
  serviceName: nccl
  replicas: 2                     # >=2 노드에서 스케줄되도록 노드 리소스 준비 필요
  selector:
    matchLabels: { app: nccl }
  template:
    metadata:
      labels: { app: nccl }
      annotations:
        k8s.v1.cni.cncf.io/networks: hostdevice-net   # HostDeviceNetwork (RDMA)
    spec:
      # RDMA 장치나 /dev/infiniband 노출이 안 보이면 네트워크 오퍼레이터의 RDMA 플러그인 구성을 확인하세요.
      # (환경에 따라 RDMA Shared Device Plugin이 필요할 수 있음) 
      containers:
        - name: worker
          image: nvcr.io/nvidia/pytorch:24.10-py3
          imagePullPolicy: IfNotPresent
          securityContext:
            capabilities: { add: ["IPC_LOCK"] }
          command: [ "bash", "-lc", "cp /config_scripts/* /workspace/ && chmod +x /workspace/run.sh && /workspace/run.sh || sleep infinity" ]
          env:
            # --- (선택) 특정 GPU만 사용하려면 아래 값을 물리 인덱스로 설정 (예: "2") ---
            # - name: CUDA_VISIBLE_DEVICES_OVERRIDE
            #   value: "2"
            - name: WORLD_SIZE
              value: "2"                # StatefulSet replicas와 동일
            - name: NUM_GPUS
              value: "8"                # 파드당 사용할 GPU 개수
            - name: MASTER_ADDR
              value: "nccl-0.nccl"      # 첫 파드의 DNS 이름
            - name: MASTER_PORT
              value: "23456"
            # RDMA 관련 기본값은 run.sh에서 설정
          resources:
            limits:
              nvidia.com/gpu: "8"
              nvidia.com/hostdev: "8"   # SR-IOV Device Plugin이 광고한 RDMA NIC 리소스
            requests:
              nvidia.com/gpu: "8"
              nvidia.com/hostdev: "8"
          volumeMounts:
            - name: scripts
              mountPath: /config_scripts
      volumes:
        - name: scripts
          configMap:
            name: nccl-scripts
            defaultMode: 0755