AWS SageMaker — Training Jobs and Pipelines

Think of SageMaker as a managed MLOps platform — it is like having AWS automatically run your MLflow tracking server, Kubernetes cluster, model registry, and feature store so you only pay for what you use and never manage the servers. Your training code stays the same; SageMaker is just the environment it runs in.

# Install
pip install sagemaker boto3

# ── Connect ──────────────────────────────────────────────────────────
import boto3
import sagemaker
from sagemaker import get_execution_role

# Works automatically inside SageMaker Studio or a SageMaker Notebook.
# Locally, set AWS_DEFAULT_REGION + credentials in ~/.aws/credentials
sess   = sagemaker.Session()
role   = get_execution_role()          # IAM role ARN
bucket = sess.default_bucket()         # auto-created S3 bucket
region = boto3.Session().region_name

print(f"Role:   {role}")
print(f"Bucket: {bucket}")
print(f"Region: {region}")

Every SageMaker job runs as an IAM role, not as your personal AWS user. The role needs AmazonSageMakerFullAccess plus S3 read/write on your training bucket. In production, create a least-privilege role that only grants access to the specific S3 paths and ECR repositories your jobs need.

import os
from sagemaker.s3 import S3Uploader

# Upload local data directory to S3
local_data = './data/processed/'
s3_data_uri = S3Uploader.upload(
    local_path=local_data,
    desired_s3_uri=f's3://{bucket}/freshmart/processed',
    sagemaker_session=sess,
)
print(f"Data uploaded to: {s3_data_uri}")
# → s3://sagemaker-us-east-1-123456789/freshmart/processed

# train.py — runs identically locally or inside SageMaker
import argparse, os, pickle, json
import pandas as pd
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.metrics import mean_absolute_error, r2_score
import mlflow, mlflow.sklearn

# SageMaker passes hyperparameters as CLI args
parser = argparse.ArgumentParser()
parser.add_argument('--n-estimators',  type=int,   default=200)
parser.add_argument('--max-depth',     type=int,   default=4)
parser.add_argument('--learning-rate', type=float, default=0.05)

# SageMaker injects these environment variables automatically
parser.add_argument('--model-dir',  default=os.environ.get('SM_MODEL_DIR',    './model'))
parser.add_argument('--train',      default=os.environ.get('SM_CHANNEL_TRAIN', './data'))
args = parser.parse_args()

# Load data from the SM_CHANNEL_TRAIN directory (mounted from S3)
df       = pd.read_parquet(os.path.join(args.train, 'train.parquet'))
feat_cols = [c for c in df.columns if c != 'delivery_minutes']
X, y     = df[feat_cols].values, df['delivery_minutes'].values

# Train
model = GradientBoostingRegressor(
    n_estimators=args.n_estimators,
    max_depth=args.max_depth,
    learning_rate=args.learning_rate,
    random_state=42,
)
model.fit(X, y)
mae = mean_absolute_error(y, model.predict(X))
r2  = r2_score(y, model.predict(X))

# Log to MLflow (SageMaker Experiments auto-patches mlflow calls)
mlflow.log_metric('train_mae', mae)
mlflow.log_metric('train_r2',  r2)

# Save — SageMaker uploads SM_MODEL_DIR to S3 automatically
os.makedirs(args.model_dir, exist_ok=True)
with open(os.path.join(args.model_dir, 'model.pkl'), 'wb') as f:
    pickle.dump({'model': model, 'feature_cols': feat_cols}, f)

print(f"Done: MAE={mae:.4f}  R2={r2:.4f}")

SageMaker mounts each input channel as a directory and setsSM_CHANNEL_<NAME>. Pass inputs={{'train': s3_uri}} in the Estimator and your script sees SM_CHANNEL_TRAIN pointing to that data — no boto3 download code needed.

from sagemaker.sklearn import SKLearn

estimator = SKLearn(
    entry_point='train.py',       # your script
    source_dir='./src',           # directory containing train.py
    role=role,
    instance_type='ml.m5.xlarge', # 4 vCPU, 16 GB — SageMaker provisions this
    instance_count=1,
    framework_version='1.2-1',
    py_version='py3',

    # Hyperparameters — passed as CLI args to train.py
    hyperparameters={
        'n-estimators':  200,
        'max-depth':     4,
        'learning-rate': 0.05,
    },

    # Experiment tracking — auto-logs metrics to SageMaker Experiments
    enable_sagemaker_metrics=True,
    sagemaker_session=sess,
)

# Submit — provisions instance, runs script, uploads model to S3
estimator.fit(
    inputs={'train': s3_data_uri},
    job_name='freshmart-delivery-v1',
    wait=True,   # block until complete; set False for async
    logs=True,
)

