Switch train_tf_ps.py to ClusterCoordinator custom training.

Problems with the callback of chief

Author: greg-ogs

workloads/local_cluster_workloads/run_tf_training_from_bastion.sh

@@ -1,13 +1,13 @@
-#!/usr/bin/env bash
-set -euo pipefail
+#!/usr/bin/env sh
+set -eu
 
 # 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,
 # and launches the training coordinator that parallelizes across the K8s pods.
 
 # Config
-WORKER_SERVICES=(tf-trainer-0 tf-trainer-1)
-PS_SERVICE=tf-trainer-ps-0
+WORKER_SERVICES="tf-trainer-0 tf-trainer-1"
+PS_SERVICE="tf-trainer-ps-0"
 PORT=2222
 DATA_PATH=${DATA_PATH:-/data/health.csv}
 OUTPUT_DIR=${OUTPUT_DIR:-/workloads/output/$(date +%Y%m%d_%H%M%S)}
@@ -22,55 +22,65 @@ fi
 
 # Resolve IPs from LoadBalancer services
 get_lb_ip() {
-  local svc="$1"
-  local ip
+  svc="$1"
   ip=$(kubectl get svc "$svc" -o jsonpath='{.status.loadBalancer.ingress[0].ip}' 2>/dev/null || true)
-  if [[ -z "$ip" ]]; then
+  if [ -z "$ip" ]; then
     # Some environments set hostname instead of ip
     ip=$(kubectl get svc "$svc" -o jsonpath='{.status.loadBalancer.ingress[0].hostname}' 2>/dev/null || true)
   fi
-  if [[ -z "$ip" ]]; then
+  if [ -z "$ip" ]; then
     echo "Service $svc has no LoadBalancer IP/hostname assigned yet. Ensure MetalLB or an LB is configured." >&2
     exit 2
   fi
   echo "$ip"
 }
 
-WORKER_ADDRS=()
-for svc in "${WORKER_SERVICES[@]}"; do
+WORKER_ADDRS_CSV=""
+WORKER_COUNT=0
+for svc in $WORKER_SERVICES; do
   ip=$(get_lb_ip "$svc")
-  WORKER_ADDRS+=("${ip}:${PORT}")
-  echo "Worker $svc -> ${ip}:${PORT}"
+  addr="${ip}:${PORT}"
+  if [ -z "$WORKER_ADDRS_CSV" ]; then
+    WORKER_ADDRS_CSV="$addr"
+  else
+    WORKER_ADDRS_CSV="$WORKER_ADDRS_CSV,$addr"
+  fi
+  WORKER_COUNT=$((WORKER_COUNT + 1))
+  echo "Worker $svc -> $addr"
 done
+
 PS_IP=$(get_lb_ip "$PS_SERVICE")
-PS_ADDRS=("${PS_IP}:${PORT}")
+PS_ADDRS_CSV="${PS_IP}:${PORT}"
+PS_COUNT=1
 echo "PS $PS_SERVICE -> ${PS_IP}:${PORT}"
 
-# Ensure python and pip exist
-if ! command -v python >/dev/null 2>&1 && command -v python3 >/dev/null 2>&1; then
-  alias python=python3
-fi
-if ! command -v python >/dev/null 2>&1; then
+# Ensure python exists; prefer python then python3
+PYTHON=""
+if command -v python >/dev/null 2>&1; then
+  PYTHON=python
+elif command -v python3 >/dev/null 2>&1; then
+  PYTHON=python3
+else
   echo "Python is required inside bastion. Aborting." >&2
   exit 3
 fi
 
 # Install TensorFlow if missing
-if ! python -c "import tensorflow as tf; print(tf.__version__)" >/dev/null 2>&1; then
+if ! "$PYTHON" -c "import tensorflow as tf; print(tf.__version__)" >/dev/null 2>&1; then
   echo "Installing TensorFlow (CPU) inside bastion container..."
-  pip install --no-cache-dir --upgrade pip >/dev/null
-  pip install --no-cache-dir tensorflow >/dev/null
+  "$PYTHON" -m pip install --no-cache-dir --upgrade pip >/dev/null 2>&1 || true
+  "$PYTHON" -m pip install --no-cache-dir tensorflow >/dev/null
 fi
 
 mkdir -p "$OUTPUT_DIR"
 
-python /workloads/local_cluster_workloads/train_tf_ps.py \
+"$PYTHON" /workloads/local_cluster_workloads/train_tf_ps.py \
   --data-path "$DATA_PATH" \
   --output-dir "$OUTPUT_DIR" \
   --epochs "$EPOCHS" \
   --batch-size "$BATCH_SIZE" \
   --use-ps \
