Hyperpod Node Debugger

#!/usr/bin/env bash
# check-efa-sg.sh
#
# Identify and diagnose EFA security group rules for a HyperPod cluster.
# Automatically extracts the cluster's exact VPC, subnets, and security groups
# from the cluster ARN — works correctly even in accounts with 1000s of resources.
#
# Usage (preferred — cluster-centric, auto-discovers resources):
#   bash check-efa-sg.sh --cluster <cluster-name-or-arn> --region <region>
#
# Usage (direct SG mode — when SG is already known):
#   bash check-efa-sg.sh --sg-id <sg-id> --region <region>
#
# Exit codes:
#   0 — all required rules in place
#   1 — one or more required rules missing

set -euo pipefail

for cmd in aws python3; do
  command -v "$cmd" &>/dev/null || {
    echo "ERROR: '$cmd' is required but not found. Install it and retry."
    exit 1
  }
done

CLUSTER=""
SG_ID=""
REGION="${AWS_DEFAULT_REGION:-}"
USE_COLOR=true

usage() {
  cat <<EOF
Usage:
  $0 --cluster <cluster-name-or-arn> --region <region> [--no-color]
  $0 --sg-id   <sg-id>               --region <region> [--no-color]

Read-only diagnostic for EFA-related security group rules on a HyperPod
cluster. Reports inbound/outbound self-referencing rules and warns on
0.0.0.0/0 outbound (which the HyperPod docs advise against on the EFA SG).
On any [FAIL] the script ends with a pointer to
"references/node-diagnostics-detail.md § A (EFA / Security Group)".

Options:
  --cluster   Auto-discovers SGs, subnets, VPC from the cluster (preferred).
  --sg-id     Check a specific security group directly.
  --region    AWS region (required unless \$AWS_DEFAULT_REGION is set).
  --no-color  Disable ANSI colors.
  -h, --help  Show this message.

Exit codes:
  0  All required rules present.
  1  One or more required rules missing.
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --cluster)  CLUSTER="$2";    shift 2 ;;
    --sg-id)    SG_ID="$2";      shift 2 ;;
    --region)   REGION="$2";     shift 2 ;;
    --no-color) USE_COLOR=false; shift ;;
    -h|--help)  usage; exit 0 ;;
    *) echo "Unknown argument: $1" >&2; usage >&2; exit 1 ;;
  esac
done

if [[ -z "$CLUSTER" && -z "$SG_ID" ]]; then
  usage >&2
  exit 1
fi

if [[ -z "$REGION" ]]; then
  echo "ERROR: --region is required (or set AWS_DEFAULT_REGION before running)." >&2
  exit 2
fi

# Mutually exclusive: --cluster auto-discovers SGs, --sg-id targets one specific SG.
# Passing both was silently ignoring --sg-id — error instead so the caller notices.
if [[ -n "$CLUSTER" && -n "$SG_ID" ]]; then
  echo "ERROR: --cluster and --sg-id are mutually exclusive (pick one)" >&2
  exit 2
fi

if [[ -n "$SG_ID" && ! "$SG_ID" =~ ^sg-[a-fA-F0-9]{8,17}$ ]]; then
  echo "ERROR: Invalid security group ID format: '$SG_ID' (expected sg-<hex>, e.g. sg-0abc1234def56789a)"
  exit 1
fi

if ! [ -t 1 ] || [ "${TERM:-}" = "dumb" ]; then
  USE_COLOR=false
fi
if "$USE_COLOR"; then
  RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
  BOLD='\033[1m'; NC='\033[0m'
else
  RED=''; GREEN=''; YELLOW=''; BOLD=''; NC=''
fi

