@aws/ml-container-creator 0.9.1 → 0.10.3

package/templates/do/optimize

@@ -106,6 +106,30 @@ elif [ -n "${INSTANCE_POOLS:-}" ]; then
 elif [ -n "${INSTANCE_TYPE:-}" ]; then
     INSTANCE_TYPES="${INSTANCE_TYPE}"
     echo "   Instances (from INSTANCE_TYPE): ${INSTANCE_TYPES}"
+elif [ "${ENDPOINT_EXTERNAL:-}" = "true" ] && [ -n "${ENDPOINT_NAME:-}" ]; then
+    # External endpoint — query the live endpoint config for instance type
+    echo "   Discovering instance type from external endpoint: ${ENDPOINT_NAME}"
+    ENDPOINT_CONFIG_NAME=$(aws sagemaker describe-endpoint \
+        --endpoint-name "${ENDPOINT_NAME}" \
+        --region "${AWS_REGION}" \
+        --query 'EndpointConfigName' \
+        --output text 2>/dev/null) || ENDPOINT_CONFIG_NAME=""
+
+    if [ -n "${ENDPOINT_CONFIG_NAME}" ]; then
+        INSTANCE_TYPES=$(aws sagemaker describe-endpoint-config \
+            --endpoint-config-name "${ENDPOINT_CONFIG_NAME}" \
+            --region "${AWS_REGION}" \
+            --query 'ProductionVariants[0].InstanceType' \
+            --output text 2>/dev/null) || INSTANCE_TYPES=""
+    fi
+
+    if [ -n "${INSTANCE_TYPES}" ] && [ "${INSTANCE_TYPES}" != "None" ]; then
+        echo "   Instances (from endpoint): ${INSTANCE_TYPES}"
+    else
+        echo "❌ Could not discover instance type from endpoint: ${ENDPOINT_NAME}"
+        echo "   Provide --instances flag, or set INSTANCE_TYPE in do/config."
+        exit 1
+    fi
 else
     echo "❌ No instance types available."
     echo "   Provide --instances flag, or set INSTANCE_POOLS or INSTANCE_TYPE in do/config."
@@ -132,9 +156,9 @@ RESUME_EXISTING=false
 
 if [ "${FORCE}" = false ] && [ -n "${OPTIMIZE_JOB_NAME:-}" ]; then
     EXISTING_STATUS=$(aws sagemaker describe-ai-recommendation-job \
-        --job-name "${OPTIMIZE_JOB_NAME}" \
+        --ai-recommendation-job-name "${OPTIMIZE_JOB_NAME}" \
         --region "${AWS_REGION}" \
-        --query 'Status' \
+        --query 'AIRecommendationJobStatus' \
         --output text 2>/dev/null) || EXISTING_STATUS=""
 
     case "${EXISTING_STATUS}" in
@@ -154,7 +178,7 @@ if [ "${FORCE}" = false ] && [ -n "${OPTIMIZE_JOB_NAME:-}" ]; then
             ;;
         FAILED|STOPPED)
             FAILURE_REASON=$(aws sagemaker describe-ai-recommendation-job \
-                --job-name "${OPTIMIZE_JOB_NAME}" \
+                --ai-recommendation-job-name "${OPTIMIZE_JOB_NAME}" \
                 --region "${AWS_REGION}" \
                 --query 'FailureReason' \
                 --output text 2>/dev/null) || FAILURE_REASON="unknown"
@@ -174,44 +198,86 @@ fi
 # ── Create recommendation job ─────────────────────────────────────────────────
 if [ "${RESUME_EXISTING}" = false ]; then
     OPTIMIZE_JOB_NAME="${PROJECT_NAME}-optimize-$(date +%Y%m%d-%H%M%S)"
+    WORKLOAD_CONFIG_NAME="${OPTIMIZE_JOB_NAME}-workload"
 
     echo "🚀 Creating AI Recommendation Job: ${OPTIMIZE_JOB_NAME}"
 