-  --worker-replicas ${#WORKER_ADDRS[@]} \
-  --ps-replicas ${#PS_ADDRS[@]} \
-  --worker-addrs "$(IFS=,; echo "${WORKER_ADDRS[*]}")" \
-  --ps-addrs "$(IFS=,; echo "${PS_ADDRS[*]}")"
+  --worker-replicas "$WORKER_COUNT" \
+  --ps-replicas "$PS_COUNT" \
+  --worker-addrs "$WORKER_ADDRS_CSV" \
+  --ps-addrs "$PS_ADDRS_CSV"

workloads/local_cluster_workloads/train_tf_ps.py

@@ -122,8 +122,6 @@ def make_parameter_server_strategy(worker_replicas: int, ps_replicas: int, port:
 
     resolver = tf.distribute.cluster_resolver.SimpleClusterResolver(
         cluster_spec=tf.train.ClusterSpec(cluster_def),
-        task_type="chief",
-        task_id=0,
         rpc_layer="grpc",
     )
     variable_partitioner = tf.distribute.experimental.partitioners.MinSizePartitioner(
@@ -177,17 +175,62 @@ def run_training(
             print("PS addrs:", ps_addrs)
         strategy = make_parameter_server_strategy(worker_replicas, ps_replicas, port, worker_addrs, ps_addrs)
 
-        # DatasetCreator lets Keras coordinate dataset distribution for PS strategy
-        def dataset_fn(input_context: tf.distribute.InputContext):
+        # Switch to ClusterCoordinator-based custom training loop (DatasetCreator removed)
+        def per_worker_dataset_fn(input_context: Optional[tf.distribute.InputContext] = None):
             local_ds = tf.data.Dataset.from_tensor_slices((X, y))
-            local_ds = local_ds.shard(input_context.num_input_pipelines, input_context.input_pipeline_id)
+            if input_context is not None:
+                local_ds = local_ds.shard(input_context.num_input_pipelines, input_context.input_pipeline_id)
             local_ds = local_ds.shuffle(buffer_size=min(10000, len(X))).batch(batch_size).repeat()
             return local_ds
 
         with strategy.scope():
             model = build_sequential_model(input_dim, num_classes)
-        creator = tf.keras.utils.experimental.DatasetCreator(dataset_fn)
-        history = model.fit(creator, epochs=epochs, steps_per_epoch=steps_per_epoch)
+            # Build losses/optimizer/metrics explicitly for custom loop
+            optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
+            loss_obj = tf.keras.losses.SparseCategoricalCrossentropy()
+            train_acc = tf.keras.metrics.SparseCategoricalAccuracy()
+            train_loss = tf.keras.metrics.Mean()
+
+        # Create a ClusterCoordinator to drive training from the chief/coordinator process
+        coordinator = tf.distribute.coordinator.ClusterCoordinator(strategy)
+        per_worker_ds = coordinator.create_per_worker_dataset(per_worker_dataset_fn)
+        per_worker_iter = iter(per_worker_ds)
+
+        @tf.function
+        def per_worker_train_step(iterator):
+            def step_fn(inputs):
+                features, labels = inputs
+                with tf.GradientTape() as tape:
+                    logits = model(features, training=True)
+                    loss = loss_obj(labels, logits)
+                    # Add possible regularization losses
+                    loss += tf.add_n(model.losses) if model.losses else 0.0
+                grads = tape.gradient(loss, model.trainable_variables)
+                optimizer.apply_gradients(zip(grads, model.trainable_variables))
+                train_acc.update_state(labels, logits)
+                train_loss.update_state(loss)
+                return loss
+
+            return strategy.run(step_fn, args=(next(iterator),))
+
+        # Training loop
+        for epoch in range(epochs):
+            print(f"Starting epoch {epoch+1}/{epochs}...")
+            train_acc.reset_state()
+            train_loss.reset_state()
+
+            # Schedule one step per required step_per_epoch; wait for completion
+            futures = []
+            for _ in range(steps_per_epoch):
+                f = coordinator.schedule(per_worker_train_step, args=(per_worker_iter,))
+                futures.append(f)
+            # Block until all scheduled steps finish
+            coordinator.join()
+
+            print(f"Epoch {epoch+1} - loss: {train_loss.result().numpy():.4f} - accuracy: {train_acc.result().numpy():.4f}")
+
+        # Mimic Keras History-like output for downstream logging
+        history = type("_H", (), {"history": {"accuracy": [train_acc.result().numpy()]}})()
     else:
         print("Running single-process (no distributed strategy).")
         model = build_sequential_model(input_dim, num_classes)