apiVersion: v1
kind: ConfigMap
metadata:
  # 이 이름은 파드 YAML의 volumeMounts에서 참조해야 합니다.
  name: ddp-test-script
  namespace: default
data:
  ddp_rdma_test.py: |
    # 파일명: ddp_rdma_test.py
    import os
    import torch
    import torch.distributed as dist
    from torch.nn.parallel import DistributedDataParallel as DDP

    # torchrun에 의해 자동 주입되는 환경 변수들을 가져옵니다.
    def setup():
        # 1. 환경 변수를 읽어 분산 통신 초기화
        # torchrun은 MASTER_ADDR, MASTER_PORT, RANK, WORLD_SIZE를 설정해줍니다.
        
        # 1.1. NCCL 백엔드 사용
        dist.init_process_group("nccl")

        # 1.2. torchrun이 설정한 로컬 랭크(해당 파드 내 GPU ID)를 가져옵니다.
        # LOCAL_RANK 환경 변수는 torchrun에 의해 자동 설정됩니다.
        local_rank = int(os.environ["LOCAL_RANK"])
        
        # 1.3. 해당 프로세스에 GPU를 할당합니다.
        torch.cuda.set_device(local_rank)
        
        return local_rank

    def run_ddp():
        local_rank = setup()
        
        # torchrun이 설정한 전체 정보를 가져옵니다.
        global_rank = dist.get_rank()
        world_size = dist.get_world_size()
        # NODE_RANK는 파드 환경 변수로 설정되어야 합니다.
        node_rank = int(os.environ.get("NODE_RANK", 0))
        
        print(f"--- Process Initialized ---")
        print(f"Global Rank: {global_rank} | Local Rank (GPU ID): {local_rank} | Node Rank: {node_rank}")
        
        # 간단한 모델 정의 및 GPU 로드
        model = torch.nn.Linear(10, 1).cuda()
        ddp_model = DDP(model, device_ids=[local_rank])
        
        # ---- RDMA 통신 검증을 위한 All-Reduce 연산 수행 ----
        tensor = torch.ones(10, dtype=torch.float32).cuda()
        
        # All-Reduce 실행: NCCL의 GPU Direct RDMA 기능이 활용됩니다.
        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

        # ----------------------------------------------------

        expected_sum = world_size * 10 
        
        # 결과 검증 및 출력
        if global_rank == 0:
            print(f"\n--- Validation Check ---")
            print(f"Total World Size: {world_size}")
        
        print(f"Global Rank {global_rank} / GPU {local_rank}: All-Reduce 결과 합계: {tensor.sum().item()}")

        # 모든 GPU에서 결과가 일치하는지 확인
        assert torch.allclose(tensor.sum().item(), torch.tensor(expected_sum, dtype=torch.float32).item()), "All-Reduce 결과 오류! 분산 통신 실패 또는 잘못된 계산."
        
        dist.destroy_process_group()

    if __name__ == '__main__':
        # torchrun은 이 파일을 실행하며, 내부에서 run_ddp 로직을 반복 실행합니다.
        run_ddp()
---
apiVersion: v1
kind: Pod
metadata:
  name: nccl-gdrdma-pod-gcore-1
  annotations:
    # 1. Multus CNI를 통해 'hostdevice-net' 네트워크에 연결합니다.
    k8s.v1.cni.cncf.io/networks: hostdevice-net
    # 2. 파드를 Node 0에 강제 배포합니다. (실제 클러스터에서는 스케줄러가 처리)
    scheduler.alpha.kubernetes.io/node-selector: kubernetes.io/hostname=nhn-aideveloper-d40ce770
spec:
  # serviceAccountName: <적절한 SA 이름> # (필요시)
  containers:
  - name: gdrdma-client
    # NVIDIA 공식 PyTorch 이미지 사용
    image: nvcr.io/nvidia/pytorch:24.03-py3
    imagePullPolicy: IfNotPresent
    command: ["/bin/bash", "-c"]
    args:
      - |
        # 1. RDMA 디바이스 정보 확인 및 테스트 도구 설치 (검증을 위해)
        apt update && apt install -y rdma-core perftest
        
        # 2. 분산 학습 환경 변수 설정
        export NCCL_DEBUG=INFO             # NCCL 디버그 레벨 (RDMA 사용 검증)
        export NCCL_DEBUG_SUBSYS=NET       # 네트워크 통신 계층 디버깅 활성화
        export MASTER_ADDR=192.168.240.1  # Node 0 파드의 hostdevice-net IP
        export MASTER_PORT=29500
        export WORLD_SIZE=16               # 총 GPU 개수 (8 * 2)
        export NODE_RANK=0                 # 현재 파드의 노드 순위 (Node 0)
        export LOCAL_WORLD_SIZE=8          # 이 파드에 할당된 GPU 개수
        
        # 3. ddp_rdma_test.py 스크립트 실행 (미리 파드에 복사되어 있어야 함)
        # torch.distributed.launch 또는 torchrun을 사용하여 8개 프로세스 실행
        torchrun \
          --nnodes=$((WORLD_SIZE / LOCAL_WORLD_SIZE)) \
          --nproc_per_node=$LOCAL_WORLD_SIZE \
          --rdzv_id=100 \
          --rdzv_backend=c10d \
          --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
          ddp_rdma_test.py

        # 3. 중요: torchrun 실행 대신, 무한 대기합니다.
        echo "Setup complete. Pod is now in a waiting state. Use 'kubectl exec' to run torchrun."
        sleep infinity  # <-- 이 명령어가 컨테이너가 종료되지 않도록 유지합니다.

    securityContext:
      capabilities:
        add: ["IPC_LOCK"] # GPU Direct RDMA를 위한 필수 설정
    resources:
      limits:
        # GPU 8개와 Host Device NIC 1개를 요청합니다.
        # 이 Host Device는 SR-IOV Device Plugin에 의해 RDMA 기능이 활성화됩니다.
        nvidia.com/gpu: "8"
        nvidia.com/hostdev: "8"
      requests:
        nvidia.com/gpu: "8"
        nvidia.com/hostdev: "8"
    volumeMounts:
    - name: ddp-script-volume
      mountPath: /app
  volumes:
  - name: ddp-script-volume
    configMap:
      name: ddp-test-script