apiVersion: v1
kind: ConfigMap
metadata:
  name: pytorch-train-script
data:
  train.py: |
    import os
    import time
    import torch
    import torch.distributed as dist

    def main():
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])

        print(f"[Rank {rank}] Starting process_group init...", flush=True)

        dist.init_process_group(
            backend="nccl",
            init_method="env://",
            rank=rank,
            world_size=world_size
        )

        print(f"[Rank {rank}] process_group initialized!", flush=True)

        # GPU 텐서 생성
        torch.cuda.set_device(rank % torch.cuda.device_count())
        tensor = torch.ones(1).cuda()

        # NCCL broadcast test
        print(f"[Rank {rank}] tensor before: {tensor}", flush=True)

        dist.broadcast(tensor, src=0)

        print(f"[Rank {rank}] tensor after broadcast: {tensor}", flush=True)

        # 종료되지 않도록 대기
        print(f"[Rank {rank}] Sleeping for 30 seconds so logs can be seen...", flush=True)
        time.sleep(600)

    if __name__ == "__main__":
        main()

---
apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: pytorch-multinode-rdma
spec:
  runPolicy:
    cleanPodPolicy: None

  pytorchReplicaSpecs:
    Master:
      replicas: 1         # 노드 수 = 2 (원하면 증가 가능)
      restartPolicy: OnFailure
      template:
        metadata:
          annotations:
            k8s.v1.cni.cncf.io/networks: hostdevice-net
        spec:
          nodeSelector:
            nodegroup: infini
          initContainers:
          - name: copy-train-script
            image: busybox
            command:
              - sh
              - -c
              - |
                echo "Copying train.py to workspace..."
                cp /config/train.py /workspace/train.py
                chmod +x /workspace/train.py
            volumeMounts:
              - name: train-config
                mountPath: /config
              - name: workspace
                mountPath: /workspace
          containers:
          - name: pytorch
            image: nvcr.io/nvidia/pytorch:24.01-py3   # RDMA 지원되는 NVIDIA PT 이미지 권장
            imagePullPolicy: IfNotPresent
            command:
              - "python"
              - "/workspace/train.py"
            env:
              # --- PyTorch DDP ---
              - name: MASTER_ADDR
                value: "pytorch-multinode-rdma-worker-0"
              - name: MASTER_PORT
                value: "23456"

              # --- NCCL RDMA 설정 ---
              - name: NCCL_DEBUG
                value: "INFO"
              - name: NCCL_IB_HCA
                value: "mlx5_*"
              - name: NCCL_IB_DISABLE
                value: "0"
              - name: NCCL_NET_GDR_LEVEL
                value: "2"
              - name: NCCL_SOCKET_IFNAME
                value: "net1"         # Multus의 RDMA 인터페이스를 여기에 넣어도 됨 (ipoib: ib0)
              - name: NCCL_IB_PCI_RELAXED_ORDERING
                value: "1"

            resources:
              limits:
                nvidia.com/gpu: 8
                nvidia.com/hostdev: "8"   # 또는 rdma/hca: 8 (리소스 명칭에 따름)
            volumeMounts:
            - name: workspace
              mountPath: /workspace
          volumes:
          - name: train-config           # ConfigMap (read-only)
            configMap:
              name: pytorch-train-script
          - name: workspace              # EmptyDir (read-write)
            emptyDir: {}
    Worker:
      replicas: 1         # 노드 수 = 2 (원하면 증가 가능)
      restartPolicy: OnFailure
      template:
        metadata:
          annotations:
            k8s.v1.cni.cncf.io/networks: hostdevice-net
        spec:
          nodeSelector:
            nodegroup: infini
          initContainers:
          - name: copy-train-script
            image: busybox
            command:
              - sh
              - -c
              - |
                echo "Copying train.py to workspace..."
                cp /config/train.py /workspace/train.py
                chmod +x /workspace/train.py
            volumeMounts:
              - name: train-config
                mountPath: /config
              - name: workspace
                mountPath: /workspace
          containers:
          - name: pytorch
            image: nvcr.io/nvidia/pytorch:24.01-py3   # RDMA 지원되는 NVIDIA PT 이미지 권장
            imagePullPolicy: IfNotPresent
            command:
              - "python"
              - "/workspace/train.py"
            env:
              # --- PyTorch DDP ---
              - name: MASTER_ADDR
                value: "pytorch-multinode-rdma-worker-0"
              - name: MASTER_PORT
                value: "23456"

              # --- NCCL RDMA 설정 ---
              - name: NCCL_DEBUG
                value: "INFO"
              - name: NCCL_IB_HCA
                value: "mlx5_*"
              - name: NCCL_IB_DISABLE
                value: "0"
              - name: NCCL_NET_GDR_LEVEL
                value: "2"
              - name: NCCL_SOCKET_IFNAME
                value: "eth0"         # Multus의 RDMA 인터페이스를 여기에 넣어도 됨 (ipoib: ib0)
              - name: NCCL_IB_PCI_RELAXED_ORDERING
                value: "1"

            resources:
              limits:
                nvidia.com/gpu: 8
                nvidia.com/hostdev: "8"   # 또는 rdma/hca: 8 (리소스 명칭에 따름)
            volumeMounts:
            - name: workspace
              mountPath: /workspace
          volumes:
          - name: train-config           # ConfigMap (read-only)
            configMap:
              name: pytorch-train-script
          - name: workspace              # EmptyDir (read-write)
            emptyDir: {}