kubeflow-trainer-v2-tf-custom-runtime/README.md

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 조기 해석)        │
│    └── 워커 수는 DNS 프로빙으로 자동 감지 (수동 설정 불필요)            │
└─────────────────────────────────────────────────────────────────────┘

TF_CONFIG 자동 구성 원리

Trainer v2에서 TrainJob을 생성하면 내부적으로 JobSet이 만들어지고, 각 Pod은 예측 가능한 hostname을 갖는다:

mlPolicy에 framework(torch/mpi) 없이 numNodes만 설정하면:
  → 단일 Job에 completions=numNodes 로 매핑
  → Pod hostname 패턴: {trainjob}-node-0-{completion_index}
  → completion_index = worker_index

참고) mlPolicy.torch 사용 시:
  → replica별 별도 Job 생성
  → Pod hostname 패턴: {trainjob}-node-{replica_index}-0

예시 (TrainJob: tf-distributed-training, numNodes: 2):
  Worker 0: tf-distributed-training-node-0-0   (completion_index=0)
  Worker 1: tf-distributed-training-node-0-1   (completion_index=1)

Headless Service DNS와 결합하면:

Worker 0: tf-distributed-training-node-0-0.tf-distributed-training.{namespace}.svc.cluster.local:12345
Worker 1: tf-distributed-training-node-0-1.tf-distributed-training.{namespace}.svc.cluster.local:12345

학습 스크립트가 DNS 프로빙으로 워커를 자동 감지하고 TF_CONFIG를 생성한다:

discover_workers() 동작:
  node-0-0 DNS 조회 → 성공 → workers에 추가
  node-0-1 DNS 조회 → 성공 → workers에 추가
  node-0-2 DNS 조회 → 실패 → 탐색 종료 → num_workers = 2

중요: TF_CONFIG는 반드시 import tensorflow 전에 설정해야 한다. TF는 import 시점에 GPU를 감지하고 TF_CONFIG가 있으면 gRPC 서버를 시작한다. import 후에 TF_CONFIG를 설정하면 "different incarnation" 에러가 발생한다.

생성되는 TF_CONFIG:

{
  "cluster": {
    "worker": [
      "tf-distributed-training-node-0-0.tf-distributed-training.{namespace}.svc.cluster.local:12345",
      "tf-distributed-training-node-0-1.tf-distributed-training.{namespace}.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 런타임만 필요한 경우:

kubectl apply -f tensorflow-custom-runtime.yaml

확인:

kubectl get clustertrainingruntime
# tensorflow-distributed 가 목록에 나타나야 함

2단계: Volcano 연동 전체 적용

Volcano gang scheduling + 토폴로지 인식 스케줄링이 포함된 전체 버전:

# ConfigMap (학습 스크립트) + Queue + Runtime + TrainJob 한번에 적용
kubectl apply -f tensorflow-volcano-trainjob-integration.yaml

확인:

# 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만 변경할 경우:

# 기존 TrainJob 삭제
kubectl delete trainjob tf-distributed-training

# 수정 후 재적용
kubectl apply -f tensorflow-volcano-trainjob-integration.yaml

주의사항

워커 수 자동 감지 (DNS 프로빙)

학습 스크립트의 discover_workers() 함수가 DNS 프로빙으로 워커 수를 자동 감지한다. NUM_WORKERS 같은 환경변수를 수동으로 관리할 필요가 없으므로, TrainJob에서 numNodes만 변경하면 된다:

apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
spec:
  runtimeRef:
    name: tensorflow-distributed-volcano
  trainer:
    numNodes: 4   # 이것만 변경하면 됨. 스크립트가 자동으로 4개 워커 감지.

비교:

	PyTorch Runtime	TensorFlow Custom Runtime
워커 수 감지	컨트롤러가 WORLD_SIZE 자동 주입	스크립트가 DNS 프로빙으로 자동 감지
numNodes 변경 시	추가 작업 없음	추가 작업 없음

InfiniBand 설정

클러스터에 InfiniBand가 없는 경우 다음 항목을 제거해야 한다:

# 제거 대상:
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가 해석되어야 하므로 이 설정이 필수다:

spec:
  template:
    spec:
      network:
        publishNotReadyAddresses: true   # 반드시 true

PyTorch runtime에서는 torchrun의 rdzv(rendezvous) 메커니즘이 이를 처리하므로 불필요하다.