-    # Build instance type list as JSON array
-    INSTANCE_TYPES_JSON="["
-    FIRST=true
+    # Step 1: Create workload config
+    echo "   Creating workload config: ${WORKLOAD_CONFIG_NAME}"
+
+    WORKLOAD_SPEC_INNER="{\"benchmark\": {\"type\": \"aiperf\"}, \"parameters\": {\"prompt_input_tokens_mean\": ${INPUT_TOKENS}, \"prompt_input_tokens_stddev\": 150, \"output_tokens_mean\": ${OUTPUT_TOKENS}, \"output_tokens_stddev\": 50, \"concurrency\": ${CONCURRENCY}, \"streaming\": true}}"
+    WORKLOAD_SPEC_OUTER="{\"WorkloadSpec\": {\"Inline\": $(python3 -c "import json; print(json.dumps('${WORKLOAD_SPEC_INNER}'))")}}"
+
+    if ! aws sagemaker create-ai-workload-config \
+        --ai-workload-config-name "${WORKLOAD_CONFIG_NAME}" \
+        --ai-workload-configs "${WORKLOAD_SPEC_OUTER}" \
+        --region "${AWS_REGION}" 2>&1 | grep -q "AIWorkloadConfigArn"; then
+        echo "❌ Failed to create workload config: ${WORKLOAD_CONFIG_NAME}"
+        echo "   Check that the execution role has sagemaker:CreateAIWorkloadConfig permission."
+        # Show actual error
+        aws sagemaker create-ai-workload-config \
+            --ai-workload-config-name "${WORKLOAD_CONFIG_NAME}" \
+            --ai-workload-configs "${WORKLOAD_SPEC_OUTER}" \
+            --region "${AWS_REGION}" 2>&1 || true
+        exit 1
+    fi
+    echo "   ✅ Workload config created"
+
+    # Step 2: Build compute spec (instance types, max 3)
+    COMPUTE_SPEC_JSON=""
     IFS=',' read -ra TYPES <<< "${INSTANCE_TYPES}"
+    INSTANCE_LIST=""
+    COUNT=0
     for itype in "${TYPES[@]}"; do
         itype=$(echo "${itype}" | xargs)  # trim whitespace
-        if [ "${FIRST}" = true ]; then
-            INSTANCE_TYPES_JSON="${INSTANCE_TYPES_JSON}\"${itype}\""
-            FIRST=false
+        if [ ${COUNT} -ge 3 ]; then
+            echo "   ⚠️  Max 3 instance types supported — truncating"
+            break
+        fi
+        if [ -n "${INSTANCE_LIST}" ]; then
+            INSTANCE_LIST="${INSTANCE_LIST},\"${itype}\""
         else
-            INSTANCE_TYPES_JSON="${INSTANCE_TYPES_JSON},\"${itype}\""
+            INSTANCE_LIST="\"${itype}\""
         fi
+        COUNT=$((COUNT + 1))
     done
-    INSTANCE_TYPES_JSON="${INSTANCE_TYPES_JSON}]"
+    COMPUTE_SPEC_JSON="InstanceTypes=[${INSTANCE_LIST}]"
+
+    # Step 3: Map goal to performance target metric
+    case "${GOAL}" in
+        latency) PERF_METRIC="ttft-ms" ;;
+        throughput) PERF_METRIC="throughput" ;;
+        cost) PERF_METRIC="cost" ;;
+    esac
 
