Finally is alive!!

Add support for chief (coordinator) node in TensorFlow distributed training

The actual train_tf_ps.py is only a placeholder for development testings and  must be refined and correctly commented for a real workload.

- Extend `train_tf_ps.py` to include a chief address and port in cluster configuration.
- Update `make_parameter_server_strategy` to validate and configure the chief node, including `TF_CONFIG` setup.
- Modify `run_tf_training_from_bastion.sh` to detect and validate a routable IPv4 address for the chief node.
- Adjust `.gitignore` to exclude `output/`.
- Add gRPC port mapping for the chief node in `docker-compose.yml`.
- Update model saving path in `train_tf_ps.py` to include `.keras` extension.

Author: greg-ogs

.gitignore

@@ -3,4 +3,5 @@
 *.terraform*
 *tfstate*
 
-*config-kind-in-container*
+*config-kind-in-container*
+*output/

infra/local/external_workloads/docker-compose.yml

@@ -33,6 +33,8 @@ services:
       - ../tf/config-kind-in-container:/root/.kube/config:ro
     environment:
       - KUBECONFIG=/root/.kube/config
+    ports:
+      - "2223:2223"   # TensorFlow coordinator gRPC port
     command: [ "bash", "-lc", "sleep infinity" ]
     networks:
       - kind

infra/local/tf/Dockerfile

@@ -1,5 +1,6 @@
 FROM tensorflow/tensorflow:latest-gpu
 LABEL authors="greg-ogs"
+# For bastion container outside the kubernetes cluster
 
 RUN apt-get update && apt-get install -y apt-transport-https ca-certificates curl gnupg
 

workloads/local_cluster_workloads/run_tf_training_from_bastion.sh

@@ -1,5 +1,6 @@
 #!/usr/bin/env sh
-set -eu
+set -e
+set -u
 
 # This script is intended to run inside the bastion container defined in infra/local/external_workloads/docker-compose.yml
 # It discovers LoadBalancer IPs for TF worker/ps services, ensures TensorFlow is installed,
@@ -13,6 +14,35 @@ DATA_PATH=${DATA_PATH:-/data/health.csv}
 OUTPUT_DIR=${OUTPUT_DIR:-/workloads/output/$(date +%Y%m%d_%H%M%S)}
 EPOCHS=${EPOCHS:-3}
 BATCH_SIZE=${BATCH_SIZE:-64}
+CHIEF_PORT=${CHIEF_PORT:-2223}
+# Determine a routable IPv4 for the coordinator (bastion). Respect CHIEF_ADDR if provided.
+if [ -z "${CHIEF_ADDR:-}" ]; then
+  # Try to detect default IPv4 via routing table
+  if command -v ip >/dev/null 2>&1; then
+    CANDIDATE=$(ip -4 route get 8.8.8.8 2>/dev/null | awk '{for (i=1;i<=NF;i++) if ($i=="src") {print $(i+1); exit}}')
+    if [ -n "$CANDIDATE" ]; then
+      CHIEF_ADDR="$CANDIDATE"
+    fi
+  fi
+fi
+if [ -z "${CHIEF_ADDR:-}" ]; then
+  # Fallback: hostname -I (capital i) and pick first IPv4 token
+  if command -v hostname >/dev/null 2>&1; then
+    for tok in $(hostname -I 2>/dev/null); do
+      case "$tok" in
+        *:*) ;; # skip IPv6 tokens
+        *.*) CHIEF_ADDR="$tok"; break ;;
+      esac
+    done
+  fi
+fi
+# Final validation: must be IPv4 (avoid IPv6 which breaks TF_CONFIG host:port handling)
+if [ -z "${CHIEF_ADDR:-}" ] || echo "$CHIEF_ADDR" | grep -q ":"; then
+  echo "Unable to auto-detect a valid IPv4 CHIEF_ADDR. Set CHIEF_ADDR to an IPv4 reachable from K8s pods (e.g., 172.x/192.168.x)." >&2
+  exit 4
+fi
+
+echo "Chief (coordinator) will advertise ${CHIEF_ADDR}:${CHIEF_PORT}"
 
 # Ensure kubectl available
 if ! command -v kubectl >/dev/null 2>&1; then
@@ -74,6 +104,19 @@ fi
 
 mkdir -p "$OUTPUT_DIR"
 