check_single_sg() {
  local sg_id="$1"
  local region="$2"
  local issues=0

  echo ""
  echo -e "${BOLD}=== EFA Security Group Diagnostic ===${NC}"
  echo -e "Security Group: ${BOLD}${sg_id}${NC}  Region: ${BOLD}${region}${NC}"
  echo ""

  local sg_json
  sg_json=$(aws ec2 describe-security-groups \
    --group-ids "$sg_id" \
    --region "$region" \
    --cli-read-timeout 30 \
    --output json 2>&1) || {
    echo -e "${RED}ERROR: Cannot describe security group '$sg_id' in region '$region'${NC}"
    echo "$sg_json"
    return 1
  }

  # Distinguish "API succeeded but returned empty" (auth-denied or malformed JSON
  # still yielding exit 0) from "SG genuinely has no rules". Without this, the
  # three rule checks below would each emit [FAIL], misleading the customer
  # into thinking rules are missing when the check itself could not run.
  local sg_count
  sg_count=$(echo "$sg_json" | python3 -c "import sys,json; d=json.load(sys.stdin); print(len(d.get('SecurityGroups',[])))" 2>/dev/null || echo 0)
  if [[ "$sg_count" -eq 0 ]]; then
    echo -e "  ${YELLOW}[WARN]${NC} Unable to check SG rules — describe-security-groups returned no data for '$sg_id' (possible IAM denial or stale ID)"
    echo -e "         → references/node-diagnostics-detail.md § A (EFA / Security Group)"
    return 0
  fi

  local sg_name vpc_id
  sg_name=$(echo "$sg_json" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['SecurityGroups'][0].get('GroupName','unknown'))" 2>/dev/null || echo "unknown")
  vpc_id=$(echo "$sg_json"  | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['SecurityGroups'][0].get('VpcId','unknown'))"   2>/dev/null || echo "unknown")
  echo -e "Name: ${sg_name}  |  VPC: ${vpc_id}"
  echo ""

  echo -e "${BOLD}--- Inbound Rules ---${NC}"
  echo "$sg_json" | python3 -c "
import sys, json
d = json.load(sys.stdin)['SecurityGroups'][0]
rules = d.get('IpPermissions', [])
if not rules:
    print('  (none)')
for r in rules:
    proto = r.get('IpProtocol', '?')
    srcs  = [g.get('GroupId','') for g in r.get('UserIdGroupPairs', [])]
    cidrs = [c.get('CidrIp','') for c in r.get('IpRanges', [])]
    for s in srcs:  print(f'  proto={proto} source=sg:{s}')
    for c in cidrs: print(f'  proto={proto} source={c}')
" 2>/dev/null

  echo ""
  echo -e "${BOLD}--- Outbound Rules ---${NC}"
  echo "$sg_json" | python3 -c "
import sys, json
d = json.load(sys.stdin)['SecurityGroups'][0]
rules = d.get('IpPermissionsEgress', [])
if not rules:
    print('  (none)')
for r in rules:
    proto = r.get('IpProtocol', '?')
    dests = [g.get('GroupId','') for g in r.get('UserIdGroupPairs', [])]
    cidrs = [c.get('CidrIp','') for c in r.get('IpRanges', [])]
    for s in dests: print(f'  proto={proto} dest=sg:{s}')
    for c in cidrs: print(f'  proto={proto} dest={c}')
" 2>/dev/null

  echo ""
  echo -e "${BOLD}--- Rule Check Results ---${NC}"

  local inbound_self outbound_self outbound_inet
  inbound_self=$(echo "$sg_json" | SG_CHECK_ID="$sg_id" python3 -c "
import sys, json, os
sg=os.environ['SG_CHECK_ID']
d = json.load(sys.stdin)['SecurityGroups'][0]
for r in d.get('IpPermissions', []):
    if r.get('IpProtocol') == '-1':
        if any(g.get('GroupId') == sg for g in r.get('UserIdGroupPairs', [])):
            print('found'); exit(0)
" 2>/dev/null || echo "")

  outbound_self=$(echo "$sg_json" | SG_CHECK_ID="$sg_id" python3 -c "
import sys, json, os
sg=os.environ['SG_CHECK_ID']
d = json.load(sys.stdin)['SecurityGroups'][0]
for r in d.get('IpPermissionsEgress', []):
    if r.get('IpProtocol') == '-1':
        if any(g.get('GroupId') == sg for g in r.get('UserIdGroupPairs', [])):
            print('found'); exit(0)
" 2>/dev/null || echo "")

  outbound_inet=$(echo "$sg_json" | python3 -c "
import sys, json
d = json.load(sys.stdin)['SecurityGroups'][0]
for r in d.get('IpPermissionsEgress', []):
    if r.get('IpProtocol') == '-1':
        if any(c.get('CidrIp') == '0.0.0.0/0' for c in r.get('IpRanges', [])):
            print('found'); exit(0)
" 2>/dev/null || echo "")

  if [[ "$inbound_self" == "found" ]]; then
    echo -e "  ${GREEN}[PASS]${NC} Inbound self-referencing rule (all traffic from ${sg_id})"
  else
    echo -e "  ${RED}[FAIL]${NC} Missing inbound self-referencing rule (all traffic from ${sg_id})"
    issues=$((issues+1))
  fi

  if [[ "$outbound_self" == "found" ]]; then
    echo -e "  ${GREEN}[PASS]${NC} Outbound self-referencing rule (all traffic to ${sg_id}) ← required for EFA"
  else
    echo -e "  ${RED}[FAIL]${NC} Missing outbound self-referencing rule ← ${BOLD}PRIMARY cause of EFA health check failure${NC}"
    issues=$((issues+1))
  fi

  if [[ "$outbound_inet" == "found" ]]; then
    echo -e "  ${YELLOW}[WARN]${NC} Outbound 0.0.0.0/0 rule present — HyperPod docs advise against this on the EFA SG (can cause EFA health check failures). Move internet egress to the subnet (NAT or VPC endpoints)."
  else
    echo -e "  ${GREEN}[PASS]${NC} No outbound 0.0.0.0/0 on EFA SG (correct per HyperPod prerequisites)"
  fi

  if [[ $issues -gt 0 ]]; then
    echo ""
    echo -e "  ${YELLOW}→ See references/node-diagnostics-detail.md § A (EFA / Security Group) for remediation.${NC}"
  fi

  return "$issues"
}

if [[ -n "$CLUSTER" ]]; then
  echo ""
  echo -e "${BOLD}=== HyperPod Cluster Resource Discovery ===${NC}"
  echo -e "Cluster: ${BOLD}${CLUSTER}${NC}"
  echo -e "Region:  ${BOLD}${REGION}${NC}"
  echo ""

  CLUSTER_JSON=$(aws sagemaker describe-cluster \
    --cluster-name "$CLUSTER" \
    --region "$REGION" \
    --cli-read-timeout 30 \
    --output json 2>&1) || {
    echo -e "${RED}ERROR: Cannot find cluster '$CLUSTER' in region '$REGION'${NC}"
    echo ""
    echo "Available clusters in this region:"
    aws sagemaker list-clusters --region "$REGION" \
      --query 'ClusterSummaries[*].{Name:ClusterName,Status:ClusterStatus,ARN:ClusterArn}' \
      --output table 2>/dev/null || echo "  (unable to list clusters)"
    echo "$CLUSTER_JSON"
    exit 1
  }

  CLUSTER_ARN=$(echo "$CLUSTER_JSON"    | python3 -c "import sys,json; print(json.load(sys.stdin).get('ClusterArn',''))"    2>/dev/null || echo "")
  CLUSTER_STATUS=$(echo "$CLUSTER_JSON" | python3 -c "import sys,json; print(json.load(sys.stdin).get('ClusterStatus',''))" 2>/dev/null || echo "")
  ORCHESTRATOR=$(echo "$CLUSTER_JSON"   | python3 -c "
import sys,json
d=json.load(sys.stdin)
print('EKS' if 'Eks' in d.get('Orchestrator',{}) else 'Slurm')
" 2>/dev/null || echo "Unknown")

  echo -e "  ARN:          ${CLUSTER_ARN}"
  echo -e "  Status:       ${CLUSTER_STATUS}"
  echo -e "  Orchestrator: ${ORCHESTRATOR}"

  RESOURCES=$(echo "$CLUSTER_JSON" | python3 -c "
import sys,json
d=json.load(sys.stdin)
vpc=d.get('VpcConfig',{})
sgs=vpc.get('SecurityGroupIds',[])
subnets=vpc.get('Subnets',[])
print('SGs='     + ','.join(sgs))
print('Subnets=' + ','.join(subnets))
" 2>/dev/null || echo "")

  CLUSTER_SGS=$(echo "$RESOURCES"     | grep "^SGs="     | cut -d= -f2)
  CLUSTER_SUBNETS=$(echo "$RESOURCES" | grep "^Subnets=" | cut -d= -f2)

  if [[ -z "$CLUSTER_SGS" ]]; then
    echo -e "${YELLOW}[WARN]${NC} No SecurityGroupIds in cluster VpcConfig — cluster may not have customer VPC"
    exit 0
  fi

  VPC_ID="unknown"
  if [[ -n "$CLUSTER_SUBNETS" ]]; then
    FIRST_SUBNET=$(echo "$CLUSTER_SUBNETS" | tr ',' '\n' | head -1)
    VPC_ID=$(aws ec2 describe-subnets \
      --subnet-ids "$FIRST_SUBNET" \
      --region "$REGION" \
      --query 'Subnets[0].VpcId' \
      --output text 2>/dev/null || echo "unknown")
  fi

  echo ""
  echo -e "${BOLD}  Resources owned by cluster '${CLUSTER}':${NC}"
  echo -e "  VPC:              ${VPC_ID}"
  echo -e "  Security Groups:  ${CLUSTER_SGS}"
  echo -e "  Subnets:          ${CLUSTER_SUBNETS}"

  if [[ -n "$CLUSTER_SUBNETS" ]]; then
    echo ""
    echo -e "${BOLD}  Subnet details:${NC}"
    IFS=',' read -ra _subnet_arr <<< "$CLUSTER_SUBNETS"
    aws ec2 describe-subnets \
      --subnet-ids "${_subnet_arr[@]}" \
      --region "$REGION" \
      --query 'Subnets[*].{SubnetId:SubnetId,AZ:AvailabilityZone,FreeIPs:AvailableIpAddressCount,VpcId:VpcId}' \
      --output table 2>/dev/null || echo "  (unable to describe subnets)"
  fi

  echo ""
  TOTAL_ISSUES=0
  # CLUSTER_SGS is guaranteed non-empty at the -z guard above, but defend anyway.
  # grep -c returns exit 1 on zero matches under pipefail, so suppress and then
  # explicitly branch on the count rather than letting 0 silently fall through.
  SG_COUNT=$(echo "$CLUSTER_SGS" | tr ',' '\n' | grep -c . || true)
  if [[ "${SG_COUNT:-0}" -eq 0 ]]; then
    echo -e "  ${YELLOW}[WARN]${NC} No security groups resolved from CLUSTER_SGS — cannot run EFA rule check"
    echo -e "         → references/node-diagnostics-detail.md § A (EFA / Security Group)"
    exit 0
  fi
  echo -e "${BOLD}Checking ${SG_COUNT} security group(s) for cluster '${CLUSTER}'...${NC}"

  for SG in $(echo "$CLUSTER_SGS" | tr ',' ' '); do
    echo ""
    echo -e "${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    # Capture rc in a subshell pattern that survives `set -e` — otherwise
    # the first SG with issues aborts the loop and later SGs are never checked.
    sg_rc=0
    check_single_sg "$SG" "$REGION" || sg_rc=$?
    TOTAL_ISSUES=$((TOTAL_ISSUES + sg_rc))
  done

  echo ""
  echo -e "${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
  if [[ $TOTAL_ISSUES -gt 0 ]]; then
    echo -e "${RED}${BOLD}RESULT: ${TOTAL_ISSUES} security group rule issue(s) found for cluster '${CLUSTER}'${NC}"
    echo "Fix the [FAIL] rules above (see references/node-diagnostics-detail.md § A for the Suggested-command block); if cluster creation was failing on EFA health checks, retry creation after fixing."
    echo ""
    echo "Verify after fixing:"
    echo "  bash check-efa-sg.sh --cluster ${CLUSTER} --region ${REGION}"
    exit 1
  else
    echo -e "${GREEN}${BOLD}RESULT: All EFA security group rules correctly configured for cluster '${CLUSTER}'${NC}"
    echo ""
    echo "If EFA health checks still fail:"
    echo "  1. Verify all instance groups use one of these SGs: ${CLUSTER_SGS}"
    echo "  2. Run check-node-reachability.sh on affected nodes via hyperpod-ssm skill"
    exit 0
  fi
fi

if [[ -n "$SG_ID" ]]; then
  check_single_sg "$SG_ID" "$REGION"
  exit $?
fi

#!/usr/bin/env bash
# check-node-reachability.sh
#
# Diagnose EFA reachability and inter-node communication health on a single
# HyperPod node. Run this on each node via the hyperpod-ssm skill.
#
# Usage (via ssm-exec.sh):
#   ssm-exec.sh --target <TARGET> --upload scripts/check-node-reachability.sh /tmp/check-node-reachability.sh
#   ssm-exec.sh --target <TARGET> 'bash /tmp/check-node-reachability.sh'
#
# Usage (direct on node):
#   bash check-node-reachability.sh [--json] [--no-color]
#
# Exit codes:
#   0 — all critical checks passed
#   1 — one or more critical checks failed

set -euo pipefail

# Note: this script runs ON the node (via SSM), so aws CLI may not be present.
# Only python3 is checked here; other tools are checked individually per section.

JSON_MODE=false
USE_COLOR=true

usage() {
  cat <<EOF
Usage: bash check-node-reachability.sh [--json] [--no-color]

Read-only on-node diagnostic for EFA reachability and inter-node communication
health. Must be executed on a HyperPod compute node (typically via the
hyperpod-ssm skill). Checks EFA interfaces, /dev/infiniband devices, GPU
count and Neuron device count against the expected counts for the node's
instance type.

Options:
  --json       Emit findings as JSON instead of human-readable output.
  --no-color   Disable ANSI colors.
  -h, --help   Show this message.

Exit codes:
  0  All critical checks passed.
  1  One or more critical checks failed.
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --json)     JSON_MODE=true;  shift ;;
    --no-color) USE_COLOR=false; shift ;;
    -h|--help)  usage; exit 0 ;;
    *) echo "Unknown argument: $1" >&2; usage >&2; exit 1 ;;
  esac
done

# Colors — auto-disable when stdout isn't a TTY.
if ! [ -t 1 ] || [ "${TERM:-}" = "dumb" ]; then
  USE_COLOR=false
fi
if "$USE_COLOR"; then
  RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
  BOLD='\033[1m';   NC='\033[0m'
else
  RED=''; GREEN=''; YELLOW=''; BOLD=''; NC=''
fi

HOSTNAME=$(hostname 2>/dev/null || echo "unknown")
TIMESTAMP=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
CRITICAL_FAILURES=0
declare -A RESULTS   # associative array: check_name → pass|fail|warn|skip

pass()  { RESULTS["$1"]="pass";  [[ "$JSON_MODE" == false ]] && echo -e "  ${GREEN}[PASS]${NC}  $1${2:+ — $2}"; }
fail()  { RESULTS["$1"]="fail";  CRITICAL_FAILURES=$((CRITICAL_FAILURES+1)); \
           [[ "$JSON_MODE" == false ]] && echo -e "  ${RED}[FAIL]${NC}  $1${2:+ — $2}"; }
warn()  { RESULTS["$1"]="warn";  [[ "$JSON_MODE" == false ]] && echo -e "  ${YELLOW}[WARN]${NC}  $1${2:+ — $2}"; }
skip()  { RESULTS["$1"]="skip";  [[ "$JSON_MODE" == false ]] && echo -e "         [SKIP]  $1${2:+ — $2}"; }
info()  { [[ "$JSON_MODE" == false ]] && echo -e "         $1"; }

if [[ "$JSON_MODE" == false ]]; then
  echo ""
  echo -e "${BOLD}=== HyperPod Node EFA Reachability Check ===${NC}"
  echo -e "Host:      ${BOLD}${HOSTNAME}${NC}"
  echo -e "Timestamp: ${TIMESTAMP}"
  echo ""
fi

if [[ "$JSON_MODE" == false ]]; then echo -e "${BOLD}--- EFA Kernel Module ---${NC}"; fi

EFA_MODULE=$(lsmod 2>/dev/null | grep -E '^efa\b' | awk '{print $1}' || true)
if [[ -n "$EFA_MODULE" ]]; then
  EFA_MODULE_VER=$(modinfo efa 2>/dev/null | grep -E '^version:' | awk '{print $2}' || echo "unknown")
  pass "efa_kernel_module" "loaded (version: ${EFA_MODULE_VER})"
else
  # Read-only invariant: detect only, never `sudo modprobe efa` — loading kernel
  # modules mutates node state, which the hyperpod-ssm skill's approval flow owns.
  fail "efa_kernel_module" "not loaded — see references/node-diagnostics-detail.md § A (EFA / Security Group)"
fi

if [[ "$JSON_MODE" == false ]]; then echo ""; echo -e "${BOLD}--- EFA Devices ---${NC}"; fi

# shellcheck disable=SC2010  # /dev/ entries are kernel-named, safe to ls|grep
EFA_DEVICES=$(ls /dev/infiniband/ 2>/dev/null | grep -E 'rdma_cm|uverbs|efa' || true)

if [[ -n "$EFA_DEVICES" ]]; then
  pass "efa_devices_present" "found in /dev/infiniband/: $(echo "$EFA_DEVICES" | tr '\n' ' ')"
else
  fail "efa_devices_present" "/dev/infiniband/ is empty or missing — EFA hardware not detected"
fi

if [[ "$JSON_MODE" == false ]]; then echo ""; echo -e "${BOLD}--- libfabric EFA Provider ---${NC}"; fi

if command -v fi_info &>/dev/null; then
  # If the previous section found no EFA hardware, fi_info failing is expected —
  # don't emit [FAIL] on top of the hardware [FAIL], which would double-count and
  # conflate "libfabric can't see EFA" with "node has no EFA at all".
  if [[ -z "$EFA_DEVICES" ]]; then
    skip "fi_info_efa_provider" "no EFA devices detected upstream — see efa_devices_present"
  else
    FI_EXIT=0
    FI_OUTPUT=$(fi_info -p efa 2>&1) || FI_EXIT=$?
    if echo "$FI_OUTPUT" | grep -q "provider: efa"; then
      EFA_PROVIDER_COUNT=$(echo "$FI_OUTPUT" | { grep -c "provider: efa" 2>/dev/null; true; })
      pass "fi_info_efa_provider" "EFA provider found (${EFA_PROVIDER_COUNT} endpoint(s))"
      info "$(echo "$FI_OUTPUT" | grep -E 'provider:|fabric:|domain:|version:' | head -8 | sed 's/^/    /')"
    else
      fail "fi_info_efa_provider" "fi_info -p efa returned no EFA provider (exit code ${FI_EXIT}) — libfabric cannot enumerate EFA devices. See references/node-diagnostics-detail.md § A (EFA / Security Group)"
      info "fi_info output: ${FI_OUTPUT:0:200}"
    fi
  fi
else
  warn "fi_info_efa_provider" "fi_info not found — install libfabric to run this check (fi_info comes with EFA installer)"
fi

if [[ "$JSON_MODE" == false ]]; then echo ""; echo -e "${BOLD}--- EFA Network Interfaces ---${NC}"; fi

# EFA interfaces typically appear as eth0/ens* for primary + rdmaX or efa* for EFA devices
# EFA ifaces on p5/p5en use regular kernel names (ens*) — filter by driver via ethtool
# rather than by name pattern (the old 'rdma|efa' name grep misses ens* on p5).
EFA_IFACES=""
if command -v ethtool &>/dev/null; then
  while IFS= read -r iface; do
    [[ -z "$iface" ]] && continue
    DRIVER=$(ethtool -i "$iface" 2>/dev/null | awk -F': ' '/^driver:/{print $2}')
    if [[ "$DRIVER" == "efa" ]]; then
      EFA_IFACES+="${iface}"$'\n'
    fi
  done < <(ip -o link show 2>/dev/null | awk -F': ' '{print $2}' | awk -F'@' '{print $1}' | grep -v '^lo$')
fi
# Fallback to name-based detection for older kernels / containers without ethtool
if [[ -z "$EFA_IFACES" ]]; then
  EFA_IFACES=$(ip link show 2>/dev/null | grep -E 'rdma|efa' | awk -F': ' '{print $2}' | tr -d '@' || true)
fi
REGULAR_IFACES=$(ip link show 2>/dev/null | grep -E 'state UP' | awk -F': ' '{print $2}' | tr -d '@' || true)

if [[ -n "$EFA_IFACES" ]]; then
  pass "efa_interfaces_up" "EFA interfaces found: $(echo "$EFA_IFACES" | tr '\n' ' ')"
  while IFS= read -r iface; do
    [[ -z "$iface" ]] && continue
    IP=$(ip addr show "$iface" 2>/dev/null | grep 'inet ' | awk '{print $2}' || true)
    if [[ -n "$IP" ]]; then
      info "  $iface → $IP"
    else
      warn "efa_interface_ip_${iface}" "interface $iface has no IP address — check DHCP/subnet config"
    fi
  done <<< "$EFA_IFACES"
else
  info "No EFA interfaces detected (by driver or name)"
  if [[ -n "$REGULAR_IFACES" ]]; then
    skip "efa_interfaces_up" "no separate EFA interface — primary interfaces: $(echo "$REGULAR_IFACES" | tr '\n' ' ' | head -c 80)"
  else
    warn "efa_interfaces_up" "no UP network interfaces found"
  fi
fi

if [[ "$JSON_MODE" == false ]]; then echo ""; echo -e "${BOLD}--- EFA Installation ---${NC}"; fi

EFA_VER_FILE="/opt/amazon/efa_installed_packages"
if [[ -f "$EFA_VER_FILE" ]]; then
  # Format is "EFA installer version: 1.30.0" — grab only the version token.
  EFA_VER=$(grep -iE '^EFA installer version' "$EFA_VER_FILE" 2>/dev/null \
              | head -1 \
              | grep -oE '[0-9]+\.[0-9]+(\.[0-9]+)?' \
              | head -1 || echo "")
  if [[ -z "$EFA_VER" ]]; then
    warn "efa_installer_present" "EFA installer file present but version line not parsed"
  else
    pass "efa_installer_present" "EFA installer version: ${EFA_VER}"
  fi
else
  warn "efa_installer_present" "EFA installer marker not found at ${EFA_VER_FILE} — EFA may not be installed via standard method"
fi

if [[ "$JSON_MODE" == false ]]; then echo ""; echo -e "${BOLD}--- NCCL / OFI Configuration ---${NC}"; fi

NCCL_VARS=("FI_PROVIDER" "FI_EFA_USE_DEVICE_RDMA" "NCCL_SOCKET_IFNAME" "NCCL_ALGO" "LD_LIBRARY_PATH")
ANY_NCCL_SET=false
for var in "${NCCL_VARS[@]}"; do
  val="${!var:-}"
  if [[ -n "$val" ]]; then
    info "  ${var}=${val}"
    ANY_NCCL_SET=true
  fi
done

if "$ANY_NCCL_SET"; then
  FI_PROVIDER_VAL="${FI_PROVIDER:-}"
  if [[ -n "$FI_PROVIDER_VAL" && "$FI_PROVIDER_VAL" != "efa" ]]; then
    warn "nccl_fi_provider" "FI_PROVIDER=${FI_PROVIDER_VAL} — for EFA workloads this should be 'efa'"
  elif [[ "$FI_PROVIDER_VAL" == "efa" ]]; then
    pass "nccl_fi_provider" "FI_PROVIDER=efa"
  fi
else
  skip "nccl_env_vars" "no NCCL/OFI env vars set in current shell — may be set in job launcher environment"
fi

if [[ "$JSON_MODE" == false ]]; then echo ""; echo -e "${BOLD}--- AWS OFI NCCL Plugin ---${NC}"; fi

OFI_LIB=$(find /opt/amazon/efa /opt/aws-ofi-nccl /usr/local/lib /usr/lib \
  -name "libnccl-net.so*" -o -name "aws-ofi-nccl.so*" 2>/dev/null | head -1 || true)

if [[ -n "$OFI_LIB" ]]; then
  pass "aws_ofi_nccl_plugin" "found: ${OFI_LIB}"
else
  if [[ -f "$EFA_VER_FILE" ]] && grep -q "ofi\|OFI" "$EFA_VER_FILE" 2>/dev/null; then
    pass "aws_ofi_nccl_plugin" "referenced in ${EFA_VER_FILE}"
  else
    warn "aws_ofi_nccl_plugin" "libnccl-net.so not found — required for EFA-accelerated NCCL (distributed training)"
  fi
fi

if [[ "$JSON_MODE" == false ]]; then echo ""; echo -e "${BOLD}--- Instance Metadata Reachability ---${NC}"; fi

IMDS_TOKEN=$(curl -s -X PUT "http://169.254.169.254/latest/api/token" \
  -H "X-aws-ec2-metadata-token-ttl-seconds: 60" --connect-timeout 3 -m 5 2>/dev/null || true)

if [[ -n "$IMDS_TOKEN" ]]; then
  INSTANCE_TYPE=$(curl -s -H "X-aws-ec2-metadata-token: $IMDS_TOKEN" \
    http://169.254.169.254/latest/meta-data/instance-type --connect-timeout 3 -m 5 2>/dev/null || echo "unknown")
  LOCAL_IP=$(curl -s -H "X-aws-ec2-metadata-token: $IMDS_TOKEN" \
    http://169.254.169.254/latest/meta-data/local-ipv4 --connect-timeout 3 -m 5 2>/dev/null || echo "unknown")
  pass "imds_reachable" "instance-type=${INSTANCE_TYPE}, local-ipv4=${LOCAL_IP}"

  # Static list of EFA-capable families; unknown types fall through to the
  # EC2 API check. aws CLI may not be present on-node, so the static path
  # covers the common case.
  case "$INSTANCE_TYPE" in
    p4de*|p4d*|p5en*|p5e*|p5*|p6*|trn1*|trn2*|inf2*|g5.48xlarge|g6e.48xlarge|g6.48xlarge|hpc6a*|hpc6id*|hpc7a*|hpc7g*|dl1*|dl2q*)
      pass "efa_capable_instance" "${INSTANCE_TYPE} supports EFA" ;;
    *)
      if command -v aws &>/dev/null; then
        EFA_CHECK=$(aws ec2 describe-instance-types \
          --instance-types "${INSTANCE_TYPE}" \
          --query 'InstanceTypes[0].NetworkInfo.EfaSupported' \
          --output text 2>/dev/null || echo "unknown")
        if [[ "$EFA_CHECK" == "True" ]]; then
          pass "efa_capable_instance" "${INSTANCE_TYPE} supports EFA (verified via API)"
        elif [[ "$EFA_CHECK" == "False" ]]; then
          warn "efa_capable_instance" "${INSTANCE_TYPE} does NOT support EFA"
        else
          warn "efa_capable_instance" "${INSTANCE_TYPE} — could not verify EFA support"
        fi
      else
        warn "efa_capable_instance" "${INSTANCE_TYPE} — not in known EFA list; verify with: aws ec2 describe-instance-types --instance-types ${INSTANCE_TYPE} --query 'InstanceTypes[0].NetworkInfo.EfaSupported'"
      fi
      ;;
  esac
  # Multi-EFA validation — counts per EC2 instance-type documentation.
  # NOTE: EFA counts vary between instance families (p5en has fewer than p5/p5e).
  EXPECTED_EFA=0
  case "$INSTANCE_TYPE" in
    p5.48xlarge|p5e.48xlarge)   EXPECTED_EFA=32 ;;
    p5en.48xlarge)              EXPECTED_EFA=16 ;;
    p4d.24xlarge|p4de.24xlarge) EXPECTED_EFA=4 ;;
    trn1.32xlarge)              EXPECTED_EFA=8 ;;
    trn2.48xlarge)              EXPECTED_EFA=16 ;;
    # p6 family and newer: don't hardcode counts; discover via ethtool to avoid false FAILs.
  esac

  if [[ "$EXPECTED_EFA" -gt 0 ]]; then
    # Count actual EFA devices — avoid grep -c pattern that returns "0\n0" fallthrough.
    ACTUAL_EFA=$(find /dev/infiniband -maxdepth 1 -name 'uverbs*' 2>/dev/null | wc -l)
    [[ -z "$ACTUAL_EFA" ]] && ACTUAL_EFA=0
    if [[ "$ACTUAL_EFA" -ge "$EXPECTED_EFA" ]]; then
      pass "multi_efa_interfaces" "${ACTUAL_EFA}/${EXPECTED_EFA} EFA interfaces present for ${INSTANCE_TYPE}"
    elif [[ "$ACTUAL_EFA" -gt 0 ]]; then
      warn "multi_efa_interfaces" "only ${ACTUAL_EFA}/${EXPECTED_EFA} EFA interfaces — some may not be attached or driver issue"
    else
      fail "multi_efa_interfaces" "0/${EXPECTED_EFA} EFA interfaces on ${INSTANCE_TYPE} — EFA driver or attachment issue"
    fi
  fi
