415 lines
15 KiB
YAML
415 lines
15 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를 tensorflow import 전에 설정 (필수)
|
|
"""
|
|
import os
|
|
import json
|
|
import socket
|
|
import re
|
|
import time
|
|
import pickle
|
|
|
|
# =============================================================
|
|
# 1단계: TF_CONFIG 설정 (tensorflow import 전에 반드시 실행)
|
|
#
|
|
# TensorFlow는 import 시점에 GPU를 감지하고, TF_CONFIG가 있으면
|
|
# gRPC 서버를 시작한다. import 후에 TF_CONFIG를 설정하면
|
|
# "different incarnation" 에러가 발생한다.
|
|
# =============================================================
|
|
|
|
def discover_workers(job_name, namespace, port, max_retries=60, retry_interval=5):
|
|
"""
|
|
DNS 프로빙으로 워커 수를 자동 감지.
|
|
|
|
동작 방식:
|
|
mlPolicy에 framework 없으면 numNodes가 completions로 매핑됨.
|
|
Pod hostname 패턴: {trainjob}-node-0-{completion_index}
|
|
- node-0-0, node-0-1, node-0-2, ... 순서로 DNS 조회
|
|
- 해석되지 않는 인덱스가 나오면 그 직전까지를 워커 목록으로 사용.
|
|
Volcano gang scheduling에서도 DNS 전파 지연이 있으므로 retry 포함.
|
|
"""
|
|
prev_count = 0
|
|
stable_rounds = 0
|
|
|
|
for attempt in range(max_retries):
|
|
workers = []
|
|
i = 0
|
|
while True:
|
|
# numNodes → completions 매핑이므로 node-0-{i} 패턴
|
|
worker_host = (
|
|
f"{job_name}-node-0-{i}.{job_name}"
|
|
f".{namespace}.svc.cluster.local"
|
|
)
|
|
try:
|
|
socket.getaddrinfo(worker_host, int(port))
|
|
workers.append(f"{worker_host}:{port}")
|
|
i += 1
|
|
except socket.gaierror:
|
|
break
|
|
|
|
current_count = len(workers)
|
|
|
|
# 2개 이상 발견 + 연속 2회 동일하면 안정된 것으로 판단
|
|
if current_count >= 2:
|
|
if current_count == prev_count:
|
|
stable_rounds += 1
|
|
else:
|
|
stable_rounds = 0
|
|
if stable_rounds >= 2:
|
|
return workers
|
|
|
|
prev_count = current_count
|
|
print(f"[Discovery] attempt {attempt+1}/{max_retries}: "
|
|
f"found {current_count} workers, retrying in "
|
|
f"{retry_interval}s...")
|
|
time.sleep(retry_interval)
|
|
|
|
return workers
|
|
|
|
|
|
def setup_tf_config():
|
|
"""
|
|
Kubeflow Trainer v2의 Pod 네이밍 규칙을 이용하여 TF_CONFIG 설정.
|
|
반드시 import tensorflow 전에 호출해야 한다.
|
|
|
|
mlPolicy에 framework(torch/mpi) 없이 numNodes만 설정하면:
|
|
- 단일 Job에 completions=numNodes 로 매핑됨
|
|
- Pod hostname: {trainjob}-node-0-{completion_index}
|
|
- completion_index가 worker_index
|
|
Headless Service DNS: {hostname}.{trainjob}.{namespace}.svc.cluster.local
|
|
"""
|
|
hostname = socket.gethostname()
|
|
port = os.environ.get('TF_WORKER_PORT', '12345')
|
|
namespace = os.environ.get('NAMESPACE', 'default')
|
|
|
|
# hostname 패턴: {trainjob}-node-0-{completion_index}
|
|
# completion_index = worker_index
|
|
match = re.match(r'^(.+)-node-(\d+)-(\d+)$', hostname)
|
|
if match:
|
|
job_name = match.group(1)
|
|
worker_index = int(match.group(3)) # completion_index가 worker
|
|
else:
|
|
print(f"[WARNING] Cannot parse hostname '{hostname}'")
|
|
job_name = os.environ.get('TRAINJOB_NAME', 'unknown')
|
|
worker_index = 0
|
|
|
|
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
|
|
|
|
|
|
# TF_CONFIG 설정 후 tensorflow import
|
|
_worker_index, _num_workers = setup_tf_config()
|
|
|
|
import tensorflow as tf
|
|
import numpy as np
|
|
|
|
# =============================================================
|
|
# 2단계: 데이터 로딩 + 모델 정의 + 학습
|
|
# =============================================================
|
|
|
|
def load_cifar10_local(data_dir='/workspace/data/cifar-10-batches-py'):
|
|
"""
|
|
로컬 hostPath에서 CIFAR-10 데이터 로드.
|
|
tf.keras.datasets.cifar10.load_data()는 외부 다운로드를 시도하므로
|
|
에어갭 환경에서는 이 함수를 사용.
|
|
"""
|
|
x_train_list, y_train_list = [], []
|
|
for i in range(1, 6):
|
|
path = os.path.join(data_dir, f'data_batch_{i}')
|
|
with open(path, 'rb') as f:
|
|
batch = pickle.load(f, encoding='bytes')
|
|
x_train_list.append(batch[b'data'])
|
|
y_train_list.append(batch[b'labels'])
|
|
|
|
x_train = (np.concatenate(x_train_list)
|
|
.reshape(-1, 3, 32, 32)
|
|
.transpose(0, 2, 3, 1))
|
|
y_train = np.array(
|
|
np.concatenate(y_train_list), dtype=np.int64
|
|
).reshape(-1, 1)
|
|
|
|
with open(os.path.join(data_dir, 'test_batch'), 'rb') as f:
|
|
batch = pickle.load(f, encoding='bytes')
|
|
x_test = (batch[b'data']
|
|
.reshape(-1, 3, 32, 32)
|
|
.transpose(0, 2, 3, 1))
|
|
y_test = np.array(
|
|
batch[b'labels'], dtype=np.int64
|
|
).reshape(-1, 1)
|
|
|
|
return (x_train, y_train), (x_test, y_test)
|
|
|
|
|
|
def build_model():
|
|
"""CIFAR-10용 CNN 모델"""
|
|
return 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)
|
|
])
|
|
|
|
|
|
def main():
|
|
worker_index, num_workers = _worker_index, _num_workers
|
|
|
|
# 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 데이터셋 로드 (로컬 hostPath)
|
|
(x_train, y_train), (x_test, y_test) = load_cifar10_local()
|
|
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
|
|
|
|
model.fit(
|
|
train_dataset,
|
|
epochs=EPOCHS,
|
|
validation_data=test_dataset,
|
|
verbose=verbose,
|
|
)
|
|
|
|
# 평가 (모든 워커가 참여해야 함 - MultiWorkerMirroredStrategy 요구사항)
|
|
eval_verbose = 2 if worker_index == 0 else 0
|
|
loss, accuracy = model.evaluate(
|
|
test_dataset, verbose=eval_verbose
|
|
)
|
|
|
|
if worker_index == 0:
|
|
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: false
|
|
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: 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: 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"
|