+# Ensure gRPC to chief bypasses proxies (avoid unexpected :443 redirection)
+if [ -n "${no_proxy:-}" ]; then
+  no_proxy="${no_proxy},${CHIEF_ADDR}"
+else
+  no_proxy="${CHIEF_ADDR}"
+fi
+if [ -n "${NO_PROXY:-}" ]; then
+  NO_PROXY="${NO_PROXY},${CHIEF_ADDR}"
+else
+  NO_PROXY="${CHIEF_ADDR}"
+fi
+export no_proxy NO_PROXY
+
 "$PYTHON" /workloads/local_cluster_workloads/train_tf_ps.py \
   --data-path "$DATA_PATH" \
   --output-dir "$OUTPUT_DIR" \
@@ -83,4 +126,6 @@ mkdir -p "$OUTPUT_DIR"
   --worker-replicas "$WORKER_COUNT" \
   --ps-replicas "$PS_COUNT" \
   --worker-addrs "$WORKER_ADDRS_CSV" \
-  --ps-addrs "$PS_ADDRS_CSV"
+  --ps-addrs "$PS_ADDRS_CSV" \
+  --chief-addr "$CHIEF_ADDR" \
+  --chief-port "$CHIEF_PORT"

workloads/local_cluster_workloads/train_tf_ps.py

@@ -7,6 +7,10 @@
 import time
 from typing import List, Tuple, Optional
 
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning, module='google.protobuf.runtime_version')
+
 import numpy as np
 import tensorflow as tf
 from urllib.request import urlopen
@@ -100,7 +104,7 @@ def build_sequential_model(input_dim: int, num_classes: int) -> tf.keras.Model:
 # Distributed strategy helpers
 # ----------------------------
 
-def build_cluster_def(worker_replicas: int, ps_replicas: int, port: int, worker_addrs: Optional[List[str]] = None, ps_addrs: Optional[List[str]] = None) -> dict:
+def build_cluster_def(worker_replicas: int, ps_replicas: int, port: int, worker_addrs: Optional[List[str]] = None, ps_addrs: Optional[List[str]] = None, chief_addr: Optional[str] = None, chief_port: int = 2223) -> dict:
     # If explicit addresses provided (e.g., from bastion via LoadBalancer IPs), use them
     if worker_addrs:
         workers = worker_addrs