else
  warn "imds_reachable" "IMDS not reachable. If running inside a container: check IMDSv2 HttpPutResponseHopLimit on the instance (default 1 is often too low for container networking — set to 2 or higher). Otherwise: verify the instance metadata service is enabled (HttpEndpoint != disabled) and that no local iptables / nftables rules block 169.254.169.254. Note: SGs do not filter link-local addresses."
fi

if [[ "$JSON_MODE" == false ]]; then echo ""; echo -e "${BOLD}--- Network Interface Statistics ---${NC}"; fi

if command -v ip &>/dev/null; then
  IFACE_ERRORS=$(ip -s link show 2>/dev/null | awk '
    BEGIN { rx_err=0; tx_err=0; iface="" }
    /^[0-9]+:/ {
      if (iface != "" && (rx_err > 0 || tx_err > 0))
        print "  " iface ": RX errors=" rx_err " TX errors=" tx_err
      iface=$2; gsub(/:$/, "", iface)
      rx_err=0; tx_err=0
    }
    /RX:/ { getline; rx_err=$3+0 }
    /TX:/ { getline; tx_err=$3+0 }
    END {
      if (iface != "" && (rx_err > 0 || tx_err > 0))
        print "  " iface ": RX errors=" rx_err " TX errors=" tx_err
    }
  ' || true)

  if [[ -n "$IFACE_ERRORS" ]]; then
    warn "network_interface_errors" "interfaces with errors detected:"
    info "$IFACE_ERRORS"
  else
    pass "network_interface_errors" "no RX/TX errors on active interfaces"
  fi
else
  skip "network_interface_errors" "ip command not available"
fi

if [[ "$JSON_MODE" == false ]]; then echo ""; echo -e "${BOLD}--- Neuron Devices (Trainium/Inferentia) ---${NC}"; fi

if command -v neuron-ls &>/dev/null; then
  NEURON_OUTPUT=$(neuron-ls 2>&1 || true)
  NEURON_DEVICE_COUNT=$(echo "$NEURON_OUTPUT" | { grep -c "neuron_device" 2>/dev/null; true; })
  if [[ "$NEURON_DEVICE_COUNT" -gt 0 ]]; then
    pass "neuron_devices" "${NEURON_DEVICE_COUNT} Neuron device(s) detected"
    info "$(echo "$NEURON_OUTPUT" | head -10 | sed 's/^/    /')"
  else
    NEURON_MOD=$(lsmod 2>/dev/null | grep -E '^neuron' || true)
    if [[ -n "$NEURON_MOD" ]]; then
      warn "neuron_devices" "Neuron driver loaded but neuron-ls shows 0 devices → references/node-diagnostics-detail.md § G.2 (Trainium/Inferentia)"
    else
      fail "neuron_devices" "Neuron driver not loaded → references/node-diagnostics-detail.md § G.2 (Trainium/Inferentia)"
    fi
  fi
elif ls /dev/neuron* &>/dev/null 2>&1; then
  NEURON_DEV_COUNT=$(find /dev -maxdepth 1 -name 'neuron*' 2>/dev/null | wc -l)
  NEURON_DEV_COUNT=${NEURON_DEV_COUNT:-0}
  warn "neuron_devices" "${NEURON_DEV_COUNT} /dev/neuron* device(s) found but neuron-ls not installed → references/node-diagnostics-detail.md § G.2 (Trainium/Inferentia)"
else
  skip "neuron_devices" "not a Trainium/Inferentia instance (no Neuron devices)"
fi

if [[ "$JSON_MODE" == false ]]; then
  echo ""
  echo -e "${BOLD}--- Summary ---${NC}"
  TOTAL=${#RESULTS[@]}
  PASSED=$(printf '%s\n' "${RESULTS[@]}" | { grep -c "^pass$" 2>/dev/null; true; })
  WARNED=$(printf '%s\n' "${RESULTS[@]}" | { grep -c "^warn$" 2>/dev/null; true; })
  FAILED=$(printf '%s\n' "${RESULTS[@]}" | { grep -c "^fail$" 2>/dev/null; true; })
  SKIPPED=$(printf '%s\n' "${RESULTS[@]}" | { grep -c "^skip$" 2>/dev/null; true; })
  echo -e "  Host: ${HOSTNAME}"
  echo -e "  Checks: ${TOTAL} total | ${GREEN}${PASSED} passed${NC} | ${YELLOW}${WARNED} warnings${NC} | ${RED}${FAILED} failed${NC} | ${SKIPPED} skipped"

  if [[ $CRITICAL_FAILURES -eq 0 ]]; then
    echo -e "\n  ${GREEN}${BOLD}Node EFA reachability checks PASSED.${NC}"
    echo "  If inter-node communication still fails, verify security group rules with check-efa-sg.sh"
    echo "  and compare EFA versions across nodes with the hyperpod-version-checker skill."
  else
    echo -e "\n  ${RED}${BOLD}Node EFA reachability checks FAILED (${CRITICAL_FAILURES} critical issue(s)).${NC}"
    echo "  See [FAIL] items above. Each finding ends with a pointer of the form"
    echo "  '→ references/node-diagnostics-detail.md § <section>' — open that section"
    echo "  for root cause and remediation. Remediation lives in references, not in scripts."
  fi
  echo ""
else
  CHECKS_JSON=""
  for key in "${!RESULTS[@]}"; do
    val="${RESULTS[$key]}"
    CHECKS_JSON+="\"${key}\": \"${val}\","
  done
  CHECKS_JSON="${CHECKS_JSON%,}"  # remove trailing comma

  cat <<EOF
{
  "hostname": "${HOSTNAME}",
  "timestamp": "${TIMESTAMP}",
  "critical_failures": ${CRITICAL_FAILURES},
  "overall_pass": $([ $CRITICAL_FAILURES -eq 0 ] && echo true || echo false),
  "checks": { ${CHECKS_JSON} }
}
EOF
fi

exit "$([[ $CRITICAL_FAILURES -eq 0 ]] && echo 0 || echo 1)"

#!/usr/bin/env bash
# check-vpc-config.sh
#
# Diagnose VPC, subnet, and EKS configuration for a HyperPod cluster.
# Automatically extracts ALL resources (VPC, subnets, SGs) from the cluster —
# no need to know resource IDs in advance, even in accounts with 1000s of resources.
#
# Checks: VPC alignment, subnet AZ, IP availability, ENI limits,
#         EKS auth mode, HyperPod namespace, VPC endpoints.
#
# Usage (cluster-centric — preferred):
#   bash check-vpc-config.sh --cluster <name-or-arn> --region <region>
#   bash check-vpc-config.sh --cluster <name-or-arn> --region <region> --eks-name <eks-cluster>
#
# Exit codes:
#   0 — all checks passed (warnings may still be present)
#   1 — one or more critical checks failed

set -euo pipefail

for cmd in aws python3; do
  command -v "$cmd" &>/dev/null || {
    echo "ERROR: '$cmd' is required but not found. Install it and retry."
    exit 1
  }
done

CLUSTER=""
REGION="${AWS_DEFAULT_REGION:-}"
EKS_NAME=""
USE_COLOR=true

usage() {
  cat <<EOF
Usage: $0 --cluster <name-or-arn> --region <region> [options]

Read-only diagnostic for VPC / subnet / EKS configuration on a HyperPod
cluster. Reports VPC alignment, subnet AZ, IP availability, ENI limits,
EKS auth mode, HyperPod namespace presence, and VPC endpoint presence.
Each [FAIL] line includes a pointer of the form
"→ references/node-diagnostics-detail.md § B (VPC / Routing)".

Options:
  --cluster     HyperPod cluster name or ARN (required).
  --region      AWS region (required unless \$AWS_DEFAULT_REGION is set).
  --eks-name    EKS cluster name if different from the HyperPod cluster name.
  --no-color    Disable ANSI colors.
  -h, --help    Show this message.

Exit codes:
  0  All checks passed (warnings may still be present).
  1  One or more critical checks failed.
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --cluster)   CLUSTER="$2";   shift 2 ;;
    --region)    REGION="$2";    shift 2 ;;
    --eks-name)  EKS_NAME="$2";  shift 2 ;;
    --no-color)  USE_COLOR=false; shift ;;
    -h|--help)   usage; exit 0 ;;
    *) echo "Unknown argument: $1" >&2; usage >&2; exit 1 ;;
  esac
