Integrating Model Training Steps

Introduction & Lesson Context

Welcome back! You've made excellent progress in your journey to master SageMaker Pipelines . In the first lesson, you successfully built your first pipeline with a data preprocessing step that transforms raw California housing data into clean, processed datasets ready for machine learning. Then, in the second lesson, you learned the critical skill of monitoring pipeline executions, tracking their progress, and diagnosing any issues that might arise during execution. Now you're ready to take the next significant step in building automated ML workflows. In this lesson, you'll expand your existing pipeline by adding a model training step that seamlessly connects to your preprocessing output. This integration represents a fundamental concept in MLOps: creating end-to-end workflows where data flows automatically from one stage to the next without manual intervention. By the end of this lesson, you'll have a complete pipeline that takes raw data, processes it, and trains a machine learning model — all in a single, automated workflow. This expanded pipeline will demonstrate how SageMaker Pipelines can orchestrate complex ML workflows while maintaining clear dependencies between steps. You'll learn how to configure training steps, connect pipeline outputs to inputs, and ensure your workflow runs smoothly from start to finish.

What We've Built So Far

Let's quickly recap what we've already accomplished. Our current pipeline contains a single data preprocessing step that handles the California housing dataset beautifully. Our preprocessing step uses an SKLearnProcessor to run our data_processing.py script, which loads the raw housing data, performs feature scaling and encoding, splits the data into training and test sets, and saves the processed datasets to designated S3 locations. This step takes raw CSV data as input and produces two outputs: processed training data and processed test data, both stored in S3 with clear reference names. This preprocessing foundation is perfect for adding model training because it produces clean, processed data in exactly the format that a training algorithm expects. The output structure we established — with separate training and test datasets — follows ML best practices and makes it straightforward to connect additional pipeline steps.

Our Training Script

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

import os import argparse import numpy as np import joblib import pandas as pd from sklearn.linear_model import LinearRegression from sklearn.metrics import mean_squared_error, r2_score if __name__ == '__main__': parser = argparse.ArgumentParser() 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')) args = parser.parse_args() # Load training data from SageMaker's input channel train_data_path = os.path.join(args.train, 'train.csv') df = pd.read_csv(train_data_path) # Apply the same machine learning steps you've done before: # 1. Separate features and target (same as you've done before) # 2. Create and train a Linear Regression model (standard scikit-learn approach) # 3. Evaluate model performance on training data # Save the trained model to SageMaker's model directory joblib.dump(model, os.path.join(args.model_dir, 'model.joblib')) print("Training completed!")

Configuring a SKLearn Estimator

Now let's add the training components to our existing pipeline. First, we need to create an estimator that defines how our model training will be executed. Our SKLearn estimator will specify the training environment, including our entry point script, computational resources, and framework version:

Python

from sagemaker.sklearn.estimator import SKLearn

# Create an estimator that defines how our model will be trained
estimator = SKLearn(
    entry_point="train.py",             # Our training script
    role=SAGEMAKER_ROLE,                # IAM role with necessary permissions
    instance_type="ml.m5.large",        # Compute instance type for training
    instance_count=1,                   # Number of instances to use
    framework_version="1.2-1",          # Specify scikit-learn version
    py_version="py3",                   # Python version
    sagemaker_session=pipeline_session  # Use pipeline session for deferred execution
)

from sagemaker.sklearn.estimator import SKLearn

# Create an estimator that defines how our model will be trained
estimator = SKLearn(
    entry_point="train.py",             # Our training script
    role=SAGEMAKER_ROLE,                # IAM role with necessary permissions
    instance_type="ml.m5.large",        # Compute instance type for training
    instance_count=1,                   # Number of instances to use
    framework_version="1.2-1",          # Specify scikit-learn version
    py_version="py3",                   # Python version
    sagemaker_session=pipeline_session  # Use pipeline session for deferred execution
)

This estimator configuration tells SageMaker to use our train.py script in a scikit-learn environment with the specified computational resources. Just as we used the pipeline_session with our processor, we use it here to ensure the training job will be executed as part of our pipeline workflow rather than immediately when the estimator is created. This estimator acts as a blueprint for how the training job should be executed when our pipeline runs.

Creating the TrainingStep

With our estimator configured, we can now create the TrainingStep that orchestrates model training within our pipeline. A TrainingStep is a pipeline component that combines an estimator (which defines the training environment) with input data sources. The key to building effective pipelines lies in properly connecting outputs from one step to inputs of another using SageMaker's step properties system.

Python

from sagemaker.workflow.steps import TrainingStep

