Revert "Revert "Remove NUM_WORKERS env var, auto-detect workers via DNS probing""

This reverts commit 5f6dec6799.

README.md

@@ -43,7 +43,7 @@ Kubeflow Trainer v2 (v1alpha1)에는 다음 런타임만 기본 제공된다:
 │                                                                     │
 │  필수 설정:                                                          │
 │    ├── network.publishNotReadyAddresses: true (DNS 조기 해석)        │
-│    └── NUM_WORKERS 환경변수 = mlPolicy.numNodes                      │
+│    └── 워커 수는 DNS 프로빙으로 자동 감지 (수동 설정 불필요)            │
 └─────────────────────────────────────────────────────────────────────┘
 ```
 
@@ -66,7 +66,16 @@ Worker 0: tf-distributed-training-node-0-0.tf-distributed-training.default.svc.c
 Worker 1: tf-distributed-training-node-1-0.tf-distributed-training.default.svc.cluster.local:12345
 ```
 
-학습 스크립트가 이를 파싱하여 자동으로 TF_CONFIG를 생성한다:
+학습 스크립트가 DNS 프로빙으로 워커를 자동 감지하고 TF_CONFIG를 생성한다:
+
+```
+discover_workers() 동작:
+  node-0 DNS 조회 → 성공 → workers에 추가
+  node-1 DNS 조회 → 성공 → workers에 추가
+  node-2 DNS 조회 → 실패 → 탐색 종료 → num_workers = 2
+```
+
+생성되는 TF_CONFIG:
 
 ```json
 {
@@ -147,24 +156,10 @@ kubectl apply -f tensorflow-volcano-trainjob-integration.yaml
 
 ## 주의사항
 
-### NUM_WORKERS와 numNodes 동기화
+### 워커 수 자동 감지 (DNS 프로빙)
 
-`NUM_WORKERS` 환경변수는 `mlPolicy.numNodes` (또는 TrainJob의 `trainer.numNodes`)와 **반드시 일치**해야 한다.
+학습 스크립트의 `discover_workers()` 함수가 DNS 프로빙으로 워커 수를 자동 감지한다.
-
+`NUM_WORKERS` 같은 환경변수를 수동으로 관리할 필요가 없으므로, TrainJob에서 `numNodes`만 변경하면 된다:
-PyTorch runtime은 컨트롤러가 `WORLD_SIZE`를 자동 주입하지만, 커스텀 TF runtime은 이를 수동으로 관리해야 한다.
-
-```yaml
-# Runtime에서 기본값 설정
-spec:
-  mlPolicy:
-    numNodes: 2      # ← 이 값과
-  ...
-  env:
-    - name: NUM_WORKERS
-      value: "2"     # ← 이 값이 일치해야 함
-```
-
-TrainJob에서 numNodes를 오버라이드할 경우, env도 함께 오버라이드해야 한다:
 
 ```yaml
 apiVersion: trainer.kubeflow.org/v1alpha1
@@ -173,12 +168,15 @@ spec:
   runtimeRef:
     name: tensorflow-distributed-volcano
   trainer:
-    numNodes: 4         # 4노드로 변경
+    numNodes: 4   # 이것만 변경하면 됨. 스크립트가 자동으로 4개 워커 감지.
-    env:
-      - name: NUM_WORKERS
-        value: "4"      # 반드시 함께 변경
 ```
 
+비교:
+| | PyTorch Runtime | TensorFlow Custom Runtime |
+|---|---|---|
+| 워커 수 감지 | 컨트롤러가 `WORLD_SIZE` 자동 주입 | 스크립트가 DNS 프로빙으로 자동 감지 |
+| numNodes 변경 시 | 추가 작업 없음 | 추가 작업 없음 |
+
 ### InfiniBand 설정
 
 클러스터에 InfiniBand가 없는 경우 다음 항목을 제거해야 한다:

tensorflow-custom-runtime.yaml

@@ -36,8 +36,6 @@ spec:
                 - name: node
                   image: tensorflow/tensorflow:2.16.1-gpu
                   env:
-                    - name: NUM_WORKERS
-                      value: "1"
                     - name: TF_WORKER_PORT
                       value: "12345"
                     - name: NAMESPACE

tensorflow-volcano-trainjob-integration.yaml

@@ -25,6 +25,42 @@ data:
     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를 자동 구성.
@@ -32,13 +68,10 @@ data:
         Pod hostname 패턴: {trainjob-name}-node-{replica_index}-{completion_index}
         Headless Service DNS: {hostname}.{trainjob-name}.{namespace}.svc.cluster.local
 
-        환경변수:
+        워커 수는 DNS 프로빙으로 자동 감지하므로 NUM_WORKERS 환경변수 불필요.
-          - NUM_WORKERS: 총 워커 수 (runtime의 numNodes와 일치해야 함)
+        numNodes가 변경되어도 학습 스크립트 수정 없이 자동 대응.
-          - TF_WORKER_PORT: gRPC 통신 포트 (기본값: 12345)
-          - NAMESPACE: Pod이 실행되는 namespace (Downward API)
         """
         hostname = socket.gethostname()
-        num_workers = int(os.environ.get('NUM_WORKERS', '1'))
         port = os.environ.get('TF_WORKER_PORT', '12345')
         namespace = os.environ.get('NAMESPACE', 'default')
 
@@ -53,14 +86,10 @@ data:
             job_name = os.environ.get('TRAINJOB_NAME', 'unknown')
             worker_index = 0
 
-        # Headless Service DNS를 이용한 worker 주소 목록 생성
+        # DNS 프로빙으로 워커 목록 자동 감지
-        workers = []
+        workers = discover_workers(job_name, namespace, port)
-        for i in range(num_workers):
+        num_workers = len(workers)
-            worker_host = (
+        print(f"[Discovery] Discovered {num_workers} workers via DNS")
-                f"{job_name}-node-{i}-0.{job_name}.{namespace}"
-                f".svc.cluster.local:{port}"
-            )
-            workers.append(worker_host)
 
         tf_config = {
             "cluster": {"worker": workers},
@@ -200,7 +229,7 @@ spec:
 #   - mlPolicy에 torch/mpi 필드 없음 → 컨트롤러가 torchrun/MPI 주입하지 않음
 #   - network.publishNotReadyAddresses: true → Pod간 DNS 조기 해석 가능
 #   - TF_CONFIG는 학습 스크립트에서 hostname + Headless Service DNS로 자동 구성
-#   - NUM_WORKERS 환경변수가 numNodes와 반드시 일치해야 함
+#   - 워커 수는 DNS 프로빙으로 자동 감지 (NUM_WORKERS 환경변수 불필요)
 apiVersion: trainer.kubeflow.org/v1alpha1
 kind: ClusterTrainingRuntime
 metadata:
@@ -258,9 +287,7 @@ spec:
                     - name: NCCL_SOCKET_IFNAME
                       value: "net"
                     # --- TF_CONFIG 자동 구성용 환경변수 ---
-                    # NUM_WORKERS는 mlPolicy.numNodes와 반드시 일치해야 함
+                    # 워커 수는 DNS 프로빙으로 자동 감지 (NUM_WORKERS 불필요)
-                    - name: NUM_WORKERS
-                      value: "2"
                     - name: TF_WORKER_PORT
                       value: "12345"
                     - name: NAMESPACE