-    # Build job input config
-    # The model is specified as either a HuggingFace model ID or S3 path
-    MODEL_SOURCE_JSON=""
+    # Step 4: Determine model source
+    # The recommendation API requires model artifacts as s3:// or https:// URI.
+    MODEL_SOURCE_ARG=""
     if [[ "${MODEL_NAME}" == s3://* ]]; then
-        MODEL_SOURCE_JSON="{\"S3DataSource\":{\"S3Uri\":\"${MODEL_NAME}\"}}"
+        MODEL_SOURCE_ARG="S3={S3Uri=${MODEL_NAME}}"
     else
-        MODEL_SOURCE_JSON="{\"ModelName\":\"${MODEL_NAME}\"}"
+        # HuggingFace model — use the HTTPS URL for the model on HuggingFace Hub
+        MODEL_SOURCE_ARG="S3={S3Uri=https://huggingface.co/${MODEL_NAME}}"
     fi
 
-    # Build workload config
-    WORKLOAD_JSON="{\"Concurrency\":${CONCURRENCY},\"InputTokens\":${INPUT_TOKENS},\"OutputTokens\":${OUTPUT_TOKENS}}"
-
-    # Build the full job specification
-    JOB_INPUT="{\"ModelSource\":${MODEL_SOURCE_JSON},\"Workload\":${WORKLOAD_JSON},\"InstanceTypes\":${INSTANCE_TYPES_JSON},\"OptimizationGoal\":\"${GOAL}\"}"
-
-    if ! aws sagemaker create-ai-recommendation-job \
-        --job-name "${OPTIMIZE_JOB_NAME}" \
-        --role-arn "${ROLE_ARN}" \
-        --input-config "${JOB_INPUT}" \
-        --region "${AWS_REGION}"; then
+    # Step 5: Create the recommendation job
+    S3_OUTPUT="s3://${TUNE_S3_BUCKET:-mlcc-tune-$(aws sts get-caller-identity --query Account --output text 2>/dev/null)-${AWS_REGION}}/${PROJECT_NAME}/optimize/"
+
+    RECOMMEND_CMD=(
+        aws sagemaker create-ai-recommendation-job
+        --ai-recommendation-job-name "${OPTIMIZE_JOB_NAME}"
+        --model-source "${MODEL_SOURCE_ARG}"
+        --output-config "S3OutputLocation=${S3_OUTPUT}"
+        --ai-workload-config-identifier "${WORKLOAD_CONFIG_NAME}"
+        --performance-target "Constraints=[{Metric=${PERF_METRIC}}]"
+        --role-arn "${ROLE_ARN}"
+        --compute-spec "${COMPUTE_SPEC_JSON}"
+        --inference-specification "Framework=VLLM"
+        --region "${AWS_REGION}"
+    )
+
+    if ! "${RECOMMEND_CMD[@]}" 2>&1; then
+        echo ""
         echo "❌ Failed to create AI Recommendation Job"
         echo "   Check that:"
         echo "   • The execution role has sagemaker:CreateAIRecommendationJob permission"
@@ -232,7 +298,7 @@ fi
 POLL_INTERVAL=30
 MAX_POLL_ATTEMPTS=120  # 60 minutes max (120 * 30s)
 
-if [ "${JOB_STATUS:-}" != "COMPLETED" ] && [ "${JOB_STATUS:-}" != "FAILED" ] && [ "${JOB_STATUS:-}" != "STOPPED" ]; then
+if [ "${JOB_STATUS:-}" != "COMPLETED" ] && [ "${JOB_STATUS:-}" != "FAILED" ] && [ "${JOB_STATUS:-}" != "STOPPED" ] && [ "${JOB_STATUS:-}" != "Completed" ] && [ "${JOB_STATUS:-}" != "Failed" ] && [ "${JOB_STATUS:-}" != "Stopped" ]; then
 
 echo "⏳ Waiting for recommendation job to complete..."
 echo "   Polling every ${POLL_INTERVAL}s (max ${MAX_POLL_ATTEMPTS} attempts = 60 min)"
@@ -243,27 +309,30 @@ JOB_STATUS=""
 
 while [ ${POLL_COUNT} -lt ${MAX_POLL_ATTEMPTS} ]; do
     JOB_STATUS=$(aws sagemaker describe-ai-recommendation-job \
-        --job-name "${OPTIMIZE_JOB_NAME}" \
+        --ai-recommendation-job-name "${OPTIMIZE_JOB_NAME}" \
         --region "${AWS_REGION}" \
-        --query 'Status' \
+        --query 'AIRecommendationJobStatus' \
         --output text 2>/dev/null) || {
         echo "⚠️  Failed to describe recommendation job (credentials may have expired)"
         echo "   Re-run to check status:"
-        echo "   aws sagemaker describe-ai-recommendation-job --job-name ${OPTIMIZE_JOB_NAME} --region ${AWS_REGION}"
+        echo "   aws sagemaker describe-ai-recommendation-job --ai-recommendation-job-name ${OPTIMIZE_JOB_NAME} --region ${AWS_REGION}"
         exit 1
     }
 
     case "${JOB_STATUS}" in
-        COMPLETED)
+        COMPLETED|Completed)
             echo "✅ Recommendation job completed!"
+            JOB_STATUS="COMPLETED"
             break
             ;;
-        FAILED)
+        FAILED|Failed)
             echo "❌ Recommendation job failed"
+            JOB_STATUS="FAILED"
             break
             ;;
-        STOPPED)
+        STOPPED|Stopped)
             echo "⚠️  Recommendation job was stopped"
+            JOB_STATUS="STOPPED"
             break
             ;;
         *)
