start_dynamo_unified.sh

#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Dynamo SGLang FULL STACK with Unified Worker
# Architecture: ETCD + NATS + Dynamo Frontend (API) → SGLang Backend Worker (Unified)
#
# This script manages ALL required components:
#   - ETCD (metadata and worker discovery)
#   - NATS (message queue for requests)
#   - Dynamo Frontend (HTTP API with built-in processor + router)
#   - Unified Worker (GPUs 0,1,2,3, TP=4, no disaggregation)
#
# Frontend: Port 8099 (HTTP API)
# ETCD: localhost:2379 (container: etcd-dynamo) - default port, override with DYNAMO_ETCD_PORT
# NATS: localhost:4222 (container: nats-dynamo) - default port, override with DYNAMO_NATS_PORT
# Worker runs in container: dynamo-sglang
#
# To stop all components: bash stop_dynamo.sh

# Configuration Variables (can be overridden via environment variables)
CONTAINER_NAME="dynamo-sglang"
WORKER_GPUS="${DYNAMO_GPU_DEVICES:-0,1,2,3,4,5,6,7}"
TP_SIZE="${DYNAMO_TP_SIZE:-2}"
HTTP_PORT="${DYNAMO_HTTP_PORT:-8000}"
SERVED_MODEL_NAME=""  # set after validation
IMAGE="nvcr.io/nvidia/ai-dynamo/sglang-runtime:0.9.0"
SHM_SIZE="${DYNAMO_SHM_SIZE:-16g}"

# Infrastructure ports (can be overridden via environment variables)
ETCD_CLIENT_PORT="${DYNAMO_ETCD_PORT:-2379}"
ETCD_PEER_PORT="${DYNAMO_ETCD_PEER_PORT:-2390}"
NATS_PORT="${DYNAMO_NATS_PORT:-4222}"
WORKER_INIT_TIMEOUT_S="${DYNAMO_WORKER_INIT_TIMEOUT_S:-1800}"

# KV Cache-Aware Routing (optional)
# Set ENABLE_KV_AWARE_ROUTING=true to enable KV cache-aware routing.
# This adds --kv-cache-block-size to the frontend and --page-size to each worker
# so the frontend can make routing decisions based on KV cache overlap.
# The block size (in tokens) must match between the frontend and all workers.
ENABLE_KV_AWARE_ROUTING="${ENABLE_KV_AWARE_ROUTING:-false}"
KV_BLOCK_SIZE="${DYNAMO_KV_BLOCK_SIZE:-64}"
# Prometheus metrics base port for workers (each worker gets WORKER_METRICS_PORT+i).
# --enable-metrics is always on so Prometheus/Grafana can scrape worker metrics.
WORKER_METRICS_PORT="${DYNAMO_WORKER_METRICS_PORT:-18081}"

# Worker performance tuning (can be overridden via environment variables)
# Fraction of GPU memory reserved for the KV cache (0.0-1.0)
MEM_FRACTION_STATIC="${DYNAMO_MEM_FRACTION_STATIC:-0.9}"
echo "MEM_FRACTION_STATIC=${MEM_FRACTION_STATIC} (from DYNAMO_MEM_FRACTION_STATIC=${DYNAMO_MEM_FRACTION_STATIC:-<unset, using default 0.9>})"
# Maximum sequence length the model will handle (unset = model default)
MAX_MODEL_LEN="${DYNAMO_MAX_MODEL_LEN:-}"
# Hard override for the number of GPU KV cache blocks (unset = auto)
NUM_GPU_BLOCKS_OVERRIDE="${DYNAMO_NUM_GPU_BLOCKS_OVERRIDE:-}"

# HiCache (hierarchical KV cache) configuration
# Enables CPU-backed overflow cache for the SGLang KV cache.
ENABLE_HIERARCHICAL_CACHE="${DYNAMO_ENABLE_HIERARCHICAL_CACHE:-false}"
HICACHE_RATIO="${DYNAMO_HICACHE_RATIO:-1.0}"
HICACHE_POLICY="${DYNAMO_HICACHE_POLICY:-write_through}"

# Validate HiCache settings when enabled
if [ "${ENABLE_HIERARCHICAL_CACHE}" = "true" ]; then
    if ! printf '%s' "$HICACHE_RATIO" | grep -qE '^[0-9]*\.?[0-9]+$' || \
       [ "$(awk -v v="$HICACHE_RATIO" 'BEGIN{print (v+0)<=0 ? 1 : 0}')" = "1" ]; then
        echo "ERROR: HICACHE_RATIO must be a positive number (got: '$HICACHE_RATIO')" >&2
        echo "  Set via DYNAMO_HICACHE_RATIO (e.g., 1.0)" >&2
        exit 1
    fi
    case "$HICACHE_POLICY" in
        write_through|write_back) ;;
        *)
            echo "ERROR: HICACHE_POLICY must be 'write_through' or 'write_back' (got: '$HICACHE_POLICY')" >&2
            echo "  Set via DYNAMO_HICACHE_POLICY" >&2
            exit 1
            ;;
    esac
