345 lines
12 KiB
YAML
345 lines
12 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 discover_workers(job_name, namespace, port, max_retries=60, retry_interval=5):
|
|
"""
|
|
DNS 프로빙으로 워커 수를 자동 감지.
|
|
NUM_WORKERS 환경변수 없이도 numNodes 변경에 자동 대응.
|
|
|
|
동작 방식:
|
|
node-0, node-1, node-2, ... 순서로 DNS를 조회하여
|
|
해석되지 않는 인덱스가 나오면 그 직전까지를 워커 목록으로 사용.
|
|
모든 워커가 아직 생성 전일 수 있으므로 retry 포함.
|
|
"""
|
|
import time
|
|
|
|
workers = []
|
|
for attempt in range(max_retries):
|
|
workers = []
|
|
i = 0
|
|
while True:
|
|
worker_host = f"{job_name}-node-{i}-0.{job_name}.{namespace}.svc.cluster.local"
|
|
try:
|
|
socket.getaddrinfo(worker_host, int(port))
|
|
workers.append(f"{worker_host}:{port}")
|
|
i += 1
|
|
except socket.gaierror:
|
|
break
|
|
|
|
if len(workers) >= 2 or (len(workers) == 1 and attempt >= 5):
|
|
# 최소 2개 워커 발견 또는 단일 노드 학습으로 판단
|
|
break
|
|
|
|
print(f"[Discovery] attempt {attempt+1}/{max_retries}: "
|
|
f"found {len(workers)} workers, retrying in {retry_interval}s...")
|
|
time.sleep(retry_interval)
|
|
|
|
return workers
|
|
|
|
|
|
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
|
|
|
|
워커 수는 DNS 프로빙으로 자동 감지하므로 NUM_WORKERS 환경변수 불필요.
|
|
numNodes가 변경되어도 학습 스크립트 수정 없이 자동 대응.
|
|
"""
|
|
hostname = socket.gethostname()
|
|
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
|
|
|
|
# DNS 프로빙으로 워커 목록 자동 감지
|
|
workers = discover_workers(job_name, namespace, port)
|
|
num_workers = len(workers)
|
|
print(f"[Discovery] Discovered {num_workers} workers via DNS")
|
|
|
|
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로 자동 구성
|
|
# - 워커 수는 DNS 프로빙으로 자동 감지 (NUM_WORKERS 환경변수 불필요)
|
|
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 자동 구성용 환경변수 ---
|
|
# 워커 수는 DNS 프로빙으로 자동 감지 (NUM_WORKERS 불필요)
|
|
- 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
|
|
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"
|