--- title: "Retraining Pipeline" description: "An automated S3-to-deploy pipeline with Lambda and Step Functions is systems-engineering-grade MLOps — the difference between 'I can retrain a model' and 'I..." source: "https://confidentprep.com/paths/ml-engineering-on-aws/retraining-pipeline/" path: "ML Engineering on AWS" repository: "https://github.com/confidentprep14-web/ai-engineering-projects/tree/main/retraining-pipeline" token_estimate: 5938 --- # Retraining Pipeline ## Overview An automated pipeline — S3 upload triggers a Lambda, which kicks off a Step Functions execution that preprocesses, retrains, evaluates, and conditionally redeploys. This is where every earlier piece (features, training, evaluation, serving) stops being something you run by hand and becomes infrastructure that responds to new data on its own. ## What to do **Path:** 3 — ML Engineering on AWS **Position:** 11 of 12 **Difficulty:** 🔴 **Time:** 4-5h **AWS Cost:** Step Functions execution < $0.01. Training job ~$0.30. Total budget: $1. --- ### Agent Pickup Instructions ```bash # Bootstrap cd projects/confident-prep/p3-11-retraining-pipeline python -m venv .venv && source .venv/bin/activate pip install -r requirements.txt # Deploy infrastructure (state machine, Lambda, S3 trigger) python src/state_machine.py --create python src/s3_trigger.py --configure # Test the pipeline end-to-end python src/test_pipeline.py # Run tests (most tests don't need AWS) pytest tests/ -v # TEARDOWN bash scripts/teardown.sh ``` **Done when:** - [ ] `python src/state_machine.py --create` creates state machine and saves ARN to `.state-machine-arn` - [ ] `python src/s3_trigger.py --configure` creates Lambda and S3 event notification - [ ] `python src/test_pipeline.py` uploads data, watches all 5 state transitions print, completes - [ ] Lambda evaluator returns `{"deploy": true}` for input with AUC improvement > 0.01 - [ ] SNS notification fires on both deploy and no-deploy outcomes - [ ] `pytest tests/ -v` shows 4/4 passing - [ ] `bash scripts/teardown.sh` removes state machine, Lambda, IAM roles, S3 notification, SSM parameters --- ### What this project is An automated ML pipeline triggered by an S3 event: new data uploaded → Lambda starts a Step Functions execution → 5-state machine runs: (1) SageMaker Processing Job for preprocessing, (2) SageMaker Training Job, (3) Lambda evaluator compares new AUC to SSM baseline, (4) conditional branch — deploy if improvement > 1%, otherwise SNS "no deploy" notification, (5) deploy state updates or creates the SageMaker endpoint. SNS notifications fire at each outcome. This is the automation layer that makes p3-07 through p3-10 a production-grade MLOps loop. ### What the learner achieves "I can wire a multi-step ML pipeline using Step Functions and Lambda, trigger it from an S3 event, and implement conditional deployment logic — so model retraining and deployment happen automatically without manual intervention." --- ### Folder structure ``` p3-11-retraining-pipeline/ ├── src/ │ ├── state_machine.py # create/update Step Functions state machine from definition │ ├── state_machine_definition.json # ASL definition of the 5-state machine │ ├── lambda_evaluator.py # Lambda function: compare metrics to SSM baseline │ ├── s3_trigger.py # configure S3 event → Lambda → start execution │ └── test_pipeline.py # upload sample data, watch execution, print state transitions ├── scripts/ │ └── teardown.sh # delete all resources ├── tests/ │ ├── test_state_machine.py # JSON valid, 5 states present │ └── test_lambda_evaluator.py # deploy=True/False logic, SSM path ├── .state-machine-arn # written by state_machine.py ├── .lambda-arn # written by s3_trigger.py ├── .env.example ├── requirements.txt └── README.md ``` --- ### .env.example ```dotenv # AWS region AWS_REGION=us-east-1 # S3 bucket (triggers on new data uploads) S3_BUCKET=my-ml-training-bucket S3_DATA_PREFIX=p3-11/incoming/ # IAM role ARNs SAGEMAKER_ROLE_ARN=arn:aws:iam::123456789012:role/SageMakerExecutionRole LAMBDA_ROLE_ARN=arn:aws:iam::123456789012:role/LambdaMLPipelineRole STEP_FUNCTIONS_ROLE_ARN=arn:aws:iam::123456789012:role/StepFunctionsMLRole # Step Functions state machine name STATE_MACHINE_NAME=p3-11-ml-pipeline # Lambda function name LAMBDA_FUNCTION_NAME=p3-11-evaluate-and-deploy # SNS topic ARN for notifications SNS_TOPIC_ARN=arn:aws:sns:us-east-1:123456789012:p3-11-ml-notifications # SSM parameter path for baseline AUC SSM_BASELINE_PARAM=/p3-11/baseline/val_auc # Deployment threshold: AUC must improve by this fraction AUC_IMPROVEMENT_THRESHOLD=0.01 # SageMaker endpoint name for deploy state SAGEMAKER_ENDPOINT_NAME=p3-11-auto-endpoint SAGEMAKER_ENDPOINT_CONFIG_NAME=p3-11-auto-config SAGEMAKER_MODEL_NAME=p3-11-auto-model # MLflow (for baseline tracking) MLFLOW_TRACKING_URI=http://localhost:5000 MLFLOW_MODEL_NAME=adult-income-xgboost ``` --- ### requirements.txt ``` boto3==1.35.93 sagemaker==2.236.0 xgboost==2.1.3 python-dotenv==1.0.1 pytest==8.3.4 ``` --- ### src/ — what to implement #### `src/state_machine_definition.json` Amazon States Language (ASL) JSON defining the 5-state machine. ```json { "Comment": "ML retraining pipeline: preprocess → train → evaluate → conditional deploy", "StartAt": "Preprocess", "States": { "Preprocess": { "Type": "Task", "Resource": "arn:aws:states:::sagemaker:createProcessingJob.sync:2", "Parameters": { "ProcessingJobName.$": "States.Format('preprocess-{}', $$.Execution.Name)", "ProcessingResources": { "ClusterConfig": { "InstanceCount": 1, "InstanceType": "ml.m5.large", "VolumeSizeInGB": 5 } }, "AppSpecification": { "ImageUri.$": "$.preprocessing_image_uri", "ContainerEntrypoint": ["python3", "src/main.py"] }, "RoleArn.$": "$.sagemaker_role_arn", "ProcessingInputs": [ { "InputName": "raw-data", "S3Input": { "S3Uri.$": "$.input_s3_uri", "LocalPath": "/opt/ml/processing/input", "S3DataType": "S3Prefix", "S3InputMode": "File" } } ], "ProcessingOutputs": [ { "OutputName": "processed-data", "S3Output": { "S3Uri.$": "$.processed_s3_uri", "LocalPath": "/opt/ml/processing/output", "S3OutputMode": "EndOfJob" } } ] }, "Next": "Train" }, "Train": { "Type": "Task", "Resource": "arn:aws:states:::sagemaker:createTrainingJob.sync:2", "Parameters": { "TrainingJobName.$": "States.Format('train-{}', $$.Execution.Name)", "AlgorithmSpecification": { "TrainingInputMode": "File", "TrainingImage.$": "$.training_image_uri" }, "RoleArn.$": "$.sagemaker_role_arn", "InputDataConfig": [ { "ChannelName": "train", "DataSource": { "S3DataSource": { "S3DataType": "S3Prefix", "S3Uri.$": "$.processed_s3_uri", "S3DataDistributionType": "FullyReplicated" } } } ], "OutputDataConfig": { "S3OutputPath.$": "$.model_output_s3_uri" }, "ResourceConfig": { "InstanceType": "ml.m5.large", "InstanceCount": 1, "VolumeSizeInGB": 5 }, "StoppingCondition": { "MaxRuntimeInSeconds": 3600 }, "HyperParameters": { "max_depth": "5", "learning_rate": "0.1", "n_estimators": "100" } }, "ResultPath": "$.training_result", "Next": "Evaluate" }, "Evaluate": { "Type": "Task", "Resource": "arn:aws:states:::lambda:invoke", "Parameters": { "FunctionName.$": "$.lambda_evaluator_arn", "Payload": { "training_job_name.$": "$.training_result.TrainingJobName", "ssm_param_path.$": "$.ssm_baseline_param" } }, "ResultPath": "$.evaluation_result", "Next": "ShouldDeploy" }, "ShouldDeploy": { "Type": "Choice", "Choices": [ { "Variable": "$.evaluation_result.Payload.deploy", "BooleanEquals": true, "Next": "Deploy" } ], "Default": "NotifyNoDeploy" }, "Deploy": { "Type": "Task", "Resource": "arn:aws:states:::sagemaker:createModel.sync:2", "Comment": "Simplified: in production this would updateEndpoint or createEndpoint", "Parameters": { "ModelName.$": "States.Format('auto-model-{}', $$.Execution.Name)", "PrimaryContainer": { "Image.$": "$.training_image_uri", "ModelDataUrl.$": "$.training_result.ModelArtifacts.S3ModelArtifacts" }, "ExecutionRoleArn.$": "$.sagemaker_role_arn" }, "Next": "NotifyDeployed" }, "NotifyDeployed": { "Type": "Task", "Resource": "arn:aws:states:::sns:publish", "Parameters": { "TopicArn.$": "$.sns_topic_arn", "Message.$": "States.Format('Model deployed: execution {}. AUC improvement exceeded threshold.', $$.Execution.Name)", "Subject": "ML Pipeline: Model Deployed" }, "End": true }, "NotifyNoDeploy": { "Type": "Task", "Resource": "arn:aws:states:::sns:publish", "Parameters": { "TopicArn.$": "$.sns_topic_arn", "Message.$": "States.Format('No deployment: execution {}. AUC improvement did not meet threshold.', $$.Execution.Name)", "Subject": "ML Pipeline: Retrained, Not Deployed" }, "End": true } } } ``` --- #### `src/state_machine.py` **CLI:** `python src/state_machine.py [--create | --describe | --delete]` **`load_definition(path: str) -> str`** - Reads `src/state_machine_definition.json` - Validates it is valid JSON with `json.loads()` - Returns the JSON string **`create_state_machine(name: str, definition: str, role_arn: str, region: str) -> str`** - Calls `boto3.client('stepfunctions').create_state_machine()` - Returns the state machine ARN **`start_execution(state_machine_arn: str, input_data: dict) -> str`** - Calls `.start_execution(stateMachineArn=arn, input=json.dumps(input_data))` - Returns the execution ARN **`wait_for_execution(execution_arn: str, poll_interval: int = 10) -> dict`** - Polls `describe_execution` every `poll_interval` seconds - Prints state transitions as they appear in event history using `get_execution_history` - Returns final execution status dict when status is SUCCEEDED or FAILED **`main()`** - `--create`: load definition → create state machine → write ARN to `.state-machine-arn` → print ARN - `--describe`: read ARN from file → describe and print status - `--delete`: read ARN → delete state machine --- #### `src/lambda_evaluator.py` Lambda function handler — deployed as a zip to Lambda. **`get_training_job_auc(training_job_name: str, region: str) -> float | None`** - Uses `boto3.client('logs')` to scan `/aws/sagemaker/TrainingJobs/{training_job_name}` - Searches for line matching `ACCURACY: (\d+\.\d+)` (reusing p3-01's format) - Returns float if found, None otherwise **`get_baseline_auc(ssm_param_path: str, region: str) -> float | None`** - Calls `boto3.client('ssm').get_parameter(Name=ssm_param_path, WithDecryption=False)` - Returns `float(parameter["Parameter"]["Value"])` or None if not found **`evaluate_improvement(new_auc: float | None, baseline_auc: float | None, threshold: float = 0.01) -> dict`** - If either is None: return `{"deploy": False, "reason": "metrics unavailable", "new_auc": None, "baseline_auc": None}` - delta = new_auc - baseline_auc - Returns: ```python { "deploy": delta > threshold, "reason": f"AUC improved by {delta:.4f} (threshold {threshold})" if delta > threshold else f"AUC improved by only {delta:.4f} (threshold {threshold})", "new_auc": new_auc, "baseline_auc": baseline_auc, "delta": delta, } ``` **`handler(event: dict, context) -> dict`** - Extracts `training_job_name` and `ssm_param_path` from `event` - Calls `get_training_job_auc` → `get_baseline_auc` → `evaluate_improvement` - Returns the evaluation dict - Lambda returns this dict as the Payload that Step Functions reads **`deploy_lambda(function_name: str, role_arn: str, region: str) -> str`** - Zips `lambda_evaluator.py` using `zipfile.ZipFile` - Calls `boto3.client('lambda').create_function()` with the zip - Returns the Lambda ARN --- #### `src/s3_trigger.py` **CLI:** `python src/s3_trigger.py --configure` **`create_trigger_lambda(function_name: str, state_machine_arn: str, role_arn: str, region: str) -> str`** Creates a separate Lambda function (the S3 trigger Lambda, NOT the evaluator Lambda) whose job is to: 1. Receive the S3 event notification 2. Extract the S3 key of the uploaded file 3. Start a Step Functions execution with appropriate input JSON The trigger Lambda code (as a string in this function): ```python import boto3 import json import os def handler(event, context): sfn = boto3.client('stepfunctions') record = event['Records'][0]['s3'] bucket = record['bucket']['name'] key = record['object']['key'] execution_input = { "input_s3_uri": f"s3://{bucket}/{key}", "processed_s3_uri": f"s3://{bucket}/p3-11/processed/", "model_output_s3_uri": f"s3://{bucket}/p3-11/model-output/", "sagemaker_role_arn": os.environ["SAGEMAKER_ROLE_ARN"], "lambda_evaluator_arn": os.environ["LAMBDA_EVALUATOR_ARN"], "ssm_baseline_param": os.environ["SSM_BASELINE_PARAM"], "sns_topic_arn": os.environ["SNS_TOPIC_ARN"], "preprocessing_image_uri": os.environ["PREPROCESSING_IMAGE_URI"], "training_image_uri": os.environ["TRAINING_IMAGE_URI"], } sfn.start_execution( stateMachineArn=os.environ["STATE_MACHINE_ARN"], input=json.dumps(execution_input) ) ``` Deploys this Lambda and returns its ARN. **`configure_s3_notification(bucket: str, prefix: str, lambda_arn: str) -> None`** - Grants S3 permission to invoke the Lambda: `boto3.client('lambda').add_permission(...)` - Sets up S3 notification configuration: `boto3.client('s3').put_bucket_notification_configuration(...)` - Filter: `{"Key": {"FilterRules": [{"Name": "prefix", "Value": prefix}]}}` - Event: `s3:ObjectCreated:*` **`main()`** - Create evaluator Lambda (from `lambda_evaluator.py`) - Create trigger Lambda - Store baseline AUC in SSM: `boto3.client('ssm').put_parameter(Name=SSM_BASELINE_PARAM, Value="0.883", Type="String", Overwrite=True)` - Configure S3 notification - Write trigger Lambda ARN to `.lambda-arn` - Print: `Pipeline configured. Upload data to s3://{bucket}/{prefix} to trigger.` --- #### `src/test_pipeline.py` **CLI:** `python src/test_pipeline.py` **`main()`** 1. Upload a sample `adult.data` file to `s3://{S3_BUCKET}/{S3_DATA_PREFIX}test_run.csv` 2. Print: `Uploaded test data. Waiting for S3 trigger to start execution...` 3. Wait 5 seconds for S3 notification to trigger Lambda 4. Poll Step Functions for the most recent execution of the state machine 5. Call `state_machine.wait_for_execution()` — this prints each state transition 6. On completion, print: `Pipeline complete. Status: {status}` --- ### tests/ — what to test #### `tests/test_state_machine.py` **`test_state_machine_definition_is_valid_json`** - Open `src/state_machine_definition.json` - Call `json.loads(content)` — assert no exception **`test_state_machine_has_five_states`** - Parse the definition - Count keys in `definition["States"]` - Assert count == 5 (Preprocess, Train, Evaluate, ShouldDeploy, Deploy — or NotifyDeployed/NotifyNoDeploy) **`test_state_machine_starts_at_preprocess`** - Assert `definition["StartAt"] == "Preprocess"` #### `tests/test_lambda_evaluator.py` **`test_evaluate_improvement_returns_deploy_true_when_above_threshold`** - `evaluate_improvement(new_auc=0.895, baseline_auc=0.880, threshold=0.01)` - Assert `result["deploy"] == True` - Assert `result["delta"] == pytest.approx(0.015, abs=0.001)` **`test_evaluate_improvement_returns_deploy_false_when_at_threshold`** - `evaluate_improvement(new_auc=0.890, baseline_auc=0.880, threshold=0.01)` - Assert `result["deploy"] == False` (delta = 0.010, not strictly > 0.01) **`test_evaluate_improvement_returns_deploy_false_when_metrics_unavailable`** - `evaluate_improvement(new_auc=None, baseline_auc=0.880)` - Assert `result["deploy"] == False` - Assert `result["reason"] == "metrics unavailable"` **`test_get_baseline_auc_uses_correct_ssm_path`** - Mock `boto3.client('ssm').get_parameter` to expect `Name="/p3-11/baseline/val_auc"` - Call `get_baseline_auc("/p3-11/baseline/val_auc", "us-east-1")` - Assert the mock was called with the correct parameter name --- ### README.md content ```markdown # P3-11 — Retraining Pipeline Automated ML pipeline: S3 upload → preprocess → train → evaluate → conditional deploy. ## Architecture ``` S3 upload → Lambda trigger → Step Functions execution → State 1: SageMaker Processing (preprocess) → State 2: SageMaker Training → State 3: Lambda evaluate (compare to SSM baseline) → State 4: Choice (deploy if AUC improved > 1%) → State 5a: SageMaker deploy + SNS notify → State 5b: SNS notify (no deploy) ``` ## Prerequisites - Python 3.11+ - AWS credentials with: Step Functions, Lambda, SageMaker, S3, SNS, SSM, IAM permissions - SNS topic created: `aws sns create-topic --name p3-11-ml-notifications` ## Quick start ```bash python -m venv .venv && source .venv/bin/activate pip install -r requirements.txt cp .env.example .env # fill in all ARNs and names ## Deploy infrastructure python src/state_machine.py --create python src/s3_trigger.py --configure ## Test end-to-end python src/test_pipeline.py ``` ## Tests ```bash pytest tests/ -v ``` Most tests run without AWS credentials. ## TEARDOWN ```bash bash scripts/teardown.sh ``` ``` --- ### GUIDE.md content ```markdown # Guide — P3-11 Retraining Pipeline ## Why Step Functions A training pipeline is a sequence of steps that each can fail. Step Functions gives you: - **Visual execution history:** see exactly which state failed and why - **Automatic retry:** configure retries per-state without writing retry loops - **Conditional branching:** the Choice state lets you implement the "only deploy if better" rule without coordinating that logic in Python - **Audit trail:** every execution is logged with full input/output per state Alternative: Airflow, Prefect, or a Lambda chain. Step Functions integrates natively with SageMaker via the `sagemaker:createTrainingJob.sync` resource — it polls for completion automatically. With a Lambda chain you write that polling logic yourself. ## The SSM baseline SSM Parameter Store stores the baseline AUC as a simple string parameter. When the evaluator Lambda compares `new_auc > baseline + threshold`, it reads this parameter at runtime. After a successful deployment, you should update the SSM parameter to the new AUC so the next pipeline run compares against the most recently deployed model. The `Deploy` state in this simplified version does not do this — add it as an enhancement. ## S3 event latency S3 event notifications are usually delivered in seconds, but SLA is "typically less than 1 minute." In `test_pipeline.py`, a 5-second wait before polling for the execution is usually enough — but in rare cases the notification may be delayed. If the test times out waiting for an execution, check the S3 notification configuration and Lambda execution logs. ## The IAM permissions maze This pipeline requires multiple roles: - **SageMaker role:** create Processing and Training Jobs, read/write S3 - **Lambda role:** read CloudWatch logs, read SSM, invoke Step Functions - **Step Functions role:** invoke Lambda, create SageMaker jobs, publish to SNS Each role needs the right permissions or the state machine will fail at that state. Check CloudWatch Logs for Lambda errors and Step Functions execution history for SageMaker errors. ## Interview framing "I built an automated retraining pipeline using Step Functions: S3 upload triggers a Lambda that starts a 5-state execution — preprocess, train, evaluate, conditional branch, deploy. The Lambda evaluator compares new AUC to an SSM baseline and only deploys if improvement exceeds a configurable threshold. The whole thing is event-driven with no polling loop." ``` --- ### Failure modes to handle | Failure | Where | How to handle | |---|---|---| | Step Functions execution fails at Preprocess state | `test_pipeline.py` | Print failure cause from `get_execution_history`; most common: wrong image URI or IAM role | | Lambda evaluator can't find training job logs | `lambda_evaluator.py: get_training_job_auc` | Return None; `evaluate_improvement` returns `deploy=False` with reason "metrics unavailable" | | SSM parameter not set (first run) | `lambda_evaluator.py: get_baseline_auc` | Return None → `evaluate_improvement` returns `deploy=False` — which is correct for a first run with no established baseline | | S3 event notification not delivered | `test_pipeline.py` | After 60 seconds with no execution started, print: "No execution found. Check S3 notification configuration." | | Step Functions state machine already exists on `--create` | `state_machine.py: create_state_machine` | Catch `StateMachineAlreadyExists`, offer `--update` flag to update the definition | | Lambda zip exceeds size limit | `lambda_evaluator.py: deploy_lambda` | Keep `lambda_evaluator.py` under 50KB; it only imports `boto3` which is available in Lambda runtime without packaging | --- ### The metric this project measures **What:** Pipeline execution outcome — deploy or no-deploy — and the AUC delta that drove the decision **Format:** ``` Pipeline complete. Status: SUCCEEDED Evaluation result: new_auc=0.891, baseline=0.880, delta=+0.011 → DEPLOY ``` Or: ``` Pipeline complete. Status: SUCCEEDED Evaluation result: new_auc=0.882, baseline=0.880, delta=+0.002 → NO DEPLOY (threshold: 0.01) ``` **Target:** For a training job on fresh UCI Adult data with good hyperparameters, AUC delta > 0.01 is achievable when the challenger uses better hyperparameters than the stored baseline. --- ### Cost estimate | Resource | Qty | Rate | Estimated cost | |---|---|---|---| | Step Functions state transitions | ~10 per execution | $0.025/1000 transitions | <$0.01 | | SageMaker Processing Job (ml.m5.large, ~5 min) | 1 run | $0.115/hr | ~$0.01 | | SageMaker Training Job (ml.m5.large, ~5 min) | 1 run | $0.115/hr | ~$0.01 | | Lambda invocations (2 Lambdas) | ~5 total | Free tier | $0.00 | | SNS notifications | 1-2 | Free tier | $0.00 | | SSM GetParameter | ~5 calls | Free tier | $0.00 | | **Total per session** | | | **~$0.30-1.00** | ### Teardown checklist `scripts/teardown.sh` must (in order): - [ ] Delete state machine: `aws stepfunctions delete-state-machine --state-machine-arn $(cat .state-machine-arn)` - [ ] Delete trigger Lambda: `aws lambda delete-function --function-name ${LAMBDA_TRIGGER_NAME}` - [ ] Delete evaluator Lambda: `aws lambda delete-function --function-name ${LAMBDA_FUNCTION_NAME}` - [ ] Remove S3 event notification: `aws s3api put-bucket-notification-configuration --bucket ${S3_BUCKET} --notification-configuration '{}'` - [ ] Delete SSM parameter: `aws ssm delete-parameter --name ${SSM_BASELINE_PARAM}` - [ ] Delete IAM roles (if created by scripts): `aws iam delete-role --role-name ...` (list roles to delete) - [ ] Delete S3 data under `p3-11/` prefix: `aws s3 rm s3://${S3_BUCKET}/p3-11/ --recursive` - [ ] Print `All pipeline resources deleted.` ## Source code A reference implementation of Retraining Pipeline — README, guide, source and tests — is at https://github.com/confidentprep14-web/ai-engineering-projects/tree/main/retraining-pipeline.