@aws/ml-container-creator 0.13.4 → 0.15.0

package/templates/do/tune

@@ -46,9 +46,13 @@ ARG_DRY_RUN=false
 ARG_LIST_MODELS=false
 ARG_NO_STALE_WARNING=false
 ARG_DISCOVER=false
+ARG_LIST_DATASETS=false
 ARG_DISCOVER_FILTER=""
 ARG_COLUMN_MAP=""
+ARG_TAKE=""
 ARG_ACCEPT_EULA=false
+ARG_DATASET_NAME=""
+ARG_EVALUATOR_NAME=""
 
 
 # ── _parse_args() ─────────────────────────────────────────────────────────────
@@ -147,6 +151,7 @@ _parse_args() {
             --help|-h) ARG_HELP=true; shift ;;
             --dry-run) ARG_DRY_RUN=true; shift ;;
             --list-models) ARG_LIST_MODELS=true; shift ;;
+            --list-datasets) ARG_LIST_DATASETS=true; shift ;;
             --no-stale-warning) ARG_NO_STALE_WARNING=true; shift ;;
             --column-map)
                 if [ -z "${2:-}" ]; then
@@ -162,6 +167,24 @@ _parse_args() {
                     shift
                 fi
                 ;;
+            --take)
+                if [ -z "${2:-}" ]; then
+                    echo "❌ --take requires an integer value"
+                    exit 1
+                fi
+                ARG_TAKE="$2"; shift 2 ;;
+            --dataset-name)
+                if [ -z "${2:-}" ]; then
+                    echo "❌ --dataset-name requires a registered dataset name"
+                    exit 1
+                fi
+                ARG_DATASET_NAME="$2"; shift 2 ;;
+            --evaluator-name)
+                if [ -z "${2:-}" ]; then
+                    echo "❌ --evaluator-name requires a registered evaluator name"
+                    exit 1
+                fi
+                ARG_EVALUATOR_NAME="$2"; shift 2 ;;
             *)
                 echo "❌ Unknown option: $1"
                 echo "   Run ./do/tune --help for usage."