print(f"Model artifact: {estimator.model_data}")
# → s3://sagemaker-.../output/freshmart-delivery-v1/output/model.tar.gz

# Describe a running or completed job
aws sagemaker describe-training-job \
  --training-job-name freshmart-delivery-v1 \
  --query 'TrainingJobStatus'

# Stream logs in real time
aws logs tail /aws/sagemaker/TrainingJobs \
  --log-stream-name-prefix freshmart-delivery-v1 \
  --follow

from sagemaker.workflow.pipeline import Pipeline
from sagemaker.workflow.steps import ProcessingStep, TrainingStep
from sagemaker.workflow.pipeline_context import PipelineSession
from sagemaker.workflow.parameters import ParameterInteger, ParameterFloat, ParameterString
from sagemaker.sklearn.processing import SKLearnProcessor
from sagemaker.processing import ProcessingInput, ProcessingOutput
from sagemaker.sklearn import SKLearn

pipeline_sess = PipelineSession()

# ── Pipeline parameters — can be overridden at runtime ───────────────
n_estimators  = ParameterInteger(name='NEstimators',  default_value=200)
max_depth     = ParameterInteger(name='MaxDepth',     default_value=4)
learning_rate = ParameterFloat(  name='LearningRate', default_value=0.05)
input_data    = ParameterString( name='InputData',    default_value=s3_data_uri)

# ── Step 1: Processing — clean and featurise raw data ────────────────
processor = SKLearnProcessor(
    framework_version='1.2-1',
    role=role,
    instance_type='ml.m5.large',
    instance_count=1,
    sagemaker_session=pipeline_sess,
)

step_process = ProcessingStep(
    name='PrepareData',
    processor=processor,
    inputs=[ProcessingInput(source=input_data, destination='/opt/ml/processing/input')],
    outputs=[ProcessingOutput(output_name='train', source='/opt/ml/processing/output/train')],
    code='./src/prepare.py',
)

# ── Step 2: Training — train on processed data ────────────────────────
estimator = SKLearn(
    entry_point='train.py',
    source_dir='./src',
    role=role,
    instance_type='ml.m5.xlarge',
    framework_version='1.2-1',
    hyperparameters={
        'n-estimators':  n_estimators,
        'max-depth':     max_depth,
        'learning-rate': learning_rate,
    },
    sagemaker_session=pipeline_sess,
)

step_train = TrainingStep(
    name='TrainModel',
    estimator=estimator,
    inputs={
        'train': sagemaker.inputs.TrainingInput(
            s3_data=step_process.properties.ProcessingOutputConfig
                    .Outputs['train'].S3Output.S3Uri,
        )
    },
)

# ── Assemble and upsert pipeline ──────────────────────────────────────
pipeline = Pipeline(
    name='freshmart-retraining-pipeline',
    parameters=[n_estimators, max_depth, learning_rate, input_data],
    steps=[step_process, step_train],
    sagemaker_session=pipeline_sess,
)

pipeline.upsert(role_arn=role)   # create or update

# ── Execute ───────────────────────────────────────────────────────────
execution = pipeline.start(
    parameters={
        'NEstimators':  300,
        'LearningRate': 0.03,
    }
)
execution.wait()
print(execution.list_steps())

Add cache_config=CacheConfig(enable_caching=True, expire_after='30d')to any step. If the step inputs (data URI + code hash) match a previous run within the expiry window, SageMaker reuses the output artifact instead of re-running. A daily retraining pipeline that detects no new data will skip the expensive training step automatically.

import boto3, time
from sagemaker.feature_store.feature_group import FeatureGroup
from sagemaker.feature_store.feature_definition import (
    FeatureDefinition, FeatureTypeEnum,
)

# ── Define a Feature Group ────────────────────────────────────────────
feature_group = FeatureGroup(
    name='freshmart-delivery-features',
    sagemaker_session=sess,
    feature_definitions=[
        FeatureDefinition('order_id',             FeatureTypeEnum.STRING),
        FeatureDefinition('postcode',             FeatureTypeEnum.STRING),
        FeatureDefinition('hour_of_day',          FeatureTypeEnum.INTEGRAL),
        FeatureDefinition('day_of_week',          FeatureTypeEnum.INTEGRAL),
        FeatureDefinition('items_in_order',       FeatureTypeEnum.INTEGRAL),
        FeatureDefinition('avg_delivery_postcode',FeatureTypeEnum.FRACTIONAL),
        FeatureDefinition('event_time',           FeatureTypeEnum.STRING),  # required
    ],
)

feature_group.create(
    s3_uri=f's3://{bucket}/feature-store',
    record_identifier_name='order_id',
    event_time_feature_name='event_time',
    role_arn=role,
    enable_online_store=True,   # DynamoDB backing for low-latency reads
)

