# files/nfs-share.sh — spark2/spark3 read spark1's checkpoint over NFSv4 on ConnectX.
# Sourced by start.sh. Same pattern as vllm-fn-nfs / GLM-5.3 (no local 476 GiB copy).
#
# Pairwise CX7: spark2 talks to head 10.0.22.1, spark3 talks to head 10.0.23.1.
# Live kernel nfsd (vllm-fn-nfs) is reused when present — a second nfsd will not start.

NFS_IMAGE="${NFS_IMAGE:-dsv41-nfs:local}"
NFS_CONTAINER="${NFS_CONTAINER:-dsv41-nfs}"
NFS_VOLUME="${NFS_VOLUME:-dsv41-weights}"
NFS_DOCKERFILE_DIR="${NFS_DOCKERFILE_DIR:-$ROOT/files/nfs-server}"
NFS_EXPORT_NAME="${NFS_EXPORT_NAME:-dsv41-native}"
HF_EXPORT_ROOT="${HF_EXPORT_ROOT:-$HOME/.cache/huggingface}"
NFS_OPTS_CLIENT="${NFS_OPTS_CLIENT:-nfsvers=4.2,ro,nconnect=8,rsize=1048576,wsize=1048576,hard,timeo=600}"

nfs_clients() {
  echo "${NFS_CLIENTS:-10.0.0.2,10.0.0.3,10.0.22.0/24,10.0.23.0/24,10.0.33.0/24}"
}

# Head ConnectX address the given worker should mount (not 10.0.0.1).
nfs_server_ip_for() {
  local host="$1" idx="$2"
  # Per-worker override: NFS_SERVER_IPS="10.0.22.1 10.0.23.1 ..." (one entry per
  # worker, or a single entry used for all, e.g. behind a switch), else the legacy
  # NFS_SERVER_IP_<n>, else NFS_SERVER_IP.
  local explicit=""
  local -a _nfs_ips=()
  if [[ -n "${NFS_SERVER_IPS:-}" ]]; then
    read -r -a _nfs_ips <<<"$(tr ',' ' ' <<<"$NFS_SERVER_IPS")"
    if [[ ${#_nfs_ips[@]} -eq 1 ]]; then explicit="${_nfs_ips[0]}"; else explicit="${_nfs_ips[$idx]:-}"; fi
  fi
  if [[ -z "$explicit" ]]; then
    local _v="NFS_SERVER_IP_$((idx + 1))"
    explicit="${!_v:-${NFS_SERVER_IP:-}}"
  fi
  if [[ -n "$explicit" ]]; then
    echo "$explicit"
    return 0
  fi
  local peer="${WORKER_IPS[$idx]}"
  # Prefer the CX7 source address toward that worker's fabric IP if set:
  # WORKER_FABRIC_IPS="10.0.22.2 10.0.23.3 ..." or legacy WORKER<n>_FABRIC_IP.
  local fabric_peer=""
  local -a _fab=()
  if [[ -n "${WORKER_FABRIC_IPS:-}" ]]; then
    read -r -a _fab <<<"$(tr ',' ' ' <<<"$WORKER_FABRIC_IPS")"
    fabric_peer="${_fab[$idx]:-}"
  fi
  if [[ -z "$fabric_peer" ]]; then
    local _f="WORKER$((idx + 1))_FABRIC_IP"
    case "$idx" in
      0) fabric_peer="${!_f:-10.0.22.2}" ;;
      1) fabric_peer="${!_f:-10.0.23.3}" ;;
      *) fabric_peer="${!_f:-$peer}" ;;
    esac
  fi
  local src
  src=$(ip route get "$fabric_peer" 2>/dev/null \
    | awk '{ for (i=1;i<=NF;i++) if ($i=="src") { print $(i+1); exit } }')
  if [[ -n "$src" && "$src" != "10.0.0.1" && "$src" != "127.0.0.1" ]]; then
    echo "$src"
    return 0
  fi
  src=$(ip route get "$peer" 2>/dev/null \
    | awk '{ for (i=1;i<=NF;i++) if ($i=="src") { print $(i+1); exit } }')
  echo "${src:-10.0.22.1}"
}

nfs_live_container() {
  docker ps --format '{{.Names}}' | grep -E '^(vllm-fn-nfs|glm53-nfs|dsv41-nfs)$' | head -1
}

nfs_rpc_ready() {
  local ip="$1"
  timeout 3 rpcinfo -t "$ip" nfs 4 >/dev/null 2>&1
}

nfs_write_exports() {
  local ctn="$1" clients="$2"
  docker exec -e NFS_CLIENTS="$clients" -e NFS_OPTS="${NFS_OPTS:-ro,sync,no_subtree_check,no_root_squash,insecure,fsid=0}" \
    "$ctn" bash -lc '
      set -e
      {
        echo "# generated by dsv41-nfs (merged clients)"
        IFS="," read -ra addrs <<< "$NFS_CLIENTS"
        for addr in "${addrs[@]}"; do
          addr="${addr// /}"
          [ -z "$addr" ] && continue
          echo "/export ${addr}(${NFS_OPTS})"
        done
      } > /etc/exports
      exportfs -rav
    '
}