# Define our training step that uses output from our processing step
training_step = TrainingStep(
    name="TrainModel",    # Unique name for this step in our pipeline
    estimator=estimator,  # Our estimator we defined above
    inputs={
        # Use the training data output from our processing step as input
        "train": sagemaker.inputs.TrainingInput(
            # Reference the S3 URI where our processing step saved the training data
            s3_data=processing_step.properties.ProcessingOutputConfig.Outputs["train_data"].S3Output.S3Uri
        )
    }
)

from sagemaker.workflow.steps import TrainingStep

# Define our training step that uses output from our processing step
training_step = TrainingStep(
    name="TrainModel",    # Unique name for this step in our pipeline
    estimator=estimator,  # Our estimator we defined above
    inputs={
        # Use the training data output from our processing step as input
        "train": sagemaker.inputs.TrainingInput(
            # Reference the S3 URI where our processing step saved the training data
            s3_data=processing_step.properties.ProcessingOutputConfig.Outputs["train_data"].S3Output.S3Uri
        )
    }
)

The critical connection happens in the inputs parameter. The "train" key in our inputs dictionary corresponds to the SM_CHANNEL_TRAIN environment variable that our training script uses to locate input data. Remember in our training script where we had args.train = os.environ.get('SM_CHANNEL_TRAIN')?

That environment variable gets populated with the local path where SageMaker downloads our training data. When the training job runs, SageMaker will automatically download the data from the S3 location we specify and make it available to our script at the path stored in SM_CHANNEL_TRAIN.

Now let's break down how we specify which S3 location to use for the s3_data parameter:

processing_step - References our preprocessing step that we defined earlier
.properties - Accesses the runtime properties of our preprocessing step
.ProcessingOutputConfig - References the output configuration of the processing job
.Outputs - Accesses the collection of all outputs from the processing step
["train_data"] - Selects the specific output we named "train_data" when we defined our processing step's outputs
.S3Output - References the S3 output configuration for this specific output
.S3Uri - Gets the actual S3 URI where this output was saved

This property chain essentially says: "Use the S3 location where the processing step saved its 'train_data' output as the input for this training step." The beauty of this approach is that SageMaker resolves these property references at runtime, so you don't need to know the exact S3 paths beforehand.

Let's Update Our Pipeline

Finally, let's update our pipeline creation to include both our preprocessing and training steps:

Python

# Create our pipeline by combining all steps
pipeline = Pipeline(
    name=PIPELINE_NAME,
    steps=[processing_step, training_step],  # Include both our steps
    sagemaker_session=sagemaker_session
)

try:
    # Create or update our pipeline (upsert = update if exists, create if not)
    pipeline.upsert(role_arn=SAGEMAKER_ROLE)
    
    # Start our pipeline execution and get execution object for monitoring
    execution = pipeline.start()

    # Print the unique ARN identifier for this execution
    print(f"Pipeline execution ARN: {execution.arn}")

    # Get detailed information about our execution
    execution_details = execution.describe()

    # Display the current status of our pipeline execution
    print(f"Status: {execution_details['PipelineExecutionStatus']}")
    
except Exception as e:
    print(f"Error: {e}")

# Create our pipeline by combining all steps
pipeline = Pipeline(
    name=PIPELINE_NAME,
    steps=[processing_step, training_step],  # Include both our steps
    sagemaker_session=sagemaker_session
)

try:
    # Create or update our pipeline (upsert = update if exists, create if not)
    pipeline.upsert(role_arn=SAGEMAKER_ROLE)
    
    # Start our pipeline execution and get execution object for monitoring
    execution = pipeline.start()

    # Print the unique ARN identifier for this execution
    print(f"Pipeline execution ARN: {execution.arn}")

    # Get detailed information about our execution
    execution_details = execution.describe()

    # Display the current status of our pipeline execution
    print(f"Status: {execution_details['PipelineExecutionStatus']}")
    
except Exception as e:
    print(f"Error: {e}")

When we execute this complete pipeline, we'll see output similar to:

text

Pipeline execution ARN: arn:aws:sagemaker:us-east-1:123456789012:pipeline/california-housing-training-pipeline/execution/xyz789-1234-5678-abcd-efghijklmnop
Status: Executing

Pipeline execution ARN: arn:aws:sagemaker:us-east-1:123456789012:pipeline/california-housing-training-pipeline/execution/xyz789-1234-5678-abcd-efghijklmnop
Status: Executing

Our pipeline will first execute the preprocessing step, then automatically proceed to the training step once preprocessing completes successfully.

Introduction & Lesson Context

What We've Built So Far

Our Training Script

Join the 1M+ learners on CodeSignal

Configuring a SKLearn Estimator

Creating the TrainingStep

Let's Update Our Pipeline

Understanding Pipeline Execution Order

Summary & Next Steps