Connecting AWS Glue and SageMaker for ML Pipelines

by

Machine learning pipelines in production require more than just model training. The reality is that data scientists spend roughly 80% of their time on data preparation, transformation, and feature engineering before they can even begin training models. This is where the combination of AWS Glue and Amazon SageMaker becomes transformative. While SageMaker excels at machine learning workloads, AWS Glue provides the robust ETL capabilities needed to prepare data at scale. Together, they form a powerful duo for building end-to-end ML pipelines that handle everything from raw data ingestion to model deployment.

In this article, we’ll explore how to effectively connect AWS Glue and SageMaker, walking through practical implementation patterns, cost considerations, and real-world orchestration strategies that data engineers and ML practitioners can apply immediately.

Understanding the Complementary Roles

Before diving into implementation, it’s crucial to understand why these services work so well together and where each shines.

AWS Glue’s Strengths:

  • Serverless ETL processing with Apache Spark under the hood
  • Automatic schema discovery and data cataloging
  • Native integration with data lakes on S3, data warehouses, and databases
  • Cost-effective for large-scale data transformations
  • Built-in data quality checks and validation
  • Support for both batch and streaming data

Amazon SageMaker’s Strengths:

  • Optimized infrastructure for ML training and inference
  • Pre-built algorithms and framework support (PyTorch, TensorFlow, XGBoost)
  • Feature Store for centralized feature management
  • Model registry and versioning
  • Built-in model monitoring and drift detection
  • Flexible deployment options from real-time endpoints to batch transforms

The key insight is that Glue prepares the data while SageMaker handles the ML workload. This separation of concerns creates cleaner architectures, better cost optimization, and more maintainable pipelines.

💡 Architecture Pattern

Glue → S3 → SageMaker → S3 → Glue

1. Glue extracts and transforms raw data
2. Cleaned data lands in S3
3. SageMaker trains or runs inference
4. Results stored back to S3
5. Glue catalogs and post-processes for consumption

Setting Up the Foundation: S3 as the Central Hub

Amazon S3 serves as the glue (pun intended) between these services. Both AWS Glue and SageMaker read from and write to S3 buckets, making it the natural integration point. Your S3 structure should reflect your ML pipeline stages:

s3://your-ml-bucket/
├── raw/                    # Raw ingested data
├── processed/              # Glue output - cleaned data
├── features/               # Feature-engineered datasets
├── training/               # SageMaker training input
├── models/                 # Trained model artifacts
├── inference-input/        # Data for batch predictions
└── inference-output/       # Prediction results

This structure provides clear data lineage and makes it easy to debug pipeline issues. When a model underperforms, you can trace back through each stage to identify where data quality issues originated.

Setting up proper IAM permissions is critical. Your Glue jobs need permissions to read raw data and write to processed locations. SageMaker needs to read processed data and write models and predictions. Create separate IAM roles for each service following the principle of least privilege:

  • Glue role: S3 read/write access, Glue catalog access, CloudWatch logs
  • SageMaker role: S3 read/write access, SageMaker full access, CloudWatch logs
  • Never share roles between services in production environments

Data Preparation with AWS Glue

The first stage in your ML pipeline is data preparation. AWS Glue excels at this through its PySpark-based transformation engine. Let’s walk through a practical example of preparing customer transaction data for a churn prediction model.

Creating a Glue Job for ML Data Prep:

Your Glue script needs to handle several common ML data preparation tasks: cleaning missing values, encoding categorical variables, feature engineering, and data splitting. Here’s a complete example:

import sys
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
from pyspark.context import SparkContext
from awsglue.context import GlueContext
from awsglue.job import Job
from pyspark.sql.functions import col, when, datediff, current_date
from pyspark.ml.feature import StringIndexer, VectorAssembler

args = getResolvedOptions(sys.argv, ['JOB_NAME', 'OUTPUT_BUCKET'])
sc = SparkContext()
glueContext = GlueContext(sc)
spark = glueContext.spark_session
job = Job(glueContext)
job.init(args['JOB_NAME'], args)