fi

# Compute container-internal GPU indices (GPUs are renumbered 0,1,2,... inside the container)
NUM_GPUS=$(echo "$WORKER_GPUS" | tr ',' '\n' | wc -l)
CONTAINER_GPU_INDICES=$(seq -s, 0 $((NUM_GPUS - 1)))

# Calculate number of workers based on available GPUs and TP size
NUM_WORKERS=$((NUM_GPUS / TP_SIZE))

# Validate GPU/TP sizing
if [ "$TP_SIZE" -le 0 ] 2>/dev/null; then
    echo "ERROR: TP_SIZE must be a positive integer (got: '$TP_SIZE')" >&2
    echo "  WORKER_GPUS=$WORKER_GPUS  NUM_GPUS=$NUM_GPUS  TP_SIZE=$TP_SIZE" >&2
    exit 1
fi
if [ "$NUM_GPUS" -lt "$TP_SIZE" ]; then
    echo "ERROR: Not enough GPUs for the requested TP size (NUM_GPUS=$NUM_GPUS < TP_SIZE=$TP_SIZE)" >&2
    echo "  WORKER_GPUS=$WORKER_GPUS  NUM_GPUS=$NUM_GPUS  TP_SIZE=$TP_SIZE" >&2
    exit 1
fi
if [ $((NUM_GPUS % TP_SIZE)) -ne 0 ]; then
    echo "ERROR: NUM_GPUS ($NUM_GPUS) is not divisible by TP_SIZE ($TP_SIZE)" >&2
    echo "  WORKER_GPUS=$WORKER_GPUS  NUM_GPUS=$NUM_GPUS  TP_SIZE=$TP_SIZE  NUM_WORKERS would be $NUM_WORKERS" >&2
    exit 1
fi
if [ "$NUM_WORKERS" -le 0 ]; then
    echo "ERROR: NUM_WORKERS is 0 — no workers can be started with this GPU/TP configuration" >&2
    echo "  WORKER_GPUS=$WORKER_GPUS  NUM_GPUS=$NUM_GPUS  TP_SIZE=$TP_SIZE" >&2
    exit 1
fi

# Local paths - DYNAMO_MODEL_DIR must be set or script will error
if [ -z "${DYNAMO_MODEL_DIR}" ]; then
    echo "ERROR: DYNAMO_MODEL_DIR environment variable must be set"
    echo ""
    echo "Example:"
    echo "  export DYNAMO_MODEL_DIR=\"/path/to/your/models/Llama-3.3-70B-Instruct\""
    echo ""
    echo "Then run this script again."
    exit 1
fi
# If directory exists, validate it's a proper model directory (NVBug 5756833)
# If it doesn't exist, the download workflow later will handle it
if [ -d "${DYNAMO_MODEL_DIR}" ]; then
    if [ ! -f "${DYNAMO_MODEL_DIR}/config.json" ]; then
        echo "ERROR: ${DYNAMO_MODEL_DIR} exists but is not a valid model directory"
        echo ""
        echo "Missing: config.json"
        echo ""
        echo "Common mistake - pointing to cache root instead of model snapshot:"
        echo "   Wrong: ~/.cache/huggingface/"
        echo "   Right: ~/.cache/huggingface/hub/models--meta-llama--Llama-3.3-70B-Instruct/snapshots/<hash>"
        echo ""
        echo "Find it: find ~/.cache/huggingface/hub -name config.json -path '*Llama-3.3-70B*'"
        exit 1
    fi

    # Verify config.json has model_type field (exact error from NVBug 5756833)
    if ! grep -q '"model_type"' "${DYNAMO_MODEL_DIR}/config.json" 2>/dev/null; then
        echo "ERROR: ${DYNAMO_MODEL_DIR}/config.json is missing 'model_type' field"
        echo ""
        echo "This usually means incomplete/corrupted download. Try:"
        echo "  rm -rf ${DYNAMO_MODEL_DIR}"
        echo "  hf download meta-llama/Llama-3.3-70B-Instruct --local-dir ${DYNAMO_MODEL_DIR}"
        exit 1
    fi
fi
LOCAL_MODEL_DIR="$(eval echo "${DYNAMO_MODEL_DIR}")"
MODEL="/workspace/models/$(basename "$LOCAL_MODEL_DIR")"
SERVED_MODEL_NAME="${DYNAMO_MODEL_NAME:-$(basename "$LOCAL_MODEL_DIR")}"

