@aws/ml-container-creator 1.0.2 → 1.0.4

package/templates/do/benchmark

@@ -13,6 +13,7 @@ set -o pipefail
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 source "${SCRIPT_DIR}/config"
 source "${SCRIPT_DIR}/lib/profile.sh"
+source "${SCRIPT_DIR}/lib/resolve-instance.sh"
 
 # ── Parse flags ───────────────────────────────────────────────────────────────
 CLEAN_AFTER=false
@@ -69,8 +70,8 @@ done
 # Query the tracked benchmark job, display status, and if completed:
 # download results, display metrics, and write to Athena (if not already done).
 if [ "${ARG_STATUS}" = true ]; then
-    # Resolve instance type: BENCHMARK_INSTANCE_TYPE (persisted by main flow) > INSTANCE_TYPE from config
-    _STATUS_INSTANCE_TYPE="${BENCHMARK_INSTANCE_TYPE:-${INSTANCE_TYPE:-}}"
+    # Resolve instance type: BENCHMARK_INSTANCE_TYPE > DEPLOYED_INSTANCE_TYPE > INSTANCE_TYPE
+    _STATUS_INSTANCE_TYPE="${BENCHMARK_INSTANCE_TYPE:-${DEPLOYED_INSTANCE_TYPE:-${INSTANCE_TYPE:-}}}"
 
     JOB_NAME="${BENCHMARK_JOB_NAME:-}"
     if [ -z "${JOB_NAME}" ]; then
@@ -575,7 +576,7 @@ print(f'Combined {n_metrics} concurrency level results')
 
         echo ""
         echo "📋 Multi-level Summary:"
-        echo "   Levels tested: ${_NORMALIZED_LEVELS}"
+        echo "   Levels tested: ${#_LEVELS[@]}"
         echo "   Failures: ${_LEVEL_FAILURES} / ${#_LEVELS[@]}"
         echo "   Results: ${_ALL_RESULTS_DIR}/"
 
@@ -845,6 +846,62 @@ except:
             [ "${RESOLVED_INSTANCE_TYPE}" = "None" ] && RESOLVED_INSTANCE_TYPE=""
         fi
     fi
