Add TensorFlow custom runtime for Kubeflow Trainer v2 with Volcano integration

Trainer v2에는 TensorFlow runtime이 기본 제공되지 않으므로
Custom ClusterTrainingRuntime을 생성하여 MultiWorkerMirroredStrategy 기반
분산학습을 지원한다. Pod hostname + Headless Service DNS로 TF_CONFIG를 자동 구성.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

.gitignore

@@ -0,0 +1 @@
+.env

README.md

@@ -0,0 +1,211 @@
+# Kubeflow Trainer v2 - TensorFlow Custom Runtime
+
+Kubeflow Trainer v2에서 TensorFlow 분산학습을 위한 Custom ClusterTrainingRuntime 생성 및 Volcano 연동 가이드.
+
+## 배경
+
+Kubeflow Trainer v2 (v1alpha1)에는 다음 런타임만 기본 제공된다:
+
+| Runtime | Framework Label | mlPolicy 필드 | 엔트리포인트 주입 |
+|---------|----------------|---------------|-------------------|
+| torch-distributed | `torch` | `mlPolicy.torch` | torchrun + MASTER_ADDR/RANK 등 자동 설정 |
+| deepspeed-distributed | `deepspeed` | `mlPolicy.mpi` | MPI (OpenMPI) + SSH 자동 설정 |
+| mlx-distributed | `mlx` | `mlPolicy.mpi` | MPI (OpenMPI) + SSH 자동 설정 |
+| torchtune-* | `torchtune` | `mlPolicy.torch` | tune run + rdzv 자동 설정 |
+
+**TensorFlow runtime은 없다.** 따라서 Custom ClusterTrainingRuntime을 직접 생성해야 한다.
+
+## PyTorch vs TensorFlow Runtime 핵심 차이
+
+```
+┌─────────────────────────────────────────────────────────────────────┐
+│                    PyTorch Runtime (built-in)                       │
+│                                                                     │
+│  mlPolicy.torch → 컨트롤러가 자동으로:                               │
+│    ├── torchrun 엔트리포인트 주입                                    │
+│    ├── MASTER_ADDR, MASTER_PORT 설정                                │
+│    ├── RANK, LOCAL_RANK, WORLD_SIZE 설정                            │
+│    └── rdzv (rendezvous) 자동 구성                                   │
+│                                                                     │
+│  → 학습 코드에서 분산환경 설정 불필요                                  │
+└─────────────────────────────────────────────────────────────────────┘
+
+┌─────────────────────────────────────────────────────────────────────┐
+│                 TensorFlow Runtime (custom)                         │
+│                                                                     │
+│  mlPolicy에 framework 없음 → 컨트롤러가 주입하는 것 없음             │
+│                                                                     │
+│  학습 스크립트에서 자체 구성:                                         │
+│    ├── Pod hostname 파싱 → TrainJob 이름, worker index 추출          │
+│    ├── Headless Service DNS로 worker 주소 목록 생성                  │
+│    ├── TF_CONFIG 환경변수 자동 구성                                   │
+│    └── tf.distribute.MultiWorkerMirroredStrategy 사용                │
+│                                                                     │
+│  필수 설정:                                                          │
+│    ├── network.publishNotReadyAddresses: true (DNS 조기 해석)        │
+│    └── NUM_WORKERS 환경변수 = mlPolicy.numNodes                      │
+└─────────────────────────────────────────────────────────────────────┘
+```
+
+## TF_CONFIG 자동 구성 원리
+
+Trainer v2에서 TrainJob을 생성하면 내부적으로 JobSet이 만들어지고, 각 Pod은 예측 가능한 hostname을 갖는다:
+
+```
+Pod hostname 패턴: {trainjob-name}-node-{replica_index}-{completion_index}
+
+예시 (TrainJob: tf-distributed-training, numNodes: 2):
+  Worker 0: tf-distributed-training-node-0-0
+  Worker 1: tf-distributed-training-node-1-0
+```
+
+Headless Service DNS와 결합하면:
+
+```
+Worker 0: tf-distributed-training-node-0-0.tf-distributed-training.default.svc.cluster.local:12345
+Worker 1: tf-distributed-training-node-1-0.tf-distributed-training.default.svc.cluster.local:12345
+```
+
+학습 스크립트가 이를 파싱하여 자동으로 TF_CONFIG를 생성한다:
+
+```json
+{
+  "cluster": {
+    "worker": [
+      "tf-distributed-training-node-0-0.tf-distributed-training.default.svc.cluster.local:12345",
+      "tf-distributed-training-node-1-0.tf-distributed-training.default.svc.cluster.local:12345"
+    ]
+  },
+  "task": {"type": "worker", "index": 0}
+}
+```
+
+## 파일 구조
+
+```
+.
+├── README.md                                      # 이 문서
+├── tensorflow-custom-runtime.yaml                 # TF Custom Runtime (기본, Volcano 미포함)
+├── tensorflow-volcano-trainjob-integration.yaml   # TF + Volcano 전체 통합 (ConfigMap + Queue + Runtime + TrainJob)
+└── volcano-trainjob-integration.yaml              # (참고) PyTorch + Volcano 통합 원본
+```
+
+## 적용 방법
+
+### 1단계: 기본 TensorFlow Runtime만 적용
+
+Volcano 없이 기본 TF 런타임만 필요한 경우:
+
+```bash
+kubectl apply -f tensorflow-custom-runtime.yaml
+```
+
+확인:
+```bash
+kubectl get clustertrainingruntime
+# tensorflow-distributed 가 목록에 나타나야 함
+```
+
+### 2단계: Volcano 연동 전체 적용
+
+Volcano gang scheduling + 토폴로지 인식 스케줄링이 포함된 전체 버전:
+
+```bash
+# ConfigMap (학습 스크립트) + Queue + Runtime + TrainJob 한번에 적용
+kubectl apply -f tensorflow-volcano-trainjob-integration.yaml
+```
+
+확인:
+```bash
+# Runtime 확인
+kubectl get clustertrainingruntime tensorflow-distributed-volcano
+
+# TrainJob 상태 확인
+kubectl get trainjob tf-distributed-training
+
+# Pod 상태 확인
+kubectl get pods -l batch.kubernetes.io/job-name
+
+# Volcano PodGroup 확인
+kubectl get podgroup
+
+# 로그 확인 (worker 0)
+kubectl logs -l batch.kubernetes.io/job-name=tf-distributed-training-node-0 -f
+```
+
+### 3단계: TrainJob만 변경하여 재실행
+
+Runtime은 유지하고 TrainJob만 변경할 경우:
+
+```bash
+# 기존 TrainJob 삭제
+kubectl delete trainjob tf-distributed-training
+
+# 수정 후 재적용
+kubectl apply -f tensorflow-volcano-trainjob-integration.yaml
+```
+
+## 주의사항
+
+### NUM_WORKERS와 numNodes 동기화
+
+`NUM_WORKERS` 환경변수는 `mlPolicy.numNodes` (또는 TrainJob의 `trainer.numNodes`)와 **반드시 일치**해야 한다.
+
+PyTorch runtime은 컨트롤러가 `WORLD_SIZE`를 자동 주입하지만, 커스텀 TF runtime은 이를 수동으로 관리해야 한다.
+
+```yaml
+# Runtime에서 기본값 설정
+spec:
+  mlPolicy:
+    numNodes: 2      # ← 이 값과
+  ...
+  env:
+    - name: NUM_WORKERS
+      value: "2"     # ← 이 값이 일치해야 함
+```
+
+TrainJob에서 numNodes를 오버라이드할 경우, env도 함께 오버라이드해야 한다:
+
+```yaml
+apiVersion: trainer.kubeflow.org/v1alpha1
+kind: TrainJob
+spec:
+  runtimeRef:
+    name: tensorflow-distributed-volcano
+  trainer:
+    numNodes: 4         # 4노드로 변경
+    env:
+      - name: NUM_WORKERS
+        value: "4"      # 반드시 함께 변경
+```
+
+### InfiniBand 설정
+
+클러스터에 InfiniBand가 없는 경우 다음 항목을 제거해야 한다:
+
+```yaml
+# 제거 대상:
+metadata:
+  annotations:
+    k8s.v1.cni.cncf.io/networks: hostdevice-net   # ← 제거
+
+resources:
+  requests:
+    nvidia.com/hostdev: "2"   # ← 제거
+  limits:
+    nvidia.com/hostdev: "2"   # ← 제거
+```
+
+### network.publishNotReadyAddresses
+
+TensorFlow의 `MultiWorkerMirroredStrategy`는 시작 시 모든 worker에 gRPC 연결을 시도한다. 모든 Pod이 Ready 상태가 되기 전에도 DNS가 해석되어야 하므로 이 설정이 **필수**다:
+
+```yaml
+spec:
+  template:
+    spec:
+      network:
+        publishNotReadyAddresses: true   # 반드시 true
+```
+
+PyTorch runtime에서는 torchrun의 rdzv(rendezvous) 메커니즘이 이를 처리하므로 불필요하다.

