@aws/ml-container-creator 0.15.1 → 1.0.2

package/templates/do/benchmark

@@ -69,6 +69,9 @@ 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:-}}"
+
     JOB_NAME="${BENCHMARK_JOB_NAME:-}"
     if [ -z "${JOB_NAME}" ]; then
         echo "❌ No benchmark job tracked"
@@ -98,7 +101,7 @@ if [ "${ARG_STATUS}" = true ]; then
             # Check if results already exist locally
             PROJECT_ROOT="${SCRIPT_DIR}/.."
             LOCAL_RESULTS_DIR="${PROJECT_ROOT}/benchmarks/${JOB_NAME}"
-            RESULTS_JSONL=$(find "${LOCAL_RESULTS_DIR}" -name "profile_export.jsonl" -type f 2>/dev/null | head -1)
+            RESULTS_JSONL=$(find "${LOCAL_RESULTS_DIR}" -name "profile_export.jsonl" -type f 2>/dev/null | head -1 || true)
 
             if [ -z "${RESULTS_JSONL}" ]; then
                 echo ""
@@ -115,19 +118,21 @@ if [ "${ARG_STATUS}" = true ]; then
                         --region "${AWS_REGION}" --quiet
                     # Untar if output.tar.gz exists
                     tar_file=""
-                    tar_file=$(find "${LOCAL_RESULTS_DIR}" -name "output.tar.gz" -type f 2>/dev/null | head -1)
+                    tar_file=$(find "${LOCAL_RESULTS_DIR}" -name "output.tar.gz" -type f 2>/dev/null | head -1 || true)
                     if [ -n "${tar_file}" ]; then
-                        # Detect whether tar has a leading directory prefix
-                        _tar_first=""
-                        _tar_first=$(tar -tzf "${tar_file}" 2>/dev/null | head -1)
-                        if echo "${_tar_first}" | grep -qE '^[^/]+/$'; then
+                        # Detect whether ALL entries share a common leading directory prefix
+                        _tar_prefix_count=""
+                        _tar_prefix_count=$(tar -tzf "${tar_file}" 2>/dev/null | sed 's|/.*||' | sort -u | wc -l | tr -d ' ')
+                        _tar_first_dir=""
+                        _tar_first_dir=$(tar -tzf "${tar_file}" 2>/dev/null | head -1)
+                        if [ "${_tar_prefix_count}" = "1" ] && echo "${_tar_first_dir}" | grep -qE '^[^/]+/$'; then
                             tar -xzf "${tar_file}" --strip-components=1 -C "${LOCAL_RESULTS_DIR}/output/" 2>/dev/null || true
                         else
                             tar -xzf "${tar_file}" -C "${LOCAL_RESULTS_DIR}/output/" 2>/dev/null || true
                         fi
                     fi
                     # Re-search after extraction
-                    RESULTS_JSONL=$(find "${LOCAL_RESULTS_DIR}" -name "profile_export.jsonl" -type f 2>/dev/null | head -1)
+                    RESULTS_JSONL=$(find "${LOCAL_RESULTS_DIR}" -name "profile_export.jsonl" -type f 2>/dev/null | head -1 || true)
                     echo "   ✅ Results downloaded to: benchmarks/${JOB_NAME}/"
                 fi
             else
@@ -140,7 +145,7 @@ if [ "${ARG_STATUS}" = true ]; then
                 if [ -n "${RESULTS_JSONL}" ] && [ -f "${RESULTS_JSONL}" ]; then
                     _WRITER_INPUT="${RESULTS_JSONL}"
                 else
-                    _WRITER_INPUT=$(find "${LOCAL_RESULTS_DIR}" -name "profile_export_aiperf.json" -type f 2>/dev/null | head -1)
+                    _WRITER_INPUT=$(find "${LOCAL_RESULTS_DIR}" -name "profile_export_aiperf.json" -type f 2>/dev/null | head -1 || true)
                 fi
 
                 if [ -n "${_WRITER_INPUT}" ]; then