# Wait until active
while feature_group.describe()['FeatureGroupStatus'] != 'Created':
    time.sleep(5)

# ── Ingest features ───────────────────────────────────────────────────
import pandas as pd

df = pd.read_parquet('features.parquet')
df['event_time'] = pd.Timestamp.utcnow().isoformat()

feature_group.ingest(
    data_frame=df,
    max_workers=4,
    wait=True,
)

# ── Read at inference time (online store) ─────────────────────────────
sm_runtime = boto3.client('sagemaker-featurestore-runtime')

record = sm_runtime.get_record(
    FeatureGroupName='freshmart-delivery-features',
    RecordIdentifierValueAsString='order-12345',
)
features = {f['FeatureName']: f['ValueAsString'] for f in record['Record']}
print(features)

from sagemaker import clarify

clarify_processor = clarify.SageMakerClarifyProcessor(
    role=role,
    instance_count=1,
    instance_type='ml.m5.xlarge',
    sagemaker_session=sess,
)

# ── Data config — where is the dataset ────────────────────────────────
data_config = clarify.DataConfig(
    s3_data_input_path=s3_data_uri,
    s3_output_path=f's3://{bucket}/clarify-output',
    label='delivery_minutes',
    headers=feat_cols + ['delivery_minutes'],
    dataset_type='text/csv',
)

# ── Model config — wrap the trained model ─────────────────────────────
model_config = clarify.ModelConfig(
    model_name='freshmart-delivery-model',   # SageMaker model name
    instance_type='ml.m5.large',
    instance_count=1,
    accept_type='text/csv',
)

# ── Bias config — which column / value is the sensitive group ──────────
bias_config = clarify.BiasConfig(
    label_values_or_threshold=[1],           # positive outcome: on-time delivery
    facet_name='postcode',                   # sensitive attribute to audit
    facet_values_or_threshold=['SW1A'],      # group of interest
)

# ── SHAP config — feature importance ──────────────────────────────────
shap_config = clarify.SHAPConfig(
    baseline=[df[feat_cols].mean().tolist()],  # baseline: feature means
    num_samples=100,
    agg_method='mean_abs',
    save_local_shap_values=True,
)

# Run the analysis
clarify_processor.run_bias(
    data_config=data_config,
    bias_config=bias_config,
    model_config=model_config,
    pre_training_methods='all',
    post_training_methods='all',
)

clarify_processor.run_explainability(
    data_config=data_config,
    model_config=model_config,
    explainability_config=shap_config,
)

from sagemaker.jumpstart.model import JumpStartModel

# ── Deploy a pre-trained model ────────────────────────────────────────
# model_id list: https://sagemaker.readthedocs.io/en/stable/doc_utils/jumpstart.html
model = JumpStartModel(
    model_id='huggingface-text2text-flan-t5-xl',
    role=role,
    sagemaker_session=sess,
)

predictor = model.deploy(
    initial_instance_count=1,
    instance_type='ml.g4dn.xlarge',   # GPU instance
)

# ── Run inference ─────────────────────────────────────────────────────
response = predictor.predict({
    'text_inputs': 'Summarise this delivery exception: order arrived 45 min late due to traffic.',
    'max_length': 100,
})
print(response[0]['generated_text'])

# ── Fine-tune on your own data ────────────────────────────────────────
from sagemaker.jumpstart.estimator import JumpStartEstimator

estimator = JumpStartEstimator(
    model_id='huggingface-text2text-flan-t5-xl',
    role=role,
    instance_type='ml.g4dn.xlarge',
    hyperparameters={'epoch': 3, 'per_device_train_batch_size': 4},
)

estimator.fit({'training': s3_fine_tune_uri})   # your labelled JSONL data

# Deploy the fine-tuned model
ft_predictor = estimator.deploy(instance_type='ml.g4dn.xlarge')

# ── Clean up ──────────────────────────────────────────────────────────
predictor.delete_endpoint()
ft_predictor.delete_endpoint()

JumpStart is like PyPI for production ML models. Instead ofpip install transformers and writing a Flask server, you write JumpStartModel(model_id='...').deploy() and AWS handles the container, GPU driver, scaling, and HTTPS endpoint.

# ── Rough monthly cost comparison: SageMaker vs self-managed EC2 ──────
# Scenario: 20 training jobs/month, each on an equivalent instance for 2 hours

sagemaker_price_per_hr = 0.23    # ml.m5.xlarge, SageMaker managed training
ec2_price_per_hr       = 0.192   # equivalent on-demand EC2 instance
jobs_per_month = 20
hours_per_job  = 2

sagemaker_compute_cost = sagemaker_price_per_hr * hours_per_job * jobs_per_month
ec2_compute_cost       = ec2_price_per_hr * hours_per_job * jobs_per_month