nfs_ensure_server() {
  local clients
  clients="$(nfs_clients)"
  local live
  live="$(nfs_live_container || true)"

  if [[ -n "$live" ]] && nfs_rpc_ready 127.0.0.1; then
    local export_source
    export_source=$(docker inspect -f '{{range .Mounts}}{{if eq .Destination "/export"}}{{.Source}}{{end}}{{end}}' "$live")
    if [[ "$export_source" == "$MODEL_DIR" ]]; then
      NFS_REUSE_EXPORT=0
    elif [[ "$export_source" == "$HF_EXPORT_ROOT" ]]; then
      NFS_REUSE_EXPORT=1
    else
      die "NFS $live exports $export_source; expected $MODEL_DIR or $HF_EXPORT_ROOT"
    fi
    info "NFS already up ($live) — configuring worker clients"
    nfs_write_exports "$live" "$clients"
    NFS_LIVE_CTN="$live"
    return 0
  fi

  [[ -d "$NFS_DOCKERFILE_DIR" ]] || die "missing $NFS_DOCKERFILE_DIR"
  [[ -d "$MODEL_DIR" ]] || die "checkpoint missing: $MODEL_DIR"

  info "building NFS image $NFS_IMAGE"
  docker build -q -t "$NFS_IMAGE" "$NFS_DOCKERFILE_DIR" >/dev/null
  docker rm -f "$NFS_CONTAINER" >/dev/null 2>&1 || true
  info "exporting $MODEL_DIR via NFS (clients: $clients)"
  docker run -d --name "$NFS_CONTAINER" --restart unless-stopped \
    --privileged --network host \
    -v "$MODEL_DIR:/export:ro" \
    -e "NFS_CLIENTS=$clients" \
    "$NFS_IMAGE" >/dev/null
  local i
  for i in $(seq 1 20); do
    if nfs_rpc_ready 127.0.0.1; then
      info "NFS server ready ($NFS_CONTAINER)"
      NFS_REUSE_EXPORT=0
      NFS_LIVE_CTN="$NFS_CONTAINER"
      return 0
    fi
    sleep 0.5
  done
  docker logs "$NFS_CONTAINER" >&2 || true
  die "NFS server did not become ready. Check: docker logs $NFS_CONTAINER"
}

# When reusing vllm-fn-nfs (HF cache as NFSv4 root), publish the checkpoint
# into that tree with hardlinks (same ext4, zero extra bytes).
nfs_publish_model() {
  if [[ "${NFS_REUSE_EXPORT:-0}" != "1" ]]; then
    NFS_DEVICE=":/"
    return 0
  fi
  local dst="${HF_EXPORT_ROOT}/${NFS_EXPORT_NAME}"
  mkdir -p "$dst"
  if [[ -f "$dst/config.json" ]]; then
    local n
    n=$(find "$dst" -maxdepth 1 -name 'model-*-of-*.safetensors' 2>/dev/null | wc -l | tr -d ' ')
    if [[ "${n:-0}" -ge "$EXPECTED_SHARDS" ]]; then
      info "NFS export already has $NFS_EXPORT_NAME ($n shards)"
      NFS_DEVICE=":/${NFS_EXPORT_NAME}"
      return 0
    fi
  fi
  info "publishing $MODEL_DIR → $dst (hardlinks into the live NFS export)"
  cp -a --link "$MODEL_DIR"/. "$dst"/ 2>/dev/null \
    || cp -al "$MODEL_DIR"/. "$dst"/
  [[ -f "$dst/config.json" ]] || die "failed to publish checkpoint into $dst"
  NFS_DEVICE=":/${NFS_EXPORT_NAME}"
}

nfs_ensure_worker_volume() {
  local host="$1" idx="$2"
  local addr device
  addr="$(nfs_server_ip_for "$host" "$idx")"
  device="${NFS_DEVICE:-:/${NFS_EXPORT_NAME}}"
  info "NFS volume $NFS_VOLUME on $host → nfs://${addr}${device}"
  remote_on "$host" "
    set -e
    docker volume rm $(printf '%q' "$NFS_VOLUME") >/dev/null 2>&1 || true
    docker volume create --driver local \
      --opt type=nfs \
      --opt $(printf '%q' "o=addr=${addr},${NFS_OPTS_CLIENT}") \
      --opt $(printf '%q' "device=${device}") \
      $(printf '%q' "$NFS_VOLUME") >/dev/null
  "
}

local_model_has_weights() {
  [[ -r "$MODEL_DIR/config.json" && -s "$MODEL_DIR/config.json" ]] || return 1
  local -a shards=("$MODEL_DIR"/model-*-of-*.safetensors)
  [[ ${#shards[@]} -ge "$EXPECTED_SHARDS" ]] || return 1
  local shard
  for shard in "${shards[@]}"; do
    [[ -r "$shard" && -s "$shard" ]] || return 1
  done
}

nfs_worker_has_model() {
  local host="$1"
  # Inspect first: docker run -v would create an empty volume on a typo.
  # Use the serving image already required by serve; status must not pull an
  # unrelated mutable image, and the probe needs neither network nor GPUs.
  local check='test -r /m/config.json && test -s /m/config.json || exit 1
    expected=$1
    set -- /m/model-*-of-*.safetensors
    test "$#" -ge "$expected" || exit 1
    for shard do test -r "$shard" && test -s "$shard" || exit 1; done'
  remote_on "$host" --timeout 180 \
    "docker volume inspect $(printf '%q' "$NFS_VOLUME") >/dev/null 2>&1 &&
     docker run --rm --pull=never --network none \
       --mount $(printf '%q' "type=volume,source=$NFS_VOLUME,target=/m,readonly") \
       --entrypoint /bin/sh $(printf '%q' "$IMAGE") \
       -c $(printf '%q' "$check") sh $(printf '%q' "$EXPECTED_SHARDS")" \
    >/dev/null 2>&1
}

nfs_unmount_workers() {
  local h
  for h in "${WORKER_HOSTS[@]}"; do
    remote_on "$h" "docker volume rm $(printf '%q' "$NFS_VOLUME") >/dev/null 2>&1 || true" || true
  done
}