+
+    # If still empty (heterogeneous pools), query the inference component directly
+    if [ -z "${RESOLVED_INSTANCE_TYPE}" ] && [ -n "${IC_NAME}" ]; then
+        RESOLVED_INSTANCE_TYPE=$(aws sagemaker describe-inference-component \
+            --inference-component-name "${IC_NAME}" \
+            --region "${AWS_REGION}" \
+            --output json 2>/dev/null | python3 -c "
+import sys, json
+try:
+    ic = json.load(sys.stdin)
+    # RuntimeConfig.CurrentCopyCount tells us it's running, but instance type
+    # is in the ResourceRequirements or the endpoint's routing
+    runtime = ic.get('RuntimeConfig', {})
+    # Try DesiredCopyCount path for instance pool info
+    spec = ic.get('Specification', {}).get('ComputeResourceRequirements', {})
+    # For pool-based endpoints, check the IC's runtime host (if available)
+    # Fallback: query endpoint InstancePools
+    print('')
+except:
+    print('')
+" 2>/dev/null) || RESOLVED_INSTANCE_TYPE=""
+        [ "${RESOLVED_INSTANCE_TYPE}" = "None" ] && RESOLVED_INSTANCE_TYPE=""
+    fi
+
+    # If still empty (heterogeneous pools), query endpoint config InstancePools
+    if [ -z "${RESOLVED_INSTANCE_TYPE}" ]; then
+        _EC_NAME="${_EC_NAME:-$(echo "${_EP_JSON}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('EndpointConfigName',''))" 2>/dev/null)}"
+        if [ -n "${_EC_NAME}" ]; then
+            RESOLVED_INSTANCE_TYPE=$(aws sagemaker describe-endpoint-config \
+                --endpoint-config-name "${_EC_NAME}" \
+                --region "${AWS_REGION}" \
+                --output json 2>/dev/null | python3 -c "
+import sys, json
+try:
+    ec = json.load(sys.stdin)
+    variants = ec.get('ProductionVariants', [])
+    for v in variants:
+        # Check InstancePoolConfig for pool-based endpoints
+        pools = v.get('InstancePoolConfig', {}).get('InstancePools', [])
+        if pools:
+            # Use the first pool's instance type
+            print(pools[0].get('InstanceType', ''))
+            break
+        # Check RoutingConfig ManagedInstanceScaling
+        it = v.get('InstanceType', '')
+        if it:
+            print(it)
+            break
+    else:
+        print('')
+except:
+    print('')
+" 2>/dev/null) || RESOLVED_INSTANCE_TYPE=""
+            [ "${RESOLVED_INSTANCE_TYPE}" = "None" ] && RESOLVED_INSTANCE_TYPE=""
+        fi
+    fi
 fi
 
 # Final fallback: use INSTANCE_TYPE from do/config

package/templates/do/deploy.d/managed-inference.ejs

@@ -718,6 +718,38 @@ if [ -z "${SKIP_TO}" ] || [ "${SKIP_TO}" = "create_ic" ] || [ "${SKIP_TO}" = "wa
         # Run capacity guardrail before deploying ICs
         _check_gpu_capacity
 
+        # _check_gpu_count_mismatch
+        #   Non-fatal warning: if any IC has IC_GPU_COUNT=1 but the endpoint has
+        #   multiple GPUs available, warn the user that they may be underutilizing.
+        #   The user may intentionally use fewer GPUs (e.g., sharing endpoint across
+        #   multiple ICs), so this is advisory only — no exit, no abort.
+        _check_gpu_count_mismatch() {
+            # Skip if DEPLOYED_GPU_COUNT is not set or is ≤1
+            if [ "${DEPLOYED_GPU_COUNT:-0}" -le 1 ]; then
+                return 0
+            fi
+
+            for conf in "${SCRIPT_DIR}"/ic/*.conf; do
+                [ -f "${conf}" ] || continue
+                local ic_gpus
+                ic_gpus=$(grep "^export IC_GPU_COUNT=" "${conf}" 2>/dev/null | sed 's/^export IC_GPU_COUNT=//' | tr -d '"' || echo "1")
+                if [ -z "${ic_gpus}" ]; then
+                    ic_gpus=1
+                fi
+                if [ "${ic_gpus}" = "1" ] && [ "${DEPLOYED_GPU_COUNT:-0}" -gt 1 ]; then
+                    local conf_name
+                    conf_name=$(basename "${conf}")
+                    echo ""
+                    echo "⚠️  Warning: IC_GPU_COUNT=1 but endpoint has ${DEPLOYED_GPU_COUNT} GPUs available"
+                    echo "   Consider setting IC_GPU_COUNT=${DEPLOYED_GPU_COUNT} in do/ic/${conf_name} for tensor parallelism"
+                    echo ""
+                    break
+                fi
+            done
+        }
+
+        _check_gpu_count_mismatch
+
         # _delete_and_wait_ic <ic_name>
         #   Deletes an inference component and waits for deletion to complete.
         #   Polls until the IC is no longer found (avoids name conflicts on recreate).
@@ -910,6 +942,57 @@ elif [ "${SKIP_TO}" = "wait_ic" ]; then
     echo "✅ Inference component is InService: ${IC_DEPLOYED_NAME}"
 fi
 
+# ── Resolve and persist actual instance type (for heterogeneous pools) ─────────
+# On pool-based endpoints, INSTANCE_TYPE in do/config may be empty or reflect
+# only the generation-time default. Query the live endpoint to determine what
+# SageMaker actually provisioned, and persist as DEPLOYED_INSTANCE_TYPE.
+# Downstream scripts (do/benchmark, do/register) read this for Athena/MPG metadata.
+if [ -z "${INSTANCE_TYPE:-}" ] || [ -n "${INSTANCE_POOLS:-}" ]; then
+    _RESOLVED_INSTANCE=""
+    _EP_DESCRIBE=$(aws sagemaker describe-endpoint \
+        --endpoint-name "${ENDPOINT_NAME}" \
+        --region "${AWS_REGION}" \
+        --output json 2>/dev/null) || _EP_DESCRIBE=""
+
+    if [ -n "${_EP_DESCRIBE}" ]; then
+        _RESOLVED_INSTANCE=$(echo "${_EP_DESCRIBE}" | python3 -c "
+import sys, json
+try:
+    ep = json.load(sys.stdin)
+    variant = ep.get('ProductionVariants', [{}])[0]
+    print(variant.get('CurrentInstanceType') or variant.get('InstanceType') or '')
+except:
+    print('')
+" 2>/dev/null) || _RESOLVED_INSTANCE=""
+
+        # Fallback: query endpoint config if variant doesn't have it
+        if [ -z "${_RESOLVED_INSTANCE}" ]; then
+            _EC_NAME=$(echo "${_EP_DESCRIBE}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('EndpointConfigName',''))" 2>/dev/null) || _EC_NAME=""
+            if [ -n "${_EC_NAME}" ]; then
+                _RESOLVED_INSTANCE=$(aws sagemaker describe-endpoint-config \
+                    --endpoint-config-name "${_EC_NAME}" \
+                    --region "${AWS_REGION}" \
+                    --query 'ProductionVariants[0].InstanceType' \
+                    --output text 2>/dev/null) || _RESOLVED_INSTANCE=""
+                [ "${_RESOLVED_INSTANCE}" = "None" ] && _RESOLVED_INSTANCE=""
+            fi
+        fi
+    fi
+
+    if [ -n "${_RESOLVED_INSTANCE}" ]; then
+        # Persist to do/config for downstream scripts
+        _config_file="${SCRIPT_DIR}/config"
+        if grep -q "^export DEPLOYED_INSTANCE_TYPE=" "${_config_file}" 2>/dev/null; then
+            sed -i.bak "s|^export DEPLOYED_INSTANCE_TYPE=.*|export DEPLOYED_INSTANCE_TYPE=\"${_RESOLVED_INSTANCE}\"|" "${_config_file}"
+            rm -f "${_config_file}.bak"
+        else
+            echo "export DEPLOYED_INSTANCE_TYPE=\"${_RESOLVED_INSTANCE}\"" >> "${_config_file}"
+        fi
+        INSTANCE_TYPE="${_RESOLVED_INSTANCE}"
+        echo "   Resolved instance type: ${_RESOLVED_INSTANCE} (persisted to do/config)"
+    fi
+fi
+
 echo "✅ Deployment complete!"
 echo ""
 echo "📋 Deployment Details:"

package/templates/do/evaluate

@@ -0,0 +1,272 @@
+#!/bin/bash
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# do/evaluate — Model Quality Evaluation
+# Evaluates a deployed model/adapter via inference requests against the endpoint.
+# Computes technique-specific quality metrics (perplexity, reward accuracy, etc.).
+#
+# Project: <%= projectName %>
+
+set -e
+set -u
+set -o pipefail
+
+# ── Source project configuration ──────────────────────────────────────────────
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+source "${SCRIPT_DIR}/config"
+
+# ── CLI Variables ─────────────────────────────────────────────────────────────
+ARG_ADAPTER=""
+ARG_EVAL_DATASET=""
+ARG_SAMPLES=""
+ARG_METRICS=""
+ARG_DRY_RUN=false
+ARG_HELP=false
+
+# ── Parse arguments ───────────────────────────────────────────────────────────
+while [ $# -gt 0 ]; do
+    case "$1" in
+        --adapter)
+            if [ -z "${2:-}" ]; then echo "❌ --adapter requires a name"; exit 1; fi
+            ARG_ADAPTER="$2"; shift 2 ;;
+        --eval-dataset)
+            if [ -z "${2:-}" ]; then echo "❌ --eval-dataset requires a value"; exit 1; fi
+            ARG_EVAL_DATASET="$2"; shift 2 ;;
+        --samples)
+            if [ -z "${2:-}" ]; then echo "❌ --samples requires a number"; exit 1; fi
+            ARG_SAMPLES="$2"; shift 2 ;;
+        --metrics)
+            if [ -z "${2:-}" ]; then echo "❌ --metrics requires a value"; exit 1; fi
+            ARG_METRICS="$2"; shift 2 ;;
+        --dry-run) ARG_DRY_RUN=true; shift ;;
+        --help|-h) ARG_HELP=true; shift ;;
+        *)
+            echo "❌ Unknown option: $1"
+            echo "   Run ./do/evaluate --help for usage."
+            exit 1
+            ;;
+    esac
+done
+
+# ── Help ──────────────────────────────────────────────────────────────────────
+if [ "${ARG_HELP}" = true ]; then
+    echo "Usage: ./do/evaluate [OPTIONS]"
+    echo "       ./do/evaluate --adapter <name> --eval-dataset <source>"
+    echo ""
+    echo "Evaluate model/adapter quality via inference requests against the deployed endpoint."
+    echo "Computes technique-specific metrics: perplexity (SFT), reward accuracy (DPO)."
+    echo ""
+    echo "Options:"
+    echo "  --adapter <name>       Evaluate specific adapter (reads IC name from conf)"
+    echo "  --eval-dataset <src>   Evaluation dataset: s3://..., hf://..., or registry name"
+    echo "  --samples <n>          Limit evaluation to N samples (default: all)"
+    echo "  --metrics <list>       Comma-separated metrics to compute (default: all for technique)"
+    echo "  --dry-run              Show what would be evaluated without making requests"
+    echo "  --help, -h             Show this help message"
+    echo ""
+    echo "Examples:"
+    echo "  ./do/evaluate                                     # Evaluate default IC"
+    echo "  ./do/evaluate --adapter sft-custom-a3f2           # Evaluate specific adapter"
+    echo "  ./do/evaluate --eval-dataset \"hf://tatsu-lab/alpaca --take 50\""
+    echo "  ./do/evaluate --adapter dpo-custom --eval-dataset s3://bucket/dpo-eval.jsonl"
+    echo ""
+    echo "Results saved to: .mlcc/eval-results/<adapter-or-ic>.json"
+    echo "Results are automatically included in do/register metadata."
+    exit 0
+fi
+
+# ── Resolve endpoint ──────────────────────────────────────────────────────────
+ENDPOINT_NAME="${ENDPOINT_NAME:-}"
+if [ -z "${ENDPOINT_NAME}" ]; then
+    echo "❌ No endpoint configured."
+    echo "   Deploy first: ./do/deploy"
+    echo "   Then run: ./do/evaluate"
+    exit 1
+fi
+
+# ── Resolve IC name ───────────────────────────────────────────────────────────
+IC_NAME=""
+ADAPTER_TECHNIQUE=""
+EVAL_TARGET_NAME=""
+
+if [ -n "${ARG_ADAPTER}" ]; then
+    # Adapter specified — look up IC from adapter conf
+    ADAPTER_CONF="${SCRIPT_DIR}/adapters/${ARG_ADAPTER}.conf"
+    if [ ! -f "${ADAPTER_CONF}" ]; then
+        echo "❌ Adapter config not found: do/adapters/${ARG_ADAPTER}.conf"
+        echo "   Available adapters:"
+        if [ -d "${SCRIPT_DIR}/adapters" ]; then
+            for conf in "${SCRIPT_DIR}"/adapters/*.conf; do
+                [ -f "${conf}" ] || continue
+                echo "     • $(basename "${conf}" .conf)"
+            done
+        fi
+        exit 1
+    fi
+    source "${ADAPTER_CONF}"
+    IC_NAME="${ADAPTER_IC_NAME:-}"
+    ADAPTER_TECHNIQUE="${ADAPTER_TECHNIQUE:-${ADAPTER_TUNE_TECHNIQUE:-}}"
+    EVAL_TARGET_NAME="${ARG_ADAPTER}"
+else
+    # No adapter — use default IC
+    if [ -d "${SCRIPT_DIR}/ic" ]; then
+        for conf in "${SCRIPT_DIR}"/ic/*.conf; do
+            [ -f "${conf}" ] || continue
+            IC_DEPLOYED_NAME=""
+            source "${conf}"
+            if [ -n "${IC_DEPLOYED_NAME}" ]; then
+                IC_NAME="${IC_DEPLOYED_NAME}"
+                break
+            fi
+        done
+    fi
+    # Fallback to legacy config
+    IC_NAME="${IC_NAME:-${INFERENCE_COMPONENT_NAME:-}}"
+    EVAL_TARGET_NAME="${IC_NAME:-default}"
+fi
+
+if [ -z "${IC_NAME}" ]; then
+    echo "❌ No inference component found to evaluate."
+    echo "   Deploy first: ./do/deploy"
+    echo "   Or specify an adapter: ./do/evaluate --adapter <name>"
+    exit 1
+fi
+
+# ── Resolve technique (for metric selection) ──────────────────────────────────
+# Priority: adapter conf ADAPTER_TECHNIQUE > TRAIN_TECHNIQUE from config > default
+TECHNIQUE="${ADAPTER_TECHNIQUE:-${TRAIN_TECHNIQUE:-}}"
+
+# ── Resolve eval dataset ──────────────────────────────────────────────────────
+EVAL_DATASET="${ARG_EVAL_DATASET:-${EVAL_DATASET:-}}"
+
+# ── Display configuration ─────────────────────────────────────────────────────
+echo "🔬 Model Evaluation"
+echo ""
+echo "   Endpoint: ${ENDPOINT_NAME}"
+echo "   IC:       ${IC_NAME}"
+if [ -n "${ARG_ADAPTER}" ]; then
+    echo "   Adapter:  ${ARG_ADAPTER}"
+fi
+if [ -n "${TECHNIQUE}" ]; then
+    echo "   Technique: ${TECHNIQUE}"
+fi
+if [ -n "${EVAL_DATASET}" ]; then
+    echo "   Dataset:  ${EVAL_DATASET}"
+fi
+if [ -n "${ARG_SAMPLES}" ]; then
+    echo "   Samples:  ${ARG_SAMPLES}"
+fi
+echo ""
+
+# ── Dry run ───────────────────────────────────────────────────────────────────
+if [ "${ARG_DRY_RUN}" = true ]; then
+    echo "🔍 Dry run — would evaluate with above configuration."
+    echo "   No inference requests will be made."
+    exit 0
+fi
+
+# ── Check endpoint is InService ───────────────────────────────────────────────
+echo "🔍 Checking endpoint status..."
+EP_STATUS=$(aws sagemaker describe-endpoint \
+    --endpoint-name "${ENDPOINT_NAME}" \
+    --region "${AWS_REGION}" \
+    --query 'EndpointStatus' \
+    --output text 2>/dev/null) || EP_STATUS=""
+
+if [ "${EP_STATUS}" != "InService" ]; then
+    echo "❌ Endpoint is not InService (status: ${EP_STATUS:-unknown})"
+    echo ""
+    echo "   Deploy first: ./do/deploy"
+    echo "   Then run: ./do/evaluate"
+    exit 1
+fi
+echo "   ✅ Endpoint is InService"
+echo ""
+
+# ── Run evaluation via Python helper ─────────────────────────────────────────
+echo "🧪 Running evaluation..."
+echo ""
+
+EVAL_ARGS=(
+    --endpoint-name "${ENDPOINT_NAME}"
+    --ic-name "${IC_NAME}"
+    --region "${AWS_REGION}"
+)
+
+if [ -n "${TECHNIQUE}" ]; then
+    EVAL_ARGS+=(--technique "${TECHNIQUE}")
+fi
+if [ -n "${EVAL_DATASET}" ]; then
+    EVAL_ARGS+=(--eval-dataset "${EVAL_DATASET}")
+fi
+if [ -n "${ARG_SAMPLES}" ]; then
+    EVAL_ARGS+=(--samples "${ARG_SAMPLES}")
+fi
+if [ -n "${ARG_METRICS}" ]; then
+    EVAL_ARGS+=(--metrics "${ARG_METRICS}")
+fi
+
+EVAL_OUTPUT=$(python3 "${SCRIPT_DIR}/.eval_helper.py" evaluate "${EVAL_ARGS[@]}" 2>/dev/null | grep -E '^\{' | tail -1) || EVAL_OUTPUT=""
+
+if [ -z "${EVAL_OUTPUT}" ]; then
+    echo "❌ Evaluation failed (no output from helper)"
+    echo "   Check: endpoint accessibility, eval dataset format, Python deps (requests)"
+    exit 1
+fi
+
+# Check for error
+HAS_ERROR=$(echo "${EVAL_OUTPUT}" | python3 -c "import sys,json; d=json.load(sys.stdin); print('yes' if d.get('error') else 'no')" 2>/dev/null) || HAS_ERROR="yes"
+
+if [ "${HAS_ERROR}" = "yes" ]; then
+    ERROR_MSG=$(echo "${EVAL_OUTPUT}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('message','Unknown error'))" 2>/dev/null) || ERROR_MSG="Unknown error"
+    echo "❌ Evaluation failed: ${ERROR_MSG}"
+    exit 1
+fi
+
+# ── Save results ──────────────────────────────────────────────────────────────
+RESULTS_DIR="${SCRIPT_DIR}/../.mlcc/eval-results"
+mkdir -p "${RESULTS_DIR}"
+RESULTS_FILE="${RESULTS_DIR}/${EVAL_TARGET_NAME}.json"
+echo "${EVAL_OUTPUT}" | python3 -c "import sys,json; json.dump(json.load(sys.stdin), open('${RESULTS_FILE}','w'), indent=2)"
+
+echo "📁 Results saved to: .mlcc/eval-results/${EVAL_TARGET_NAME}.json"
+echo ""
+
+# ── Display summary ───────────────────────────────────────────────────────────
+echo "╔══════════════════════════════════════════════════════════════════╗"
+echo "║              Evaluation Results                                 ║"
+echo "╠══════════════════════════════════════════════════════════════════╣"
+echo "${EVAL_OUTPUT}" | python3 -c "
+import sys, json
+data = json.load(sys.stdin)
+print(f\"║  Target: {data.get('adapter_name', data.get('ic_name', 'unknown'))}\")
+print(f\"║  Technique: {data.get('technique', 'unknown')}\")
+print(f\"║  Samples: {data.get('samples_evaluated', 0)}\")
+print(f\"║\")
+metrics = data.get('metrics', {})
+for name, value in metrics.items():
+    if isinstance(value, float):
+        print(f'║  {name}: {value:.4f}')
+    else:
+        print(f'║  {name}: {value}')
+" 2>/dev/null
+echo "╚══════════════════════════════════════════════════════════════════╝"
+echo ""
+
+# ── Optionally write to Athena ────────────────────────────────────────────────
+if [ -n "${CI_BENCHMARK_RESULTS_BUCKET:-}" ]; then
+    echo "☁️  Persisting evaluation results to Athena..."
+    python3 "${SCRIPT_DIR}/.eval_helper.py" eval-write \
+        --results-file "${RESULTS_FILE}" \
+        --bucket "${CI_BENCHMARK_RESULTS_BUCKET}" \
+        --region "${AWS_REGION}" 2>/dev/null | grep -E '^\{' | tail -1 > /dev/null || {
+        echo "   ⚠️  Failed to persist to Athena (non-fatal)"
+    }
+fi
+
+echo "✅ Evaluation complete."
+echo ""
+echo "   Next steps:"
+echo "   • Run ./do/register to include eval metrics in model package metadata"
+echo "   • Compare adapters: ./do/evaluate --adapter <other-adapter>"

package/templates/do/lib/resolve-instance.sh

@@ -0,0 +1,155 @@
+#!/usr/bin/env bash
+# Lazy instance type resolution for heterogeneous pool endpoints.
+# Source this file after do/config + lib/profile.sh in any script that needs INSTANCE_TYPE.
+#
+# When INSTANCE_TYPE is empty (pool endpoints) and DEPLOYED_INSTANCE_TYPE hasn't been
+# persisted yet (no do/deploy run), queries the live endpoint once and persists the result.
+# Subsequent calls read from do/config without any AWS API calls.
+#
+# After sourcing, INSTANCE_TYPE is guaranteed to be set (or empty if resolution failed).
+# DEPLOYED_GPU_COUNT is also resolved from a static lookup table (instances.json catalog)
+# and persisted to do/config alongside the instance type.
+#
+# Usage:
+#   source "${SCRIPT_DIR}/config"
+#   source "${SCRIPT_DIR}/lib/profile.sh"
+#   source "${SCRIPT_DIR}/lib/resolve-instance.sh"
+#   # INSTANCE_TYPE and DEPLOYED_GPU_COUNT are now resolved
+
+# Resolve SCRIPT_DIR if not already set (defensive — normally inherited from caller)
+if [ -z "${SCRIPT_DIR:-}" ]; then
+    SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+fi
+
+# Skip if INSTANCE_TYPE is already set (single-instance endpoints)
+if [ -n "${INSTANCE_TYPE:-}" ]; then
+    return 0 2>/dev/null || true
+fi
+
+# Check if DEPLOYED_INSTANCE_TYPE was previously persisted
+if [ -n "${DEPLOYED_INSTANCE_TYPE:-}" ]; then
+    INSTANCE_TYPE="${DEPLOYED_INSTANCE_TYPE}"
+    export INSTANCE_TYPE
+    # Ensure DEPLOYED_GPU_COUNT is also exported (may already be in do/config)
+    if [ -n "${DEPLOYED_GPU_COUNT:-}" ]; then
+        export DEPLOYED_GPU_COUNT
+    fi
+    return 0 2>/dev/null || true
+fi
+
+# Check if BENCHMARK_INSTANCE_TYPE was previously persisted (by do/benchmark)
+if [ -n "${BENCHMARK_INSTANCE_TYPE:-}" ]; then
+    INSTANCE_TYPE="${BENCHMARK_INSTANCE_TYPE}"
+    export INSTANCE_TYPE
+    return 0 2>/dev/null || true
+fi
+
+# ── Live resolution from endpoint (one-time, persisted) ──────────────────────
+# Only attempt if ENDPOINT_NAME is configured and AWS credentials are available.
+if [ -z "${ENDPOINT_NAME:-}" ]; then
+    return 0 2>/dev/null || true
+fi
+
+_RESOLVED_INSTANCE=""
+_EP_DESCRIBE=$(aws sagemaker describe-endpoint \
+    --endpoint-name "${ENDPOINT_NAME}" \
+    --region "${AWS_REGION:-us-east-1}" \
+    --output json 2>/dev/null) || _EP_DESCRIBE=""
+
+if [ -n "${_EP_DESCRIBE}" ]; then
+    _RESOLVED_INSTANCE=$(echo "${_EP_DESCRIBE}" | python3 -c "
+import sys, json
+try:
+    ep = json.load(sys.stdin)
+    variant = ep.get('ProductionVariants', [{}])[0]
+    print(variant.get('CurrentInstanceType') or variant.get('InstanceType') or '')
+except:
+    print('')
+" 2>/dev/null) || _RESOLVED_INSTANCE=""
+
+    # Fallback: query endpoint config for InstanceType or first pool entry
+    if [ -z "${_RESOLVED_INSTANCE}" ]; then
+        _EC_NAME=$(echo "${_EP_DESCRIBE}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('EndpointConfigName',''))" 2>/dev/null) || _EC_NAME=""
+        if [ -n "${_EC_NAME}" ]; then
+            _RESOLVED_INSTANCE=$(aws sagemaker describe-endpoint-config \
+                --endpoint-config-name "${_EC_NAME}" \
+                --region "${AWS_REGION:-us-east-1}" \
+                --query 'ProductionVariants[0].InstanceType' \
+                --output text 2>/dev/null) || _RESOLVED_INSTANCE=""
+            [ "${_RESOLVED_INSTANCE}" = "None" ] && _RESOLVED_INSTANCE=""
+
+            # Final fallback: first entry in InstancePools
+            if [ -z "${_RESOLVED_INSTANCE}" ]; then
+                _RESOLVED_INSTANCE=$(aws sagemaker describe-endpoint-config \
+                    --endpoint-config-name "${_EC_NAME}" \
+                    --region "${AWS_REGION:-us-east-1}" \
+                    --output json 2>/dev/null | python3 -c "
+import sys, json
+try:
+    ec = json.load(sys.stdin)
+    pools = ec.get('ProductionVariants', [{}])[0].get('InstancePools', [])
+    if pools:
+        best = min(pools, key=lambda p: p.get('Priority', 999))
+        print(best.get('InstanceType', ''))
+    else:
+        print('')
+except:
+    print('')
+" 2>/dev/null) || _RESOLVED_INSTANCE=""
+            fi
+        fi
+    fi
+fi
+
+# Persist to do/config (one-time write — subsequent sources read it directly)
+if [ -n "${_RESOLVED_INSTANCE}" ]; then
+    _config_file="${SCRIPT_DIR}/config"
+    if grep -q "^export DEPLOYED_INSTANCE_TYPE=" "${_config_file}" 2>/dev/null; then
+        sed -i.bak "s|^export DEPLOYED_INSTANCE_TYPE=.*|export DEPLOYED_INSTANCE_TYPE=\"${_RESOLVED_INSTANCE}\"|" "${_config_file}"
+        rm -f "${_config_file}.bak"
+    else
+        echo "export DEPLOYED_INSTANCE_TYPE=\"${_RESOLVED_INSTANCE}\"" >> "${_config_file}"
+    fi
+    INSTANCE_TYPE="${_RESOLVED_INSTANCE}"
+    DEPLOYED_INSTANCE_TYPE="${_RESOLVED_INSTANCE}"
+    export INSTANCE_TYPE DEPLOYED_INSTANCE_TYPE
+
+    # ── Resolve GPU count from instance type ─────────────────────────────────
+    # Static lookup table derived from servers/lib/catalogs/instances.json.
+    # Maps known SageMaker instance types to their GPU count.
+    _resolve_gpu_count() {
+        case "$1" in
+            ml.g5.xlarge|ml.g5.2xlarge|ml.g5.4xlarge|ml.g5.8xlarge|ml.g5.16xlarge) echo 1 ;;
+            ml.g5.12xlarge|ml.g5.24xlarge) echo 4 ;;
+            ml.g5.48xlarge) echo 8 ;;
+            ml.g4dn.xlarge|ml.g4dn.2xlarge|ml.g4dn.4xlarge|ml.g4dn.8xlarge|ml.g4dn.16xlarge) echo 1 ;;
+            ml.g4dn.12xlarge) echo 4 ;;
+            ml.g6.xlarge|ml.g6.2xlarge|ml.g6.4xlarge|ml.g6.8xlarge|ml.g6.16xlarge) echo 1 ;;
+            ml.g6.12xlarge|ml.g6.24xlarge) echo 4 ;;
+            ml.g6.48xlarge) echo 8 ;;
+            ml.g6e.xlarge|ml.g6e.2xlarge|ml.g6e.4xlarge|ml.g6e.8xlarge|ml.g6e.16xlarge) echo 1 ;;
+            ml.g6e.12xlarge|ml.g6e.24xlarge) echo 4 ;;
+            ml.g6e.48xlarge) echo 8 ;;
+            ml.p4d.24xlarge|ml.p4de.24xlarge) echo 8 ;;
+            ml.p5.48xlarge|ml.p5e.48xlarge) echo 8 ;;
+            *) echo "" ;;
+        esac
+    }
+
+    _GPU_COUNT=$(_resolve_gpu_count "${_RESOLVED_INSTANCE}")
+    if [ -n "${_GPU_COUNT}" ]; then
+        if grep -q "^export DEPLOYED_GPU_COUNT=" "${_config_file}" 2>/dev/null; then
+            sed -i.bak "s|^export DEPLOYED_GPU_COUNT=.*|export DEPLOYED_GPU_COUNT=\"${_GPU_COUNT}\"|" "${_config_file}"
+            rm -f "${_config_file}.bak"
+        else
+            echo "export DEPLOYED_GPU_COUNT=\"${_GPU_COUNT}\"" >> "${_config_file}"
+        fi
+        DEPLOYED_GPU_COUNT="${_GPU_COUNT}"
+        export DEPLOYED_GPU_COUNT
+    fi
+    unset _GPU_COUNT
+    unset -f _resolve_gpu_count
+fi
+
+# Clean up internal vars
+unset _RESOLVED_INSTANCE _EP_DESCRIBE _EC_NAME

package/templates/do/register

@@ -10,11 +10,16 @@ set -o pipefail
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 source "${SCRIPT_DIR}/config"
 source "${SCRIPT_DIR}/lib/profile.sh"
+source "${SCRIPT_DIR}/lib/resolve-instance.sh"
 
 # ── Profile-resolved variables (env var > profile > default) ──────────────────
 ROLE_ARN="${ROLE_ARN:-${_PROFILE_roleArn:-}}"
 ECR_REPOSITORY_NAME="${ECR_REPOSITORY_NAME:-${_PROFILE_ecrRepositoryName:-ml-container-creator}}"
 
+# Resolve INSTANCE_TYPE for heterogeneous pool endpoints (may be empty from config,
+# filled by do/deploy after InService resolution)
+INSTANCE_TYPE="${INSTANCE_TYPE:-${DEPLOYED_INSTANCE_TYPE:-${BENCHMARK_INSTANCE_TYPE:-}}}"
+
 # ============================================================
 # Register deployment to the deployment registry
 # ============================================================

package/templates/do/test

@@ -10,6 +10,7 @@ set -o pipefail
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 source "${SCRIPT_DIR}/config"
 source "${SCRIPT_DIR}/lib/profile.sh"
+source "${SCRIPT_DIR}/lib/resolve-instance.sh"
 
 <% if (deploymentTarget === 'realtime-inference') { %>
 # ============================================================