print(f"SageMaker compute:  ${sagemaker_compute_cost:.2f}/month")
print(f"Raw EC2 compute:    ${ec2_compute_cost:.2f}/month")
print(f"SageMaker premium:  ${sagemaker_compute_cost - ec2_compute_cost:.2f}/month for managed provisioning")

# But raw EC2 does not include, out of the box:
#  - Job scheduling and queueing        (someone builds this)
#  - Experiment tracking                (someone runs MLflow)
#  - Auto-shutdown on idle              (someone writes this, or pays for idle time)
#  - Endpoint autoscaling               (someone configures it)
#
# Engineering time to build and maintain that stack: typically 0.25-0.5 FTE ongoing.
# At a loaded engineer cost of roughly $150k/year, 0.25 FTE is about $3,125/month —
# which dwarfs the tiny managed premium at this job volume.

print("\nAt low job volume, the managed premium is negligible next to the")
print("engineering time it saves. At very high, steady-state volume — hundreds")
print("of long-running jobs a month — the premium compounds, and self-managed")
print("infra can become cheaper. The crossover point depends on job volume AND")
print("how much of the surrounding MLOps stack you would otherwise build.")

Nothing stops you from running training scripts on raw EC2 instances, or on self-managed Kubernetes with open-source MLflow, Airflow, and KServe. SageMaker is a convenience layer bundling those tools as a managed service — it is optional. Plenty of AWS-based ML teams, especially ones that already run Kubernetes for their main product, never touch SageMaker at all.

SageMaker removes server provisioning and patching, but you still own IAM role scoping, cost monitoring — a forgotten endpoint keeps billing whether or not it receives traffic — quota management (GPU instance types need an explicit quota request), pipeline debugging, and model quality. None of that is automated away. Managed narrows the operational surface; it does not remove it.

SageMaker's per-second managed compute pricing carries a premium over the identical raw EC2 instance type, and that premium compounds at high, sustained volume. It usually wins for teams that would otherwise spend engineering time building scheduling, tracking, and autoscaling from scratch — but for an org already running that infrastructure for other workloads, adding SageMaker on top can mean paying for infrastructure twice.

Every SageMaker component is independently addressable. Plenty of teams use only Training Jobs for managed compute and keep their existing MLflow tracking, Airflow scheduling, and custom feature pipeline exactly as they were. Treating SageMaker as an all-or- nothing platform migration is a self-imposed constraint, not a technical one.

JumpStart and AutoML remove the work of writing training code and picking an architecture, not the work of evaluating whether the resulting model is actually good enough, unbiased, and safe for your specific data and use case. Clarify's bias and explainability reports exist precisely because a model that trained without errors is not the same thing as a model that is fit to ship.

Managed provisioning — request a training job, get an instance, pay per second, the instance terminates automatically with no idle server — a built-in experiment tracking and model registry integrated with the console, pipeline orchestration with step-level caching, and Model Monitor and Clarify for drift and bias detection without standing up your own Evidently or SHAP pipeline. None of these are capabilities you cannot build yourself — the value is not having to re-engineer and maintain them.

When a team already runs Kubernetes for its core product with a platform team maintaining it — adding SageMaker means a second operational surface instead of extending the one already paid for. Also when there is a hard compliance requirement SageMaker's managed defaults do not satisfy, or when job volume is high and steady enough that the managed compute premium meaningfully exceeds the engineering cost of running it yourself.

Use Managed Spot Training for jobs that can tolerate interruption — up to 90 percent off the on-demand price, with SageMaker automatically resuming from the last checkpoint. Set short max_run limits so a bug cannot leave a job running for days. Right-size instance types by profiling actual CPU and memory usage instead of guessing large. Check whether Serverless Inference or a smaller endpoint instance covers a low-traffic model instead of an always-on instance. Delete idle endpoints — a forgotten real-time endpoint bills continuously whether or not it receives traffic.

I would ask what specific gap the existing stack has before adding a second orchestrator. Pipelines' step caching and visual DAG are genuinely nice, but running two orchestration systems means two places to debug when a retraining job fails, two credential surfaces, and no shared history across old and new pipelines. Unless there is a concrete reason the existing stack cannot meet, extending the existing Airflow and MLflow setup usually costs less than standing up and maintaining a parallel one.

Keep training and inference code framework-native — plain scikit-learn or PyTorch scripts, not code that calls SageMaker-specific APIs inside the model logic — so the same script runs unchanged on EC2, on-prem, or another cloud's equivalent service. SageMaker's SDK wraps the script rather than replacing it, which is exactly why the training scripts in this module look identical to local development. Keep model artifacts in open formats rather than a proprietary one, and treat Feature Store and Pipelines as replaceable orchestration layers around portable code, not as the source of truth for your ML logic.