@@ -153,7 +158,8 @@ if [ "${ARG_STATUS}" = true ]; then
                         --workload "${BENCHMARK_WORKLOAD:-manual}" \
                         --concurrency "${BENCHMARK_CONCURRENCY:-2}" \
                         --bucket "${CI_BENCHMARK_RESULTS_BUCKET}" \
-                        --region "${AWS_REGION:-${REGION}}" \
+                        --region "${AWS_REGION:-${REGION:-us-east-1}}" \
+                        ${_STATUS_INSTANCE_TYPE:+--instance-type "${_STATUS_INSTANCE_TYPE}"} \
                         ${ADAPTER_ARG:+--adapter-name "${ADAPTER_ARG}"}; then
                         echo "   ✅ Results persisted to Athena"
                     else
@@ -559,6 +565,7 @@ print(f'Combined {n_metrics} concurrency level results')
                 --workload "${BENCHMARK_WORKLOAD:-manual}" \
                 --bucket "${CI_BENCHMARK_RESULTS_BUCKET}" \
                 --region "${AWS_REGION:-${REGION}}" \
+                ${RESOLVED_INSTANCE_TYPE:+--instance-type "${RESOLVED_INSTANCE_TYPE}"} \
                 ${ADAPTER_ARG:+--adapter-name "${ADAPTER_ARG}"}; then
                 echo "✅ Multi-level benchmark results persisted to S3"
             else
@@ -799,6 +806,55 @@ fi
 
 echo "✅ Endpoint is InService: ${ENDPOINT_NAME}"
 
