|
| 1 | +#!/usr/bin/env bash |
| 2 | +set -euo pipefail |
| 3 | + |
| 4 | +# ------------------------- |
| 5 | +# User-configurable defaults |
| 6 | +# ------------------------- |
| 7 | + |
| 8 | +# Default mount directories on the host machine |
| 9 | +DATASETS_DIR="${DATASETS_DIR:-$HOME/datasets}" |
| 10 | +MODELS_DIR="${MODELS_DIR:-$HOME/models}" |
| 11 | +EVAL_DIR="${EVAL_DIR:-$HOME/eval}" |
| 12 | + |
| 13 | +# Docker image name and tag for the VLN policy server |
| 14 | +DOCKER_IMAGE_NAME="${DOCKER_IMAGE_NAME:-vln_policy_server}" |
| 15 | +DOCKER_VERSION_TAG="${DOCKER_VERSION_TAG:-latest}" |
| 16 | + |
| 17 | +# Rebuild controls |
| 18 | +FORCE_REBUILD="${FORCE_REBUILD:-false}" |
| 19 | +NO_CACHE="" |
| 20 | + |
| 21 | +# Server parameters (can also be overridden via environment variables) |
| 22 | +HOST="${HOST:-0.0.0.0}" |
| 23 | +PORT="${PORT:-5555}" |
| 24 | +API_TOKEN="${API_TOKEN:-}" |
| 25 | +TIMEOUT_MS="${TIMEOUT_MS:-15000}" |
| 26 | +POLICY_TYPE="${POLICY_TYPE:-isaaclab_arena_navila.navila_server_policy.NaVilaServerPolicy}" |
| 27 | + |
| 28 | +# GPU selection for docker --gpus (can also be overridden via environment variables) |
| 29 | +# Examples: |
| 30 | +# all -> use all GPUs |
| 31 | +# 1 -> use 1 GPU (count) |
| 32 | +# "device=0" -> use GPU 0 |
| 33 | +# "device=0,1" -> use GPU 0 and 1 |
| 34 | +GPUS="${GPUS:-all}" |
| 35 | + |
| 36 | +# ------------------------- |
| 37 | +# Help message |
| 38 | +# ------------------------- |
| 39 | +usage() { |
| 40 | + script_name=$(basename "$0") |
| 41 | + cat <<EOF |
| 42 | +Helper script to build and run the VLN policy server Docker environment. |
| 43 | +
|
| 44 | +Usage: |
| 45 | + $script_name [options] [-- server-args...] |
| 46 | +
|
| 47 | +Options (Docker / paths; env vars with the same name take precedence): |
| 48 | + -v Verbose output (set -x). |
| 49 | + -d <datasets directory> Path to datasets on the host. Default: "$DATASETS_DIR". |
| 50 | + -m <models directory> Path to models on the host. Default: "$MODELS_DIR". |
| 51 | + -e <eval directory> Path to evaluation data on the host. Default: "$EVAL_DIR". |
| 52 | + -n <docker name> Docker image name. Default: "$DOCKER_IMAGE_NAME". |
| 53 | + -g <gpus> GPU selection for docker --gpus. Default: "all". |
| 54 | + Examples: "all", "1", "device=0", "device=0,1". |
| 55 | + -r Force rebuilding of the Docker image. |
| 56 | + -R Force rebuilding of the Docker image, without cache. |
| 57 | +
|
| 58 | +Server-specific options (passed through to the policy server entrypoint): |
| 59 | + --host HOST |
| 60 | + --port PORT |
| 61 | + --api_token TOKEN |
| 62 | + --timeout_ms MS |
| 63 | + --policy_type TYPE |
| 64 | + --model_path PATH |
| 65 | + --num_video_frames N |
| 66 | + --conv_mode MODE |
| 67 | +
|
| 68 | +Examples: |
| 69 | + # Minimal: use all defaults (model at default path, all GPUs, port 5555) |
| 70 | + bash $script_name |
| 71 | +
|
| 72 | + # Custom model path and single GPU |
| 73 | + bash $script_name -m /path/to/navila-checkpoint -g "device=0" |
| 74 | +
|
| 75 | + # Custom port |
| 76 | + bash $script_name --port 6000 |
| 77 | +EOF |
| 78 | +} |
| 79 | + |
| 80 | +# ------------------------- |
| 81 | +# Parse all options |
| 82 | +# ------------------------- |
| 83 | +# Server parameters that can be overridden individually via --flags. |
| 84 | +# Unset values will be filled with defaults after parsing. |
| 85 | +MODEL_PATH="" |
| 86 | + |
| 87 | +while [[ $# -gt 0 ]]; do |
| 88 | + case "$1" in |
| 89 | + -v) set -x; shift 1 ;; |
| 90 | + -d) DATASETS_DIR="$2"; shift 2 ;; |
| 91 | + -m) MODELS_DIR="$2"; shift 2 ;; |
| 92 | + -e) EVAL_DIR="$2"; shift 2 ;; |
| 93 | + -n) DOCKER_IMAGE_NAME="$2"; shift 2 ;; |
| 94 | + -g) GPUS="$2"; shift 2 ;; |
| 95 | + -r) FORCE_REBUILD="true"; shift 1 ;; |
| 96 | + -R) FORCE_REBUILD="true"; NO_CACHE="--no-cache"; shift 1 ;; |
| 97 | + -h|--help) usage; exit 0 ;; |
| 98 | + --host) HOST="$2"; shift 2 ;; |
| 99 | + --port) PORT="$2"; shift 2 ;; |
| 100 | + --api_token) API_TOKEN="$2"; shift 2 ;; |
| 101 | + --timeout_ms) TIMEOUT_MS="$2"; shift 2 ;; |
| 102 | + --policy_type) POLICY_TYPE="$2"; shift 2 ;; |
| 103 | + --model_path) MODEL_PATH="$2"; shift 2 ;; |
| 104 | + --num_video_frames) NUM_VIDEO_FRAMES="$2"; shift 2 ;; |
| 105 | + --conv_mode) CONV_MODE="$2"; shift 2 ;; |
| 106 | + --policy_device) POLICY_DEVICE="$2"; shift 2 ;; |
| 107 | + *) |
| 108 | + echo "Unknown option: $1" >&2 |
| 109 | + usage |
| 110 | + exit 1 |
| 111 | + ;; |
| 112 | + esac |
| 113 | +done |
| 114 | + |
| 115 | +# Build server args: always include all parameters, using defaults for |
| 116 | +# anything the user did not override. |
| 117 | +SERVER_ARGS=( |
| 118 | + --host "${HOST}" |
| 119 | + --port "${PORT}" |
| 120 | + --timeout_ms "${TIMEOUT_MS}" |
| 121 | + --policy_type "${POLICY_TYPE}" |
| 122 | + --model_path "${MODEL_PATH:-/models}" |
| 123 | +) |
| 124 | +[ -n "${API_TOKEN}" ] && SERVER_ARGS+=(--api_token "${API_TOKEN}") |
| 125 | +[ -n "${NUM_VIDEO_FRAMES+x}" ] && SERVER_ARGS+=(--num_video_frames "${NUM_VIDEO_FRAMES}") |
| 126 | +[ -n "${CONV_MODE+x}" ] && SERVER_ARGS+=(--conv_mode "${CONV_MODE}") |
| 127 | +[ -n "${POLICY_DEVICE+x}" ] && SERVER_ARGS+=(--policy_device "${POLICY_DEVICE}") |
| 128 | + |
| 129 | +echo "Host paths:" |
| 130 | +echo " DATASETS_DIR = ${DATASETS_DIR}" |
| 131 | +echo " MODELS_DIR = ${MODELS_DIR}" |
| 132 | +echo " EVAL_DIR = ${EVAL_DIR}" |
| 133 | +echo "Docker image:" |
| 134 | +echo " ${DOCKER_IMAGE_NAME}:${DOCKER_VERSION_TAG}" |
| 135 | +echo "GPU:" |
| 136 | +echo " --gpus ${GPUS}" |
| 137 | +echo "Rebuild:" |
| 138 | +echo " FORCE_REBUILD = ${FORCE_REBUILD}, NO_CACHE = '${NO_CACHE}'" |
| 139 | +echo "Server args:" |
| 140 | +printf ' %q ' "${SERVER_ARGS[@]}"; echo |
| 141 | + |
| 142 | +# ------------------------- |
| 143 | +# 1) Build the Docker image |
| 144 | +# ------------------------- |
| 145 | + |
| 146 | +IMAGE_TAG_FULL="${DOCKER_IMAGE_NAME}:${DOCKER_VERSION_TAG}" |
| 147 | + |
| 148 | +SHOULD_BUILD=false |
| 149 | + |
| 150 | +if [ "${FORCE_REBUILD}" = "true" ]; then |
| 151 | + SHOULD_BUILD=true |
| 152 | +else |
| 153 | + if [ -z "$(docker images -q "${IMAGE_TAG_FULL}")" ]; then |
| 154 | + SHOULD_BUILD=true |
| 155 | + fi |
| 156 | +fi |
| 157 | + |
| 158 | +if [ "${SHOULD_BUILD}" = "true" ]; then |
| 159 | + echo "Building Docker image ${IMAGE_TAG_FULL}..." |
| 160 | + # Use existing image layers as cache source (BuildKit may GC intermediate |
| 161 | + # layer cache, but the final image layers can still be reused). |
| 162 | + CACHE_FROM_ARGS="" |
| 163 | + if [ -n "$(docker images -q "${IMAGE_TAG_FULL}" 2>/dev/null)" ]; then |
| 164 | + CACHE_FROM_ARGS="--cache-from ${IMAGE_TAG_FULL}" |
| 165 | + fi |
| 166 | + docker build \ |
| 167 | + ${NO_CACHE} \ |
| 168 | + ${CACHE_FROM_ARGS} \ |
| 169 | + --network host \ |
| 170 | + -f docker/Dockerfile.vln_server \ |
| 171 | + -t "${IMAGE_TAG_FULL}" \ |
| 172 | + . |
| 173 | +else |
| 174 | + echo "Docker image ${IMAGE_TAG_FULL} already exists. Skipping rebuild." |
| 175 | + echo "Use -r or -R to force rebuilding the image." |
| 176 | +fi |
| 177 | + |
| 178 | +# ------------------------- |
| 179 | +# 2) Run the container |
| 180 | +# ------------------------- |
| 181 | + |
| 182 | +DOCKER_RUN_ARGS=( |
| 183 | + --rm |
| 184 | + --gpus "${GPUS}" |
| 185 | + --net host |
| 186 | + --name vln_policy_server_container |
| 187 | + -v "${MODELS_DIR}":/models |
| 188 | +) |
| 189 | + |
| 190 | +if [ -d "${DATASETS_DIR}" ]; then |
| 191 | + DOCKER_RUN_ARGS+=(-v "${DATASETS_DIR}":/datasets) |
| 192 | +fi |
| 193 | + |
| 194 | +if [ -d "${EVAL_DIR}" ]; then |
| 195 | + DOCKER_RUN_ARGS+=(-v "${EVAL_DIR}":/eval) |
| 196 | +fi |
| 197 | + |
| 198 | +docker run "${DOCKER_RUN_ARGS[@]}" \ |
| 199 | + "${IMAGE_TAG_FULL}" \ |
| 200 | + "${SERVER_ARGS[@]}" |
0 commit comments