# Read from Glue Data Catalog
customer_df = glueContext.create_dynamic_frame.from_catalog(
    database="ml_database",
    table_name="customer_transactions"
).toDF()

# Feature engineering for churn prediction
prepared_df = customer_df \
    .withColumn('account_age_days', datediff(current_date(), col('signup_date'))) \
    .withColumn('avg_transaction_value', col('total_spent') / col('transaction_count')) \
    .withColumn('is_premium', when(col('subscription_tier') == 'premium', 1).otherwise(0)) \
    .fillna({'transaction_count': 0, 'total_spent': 0}) \
    .drop('signup_date', 'customer_email')  # Remove PII and non-features

# Encode categorical variables
indexer = StringIndexer(inputCol="region", outputCol="region_encoded")
prepared_df = indexer.fit(prepared_df).transform(prepared_df)

# Split data for training and validation (80/20 split)
train_df, validation_df = prepared_df.randomSplit([0.8, 0.2], seed=42)

# Write to S3 in format optimized for SageMaker
train_df.write.mode('overwrite').parquet(f"s3://{args['OUTPUT_BUCKET']}/training/train/")
validation_df.write.mode('overwrite').parquet(f"s3://{args['OUTPUT_BUCKET']}/training/validation/")

# Update Glue Data Catalog with output schema
glueContext.write_dynamic_frame.from_options(
    frame=DynamicFrame.fromDF(train_df, glueContext, "train"),
    connection_type="s3",
    connection_options={"path": f"s3://{args['OUTPUT_BUCKET']}/training/train/"},
    format="parquet"
)

job.commit()

This script demonstrates several important practices. First, it uses the Glue Data Catalog to read data, which provides automatic schema management. Second, it performs feature engineering that’s specific to your ML problem—in this case, calculating account age and average transaction values for churn prediction. Third, it properly splits data before training, ensuring your validation set is completely separate. Finally, it writes in Parquet format, which is columnar and highly efficient for SageMaker to read.

Data Quality Validation:

Before handing data off to SageMaker, implement quality checks in your Glue job. AWS Glue Data Quality rules can catch issues early:

from awsglue.dataqualitylib import DataQualityEvaluator

# Define quality rules
quality_rules = """
    Rules = [
        RowCount > 1000,
        ColumnValues "transaction_count" >= 0,
        ColumnValues "total_spent" >= 0,
        Completeness "customer_id" > 0.99,
        Uniqueness "customer_id" > 0.95
    ]
"""

# Evaluate data quality
evaluator = DataQualityEvaluator()
result = evaluator.evaluate(prepared_df, quality_rules)

if result.overall_status == "Failed":
    raise Exception(f"Data quality check failed: {result.failures}")

These checks ensure that your training data meets minimum quality standards. If transaction counts are negative or customer IDs aren’t unique, you’ll catch it before wasting compute resources on SageMaker training.

Training Models with SageMaker

Once Glue has prepared your data, SageMaker takes over for the ML workload. The handoff happens through S3—Glue writes the prepared data to the training location, and SageMaker reads from there.

Using SageMaker’s Built-in Algorithms:

For many use cases, SageMaker’s built-in algorithms provide excellent performance without requiring custom code. The XGBoost algorithm works well for the churn prediction scenario:

import boto3
import sagemaker
from sagemaker import image_uris
from sagemaker.inputs import TrainingInput

session = sagemaker.Session()
role = "arn:aws:iam::YOUR_ACCOUNT:role/SageMakerExecutionRole"
region = boto3.Session().region_name

# Get the XGBoost container
container = image_uris.retrieve("xgboost", region, "1.7-1")

# Define training configuration
xgb = sagemaker.estimator.Estimator(
    container,
    role,
    instance_count=1,
    instance_type='ml.m5.xlarge',
    output_path=f's3://your-ml-bucket/models/',
    sagemaker_session=session
)

# Set hyperparameters
xgb.set_hyperparameters(
    objective='binary:logistic',
    num_round=100,
    max_depth=5,
    eta=0.2,
    subsample=0.8,
    colsample_bytree=0.8,
    eval_metric='auc'
)