+# ── Resolve actual instance type from endpoint ────────────────────────────────
+# For heterogeneous instance pools, INSTANCE_TYPE in do/config may not reflect
+# the actual provisioned instance. Query the endpoint to determine the real type.
+RESOLVED_INSTANCE_TYPE=""
+_EP_JSON=$(aws sagemaker describe-endpoint \
+    --endpoint-name "${ENDPOINT_NAME}" \
+    --region "${AWS_REGION}" \
+    --output json 2>/dev/null) || _EP_JSON=""
+
+if [ -n "${_EP_JSON}" ]; then
+    # Try InstanceType from the primary variant runtime response
+    RESOLVED_INSTANCE_TYPE=$(echo "${_EP_JSON}" | python3 -c "
+import sys, json
+try:
+    ep = json.load(sys.stdin)
+    variant = ep.get('ProductionVariants', [{}])[0]
+    it = variant.get('CurrentInstanceType') or variant.get('InstanceType') or ''
+    if it:
+        print(it)
+    else:
+        # Fall back to endpoint config for pool-based endpoints
+        print('')
+except:
+    print('')
+" 2>/dev/null) || RESOLVED_INSTANCE_TYPE=""
+
+    # If still empty, query endpoint config for InstancePools
+    if [ -z "${RESOLVED_INSTANCE_TYPE}" ]; then
+        _EC_NAME=$(echo "${_EP_JSON}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('EndpointConfigName',''))" 2>/dev/null) || _EC_NAME=""
+        if [ -n "${_EC_NAME}" ]; then
+            RESOLVED_INSTANCE_TYPE=$(aws sagemaker describe-endpoint-config \
+                --endpoint-config-name "${_EC_NAME}" \
+                --region "${AWS_REGION}" \
+                --query 'ProductionVariants[0].InstanceType' \
+                --output text 2>/dev/null) || RESOLVED_INSTANCE_TYPE=""
+            # 'None' is returned as literal text when the field is null
+            [ "${RESOLVED_INSTANCE_TYPE}" = "None" ] && RESOLVED_INSTANCE_TYPE=""
+        fi
+    fi
+fi
+
+# Final fallback: use INSTANCE_TYPE from do/config
+RESOLVED_INSTANCE_TYPE="${RESOLVED_INSTANCE_TYPE:-${INSTANCE_TYPE:-}}"
+
+# Persist to do/config for --status (endpoint may be gone by then)
+if [ -n "${RESOLVED_INSTANCE_TYPE}" ]; then
+    _update_benchmark_var "BENCHMARK_INSTANCE_TYPE" "${RESOLVED_INSTANCE_TYPE}"
+fi
+
 # ── Pre-flight check: Ensure S3 output bucket exists ──────────────────────────
 echo "🔍 Pre-flight: Checking S3 output bucket..."
 
@@ -1097,11 +1153,14 @@ if [ "${JOB_STATUS}" = "Completed" ]; then
             # Extract any tar.gz archives (benchmark service packages results as output.tar.gz)
             for ARCHIVE in $(find "${LOCAL_RESULTS_DIR}" -name "*.tar.gz" -type f 2>/dev/null); do
                 ARCHIVE_DIR=$(dirname "${ARCHIVE}")
-                # Detect whether tar has a leading directory prefix to strip.
-                # Some AIPerf versions wrap in output/, others are flat.
-                _TAR_FIRST=$(tar -tzf "${ARCHIVE}" 2>/dev/null | head -1)
-                if echo "${_TAR_FIRST}" | grep -qE '^[^/]+/$'; then
-                    # Leading directory (e.g., "output/") — strip it
+                # Detect whether ALL entries share a common leading directory prefix.
+                # Only strip if every entry starts with the same dir (e.g., "output/file1", "output/file2").
+                # A flat archive with mixed top-level files/dirs (e.g., "plots/", "profile_export.jsonl")
+                # must NOT be stripped.
+                _TAR_PREFIX=$(tar -tzf "${ARCHIVE}" 2>/dev/null | sed 's|/.*||' | sort -u | wc -l | tr -d ' ')
+                _TAR_FIRST_DIR=$(tar -tzf "${ARCHIVE}" 2>/dev/null | head -1)
+                if [ "${_TAR_PREFIX}" = "1" ] && echo "${_TAR_FIRST_DIR}" | grep -qE '^[^/]+/$'; then
+                    # Single common leading directory (e.g., all under "output/") — strip it
                     tar -xzf "${ARCHIVE}" --strip-components=1 -C "${ARCHIVE_DIR}" 2>/dev/null || true
                 else
                     # Flat archive — extract as-is
@@ -1366,6 +1425,7 @@ except Exception as e:
             --concurrency "${BENCHMARK_CONCURRENCY}" \
             --bucket "${CI_BENCHMARK_RESULTS_BUCKET}" \
             --region "${AWS_REGION:-${REGION}}" \
+            ${RESOLVED_INSTANCE_TYPE:+--instance-type "${RESOLVED_INSTANCE_TYPE}"} \
             ${ADAPTER_ARG:+--adapter-name "${ADAPTER_ARG}"}; then
             echo "✅ Benchmark results persisted to S3"
         else

package/templates/do/config

@@ -46,7 +46,7 @@ export INSTANCE_TYPE="<%= instanceType %>"
 <% if (typeof instancePools !== 'undefined' && instancePools && instancePools.length > 1) { %>
 # Instance pools: heterogeneous instance types with priority-based fallback
 # Priority = selection order (1 = preferred, higher = fallback)
-export INSTANCE_POOLS='<%= JSON.stringify(instancePools) %>'
+export INSTANCE_POOLS='<%- JSON.stringify(instancePools) %>'
 <% } else { %>
 # Instance pools: heterogeneous instance types with priority-based fallback (uncomment to enable)
 # Format: [{"InstanceType":"ml.g6e.48xlarge","Priority":1},{"InstanceType":"ml.g5.48xlarge","Priority":2}]

package/templates/do/lib/inference-component.sh

@@ -112,31 +112,12 @@ create_inference_component() {
 
     # Build specification JSON — multi-spec (Specifications array) or single (Specification object)
     local spec_json
-    if [ "${IC_MULTI_SPEC:-false}" = "true" ] && [ "${IC_SPEC_COUNT:-0}" -gt 0 ]; then
-        # Multi-spec: build Specifications array with per-instance-type compute resources
-        spec_json="{\"Specifications\":["
-        local i=1
-        while [ "${i}" -le "${IC_SPEC_COUNT}" ]; do
-            local spec_instance_type_var="IC_SPEC_${i}_INSTANCE_TYPE"
-            local spec_gpu_count_var="IC_SPEC_${i}_GPU_COUNT"
-            local spec_min_memory_var="IC_SPEC_${i}_MIN_MEMORY_MB"
-
-            local spec_instance_type="${!spec_instance_type_var}"
-            local spec_gpu_count="${!spec_gpu_count_var:-1}"
-            local spec_min_memory="${!spec_min_memory_var:-1024}"
-
-            if [ "${i}" -gt 1 ]; then
-                spec_json="${spec_json},"
-            fi
-            spec_json="${spec_json}{\"Container\":${container_spec},\"StartupParameters\":{\"ContainerStartupHealthCheckTimeoutInSeconds\":${IC_STARTUP_TIMEOUT:-900}},\"ComputeResourceRequirements\":{\"NumberOfAcceleratorDevicesRequired\":${spec_gpu_count},\"MinMemoryRequiredInMb\":${spec_min_memory}}}"
-
-            i=$((i + 1))
-        done
-        spec_json="${spec_json}]}"
-    else
-        # Single spec: standard Specification object (existing behavior)
-        spec_json="{\"Container\":${container_spec},\"StartupParameters\":{\"ContainerStartupHealthCheckTimeoutInSeconds\":${IC_STARTUP_TIMEOUT:-900}},\"ComputeResourceRequirements\":{\"NumberOfAcceleratorDevicesRequired\":${IC_GPU_COUNT:-1},\"MinMemoryRequiredInMb\":${IC_MIN_MEMORY_MB:-1024}}}"
-    fi
+    # Always use singular Specification. For heterogeneous instance pools, the IC
+    # declares its minimum resource requirements and SageMaker places it on whatever
+    # instance was provisioned from the pool. Multi-spec (Specifications plural) is
+    # only needed when you want different configurations per instance type (e.g.,
+    # different TP, different model artifact) — a future optimization.
+    spec_json="{\"Container\":${container_spec},\"StartupParameters\":{\"ContainerStartupHealthCheckTimeoutInSeconds\":${IC_STARTUP_TIMEOUT:-900}},\"ComputeResourceRequirements\":{\"NumberOfAcceleratorDevicesRequired\":${IC_GPU_COUNT:-1},\"MinMemoryRequiredInMb\":${IC_MIN_MEMORY_MB:-1024}}}"
 
     echo "📦 Creating inference component: ${ic_name}"
     if ! aws sagemaker create-inference-component \

package/templates/do/register

@@ -41,6 +41,7 @@ _show_usage() {
     echo "  --technique <tech>  Technique: sft, dpo, rlaif, rlvr (default: sft)"
     echo "  --row-count <n>     Number of records"
     echo "  --column-schema <j> Column schema as JSON string"
+    echo "  --force             Force new version even if content is unchanged"
     echo ""
     echo "Evaluator options:"
     echo "  <name>              Evaluator name (required, positional)"
@@ -89,6 +90,7 @@ DATASET_FORMAT="jsonl"
 DATASET_TECHNIQUE="sft"
 DATASET_ROW_COUNT=""
 DATASET_COLUMN_SCHEMA=""
+DATASET_FORCE=false
 EVALUATOR_NAME=""
 EVALUATOR_TYPE=""
 EVALUATOR_ARN_OR_URI=""
@@ -135,6 +137,7 @@ if [ "${SUBCOMMAND}" = "dataset" ]; then
             --dataset-technique) DATASET_TECHNIQUE="$2"; shift 2 ;;
             --dataset-row-count) DATASET_ROW_COUNT="$2"; shift 2 ;;
             --dataset-column-schema) DATASET_COLUMN_SCHEMA="$2"; shift 2 ;;
+            --force)        DATASET_FORCE=true; shift ;;
             --help|-h) _show_usage; exit 0 ;;
             *) echo "⚠️  Unknown dataset option: $1"; _show_usage; exit 1 ;;
         esac