tensorflow-custom-runtime.yaml

@@ -0,0 +1,46 @@
+# Kubeflow Trainer v2 - TensorFlow Custom ClusterTrainingRuntime
+# Trainer v2에는 TensorFlow runtime이 기본 제공되지 않으므로 Custom Runtime으로 생성
+#
+# 핵심 차이점:
+#   - PyTorch runtime: mlPolicy.torch 사용 → torchrun 자동 주입, MASTER_ADDR/RANK 등 자동 설정
+#   - TensorFlow runtime (custom): mlPolicy에 framework 없음 → TF_CONFIG를 학습 스크립트에서 자체 구성
+#
+# TF_CONFIG 자동 구성 방식:
+#   1. Pod hostname에서 TrainJob 이름과 replica index 추출
+#   2. Headless Service DNS를 이용한 worker 주소 목록 생성
+#   3. tf.distribute.MultiWorkerMirroredStrategy로 분산학습 실행
+
+apiVersion: trainer.kubeflow.org/v1alpha1
+kind: ClusterTrainingRuntime
+metadata:
+  name: tensorflow-distributed
+  labels:
+    trainer.kubeflow.org/framework: tensorflow
+spec:
+  mlPolicy:
+    numNodes: 1
+  template:
+    spec:
+      network:
+        publishNotReadyAddresses: true
+      replicatedJobs:
+      - name: node
+        template:
+          metadata:
+            labels:
+              trainer.kubeflow.org/trainjob-ancestor-step: trainer
+          spec:
+            template:
+              spec:
+                containers:
+                - name: node
+                  image: tensorflow/tensorflow:2.16.1-gpu
+                  env:
+                    - name: NUM_WORKERS
+                      value: "1"
+                    - name: TF_WORKER_PORT
+                      value: "12345"
+                    - name: NAMESPACE
+                      valueFrom:
+                        fieldRef:
+                          fieldPath: metadata.namespace