@@ -292,7 +361,7 @@ if [ "${JOB_STATUS}" = "COMPLETED" ]; then
 
     # Get the full job description with results
     JOB_DESCRIPTION=$(aws sagemaker describe-ai-recommendation-job \
-        --job-name "${OPTIMIZE_JOB_NAME}" \
+        --ai-recommendation-job-name "${OPTIMIZE_JOB_NAME}" \
         --region "${AWS_REGION}" \
         --output json 2>/dev/null) || {
         echo "❌ Failed to fetch recommendation results"
@@ -503,7 +572,7 @@ except:
 
     if [ "${FAILURE_REASON}" = "unknown" ]; then
         FAILURE_REASON=$(aws sagemaker describe-ai-recommendation-job \
-            --job-name "${OPTIMIZE_JOB_NAME}" \
+            --ai-recommendation-job-name "${OPTIMIZE_JOB_NAME}" \
             --region "${AWS_REGION}" \
             --query 'FailureReason' \
             --output text 2>/dev/null) || FAILURE_REASON="unknown"
@@ -513,7 +582,7 @@ except:
     echo "   Reason: ${FAILURE_REASON}"
     echo ""
     echo "   Debug:"
-    echo "   aws sagemaker describe-ai-recommendation-job --job-name ${OPTIMIZE_JOB_NAME} --region ${AWS_REGION}"
+    echo "   aws sagemaker describe-ai-recommendation-job --ai-recommendation-job-name ${OPTIMIZE_JOB_NAME} --region ${AWS_REGION}"
 
 elif [ "${JOB_STATUS}" = "STOPPED" ]; then
     echo "⚠️  Recommendation job was stopped before completion"

package/templates/do/register

@@ -22,6 +22,8 @@ JSON_OUTPUT=false
 CI_MODE=false
 CI_TABLE_NAME="${CI_TABLE_NAME:-mlcc-ci-table}"
 CI_BUILD_STRATEGY="codebuild-submit"
+BENCHMARK_STATUS=""
+BENCHMARK_RUN_ID=""
 
 while [[ $# -gt 0 ]]; do
     case "$1" in
@@ -70,15 +72,102 @@ while [[ $# -gt 0 ]]; do
             CI_BUILD_STRATEGY="${1#*=}"
             shift
             ;;
+        --benchmark-status)
+            BENCHMARK_STATUS="$2"
+            shift 2
+            ;;
+        --benchmark-status=*)
+            BENCHMARK_STATUS="${1#*=}"
+            shift
+            ;;
+        --benchmark-run-id)
+            BENCHMARK_RUN_ID="$2"
+            shift 2
+            ;;
+        --benchmark-run-id=*)
+            BENCHMARK_RUN_ID="${1#*=}"
+            shift
+            ;;
         *)
             echo "⚠️  Unknown option: $1"
             echo ""
             echo "Usage: ./do/register [--notes \"text\"] [--status success|partial|failed] [--project] [--json] [--ci] [--ci-table <name>] [--build-strategy <strategy>]"
+            echo "       ./do/register --benchmark-status <completed|failed|in-progress> --benchmark-run-id <run-id>"
             exit 1
             ;;
     esac
 done
 