echo "========================================================="
echo "Dynamo SGLang FULL STACK (UNIFIED MODE)"
echo "========================================================="
echo "Model: $SERVED_MODEL_NAME"
echo "Container: $CONTAINER_NAME"
echo "HTTP Port: $HTTP_PORT"
echo ""
echo "Components:"
echo "  - ETCD (metadata and discovery)"
echo "  - NATS (message queue for requests)"
echo "  - Dynamo Frontend (HTTP API on port $HTTP_PORT)"
echo "  - SGLang Worker (unified mode)"
echo ""
echo "Backend Workers:"
echo "  Workers: $NUM_WORKERS (GPUs: $NUM_GPUS, TP=$TP_SIZE per worker)"
echo "  GPUs: $WORKER_GPUS"
echo "  Mode: UNIFIED (no prefill/decode disaggregation)"
echo ""
echo "Routing Mode:"
if [ "${ENABLE_KV_AWARE_ROUTING}" = "true" ]; then
    echo "  KV Cache-Aware (ENABLE_KV_AWARE_ROUTING=true)"
    echo "  KV Block Size: $KV_BLOCK_SIZE tokens"
else
    echo "  Round-Robin (default)"
    echo "  Set ENABLE_KV_AWARE_ROUTING=true to enable KV cache-aware routing"
fi
echo ""
echo "HiCache:"
if [ "${ENABLE_HIERARCHICAL_CACHE}" = "true" ]; then
    echo "  Enabled (DYNAMO_ENABLE_HIERARCHICAL_CACHE=true)"
    echo "  Ratio: $HICACHE_RATIO, Policy: $HICACHE_POLICY"
else
    echo "  Disabled (default)"
    echo "  Set DYNAMO_ENABLE_HIERARCHICAL_CACHE=true to enable"
fi
echo ""
echo "========================================================="

# Start ETCD if not running
if docker ps -a --format '{{.Names}}' | grep -q "^etcd-dynamo$"; then
    # Container exists (running or stopped), remove it first
    echo ""
    echo "Removing existing ETCD container..."
    docker rm -f etcd-dynamo
fi

echo "Starting ETCD container..."
docker run -d \
  --name etcd-dynamo \
  --network host \
  -e ALLOW_NONE_AUTHENTICATION=yes \
  -e ETCD_LISTEN_CLIENT_URLS=http://0.0.0.0:$ETCD_CLIENT_PORT \
  -e ETCD_ADVERTISE_CLIENT_URLS=http://localhost:$ETCD_CLIENT_PORT \
  -e ETCD_LISTEN_PEER_URLS=http://0.0.0.0:$ETCD_PEER_PORT \
  -e ETCD_INITIAL_ADVERTISE_PEER_URLS=http://localhost:$ETCD_PEER_PORT \
  -e ETCD_INITIAL_CLUSTER=default=http://localhost:$ETCD_PEER_PORT \
  bitnamilegacy/etcd:3.6.1

# Wait for ETCD to be ready
echo "Waiting for ETCD to be ready..."
for i in {1..30}; do
    if curl -s http://localhost:$ETCD_CLIENT_PORT/health > /dev/null 2>&1; then
        echo "✓ ETCD is ready on port $ETCD_CLIENT_PORT"
        sleep 2  # Extra settling time
        break
    fi
    if [ $i -eq 30 ]; then
        echo "✗ ERROR: ETCD failed to start within 30 seconds"
        docker logs etcd-dynamo
        exit 1
    fi
    sleep 1
done

# Start NATS if not running
if docker ps -a --format '{{.Names}}' | grep -q "^nats-dynamo$"; then
    # Container exists (running or stopped), remove it first
    echo ""
    echo "Removing existing NATS container..."
    docker rm -f nats-dynamo
fi

echo "Starting NATS container..."
docker run -d \
  --name nats-dynamo \
  --network host \
  nats:2.11.4 \
  -js -p $NATS_PORT

# Wait for NATS to be ready
echo "Waiting for NATS to be ready..."
for i in {1..30}; do
    if timeout 2 bash -c "cat < /dev/null > /dev/tcp/localhost/$NATS_PORT" 2>/dev/null; then
        echo "✓ NATS is ready on port $NATS_PORT"
        break
    fi
    if [ $i -eq 30 ]; then
        echo "✗ ERROR: NATS failed to start within 30 seconds"
        docker logs nats-dynamo
        exit 1
    fi
    sleep 1
done

echo ""

