5.2 Custom Models with Script Mode and BYOC

AWS Certified Machine Learning Engineer — Associate MLA-C01

Track your progress as you read

Finished sections get a check mark.

1 Chapter 1: Introduction & Exam Strategy

1.1 AWS MLA-C01 Blueprint, Format & Strategy 1.2 The Machine Learning Engineering Lifecycle on AWS

2 Domain 1: Data Ingestion & Storage Architecture for ML

2.1 Ingesting Batch and Streaming Data for ML 2.2 High-Performance Storage for Training and Inference 2.3 Data Cataloging, Governance, and Querying

3 Domain 1: Feature Engineering & Data Transformation

3.1 Distributed Feature Engineering with AWS Glue & DataBrew 3.2 Interactive Data Preparation with SageMaker Data Wrangler 3.3 Feature Management with SageMaker Feature Store

4 Domain 1: Data Quality, Integrity & Ground Truth Labeling

4.1 Data Validation, Schema Drift, and Quality Rules 4.2 Handling Missing Values, Outliers & Class Imbalance 4.3 Data Labeling with SageMaker Ground Truth & Ground Truth Plus

5 Domain 2: Model Selection & Framework Architecture

5.1 SageMaker Built-in Algorithms & Selection Criteria 5.2 Custom Models with Script Mode and BYOC 5.3 Foundation Models, JumpStart, and Amazon Bedrock

6 Domain 2: SageMaker Model Training & Hyperparameter Optimization

6.1 SageMaker Training Jobs & Cost Optimization 6.2 Distributed Training Strategies & Architectures 6.3 Hyperparameter Optimization with Automatic Model Tuning (AMT)

7 Domain 2: Model Evaluation, Tracking & Model Registry

7.1 ML Evaluation Metrics & Validation Strategies 7.2 Experiment Tracking with SageMaker Experiments 7.3 Model Governance with SageMaker Model Registry

8 Domain 3: SageMaker Inference Options & Endpoint Architecture

8.1 Real-Time, Serverless, and Asynchronous Inference 8.2 High-Throughput Offline Inference with Batch Transform 8.3 Multi-Model Endpoints (MME) & Multi-Container Endpoints (MCE)

9 Domain 3: Deployment Strategies & Infrastructure Scaling

9.1 Production Deployment Patterns & Traffic Shifting 9.2 Endpoint Auto-Scaling & Hardware Acceleration

10 Domain 3: MLOps Pipelines & Workflow Orchestration

10.1 SageMaker Pipelines Core Architecture 10.2 Advanced Pipeline Patterns & Integration 10.3 End-to-End MLOps CI/CD & Automated Retraining

11 Domain 4: Production Model Monitoring with SageMaker Model Monitor

11.1 Data Quality & Model Quality Monitoring 11.2 Model Bias & Model Explainability Drift 11.3 Operational Observability with Amazon CloudWatch

12 Domain 4: Responsible AI, Bias Detection & Governance

12.1 Pre-Training and Post-Training Bias with SageMaker Clarify 12.2 Model Documentation & Auditing with Model Cards & Dashboard 12.3 Safety, Guardrails & Generative AI Governance

13 Domain 4: ML Security, Network Isolation & Cost Optimization

13.1 IAM Roles and Least Privilege for ML 13.2 Network Isolation, VPC Endpoints & Data Encryption 13.3 ML Cost Optimization & Resource Governance

5.2 Custom Models with Script Mode and BYOC

Key Takeaways Explain This Topic AI

SageMaker Script Mode allows ML engineers to run standard PyTorch, TensorFlow, HuggingFace, and Scikit-learn scripts inside AWS-managed Deep Learning Containers without authoring Dockerfiles.

In Script Mode, custom third-party Python dependencies are automatically installed at training runtime by placing a standard requirements.txt file in the directory specified by source_dir .

SageMaker passes training environment configurations via standard environment variables: SM_MODEL_DIR (/opt/ml/model), SM_CHANNEL_<name> (/opt/ml/input/data/<name>), SM_NUM_GPUS , and SM_HOSTS .

Bring Your Own Container (BYOC) for training requires an executable train script reading from /opt/ml/input/data/ and writing model artifacts to /opt/ml/model/ with failure logs in /opt/ml/output/failure .

BYOC inference serving requires an HTTP web server listening on port 8080 responding to GET /ping (health check returning 200 OK within 2s) and POST /invocations (payload inference).

Last updated: August 2026

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

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 .

Reading Inputs:

Hyperparameters are read from /opt/ml/input/config/hyperparameters.json .

Input channels are mounted at /opt/ml/input/data/<channel_name>/ .

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.

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

Criteria Built-in Algorithms Script Mode Extend AWS Container Full BYOC

Development Effort Lowest (No code) Low (Python script only) Medium (Small Dockerfile) High (Full Docker management)

Framework Version AWS fixed versions AWS DLC versions AWS DLC + custom packages Any framework / runtime

Custom Python Dependencies Not supported Supported via requirements.txt Baked into Docker image Baked into Docker image

System Libraries (apt/yum) Not supported Not supported Supported ( RUN apt-get... ) Fully customizable OS/kernel

Inference Server Management Handled by AWS Handled by AWS DLC handlers Handled by AWS DLC handlers Must implement /ping & /invocations on 8080

ECR Image Maintenance Managed by AWS Managed by AWS Customer manages ECR image Customer manages ECR image

Loading diagram... Custom Model Architecture Decision Tree