done

if [[ -z "$CLUSTER" ]]; then
  usage >&2
  exit 1
fi

if [[ -z "$REGION" ]]; then
  echo "ERROR: --region is required (or set AWS_DEFAULT_REGION before running)." >&2
  exit 2
fi

if ! [ -t 1 ] || [ "${TERM:-}" = "dumb" ]; then
  USE_COLOR=false
fi
if "$USE_COLOR"; then
  RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
  BOLD='\033[1m'; NC='\033[0m'
else
  RED=''; GREEN=''; YELLOW=''; BOLD=''; NC=''
fi

ENI_QUOTA_CODE="L-DF5E4CA3"   # AWS Service Quotas code for "Network interfaces per Region"

CRITICAL_FAILURES=0

pass()  { echo -e "  ${GREEN}[PASS]${NC}  $1${2:+ — $2}"; }
fail()  { CRITICAL_FAILURES=$((CRITICAL_FAILURES+1)); echo -e "  ${RED}[FAIL]${NC}  $1${2:+ — $2}"; }
warn()  { echo -e "  ${YELLOW}[WARN]${NC}  $1${2:+ — $2}"; }
info()  { echo -e "         $1"; }
header(){ echo ""; echo -e "${BOLD}--- $1 ---${NC}"; }

echo ""
echo -e "${BOLD}=== HyperPod VPC Configuration Check ===${NC}"
echo -e "Cluster: ${BOLD}${CLUSTER}${NC}"
echo -e "Region:  ${BOLD}${REGION}${NC}"

