selee
/
kubeflow-trainer-v2-tf-custom-runtime
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix Pod naming pattern and TF_CONFIG initialization order 

1. Pod hostname 패턴 수정: node-{i}-0 → node-0-{i}
   mlPolicy에 framework 없으면 numNodes가 completions로 매핑됨
   (단일 Job, completions=numNodes)

2. TF_CONFIG를 import tensorflow 전에 설정
   TF는 import 시 GPU 감지 + gRPC 서버 시작하므로
   TF_CONFIG가 없는 상태에서 import하면 incarnation 에러 발생

3. CIFAR-10 로컬 hostPath 볼륨 마운트 추가
   tf.keras.datasets.cifar10.load_data() 대신 로컬 pickle 파일 로드

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Cloud User 2026-03-03 15:07:57 +09:00
parent a96f0fbc22
commit 9b1cd02a87
3 changed files with 134 additions and 449 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
README.md
View File

@ -52,37 +52,48 @@ Kubeflow Trainer v2 (v1alpha1)에는 다음 런타임만 기본 제공된다:
Trainer v2에서 TrainJob을 생성하면 내부적으로 JobSet이 만들어지고, 각 Pod은 예측 가능한 hostname을 갖는다:
```
Pod hostname 패턴: {trainjob-name}-node-{replica_index}-{completion_index}
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
Worker 1: tf-distributed-training-node-1-0
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.default.svc.cluster.local:12345
Worker 1: tf-distributed-training-node-1-0.tf-distributed-training.default.svc.cluster.local:12345
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 DNS 조회 → 성공 → workers에 추가
node-1 DNS 조회 → 성공 → workers에 추가
node-2 DNS 조회 → 실패 → 탐색 종료 → num_workers = 2
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:
```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"
"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}

160
tensorflow-volcano-trainjob-integration.yaml
View File

@ -15,34 +15,46 @@ data:
"""
TensorFlow Distributed Training Example for Kubeflow Trainer v2
- CNN (CIFAR-10) + MultiWorkerMirroredStrategy + NCCL
- TF_CONFIG를 Pod hostname과 환경변수로 자동 구성
- TF_CONFIG를 tensorflow import 전에 설정 (필수)
"""
import os
import json
import socket
import re
import tensorflow as tf
import numpy as np
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 프로빙으로 워커 수를 자동 감지.
NUM_WORKERS 환경변수 없이도 numNodes 변경에 자동 대응.
동작 방식:
node-0, node-1, node-2, ... 순서로 DNS를 조회하여
해석되지 않는 인덱스가 나오면 그 직전까지를 워커 목록으로 사용.
모든 워커가 아직 생성 전일 수 있으므로 retry 포함.
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 포함.
"""
import time
prev_count = 0
stable_rounds = 0
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"
# 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}")
@ -50,12 +62,21 @@ data:
except socket.gaierror:
break
if len(workers) >= 2 or (len(workers) == 1 and attempt >= 5):
# 최소 2개 워커 발견 또는 단일 노드 학습으로 판단
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 {len(workers)} workers, retrying in {retry_interval}s...")
f"found {current_count} workers, retrying in "
f"{retry_interval}s...")
time.sleep(retry_interval)
return workers
@ -63,30 +84,30 @@ data:
def setup_tf_config():
"""
Kubeflow Trainer v2의 Pod 네이밍 규칙을 이용하여 TF_CONFIG를 자동 구성.
Kubeflow Trainer v2의 Pod 네이밍 규칙을 이용하여 TF_CONFIG 설정.
반드시 import tensorflow 전에 호출해야 한다.
Pod hostname 패턴: {trainjob-name}-node-{replica_index}-{completion_index}
Headless Service DNS: {hostname}.{trainjob-name}.{namespace}.svc.cluster.local
워커 수는 DNS 프로빙으로 자동 감지하므로 NUM_WORKERS 환경변수 불필요.
numNodes가 변경되어도 학습 스크립트 수정 없이 자동 대응.
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 이름과 replica index 추출
# 패턴: {trainjob-name}-node-{replica_index}-{completion_index}
# 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(2))
worker_index = int(match.group(3)) # completion_index가 worker
else:
print(f"[WARNING] Cannot parse hostname '{hostname}', using defaults")
print(f"[WARNING] Cannot parse hostname '{hostname}'")
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")
@ -104,9 +125,52 @@ data:
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 모델"""
model = tf.keras.Sequential([
return tf.keras.Sequential([
tf.keras.layers.Conv2D(
32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
tf.keras.layers.BatchNormalization(),
@ -126,15 +190,16 @@ data:
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(10)
])
return model
def main():
worker_index, num_workers = setup_tf_config()
worker_index, num_workers = _worker_index, _num_workers
# MultiWorkerMirroredStrategy (GPU간 NCCL 통신)
communication_options = tf.distribute.experimental.CommunicationOptions(
implementation=tf.distribute.experimental.CommunicationImplementation.NCCL
implementation=(
tf.distribute.experimental
.CommunicationImplementation.NCCL)
)
strategy = tf.distribute.MultiWorkerMirroredStrategy(
communication_options=communication_options
@ -143,10 +208,8 @@ data:
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()
)
# 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
@ -157,13 +220,17 @@ data:
)
# 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)
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():
@ -179,18 +246,11 @@ data:
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만 출력)
@ -307,6 +367,8 @@ spec:
volumeMounts:
- name: shared-memory
mountPath: /dev/shm
- name: cifar-data
mountPath: /workspace/data/cifar-10-batches-py
- name: training-scripts
mountPath: /workspace/scripts
volumes:
@ -314,6 +376,10 @@ spec:
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

392
volcano-trainjob-integration.yaml
View File

@ -1,392 +0,0 @@
# 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"