tensorflow-volcano-trainjob-integration.yaml

@@ -0,0 +1,318 @@
+# Volcano <-> Kubeflow Trainer v2 TensorFlow Integration
+# Reference: https://www.kubeflow.org/docs/components/trainer/gang-scheduling/volcano/
+#
+# Trainer v2에는 TensorFlow runtime이 없으므로 Custom Runtime을 생성하여 사용
+# TF_CONFIG는 학습 스크립트에서 Pod hostname + Headless Service DNS로 자동 구성
+
+---
+# 0. ConfigMap for TensorFlow training script (CNN + CIFAR-10 + MultiWorkerMirroredStrategy)
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: tf-training-scripts
+data:
+  train_tf.py: |
+    """
+    TensorFlow Distributed Training Example for Kubeflow Trainer v2
+    - CNN (CIFAR-10) + MultiWorkerMirroredStrategy + NCCL
+    - TF_CONFIG를 Pod hostname과 환경변수로 자동 구성
+    """
+    import os
+    import json
+    import socket
+    import re
+    import tensorflow as tf
+    import numpy as np
+
+
+    def setup_tf_config():
+        """
+        Kubeflow Trainer v2의 Pod 네이밍 규칙을 이용하여 TF_CONFIG를 자동 구성.
+
+        Pod hostname 패턴: {trainjob-name}-node-{replica_index}-{completion_index}
+        Headless Service DNS: {hostname}.{trainjob-name}.{namespace}.svc.cluster.local
+
+        환경변수:
+          - NUM_WORKERS: 총 워커 수 (runtime의 numNodes와 일치해야 함)
+          - TF_WORKER_PORT: gRPC 통신 포트 (기본값: 12345)
+          - NAMESPACE: Pod이 실행되는 namespace (Downward API)
+        """
+        hostname = socket.gethostname()
+        num_workers = int(os.environ.get('NUM_WORKERS', '1'))
+        port = os.environ.get('TF_WORKER_PORT', '12345')
+        namespace = os.environ.get('NAMESPACE', 'default')
+
+        # hostname에서 TrainJob 이름과 replica index 추출
+        # 패턴: {trainjob-name}-node-{replica_index}-{completion_index}
+        match = re.match(r'^(.+)-node-(\d+)-(\d+)$', hostname)
+        if match:
+            job_name = match.group(1)
+            worker_index = int(match.group(2))
+        else:
+            print(f"[WARNING] Cannot parse hostname '{hostname}', using defaults")
+            job_name = os.environ.get('TRAINJOB_NAME', 'unknown')
+            worker_index = 0
+
+        # Headless Service DNS를 이용한 worker 주소 목록 생성
+        workers = []
+        for i in range(num_workers):
+            worker_host = (
+                f"{job_name}-node-{i}-0.{job_name}.{namespace}"
+                f".svc.cluster.local:{port}"
+            )
+            workers.append(worker_host)
+
+        tf_config = {
+            "cluster": {"worker": workers},
+            "task": {"type": "worker", "index": worker_index}
+        }
+
+        os.environ['TF_CONFIG'] = json.dumps(tf_config)
+        print(f"[Worker {worker_index}/{num_workers}] hostname: {hostname}")
+        print(f"[Worker {worker_index}/{num_workers}] TF_CONFIG:")
+        print(json.dumps(tf_config, indent=2))
+
+        return worker_index, num_workers
+
+
+    def build_model():
+        """CIFAR-10용 CNN 모델"""
+        model = tf.keras.Sequential([
+            tf.keras.layers.Conv2D(
+                32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
+            tf.keras.layers.BatchNormalization(),
+            tf.keras.layers.Conv2D(32, (3, 3), activation='relu'),
+            tf.keras.layers.MaxPooling2D((2, 2)),
+            tf.keras.layers.Dropout(0.25),
+
+            tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
+            tf.keras.layers.BatchNormalization(),
+            tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
+            tf.keras.layers.MaxPooling2D((2, 2)),
+            tf.keras.layers.Dropout(0.25),
+
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(512, activation='relu'),
+            tf.keras.layers.BatchNormalization(),
+            tf.keras.layers.Dropout(0.5),
+            tf.keras.layers.Dense(10)
+        ])
+        return model
+
+
+    def main():
+        worker_index, num_workers = setup_tf_config()
+
+        # MultiWorkerMirroredStrategy (GPU간 NCCL 통신)
+        communication_options = tf.distribute.experimental.CommunicationOptions(
+            implementation=tf.distribute.experimental.CommunicationImplementation.NCCL
+        )
+        strategy = tf.distribute.MultiWorkerMirroredStrategy(
+            communication_options=communication_options
+        )
+
+        print(f"[Worker {worker_index}] num_replicas_in_sync: "
+              f"{strategy.num_replicas_in_sync}")
+
+        # CIFAR-10 데이터셋 로드
+        (x_train, y_train), (x_test, y_test) = (
+            tf.keras.datasets.cifar10.load_data()
+        )
+        x_train = x_train.astype('float32') / 255.0
+        x_test = x_test.astype('float32') / 255.0
+
+        # 배치 크기 설정
+        BATCH_SIZE_PER_REPLICA = 64
+        GLOBAL_BATCH_SIZE = (
+            BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync
+        )
+
+        # tf.data.Dataset 생성
+        train_dataset = tf.data.Dataset.from_tensor_slices(
+            (x_train, y_train)
+        ).shuffle(50000).batch(GLOBAL_BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
+
+        test_dataset = tf.data.Dataset.from_tensor_slices(
+            (x_test, y_test)
+        ).batch(GLOBAL_BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
+
+        # Strategy scope 내에서 모델 생성 및 컴파일
+        with strategy.scope():
+            model = build_model()
+            model.compile(
+                optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
+                loss=tf.keras.losses.SparseCategoricalCrossentropy(
+                    from_logits=True),
+                metrics=['accuracy']
+            )
+
+        # 학습
+        EPOCHS = int(os.environ.get('EPOCHS', '10'))
+        verbose = 2 if worker_index == 0 else 0
+
+        callbacks = []
+        if worker_index == 0:
+            callbacks.append(
+                tf.keras.callbacks.TensorBoard(log_dir='/workspace/logs')
+            )
+
+        model.fit(
+            train_dataset,
+            epochs=EPOCHS,
+            validation_data=test_dataset,
+            verbose=verbose,
+            callbacks=callbacks
+        )
+
+        # 평가 (worker 0만 출력)
+        if worker_index == 0:
+            loss, accuracy = model.evaluate(test_dataset, verbose=2)
+            print(f"\n===== Final Results =====")
+            print(f"Test Loss:     {loss:.4f}")
+            print(f"Test Accuracy: {accuracy:.4f}")
+            print(f"=========================")
+
+
+    if __name__ == '__main__':
+        main()
+
+---
+# 1. Volcano Queue for training workloads
+# Queue는 리소스 할당 단위. capability는 클러스터 실제 용량에 맞게 조정 필요.
+apiVersion: scheduling.volcano.sh/v1beta1
+kind: Queue
+metadata:
+  name: training-queue
+spec:
+  weight: 1
+  reclaimable: true
+  capability:
+    # 클러스터 리소스에 맞게 조정
+    cpu: "100"
+    memory: "512Gi"
+    nvidia.com/gpu: "24"
+    nvidia.com/hostdev: "6"
+
+---
+# 2. TensorFlow Custom ClusterTrainingRuntime with Volcano gang scheduling
+#
+# PyTorch runtime과의 핵심 차이:
+#   - mlPolicy에 torch/mpi 필드 없음 → 컨트롤러가 torchrun/MPI 주입하지 않음
+#   - network.publishNotReadyAddresses: true → Pod간 DNS 조기 해석 가능
+#   - TF_CONFIG는 학습 스크립트에서 hostname + Headless Service DNS로 자동 구성
+#   - NUM_WORKERS 환경변수가 numNodes와 반드시 일치해야 함
+apiVersion: trainer.kubeflow.org/v1alpha1
+kind: ClusterTrainingRuntime
+metadata:
+  name: tensorflow-distributed-volcano
+  labels:
+    trainer.kubeflow.org/framework: tensorflow
+spec:
+  mlPolicy:
+    numNodes: 2
+  # Volcano gang scheduling 활성화 - PodGroup 자동 생성
+  podGroupPolicy:
+    volcano:
+      networkTopology:
+        mode: hard
+        highestTierAllowed: 1
+  template:
+    metadata:
+      annotations:
+        # Queue 지정 (runtime level)
+        scheduling.volcano.sh/queue-name: training-queue
+    spec:
+      network:
+        # TensorFlow MultiWorkerMirroredStrategy가 gRPC로 통신하므로
+        # 모든 Pod이 Ready 전에도 DNS 해석 가능해야 함
+        publishNotReadyAddresses: true
+      replicatedJobs:
+      - name: node
+        template:
+          metadata:
+            labels:
+              trainer.kubeflow.org/trainjob-ancestor-step: trainer
+          spec:
+            template:
+              metadata:
+                annotations:
+                  # InfiniBand CNI (클러스터에 InfiniBand가 없으면 이 줄 제거)
+                  k8s.v1.cni.cncf.io/networks: hostdevice-net
+              spec:
+                containers:
+                - name: node
+                  image: nvcr.io/nvidia/tensorflow:24.03-tf2-py3
+                  securityContext:
+                    privileged: true
+                    capabilities:
+                      add:
+                        - IPC_LOCK
+                  env:
+                    # --- NCCL 설정 (TF MultiWorkerMirroredStrategy + NCCL backend) ---
+                    - name: NCCL_DEBUG
+                      value: "INFO"
+                    - name: NCCL_DEBUG_SUBSYS
+                      value: "INIT,NET,IB"
+                    - name: NCCL_IB_DISABLE
+                      value: "0"
+                    - name: NCCL_SOCKET_IFNAME
+                      value: "net"
+                    # --- TF_CONFIG 자동 구성용 환경변수 ---
+                    # NUM_WORKERS는 mlPolicy.numNodes와 반드시 일치해야 함
+                    - name: NUM_WORKERS
+                      value: "2"
+                    - name: TF_WORKER_PORT
+                      value: "12345"
+                    - name: NAMESPACE
+                      valueFrom:
+                        fieldRef:
+                          fieldPath: metadata.namespace
+                    # 학습 에포크 수
+                    - name: EPOCHS
+                      value: "10"
+                  resources:
+                    requests:
+                      nvidia.com/gpu: "8"
+                      nvidia.com/hostdev: "2"
+                    limits:
+                      nvidia.com/gpu: "8"
+                      nvidia.com/hostdev: "2"
+                  volumeMounts:
+                    - name: shared-memory
+                      mountPath: /dev/shm
+                    - name: training-scripts
+                      mountPath: /workspace/scripts
+                volumes:
+                  - name: shared-memory
+                    emptyDir:
+                      medium: Memory
+                      sizeLimit: 128Gi
+                  - name: training-scripts
+                    configMap:
+                      name: tf-training-scripts
+                      defaultMode: 0755
+
+---
+# 3. Example TrainJob (TensorFlow 2노드 분산학습)
+# podGroupPolicy, env, volumes, command 등은 runtime에서 상속됨
+apiVersion: trainer.kubeflow.org/v1alpha1
+kind: TrainJob
+metadata:
+  name: tf-distributed-training
+  namespace: default
+spec:
+  runtimeRef:
+    name: tensorflow-distributed-volcano
+  trainer:
+    image: nvcr.io/nvidia/tensorflow:24.03-tf2-py3
+    command:
+      - python
+      - /workspace/scripts/train_tf.py
+    numNodes: 2
+    resourcesPerNode:
+      requests:
+        nvidia.com/gpu: "8"
+        nvidia.com/hostdev: "2"
+      limits:
+        nvidia.com/gpu: "8"
+        nvidia.com/hostdev: "2"

volcano-trainjob-integration.yaml

@@ -0,0 +1,392 @@
+# Volcano <-> Kubeflow Trainer Integration Configuration Sample
+# Reference: https://www.kubeflow.org/docs/components/trainer/gang-scheduling/volcano/
+
+---
+# 0. ConfigMap for test training script (VGG11 + CIFAR10 + NCCL)
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: training-scripts
+data:
+  train_nccl.py: |
+    import os
+    import time
+    import torch
+    import argparse
+    import torch.distributed as dist
+    import torch.multiprocessing as mp
+    from torch.optim.lr_scheduler import StepLR
+
+    # for dataset
+    from torchvision.datasets.cifar import CIFAR10
+    import torchvision.transforms as tfs
+    from torch.utils.data import DataLoader
+    from torch.utils.data.distributed import DistributedSampler
+
+    # for model
+    from torchvision.models import vgg11
+    from torch.nn.parallel import DistributedDataParallel as DDP
+
+    import numpy as np
+    import random
+    import datetime
+
+
+    def set_random_seeds(random_seed=0):
+        torch.manual_seed(random_seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        np.random.seed(random_seed)
+        random.seed(random_seed)
+
+
+    def get_args_parser():
+        parser = argparse.ArgumentParser(add_help=False)
+        parser.add_argument('--lr', type=float, default=0.01)
+        parser.add_argument('--epoch', type=int, default=90)
+        parser.add_argument('--batch_size', type=int, default=1200)
+        parser.add_argument('--global_rank', type=int, default=0)
+        parser.add_argument('--vis_step', type=int, default=10)
+        parser.add_argument('--num_workers', type=int, default=24)
+        parser.add_argument("--local_rank", type=int,
+                            help="Local rank. Necessary for using the torch.distributed.launch utility.")
+        parser.add_argument('--world_size', type=int, default=0)
+        parser.add_argument('--port', type=int, default=2022)
+        parser.add_argument('--root', type=str, default='data')
+        parser.add_argument('--start_epoch', type=int, default=0)
+        parser.add_argument('--save_path', type=str, default='./save')
+        parser.add_argument('--save_file_name', type=str, default='vgg_cifar')
+        return parser
+
+
+    def main(opts):
+        # 1. set random seeds
+        set_random_seeds(random_seed=0)
+
+        # 2. initialization
+        init_for_distributed(opts)
+
+        # 3. visdom
+        vis = None
+
+        # 4. data set
+        transform_train = tfs.Compose([
+            tfs.Resize(256),
+            tfs.RandomCrop(224),
+            tfs.RandomHorizontalFlip(),
+            tfs.ToTensor(),
+            tfs.Normalize(mean=(0.4914, 0.4822, 0.4465),
+                          std=(0.2023, 0.1994, 0.2010)),
+        ])
+
+        transform_test = tfs.Compose([
+                                      tfs.Resize(256),
+                                      tfs.CenterCrop(224),
+                                      tfs.ToTensor(),
+                                      tfs.Normalize(mean=(0.4914, 0.4822, 0.4465),
+                                                    std=(0.2023, 0.1994, 0.2010)),
+                                            ])
+
+        train_set = CIFAR10(root=opts.root,
+                            train=True,
+                            transform=transform_train,
+                            download=True)
+
+        test_set = CIFAR10(root=opts.root,
+                           train=False,
+                           transform=transform_test,
+                           download=True)
+
+        train_sampler = DistributedSampler(dataset=train_set, shuffle=True)
+        test_sampler = DistributedSampler(dataset=test_set, shuffle=False)
+
+        train_loader = DataLoader(dataset=train_set,
+                                  batch_size=int(opts.batch_size / opts.world_size),
+                                  shuffle=False,
+                                  num_workers=int(opts.num_workers / opts.world_size),
+                                  sampler=train_sampler,
+                                  pin_memory=True)
+
+        test_loader = DataLoader(dataset=test_set,
+                                 batch_size=int(opts.batch_size / opts.world_size),
+                                 shuffle=False,
+                                 num_workers=int(opts.num_workers / opts.world_size),
+                                 sampler=test_sampler,
+                                 pin_memory=True)
+
+        # 5. model
+        model = vgg11(pretrained=False)
+        model = model.cuda(opts.local_rank)
+        model = DDP(module=model,
+                    device_ids=[opts.local_rank])
+
+        # 6. criterion
+        criterion = torch.nn.CrossEntropyLoss().to(opts.local_rank)
+
+        # 7. optimizer
+        optimizer = torch.optim.SGD(params=model.parameters(),
+                                    lr=0.01,
+                                    weight_decay=0.0005,
+                                    momentum=0.9)
+
+        # 8. scheduler
+        scheduler = StepLR(optimizer=optimizer,
+                           step_size=30,
+                           gamma=0.1)
+
+        if opts.start_epoch != 0:
+
+            checkpoint = torch.load(os.path.join(opts.save_path, opts.save_file_name) + '.{}.pth.tar'
+                                    .format(opts.start_epoch - 1),
+                                    map_location=torch.device('cuda:{}'.format(opts.local_rank)))
+            model.load_state_dict(checkpoint['model_state_dict'])
+            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+            scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+            if opts.global_rank == 0:
+                print('\nLoaded checkpoint from epoch %d.\n' % (int(opts.start_epoch) - 1))
+
+        for epoch in range(opts.start_epoch, opts.epoch):
+
+            # 9. train
+            tic = time.time()
+            model.train()
+            train_sampler.set_epoch(epoch)
+
+            for i, (images, labels) in enumerate(train_loader):
+                images = images.to(opts.local_rank)
+                labels = labels.to(opts.local_rank)
+                outputs = model(images)
+
+                # ----------- update -----------
+                optimizer.zero_grad()
+                loss = criterion(outputs, labels)
+                loss.backward()
+                optimizer.step()
+
+                # get lr
+                for param_group in optimizer.param_groups:
+                    lr = param_group['lr']
+
+                # time
+                toc = time.time()
+
+                # visualization
+                if (i % opts.vis_step == 0 or i == len(train_loader) - 1):
+                    print('GPU[{0}] Epoch [{1}/{2}], Iter [{3}/{4}], Loss: {5:.4f}, LR: {6:.5f}, Time: {7:.2f}'.format(opts.global_rank,
+                                                                                                                       epoch,
+                                                                                                                       opts.epoch,
+                                                                                                                       i,
+                                                                                                                       len(train_loader),
+                                                                                                                       loss.item(),
+                                                                                                                       lr,
+                                                                                                                       toc - tic))
+
+            # save pth file
+            if opts.local_rank == 0:
+                if not os.path.exists(opts.save_path):
+                    os.mkdir(opts.save_path)
+
+                checkpoint = {'epoch': epoch,
+                              'model_state_dict': model.state_dict(),
+                              'optimizer_state_dict': optimizer.state_dict(),
+                              'scheduler_state_dict': scheduler.state_dict()}
+
+                torch.save(checkpoint, os.path.join(opts.save_path, opts.save_file_name + '.{}.pth.tar'.format(epoch)))
+                print("save pth.tar {} epoch!".format(epoch))
+
+                # 10. test
+                model.eval()
+
+                val_avg_loss = 0
+                correct_top1 = 0
+                correct_top5 = 0
+                total = 0
+
+                with torch.no_grad():
+                    for i, (images, labels) in enumerate(test_loader):
+                        images = images.to(opts.local_rank)
+                        labels = labels.to(opts.local_rank)
+                        outputs = model(images)
+                        loss = criterion(outputs, labels)
+                        val_avg_loss += loss.item()
+                        # rank 1
+                        _, pred = torch.max(outputs, 1)
+                        total += labels.size(0)
+                        correct_top1 += (pred == labels).sum().item()
+
+                        # rank 5
+                        _, rank5 = outputs.topk(5, 1, True, True)
+                        rank5 = rank5.t()
+                        correct5 = rank5.eq(labels.view(1, -1).expand_as(rank5))
+
+                        for k in range(5):
+                            correct_k = correct5[:k+1].reshape(-1).float().sum(0, keepdim=True)
+                        correct_top5 += correct_k.item()
+
+                accuracy_top1 = correct_top1 / total
+                accuracy_top5 = correct_top5 / total
+
+                val_avg_loss = val_avg_loss / len(test_loader)
+
+                print("top-1 percentage :  {0:0.3f}%".format(correct_top1 / total * 100))
+                print("top-5 percentage :  {0:0.3f}%".format(correct_top5 / total * 100))
+                scheduler.step()
+
+        return 0
+
+
+    def init_for_distributed(opts):
+
+        # 1. setting for distributed training
+        opts.global_rank = int(os.environ['RANK'])
+        opts.local_rank = int(os.environ['LOCAL_RANK'])
+        opts.world_size = int(os.environ['WORLD_SIZE'])
+        torch.cuda.set_device(opts.local_rank)
+        if opts.global_rank is not None and opts.local_rank is not None:
+            print("Use GPU: [{}/{}] for training".format(opts.global_rank, opts.local_rank))
+
+        # 2. init_process_group
+        dist.init_process_group(
+            backend="nccl",
+            rank=opts.global_rank,
+            world_size=opts.world_size,
+            device_id=torch.device(f"cuda:{opts.local_rank}"),
+            timeout=datetime.timedelta(seconds=60)
+        )
+
+        return
+
+
+    if __name__ == '__main__':
+
+        parser = argparse.ArgumentParser('vgg11 cifar training', parents=[get_args_parser()])
+        opts = parser.parse_args()
+        main(opts)
+
+---
+# 1. Volcano Queue for training workloads
+# Queue는 리소스 할당 단위. capability는 클러스터 실제 용량에 맞게 조정 필요.
+apiVersion: scheduling.volcano.sh/v1beta1
+kind: Queue
+metadata:
+  name: training-queue
+spec:
+  weight: 1
+  reclaimable: true
+  capability:
+    # 3노드 x GPU 8장 = 24, hostdev 2개 = 6
+    cpu: "100"
+    memory: "512Gi"
+    nvidia.com/gpu: "24"
+    nvidia.com/hostdev: "6"
+
+---
+# 2. ClusterTrainingRuntime with Volcano gang scheduling + topology-aware scheduling
+# podGroupPolicy.volcano를 사용하면 PodGroup이 자동 생성됨 (수동 생성 불필요)
+# MASTER_ADDR / MASTER_PORT는 Kubeflow Trainer torch runtime이 자동 설정 (torchrun rdzv)
+apiVersion: trainer.kubeflow.org/v1alpha1
+kind: ClusterTrainingRuntime
+metadata:
+  name: torch-distributed-volcano
+  labels:
+    trainer.kubeflow.org/framework: torch
+spec:
+  mlPolicy:
+    torch:
+      numProcPerNode: 8
+    numNodes: 2
+  # Volcano gang scheduling 활성화 - PodGroup 자동 생성
+  podGroupPolicy:
+    volcano:
+      networkTopology:
+        mode: hard
+        highestTierAllowed: 1
+  template:
+    metadata:
+      annotations:
+        # Queue 지정 (runtime level)
+        scheduling.volcano.sh/queue-name: training-queue
+    spec:
+      replicatedJobs:
+      - name: node
+        template:
+          metadata:
+            labels:
+              trainer.kubeflow.org/trainjob-ancestor-step: trainer
+          spec:
+            template:
+              metadata:
+                annotations:
+                  # InfiniBand CNI
+                  k8s.v1.cni.cncf.io/networks: hostdevice-net
+              spec:
+                containers:
+                - name: node
+                  image: nvcr.io/nvidia/pytorch:24.10-py3
+                  securityContext:
+                    privileged: true
+                    capabilities:
+                      add:
+                        - IPC_LOCK
+                  env:
+                    - name: NCCL_DEBUG
+                      value: "INFO"
+                    - name: NCCL_DEBUG_SUBSYS
+                      value: "INIT,NET,IB"
+                    - name: NCCL_IB_DISABLE
+                      value: "0"
+                    - name: NCCL_SOCKET_IFNAME
+                      value: "net"
+                  resources:
+                    requests:
+                      nvidia.com/gpu: "8"
+                      nvidia.com/hostdev: "2"
+                    limits:
+                      nvidia.com/gpu: "8"
+                      nvidia.com/hostdev: "2"
+                  volumeMounts:
+                    - name: shared-memory
+                      mountPath: /dev/shm
+                    - name: cifar-data
+                      mountPath: /workspace/data/cifar-10-batches-py
+                    - name: training-scripts
+                      mountPath: /workspace/scripts
+                volumes:
+                  - name: shared-memory
+                    emptyDir:
+                      medium: Memory
+                      sizeLimit: 128Gi
+                  - name: cifar-data
+                    hostPath:
+                      path: /home/ubuntu/cifar-10-batches-py
+                      type: Directory
+                  - name: training-scripts
+                    configMap:
+                      name: training-scripts
+                      defaultMode: 0755
+
+---
+# 3. Example TrainJob (테스트용 - 2노드 분산학습)
+# podGroupPolicy, env, volumes, command 등은 runtime에서 상속됨
+apiVersion: trainer.kubeflow.org/v1alpha1
+kind: TrainJob
+metadata:
+  name: example-distributed-training
+  namespace: default
+spec:
+  runtimeRef:
+    name: torch-distributed-volcano
+  trainer:
+    image: nvcr.io/nvidia/pytorch:24.10-py3
+    command:
+      - torchrun
+      - /workspace/scripts/train_nccl.py
+    numNodes: 2
+    resourcesPerNode:
+      requests:
+        nvidia.com/gpu: "8"
+        nvidia.com/hostdev: "2"
+      limits:
+        nvidia.com/gpu: "8"
+        nvidia.com/hostdev: "2"
+