319 lines
11 KiB
YAML
319 lines
11 KiB
YAML
# 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"
|