header "1. Cluster VPC Configuration"

CLUSTER_JSON=$(aws sagemaker describe-cluster \
  --cluster-name "$CLUSTER" \
  --region "$REGION" \
  --cli-read-timeout 30 \
  --output json 2>&1) || {
  echo -e "${RED}ERROR: Could not describe cluster '$CLUSTER' in region '$REGION'${NC}"
  echo "$CLUSTER_JSON"
  exit 1
}

CLUSTER_STATUS=$(echo "$CLUSTER_JSON" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('ClusterStatus','unknown'))" 2>/dev/null || echo "unknown")
ORCHESTRATOR=$(echo "$CLUSTER_JSON" | python3 -c "import sys,json; d=json.load(sys.stdin); o=d.get('Orchestrator',{}); print('EKS' if 'Eks' in o else 'Slurm')" 2>/dev/null || echo "unknown")
NODE_RECOVERY=$(echo "$CLUSTER_JSON" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('NodeRecovery','Unknown'))" 2>/dev/null || echo "Unknown")

info "Status: $CLUSTER_STATUS | Orchestrator: $ORCHESTRATOR | NodeRecovery: $NODE_RECOVERY"

SUBNET_IDS=$(echo "$CLUSTER_JSON" | python3 -c "
import sys,json
d=json.load(sys.stdin)
subnets=d.get('VpcConfig',{}).get('Subnets',[])
print(' '.join(subnets))
" 2>/dev/null || echo "")

SG_IDS=$(echo "$CLUSTER_JSON" | python3 -c "
import sys,json
d=json.load(sys.stdin)
sgs=d.get('VpcConfig',{}).get('SecurityGroupIds',[])
print(' '.join(sgs))
" 2>/dev/null || echo "")

if [[ -n "$SUBNET_IDS" ]]; then
  pass "VpcConfig found"
  info "Subnets: $SUBNET_IDS"
  info "SecurityGroups: $SG_IDS"
else
  warn "VpcConfig" "no VpcConfig found in cluster"
fi

if [[ "$ORCHESTRATOR" == "EKS" && -z "$EKS_NAME" ]]; then
  EKS_NAME=$(echo "$CLUSTER_JSON" | python3 -c "
import sys,json
d=json.load(sys.stdin)
arn=d.get('Orchestrator',{}).get('Eks',{}).get('ClusterArn','')
print(arn.split('/')[-1] if arn else '')
" 2>/dev/null || echo "")
fi

if [[ -n "$SUBNET_IDS" ]]; then
  header "2. Subnet VPC Alignment"

  read -ra _subnet_arr <<< "$SUBNET_IDS"
  SUBNET_JSON=$(aws ec2 describe-subnets \
    --subnet-ids "${_subnet_arr[@]}" \
    --region "$REGION" \
    --cli-read-timeout 30 \
    --output json 2>/dev/null || echo '{"Subnets":[]}')

  VPC_IDS=$(echo "$SUBNET_JSON" | python3 -c "
import sys,json
subnets=json.load(sys.stdin).get('Subnets',[])
vpc_ids=set(s.get('VpcId','?') for s in subnets)
for s in subnets:
    free=s.get('AvailableIpAddressCount',0)
    az=s.get('AvailabilityZone','?')
    sid=s.get('SubnetId','?')
    vpc=s.get('VpcId','?')
    flag='LOW IPs' if free < 10 else ''
    print(f'  {sid}: VPC={vpc} AZ={az} FreeIPs={free} {flag}')
print('VPCS=' + ','.join(vpc_ids))
" 2>/dev/null || echo "")

  echo "$VPC_IDS" | grep -v "^VPCS=" || true

  UNIQUE_VPCS=$(echo "$VPC_IDS" | grep "^VPCS=" | cut -d= -f2 | tr ',' '\n' | sort -u | tr '\n' ',' | sed 's/,$//')
  VPC_COUNT=$(echo "$UNIQUE_VPCS" | tr ',' '\n' | { grep -c . 2>/dev/null; true; })

  if [[ "$VPC_COUNT" -gt 1 ]]; then
    fail "Subnet VPC alignment" "Subnets are in DIFFERENT VPCs: $UNIQUE_VPCS — all must be in the same VPC → references/node-diagnostics-detail.md § B (VPC / Routing)"
  elif [[ "$VPC_COUNT" -eq 1 ]]; then
    pass "Subnet VPC alignment" "All subnets in VPC: $UNIQUE_VPCS"
  else
    # VPC_COUNT=0 means describe-subnets returned empty — usually an IAM denial
    # on ec2:DescribeSubnets or a stale subnet ID. Without this branch the
    # check would silently fall through and the customer sees no line at all.
    warn "Subnet VPC alignment" "Unable to determine VPC — describe-subnets returned no data (check IAM ec2:DescribeSubnets) → references/node-diagnostics-detail.md § B (VPC / Routing)"
  fi

  if [[ -n "$SG_IDS" ]]; then
    read -ra _sg_arr <<< "$SG_IDS"
    SG_JSON=$(aws ec2 describe-security-groups \
      --group-ids "${_sg_arr[@]}" \
      --region "$REGION" \
      --output json 2>/dev/null || echo '{"SecurityGroups":[]}')

    SG_VPC_CHECK=$(echo "$SG_JSON" | SUBNET_VPC="$UNIQUE_VPCS" python3 -c "
import sys, json, os
sgs=json.load(sys.stdin).get('SecurityGroups',[])
subnet_vpc=os.environ.get('SUBNET_VPC','')
subnet_vpc_set=set(subnet_vpc.split(',')) if subnet_vpc else set()
all_ok=True
for sg in sgs:
    sgid=sg.get('GroupId','?')
    vpc=sg.get('VpcId','?')
    if vpc not in subnet_vpc_set:
        print(f'MISMATCH:{sgid} is in VPC {vpc} but subnets are in {subnet_vpc}')
        all_ok=False
    else:
        print(f'OK:{sgid} in {vpc}')
print('RESULT=' + ('PASS' if all_ok else 'FAIL'))
" 2>/dev/null || echo "RESULT=SKIP")

    echo "$SG_VPC_CHECK" | grep -v "^RESULT=" | sed 's/^OK:/  [OK]   SG /;s/^MISMATCH:/  [FAIL] SG /' || true
    SG_RESULT=$(echo "$SG_VPC_CHECK" | grep "^RESULT=" | cut -d= -f2)
    if [[ "$SG_RESULT" == "PASS" ]]; then
      pass "SecurityGroup VPC alignment"
    elif [[ "$SG_RESULT" == "FAIL" ]]; then
      fail "SecurityGroup VPC alignment" "SG and subnet must be in the same VPC → references/node-diagnostics-detail.md § B (VPC / Routing)"
    else
      # SG_RESULT is "SKIP" (json parse error) or empty (describe-security-groups
      # returned nothing). Either way the check did not run — say so, don't
      # leave the customer staring at a missing line.
      warn "SecurityGroup VPC alignment" "Unable to verify — describe-security-groups returned no usable data (check IAM ec2:DescribeSecurityGroups) → references/node-diagnostics-detail.md § B (VPC / Routing)"
    fi
  fi
fi

header "2a. VPC DNS Support & Hostnames"

# HyperPod requires enableDnsSupport + enableDnsHostnames on the VPC so that
# EKS service DNS and node internal hostnames resolve correctly.
if [[ -n "$UNIQUE_VPCS" && "$UNIQUE_VPCS" != *,* ]]; then
  DNS_SUPPORT=$(aws ec2 describe-vpc-attribute \
    --vpc-id "$UNIQUE_VPCS" --attribute enableDnsSupport \
    --region "$REGION" \
    --query 'EnableDnsSupport.Value' --output text 2>/dev/null || echo "unknown")
  DNS_HOSTNAMES=$(aws ec2 describe-vpc-attribute \
    --vpc-id "$UNIQUE_VPCS" --attribute enableDnsHostnames \
    --region "$REGION" \
    --query 'EnableDnsHostnames.Value' --output text 2>/dev/null || echo "unknown")

  if [[ "$DNS_SUPPORT" == "True" ]]; then
    pass "VPC enableDnsSupport" "enabled"
  else
    fail "VPC enableDnsSupport" "must be True — EKS internal DNS and node hostname resolution will fail. → references/node-diagnostics-detail.md § B (VPC / Routing)"
  fi
  if [[ "$DNS_HOSTNAMES" == "True" ]]; then
    pass "VPC enableDnsHostnames" "enabled"
  else
    fail "VPC enableDnsHostnames" "must be True — EKS internal DNS and node hostname resolution will fail. → references/node-diagnostics-detail.md § B (VPC / Routing)"
  fi
else
  warn "VPC DNS attributes" "skipped — subnets span multiple VPCs or no VPC resolved"
fi

header "2b. Private Subnet / Routing"

# HyperPod requires private subnets — a subnet is "public" if its route table has
# a default route (0.0.0.0/0) pointing at an internet gateway. For outbound
# access from a private subnet, the default route must point at a NAT gateway
# (or be absent in a fully air-gapped VPC that relies on VPC endpoints).
if [[ -n "$SUBNET_IDS" ]]; then
  PRIVATE_CHECK=$(aws ec2 describe-route-tables \
    --filters "Name=association.subnet-id,Values=$(echo "$SUBNET_IDS" | tr ' ' ',')" \
    --region "$REGION" \
    --query "RouteTables[*].{SubnetAssoc:Associations[?SubnetId!=\`null\`].SubnetId,Routes:Routes[?DestinationCidrBlock==\`0.0.0.0/0\`].{Target:GatewayId,NatGw:NatGatewayId}}" \
    --output json 2>/dev/null || echo '[]')

  echo "$PRIVATE_CHECK" | python3 -c "
import sys, json
rts = json.load(sys.stdin)
if not rts:
    print('INFO:no route tables associated — subnets likely use the main route table')
    sys.exit(0)
for rt in rts:
    subs = rt.get('SubnetAssoc', []) or []
    routes = rt.get('Routes', []) or []
    for r in routes:
        tgt = (r.get('Target') or '') or ''
        nat = (r.get('NatGw') or '') or ''
        subs_str = ','.join(subs) if subs else '(main)'
        if tgt.startswith('igw-'):
            print(f'FAIL:Subnet(s) {subs_str} route 0.0.0.0/0 -> Internet Gateway ({tgt}). HyperPod requires PRIVATE subnets; use a NAT gateway instead.')
        elif nat.startswith('nat-'):
            print(f'PASS:Subnet(s) {subs_str} route 0.0.0.0/0 -> NAT Gateway ({nat}) — private subnet, outbound via NAT.')
        elif tgt.startswith('vpce-'):
            print(f'INFO:Subnet(s) {subs_str} route 0.0.0.0/0 -> VPC endpoint ({tgt})')
        else:
            print(f'INFO:Subnet(s) {subs_str} route 0.0.0.0/0 -> {tgt or nat or \"unknown\"}')
" 2>/dev/null | while IFS=: read -r level msg; do
    case "$level" in
      PASS) pass "Private subnet routing" "$msg" ;;
      FAIL) fail "Private subnet routing" "$msg → references/node-diagnostics-detail.md § B (VPC / Routing)" ;;
      WARN) warn "Private subnet routing" "$msg" ;;
      INFO) info "$msg" ;;
    esac
  done
fi

header "3. IP Address Availability"

if [[ -n "$SUBNET_IDS" ]]; then
  _IP_CHECK=$(echo "$SUBNET_JSON" | python3 -c "
import sys,json
subnets=json.load(sys.stdin).get('Subnets',[])
for s in subnets:
    free=s.get('AvailableIpAddressCount',0)
    sid=s.get('SubnetId','?')
    if free < 5:
        print(f'FAIL:{sid} only {free} free IPs — CRITICALLY LOW')
    elif free < 50:
        print(f'WARN:{sid} only {free} free IPs — consider expanding CIDR')
    else:
        print(f'PASS:{sid} has {free} free IPs')
" 2>/dev/null || echo "")

  while IFS= read -r line; do
    [[ -z "$line" ]] && continue
    level=$(echo "$line" | cut -d: -f1)
    msg=$(echo "$line" | cut -d: -f2-)
    case "$level" in
      FAIL) fail "IP availability" "$msg → references/node-diagnostics-detail.md § B (VPC / Routing)" ;;
      WARN) warn "IP availability" "$msg" ;;
      PASS) pass "IP availability" "$msg" ;;
    esac
  done <<< "$_IP_CHECK"
