command to start sagemaker training

include sample training
2026-05-25 16:48:31 -04:00
parent 62ffe163e8
commit 0e728cc193
13 changed files with 796 additions and 5 deletions
--- a/examples/training/README.md
+++ b/examples/training/README.md
@@ -0,0 +1,90 @@
+# SageMaker Training Example
+
+This example downloads a small image-classification dataset, uploads it through `qc-cli`, and submits a live SageMaker training job.
+
+## Prerequisites
+
+- AWS credentials configured for the profile in `config.yaml`
+- Infrastructure already deployed with `qc-cli infra setup`
+- `config.yaml` updated with:
+
+```yaml
+s3:
+  bucket: your-bucket-name
+
+sagemaker:
+  role_name: <role-name>
+  training:
+    image_uri: 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.6-cpu-py312-ubuntu22.04-sagemaker-v1
+    instance_type: ml.m4.xlarge
+    instance_count: 1
+    source_dir: examples/training/source
+    entry_point: train.py
+    hyperparameters:
+      epochs: 1
+      batch-size: 32
+      learning-rate: 0.001
+      image-size: 160
+      validation-split: 0.2
+```
+
+## Training Hyperparameters
+
+Values under `sagemaker.training.hyperparameters` are passed to the training entry point as command-line arguments. For this example, they map to arguments defined in [source/train.py](source/train.py).
+
+Supported by this example:
+
+| Name | Type | Default | Description |
+|---|---:|---:|---|
+| `epochs` | int | `1` | Number of training epochs. |
+| `batch-size` | int | `32` | Images per training batch. |
+| `learning-rate` | float | `0.001` | Adam optimizer learning rate. |
+| `image-size` | int | `160` | Resize images to square `image-size x image-size`. |
+| `validation-split` | float | `0.2` | Fraction of data used for validation. |
+| `max-samples` | int | `0` | Optional cap for smoke tests; `0` means use all images. |
+| `seed` | int | `13` | Random seed for reproducible splitting. |
+| `num-workers` | int | `2` | DataLoader worker count. |
+
+Do not set `train-dir` or `model-dir` in normal SageMaker runs. SageMaker sets those automatically through `SM_CHANNEL_TRAIN` and `SM_MODEL_DIR`.
+
+## 1. Download The Dataset
+
+```bash
+bash examples/training/download_flower_photos.sh
+```
+
+This creates:
+
+```text
+examples/training/data/flower_photos_sagemaker/
+  daisy/
+  dandelion/
+  roses/
+  sunflowers/
+  tulips/
+```
+
+## 2. Run Training
+
+Run the training script and wait until it finishes:
+
+```bash
+bash examples/training/run_training.sh --config config.yaml --wait
+```
+
+Use a dataset that is already uploaded to `s3.data_prefix`:
+
+```bash
+bash examples/training/run_training.sh \
+  --config config.yaml \
+  --skip-upload \
+  --wait
+```
+
+## Notes
+
+- The default dataset path is `examples/training/data/flower_photos_sagemaker`.
+- Uploaded data uses the `s3.bucket` and `s3.data_prefix` values from `config.yaml`.
+- Training artifacts are written under `s3://<bucket>/<model_prefix>/`.
+- The SageMaker `model.tar.gz` contains `model.onnx`, `model.pt`, `class_to_idx.json`, and `metrics.json`.
+- SageMaker packages `examples/training/source`, installs `requirements.txt`, and runs `train.py`.
--- a/examples/training/download_flower_photos.sh
+++ b/examples/training/download_flower_photos.sh
@@ -0,0 +1,40 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+DATASET_URL="https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz"
+DEST_DIR="${1:-examples/training/data}"
+ARCHIVE_PATH="${DEST_DIR}/flower_photos.tgz"
+RAW_DATASET_DIR="${DEST_DIR}/flower_photos"
+DATASET_DIR="${DEST_DIR}/flower_photos_sagemaker"
+CLASS_NAMES=("daisy" "dandelion" "roses" "sunflowers" "tulips")
+
+mkdir -p "${DEST_DIR}"
+
+if [[ -d "${DATASET_DIR}" ]]; then
+  echo "Dataset already exists: ${DATASET_DIR}"
+  echo "Use this path with run_training.py:"
+  echo "  ${DATASET_DIR}"
+  exit 0
+fi
+
+echo "Downloading TensorFlow flower_photos dataset..."
+if command -v curl >/dev/null 2>&1; then
+  curl -L "${DATASET_URL}" -o "${ARCHIVE_PATH}"
+elif command -v wget >/dev/null 2>&1; then
+  wget -O "${ARCHIVE_PATH}" "${DATASET_URL}"
+else
+  echo "Either curl or wget is required." >&2
+  exit 1
+fi
+
+echo "Extracting dataset..."
+tar -xzf "${ARCHIVE_PATH}" -C "${DEST_DIR}"
+
+echo "Preparing SageMaker directory layout..."
+mkdir -p "${DATASET_DIR}"
+for class_name in "${CLASS_NAMES[@]}"; do
+  cp -R "${RAW_DATASET_DIR}/${class_name}" "${DATASET_DIR}/${class_name}"
+done
+
+echo "Dataset ready: ${DATASET_DIR}"
+find "${DATASET_DIR}" -mindepth 1 -maxdepth 1 -type d -print | sort
--- a/examples/training/run_training.sh
+++ b/examples/training/run_training.sh
@@ -0,0 +1,111 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+CONFIG_PATH="config.yaml"
+DATASET_DIR="examples/training/data/flower_photos_sagemaker"
+WAIT=false
+SKIP_UPLOAD=false
+POLL_SECONDS=60
+
+usage() {
+  cat <<EOF
+Usage: $0 [options]
+
+Options:
+  --config PATH          Path to qc-cli config file. Default: config.yaml
+  --dataset-dir PATH     Dataset directory to upload. Default: ${DATASET_DIR}
+  --skip-upload          Train against data already uploaded to s3.data_prefix.
+  --wait                 Poll until training completes.
+  -h, --help             Show this help.
+EOF
+}
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --config)
+      CONFIG_PATH="$2"
+      shift 2
+      ;;
+    --dataset-dir)
+      DATASET_DIR="$2"
+      shift 2
+      ;;
+    --skip-upload)
+      SKIP_UPLOAD=true
+      shift
+      ;;
+    --wait)
+      WAIT=true
+      shift
+      ;;
+    -h|--help)
+      usage
+      exit 0
+      ;;
+    *)
+      echo "Unknown option: $1" >&2
+      usage >&2
+      exit 1
+      ;;
+  esac
+done
+
+if [[ ! -f "${CONFIG_PATH}" ]]; then
+  echo "Config not found: ${CONFIG_PATH}" >&2
+  exit 1
+fi
+
+if [[ "${SKIP_UPLOAD}" == false && ! -d "${DATASET_DIR}" ]]; then
+  echo "Dataset not found: ${DATASET_DIR}" >&2
+  echo "Run: bash examples/training/download_flower_photos.sh" >&2
+  exit 1
+fi
+
+run() {
+  echo "+ $*"
+  "$@"
+}
+
+run uv run qc-cli infra status --config "${CONFIG_PATH}"
+
+if [[ "${SKIP_UPLOAD}" == false ]]; then
+  run uv run qc-cli upload "${DATASET_DIR}" --config "${CONFIG_PATH}"
+fi
+
+TRAIN_OUTPUT="$(uv run qc-cli train start --config "${CONFIG_PATH}")"
+echo "${TRAIN_OUTPUT}"
+
+JOB_NAME="$(printf '%s\n' "${TRAIN_OUTPUT}" | grep -Eo 'qc-cli-[0-9]{8}-[0-9]{6}' | tail -n 1)"
+if [[ -z "${JOB_NAME}" ]]; then
+  echo "Could not find training job name in qc-cli output." >&2
+  exit 1
+fi
+
+echo "Submitted SageMaker training job: ${JOB_NAME}"
+
+if [[ "${WAIT}" == false ]]; then
+  run uv run qc-cli train status "${JOB_NAME}" --config "${CONFIG_PATH}"
+  exit 0
+fi
+
+while true; do
+  STATUS_OUTPUT="$(uv run qc-cli train status "${JOB_NAME}" --config "${CONFIG_PATH}")"
+  echo "${STATUS_OUTPUT}"
+
+  if printf '%s\n' "${STATUS_OUTPUT}" | grep -q 'Status:.*Completed'; then
+    echo "Training completed successfully."
+    exit 0
+  fi
+
+  if printf '%s\n' "${STATUS_OUTPUT}" | grep -q 'Status:.*Failed'; then
+    echo "Training failed." >&2
+    exit 1
+  fi
+
+  if printf '%s\n' "${STATUS_OUTPUT}" | grep -q 'Status:.*Stopped'; then
+    echo "Training stopped." >&2
+    exit 1
+  fi
+
+  sleep "${POLL_SECONDS}"
+done
--- a/examples/training/source/requirements.txt
+++ b/examples/training/source/requirements.txt
@@ -0,0 +1 @@
+onnx==1.21.0
--- a/examples/training/source/train.py
+++ b/examples/training/source/train.py
@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+"""SageMaker entry point for CPU image-classification training."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, Subset, random_split
+from torchvision import datasets, transforms
+
+
+class SmallImageClassifier(nn.Module):
+    def __init__(self, class_count: int) -> None:
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 16, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            nn.Conv2d(16, 32, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),
+            nn.AdaptiveAvgPool2d((1, 1)),
+        )
+        self.classifier = nn.Linear(64, class_count)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = torch.flatten(x, 1)
+        return self.classifier(x)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--learning-rate", type=float, default=0.001)
+    parser.add_argument("--image-size", type=int, default=160)
+    parser.add_argument("--validation-split", type=float, default=0.2)
+    parser.add_argument("--max-samples", type=int, default=0)
+    parser.add_argument("--seed", type=int, default=13)
+    parser.add_argument("--num-workers", type=int, default=2)
+    parser.add_argument("--train-dir", default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
+    parser.add_argument("--model-dir", default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
+    return parser.parse_args()
+
+
+def build_datasets(args: argparse.Namespace) -> tuple[Subset, Subset, dict[str, int]]:
+    transform = transforms.Compose(
+        [
+            transforms.Resize((args.image_size, args.image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+        ]
+    )
+    dataset = datasets.ImageFolder(args.train_dir, transform=transform)
+    if len(dataset.classes) < 2:
+        raise ValueError(f"Expected at least two classes in {args.train_dir}. Found: {dataset.classes}")
+
+    if args.max_samples > 0 and args.max_samples < len(dataset):
+        indices = list(range(len(dataset)))
+        random.Random(args.seed).shuffle(indices)
+        dataset = Subset(dataset, indices[: args.max_samples])
+
+    validation_size = max(1, int(len(dataset) * args.validation_split))
+    train_size = len(dataset) - validation_size
+    if train_size < 1:
+        raise ValueError("Not enough images to create a train/validation split.")
+
+    generator = torch.Generator().manual_seed(args.seed)
+    train_dataset, validation_dataset = random_split(dataset, [train_size, validation_size], generator=generator)
+    return train_dataset, validation_dataset, getattr(dataset, "dataset", dataset).class_to_idx
+
+
+def run_epoch(
+    model: nn.Module,
+    data_loader: DataLoader,
+    criterion: nn.Module,
+    optimizer: torch.optim.Optimizer | None,
+    device: torch.device,
+) -> tuple[float, float]:
+    training = optimizer is not None
+    model.train(training)
+
+    total_loss = 0.0
+    total_correct = 0
+    total_examples = 0
+
+    for images, labels in data_loader:
+        images = images.to(device)
+        labels = labels.to(device)
+
+        with torch.set_grad_enabled(training):
+            logits = model(images)
+            loss = criterion(logits, labels)
+
+            if training:
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+        total_loss += loss.item() * images.size(0)
+        total_correct += (logits.argmax(dim=1) == labels).sum().item()
+        total_examples += images.size(0)
+
+    return total_loss / total_examples, total_correct / total_examples
+
+
+def export_onnx(model: nn.Module, model_dir: Path, image_size: int) -> None:
+    model.eval()
+    dummy_input = torch.randn(1, 3, image_size, image_size)
+    torch.onnx.export(
+        model,
+        dummy_input,
+        model_dir / "model.onnx",
+        export_params=True,
+        opset_version=17,
+        do_constant_folding=True,
+        input_names=["input"],
+        output_names=["logits"],
+        dynamic_axes={
+            "input": {0: "batch_size"},
+            "logits": {0: "batch_size"},
+        },
+    )
+
+
+def main() -> None:
+    args = parse_args()
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+
+    train_dataset, validation_dataset, class_to_idx = build_datasets(args)
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+    )
+    validation_loader = DataLoader(
+        validation_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = SmallImageClassifier(class_count=len(class_to_idx)).to(device)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
+
+    print(f"Training on {device}. Classes: {sorted(class_to_idx)}")
+    metrics = []
+    for epoch in range(1, args.epochs + 1):
+        train_loss, train_accuracy = run_epoch(model, train_loader, criterion, optimizer, device)
+        validation_loss, validation_accuracy = run_epoch(model, validation_loader, criterion, None, device)
+        epoch_metrics = {
+            "epoch": epoch,
+            "train_loss": train_loss,
+            "train_accuracy": train_accuracy,
+            "validation_loss": validation_loss,
+            "validation_accuracy": validation_accuracy,
+        }
+        metrics.append(epoch_metrics)
+        print(json.dumps(epoch_metrics, sort_keys=True))
+
+    model_dir = Path(args.model_dir)
+    model_dir.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "model_state_dict": model.cpu().state_dict(),
+            "class_to_idx": class_to_idx,
+            "image_size": args.image_size,
+        },
+        model_dir / "model.pt",
+    )
+    export_onnx(model, model_dir, args.image_size)
+    (model_dir / "class_to_idx.json").write_text(json.dumps(class_to_idx, indent=2), encoding="utf-8")
+    (model_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
+    print(f"Saved model artifacts to {model_dir}")
+
+
+if __name__ == "__main__":
+    main()