# Point to data prepared by Glue
train_input = TrainingInput(
    s3_data='s3://your-ml-bucket/training/train/',
    content_type='application/x-parquet'
)
validation_input = TrainingInput(
    s3_data='s3://your-ml-bucket/training/validation/',
    content_type='application/x-parquet'
)

# Start training
xgb.fit({'train': train_input, 'validation': validation_input})

Notice how we specify the content type as Parquet—this tells SageMaker how to parse the data that Glue produced. The training job will pull data from S3, train the model, and save the trained artifact back to S3 automatically.

Batch Inference for Predictions:

After training, you often need to score large datasets. SageMaker Batch Transform is designed for this, and it integrates seamlessly with Glue-prepared data:

from sagemaker.transformer import Transformer

# Create transformer from trained model
transformer = Transformer(
    model_name='your-churn-model',
    instance_count=1,
    instance_type='ml.c5.xlarge',
    output_path='s3://your-ml-bucket/inference-output/',
    sagemaker_session=session
)

# Run batch predictions on Glue-prepared data
transformer.transform(
    data='s3://your-ml-bucket/inference-input/',
    content_type='application/x-parquet',
    split_type='Line',
    join_source='Input'  # Include input features with predictions
)

transformer.wait()

The join_source='Input' parameter is particularly useful—it appends predictions to your original features, making it easy for downstream Glue jobs to enrich the results with additional context.

Orchestrating with Apache Airflow

In production, you need orchestration to run these services in sequence, handle failures, and manage dependencies. Apache Airflow has become the de facto standard, and it has native operators for both Glue and SageMaker.

Building an End-to-End DAG:

Here’s a complete Airflow DAG that orchestrates the entire pipeline:

from datetime import datetime, timedelta
from airflow import DAG
from airflow.providers.amazon.aws.operators.glue import GlueJobOperator
from airflow.providers.amazon.aws.operators.sagemaker import (
    SageMakerTrainingOperator,
    SageMakerTransformOperator
)
from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor

default_args = {
    'owner': 'data-team',
    'depends_on_past': False,
    'start_date': datetime(2024, 1, 1),
    'email_on_failure': True,
    'email_on_retry': False,
    'retries': 2,
    'retry_delay': timedelta(minutes=5)
}

dag = DAG(
    'ml_pipeline_glue_sagemaker',
    default_args=default_args,
    description='End-to-end ML pipeline with Glue and SageMaker',
    schedule_interval='@daily',
    catchup=False
)

# Wait for raw data to arrive
wait_for_data = S3KeySensor(
    task_id='wait_for_raw_data',
    bucket_name='your-ml-bucket',
    bucket_key='raw/transactions/dt={{ ds }}/',
    aws_conn_id='aws_default',
    timeout=3600,
    poke_interval=300,
    dag=dag
)

# Run Glue ETL job
prepare_data = GlueJobOperator(
    task_id='prepare_training_data',
    job_name='ML_Data_Preparation',
    script_args={
        '--OUTPUT_BUCKET': 'your-ml-bucket',
        '--PROCESSING_DATE': '{{ ds }}'
    },
    aws_conn_id='aws_default',
    dag=dag
)

# Train model (only run weekly)
train_model = SageMakerTrainingOperator(
    task_id='train_churn_model',
    config={
        'AlgorithmSpecification': {
            'TrainingImage': '683313688378.dkr.ecr.us-east-1.amazonaws.com/xgboost:1',
            'TrainingInputMode': 'File'
        },
        'RoleArn': 'arn:aws:iam::YOUR_ACCOUNT:role/SageMakerRole',
        'InputDataConfig': [{
            'ChannelName': 'train',
            'DataSource': {
                'S3DataSource': {
                    'S3DataType': 'S3Prefix',
                    'S3Uri': 's3://your-ml-bucket/training/train/',
                    'S3DataDistributionType': 'FullyReplicated'
                }
            },
            'ContentType': 'application/x-parquet'
        }],
        'OutputDataConfig': {
            'S3OutputPath': 's3://your-ml-bucket/models/'
        },
        'ResourceConfig': {
            'InstanceType': 'ml.m5.xlarge',
            'InstanceCount': 1,
            'VolumeSizeInGB': 30
        },
        'StoppingCondition': {
            'MaxRuntimeInSeconds': 3600
        },
        'TrainingJobName': 'churn-training-{{ ds_nodash }}'
    },
    wait_for_completion=True,
    aws_conn_id='aws_default',
    dag=dag
)