@@ -218,7 +241,8 @@ _show_help() {
     echo ""
     echo "Required:"
     echo "  --technique <t>       Customization technique: sft, dpo, rlaif, rlvr"
-    echo "  --dataset <source>    Dataset: s3://bucket/path.jsonl or hf://org/name[/split]"
+    echo "  --dataset <source>    Dataset: s3://bucket/path.jsonl or hf://org/name[/split][?file=pattern]"
+    echo "                        ⚠️  Quote the URI if it contains ? or * to prevent shell expansion"
     echo ""
     echo "Model selection:"
     echo "  --model <id>          JumpStart Hub content name to use for fine-tuning."
@@ -251,6 +275,13 @@ _show_help() {
     echo "  --no-wait             Submit and exit without polling for completion"
     echo "  --status              Show status of all tracked tune jobs"
     echo ""
+    echo "Dataset options:"
+    echo "  --column-map <map>    Rename columns (e.g., prompt=question,completion=answer)"
+    echo "  --take <n>            Take only the first N records from the dataset"
+    echo ""
+    echo "  Note: Always quote --dataset values containing ? or * characters."
+    echo "  Unquoted, bash may interpret ? as a glob and * as a wildcard expansion."
+    echo ""
     echo "Discovery and diagnostics:"
     echo "  --discover [filter]   Query JumpStart Hub for tune-eligible models."
     echo "                        Without a filter, shows models for the current family."
@@ -276,6 +307,12 @@ _show_help() {
     echo "  # Fine-tune a gated model (Meta Llama) — requires EULA acceptance:"
     echo "  ./do/tune --technique dpo --dataset hf://argilla/ultrafeedback-binarized-preferences-cleaned --accept-eula"
     echo ""
+    echo "  # Take only 500 records for a quick test run:"
+    echo "  ./do/tune --technique sft --dataset hf://timdettmers/openassistant-guanaco --take 500"
+    echo ""
+    echo "  # Filter files in a multi-file dataset (quote to prevent shell glob expansion):"
+    echo "  ./do/tune --technique sft --dataset \"hf://Open-Orca/OpenOrca?file=1M-GPT4\""
+    echo ""
     echo "  # Discover available models:"
     echo "  ./do/tune --discover                    # Models for current family"
     echo "  ./do/tune --discover qwen               # Filter by keyword"
@@ -329,13 +366,47 @@ _show_status() {
                 echo "     Elapsed: ${mins}m ${secs}s"
             fi
 
-            # Show output path if completed
+            # Show output path if completed — resolve if not yet set
             local output_var="TUNE_ADAPTER_PATH_$(echo "${technique}" | tr '[:lower:]' '[:upper:]')"
             local model_var="TUNE_MODEL_PATH_$(echo "${technique}" | tr '[:lower:]' '[:upper:]')"
             if [ -n "${!output_var:-}" ]; then
                 echo "     Output (adapter): ${!output_var}"
             elif [ -n "${!model_var:-}" ]; then
                 echo "     Output (model): ${!model_var}"
+            elif [ "${status}" = "Completed" ]; then
+                # Job is complete but output path not set — resolve now
+                echo "     🔄 Resolving artifacts..."
+                local training_type="${TUNE_TRAINING_TYPE:-lora}"
+                local resolve_result
+                resolve_result=$(python3 "${HELPER_SCRIPT}" resolve \
+                    --job-name "${job_name}" \
+                    --region "${AWS_REGION}" \
+                    --training-type "${training_type}" \
+                    --model-package-group "${PROJECT_NAME}-tune-models" 2>/dev/null) || resolve_result=""
+
+                if [ -n "${resolve_result}" ]; then
+                    local artifact_path
+                    artifact_path=$(echo "${resolve_result}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('artifact_path',''))" 2>/dev/null) || artifact_path=""
+                    local output_type
+                    output_type=$(echo "${resolve_result}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('output_type',''))" 2>/dev/null) || output_type=""
+
+                    if [ -n "${artifact_path}" ]; then
+                        local technique_upper
+                        technique_upper=$(echo "${technique}" | tr '[:lower:]' '[:upper:]')
+
+                        # Update config
+                        if [ "${output_type}" = "adapter" ]; then
+                            _update_config_var "TUNE_ADAPTER_PATH_${technique_upper}" "${artifact_path}"
+                            echo "     Output (adapter): ${artifact_path}"
+                        else
+                            _update_config_var "TUNE_MODEL_PATH_${technique_upper}" "${artifact_path}"
+                            echo "     Output (model): ${artifact_path}"
+                        fi
+                        _update_config_var "TUNE_OUTPUT_PATH_LATEST" "${artifact_path}"
+                        _update_config_var "TUNE_OUTPUT_TYPE_LATEST" "${output_type}"
+                        echo "     ✅ Updated do/config with output paths"
+                    fi
+                fi
             fi
             echo ""
         fi
@@ -720,10 +791,51 @@ else:
 _validate_dataset() {
     local dataset="${ARG_DATASET}"
 
+    # 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_result
+        resolve_result=$(python3 "${SCRIPT_DIR}/.register_helper.py" resolve-dataset \
+            --name "${ARG_DATASET_NAME}" 2>/dev/null) || resolve_result=""
+
+        if [ -n "${resolve_result}" ]; then
+            local resolved_uri
+            resolved_uri=$(echo "${resolve_result}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('s3_uri',''))" 2>/dev/null) || resolved_uri=""
+            if [ -n "${resolved_uri}" ]; then
+                echo "   Resolved to: ${resolved_uri}"
+                dataset="${resolved_uri}"
+                ARG_DATASET="${resolved_uri}"
+            else
+                echo "❌ Dataset '${ARG_DATASET_NAME}' not found in registry"
+                echo "   Register it first: ./do/register --dataset --dataset-name ${ARG_DATASET_NAME} --dataset-s3-uri s3://..."
+                exit 1
+            fi
+        else
+            echo "❌ Failed to resolve dataset '${ARG_DATASET_NAME}' from registry"
+            echo "   Register it first: ./do/register --dataset --dataset-name ${ARG_DATASET_NAME} --dataset-s3-uri s3://..."
+            exit 1
+        fi
+    fi
+
+    # If --dataset value is not an S3 URI or HF reference, treat as a registry name
+    if [ -z "${ARG_DATASET_NAME}" ] && [ -n "${dataset}" ] && \
+       [[ "${dataset}" != s3://* ]] && [[ "${dataset}" != hf://* ]]; then
+        # Looks like a name — resolve from registry
+        ARG_DATASET_NAME="${dataset}"
+        dataset=""
+    fi
+
     if [ -z "${dataset}" ]; then
-        echo "❌ --dataset is required"
-        echo "   Provide an S3 URI (s3://bucket/path.jsonl) or HF reference (hf://org/name)"
-        exit 1
+        if [ -n "${ARG_DATASET_NAME}" ]; then
+            # Name-based resolution happens below via resolve-dataset
+            :
+        else
+            echo "❌ --dataset is required"
+            echo "   Provide an S3 URI (s3://bucket/path.jsonl), HF reference (hf://org/name), or registered name"
+            echo "   Run ./do/tune --list-datasets to see available registered datasets."
+            exit 1
+        fi
     fi
 
     # Determine dataset type
@@ -826,6 +938,9 @@ _validate_dataset() {
             stage_args+=(--column-map "${ARG_COLUMN_MAP}")
         fi
         stage_args+=(--technique "${ARG_TECHNIQUE}")
+        if [ -n "${ARG_TAKE}" ]; then
+            stage_args+=(--take "${ARG_TAKE}")
+        fi
         if [ -n "${hf_file}" ]; then
             stage_args+=(--hf-file "${hf_file}")
         fi
@@ -850,16 +965,16 @@ _validate_dataset() {
         fi
 
         RESOLVED_DATASET_S3_URI=$(echo "${stage_result}" | python3 -c "import sys,json; print(json.load(sys.stdin)['s3_uri'])" 2>/dev/null)
-        local num_records
-        num_records=$(echo "${stage_result}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('num_records',0))" 2>/dev/null) || num_records="0"
+        RESOLVED_DATASET_ROW_COUNT=$(echo "${stage_result}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('num_records',0))" 2>/dev/null) || RESOLVED_DATASET_ROW_COUNT="0"
 
         echo "   ✅ Staged to: ${RESOLVED_DATASET_S3_URI}"
-        echo "   Records: ${num_records}"
+        echo "   Records: ${RESOLVED_DATASET_ROW_COUNT}"
         echo ""
 
     else
         echo "❌ Invalid dataset format: ${dataset}"
-        echo "   Expected: s3://bucket/path.jsonl or hf://org/name[/split]"
+        echo "   Expected: s3://bucket/path.jsonl or hf://org/name[/split][?file=pattern]"
+        echo "   Hint: Quote the value if it contains ? or * (e.g., \"hf://org/name?file=pattern\")"
         exit 1
     fi
 }
@@ -1067,6 +1182,37 @@ print(entry.get('provider', ''))
     if [ -n "${ARG_REWARD_PROMPT}" ]; then
         submit_args+=(--reward-prompt "${ARG_REWARD_PROMPT}")
     fi
+
+    # Resolve evaluator from registry if --evaluator-name is set (AC-2c.3, AC-2c.4)
+    if [ -n "${ARG_EVALUATOR_NAME}" ] && [ -z "${ARG_REWARD_FUNCTION}" ] && [ -z "${ARG_REWARD_PROMPT}" ]; then
+        echo "🔍 Resolving evaluator '${ARG_EVALUATOR_NAME}' from registry..."
+        local ev_resolve_result
+        ev_resolve_result=$(python3 "${SCRIPT_DIR}/.register_helper.py" resolve-evaluator \
+            --name "${ARG_EVALUATOR_NAME}" 2>/dev/null) || ev_resolve_result=""
+
+        if [ -n "${ev_resolve_result}" ]; then
+            local ev_type ev_arn_or_uri
+            ev_type=$(echo "${ev_resolve_result}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('type',''))" 2>/dev/null) || ev_type=""
+            ev_arn_or_uri=$(echo "${ev_resolve_result}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('arn_or_uri',''))" 2>/dev/null) || ev_arn_or_uri=""
+
+            if [ -n "${ev_arn_or_uri}" ]; then
+                echo "   Resolved evaluator: ${ev_type} → ${ev_arn_or_uri}"
+                if [ "${ev_type}" = "lambda" ]; then
+                    submit_args+=(--reward-function "${ev_arn_or_uri}")
+                else
+                    submit_args+=(--reward-prompt "${ev_arn_or_uri}")
+                fi
+            else
+                echo "⚠️  Evaluator '${ARG_EVALUATOR_NAME}' not found in registry"
+                echo "   Register it first: ./do/register --evaluator --evaluator-name ${ARG_EVALUATOR_NAME} ..."
+                exit 1
+            fi
+        else
+            echo "⚠️  Failed to resolve evaluator '${ARG_EVALUATOR_NAME}' from registry"
+            echo "   Register it first: ./do/register --evaluator --evaluator-name ${ARG_EVALUATOR_NAME} ..."
+            exit 1
+        fi
+    fi
     if [ "${ARG_ACCEPT_EULA}" = true ]; then
         submit_args+=(--accept-eula)
     fi
@@ -1158,6 +1304,14 @@ print(entry.get('provider', ''))
     _update_config_var "TUNE_TECHNIQUE" "${ARG_TECHNIQUE}"
     _update_config_var "TUNE_TRAINING_TYPE" "${ARG_TRAINING_TYPE}"
     _update_config_var "TUNE_DATASET_PATH" "${ARG_DATASET}"
+    _update_config_var "TUNE_DATASET_S3_URI" "${RESOLVED_DATASET_S3_URI:-}"
+    _update_config_var "TUNE_DATASET_ROW_COUNT" "${RESOLVED_DATASET_ROW_COUNT:-0}"
+    _update_config_var "TUNE_DATASET_SOURCE" "${ARG_DATASET}"
+    _update_config_var "TUNE_LAST_JOB_NAME" "${JOB_NAME}"
+    # Technique-specific dataset vars (allows --from-tune <technique> resolution)
+    _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}"
 }
 
 
@@ -1229,12 +1383,60 @@ _handle_interrupt() {
     echo ""
     echo "⚠️  Interrupted — job continues running in background"
     echo "   Job: ${job_name}"
-    echo ""
-    echo "   Resume monitoring: ./do/tune --technique ${ARG_TECHNIQUE} --dataset ${ARG_DATASET}"
-    echo "   Check status:      ./do/tune --status"
+    echo "   Check status: ./do/tune --status"
     exit 130
 }
 
+# ── _derive_dataset_slug() ────────────────────────────────────────────────────
+# Derive a short slug from the dataset argument for config variable naming.
+# Rules: lowercase, strip non-alphanumeric (keep hyphens), truncate to 20 chars,
+# collapse consecutive hyphens, strip leading/trailing hyphens.
+#
+# Examples:
+#   hf://tatsu-lab/alpaca           → alpaca
+#   hf://Open-Orca/OpenOrca         → openorca
+#   s3://bucket/path/train.jsonl    → train
+#   s3://bucket/path/file.parquet   → file
+_derive_dataset_slug() {
+    local dataset="${1:-}"
+    local slug=""
+
+    if [ -z "${dataset}" ]; then
+        echo ""
+        return
+    fi
+
+    if [[ "${dataset}" == hf://* ]]; then
+        # HuggingFace format: hf://org/name or hf://org/name?file=...
+        # Strip query params
+        local hf_path="${dataset#hf://}"
+        hf_path="${hf_path%%\?*}"
+        # Take the last component (dataset name, not org)
+        slug="${hf_path##*/}"
+    elif [[ "${dataset}" == s3://* ]]; then
+        # S3 format: s3://bucket/path/file.ext → slug from filename without extension
+        local filename="${dataset##*/}"
+        slug="${filename%.*}"
+    else
+        # Fallback: use the last path component without extension
+        local filename="${dataset##*/}"
+        slug="${filename%.*}"
+    fi
+
+    # Slugification: lowercase, strip non-alphanumeric (keep hyphens)
+    slug=$(echo "${slug}" | tr '[:upper:]' '[:lower:]' | sed 's/[^a-z0-9-]//g')
+    # Collapse consecutive hyphens
+    slug=$(echo "${slug}" | sed 's/-\{2,\}/-/g')
+    # Strip leading/trailing hyphens
+    slug=$(echo "${slug}" | sed 's/^-//;s/-$//')
+    # Truncate to 20 chars
+    slug="${slug:0:20}"
+    # Strip trailing hyphen after truncation
+    slug=$(echo "${slug}" | sed 's/-$//')
+
+    echo "${slug}"
+}
+
 # ── _handle_completion() ──────────────────────────────────────────────────────
 # Store output paths, detect output type, print next-step commands.
 _handle_completion() {
@@ -1294,6 +1496,14 @@ _handle_completion() {
     # Store output paths in config
     if [ "${output_type}" = "adapter" ]; then
         _update_config_var "TUNE_ADAPTER_PATH_${technique_upper}" "${artifact_path}"
+        # Write dataset-specific adapter path if dataset slug is available
+        local dataset_slug
+        dataset_slug=$(_derive_dataset_slug "${ARG_DATASET:-}")
+        if [ -n "${dataset_slug}" ]; then
+            local slug_upper
+            slug_upper=$(echo "${dataset_slug}" | tr '[:lower:]' '[:upper:]' | sed 's/-/_/g')
+            _update_config_var "TUNE_ADAPTER_PATH_${technique_upper}_${slug_upper}" "${artifact_path}"
+        fi
     else
         _update_config_var "TUNE_MODEL_PATH_${technique_upper}" "${artifact_path}"
     fi
@@ -1304,9 +1514,14 @@ _handle_completion() {
     echo "📋 Next steps:"
     echo ""
     if [ "${output_type}" = "adapter" ]; then
+        local dataset_slug
+        dataset_slug=$(_derive_dataset_slug "${ARG_DATASET:-}")
         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}"
+        if [ -n "${dataset_slug}" ]; then
+            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}"
     else
         echo "   Deploy as new inference component:"
@@ -1403,6 +1618,43 @@ if [ "${ARG_DISCOVER}" = true ]; then
     exit 0
 fi
 
+# Handle --list-datasets (before requiring --technique and --dataset)
+if [ "${ARG_LIST_DATASETS}" = true ]; then
+    echo ""
+    echo "📦 Registered datasets:"
+    echo ""
+
+    _ds_output=$(python3 "${SCRIPT_DIR}/.register_helper.py" list-datasets 2>/dev/null) || _ds_output=""
+    _ds_json=$(echo "${_ds_output}" | grep -E '^\{' | tail -1)
+
+    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" "----" "---------" "----" "------"
+            echo "${_ds_json}" | python3 -c "
+import sys, json
+data = json.load(sys.stdin)
+for ds in data.get('datasets', []):
+    name = ds.get('name','')[:25]
+    tech = ds.get('technique','')[:10]
+    rows = str(ds.get('row_count',''))[:8]
+    uri = ds.get('s3_uri','')
+    print(f'  {name:<25} {tech:<10} {rows:<8} {uri}')
+" 2>/dev/null
+        else
+            echo "  (none registered)"
+        fi
+    else
+        echo "  (none registered)"
+    fi
+    echo ""
+    echo "  Register: ./do/register dataset <name> --s3-uri <uri> --technique <sft|dpo>"
+    echo "  Use:      ./do/tune --technique sft --dataset <name>"
+    echo ""
+    exit 0
+fi
+
 # Validate required arguments for job submission
 if [ -z "${ARG_TECHNIQUE}" ]; then
     echo "❌ --technique is required"