fi

header "4. ENI Limits"

if [[ -n "$UNIQUE_VPCS" ]]; then
  VPC_ID=$(echo "$UNIQUE_VPCS" | tr ',' '\n' | head -1)
  ENI_COUNT=$(aws ec2 describe-network-interfaces \
    --filters "Name=vpc-id,Values=$VPC_ID" \
    --region "$REGION" \
    --query 'length(NetworkInterfaces)' \
    --output text 2>/dev/null || echo "unknown")

  ENI_QUOTA=$(aws service-quotas get-service-quota \
    --service-code ec2 \
    --quota-code "$ENI_QUOTA_CODE" \
    --region "$REGION" \
    --query 'Quota.Value' \
    --output text 2>/dev/null || echo "unknown")

  info "Current ENI count in VPC $VPC_ID: $ENI_COUNT"
  info "ENI quota for region: $ENI_QUOTA"

  if [[ "$ENI_COUNT" != "unknown" && "$ENI_QUOTA" != "unknown" ]]; then
    USAGE_PCT=$(python3 -c "q=int(${ENI_QUOTA}); print(int(${ENI_COUNT}/q*100) if q > 0 else '?')" 2>/dev/null || echo "?")
    if [[ "$USAGE_PCT" != "?" && "$USAGE_PCT" -gt 80 ]]; then
      warn "ENI limits" "${USAGE_PCT}% of quota used — request increase via Service Quotas if provisioning fails → references/node-diagnostics-detail.md § B (VPC / Routing)"
    else
      pass "ENI limits" "${ENI_COUNT}/${ENI_QUOTA} ENIs used (${USAGE_PCT}%)"
    fi
  else
    warn "ENI limits" "Could not determine ENI usage — verify manually → references/node-diagnostics-detail.md § B (VPC / Routing)"
  fi
