5.2 Custom Models with Script Mode and BYOC

5.2 Custom Models with Script Mode and BYOC

While Amazon SageMaker built-in algorithms provide high performance for standard ML workloads, enterprise machine learning applications frequently require custom deep learning architectures, proprietary loss functions, cutting-edge transformer models from HuggingFace, or legacy Scikit-learn pipelines. To support any framework or runtime, SageMaker offers a flexible continuum of customization options: Script Mode, Extending Pre-built Deep Learning Containers, and Bring Your Own Container (BYOC).

On the AWS Certified Machine Learning Engineer — Associate (MLA-C01) exam, you must know when and how to implement Script Mode, configure training environment variables, manage custom dependencies via requirements.txt, and adhere to the strict POSIX file-structure contracts required for BYOC training and inference.

+------------------------------------------------------------------------------------------------+
|                       SAGEMAKER MODEL CUSTOMIZATION SPECTRUM                                   |
|                                                                                                |
|   [LOWEST OPERATIONAL OVERHEAD]                                   [MAXIMUM CONTROL & FLEX]     |
|                                                                                                |
|   1. Built-in Algos   --->   2. Script Mode    --->   3. Extend Pre-built  ---> 4. Full BYOC   |
|   - Zero code                - Custom script         - Dockerfile FROM         - Custom Docker |
|   - AWS Managed DLC          - AWS Managed DLC         AWS DLC image             from scratch  |
|   - Hyperparams only         - requirements.txt      - Custom OS libs          - C++/Rust/R    |
+------------------------------------------------------------------------------------------------+

---

1. SageMaker Script Mode Architecture

Script Mode allows engineers to execute native Python training code (e.g., standard PyTorch .py scripts or TensorFlow models) using pre-built, AWS-maintained Deep Learning Containers (DLCs). AWS manages CUDA drivers, cuDNN, Python runtimes, NCCL communication libraries, and deep learning framework optimizations, while the engineer simply authors the training script.

+------------------------------------------------------------------------------------------------+
|                               SCRIPT MODE EXECUTION ARCHITECTURE                               |
|                                                                                                |
|   Local Environment / SageMaker Studio Notebook                                                |
|   ├── my_project/                                                                              |
|   │   ├── train.py          <--- Entry point script                                            |
|   │   ├── model_def.py      <--- Custom neural network architecture                            |
|   │   ├── utils.py          <--- Data loading utilities                                        |
|   └── requirements.txt  <--- Additional pip dependencies (e.g., albumentations, timm)       |
|                                                                                                |
|             | (SageMaker Python SDK: estimator.fit({'train': 's3://...'}))                     |
|             v                                                                                  |
|   Managed SageMaker Training Cluster (EC2 Instance: ml.g5.2xlarge)                             |
|   +----------------------------------------------------------------------------------------+   |
|   | AWS Deep Learning Container (Pre-built PyTorch 2.1 + CUDA 12.1 + Ubuntu 22.04)         |   |
|   |                                                                                        |   |
|   | 1. Downloads tar.gz of my_project/ from S3 to container working directory              |   |
|   | 2. Executes: pip install -r requirements.txt                                           |   |
|   | 3. Injects S3 channels into /opt/ml/input/data/train & sets SM_CHANNEL_* env vars      |   |
|   | 4. Invokes: python train.py --epochs 50 --batch-size 64 --lr 0.001                     |   |
|   | 5. Script reads data, trains model, saves final weights to /opt/ml/model               |   |
|   | 6. SageMaker automatically archives /opt/ml/model to s3://.../model.tar.gz             |   |
|   +----------------------------------------------------------------------------------------+   |
+------------------------------------------------------------------------------------------------+

1.1 Script Mode Estimator Configuration

In the SageMaker Python SDK, framework-specific Estimator classes (PyTorch, TensorFlow, HuggingFace, SKLearn) encapsulate Script Mode execution:

import sagemaker
from sagemaker.pytorch import PyTorch

sagemaker_session = sagemaker.Session()
role = sagemaker.get_execution_role()

# Configure PyTorch Estimator in Script Mode
estimator = PyTorch(
    entry_point="train.py",             # Python script containing the training loop
    source_dir="src",                   # Directory containing train.py, utils.py, requirements.txt
    role=role,
    framework_version="2.1.0",          # Target PyTorch framework version
    py_version="py310",                 # Python runtime version
    instance_count=2,                   # Distributed multi-node training
    instance_type="ml.g5.12xlarge",     # Multi-GPU compute instance
    hyperparameters={
        "epochs": 50,
        "batch-size": 64,
        "learning-rate": 0.0005,
        "backbone": "resnet50"
    },
    metric_definitions=[
        {"Name": "train:loss", "Regex": "loss: ([0-9\\.]+)"},
        {"Name": "val:acc", "Regex": "val_acc: ([0-9\\.]+)"}
    ]
)

