@aws/ml-container-creator 0.2.6 → 0.4.0

package/templates/do/clean

@@ -740,6 +740,49 @@ case "${CLEANUP_TARGET}" in
     codebuild)
         clean_codebuild
         ;;
+<% if (typeof includeBenchmark !== 'undefined' && includeBenchmark) { %>
+    benchmark)
+        echo "🧹 Cleaning benchmark resources..."
+        WORKLOAD_CONFIG_NAME="${PROJECT_NAME}-benchmark-config"
+
+        # Delete workload config if exists
+        if aws sagemaker describe-ai-workload-config \
+            --ai-workload-config-name "$WORKLOAD_CONFIG_NAME" \
+            --region "$AWS_REGION" 2>/dev/null; then
+            aws sagemaker delete-ai-workload-config \
+                --ai-workload-config-name "$WORKLOAD_CONFIG_NAME" \
+                --region "$AWS_REGION"
+            echo "   ✓ Deleted workload config: $WORKLOAD_CONFIG_NAME"
+        fi
+
+        # Delete terminal benchmark jobs matching project prefix
+        aws sagemaker list-ai-benchmark-jobs \
+            --name-contains "${PROJECT_NAME}-benchmark-" \
+            --region "$AWS_REGION" \
+            --query 'AIBenchmarkJobs[?AIBenchmarkJobStatus!=`InProgress`].AIBenchmarkJobName' \
+            --output text | tr '\t' '\n' | while read -r job; do
+            [ -z "$job" ] && continue
+            aws sagemaker delete-ai-benchmark-job \
+                --ai-benchmark-job-name "$job" \
+                --region "$AWS_REGION"
+            echo "   ✓ Deleted benchmark job: $job"
+        done
+
+        # Delete local benchmark results
+        if [ -d "${SCRIPT_DIR}/../benchmarks" ]; then
+            read -p "Delete local benchmark results? (Y/n) " CONFIRM_DELETE
+            CONFIRM_DELETE="${CONFIRM_DELETE:-Y}"
+            if [[ "${CONFIRM_DELETE}" =~ ^[Yy]$ ]]; then
+                rm -rf "${SCRIPT_DIR}/../benchmarks"
+                echo "   ✓ Deleted local benchmarks/ directory"
+            else
+                echo "   ⏭ Skipped local benchmarks/ deletion"
+            fi
+        fi
+
+        echo "✅ Benchmark cleanup complete"
+        ;;
+<% } %>
     all)
         echo "🧹 Performing complete cleanup"
         echo ""
@@ -790,6 +833,49 @@ case "${CLEANUP_TARGET}" in
             CLEANED_ITEMS+=("CodeBuild resources")
         fi
         
+<% if (typeof includeBenchmark !== 'undefined' && includeBenchmark) { %>
+        echo ""
+        
+        # Clean benchmark resources
+        WORKLOAD_CONFIG_NAME="${PROJECT_NAME}-benchmark-config"
+
+        # Delete workload config if exists
+        if aws sagemaker describe-ai-workload-config \
+            --ai-workload-config-name "$WORKLOAD_CONFIG_NAME" \
+            --region "$AWS_REGION" 2>/dev/null; then
+            aws sagemaker delete-ai-workload-config \
+                --ai-workload-config-name "$WORKLOAD_CONFIG_NAME" \
+                --region "$AWS_REGION"
+            echo "   ✓ Deleted workload config: $WORKLOAD_CONFIG_NAME"
+        fi
+
+        # Delete terminal benchmark jobs matching project prefix
+        aws sagemaker list-ai-benchmark-jobs \
+            --name-contains "${PROJECT_NAME}-benchmark-" \
+            --region "$AWS_REGION" \
+            --query 'AIBenchmarkJobs[?AIBenchmarkJobStatus!=`InProgress`].AIBenchmarkJobName' \
+            --output text | tr '\t' '\n' | while read -r job; do
+            [ -z "$job" ] && continue
+            aws sagemaker delete-ai-benchmark-job \
+                --ai-benchmark-job-name "$job" \
+                --region "$AWS_REGION"
+            echo "   ✓ Deleted benchmark job: $job"
+        done
+
+        # Delete local benchmark results
+        if [ -d "${SCRIPT_DIR}/../benchmarks" ]; then
+            read -p "Delete local benchmark results? (Y/n) " CONFIRM_DELETE
+            CONFIRM_DELETE="${CONFIRM_DELETE:-Y}"
+            if [[ "${CONFIRM_DELETE}" =~ ^[Yy]$ ]]; then
+                rm -rf "${SCRIPT_DIR}/../benchmarks"
+                echo "   ✓ Deleted local benchmarks/ directory"
+            else
+                echo "   ⏭ Skipped local benchmarks/ deletion"
+            fi
+        fi
+
+        CLEANED_ITEMS+=("Benchmark resources")
+<% } %>
         # Display summary
         echo ""
         echo "✅ Cleanup complete!"