@@ -196,7 +199,10 @@ if [ "${SUBCOMMAND}" = "dataset" ]; then
                 _slug=$(basename "${_source}" .jsonl)
             fi
             _slug=$(echo "${_slug}" | tr '[:upper:]' '[:lower:]' | sed 's/[^a-z0-9]/-/g' | sed 's/--*/-/g' | sed 's/^-//' | sed 's/-$//')
-            DATASET_NAME="${_slug:-dataset}-${TUNE_TECHNIQUE:-sft}-${TUNE_DATASET_ROW_COUNT:-0}"
+            # Salt with 4-char hash of S3 URI to prevent slug conflicts for
+            # same repo with different preprocessing/split/technique
+            _salt=$(echo "${DATASET_S3_URI}" | shasum | cut -c1-4)
+            DATASET_NAME="${_slug:-dataset}-${_salt}"
         fi
     fi
 
@@ -427,6 +433,11 @@ if [ "${SUBCOMMAND}" = "dataset" ]; then
     fi
 
     DS_ARGS+=("--project-name" "${PROJECT_NAME}")
+    DS_ARGS+=("--region" "${AWS_REGION}")
+
+    if [ "${DATASET_FORCE}" = true ]; then
+        DS_ARGS+=("--force")
+    fi
 
     # Call .register_helper.py register-dataset
     if ds_output=$(python3 "${SCRIPT_DIR}/.register_helper.py" "${DS_ARGS[@]}" 2>/dev/null); then