# Start monitoring stack (Prometheus + Grafana) if not running
MONITORING_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/monitoring"
if [ -f "$MONITORING_DIR/docker-compose.yml" ]; then
    PROMETHEUS_RUNNING=$(docker ps --format '{{.Names}}' | grep -q "^dynamo-prometheus$" && echo "true" || echo "false")
    GRAFANA_RUNNING=$(docker ps --format '{{.Names}}' | grep -q "^dynamo-grafana$" && echo "true" || echo "false")

    if [ "$PROMETHEUS_RUNNING" = "false" ] || [ "$GRAFANA_RUNNING" = "false" ]; then
        echo "Starting monitoring stack (Prometheus + Grafana)..."
        cd "$MONITORING_DIR"
        docker compose up -d
        cd - > /dev/null

        # Wait for Prometheus to be ready
        echo "Waiting for Prometheus to be ready..."
        for i in {1..30}; do
            if curl -s http://localhost:9090/-/ready > /dev/null 2>&1; then
                echo "✓ Prometheus is ready (http://localhost:9090)"
                break
            fi
            if [ $i -eq 30 ]; then
                echo "⚠ WARNING: Prometheus may not be fully ready yet"
            fi
            sleep 1
        done

        # Wait for Grafana to be ready
        echo "Waiting for Grafana to be ready..."
        for i in {1..30}; do
            if curl -s http://localhost:3000/api/health > /dev/null 2>&1; then
                echo "✓ Grafana is ready (http://localhost:3000)"
                break
            fi
            if [ $i -eq 30 ]; then
                echo "⚠ WARNING: Grafana may not be fully ready yet"
            fi
            sleep 1
        done
        echo ""
    else
        echo "✓ Monitoring stack already running"
        echo "  Prometheus: http://localhost:9090"
        echo "  Grafana:    http://localhost:3000"
        echo ""
    fi
else
    echo "⚠ Monitoring docker-compose.yml not found at: $MONITORING_DIR"
    echo "  Skipping monitoring stack startup"
    echo ""
fi

# Clean up existing Dynamo container if it exists
if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
    echo "Removing existing Dynamo container: $CONTAINER_NAME"
    docker rm -f $CONTAINER_NAME
fi

# Verify HF_TOKEN is set
if [ -z "$HF_TOKEN" ]; then
    echo ""
    echo "⚠ HF_TOKEN environment variable is not set."
    echo ""
    echo "The model is cached locally at: $LOCAL_MODEL_DIR"
    if [ -d "$LOCAL_MODEL_DIR" ]; then
        echo "✓ Local model found - proceeding without HF_TOKEN"
        echo "  Note: Set HF_TOKEN if you need to download models from HuggingFace"
        HF_TOKEN="dummy"  # Set dummy token since model is cached
    else
        echo "✗ Local model NOT found and no HF_TOKEN to download it"
        echo ""
        read -p "Please enter your HuggingFace token (or press Enter to skip): " HF_TOKEN

        if [ -z "$HF_TOKEN" ]; then
            echo ""
            echo "WARNING: Proceeding without HF_TOKEN. This may fail if the model needs to be downloaded."
            echo "To set HF_TOKEN: export HF_TOKEN='your_token_here'"
            HF_TOKEN="dummy"
        else
            echo ""
            echo "✓ HuggingFace token received"
        fi
    fi
else
    echo "✓ HuggingFace token is set"
fi
echo ""


# Verify model exists locally
if [ ! -d "$LOCAL_MODEL_DIR" ]; then
    echo "WARNING: Model directory not found at: $LOCAL_MODEL_DIR"
    echo ""
    echo "To download the model, run:"
    echo "  hf download meta-llama/Llama-3.3-70B-Instruct --local-dir $LOCAL_MODEL_DIR"
    echo ""
    read -p "Continue anyway (model will be downloaded from HuggingFace)? [y/N] " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        exit 1
    fi
fi

