Module 5 — Custom containers and training scripts

Built-in XGBoost is convenient, but the moment you want to add a preprocessing step, a stratified split or a metric that AWS does not export, you leave that comfort. This module ports the churn model to a scikit-learn script, then shows the two escape hatches for cases the script mode does not cover.

Script mode, the default path

The script mode is the middle ground: AWS provides a container with a framework already installed, you provide the training script. Nine cases out of ten this is what you want.

from sagemaker.sklearn.estimator import SKLearn

sk = SKLearn(
    entry_point="train.py",
    source_dir="src/",
    role=role,
    instance_type="ml.m5.xlarge",
    framework_version="1.2-1",
    py_version="py3",
    output_path=f"s3://{bucket}/churn/models/",
    hyperparameters={
        "n-estimators": 200,
        "max-depth": 8,
    },
)
sk.fit({"train": train_input, "validation": val_input})

The source_dir is a folder that gets bundled and uploaded to S3, then extracted inside the container. Add a requirements.txt in that folder and SageMaker installs those dependencies before running the script. That is the polite way to add imbalanced-learn or shap without touching an image.

What the script sees

Inside the container, SageMaker exposes the same interface for every framework, through environment variables and fixed paths. Learning them is a fifteen-minute investment that pays off every project.

Variable	Points to	Contents
SM_CHANNEL_TRAIN	/opt/ml/input/data/train	Files from the train channel
SM_CHANNEL_VALIDATION	/opt/ml/input/data/validation	Files from the validation channel
SM_MODEL_DIR	/opt/ml/model	Where to save the model; auto-uploaded to S3
SM_OUTPUT_DATA_DIR	/opt/ml/output/data	Extra outputs, tarballed to S3
SM_NUM_GPUS, SM_NUM_CPUS	—	Instance capabilities
SM_HPS	JSON	All hyperparameters as a single object

The corresponding train.py, dead-simple:

import argparse, os, json
import pandas as pd
import joblib
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import roc_auc_score

def parse():
    p = argparse.ArgumentParser()
    p.add_argument("--n-estimators", type=int, default=100)
    p.add_argument("--max-depth", type=int, default=6)
    p.add_argument("--train", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
    p.add_argument("--validation", type=str, default=os.environ["SM_CHANNEL_VALIDATION"])
    p.add_argument("--model-dir", type=str, default=os.environ["SM_MODEL_DIR"])
    return p.parse_args()

def load(path):
    frames = [pd.read_parquet(os.path.join(path, f))
              for f in os.listdir(path) if f.endswith(".parquet")]
    df = pd.concat(frames, ignore_index=True)
    return df.drop(columns=["churn"]), df["churn"]

if __name__ == "__main__":
    args = parse()
    X_tr, y_tr = load(args.train)
    X_val, y_val = load(args.validation)

    model = GradientBoostingClassifier(
        n_estimators=args.n_estimators, max_depth=args.max_depth
    ).fit(X_tr, y_tr)

    auc = roc_auc_score(y_val, model.predict_proba(X_val)[:, 1])
    print(f"validation-auc: {auc:.4f}")   # picked up by the log metric parser

    joblib.dump(model, os.path.join(args.model_dir, "model.joblib"))

Two points to notice. Hyperparameters passed to the estimator arrive as command-line arguments, with dashes converted from underscores — n-estimators from the SDK becomes --n-estimators for argparse. And the model must be written to SM_MODEL_DIR; anything left elsewhere is discarded when the container stops. Files written to SM_MODEL_DIR are tarred into model.tar.gz and uploaded to the estimator's output_path.

Exposing metrics for the console and tuning

The training console shows exactly the metrics you tell it about. Declare them once with a regex on the standard output:

sk = SKLearn(
    ...,
    metric_definitions=[
        {"Name": "validation:auc", "Regex": r"validation-auc:\s*([0-9\.]+)"},
    ],
)

The Regex group ([0-9\.]+) captures the number that follows validation-auc: in the script's print. That value is what module 6 uses as the objective for tuning; without a metric definition, the tuning job has nothing to optimize.

When script mode is not enough: the fully custom image

There is a small set of cases the framework containers do not cover: a system library needed at runtime (libgeos for geopandas), a compiled binary (OR-Tools), an unusual framework version, or a corporate wrapper that must be installed the same way in dev and in prod. The answer is a fully custom image.

FROM public.ecr.aws/deep-learning-containers/sklearn:1.2-1-cpu-py3

USER root
RUN apt-get update && apt-get install -y --no-install-recommends libgomp1 \
    && rm -rf /var/lib/apt/lists/*

COPY train.py /opt/ml/code/train.py
ENV SAGEMAKER_PROGRAM=train.py

Push the image to ECR, pass its URI as image_uri= to a generic Estimator, and SageMaker treats it like a built-in. The interface is the same — the /opt/ml/... paths and the SM_* variables — because it is what SageMaker inspects; the base image only has to respect it.

Local mode: debugging without paying

Building an image and pushing it to ECR each iteration is a slow feedback loop. SageMaker's local mode runs the same container on your laptop's Docker, with the exact same environment variables, before you launch anything on AWS.

sk = SKLearn(
    entry_point="train.py",
    source_dir="src/",
    role=role,
    instance_type="local",           # the magic value
    framework_version="1.2-1",
)
sk.fit({"train": "file://./data/train", "validation": "file://./data/validation"})

file:// paths and instance_type="local" are the only two changes. Loop on local mode until the script exits with the metric printed, then switch back to ml.m5.xlarge for the real run. A pushed-to-cloud iteration that would take five minutes takes fifteen seconds locally; that is not a minor difference on a script you are debugging.

Every hyperparameter has to arrive as a string

SageMaker sends hyperparameters to the container as strings; argparse with type=int casts them back. A script that writes if args.max_depth == 6 when it was passed "6" will fail silently. This surfaces most often on booleans: type=bool accepts any non-empty string as True. Use type=lambda s: s.lower() == "true" explicitly.

Summary

• Script mode is the default: your script inside AWS's framework container, with requirements.txt for extra dependencies.
• The script talks to SageMaker through a fixed set of environment variables and paths — SM_CHANNEL_*, SM_MODEL_DIR.
• metric_definitions regexes turn print() lines into structured metrics for the console and for tuning.
• A fully custom image covers system dependencies; local mode shortens the debug loop from minutes to seconds.

Next module: let SageMaker choose the hyperparameters for us.