@@ -1288,9 +1299,10 @@ elif [ -n "${MODEL_PKG_ARN:-}" ] && [ -d "${SCRIPT_DIR}/adapters" ]; then
         ADAPTER_TECHNIQUE=""
         eval "$(grep '^export ADAPTER_WEIGHTS_URI=' "${conf}" 2>/dev/null)" 2>/dev/null || true
         eval "$(grep '^export ADAPTER_TECHNIQUE=' "${conf}" 2>/dev/null)" 2>/dev/null || true
+        eval "$(grep '^export ADAPTER_TUNE_TECHNIQUE=' "${conf}" 2>/dev/null)" 2>/dev/null || true
 
         _ADAPTER_DATA_URL="${ADAPTER_WEIGHTS_URI:-}"
-        _ADAPTER_TECHNIQUE="${ADAPTER_TECHNIQUE:-${TUNE_TECHNIQUE:-}}"
+        _ADAPTER_TECHNIQUE="${ADAPTER_TECHNIQUE:-${ADAPTER_TUNE_TECHNIQUE:-${TUNE_TECHNIQUE:-}}}"
 
         echo ""
         echo "📦 Registering adapter: ${_ADAPTER_NAME}"
@@ -1316,6 +1328,21 @@ elif [ -n "${MODEL_PKG_ARN:-}" ] && [ -d "${SCRIPT_DIR}/adapters" ]; then
         [ -n "${AWS_REGION:-}" ] && ADAPTER_REG_ARGS+=("--region" "${AWS_REGION}")
         [ -n "${ROLE_ARN:-}" ] && ADAPTER_REG_ARGS+=("--role-arn" "${ROLE_ARN}")
 