# Launch training job
estimator.fit({
    "train": "s3://my-ml-bucket/data/train/",
    "validation": "s3://my-ml-bucket/data/validation/"
})

1.2 Managing Python Dependencies with requirements.txt

If your custom script depends on third-party libraries not included in the pre-built AWS Deep Learning Container (such as timm, albumentations, evaluate, or wandb), place a standard requirements.txt file inside the directory referenced by source_dir.

• When the training container boots up, SageMaker automatically executes pip install -r requirements.txt before invoking the entry_point script.
• Network Note: The training instance must have outbound internet access (via NAT Gateway or direct VPC routing) to reach the PyPI package repository, unless packages are bundled in a private wheel directory or private ECR/VPC endpoint.

1.3 Script Mode Environment Variables & Command-Line Arguments

SageMaker passes input data channels, output paths, hyperparameters, and cluster configuration into the training container via environment variables and command-line arguments (argparse).

+------------------------------------------------------------------------------------------------+
|                       SAGEMAKER TRAINING ENVIRONMENT VARIABLES                                 |
|                                                                                                |
|   Environment Variable          Local Container Path                 Description               |
|   ----------------------------  ----------------------------------   ------------------------- |
|   `SM_MODEL_DIR`                `/opt/ml/model`                      Destination for model     |
|                                                                      artifacts (saved to S3)   |
|   `SM_CHANNEL_<CHANNEL_NAME>`   `/opt/ml/input/data/<channel_name>`  Input data channel path   |
|   `SM_OUTPUT_DATA_DIR`          `/opt/ml/output/data`                Auxiliary output files    |
|   `SM_HOSTS`                    `["algo-1", "algo-2"]`              JSON list of all hosts    |
|   `SM_CURRENT_HOST`             `"algo-1"`                           Hostname of current node  |
|   `SM_NUM_GPUS`                 `4`                                  Number of GPUs available  |
|   `SM_NUM_CPUS`                 `32`                                 Number of CPU cores       |
|   `SM_HPS`                      `'{"epochs": 50, "lr": 0.001}'`     JSON string of hyperparams|
+------------------------------------------------------------------------------------------------+

1.4 Production Script Mode train.py Template

import argparse
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

def parse_args():
    parser = argparse.ArgumentParser()
    # Hyperparameters passed via estimator.set_hyperparameters() are CLI arguments
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--learning-rate", type=float, default=0.001)
    
    # SageMaker passes environment variables pointing to directories
    parser.add_argument("--model-dir", type=str, default=os.environ.get("SM_MODEL_DIR"))
    parser.add_argument("--train", type=str, default=os.environ.get("SM_CHANNEL_TRAIN"))
    parser.add_argument("--val", type=str, default=os.environ.get("SM_CHANNEL_VALIDATION"))
    parser.add_argument("--num-gpus", type=int, default=os.environ.get("SM_NUM_GPUS", 0))
    return parser.parse_args()

def train():
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() and args.num_gpus > 0 else "cpu")
    
    # 1. Load data from args.train (/opt/ml/input/data/train)
    # dataset = CustomDataset(data_dir=args.train)
    # dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
    
    # 2. Build model
    model = nn.Sequential(nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 2)).to(device)
    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
    criterion = nn.CrossEntropyLoss()
    
    # 3. Training Loop
    for epoch in range(args.epochs):
        model.train()
        # Forward, backward, optimizer step...
        print(f"Epoch [{epoch+1}/{args.epochs}] - loss: 0.245 - val_acc: 0.942")
        
    # 4. Save model artifacts to args.model_dir (/opt/ml/model)
    # SageMaker automatically archives this directory into model.tar.gz upon completion
    model_save_path = os.path.join(args.model_dir, "model.pth")
    torch.save(model.state_dict(), model_save_path)
    print(f"Model successfully saved to {model_save_path}")

if __name__ == "__main__":
    train()

---

2. Custom Inference Handlers in Script Mode

When deploying a trained model from Script Mode to a SageMaker Real-Time or Serverless Endpoint, SageMaker's hosting container (TorchServe, TensorFlow Serving, or Multi Model Server) invokes an inference.py script containing four standard handler functions:

+------------------------------------------------------------------------------------------------+
|                       SAGEMAKER INFERENCE HANDLER LIFECYCLE                                    |
|                                                                                                |
|   1. Container Startup: `model_fn(model_dir)`                                                  |
|      Loads weights from `/opt/ml/model/` into GPU/CPU memory. Executed once.                   |
|                                                                                                |
|   2. Inference Request: Client sends HTTP POST payload to `/invocations`                      |
|                                                                                                |
|      +---> `input_fn(request_body, request_content_type)`                                      |
|      |     Deserializes incoming bytes (JSON, CSV, NPY, image) into tensor/array.              |
|      |                                                                                         |
|      +---> `predict_fn(input_data, model)`                                                     |
|      |     Executes forward pass using loaded model.                                           |
|      |                                                                                         |
|      +---> `output_fn(prediction, accept_type)`                                                |
|            Serializes tensor output to client-requested MIME type (e.g., application/json).    |
+------------------------------------------------------------------------------------------------+

# inference.py - Custom serving handler script
import json
import os
import torch

def model_fn(model_dir):
    """Loads model from disk into memory upon container boot."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = CustomNetwork()
    with open(os.path.join(model_dir, "model.pth"), "rb") as f:
        model.load_state_dict(torch.load(f, map_location=device))
    model.to(device).eval()
    return model

def input_fn(request_body, request_content_type):
    """Deserializes incoming HTTP payload."""
    if request_content_type == "application/json":
        data = json.loads(request_body)
        return torch.tensor(data["inputs"], dtype=torch.float32)
    raise ValueError(f"Unsupported content type: {request_content_type}")

def predict_fn(input_data, model):
    """Executes forward inference pass."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    with torch.no_grad():
        return model(input_data.to(device))

def output_fn(prediction, accept):
    """Formats model output for client HTTP response."""
    if accept == "application/json":
        return json.dumps({"predictions": prediction.cpu().tolist()}), accept
    raise ValueError(f"Unsupported accept type: {accept}")

---

3. Bring Your Own Container (BYOC) for Training

When your workload requires an unsupported runtime (e.g., C++, Julia, Rust, Go), proprietary compiled shared libraries (.so), specialized Linux system dependencies, or strict security-hardened golden Docker images, you must author a custom Docker container and push it to Amazon Elastic Container Registry (Amazon ECR).

+------------------------------------------------------------------------------------------------+
|                          BYOC TRAINING DIRECTORY STRUCTURE CONTRACT                            |
|                                                                                                |
|   /opt/ml/                                                                                     |
|   ├── input/                                                                                   |
|   │   ├── config/                                                                              |
|   │   │   ├── hyperparameters.json   <--- JSON dictionary of user hyperparameters              |
|   │   │   ├── inputdataconfig.json   <--- Data channel specifications                          |
|   │   │   └── resourceconfig.json    <--- Cluster hosts, current host, network info            |
|   │   └── data/                                                                                |
|   │       ├── train/                 <--- Ingested S3 training channel files                   |
|   │       └── validation/            <--- Ingested S3 validation channel files                 |
|   ├── model/                         <--- YOUR SCRIPT WRITES FINAL ARTIFACTS HERE              |
|   └── output/                                                                                  |
|       ├── failure                    <--- Write error message string here if job fails         |
|       └── data/                      <--- Auxiliary output metrics                             |
+------------------------------------------------------------------------------------------------+

3.1 The BYOC Training Container Contract

1. Executable Program: The container must define an executable named train (located in the system PATH, or specified via ENTRYPOINT or CMD in the Dockerfile). SageMaker runs docker run <image> train.
2. Reading Inputs:
   • Hyperparameters are read from /opt/ml/input/config/hyperparameters.json.
   • Input channels are mounted at /opt/ml/input/data/<channel_name>/.
3. Writing Outputs:
   • Final model weights and checkpoints must be written to /opt/ml/model/. Everything in this directory is compressed into model.tar.gz and uploaded to the training job's S3 output location.
4. Failure Handling:
   • If the training job fails, the process must exit with a non-zero exit code.
   • To display a human-readable failure reason in the SageMaker Console, CloudWatch Logs, and DescribeTrainingJob API, the container should write the failure traceback to /opt/ml/output/failure.

3.2 BYOC Training Dockerfile Example

# Example BYOC Dockerfile for Training
FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04