fi

if [[ "$ORCHESTRATOR" == "EKS" && -n "$EKS_NAME" ]]; then
  header "5. EKS Prerequisites"

  EKS_DESC=$(aws eks describe-cluster \
    --name "$EKS_NAME" \
    --region "$REGION" \
    --output json 2>/dev/null || echo '{}')

  # VPC alignment — the EKS cluster's VPC must match the HyperPod cluster's VPC.
  EKS_VPC=$(echo "$EKS_DESC" | python3 -c "import sys,json; print(json.load(sys.stdin).get('cluster',{}).get('resourcesVpcConfig',{}).get('vpcId',''))" 2>/dev/null || echo "")
  if [[ -n "$EKS_VPC" && -n "$UNIQUE_VPCS" ]]; then
    if [[ ",$UNIQUE_VPCS," == *",$EKS_VPC,"* ]]; then
      pass "EKS VPC alignment" "EKS cluster in same VPC as HyperPod ($EKS_VPC)"
    else
      fail "EKS VPC alignment" "EKS cluster is in VPC $EKS_VPC but HyperPod subnets are in $UNIQUE_VPCS — they must match → references/node-diagnostics-detail.md § B (VPC / Routing)"
    fi
  fi

  # SG cross-reference — the HyperPod cluster SG must either be attached to the
  # EKS cluster, OR the EKS cluster SG must allow inbound from the HyperPod SG.
  EKS_SGS=$(echo "$EKS_DESC" | python3 -c "
import sys,json
d=json.load(sys.stdin).get('cluster',{}).get('resourcesVpcConfig',{})
all_sgs = set(d.get('securityGroupIds',[]) or [])
csg = d.get('clusterSecurityGroupId','')
if csg: all_sgs.add(csg)
print(' '.join(sorted(all_sgs)))
" 2>/dev/null || echo "")

  if [[ -n "$EKS_SGS" && -n "$SG_IDS" ]]; then
    HP_SG_SET=$(echo "$SG_IDS" | tr ',' ' ')
    SG_ATTACHED=false
    for hp in $HP_SG_SET; do
      for eks in $EKS_SGS; do
        [[ "$hp" == "$eks" ]] && { SG_ATTACHED=true; break 2; }
      done
    done
    if "$SG_ATTACHED"; then
      pass "HyperPod SG on EKS" "HyperPod SG is attached to the EKS cluster"
    else
      EKS_SG_LIST=$(echo "$EKS_SGS" | tr ' ' ',' | sed 's/,$//')
      read -r -a EKS_SG_ARR <<< "$EKS_SGS"
      EKS_INGRESS=$(aws ec2 describe-security-groups \
        --group-ids "${EKS_SG_ARR[@]}" \
        --region "$REGION" --output json 2>/dev/null || echo '{"SecurityGroups":[]}')
      CROSS_OK=$(echo "$EKS_INGRESS" | HP_SGS="$SG_IDS" python3 -c "
import sys,json,os
hp=set(os.environ.get('HP_SGS','').replace(',', ' ').split())
sgs=json.load(sys.stdin).get('SecurityGroups',[])
for sg in sgs:
    for rule in sg.get('IpPermissions',[]):
        for pair in rule.get('UserIdGroupPairs',[]):
            if pair.get('GroupId','') in hp:
                print('YES'); sys.exit(0)
print('NO')
" 2>/dev/null || echo "UNKNOWN")
      if [[ "$CROSS_OK" == "YES" ]]; then
        pass "HyperPod<->EKS SG" "EKS cluster SG ($EKS_SG_LIST) allows inbound from HyperPod SG"
      else
        fail "HyperPod<->EKS SG" "HyperPod SG is NOT attached to EKS and EKS SG ($EKS_SG_LIST) does not allow inbound from HyperPod SG → references/node-diagnostics-detail.md § A (EFA / Security Group)"
      fi
    fi
  fi

  EKS_AUTH=$(echo "$EKS_DESC" | python3 -c "import sys,json; print(json.load(sys.stdin).get('cluster',{}).get('accessConfig',{}).get('authenticationMode','unknown'))" 2>/dev/null || echo "unknown")

  if [[ "$EKS_AUTH" == "CONFIG_MAP" ]]; then
    warn "EKS auth mode" "CONFIG_MAP-only; access entries require API or API_AND_CONFIG_MAP — see the EKS access-entries documentation for the switching procedure"
  elif [[ "$EKS_AUTH" == "API" || "$EKS_AUTH" == "API_AND_CONFIG_MAP" ]]; then
    pass "EKS auth mode" "$EKS_AUTH"
  else
    warn "EKS auth mode" "Could not determine ($EKS_AUTH) — verify manually"
  fi

  # EKS endpoint accessibility (reuses $EKS_DESC captured above).
  PUB=$(echo "$EKS_DESC" | python3 -c "import sys,json; print(json.load(sys.stdin).get('cluster',{}).get('resourcesVpcConfig',{}).get('endpointPublicAccess',False))" 2>/dev/null || echo "false")
  PRIV=$(echo "$EKS_DESC" | python3 -c "import sys,json; print(json.load(sys.stdin).get('cluster',{}).get('resourcesVpcConfig',{}).get('endpointPrivateAccess',False))" 2>/dev/null || echo "false")

  info "EKS endpoint: public=$PUB, private=$PRIV"
  if [[ "$PUB" == "False" && "$PRIV" == "True" ]]; then
    warn "EKS endpoint" "Private-only endpoint — ensure worker subnets can reach EKS API (port 443), create EKS VPC endpoint if needed"
  elif [[ "$PUB" == "True" ]]; then
    pass "EKS endpoint" "Public access enabled"
  fi

  if command -v kubectl &>/dev/null; then
    if kubectl get namespace aws-hyperpod &>/dev/null 2>&1; then
      pass "aws-hyperpod namespace" "exists"
    else
      fail "aws-hyperpod namespace" "Missing → references/node-diagnostics-detail.md § B (VPC / Routing)"
    fi
  else
    warn "aws-hyperpod namespace" "kubectl not found — check skipped"
  fi
fi

header "6. VPC Endpoints"

if [[ -n "$UNIQUE_VPCS" ]]; then
  VPC_ID=$(echo "$UNIQUE_VPCS" | tr ',' '\n' | head -1)
  ENDPOINTS=$(aws ec2 describe-vpc-endpoints \
    --filters "Name=vpc-id,Values=$VPC_ID" \
    --region "$REGION" \
    --query "VpcEndpoints[?State==\`available\`].ServiceName" \
    --output text 2>/dev/null || echo "")

  # Required for private/air-gapped VPCs. Port 443 is the default for every
  # interface endpoint below; S3 uses a Gateway endpoint over the route table.
  # FSx users additionally need com.amazonaws.<region>.fsx if using FSx on Lustre/OpenZFS.
  REQUIRED_ENDPOINTS=("s3" "ecr.api" "ecr.dkr" "sts" "ssm" "ssmmessages" "ec2messages" "ec2" "sagemaker.api" "sagemaker.runtime" "logs")
  for svc in "${REQUIRED_ENDPOINTS[@]}"; do

    if echo "$ENDPOINTS" | grep -qE "\.${svc}$|\.${svc}[^a-z]"; then
      pass "VPC endpoint: $svc"
    else
      warn "VPC endpoint: $svc" "not found — required for internet-disabled (private) VPCs; skip if outbound 0.0.0.0/0 via NAT is available → references/node-diagnostics-detail.md § B (VPC / Routing)"
    fi
  done

  if [[ "$ORCHESTRATOR" == "EKS" ]]; then
    if echo "$ENDPOINTS" | grep -qE "\.eks$|\.eks[^a-z]"; then
      pass "VPC endpoint: eks"
    else
      warn "VPC endpoint: eks" "not found — needed if EKS endpoint is private-only → references/node-diagnostics-detail.md § B (VPC / Routing)"
    fi
  fi

  if ! echo "$ENDPOINTS" | grep -qE "\.fsx"; then
    info "VPC endpoint: fsx — not present (only required if this cluster uses FSx for Lustre or OpenZFS in a private/air-gapped VPC)"
  fi
fi

echo ""
echo -e "${BOLD}--- Summary ---${NC}"

if [[ $CRITICAL_FAILURES -eq 0 ]]; then
  echo -e "  ${GREEN}${BOLD}VPC configuration checks PASSED (${CRITICAL_FAILURES} critical issues).${NC}"
  echo "  If cluster creation still fails, check EFA security group rules:"
  echo "  bash check-efa-sg.sh --sg-id <SG_ID> --region $REGION"
else
  echo -e "  ${RED}${BOLD}VPC configuration checks FAILED (${CRITICAL_FAILURES} critical issue(s)).${NC}"
  echo "  Fix the [FAIL] items above and retry cluster creation."
fi
echo ""

exit "$([[ $CRITICAL_FAILURES -eq 0 ]] && echo 0 || echo 1)"