@aws/ml-container-creator 1.0.3 → 1.1.0

package/templates/do/training/sft/train.py

@@ -0,0 +1,310 @@
+#!/usr/bin/env python3
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unmanaged SFT fine-tuning using TRL SFTTrainer + PEFT LoRA.
+
+This script performs LoRA-based supervised fine-tuning on a causal language
+model. It is invoked via `accelerate launch` and works on both single-GPU
+and multi-GPU instances without code changes.
+
+Portable env-var contract (works on SageMaker AI and HyperPod EKS):
+    DATA_DIR / SM_CHANNEL_TRAINING   -> training data path
+    OUTPUT_DIR / SM_MODEL_DIR        -> model artifact output
+    CHECKPOINT_DIR / SM_CHECKPOINT_DIR -> checkpoint path for spot resume
+    HF_MODEL_ID / SM_HP_MODEL_ID     -> base model HuggingFace ID
+    SM_HPS                           -> JSON blob of all hyperparameters
+
+Output:
+    LoRA adapter saved to OUTPUT_DIR (adapter_model.safetensors + adapter_config.json)
+    Metrics logged to stdout in SageMaker-parseable format
+"""
+
+import glob
+import json
+import logging
+import os
+import sys
+
+# ── Logging ───────────────────────────────────────────────────────────────────
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+)
+logger = logging.getLogger("sft-trainer")
+
+# ── Portable Path Resolution ─────────────────────────────────────────────────
+# Fallback chain: generic env var -> SageMaker env var -> default path
+# This allows the same script to run on SageMaker, HyperPod, or locally.
+
+DATA_DIR = (
+    os.environ.get("DATA_DIR")
+    or os.environ.get("SM_CHANNEL_TRAINING")
+    or "/opt/ml/input/data/training"
+)
+
+OUTPUT_DIR = (
+    os.environ.get("OUTPUT_DIR")
+    or os.environ.get("SM_MODEL_DIR")
+    or "/opt/ml/model"
+)
+
+CHECKPOINT_DIR = (
+    os.environ.get("CHECKPOINT_DIR")
+    or os.environ.get("SM_CHECKPOINT_DIR")
+    or "/opt/ml/checkpoints"
+)
+
+MODEL_ID = (
+    os.environ.get("HF_MODEL_ID")
+    or os.environ.get("SM_HP_MODEL_ID")
+    or ""
+)
+
+
+# ── Hyperparameter Loading ────────────────────────────────────────────────────
+
+def load_hyperparameters():
+    """Load hyperparameters from SageMaker SM_HPS env var or individual SM_HP_* vars.
+
+    Returns:
+        dict with typed hyperparameter values.
+    """
+    defaults = {
+        "model_id": MODEL_ID,
+        "lora_r": 16,
+        "lora_alpha": 32,
+        "lora_dropout": 0.05,
+        "learning_rate": 2e-4,
+        "epochs": 3,
+        "batch_size": 4,
+        "max_seq_length": 2048,
+        "gradient_accumulation_steps": 4,
+        "warmup_ratio": 0.03,
+        "dataset_text_field": "text",
+    }
+
+    # Try SM_HPS (JSON blob of all hyperparameters)
+    sm_hps = os.environ.get("SM_HPS")
+    if sm_hps:
+        try:
+            raw = json.loads(sm_hps)
+            # SageMaker passes all values as strings — cast them
+            for key, default_val in defaults.items():
+                if key in raw:
+                    defaults[key] = _cast(raw[key], type(default_val))
+            return defaults
+        except (json.JSONDecodeError, ValueError) as e:
+            logger.warning("Failed to parse SM_HPS: %s", e)
+
+    # Fallback: individual SM_HP_* env vars
+    for key, default_val in defaults.items():
+        env_key = f"SM_HP_{key.upper()}"
+        env_val = os.environ.get(env_key)
+        if env_val is not None:
+            defaults[key] = _cast(env_val, type(default_val))
+
+    return defaults
+
+
+def _cast(value, target_type):
+    """Cast a string value to the target type."""
+    if target_type == bool:
+        return str(value).lower() in ("true", "1", "yes")
+    if target_type == int:
+        return int(float(value))
+    if target_type == float:
+        return float(value)
+    return str(value)
+
+
+# ── Dataset Loading ───────────────────────────────────────────────────────────
+
+def load_dataset(data_dir, text_field):
+    """Load training dataset from data directory.
+
+    Supports .jsonl, .parquet, and .csv files.
+
+    Args:
+        data_dir: Path to directory containing training data files.
+        text_field: Name of the text column in the dataset.
+
+    Returns:
+        A Hugging Face Dataset object.
+    """
+    from datasets import load_dataset as hf_load_dataset
+
+    # Find data files
+    extensions = ["jsonl", "json", "parquet", "csv"]
+    data_files = []
+    for ext in extensions:
+        data_files.extend(glob.glob(os.path.join(data_dir, f"*.{ext}")))
+        data_files.extend(glob.glob(os.path.join(data_dir, f"**/*.{ext}"), recursive=True))
+
+    if not data_files:
+        logger.error("No data files found in %s (searched: %s)", data_dir, extensions)
+        sys.exit(1)
+
+    # Deduplicate and sort
+    data_files = sorted(set(data_files))
+    logger.info("Found %d data file(s) in %s", len(data_files), data_dir)
+
+    # Determine format from first file extension
+    first_ext = data_files[0].rsplit(".", 1)[-1].lower()
+    format_map = {"jsonl": "json", "json": "json", "parquet": "parquet", "csv": "csv"}
+    file_format = format_map.get(first_ext, "json")
+
+    dataset = hf_load_dataset(file_format, data_files=data_files, split="train")
+    logger.info("Loaded dataset: %d rows, columns: %s", len(dataset), dataset.column_names)
+
+    # Verify text field exists
+    if text_field not in dataset.column_names:
+        logger.error(
+            "Text field '%s' not found in dataset. Available columns: %s",
+            text_field, dataset.column_names,
+        )
+        sys.exit(1)
+
+    return dataset
+
+
+# ── Main Training Function ────────────────────────────────────────────────────
+
+def main():
+    """Run SFT training with TRL SFTTrainer + PEFT LoRA."""
+    from accelerate import Accelerator
+    from peft import LoraConfig, TaskType
+    from transformers import (
+        AutoModelForCausalLM,
+        AutoTokenizer,
+        TrainingArguments,
+    )
+    from trl import SFTTrainer
+
+    # Initialize accelerator (handles distributed setup)
+    accelerator = Accelerator()
+
+    # Load hyperparameters
+    hparams = load_hyperparameters()
+    model_id = hparams["model_id"]
+
+    if not model_id:
+        logger.error("No model ID specified. Set HF_MODEL_ID env var or model_id hyperparameter.")
+        sys.exit(1)
+
+    if accelerator.is_main_process:
+        logger.info("=" * 60)
+        logger.info("SFT Training Configuration")
+        logger.info("=" * 60)
+        logger.info("  Model:          %s", model_id)
+        logger.info("  Data dir:       %s", DATA_DIR)
+        logger.info("  Output dir:     %s", OUTPUT_DIR)
+        logger.info("  Checkpoint dir: %s", CHECKPOINT_DIR)
+        logger.info("  LoRA r:         %d", hparams["lora_r"])
+        logger.info("  LoRA alpha:     %d", hparams["lora_alpha"])
+        logger.info("  Learning rate:  %s", hparams["learning_rate"])
+        logger.info("  Epochs:         %d", hparams["epochs"])
+        logger.info("  Batch size:     %d", hparams["batch_size"])
+        logger.info("  Max seq len:    %d", hparams["max_seq_length"])
+        logger.info("  Text field:     %s", hparams["dataset_text_field"])
+        logger.info("=" * 60)
+
+    # ── Load tokenizer and model ──────────────────────────────────────────────
+    if accelerator.is_main_process:
+        logger.info("Loading tokenizer and model: %s", model_id)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype="auto",
+        trust_remote_code=True,
+    )
+
+    # ── Configure LoRA ────────────────────────────────────────────────────────
+    lora_config = LoraConfig(
+        r=hparams["lora_r"],
+        lora_alpha=hparams["lora_alpha"],
+        lora_dropout=hparams["lora_dropout"],
+        target_modules="all-linear",
+        task_type=TaskType.CAUSAL_LM,
+    )
+
+    # ── Load dataset ──────────────────────────────────────────────────────────
+    dataset = load_dataset(DATA_DIR, hparams["dataset_text_field"])
+
+    # ── Training arguments ────────────────────────────────────────────────────
+    training_args = TrainingArguments(
+        output_dir=os.path.join(CHECKPOINT_DIR, "trainer-state"),
+        num_train_epochs=hparams["epochs"],
+        per_device_train_batch_size=hparams["batch_size"],
+        gradient_accumulation_steps=hparams["gradient_accumulation_steps"],
+        learning_rate=hparams["learning_rate"],
+        warmup_ratio=hparams["warmup_ratio"],
+        bf16=True,
+        logging_steps=10,
+        save_strategy="epoch",
+        save_total_limit=2,
+        report_to="none",
+        remove_unused_columns=False,
+    )
+
+    # ── Check for existing checkpoint (spot resume) ───────────────────────────
+    resume_from_checkpoint = None
+    trainer_state_dir = os.path.join(CHECKPOINT_DIR, "trainer-state")
+    if os.path.isdir(trainer_state_dir):
+        checkpoints = sorted(
+            glob.glob(os.path.join(trainer_state_dir, "checkpoint-*")),
+            key=lambda x: int(x.rsplit("-", 1)[-1]) if x.rsplit("-", 1)[-1].isdigit() else 0,
+        )
+        if checkpoints:
+            resume_from_checkpoint = checkpoints[-1]
+            if accelerator.is_main_process:
+                logger.info("Resuming from checkpoint: %s", resume_from_checkpoint)
+
+    # ── Initialize SFTTrainer ─────────────────────────────────────────────────
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=dataset,
+        peft_config=lora_config,
+        args=training_args,
+        max_seq_length=hparams["max_seq_length"],
+        dataset_text_field=hparams["dataset_text_field"],
+    )
+
+    # ── Train ─────────────────────────────────────────────────────────────────
+    if accelerator.is_main_process:
+        logger.info("Starting SFT training...")
+
+    train_result = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+
+    # ── Save adapter (rank 0 only) ───────────────────────────────────────────
+    if accelerator.is_main_process:
+        logger.info("Saving LoRA adapter to: %s", OUTPUT_DIR)
+        trainer.save_model(OUTPUT_DIR)
+        tokenizer.save_pretrained(OUTPUT_DIR)
+
+        # ── Log final metrics ─────────────────────────────────────────────────
+        # Format: metric_name: value (SageMaker captures via regex in config.yaml)
+        metrics = train_result.metrics
+        print(f"train_loss: {metrics.get('train_loss', 0.0):.4f}")
+        print(f"train_runtime: {metrics.get('train_runtime', 0.0):.1f}")
+        print(f"train_samples_per_second: {metrics.get('train_samples_per_second', 0.0):.2f}")
+        print(f"epochs: {hparams['epochs']}")
+
+        logger.info("Training complete!")
+        logger.info("  Loss:     %.4f", metrics.get("train_loss", 0.0))
+        logger.info("  Runtime:  %.1fs", metrics.get("train_runtime", 0.0))
+        logger.info("  Samples/s: %.2f", metrics.get("train_samples_per_second", 0.0))
+
+    # Wait for all processes before exit
+    accelerator.wait_for_everyone()
+
+
+# ── Entry Point ───────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    main()

package/templates/do/tune

@@ -796,9 +796,16 @@ else:
 _validate_dataset() {
     local dataset="${ARG_DATASET}"
 
-    # ── Parse @v<N> version suffix (AC-2.1, AC-2.3) ──────────────────────────
+    # ── Parse @v<version> suffix (AC-2.1, AC-2.3) ───────────────────────────
     # Syntax: dataset-name@v2 → name="dataset-name", version ordinal=2
-    if [[ "${dataset}" =~ ^(.+)@v([0-9]+)$ ]]; then
+    #         dataset-name@v1.0.0 → name="dataset-name", version="1.0.0" (semver)
+    if [[ "${dataset}" =~ ^(.+)@v([0-9]+\.[0-9]+\.[0-9]+)$ ]]; then
+        # Semver syntax: @v1.0.0
+        ARG_DATASET_NAME="${BASH_REMATCH[1]}"
+        ARG_DATASET_VERSION="${BASH_REMATCH[2]}"
+        dataset=""  # Clear so name-based resolution takes over
+    elif [[ "${dataset}" =~ ^(.+)@v([0-9]+)$ ]]; then
+        # Ordinal syntax: @v1, @v2
         ARG_DATASET_NAME="${BASH_REMATCH[1]}"
         ARG_DATASET_VERSION="${BASH_REMATCH[2]}"
         dataset=""  # Clear so name-based resolution takes over
@@ -824,6 +831,8 @@ _validate_dataset() {
                 echo "   Resolved to: ${resolved_uri}"
                 dataset="${resolved_uri}"
                 ARG_DATASET="${resolved_uri}"
+                RESOLVED_DATASET_S3_URI="${resolved_uri}"
+                return 0
             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://..."

package/src/lib/auto-prompt-builder.js

@@ -1,172 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-
-/**
- * Auto-Prompt Builder — generates targeted prompts for missing required parameters.
- * 
- * Used by --auto-prompt mode to ask only for values that cannot be inferred
- * or defaulted from the provided CLI flags.
- */
-
-/**
- * Builds a minimal set of prompts for the given missing parameters.
- * Each prompt is self-contained and doesn't depend on multi-phase wizard state.
- *
- * @param {string[]} missingParams - Parameter names that need values
- * @param {object} currentConfig - Current configuration (with defaults filled)
- * @returns {Array} Array of prompt objects compatible with runPrompts()
- */
-export function buildAutoPrompts(missingParams, currentConfig) {
-    const prompts = [];
-
-    for (const param of missingParams) {
-        const builder = PROMPT_BUILDERS[param];
-        if (builder) {
-            const prompt = builder(currentConfig);
-            if (prompt) {
-                prompts.push(prompt);
-            }
-        } else {
-            // Fallback: generic text input for unknown parameters
-            prompts.push({
-                type: 'input',
-                name: param,
-                message: `Enter value for ${param}:`
-            });
-        }
-    }
-
-    return prompts;
-}
-
-/**
- * Map of parameter names to prompt builder functions.
- * Each builder receives the current config and returns a prompt object.
- */
-const PROMPT_BUILDERS = {
-    deploymentConfig: (_config) => ({
-        type: 'list',
-        name: 'deploymentConfig',
-        message: 'Select deployment configuration:',
-        choices: [
-            { type: 'separator', separator: '── Large Language Models ──' },
-            { name: 'Transformers with vLLM', value: 'transformers-vllm' },
-            { name: 'Transformers with SGLang', value: 'transformers-sglang' },
-            { name: 'Transformers with TensorRT-LLM', value: 'transformers-tensorrt-llm' },
-            { name: 'Transformers with LMI', value: 'transformers-lmi' },
-            { name: 'Transformers with DJL', value: 'transformers-djl' },
-            { type: 'separator', separator: '── HTTP Serving ──' },
-            { name: 'HTTP with Flask', value: 'http-flask' },
-            { name: 'HTTP with FastAPI', value: 'http-fastapi' },
-            { type: 'separator', separator: '── NVIDIA Triton ──' },
-            { name: 'Triton FIL (XGBoost, LightGBM)', value: 'triton-fil' },
-            { name: 'Triton ONNX Runtime', value: 'triton-onnxruntime' },
-            { name: 'Triton TensorFlow', value: 'triton-tensorflow' },
-            { name: 'Triton PyTorch', value: 'triton-pytorch' },
-            { name: 'Triton vLLM', value: 'triton-vllm' },
-            { name: 'Triton TensorRT-LLM', value: 'triton-tensorrtllm' },
-            { name: 'Triton Python Backend', value: 'triton-python' },
-            { type: 'separator', separator: '── Diffusion Models ──' },
-            { name: 'Diffusors with vLLM Omni', value: 'diffusors-vllm-omni' }
-        ]
-    }),
-
-    instanceType: (config) => {
-        const architecture = config.architecture || 'http';
-        const isGpu = architecture === 'transformers' || architecture === 'triton' || architecture === 'diffusors';
-
-        const gpuChoices = [
-            { name: 'ml.g5.xlarge  (1× A10G 24GB — small LLMs)', value: 'ml.g5.xlarge' },
-            { name: 'ml.g5.2xlarge (1× A10G 24GB — medium LLMs)', value: 'ml.g5.2xlarge' },
-            { name: 'ml.g5.4xlarge (1× A10G 24GB — larger models)', value: 'ml.g5.4xlarge' },
-            { name: 'ml.g5.12xlarge (4× A10G 96GB — large LLMs)', value: 'ml.g5.12xlarge' },
-            { name: 'ml.g5.48xlarge (8× A10G 192GB — very large)', value: 'ml.g5.48xlarge' },
-            { name: 'ml.g6.xlarge  (1× L4 24GB)', value: 'ml.g6.xlarge' },
-            { name: 'ml.g6.2xlarge (1× L4 24GB)', value: 'ml.g6.2xlarge' },
-            { name: 'ml.p4d.24xlarge (8× A100 320GB)', value: 'ml.p4d.24xlarge' },
-            { name: 'ml.p5.48xlarge (8× H100 640GB)', value: 'ml.p5.48xlarge' },
-            { name: 'Custom (enter manually)', value: '_custom' }
-        ];
-
-        const cpuChoices = [
-            { name: 'ml.m5.large   (2 vCPU, 8GB — lightweight)', value: 'ml.m5.large' },
-            { name: 'ml.m5.xlarge  (4 vCPU, 16GB — small models)', value: 'ml.m5.xlarge' },
-            { name: 'ml.m5.2xlarge (8 vCPU, 32GB — medium models)', value: 'ml.m5.2xlarge' },
-            { name: 'ml.m5.4xlarge (16 vCPU, 64GB — large models)', value: 'ml.m5.4xlarge' },
-            { name: 'ml.c5.xlarge  (4 vCPU, 8GB — compute-heavy)', value: 'ml.c5.xlarge' },
-            { name: 'ml.c5.2xlarge (8 vCPU, 16GB — compute-heavy)', value: 'ml.c5.2xlarge' },
-            { name: 'Custom (enter manually)', value: '_custom' }
-        ];
-
-        return {
-            type: 'list',
-            name: 'instanceType',
-            message: `Select instance type${isGpu ? ' (GPU recommended for this architecture)' : ''}:`,
-            choices: isGpu ? gpuChoices : cpuChoices
-        };
-    },
-
-    deploymentTarget: (_config) => ({
-        type: 'list',
-        name: 'deploymentTarget',
-        message: 'Select deployment target:',
-        choices: [
-            { name: 'Real-Time Inference', value: 'realtime-inference' },
-            { name: 'Async Inference', value: 'async-inference' },
-            { name: 'Batch Transform', value: 'batch-transform' },
-            { name: 'HyperPod EKS', value: 'hyperpod-eks' }
-        ]
-    }),
-
-    modelFormat: (config) => {
-        const engine = config.engine || 'sklearn';
-        const formatMap = {
-            sklearn: [
-                { name: 'pkl (pickle)', value: 'pkl' },
-                { name: 'joblib', value: 'joblib' }
-            ],
-            xgboost: [
-                { name: 'json', value: 'json' },
-                { name: 'model (binary)', value: 'model' },
-                { name: 'ubj (universal binary JSON)', value: 'ubj' }
-            ],
-            tensorflow: [
-                { name: 'keras', value: 'keras' },
-                { name: 'h5', value: 'h5' },
-                { name: 'SavedModel', value: 'SavedModel' }
-            ]
-        };
-
-        const choices = formatMap[engine] || formatMap.sklearn;
-
-        return {
-            type: 'list',
-            name: 'modelFormat',
-            message: `Select model format for ${engine}:`,
-            choices
-        };
-    },
-
-    awsRegion: (_config) => ({
-        type: 'list',
-        name: 'awsRegion',
-        message: 'Select AWS region:',
-        choices: [
-            { name: 'us-east-1 (N. Virginia)', value: 'us-east-1' },
-            { name: 'us-west-2 (Oregon)', value: 'us-west-2' },
-            { name: 'eu-west-1 (Ireland)', value: 'eu-west-1' },
-            { name: 'ap-northeast-1 (Tokyo)', value: 'ap-northeast-1' },
-            { name: 'ap-southeast-1 (Singapore)', value: 'ap-southeast-1' },
-            { name: 'Custom (enter manually)', value: '_custom' }
-        ]
-    }),
-
-    buildTarget: (_config) => ({
-        type: 'list',
-        name: 'buildTarget',
-        message: 'Select build target:',
-        choices: [
-            { name: 'CodeBuild (recommended)', value: 'codebuild' }
-        ]
-    })
-};