+        # Include dataset lineage from tune state (AC-2.7: version reference for reproducibility)
+        if [ -n "${TUNE_DATASET_S3_URI:-}" ]; then
+            ADAPTER_REG_ARGS+=("--dataset-s3-uri" "${TUNE_DATASET_S3_URI}")
+        fi
+        # Resolve dataset version from technique-specific or generic config var
+        _ds_version=""
+        if [ -n "${_ADAPTER_TECHNIQUE}" ]; then
+            _ds_ver_var="TUNE_DATASET_VERSION_$(echo "${_ADAPTER_TECHNIQUE}" | tr '[:lower:]' '[:upper:]')"
+            _ds_version="${!_ds_ver_var:-}"
+        fi
+        [ -z "${_ds_version}" ] && _ds_version="${TUNE_DATASET_VERSION:-}"
+        if [ -n "${_ds_version}" ]; then
+            ADAPTER_REG_ARGS+=("--dataset-version" "${_ds_version}")
+        fi
+
         # Call .register_helper.py register-adapter — non-fatal on failure
         if adapter_output=$(python3 "${SCRIPT_DIR}/.register_helper.py" "${ADAPTER_REG_ARGS[@]}" 2>/dev/null); then
             adapter_json=$(echo "${adapter_output}" | grep -E '^\{' | tail -1)

package/templates/do/tune

@@ -52,7 +52,9 @@ ARG_COLUMN_MAP=""
 ARG_TAKE=""
 ARG_ACCEPT_EULA=false
 ARG_DATASET_NAME=""
+ARG_DATASET_VERSION=""
 ARG_EVALUATOR_NAME=""
+ARG_NO_REGISTER=false
 
 
 # ── _parse_args() ─────────────────────────────────────────────────────────────
@@ -147,6 +149,7 @@ _parse_args() {
             --force) ARG_FORCE=true; shift ;;
             --accept-eula) ARG_ACCEPT_EULA=true; shift ;;
             --no-wait) ARG_NO_WAIT=true; shift ;;
+            --no-register) ARG_NO_REGISTER=true; shift ;;
             --status) ARG_STATUS=true; shift ;;
             --help|-h) ARG_HELP=true; shift ;;
             --dry-run) ARG_DRY_RUN=true; shift ;;
@@ -273,6 +276,8 @@ _show_help() {
     echo "  --force               Force new job even if one exists for this technique"
     echo "  --accept-eula         Accept model EULA (required for gated models like Llama)"
     echo "  --no-wait             Submit and exit without polling for completion"
+    echo "  --no-register         Skip auto-stage and auto-register after completion"
+    echo "                        (prints next-step commands instead)"
     echo "  --status              Show status of all tracked tune jobs"
     echo ""
     echo "Dataset options:"
@@ -791,13 +796,26 @@ else:
 _validate_dataset() {
     local dataset="${ARG_DATASET}"
 
+    # ── Parse @v<N> version suffix (AC-2.1, AC-2.3) ──────────────────────────
+    # Syntax: dataset-name@v2 → name="dataset-name", version ordinal=2
+    if [[ "${dataset}" =~ ^(.+)@v([0-9]+)$ ]]; then
+        ARG_DATASET_NAME="${BASH_REMATCH[1]}"
+        ARG_DATASET_VERSION="${BASH_REMATCH[2]}"
+        dataset=""  # Clear so name-based resolution takes over
+    fi
+
     # If --dataset-name is set, resolve from registry (AC-2b.4)
     # --dataset-name takes precedence over --dataset for named registry lookup
     if [ -n "${ARG_DATASET_NAME}" ]; then
         echo "🔍 Resolving dataset '${ARG_DATASET_NAME}' from registry..."
+        local resolve_args=("--name" "${ARG_DATASET_NAME}")
+        if [ -n "${ARG_DATASET_VERSION}" ]; then
+            resolve_args+=("--version" "${ARG_DATASET_VERSION}")
+            echo "   Version: v${ARG_DATASET_VERSION}"
+        fi
         local resolve_result
         resolve_result=$(python3 "${SCRIPT_DIR}/.register_helper.py" resolve-dataset \
-            --name "${ARG_DATASET_NAME}" 2>/dev/null) || resolve_result=""
+            "${resolve_args[@]}" 2>/dev/null) || resolve_result=""
 
         if [ -n "${resolve_result}" ]; then
             local resolved_uri
@@ -830,9 +848,14 @@ _validate_dataset() {
         if [ -n "${ARG_DATASET_NAME}" ]; then
             # Name-based resolution happens below via resolve-dataset
             echo "🔍 Resolving dataset '${ARG_DATASET_NAME}' from registry..."
+            local resolve_args=("--name" "${ARG_DATASET_NAME}")
+            if [ -n "${ARG_DATASET_VERSION}" ]; then
+                resolve_args+=("--version" "${ARG_DATASET_VERSION}")
+                echo "   Version: v${ARG_DATASET_VERSION}"
+            fi
             local resolve_result
             resolve_result=$(python3 "${SCRIPT_DIR}/.register_helper.py" resolve-dataset \
-                --name "${ARG_DATASET_NAME}" 2>/dev/null) || resolve_result=""
+                "${resolve_args[@]}" 2>/dev/null) || resolve_result=""
 
             if [ -n "${resolve_result}" ]; then
                 local resolved_uri
@@ -1330,6 +1353,10 @@ print(entry.get('provider', ''))
     _update_config_var "TUNE_DATASET_S3_URI_${technique_upper}" "${RESOLVED_DATASET_S3_URI:-}"
     _update_config_var "TUNE_DATASET_ROW_COUNT_${technique_upper}" "${RESOLVED_DATASET_ROW_COUNT:-0}"
     _update_config_var "TUNE_DATASET_SOURCE_${technique_upper}" "${ARG_DATASET}"
+    # Store dataset version ordinal if pinned (AC-2.6)
+    if [ -n "${ARG_DATASET_VERSION}" ]; then
+        _update_config_var "TUNE_DATASET_VERSION_${technique_upper}" "${ARG_DATASET_VERSION}"
+    fi
 }
 
 
