@aws/ml-container-creator 0.2.6 → 0.3.0

package/templates/do/config

@@ -151,17 +151,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 %>"
 <% } %>
 <% } %>

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)
@@ -400,15 +435,20 @@ 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
             },
@@ -767,10 +807,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 +1408,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/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') { %>