Integrating Model Training Steps

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.

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.

Before we add the training step to our pipeline, let's understand what our training script does. Our train.py script contains the actual machine learning code that will execute during the training step. This script follows SageMaker's training script conventions, which means it knows how to read input data from specific locations and save the trained model to the correct output directory.

Here's what our training script looks like:

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!")

The main training logic follows the same machine learning workflow you've used before: loading the processed data, separating features from the target variable, training a Linear Regression model, evaluating its performance, and saving the trained model using joblib. The key difference is that this script is structured to work within SageMaker's training environment, reading from designated input channels and saving to the correct output directory.