# Install system dependencies & Python
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3-pip python3-dev build-essential && \
    rm -rf /var/lib/apt/lists/*

# Install Python ML libraries
COPY requirements.txt /opt/program/requirements.txt
RUN pip3 install --no-cache-dir -r /opt/program/requirements.txt

# Set up code in container and make train script executable
COPY src/ /opt/program/
WORKDIR /opt/program/
RUN chmod +x train

# Add program directory to system PATH
ENV PATH="/opt/program:${PATH}"
ENTRYPOINT ["python3", "train.py"]

---

4. Bring Your Own Container (BYOC) for Inference Serving

To host a custom model container on a SageMaker Inference Endpoint (Real-Time, Serverless, or Asynchronous), the container must implement a lightweight HTTP web server contract.

+------------------------------------------------------------------------------------------------+
|                           BYOC INFERENCE WEB SERVER CONTRACT                                   |
|                                                                                                |
|   Port Requirement: Web server must listen on port 8080                                        |
|                                                                                                |
|   1. GET /ping (Health Check Endpoint)                                                         |
|      - Called by SageMaker during endpoint initialization and ongoing container monitoring.    |
|      - Must respond with HTTP status 200 OK within 2 seconds.                                  |
|      - Container is considered healthy only when /ping returns 200.                            |
|                                                                                                |
|   2. POST /invocations (Prediction Endpoint)                                                   |
|      - Called by SageMaker when a client executes InvokeEndpoint API.                          |
|      - Receives request payload in HTTP body, generates predictions, and returns HTTP 200      |
|        with the inference result.                                                              |
|      - Returns 4XX for client formatting errors; returns 5XX for internal model exceptions.    |
+------------------------------------------------------------------------------------------------+

4.1 BYOC Serving Web Server Implementation (FastAPI / Gunicorn)

# serve.py - Custom BYOC Inference Server with FastAPI
import os
import torch
from fastapi import FastAPI, Request, Response, status

app = FastAPI()
MODEL_PATH = "/opt/ml/model"
model = None

@app.on_event("startup")
def load_model():
    global model
    # Model weights are extracted by SageMaker to /opt/ml/model
    weights_file = os.path.join(MODEL_PATH, "model.pth")
    model = CustomNetwork()
    model.load_state_dict(torch.load(weights_file, map_location="cpu"))
    model.eval()
    print("Model successfully loaded into memory.")

@app.get("/ping", status_code=status.HTTP_200_OK)
def ping():
    """SageMaker Health Check: Must return 200 if container is healthy."""
    if model is not None:
        return {"status": "healthy"}
    return Response(status_code=status.HTTP_503_SERVICE_UNAVAILABLE)

@app.post("/invocations")
async def invocations(request: Request):
    """SageMaker Prediction: Receives client payload and returns predictions."""
    content_type = request.headers.get("Content-Type", "")
    if content_type != "application/json":
        return Response(content="Unsupported Media Type", status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE)
    
    body = await request.json()
    inputs = torch.tensor(body["instances"], dtype=torch.float32)
    
    with torch.no_grad():
        outputs = model(inputs)
        
    return {"predictions": outputs.tolist()}

---

5. Serial Inference Pipelines (Multi-Container Endpoints)

In enterprise systems, raw client requests often require multi-stage processing: data preprocessing/featurization (e.g., Scikit-learn or Spark ML), model scoring (e.g., PyTorch, XGBoost, or TensorFlow), and business rule post-processing. Deploying these stages as separate endpoints adds network latency and operational complexity.

A SageMaker Serial Inference Pipeline chains 2 to 15 containers sequentially behind a single HTTPS endpoint.

+------------------------------------------------------------------------------------------------+
|                            SERIAL INFERENCE PIPELINE DATA FLOW                                 |
|                                                                                                |
|   Client Request (Raw JSON/CSV)                                                                |
|         |                                                                                      |
|         v (HTTP POST /invocations)                                                             |
|   [SageMaker Endpoint: pipeline-endpoint-prod]                                                 |
|   +----------------------------------------------------------------------------------------+   |
|   | +-----------------------+   +-----------------------+   +----------------------------+ |   |
|   | | Container 1 (SKLearn) |-->| Container 2 (PyTorch) |-->| Container 3 (Post-process) | |   |
|   | | - One-hot encoding    |   | - Neural network      |   | - Probability thresholding | |   |
|   | | - Standard scaling    |   |   forward pass        |   | - JSON schema packaging    | |   |
|   | +-----------------------+   +-----------------------+   +----------------------------+ |   |
|   +----------------------------------------------------------------------------------------+   |
|         |                                                                                      |
|         v                                                                                      |
|   Client Response (Final Prediction JSON)                                                      |
+------------------------------------------------------------------------------------------------+

• Data is passed in-memory between co-located containers via localhost HTTP requests.
• Guarantees end-to-end atomic latency with zero intermediate S3 serialization overhead.

---

6. Customization Selection Decision Framework