# 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
  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"