@@ -1528,12 +1555,59 @@ _handle_completion() {
     _update_config_var "TUNE_OUTPUT_PATH_LATEST" "${artifact_path}"
     _update_config_var "TUNE_OUTPUT_TYPE_LATEST" "${output_type}"
 
-    # Print next-step commands
-    echo "📋 Next steps:"
-    echo ""
-    if [ "${output_type}" = "adapter" ]; then
+    # Auto-register or print next-step commands
+    if [ "${output_type}" = "adapter" ] && [ "${ARG_NO_REGISTER}" != true ]; then
+        # Auto-register: stage adapter and register in deployment MPG
+        local dataset_slug
+        dataset_slug=$(_derive_dataset_slug "${ARG_DATASET:-}")
+        local adapter_name="tuned-${ARG_TECHNIQUE}-${dataset_slug}"
+        if [ -z "${dataset_slug}" ]; then
+            adapter_name="tuned-${ARG_TECHNIQUE}"
+        fi
+
+        echo "🔄 Auto-registering adapter: ${adapter_name}"
+        echo ""
+
+        # Step 1: Stage the adapter via do/adapter add
+        local adapter_add_output
+        if adapter_add_output=$("${SCRIPT_DIR}/adapter" add "${adapter_name}" --from-tune "${ARG_TECHNIQUE}" 2>&1); then
+            echo "   ✅ Adapter staged: ${adapter_name}"
+
+            # Step 2: Register in deployment MPG via do/register
+            local register_output
+            if register_output=$("${SCRIPT_DIR}/register" 2>&1); then
+                echo "   ✅ Registration complete"
+
+                # Step 3: Extract adapter deployment ARN from register output
+                local adapter_deploy_arn
+                adapter_deploy_arn=$(echo "${register_output}" | grep "${adapter_name}" | grep -E '^\{' | tail -1 | jq -r '.model_package_arn' 2>/dev/null) || adapter_deploy_arn=""
+
+                if [ -n "${adapter_deploy_arn}" ] && [ "${adapter_deploy_arn}" != "null" ]; then
+                    _update_config_var "TUNE_ADAPTER_DEPLOY_ARN_${technique_upper}" "${adapter_deploy_arn}"
+                    echo "   ✅ Deployment ARN stored: ${adapter_deploy_arn}"
+                else
+                    echo "   ⚠️  Could not extract adapter deployment ARN from register output"
+                    echo "      (adapter was staged and registered — ARN can be found via do/register --status)"
+                fi
+                echo ""
+            else
+                echo "   ⚠️  Registration failed (adapter was staged successfully)"
+                echo "      Run manually: ./do/register"
+                echo ""
+            fi
+        else
+            echo "   ⚠️  Adapter staging failed"
+            echo "      Run manually:"
+            echo "        ./do/adapter add ${adapter_name} --from-tune ${ARG_TECHNIQUE}"
+            echo "        ./do/register"
+            echo ""
+        fi
+    elif [ "${output_type}" = "adapter" ]; then
+        # --no-register: print next steps as before
         local dataset_slug
         dataset_slug=$(_derive_dataset_slug "${ARG_DATASET:-}")
+        echo "📋 Next steps:"
+        echo ""
         echo "   Deploy as LoRA adapter:"
         echo "     ./do/adapter add tuned-${ARG_TECHNIQUE} --from-tune"
         echo "     ./do/adapter add tuned-${ARG_TECHNIQUE} --from-tune ${ARG_TECHNIQUE}"
@@ -1541,13 +1615,16 @@ _handle_completion() {
             echo "     ./do/adapter add tuned-${ARG_TECHNIQUE}-${dataset_slug} --from-tune ${ARG_TECHNIQUE}-${dataset_slug}"
         fi
         echo "     ./do/adapter add tuned-${ARG_TECHNIQUE} --weights ${artifact_path}"
+        echo ""
     else
+        echo "📋 Next steps:"
+        echo ""
         echo "   Deploy as new inference component:"
         echo "     ./do/add-ic tuned-v1 --from-tune"
         echo "     ./do/add-ic tuned-v1 --model-data ${artifact_path}"
         echo "     ./do/deploy --force-ic --model-data ${artifact_path}"
+        echo ""
     fi
-    echo ""
 }
 
 
@@ -1648,17 +1725,21 @@ if [ "${ARG_LIST_DATASETS}" = true ]; then
     if [ -n "${_ds_json}" ]; then
         _ds_count=$(echo "${_ds_json}" | python3 -c "import sys,json; print(len(json.load(sys.stdin).get('datasets',[])))" 2>/dev/null) || _ds_count=0
         if [ "${_ds_count}" -gt 0 ]; then
-            printf "  %-25s %-10s %-8s %s\n" "NAME" "TECHNIQUE" "ROWS" "S3 URI"
-            printf "  %-25s %-10s %-8s %s\n" "----" "---------" "----" "------"
+            printf "  %-20s %-10s %-10s %-8s %s\n" "NAME" "TECHNIQUE" "LATEST" "ROWS" "S3 URI"
+            printf "  %-20s %-10s %-10s %-8s %s\n" "----" "---------" "------" "----" "------"
             echo "${_ds_json}" | python3 -c "
 import sys, json
 data = json.load(sys.stdin)
 for ds in data.get('datasets', []):
-    name = ds.get('name','')[:25]
+    name = ds.get('name','')[:20]
     tech = ds.get('technique','')[:10]
-    rows = str(ds.get('row_count',''))[:8]
+    latest = ds.get('latest_version','')[:10]
+    ver_count = ds.get('version_count', 1)
+    if ver_count > 1:
+        latest = f'{latest} ({ver_count}v)'
+    rows = str(ds.get('row_count','') or '')[:8]
     uri = ds.get('s3_uri','')
-    print(f'  {name:<25} {tech:<10} {rows:<8} {uri}')
+    print(f'  {name:<20} {tech:<10} {latest:<10} {rows:<8} {uri}')
 " 2>/dev/null
         else
             echo "  (none registered)"
@@ -1669,6 +1750,7 @@ for ds in data.get('datasets', []):
     echo ""
     echo "  Register: ./do/register dataset <name> --s3-uri <uri> --technique <sft|dpo>"
     echo "  Use:      ./do/tune --technique sft --dataset <name>"
+    echo "  Versions: python3 .register_helper.py list-dataset-versions --name <name>"
     echo ""
     exit 0
 fi