+# ============================================================
+# Handle --benchmark-status mode (Stage 2 failure isolation)
+# This mode ONLY writes benchmark fields to DynamoDB.
+# It uses an UpdateExpression that sets ONLY:
+#   lastBenchmarkRunId, lastBenchmarkTimestamp, lastBenchmarkStatus
+# It does NOT touch testStatus, configJson, or any other fields.
+# This ensures Stage 2 failure does not change the CI gate result.
+# Requirements: 1.4, 7.3
+# ============================================================
+
+if [ -n "${BENCHMARK_STATUS}" ]; then
+    # Validate benchmark status
+    case "${BENCHMARK_STATUS}" in
+        completed|failed|in-progress) ;;
+        *)
+            echo "❌ Invalid benchmark status: ${BENCHMARK_STATUS}"
+            echo "   Valid values: completed, failed, in-progress"
+            exit 1
+            ;;
+    esac
+
+    if [ -z "${BENCHMARK_RUN_ID}" ]; then
+        echo "❌ --benchmark-run-id is required when using --benchmark-status"
+        exit 1
+    fi
+
+    # Compute configId (same hash logic used by write_ci_record below)
+    _benchmark_input="${DEPLOYMENT_CONFIG}:${MODEL_NAME:-none}:${INSTANCE_TYPE}:${AWS_REGION}:${DEPLOYMENT_TARGET}:ic${IC_COUNT:-1}:adapt${ADAPTER_COUNT:-0}"
+    if command -v sha256sum &> /dev/null; then
+        CONFIG_ID=$(echo -n "$_benchmark_input" | sha256sum | cut -c1-16)
+    else
+        CONFIG_ID=$(echo -n "$_benchmark_input" | shasum -a 256 | cut -c1-16)
+    fi
+    BENCHMARK_TIMESTAMP=$(date -u +%Y-%m-%dT%H:%M:%SZ)
+
+    echo "📊 Recording benchmark status for configId: ${CONFIG_ID}"
+    echo "   Benchmark run ID:  ${BENCHMARK_RUN_ID}"
+    echo "   Benchmark status:  ${BENCHMARK_STATUS}"
+    echo "   Timestamp:         ${BENCHMARK_TIMESTAMP}"
+    echo ""
+
+    # Check if CI_Table exists
+    if ! aws dynamodb describe-table --table-name "${CI_TABLE_NAME}" --region "${AWS_REGION}" &>/dev/null; then
+        echo "⚠️  CI infrastructure not provisioned. Run 'ml-container-creator bootstrap' with CI enabled."
+        echo "   Skipping benchmark status write."
+        exit 0
+    fi
+
+    # UpdateExpression ONLY sets the 3 benchmark fields.
+    # This is a SET operation — it does NOT overwrite the entire record (PutItem).
+    # testStatus, configJson, and all other fields remain unchanged.
+    if aws dynamodb update-item \
+        --region "${AWS_REGION}" \
+        --table-name "${CI_TABLE_NAME}" \
+        --key "{\"configId\": {\"S\": \"${CONFIG_ID}\"}}" \
+        --update-expression "SET lastBenchmarkRunId = :rid, lastBenchmarkTimestamp = :ts, lastBenchmarkStatus = :bs" \
+        --expression-attribute-values "{
+            \":rid\": {\"S\": \"${BENCHMARK_RUN_ID}\"},
+            \":ts\": {\"S\": \"${BENCHMARK_TIMESTAMP}\"},
+            \":bs\": {\"S\": \"${BENCHMARK_STATUS}\"}
+        }" 2>/dev/null; then
+        echo "✅ Benchmark status recorded: ${BENCHMARK_STATUS}"
+    else
+        echo "❌ Failed to write benchmark status for configId: ${CONFIG_ID}"
+        exit 1
+    fi
+
+    exit 0
+fi
+
 # Validate status
 case "${STATUS}" in
     success|partial|failed) ;;

package/templates/do/test

@@ -214,6 +214,7 @@ else
     # Resolve inference component name
     # Precedence: do/adapters/ → do/ic/ → legacy config
     IC_NAME=""
+    ADAPTER_MODEL_NAME=""
     if [ -n "${IC_ARG}" ] && [ -f "${SCRIPT_DIR}/adapters/${IC_ARG}.conf" ]; then
         # Argument matches an adapter name — use adapter IC
         ADAPTER_IC_NAME=""
@@ -223,6 +224,7 @@ else
             exit 1
         fi
         IC_NAME="${ADAPTER_IC_NAME}"
+        ADAPTER_MODEL_NAME="${IC_ARG}"
     elif [ -n "${IC_ARG}" ]; then
         # Explicit IC name provided as argument
         IC_CONF="${SCRIPT_DIR}/ic/${IC_ARG}.conf"
@@ -258,6 +260,17 @@ else
         IC_NAME="${INFERENCE_COMPONENT_NAME:-}"
     fi
 
+    # If testing an adapter, override the model name in the payload
+    if [ -n "${ADAPTER_MODEL_NAME}" ] && [ "${FRAMEWORK}" = "transformers" ]; then
+        case "${MODEL_SERVER}" in
+            vllm|sglang)
+                TEST_PAYLOAD='{"model": "'"${ADAPTER_MODEL_NAME}"'", "messages": [{"role": "user", "content": "What is machine learning?"}], "max_tokens": 50, "temperature": 0.7}'
+                echo "${TEST_PAYLOAD}" > "${TEMP_PAYLOAD}"
+                echo "   (Using adapter model name: ${ADAPTER_MODEL_NAME})"
+                ;;
+        esac
+    fi
+
     INVOKE_ARGS=(
         --endpoint-name "${ENDPOINT_NAME}"
         --region "${AWS_REGION}"