# Run batch predictions
batch_predict = SageMakerTransformOperator(
    task_id='batch_predictions',
    config={
        'TransformJobName': 'churn-predictions-{{ ds_nodash }}',
        'ModelName': 'churn-model-production',
        'TransformInput': {
            'DataSource': {
                'S3DataSource': {
                    'S3DataType': 'S3Prefix',
                    'S3Uri': 's3://your-ml-bucket/inference-input/'
                }
            },
            'ContentType': 'application/x-parquet'
        },
        'TransformOutput': {
            'S3OutputPath': 's3://your-ml-bucket/inference-output/'
        },
        'TransformResources': {
            'InstanceType': 'ml.c5.xlarge',
            'InstanceCount': 1
        }
    },
    wait_for_completion=True,
    aws_conn_id='aws_default',
    dag=dag
)

# Post-process predictions with Glue
postprocess_results = GlueJobOperator(
    task_id='enrich_predictions',
    job_name='ML_Results_Postprocessing',
    script_args={
        '--INPUT_PATH': 's3://your-ml-bucket/inference-output/',
        '--OUTPUT_PATH': 's3://your-ml-bucket/final-results/'
    },
    aws_conn_id='aws_default',
    dag=dag
)

# Define dependencies
wait_for_data >> prepare_data >> train_model >> batch_predict >> postprocess_results

This DAG demonstrates several production best practices. The S3KeySensor ensures data availability before processing begins—no point running expensive jobs if the input data isn’t ready. The Glue and SageMaker operators have wait_for_completion=True, which means Airflow monitors job status and only proceeds when each step succeeds. Error handling is built in through the retry mechanism, and the {{ ds }} templating allows for date-based processing.

⚡ Performance Tip

Separate training frequency from inference frequency

Don’t retrain models daily if weekly is sufficient. Use Airflow’s BranchPythonOperator to conditionally run training (e.g., only on Sundays) while running inference daily. This can reduce costs by 85% while maintaining prediction freshness.

Monitoring and Alerting:

Your Airflow DAG should include monitoring tasks that check data quality and model performance:

from airflow.operators.python import PythonOperator
import boto3

def check_prediction_quality(**context):
    s3 = boto3.client('s3')
    # Read prediction results
    # Calculate metrics like prediction distribution
    # Alert if anomalies detected
    pass

monitor_predictions = PythonOperator(
    task_id='monitor_prediction_quality',
    python_callable=check_prediction_quality,
    provide_context=True,
    dag=dag
)

postprocess_results >> monitor_predictions

Cost Optimization Strategies

Running ML pipelines can get expensive quickly if you’re not careful. Here are proven strategies for optimizing costs when connecting Glue and SageMaker.

Glue Cost Optimization:

Glue charges per DPU-hour (Data Processing Unit), currently $0.44 per DPU-hour. A G.2X worker equals 2 DPUs, so a job with 2 workers running for 40 minutes costs approximately $1.19. To optimize:

  • Use G.1X workers when possible: If your transformations don’t need 16GB of memory per worker, G.1X workers (8GB, 1 DPU) cost half as much
  • Enable Glue job metrics: Track DPU utilization to identify over-provisioning. If your job uses only 40% of available DPUs, you’re wasting 60% of costs
  • Use FLEX execution class: For non-time-sensitive jobs, FLEX costs $0.29/DPU-hour (34% savings) but may take longer to start
  • Partition your data: Process only what changed instead of full table scans. Use partition pruning: WHERE processing_date = '{{ ds }}'
  • Compress output: Writing Snappy-compressed Parquet reduces S3 storage costs and SageMaker read times