@@ -113,13 +117,43 @@ def build_cluster_def(worker_replicas: int, ps_replicas: int, port: int, worker_
         else:
             ps = [f"tf-trainer-ps-{i}.tf-trainer-ps-headless:{port}" for i in range(ps_replicas)]
         cluster_def["ps"] = ps
+    # Include chief if provided (must be routable from the K8s pods)
+    if chief_addr:
+        cluster_def["chief"] = [f"{chief_addr}:{chief_port}"]
     return cluster_def
 
 
-def make_parameter_server_strategy(worker_replicas: int, ps_replicas: int, port: int = 2222, worker_addrs: Optional[List[str]] = None, ps_addrs: Optional[List[str]] = None) -> tf.distribute.ParameterServerStrategy:
-    cluster_def = build_cluster_def(worker_replicas, ps_replicas, port, worker_addrs, ps_addrs)
+def make_parameter_server_strategy(worker_replicas: int, ps_replicas: int, port: int = 2222, worker_addrs: Optional[List[str]] = None, ps_addrs: Optional[List[str]] = None, chief_addr: Optional[str] = None, chief_port: int = 2223) -> tf.distribute.ParameterServerStrategy:
+    cluster_def = build_cluster_def(worker_replicas, ps_replicas, port, worker_addrs, ps_addrs, chief_addr, chief_port)
     print("Computed ClusterSpec:", json.dumps(cluster_def), flush=True)
 
+    # Basic validation and sanitization for chief address to avoid IPv6 and malformed inputs
+    if chief_addr:
+        # Reject IPv6 literals or bracketed addresses and any scheme prefixes
+        if ":" in chief_addr and "." not in chief_addr:
+            raise RuntimeError(
+                f"chief_addr appears to be IPv6 ('{chief_addr}'). Please provide an IPv4 address reachable from K8s pods."
+            )
+        if any(sym in chief_addr for sym in ["/", "[", "]", " "]):
+            raise RuntimeError(
+                f"chief_addr '{chief_addr}' is malformed. Provide a raw IPv4 like 192.168.1.10 without scheme or brackets."
+            )
+        # Optional strict IPv4 check
+        parts = chief_addr.split(".")
+        if len(parts) != 4 or any(not p.isdigit() or not (0 <= int(p) <= 255) for p in parts):
+            raise RuntimeError(
+                f"chief_addr '{chief_addr}' is not a valid IPv4 address."
+            )
+
+    # If a chief address is provided, declare this process as chief via TF_CONFIG
+    if chief_addr:
+        tf_config = {
+            "cluster": cluster_def,
+            "task": {"type": "chief", "index": 0},
+        }
+        os.environ["TF_CONFIG"] = json.dumps(tf_config)
+        print("TF_CONFIG set:", os.environ["TF_CONFIG"], flush=True)
+
     resolver = tf.distribute.cluster_resolver.SimpleClusterResolver(
         cluster_spec=tf.train.ClusterSpec(cluster_def),
         rpc_layer="grpc",
@@ -148,6 +182,8 @@ def run_training(
     port: int = 2222,
     worker_addrs: Optional[List[str]] = None,
     ps_addrs: Optional[List[str]] = None,
+    chief_addr: Optional[str] = None,
+    chief_port: int = 2223,
 ) -> None:
     os.makedirs(output_dir, exist_ok=True)
 
@@ -173,7 +209,7 @@ def run_training(
             print("Worker addrs:", worker_addrs)
         if ps_addrs:
             print("PS addrs:", ps_addrs)
-        strategy = make_parameter_server_strategy(worker_replicas, ps_replicas, port, worker_addrs, ps_addrs)
+        strategy = make_parameter_server_strategy(worker_replicas, ps_replicas, port, worker_addrs, ps_addrs, chief_addr, chief_port)
 
         # Switch to ClusterCoordinator-based custom training loop (DatasetCreator removed)
         def per_worker_dataset_fn(input_context: Optional[tf.distribute.InputContext] = None):
@@ -237,7 +273,7 @@ def step_fn(inputs):
         history = model.fit(ds, epochs=epochs, steps_per_epoch=steps_per_epoch)
 
     # Save model
-    save_path = os.path.join(output_dir, "model")
+    save_path = os.path.join(output_dir, "model.keras")
     model.save(save_path)
     print(f"Model saved to: {save_path}")
 
@@ -250,7 +286,7 @@ def parse_args(argv: List[str]):
     parser = argparse.ArgumentParser(description="Train TF Keras model on health.csv with optional ParameterServerStrategy")
     parser.add_argument("--data-path", default=os.environ.get("DATA_PATH", "infra/local/mysql-database/health.csv"), help="Path to CSV (when running on bastion/host)")
     parser.add_argument("--data-url", default=os.environ.get("DATA_URL", ""), help="HTTP(S) URL to CSV (used inside cluster if path not mounted)")
-    parser.add_argument("--output-dir", default=os.environ.get("OUTPUT_DIR", "/tmp/tf-model"))
+    parser.add_argument("--output-dir", default=os.environ.get("OUTPUT_DIR", "./tf-model"))
     parser.add_argument("--epochs", type=int, default=int(os.environ.get("EPOCHS", "3")))
     parser.add_argument("--batch-size", type=int, default=int(os.environ.get("BATCH_SIZE", "64")))
     parser.add_argument("--use-ps", action="store_true", help="Enable ParameterServerStrategy coordinator mode")
@@ -259,12 +295,14 @@ def parse_args(argv: List[str]):
     parser.add_argument("--port", type=int, default=int(os.environ.get("TF_GRPC_PORT", "2222")))
     parser.add_argument("--worker-addrs", default=os.environ.get("WORKER_ADDRS", ""), help="Comma-separated worker addresses (host:port) when running outside cluster")
     parser.add_argument("--ps-addrs", default=os.environ.get("PS_ADDRS", ""), help="Comma-separated ps addresses (host:port) when running outside cluster")
+    parser.add_argument("--chief-addr", default=os.environ.get("CHIEF_ADDR", ""), help="Routable IPv4 address of the coordinator (bastion) accessible from K8s pods")
+    parser.add_argument("--chief-port", type=int, default=int(os.environ.get("CHIEF_PORT", "2223")), help="Coordinator gRPC port (exposed on tf-bastion)")
     return parser.parse_args(argv)
 
 
 if __name__ == "__main__":
     args = parse_args(sys.argv[1:])
-
+    input("Press enter to continue...")
     # Resolve data source: prefer local path; if not existent and data-url provided -> use URL
     data_source = args.data_path
     if not os.path.exists(data_source):
@@ -288,4 +326,6 @@ def parse_args(argv: List[str]):
         port=args.port,
         worker_addrs=worker_addrs,
         ps_addrs=ps_addrs,
+        chief_addr=(args.chief_addr if args.chief_addr else None),
+        chief_port=args.chief_port,
     )