package/templates/do/config

@@ -30,6 +30,9 @@ export INSTANCE_TYPE="<%= instanceType %>"
 <% if (inferenceAmiVersion) { %>
 export INFERENCE_AMI_VERSION="<%= inferenceAmiVersion %>"
 <% } %>
+<% if (typeof capacityReservationArn !== 'undefined' && capacityReservationArn) { %>
+export CAPACITY_RESERVATION_ARN="<%= capacityReservationArn %>"
+<% } %>
 <% } %>
 
 <% if (deploymentTarget === 'async-inference') { %>
@@ -126,6 +129,8 @@ export IC_MEMORY_SIZE="<%= icMemorySize %>"
 <% } %>
 <% if (typeof icGpuCount !== 'undefined' && icGpuCount != null) { %>
 export IC_GPU_COUNT="<%= icGpuCount %>"
+<% } else { %>
+export IC_GPU_COUNT="${IC_GPU_COUNT:-1}"
 <% } %>
 <% if (typeof icCopyCount !== 'undefined' && icCopyCount != null) { %>
 export IC_COPY_COUNT="<%= icCopyCount %>"
@@ -151,17 +156,29 @@ export <%= key %>=${<%= key %>:-<%= value %>}
 # Framework-specific configuration
 <% if (framework === 'transformers') { %>
 export MODEL_NAME="<%= modelName %>"
-<% if (hfToken) { %>
+# Secrets Manager integration: when an ARN is configured, do-scripts resolve the
+# secret at the appropriate stage (build-time or runtime). When a plaintext value
+# is configured, it is exported directly. The _ARN suffix signals resolution is needed.
+<% if (typeof hfTokenArn !== 'undefined' && hfTokenArn) { %>
+export HF_TOKEN_ARN="<%= hfTokenArn %>"
+<% } else if (hfToken) { %>
 export HF_TOKEN="<%= hfToken %>"
 <% } %>
-<% if (ngcApiKey) { %>
+<% if (typeof ngcTokenArn !== 'undefined' && ngcTokenArn) { %>
+export NGC_API_KEY_ARN="<%= ngcTokenArn %>"
+<% } else if (ngcApiKey) { %>
 export NGC_API_KEY="<%= ngcApiKey %>"
 <% } %>
 <% } %>
 
 <% if (framework === 'diffusors') { %>
 export MODEL_NAME="<%= modelName %>"
-<% if (hfToken) { %>
+# Secrets Manager integration: when an ARN is configured, do-scripts resolve the
+# secret at the appropriate stage (build-time or runtime). When a plaintext value
+# is configured, it is exported directly. The _ARN suffix signals resolution is needed.
+<% if (typeof hfTokenArn !== 'undefined' && hfTokenArn) { %>
+export HF_TOKEN_ARN="<%= hfTokenArn %>"
+<% } else if (hfToken) { %>
 export HF_TOKEN="<%= hfToken %>"
 <% } %>
 <% } %>
@@ -174,6 +191,26 @@ export MODEL_FORMAT="<%= modelFormat %>"
 export ROLE_ARN="<%= roleArn %>"
 <% } %>
 
+<% if (typeof includeBenchmark !== 'undefined' && includeBenchmark) { %>
+# SageMaker AI Benchmarking configuration
+export BENCHMARK_CONCURRENCY="<%= benchmarkConcurrency %>"
+export BENCHMARK_INPUT_TOKENS_MEAN="<%= benchmarkInputTokensMean %>"
+export BENCHMARK_OUTPUT_TOKENS_MEAN="<%= benchmarkOutputTokensMean %>"
+export BENCHMARK_STREAMING="<%= benchmarkStreaming %>"
+<% if (benchmarkRequestCount) { %>
+export BENCHMARK_REQUEST_COUNT="<%= benchmarkRequestCount %>"
+<% } else { %>
+export BENCHMARK_REQUEST_COUNT=""
+<% } %>
+<% if (benchmarkS3OutputPath) { %>
+export BENCHMARK_S3_OUTPUT_PATH="<%= benchmarkS3OutputPath %>"
+<% } else { %>
+export BENCHMARK_S3_OUTPUT_PATH="s3://ml-container-creator-benchmark-${AWS_REGION}-$(aws sts get-caller-identity --query Account --output text)/${PROJECT_NAME}/"
+<% } %>
+export BENCHMARK_JOB_NAME=""
+export BENCHMARK_WORKLOAD_CONFIG_NAME=""
+<% } %>
+
 <% if (orderedEnvVars && orderedEnvVars.length > 0) { %>
 # Runtime environment variables (from catalog)
 <% orderedEnvVars.forEach(({ key, value }) => { %>
@@ -181,9 +218,7 @@ export <%= key %>=${<%= key %>:-<%= value %>}
 <% }); %>
 <% } %>
 
-<% if (baseImage) { %>
-export BASE_IMAGE=${BASE_IMAGE:-<%= baseImage %>}
-<% } %>
+export BASE_IMAGE=${BASE_IMAGE:-<%= baseImage || '' %>}
 
 # Allow environment variable overrides
 export AWS_REGION=${AWS_REGION:-<%= awsRegion %>}

package/templates/do/deploy

@@ -95,6 +95,41 @@ fi
 echo "✅ ECR image found: ${ECR_REPOSITORY}:${PROJECT_NAME}-latest"
 IMAGE_TAG="${PROJECT_NAME}-latest"
 
+# ============================================================
+# Shared: Resolve secrets for container environment
+# ============================================================
+CONTAINER_ENV_JSON=""
+
+if [ -n "${HF_TOKEN_ARN:-}" ]; then
+    echo "🔐 Resolving HuggingFace token from Secrets Manager..."
+    RESOLVED_HF_TOKEN=$(aws secretsmanager get-secret-value --secret-id "${HF_TOKEN_ARN}" --query SecretString --output text --region "${AWS_REGION}") || {
+        echo "❌ Failed to resolve HuggingFace token from Secrets Manager"
+        exit 3
+    }
+    CONTAINER_ENV_JSON="\"HF_TOKEN\":\"${RESOLVED_HF_TOKEN}\""
+elif [ -n "${HF_TOKEN:-}" ]; then
+    CONTAINER_ENV_JSON="\"HF_TOKEN\":\"${HF_TOKEN}\""
+fi
+
+if [ -n "${NGC_API_KEY_ARN:-}" ]; then
+    echo "🔐 Resolving NGC API key from Secrets Manager..."
+    RESOLVED_NGC_KEY=$(aws secretsmanager get-secret-value --secret-id "${NGC_API_KEY_ARN}" --query SecretString --output text --region "${AWS_REGION}") || {
+        echo "❌ Failed to resolve NGC API key from Secrets Manager"
+        exit 3
+    }
+    if [ -n "${CONTAINER_ENV_JSON}" ]; then
+        CONTAINER_ENV_JSON="${CONTAINER_ENV_JSON},\"NGC_API_KEY\":\"${RESOLVED_NGC_KEY}\""
+    else
+        CONTAINER_ENV_JSON="\"NGC_API_KEY\":\"${RESOLVED_NGC_KEY}\""
+    fi
+elif [ -n "${NGC_API_KEY:-}" ]; then
+    if [ -n "${CONTAINER_ENV_JSON}" ]; then
+        CONTAINER_ENV_JSON="${CONTAINER_ENV_JSON},\"NGC_API_KEY\":\"${NGC_API_KEY}\""
+    else
+        CONTAINER_ENV_JSON="\"NGC_API_KEY\":\"${NGC_API_KEY}\""
+    fi
+fi
+
 <% if (deploymentTarget === 'realtime-inference') { %>
 # ============================================================
 # SageMaker Real-Time Inference Deployment (Inference Components)
@@ -301,6 +336,11 @@ if [ -z "${SKIP_TO}" ]; then
         echo "   AMI version: ${INFERENCE_AMI_VERSION}"
     fi
 
+    if [ -n "${CAPACITY_RESERVATION_ARN:-}" ]; then
+        VARIANT_JSON="${VARIANT_JSON},\"CapacityReservationConfig\":{\"CapacityReservationPreference\":\"capacity-reservations-only\",\"MlReservationArn\":\"${CAPACITY_RESERVATION_ARN}\"}"
+        echo "   ⚠️  Capacity reservation (experimental): ${CAPACITY_RESERVATION_ARN}"
+    fi
+
     VARIANT_JSON="${VARIANT_JSON}}]"
 
     echo "⚙️  Creating endpoint configuration: ${ENDPOINT_CONFIG_NAME}"
@@ -400,20 +440,25 @@ if [ -z "${SKIP_TO}" ] || [ "${SKIP_TO}" = "create_ic" ] || [ "${SKIP_TO}" = "wa
         _update_config_var "INFERENCE_COMPONENT_NAME" "${IC_NAME}"
     fi
 
+    # Build container spec JSON
+    CONTAINER_SPEC="{\"Image\":\"${ECR_REPOSITORY}:${IMAGE_TAG}\""
+    if [ -n "${CONTAINER_ENV_JSON}" ]; then
+        CONTAINER_SPEC="${CONTAINER_SPEC},\"Environment\":{${CONTAINER_ENV_JSON}}"
+    fi
+    CONTAINER_SPEC="${CONTAINER_SPEC}}"
+
     echo "📦 Creating inference component: ${IC_NAME}"
     if ! aws sagemaker create-inference-component \
         --inference-component-name "${IC_NAME}" \
         --endpoint-name "${ENDPOINT_NAME}" \
         --variant-name "AllTraffic" \
         --specification "{
-            \"Container\": {
-                \"Image\": \"${ECR_REPOSITORY}:${IMAGE_TAG}\"
-            },
+            \"Container\": ${CONTAINER_SPEC},
             \"StartupParameters\": {
                 \"ContainerStartupHealthCheckTimeoutInSeconds\": 900
             },
             \"ComputeResourceRequirements\": {
-                \"NumberOfAcceleratorDevicesRequired\": 1,
+                \"NumberOfAcceleratorDevicesRequired\": ${IC_GPU_COUNT},
                 \"MinMemoryRequiredInMb\": 1024
             }
         }" \
@@ -767,10 +812,17 @@ if [ -z "${SKIP_TO}" ]; then
     _update_config_var "SAGEMAKER_MODEL_NAME" "${MODEL_NAME_SM}"
 
     # Step 1: Create SageMaker model
+    # Build primary container spec
+    PRIMARY_CONTAINER="{\"Image\":\"${ECR_REPOSITORY}:${IMAGE_TAG}\""
+    if [ -n "${CONTAINER_ENV_JSON}" ]; then
+        PRIMARY_CONTAINER="${PRIMARY_CONTAINER},\"Environment\":{${CONTAINER_ENV_JSON}}"
+    fi
+    PRIMARY_CONTAINER="${PRIMARY_CONTAINER}}"
+
     echo "📦 Creating SageMaker model: ${MODEL_NAME_SM}"
     if ! aws sagemaker create-model \
         --model-name "${MODEL_NAME_SM}" \
-        --primary-container "{\"Image\":\"${ECR_REPOSITORY}:${IMAGE_TAG}\"}" \
+        --primary-container "${PRIMARY_CONTAINER}" \
         --execution-role-arn "${ROLE_ARN}" \
         --region "${AWS_REGION}"; then
 
@@ -1361,9 +1413,17 @@ _update_config_var "SAGEMAKER_MODEL_NAME" "${MODEL_NAME_SM}"
 
 # Step 1: Create SageMaker model
 echo "📦 Creating SageMaker model: ${MODEL_NAME_SM}"
+
+# Build primary container spec
+BATCH_PRIMARY_CONTAINER="{\"Image\":\"${ECR_REPOSITORY}:${IMAGE_TAG}\""
+if [ -n "${CONTAINER_ENV_JSON}" ]; then
+    BATCH_PRIMARY_CONTAINER="${BATCH_PRIMARY_CONTAINER},\"Environment\":{${CONTAINER_ENV_JSON}}"
+fi
+BATCH_PRIMARY_CONTAINER="${BATCH_PRIMARY_CONTAINER}}"
+
 if ! aws sagemaker create-model \
     --model-name "${MODEL_NAME_SM}" \
-    --primary-container "{\"Image\":\"${ECR_REPOSITORY}:${IMAGE_TAG}\"}" \
+    --primary-container "${BATCH_PRIMARY_CONTAINER}" \
     --execution-role-arn "${ROLE_ARN}" \
     --region "${AWS_REGION}"; then
 

package/templates/do/logs

@@ -51,11 +51,15 @@ echo "━━━━━━━━━━━━━━━━━━━━━━━━
 echo ""
 
 # Wait for log group to exist before tailing
-MAX_WAIT=300
+MAX_WAIT=900
 INTERVAL=10
 ELAPSED=0
 
+# Try IC-specific log group first, fall back to endpoint log group
+FALLBACK_LOG_GROUP="/aws/sagemaker/Endpoints/${ENDPOINT}"
+
 while true; do
+    # Check IC-specific log group
     if aws logs describe-log-groups \
         --log-group-name-prefix "${LOG_GROUP}" \
         --region "${AWS_REGION}" \
@@ -64,6 +68,17 @@ while true; do
         break
     fi
 
+    # Check endpoint-level log group as fallback
+    if aws logs describe-log-groups \
+        --log-group-name-prefix "${FALLBACK_LOG_GROUP}" \
+        --region "${AWS_REGION}" \
+        --query "logGroups[?logGroupName=='${FALLBACK_LOG_GROUP}'].logGroupName" \
+        --output text 2>/dev/null | grep -q "${FALLBACK_LOG_GROUP}"; then
+        LOG_GROUP="${FALLBACK_LOG_GROUP}"
+        echo "   ℹ️  Using endpoint log group: ${LOG_GROUP}"
+        break
+    fi
+
     if [ "${ELAPSED}" -ge "${MAX_WAIT}" ]; then
         echo "❌ Timed out after ${MAX_WAIT}s waiting for log group: ${LOG_GROUP}"
         echo ""
@@ -123,7 +138,7 @@ echo "━━━━━━━━━━━━━━━━━━━━━━━━
 echo ""
 
 # Wait for log group to exist before tailing
-MAX_WAIT=300
+MAX_WAIT=900
 INTERVAL=10
 ELAPSED=0
 
@@ -195,7 +210,7 @@ echo "━━━━━━━━━━━━━━━━━━━━━━━━
 echo ""
 
 # Wait for log group to exist before tailing
-MAX_WAIT=300
+MAX_WAIT=900
 INTERVAL=10
 ELAPSED=0
 

package/templates/do/register

@@ -393,6 +393,7 @@ CJEOF
 
     # Try put-item with condition (new record)
     if aws dynamodb put-item \
+        --region "${AWS_REGION}" \
         --table-name "${CI_TABLE_NAME}" \
         --item "{
             \"configId\": {\"S\": \"${config_id}\"},
@@ -412,6 +413,7 @@ CJEOF
     else
         # Record already exists — update it (reset testStatus, update configJson, preserve createdAt)
         if aws dynamodb update-item \
+            --region "${AWS_REGION}" \
             --table-name "${CI_TABLE_NAME}" \
             --key "{\"configId\": {\"S\": \"${config_id}\"}}" \
             --update-expression "SET configJson = :cj, testStatus = :ts, deploymentConfig = :dc, baseImage = :bi, baseImageVersion = :bv, buildStrategy = :bs, projectName = :pn, schemaVersion = :sv" \
@@ -496,6 +498,11 @@ DJEOF
     echo "${DEPLOYMENT_JSON}" | python3 -c "import sys,json; print(json.dumps(json.load(sys.stdin), indent=2))" 2>/dev/null || echo "${DEPLOYMENT_JSON}"
 
     if [ "${CI_MODE}" = true ]; then
+        # Strip capacity reservation ARN for CI — force on-demand deployment
+        # CI projects must never use reserved capacity (reservations are account-specific
+        # and time-bound; CI replay should always target on-demand instances)
+        unset CAPACITY_RESERVATION_ARN 2>/dev/null || true
+
         echo ""
         echo "⚠️  CI Integration is experimental and currently only tested for"
         echo "   SageMaker Real-Time Inference endpoints."
@@ -507,7 +514,7 @@ DJEOF
         echo "🔑 configId: ${CONFIG_ID}"
 
         # Check if CI_Table exists before writing
-        if ! aws dynamodb describe-table --table-name "${CI_TABLE_NAME}" &>/dev/null; then
+        if ! aws dynamodb describe-table --table-name "${CI_TABLE_NAME}" --region "${AWS_REGION}" &>/dev/null; then
             echo ""
             echo "⚠️  CI infrastructure not provisioned. Run 'ml-container-creator bootstrap' with CI enabled."
             echo "   Skipping CI table write."

package/templates/do/run

@@ -68,6 +68,16 @@ if [ -n "${MODEL_DIR:-}" ]; then
     fi
 fi
 
+# --- Secrets Manager resolution (runtime) ---
+if [ -n "${HF_TOKEN_ARN:-}" ]; then
+    echo "🔐 Resolving HuggingFace token from Secrets Manager..."
+    HF_TOKEN=$(aws secretsmanager get-secret-value --secret-id "${HF_TOKEN_ARN}" --query SecretString --output text) || {
+        echo "❌ Failed to resolve HuggingFace token from Secrets Manager"
+        exit 3
+    }
+    export HF_TOKEN
+fi
+
 # Prepare environment variables
 ENV_VARS=""
 <% if (framework === 'transformers') { %>

package/templates/triton/Dockerfile

@@ -122,6 +122,11 @@ EXPOSE 8080
 # --http-port=8080: SageMaker requires port 8080
 # --model-repository: Path to model repository
 # --strict-model-config=false: Allow Triton to auto-complete config for some backends
+
+# CUDA compatibility: ensure compat libs are on LD_LIBRARY_PATH for newer SageMaker AMIs
+# (NVIDIA Container Toolkit 1.17.4+ no longer auto-mounts these)
+ENV LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH:-}"
+
 ENTRYPOINT ["tritonserver", \
             "--http-port=8080", \
             "--model-repository=/opt/ml/model/model_repository", \