# Start container with unified SGLang worker + Dynamo frontend
echo ""
echo "Starting Dynamo container with unified SGLang worker + frontend..."
docker run -d \
  --name $CONTAINER_NAME \
  --gpus "\"device=${WORKER_GPUS}\"" \
  --network host \
  --ipc=host \
  --shm-size=$SHM_SIZE \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  -v $LOCAL_MODEL_DIR:$MODEL:ro \
  -e HF_TOKEN="$HF_TOKEN" \
  -e HUGGING_FACE_HUB_TOKEN="$HF_TOKEN" \
  -e RUST_BACKTRACE=1 \
  -e PYTHONUNBUFFERED=1 \
  -e ETCD_ENDPOINTS=http://localhost:$ETCD_CLIENT_PORT \
  -e NATS_SERVER=nats://localhost:$NATS_PORT \
  -e WORKER_METRICS_PORT=$WORKER_METRICS_PORT \
  -e DYNAMO_WORKER_COMPONENT=backend \
  -e ENABLE_KV_AWARE_ROUTING=$ENABLE_KV_AWARE_ROUTING \
  -e KV_BLOCK_SIZE=$KV_BLOCK_SIZE \
  -e MAX_MODEL_LEN=$MAX_MODEL_LEN \
  -e NUM_GPU_BLOCKS_OVERRIDE=$NUM_GPU_BLOCKS_OVERRIDE \
  -e ENABLE_HIERARCHICAL_CACHE=$ENABLE_HIERARCHICAL_CACHE \
  -e HICACHE_RATIO=$HICACHE_RATIO \
  -e HICACHE_POLICY=$HICACHE_POLICY \
  $IMAGE \
  bash -c "
    set -e  # Exit on any error

    echo '========================================================='
    echo 'Verifying external infrastructure services...'
    echo '========================================================='

    # Verify ETCD is accessible
    if curl -s http://localhost:$ETCD_CLIENT_PORT/health > /dev/null 2>&1; then
        echo \"✓ ETCD accessible at localhost:$ETCD_CLIENT_PORT\"
    else
        echo \"✗ ERROR: ETCD not accessible at localhost:$ETCD_CLIENT_PORT\"
        echo '  Make sure ETCD container is running with --network host'
        exit 1
    fi

    # Verify NATS is accessible (basic TCP check)
    if timeout 2 bash -c '</dev/tcp/localhost/$NATS_PORT' 2>/dev/null; then
        echo \"✓ NATS accessible at localhost:$NATS_PORT\"
    else
        echo \"✗ ERROR: NATS not accessible at localhost:$NATS_PORT\"
        echo '  Make sure NATS container is running with --network host'
        exit 1
    fi

    echo ''

    # Function to wait for worker initialization by checking ETCD registration
    # Dynamo workers register with ETCD, they don't expose HTTP health endpoints
    wait_for_worker() {
        local worker_type=\$1
        local pid=\$2
        # Use WORKER_INIT_TIMEOUT_S (defaults to 1800s / 30 min)
        local max_wait=$WORKER_INIT_TIMEOUT_S
        local elapsed=0
        local poll_interval=5

        echo \"Waiting for \$worker_type worker (PID \$pid) to initialize...\"
        echo \"  Detection: ETCD worker registration\"
        echo \"  Timeout: \${max_wait}s\"

        while [ \$elapsed -lt \$max_wait ]; do
            # Check if process is still running
            if ! kill -0 \$pid 2>/dev/null; then
                echo \"ERROR: \$worker_type worker process died!\"
                return 1
            fi

            # Check ETCD for registered workers using v3 API
            # Query ALL keys to find where Dynamo registers (empty key "" with range_end "\0" = all keys)
            # Base64: "" -> AA==, "\0" -> AA==  (we use keys_only to reduce response size)
            local etcd_response=\$(curl -s --max-time 2 http://localhost:$ETCD_CLIENT_PORT/v3/kv/range \
                -X POST \
                -H \"Content-Type: application/json\" \
                -d '{\"key\":\"AA==\",\"range_end\":\"AA==\",\"keys_only\":true}' 2>&1)

            # Debug: Print ETCD response every 30s (truncated)
            if [ \$((elapsed % 30)) -eq 0 ] && [ \$elapsed -gt 0 ]; then
                echo \"  [DEBUG] ETCD keys found: \$(echo \"\$etcd_response\" | grep -o '\"key\":\"[^\"]*\"' | head -5)\"
                echo \"  [DEBUG] ETCD count: \$(echo \"\$etcd_response\" | grep -o '\"count\":\"[^\"]*\"')\"
            fi

            # Check if we got any keys back (count > 0 means workers registered)
            if echo \"\$etcd_response\" | grep -q '\"count\"' && \
               ! echo \"\$etcd_response\" | grep -q '\"count\":\"0\"'; then
                echo \"✓ \$worker_type worker is ready (registered with ETCD at \${elapsed}s)\"
                return 0
            fi

            sleep \$poll_interval
            elapsed=\$((elapsed + poll_interval))
            if [ \$((elapsed % 30)) -eq 0 ]; then
                echo \"  ... \${elapsed}s / \${max_wait}s (waiting for ETCD registration)\"
            fi
        done

        echo \"ERROR: \$worker_type worker failed to register with ETCD within \${max_wait}s\"
        echo \"  Image: $IMAGE\"
        echo \"  The model may require more time to load, or there may be a startup error.\"
        echo \"  Check worker logs for details.\"
        return 1
    }

    echo '========================================================='
    echo 'Step 1: Starting $NUM_WORKERS Unified Worker(s) (Host GPUs $WORKER_GPUS -> Container GPUs $CONTAINER_GPU_INDICES)...'
    echo '========================================================='

    # Start multiple workers, each using TP_SIZE GPUs
    WORKER_PIDS=()
    for i in \$(seq 0 \$(($NUM_WORKERS - 1))); do
        # Calculate GPU range for this worker (e.g., worker 0: 0,1; worker 1: 2,3; etc.)
        START_GPU=\$((i * $TP_SIZE))
        END_GPU=\$(((i + 1) * $TP_SIZE - 1))
        WORKER_GPU_LIST=\$(seq -s, \$START_GPU \$END_GPU)
        WORKER_PORT=\$((30000 + i))

        echo \"Starting Worker \$i: GPUs \$WORKER_GPU_LIST, Port \$WORKER_PORT\"
        # Build optional flags for the worker
        # --enable-metrics is always on so Prometheus/Grafana can scrape worker metrics
        EXTRA_WORKER_FLAGS=\"--enable-metrics\"
        if [ \"\$ENABLE_KV_AWARE_ROUTING\" = \"true\" ]; then
            EXTRA_WORKER_FLAGS=\"\$EXTRA_WORKER_FLAGS --page-size \$KV_BLOCK_SIZE\"
        fi
        if [ -n \"\$MAX_MODEL_LEN\" ]; then
            EXTRA_WORKER_FLAGS=\"\$EXTRA_WORKER_FLAGS --max-total-tokens \$MAX_MODEL_LEN\"
        fi
        if [ -n \"\$NUM_GPU_BLOCKS_OVERRIDE\" ]; then
            EXTRA_WORKER_FLAGS=\"\$EXTRA_WORKER_FLAGS --num-gpu-blocks-override \$NUM_GPU_BLOCKS_OVERRIDE\"
        fi
        if [ \"\$ENABLE_HIERARCHICAL_CACHE\" = \"true\" ]; then
            EXTRA_WORKER_FLAGS=\"\$EXTRA_WORKER_FLAGS --enable-hierarchical-cache\"
            EXTRA_WORKER_FLAGS=\"\$EXTRA_WORKER_FLAGS --hicache-ratio \$HICACHE_RATIO\"
            EXTRA_WORKER_FLAGS=\"\$EXTRA_WORKER_FLAGS --hicache-write-policy \$HICACHE_POLICY\"
        fi
        # DYN_SYSTEM_PORT: unique Prometheus metrics port per worker (required by --enable-metrics;
        # workers share the host network so each needs a distinct port).
        # DYN_NAMESPACE=workers: puts workers in the workers namespace so the Grafana dashboard
        # Request Flow panel Worker series (filtered on namespace=workers) is populated.
        # dynamo.frontend is started with --namespace workers below to match.
        CUDA_VISIBLE_DEVICES=\$WORKER_GPU_LIST \
        DYN_SYSTEM_PORT=\$((${WORKER_METRICS_PORT} + i)) \
        DYN_NAMESPACE=workers \
        python3 -m dynamo.sglang \
          --model-path $MODEL \
          --served-model-name $SERVED_MODEL_NAME \
          --host 0.0.0.0 \
          --port \$WORKER_PORT \
          --tp $TP_SIZE \
          --trust-remote-code \
          --mem-fraction-static $MEM_FRACTION_STATIC \
          \$EXTRA_WORKER_FLAGS &
        WORKER_PIDS+=(\$!)
        echo \"  Worker \$i PID: \${WORKER_PIDS[\$i]}\"
    done
    echo \"\"
    echo \"Total workers started: \${#WORKER_PIDS[@]}\"
    echo \"Worker PIDs: \${WORKER_PIDS[*]}\"
    echo \"\"

    # Wait for first worker to initialize (checks ETCD registration)
    # Once one worker is registered, the frontend can start discovering workers
    wait_for_worker \"Unified\" \${WORKER_PIDS[0]} || exit 1

    # Give additional workers time to initialize
    if [ \${#WORKER_PIDS[@]} -gt 1 ]; then
        echo \"Waiting additional 30s for remaining workers to initialize...\"
        sleep 30
    fi

    echo ''
    echo '========================================================='
    echo 'Step 2: Starting Dynamo Frontend (HTTP API on port $HTTP_PORT)...'
    echo '========================================================='
    # Build optional KV cache flag for the frontend
    # Worker metrics are always available (--enable-metrics is always on)
    echo \"Worker metrics: http://localhost:$WORKER_METRICS_PORT/metrics ... http://localhost:\$((${WORKER_METRICS_PORT} + $NUM_WORKERS - 1))/metrics\"
    KV_FRONTEND_FLAGS=\"\"
    if [ \"\$ENABLE_KV_AWARE_ROUTING\" = \"true\" ]; then
        echo \"KV Cache-Aware Routing enabled (block size: \$KV_BLOCK_SIZE tokens)\"
        # --router-mode kv: switches the frontend from default routing to KV-aware routing
        # --kv-cache-block-size: sets block size for KV overlap computation (must match worker --page-size)
        # --no-kv-events: router predicts cache state from its own routing decisions
        #   (workers in unified mode don't publish kv-events-config, so events are unavailable)
        KV_FRONTEND_FLAGS=\"--router-mode kv --kv-cache-block-size \$KV_BLOCK_SIZE --no-kv-events\"
    fi
    python3 -m dynamo.frontend \
      --http-port=$HTTP_PORT \
      --model-name $SERVED_MODEL_NAME \
      --model-path $MODEL \
      --namespace workers \
      \$KV_FRONTEND_FLAGS &
    FRONTEND_PID=\$!
    echo \"Frontend PID: \$FRONTEND_PID\"
    echo \"Waiting 15s for frontend to discover workers...\"
    sleep 15
    echo \"\"

    echo ''
    echo '========================================================='
    echo '✓ All components started successfully!'
    echo '========================================================='
    echo \"Infrastructure Services (External):\"
    echo \"  ETCD: localhost:$ETCD_CLIENT_PORT\"
    echo \"  NATS: localhost:$NATS_PORT\"
    echo \"\"
    echo \"Dynamo Components (This Container):\"
    echo \"  Unified Workers: \${#WORKER_PIDS[@]} workers (GPUs $WORKER_GPUS, TP=$TP_SIZE each)\"
    for i in \$(seq 0 \$((\${#WORKER_PIDS[@]} - 1))); do
        START_GPU=\$((i * $TP_SIZE))
        END_GPU=\$(((i + 1) * $TP_SIZE - 1))
        echo \"    Worker \$i: PID \${WORKER_PIDS[\$i]}, GPUs \$START_GPU-\$END_GPU, port \$((30000 + i))\"
    done
    echo \"  Frontend: PID \$FRONTEND_PID  (HTTP API on port $HTTP_PORT)\"
    echo ''
    echo 'Request Flow:'
    echo '  Client → Frontend API (port $HTTP_PORT)'
    echo '         ↓'
    echo '  Frontend discovers workers via ETCD'
    echo '         ↓'
    echo '  Frontend routes to Unified Worker'
    echo '         ↓'
    echo '  Response'
    echo '========================================================='

    # Monitor all processes
    while true; do
        # Check if any critical process died
        if ! kill -0 \$FRONTEND_PID 2>/dev/null; then
            echo \"ERROR: Frontend died!\"
            exit 1
        fi
        for i in \$(seq 0 \$((\${#WORKER_PIDS[@]} - 1))); do
            if ! kill -0 \${WORKER_PIDS[\$i]} 2>/dev/null; then
                echo \"ERROR: Worker \$i (PID \${WORKER_PIDS[\$i]}) died!\"
                exit 1
            fi
        done
        sleep 10
    done
  "

# Wait for container to start
echo ""
echo "Waiting for container to start..."
sleep 15

# Check if container started successfully
if docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
    echo ""
    echo "========================================================="
    echo "✓ Dynamo SGLang FULL STACK Started (UNIFIED MODE)!"
    echo "========================================================="
    echo ""
    echo "Architecture:"
    echo "  Client Request"
    echo "    ↓"
    echo "  Dynamo Frontend (port $HTTP_PORT)"
    echo "    ↓"
    echo "  Frontend discovers workers via ETCD"
    echo "    ↓"
    echo "  Frontend routes to one of $NUM_WORKERS Unified Workers"
    echo "    ↓              (localhost:$ETCD_CLIENT_PORT - worker discovery)"
    echo "  Unified Workers ($NUM_WORKERS x TP=$TP_SIZE = $NUM_GPUS GPUs total)"
    echo "    ↓"
    echo "  Response"
    echo ""
    echo "Infrastructure Services (Managed):"
    echo "  ETCD: etcd-dynamo container, localhost:$ETCD_CLIENT_PORT"
    echo "  NATS: nats-dynamo container, localhost:$NATS_PORT"
    echo ""
    echo "Dynamo Components (This Container):"
    echo "  Frontend: HTTP API on port $HTTP_PORT"
    echo "  Unified Workers: $NUM_WORKERS workers (TP=$TP_SIZE each, ports 30000-$((30000 + NUM_WORKERS - 1)))"
    echo ""
    echo "Prometheus Metrics Endpoints:"
    echo "  Frontend: http://localhost:$HTTP_PORT/metrics"
    echo "  Workers:  http://localhost:$WORKER_METRICS_PORT/metrics - $((WORKER_METRICS_PORT + NUM_WORKERS - 1))/metrics"
    echo ""
    echo "API Endpoint: http://localhost:$HTTP_PORT/v1/chat/completions"
    echo "Health Check: http://localhost:$HTTP_PORT/health"
    echo "Models Endpoint: http://localhost:$HTTP_PORT/v1/models"
    echo ""
    echo "Useful Commands:"
    echo "  Interactive shell:    docker exec -it $CONTAINER_NAME bash"
    echo "  View Dynamo logs:     docker logs -f $CONTAINER_NAME"
    echo "  View ETCD logs:       docker logs -f etcd-dynamo"
    echo "  View NATS logs:       docker logs -f nats-dynamo"
    echo "  GPU usage:            watch -n 2 nvidia-smi"
    echo "  Stop all:             bash stop_dynamo.sh"
    echo ""
    echo "Monitoring Dashboards:"
    echo "  Grafana:    http://localhost:3000 (no login required)"
    echo "  Prometheus: http://localhost:9090"
    echo ""
    echo "========================================================="
    echo "Test Request:"
    echo "========================================================="
    echo ""
    echo "# Basic test"
    echo "curl http://localhost:$HTTP_PORT/v1/chat/completions \\"
    echo "  -H 'Content-Type: application/json' \\"
    echo "  -d '{"
    echo "    \"model\": \"$SERVED_MODEL_NAME\","
    echo "    \"messages\": [{\"role\": \"user\", \"content\": \"Hello!\"}],"
    echo "    \"max_tokens\": 50"
    echo "  }'"
    echo ""
    echo "# Streaming test"
    echo "curl http://localhost:$HTTP_PORT/v1/chat/completions \\"
    echo "  -H 'Content-Type: application/json' \\"
    echo "  -d '{"
    echo "    \"model\": \"$SERVED_MODEL_NAME\","
    echo "    \"messages\": [{\"role\": \"user\", \"content\": \"Hello!\"}],"
    echo "    \"max_tokens\": 50,"
    echo "    \"stream\": true"
    echo "  }'"
    echo ""
    echo "========================================================="
    echo "NAT Integration Test:"
    echo "========================================================="
    echo ""
    echo "cd /path/to/NeMo-Agent-Toolkit"
    echo "source .venv/bin/activate"
    echo ""
    echo "nat run \\"
    echo "  --config_file examples/dynamo_integration/react_benchmark_agent/configs/config_dynamo_e2e_test.yml \\"
    echo "  --input 'Hello'"
    echo ""
    echo "========================================================="
    echo ""
    echo "Waiting for SGLang to initialize (this will likely take 5-10 minutes for a 70B model)..."
    echo "Monitoring logs (Ctrl+C to exit, container continues)..."
    echo ""

    # Wait for server to be ready (check /v1/models which only works when workers are discovered)
    echo "Checking for API availability (timeout=${WORKER_INIT_TIMEOUT_S}s)..."
    max_attempts=$WORKER_INIT_TIMEOUT_S
    attempt=0

    while [ $attempt -lt $max_attempts ]; do
        # Check /v1/models - only returns data when workers are registered
        models_response=$(curl -s http://localhost:$HTTP_PORT/v1/models 2>/dev/null)
        if echo "$models_response" | grep -q '"id"'; then
            echo "✓ SGLang API is ready! (models discovered)"
            break
        fi
        attempt=$((attempt + 1))
        if [ $((attempt % 15)) -eq 0 ]; then
            echo "  ... still waiting ($attempt/$max_attempts)"
        fi
        sleep 1
    done

    if [ $attempt -ge $max_attempts ]; then
        echo ""
        echo "⚠ Timeout waiting for API. Check logs with: docker logs $CONTAINER_NAME"
        echo ""
    else
        echo ""
        echo "Quick test:"
        echo ""
        curl -s http://localhost:$HTTP_PORT/v1/chat/completions \
          -H "Content-Type: application/json" \
          -d '{
            "model": "'$SERVED_MODEL_NAME'",
            "messages": [{"role": "user", "content": "Say hello"}],
            "max_tokens": 20
          }' | jq '.choices[0].message.content, .usage'

        echo ""
        echo "========================================================="
        echo "Container is running. View logs with:"
        echo "  docker logs -f $CONTAINER_NAME"
        echo "========================================================="
    fi
else
    echo ""
    echo "========================================================="
    echo "✗ Container failed to start!"
    echo "========================================================="
    echo ""
    echo "Check logs with: docker logs $CONTAINER_NAME"
    exit 1
fi