SageMaker Cost Optimization:

SageMaker training and inference instances can be expensive, but strategic choices make a huge difference:

  • Use Graviton-based instances: ml.c7g instances offer 30-50% cost savings over comparable x86 instances for PyTorch and XGBoost workloads
  • Spot instances for training: Use managed spot training to save up to 90% on training costs. SageMaker handles spot interruptions automatically
  • Right-size instances: Don’t default to large instances. A ml.m5.xlarge costs $0.23/hour while ml.m5.24xlarge costs $5.53/hour—use the smallest that meets your time requirements
  • Batch Transform over real-time endpoints: For non-real-time inference, Batch Transform costs significantly less than keeping an endpoint running 24/7
  • Use SageMaker Savings Plans: Commit to consistent usage for up to 64% savings on compute

Data Transfer Optimization:

Don’t overlook data transfer costs between services:

  • Keep Glue, SageMaker, and S3 in the same AWS region to avoid cross-region transfer fees
  • Use S3 Intelligent-Tiering for data you access infrequently (like old training datasets)
  • Compress data between stages—a 10GB Parquet file compressed to 3GB saves on both storage and transfer

Real-World Implementation Considerations

Theory meets reality in production environments. Here are lessons learned from implementing Glue-SageMaker pipelines at scale.

Data Format Selection:

While CSV is tempting due to familiarity, Parquet is dramatically better for ML pipelines. In testing, reading a 10GB CSV into SageMaker took 8 minutes and cost $0.32, while the same data in Parquet took 1.5 minutes and cost $0.06. Parquet’s columnar format means SageMaker only reads the features it needs, not entire rows.

Handling Schema Evolution:

ML features evolve over time. When you add new features, your pipeline needs to handle data with different schemas gracefully. Use Glue’s schema evolution capabilities:

# Enable schema merging when reading
prepared_df = spark.read \
    .option("mergeSchema", "true") \
    .parquet("s3://your-ml-bucket/training/train/")

# Fill missing columns with defaults
for expected_col in required_features:
    if expected_col not in prepared_df.columns:
        prepared_df = prepared_df.withColumn(expected_col, lit(0))

Managing Model Versions:

As you retrain models, version them properly. SageMaker Model Registry helps, but integrate it into your Airflow pipeline:

from airflow.providers.amazon.aws.operators.sagemaker import SageMakerRegisterModelVersionOperator

register_model = SageMakerRegisterModelVersionOperator(
    task_id='register_model_version',
    model_package_group_name='churn-prediction-models',
    model_metrics={
        'accuracy': 0.87,
        'auc': 0.92
    },
    approval_status='PendingManualApproval',
    aws_conn_id='aws_default',
    dag=dag
)

This creates an audit trail of which model version generated which predictions, critical for debugging and compliance.

Error Recovery:

Glue jobs can fail mid-execution. Enable Glue job bookmarks to checkpoint progress and resume from the last successful state rather than reprocessing all data. For SageMaker, use checkpointing in training jobs to save costs on spot interruptions:

xgb = sagemaker.estimator.Estimator(
    # ... other params
    checkpoint_s3_uri='s3://your-ml-bucket/checkpoints/',
    use_spot_instances=True,
    max_run=7200,
    max_wait=10800  # Allow time for spot interruption recovery
)

Conclusion

Connecting AWS Glue and SageMaker creates robust ML pipelines that separate data engineering concerns from ML workloads. Glue handles the heavy lifting of data transformation at scale, while SageMaker provides optimized infrastructure for training and inference. The key is treating S3 as your integration layer and using Airflow to orchestrate the workflow with proper error handling and monitoring.

The architecture patterns and code examples in this article provide a foundation you can adapt to your specific use cases. Start with a simple pipeline—Glue for preparation, SageMaker for training, Airflow for orchestration—and expand as your ML operations mature. Focus on data quality checks in Glue, proper instance sizing in SageMaker, and comprehensive monitoring in Airflow. These practices will help you build ML pipelines that are not only functional but also cost-effective and maintainable at scale.

Leave a Comment