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