--- title: "A/B Testing" description: "Running a real two-variant traffic split on SageMaker is a direct, hands-on answer to 'how would you safely roll out a new model' — a near-universal ML systems..." source: "https://confidentprep.com/paths/ml-engineering-on-aws/ab-testing/" path: "ML Engineering on AWS" repository: "https://github.com/confidentprep14-web/ai-engineering-projects/tree/main/ab-testing" token_estimate: 4696 --- # A/B Testing ## Overview Deploy two model variants behind one SageMaker endpoint, shift traffic between them, and pick the winner on real results. This runs two instances at once (about $0.11/hour combined) — tear it down right after testing. Same endpoint pattern as the previous project, now split to answer a question no offline metric can: which model actually wins with real traffic. ## What to do **Path:** 3 — ML Engineering on AWS **Position:** 8 of 12 **Difficulty:** 🔴 **Time:** 3-4h **AWS Cost:** 2 × ml.t2.medium × ~2 hours = ~$0.22 --- ### Agent Pickup Instructions ```bash # Bootstrap cd projects/confident-prep/p3-08-ab-testing python -m venv .venv && source .venv/bin/activate pip install -r requirements.txt # Ensure p3-05 MLflow registry has "production" alias # Prepare both model variants python src/prepare_variants.py # Deploy A/B endpoint (90/10 split) python src/deploy_ab.py # Send traffic (1000 requests, no target variant header) python src/invoke_traffic.py --endpoint-name $(cat .endpoint-name) --n-requests 1000 # Evaluate and decide winner python src/evaluate_winner.py --endpoint-name $(cat .endpoint-name) # Optionally shift traffic python src/shift_traffic.py --current-weight 5 --challenger-weight 5 # Promote winner to MLflow registry python src/promote_winner.py # TEARDOWN IMMEDIATELY bash scripts/teardown.sh ``` **Done when:** - [ ] `python src/prepare_variants.py` creates two `model.tar.gz` files and uploads both to S3 - [ ] `python src/deploy_ab.py` creates endpoint with two ProductionVariants (weights 9+1=10) - [ ] `python src/invoke_traffic.py` sends 1000 requests and tracks per-variant routing from response metadata - [ ] `python src/evaluate_winner.py` prints winner with AUC values for both variants - [ ] `python src/promote_winner.py` updates MLflow "production" alias to winner - [ ] `pytest tests/ -v` shows 4/4 passing - [ ] `bash scripts/teardown.sh` deletes endpoint, both endpoint configs, both SageMaker models --- ### What this project is You deploy two XGBoost model variants behind a single SageMaker endpoint with 90/10 traffic weighting — the current "production" model from p3-05 gets 90% of traffic, a new "challenger" (trained with different hyperparameters) gets 10%. You send 1000 requests with no target variant header and observe SageMaker's traffic routing. You shift the split to 50/50 then 0/100. You evaluate per-variant CloudWatch metrics, apply a decision rule (challenger beats baseline if AUC improves by > 0.01), and promote the winner to the MLflow registry. ### What the learner achieves "I can implement A/B testing for ML models using SageMaker's multi-variant endpoint, apply a quantitative decision rule, and promote the winner — which is how you safely roll out model updates without a big-bang deployment." --- ### Folder structure ``` p3-08-ab-testing/ ├── src/ │ ├── prepare_variants.py # load current + train challenger, package both, upload to S3 │ ├── deploy_ab.py # create endpoint with 2 ProductionVariants at 90/10 weights │ ├── invoke_traffic.py # send N requests, no variant header, record routing from response │ ├── shift_traffic.py # CLI to update variant weights │ ├── evaluate_winner.py # fetch CloudWatch per-variant metrics, apply decision rule │ └── promote_winner.py # update MLflow registry alias to winning variant ├── scripts/ │ └── teardown.sh # delete endpoint + both configs + both models ├── tests/ │ ├── test_weights.py # weights sum to 10, decision rule logic │ └── test_variants.py # prepare_variants creates distinct artifacts ├── .endpoint-name # written by deploy_ab.py ├── .env.example ├── requirements.txt └── README.md ``` --- ### .env.example ```dotenv # AWS region AWS_REGION=us-east-1 # S3 bucket S3_BUCKET=my-ml-training-bucket # IAM role ARN SAGEMAKER_ROLE_ARN=arn:aws:iam::123456789012:role/SageMakerExecutionRole # MLflow MLFLOW_TRACKING_URI=http://localhost:5000 MLFLOW_MODEL_NAME=adult-income-xgboost # SageMaker resource names SAGEMAKER_CURRENT_MODEL_NAME=p3-08-current-model SAGEMAKER_CHALLENGER_MODEL_NAME=p3-08-challenger-model SAGEMAKER_ENDPOINT_CONFIG_NAME=p3-08-ab-config SAGEMAKER_ENDPOINT_NAME=p3-08-ab-endpoint # Variant names (used in ProductionVariant config and CloudWatch dimensions) VARIANT_CURRENT=current VARIANT_CHALLENGER=challenger # Initial traffic split (must sum to 10) INITIAL_CURRENT_WEIGHT=9 INITIAL_CHALLENGER_WEIGHT=1 # Decision rule: challenger must improve AUC by at least this much AUC_IMPROVEMENT_THRESHOLD=0.01 # Challenger hyperparameters (different from production) CHALLENGER_MAX_DEPTH=7 CHALLENGER_LEARNING_RATE=0.05 CHALLENGER_N_ESTIMATORS=200 ``` --- ### requirements.txt ``` boto3==1.35.93 sagemaker==2.236.0 xgboost==2.1.3 mlflow==2.19.0 pandas==2.2.3 numpy==2.2.1 scikit-learn==1.6.0 python-dotenv==1.0.1 pytest==8.3.4 ``` --- ### src/ — what to implement #### `src/prepare_variants.py` **CLI:** `python src/prepare_variants.py` **`load_current_model(model_name: str, alias: str) -> xgb.XGBClassifier`** - Loads from MLflow registry: `mlflow.xgboost.load_model(f"models:/{model_name}@{alias}")` **`train_challenger(data_path: str, hyperparams: dict) -> tuple[xgb.XGBClassifier, float]`** - Trains a new XGBClassifier with `hyperparams` (read from env: `CHALLENGER_MAX_DEPTH`, `CHALLENGER_LEARNING_RATE`, `CHALLENGER_N_ESTIMATORS`) - Computes val_AUC on 20% holdout (same split as p3-01: random_state=42) - Returns `(challenger_model, challenger_auc)` **`package_and_upload(model: xgb.XGBClassifier, variant_name: str, bucket: str) -> str`** - Saves model to temp dir as `model.xgb` - Copies a SageMaker-compatible `inference.py` (same structure as p3-07) to the same temp dir - Creates `model.tar.gz` with both files at root - Uploads to `s3://{bucket}/p3-08/{variant_name}/model.tar.gz` - Returns S3 URI **`main()`** 1. Load current model from registry 2. Train challenger with env-specified hyperparams 3. Package and upload current → `s3://{bucket}/p3-08/current/model.tar.gz` 4. Package and upload challenger → `s3://{bucket}/p3-08/challenger/model.tar.gz` 5. Print: `Prepared current and challenger variants. Challenger val_AUC: {auc:.4f}` **Edge cases:** - Verify the two S3 URIs are different strings before returning — they must point to distinct objects --- #### `src/deploy_ab.py` **CLI:** `python src/deploy_ab.py` **`create_variant_model(model_s3_uri: str, model_name: str, role_arn: str, region: str) -> None`** - Creates a SageMaker Model resource using `boto3.client('sagemaker').create_model()` - Same container image as p3-07 (XGBoost 1.7-1) **`create_ab_endpoint_config(current_model: str, challenger_model: str, config_name: str, instance_type: str, current_weight: int, challenger_weight: int) -> None`** - Calls `create_endpoint_config` with TWO ProductionVariants: ```python [ { "VariantName": "current", "ModelName": current_model, "InstanceType": instance_type, "InitialInstanceCount": 1, "InitialVariantWeight": current_weight, # e.g., 9 }, { "VariantName": "challenger", "ModelName": challenger_model, "InstanceType": instance_type, "InitialInstanceCount": 1, "InitialVariantWeight": challenger_weight, # e.g., 1 } ] ``` - Validates: `current_weight + challenger_weight == 10`; raises `ValueError` if not **`create_and_wait_endpoint(config_name: str, endpoint_name: str) -> None`** - Creates endpoint, waits for `InService` using waiter (15-min timeout) **`main()`** 1. Create current model resource 2. Create challenger model resource 3. `create_ab_endpoint_config` with weights from env 4. `create_and_wait_endpoint` 5. Write endpoint name to `.endpoint-name` 6. Print: `A/B endpoint {endpoint_name} is InService. Traffic: {current_weight*10}% current / {challenger_weight*10}% challenger` --- #### `src/invoke_traffic.py` **CLI:** `python src/invoke_traffic.py --endpoint-name [--n-requests 1000] [--data data/adult.data]` **`send_requests(endpoint_name: str, payload: str, n_requests: int) -> dict`** - Uses `boto3.client('sagemaker-runtime')` - Sends `n_requests` requests with NO `TargetVariant` header — SageMaker routes internally - From each response, reads `response["InvokedProductionVariant"]` to track which variant served - Returns `{"current": int, "challenger": int, "total": int, "ratio": float}` where ratio = challenger/total **`print_traffic_report(routing: dict) -> None`** - Prints: ``` Traffic routing over {total} requests: current: {current} ({current/total*100:.1f}%) challenger: {challenger} ({challenger/total*100:.1f}%) ``` **`main()`** - Create sample payload (1 row from adult.data, preprocessed) - `send_requests` → `print_traffic_report` - Save routing dict to `output/traffic_routing.json` --- #### `src/shift_traffic.py` **CLI:** `python src/shift_traffic.py --current-weight N --challenger-weight M [--endpoint-name ]` **`validate_weights(current: int, challenger: int) -> None`** - Raises `ValueError(f"Weights must sum to 10, got {current + challenger}")` if not **`update_traffic_weights(endpoint_name: str, current_weight: int, challenger_weight: int) -> None`** - Validates weights sum to 10 - Calls `boto3.client('sagemaker').update_endpoint_weights_and_capacities()` with new variant weights - No need to recreate endpoint — this is a live traffic shift - Prints: `Traffic updated: {current_weight*10}% current / {challenger_weight*10}% challenger` **`main()`** - Parse args → `validate_weights` → `update_traffic_weights` --- #### `src/evaluate_winner.py` **CLI:** `python src/evaluate_winner.py --endpoint-name [--baseline-auc 0.88]` **`fetch_variant_invocations(endpoint_name: str, variant_name: str, hours: int = 2) -> dict`** - Uses `boto3.client('cloudwatch').get_metric_statistics()` - Namespace: `"AWS/SageMaker"`, MetricName: `"Invocations"` - Dimensions: `[{"Name": "EndpointName", ...}, {"Name": "VariantName", ...}]` - Returns `{"variant": variant_name, "invocation_count": int}` **`apply_decision_rule(challenger_auc: float, baseline_auc: float, threshold: float = 0.01) -> dict`** - Returns `{"deploy_challenger": bool, "reason": str, "delta": float}` - Example reasons: - `"Challenger AUC 0.893 exceeds baseline 0.880 by 0.013 (threshold 0.010) → deploy challenger"` - `"Challenger AUC 0.875 does not exceed baseline 0.880 by threshold 0.010 → keep current"` **`main()`** - Read baseline AUC from `--baseline-auc` arg or from MLflow registry (logged metric on production run) - Read challenger AUC from `.challenger-auc` file written by `prepare_variants.py` - Fetch per-variant invocation counts from CloudWatch - `apply_decision_rule` - Print decision and reasons - Write result to `output/ab_decision.json` --- #### `src/promote_winner.py` **CLI:** `python src/promote_winner.py [--winner current|challenger]` **`promote_to_production(winner: str, model_name: str) -> None`** - If `winner == "challenger"`: - Log the challenger model as a new MLflow run and register it with `mlflow.register_model()` - Set alias `"production"` to the new version with `mlflow.MlflowClient().set_registered_model_alias()` - Print: `Challenger promoted to production. Registry updated.` - If `winner == "current"`: - Print: `Current model retained. No registry change.` **`main()`** - Read winner from `output/ab_decision.json` or `--winner` arg - `promote_to_production` --- ### tests/ — what to test #### `tests/test_weights.py` **`test_traffic_weights_must_sum_to_10`** - Call `deploy_ab.create_ab_endpoint_config(...)` with weights 6+6=12 (via mock boto3) - Assert `ValueError` is raised before any AWS call **`test_decision_rule_picks_challenger_when_improvement_exceeds_threshold`** - `apply_decision_rule(challenger_auc=0.895, baseline_auc=0.880, threshold=0.01)` - Assert `result["deploy_challenger"] == True` - Assert `result["delta"] == pytest.approx(0.015, abs=0.001)` **`test_decision_rule_keeps_current_when_improvement_below_threshold`** - `apply_decision_rule(challenger_auc=0.885, baseline_auc=0.880, threshold=0.01)` - Assert `result["deploy_challenger"] == False` **`test_shift_traffic_validates_weights`** - `shift_traffic.validate_weights(7, 2)` — sum is 9, not 10 - Assert `ValueError` is raised #### `tests/test_variants.py` **`test_prepare_variants_creates_distinct_s3_uris`** - Mock `package_and_upload` to record the `variant_name` arg - Run `prepare_variants.main()` with a tiny XGBClassifier and mock boto3 - Assert the two recorded variant names are `"current"` and `"challenger"` (distinct) --- ### README.md content ```markdown # P3-08 — A/B Testing Deploy two model variants behind one SageMaker endpoint. Shift traffic. Pick the winner. ## WARNING This project creates a SageMaker endpoint with TWO instances (~$0.11/hr total). Run `bash scripts/teardown.sh` immediately after you finish testing. ## Prerequisites - Python 3.11+ - AWS credentials configured - p3-05 completed: MLflow registry has "production" alias - MLflow server running at http://localhost:5000 ## Quick start ```bash python -m venv .venv && source .venv/bin/activate pip install -r requirements.txt cp .env.example .env # fill in S3_BUCKET, SAGEMAKER_ROLE_ARN python src/prepare_variants.py python src/deploy_ab.py python src/invoke_traffic.py --endpoint-name $(cat .endpoint-name) --n-requests 1000 python src/evaluate_winner.py --endpoint-name $(cat .endpoint-name) python src/promote_winner.py ``` ## Traffic shifting ```bash ## 50/50 split python src/shift_traffic.py --current-weight 5 --challenger-weight 5 ## Full challenger rollout python src/shift_traffic.py --current-weight 0 --challenger-weight 10 ``` ## TEARDOWN ```bash bash scripts/teardown.sh ``` ## Tests ```bash pytest tests/ -v ``` ``` --- ### GUIDE.md content ```markdown # Guide — P3-08 A/B Testing ## Why multi-variant endpoints A traditional deployment flips 100% of traffic from old model to new model at once. If the new model is worse, every user is affected before you notice. A multi-variant endpoint lets you route a small percentage of traffic to the new model first, measure its performance in production, and only roll it out fully once you have evidence it works. ## How SageMaker routes traffic without TargetVariant When you send requests with no `TargetVariant` header, SageMaker uses the `InitialVariantWeight` values to do probabilistic routing. With weights 9 and 1 (sum=10), approximately 90% of requests go to the current variant and 10% to the challenger. It is not perfectly deterministic per-request, but over 1000 requests you will see the ratio converge. ## The decision rule: AUC delta threshold A challenger that is only 0.001 AUC better may be within statistical noise. The 0.01 threshold is a business decision: is it worth the deployment risk and operational cost for a 0.01 AUC improvement? In practice this threshold should come from: what AUC delta translates to meaningful business outcomes (e.g., additional revenue, fewer false positives). ## CloudWatch per-variant metrics SageMaker automatically publishes invocation counts broken down by `VariantName` dimension to CloudWatch. You can also see per-variant `ModelLatency` and `OverheadLatency`. The `evaluate_winner.py` script reads these to understand actual traffic distribution, but for AUC you need to run inference on your hold-out set — CloudWatch does not know your model's accuracy. ## The shift_traffic flow `update_endpoint_weights_and_capacities()` is a live update — no endpoint recreation needed. The shift takes effect within seconds. This is the correct way to do a gradual rollout: 90/10 → 50/50 → 0/100, pausing to evaluate at each step. ## Interview framing "I implement model A/B testing using SageMaker's multi-variant endpoints: start with a 90/10 traffic split, measure the challenger's performance under real traffic, apply a quantitative AUC improvement threshold to make the promotion decision, and use live traffic shifts to roll out gradually. No big-bang deployments." ``` --- ### Failure modes to handle | Failure | Where | How to handle | |---|---|---| | Weights don't sum to 10 | `deploy_ab.py: create_ab_endpoint_config` | Validate before any AWS call; raise `ValueError` | | CloudWatch metrics not yet available (< 5 min of traffic) | `evaluate_winner.py` | If invocation count is 0, print warning: "Not enough traffic yet — run invoke_traffic.py first" | | Challenger AUC not recorded | `evaluate_winner.py` | `prepare_variants.py` writes `.challenger-auc` file; if missing, raise with message | | Endpoint creation fails (ResourceLimitExceeded) | `deploy_ab.py` | Catch `botocore.exceptions.ClientError`, print the limit type and a link to quota increase | | `update_endpoint_weights_and_capacities` with weight 0 for current | `shift_traffic.py` | Weight 0 is valid (0/10 = 0% traffic); document this is how you complete rollout | | Teardown order matters | `scripts/teardown.sh` | Delete endpoint FIRST, then endpoint config, then models — reverse creation order | --- ### The metric this project measures **What:** Per-variant traffic routing ratio (observed vs configured), and challenger AUC delta **Format:** ``` Traffic routing: current=901 (90.1%), challenger=99 (9.9%) Configured: 90% / 10% — routing within expected range. Challenger delta: +0.013 AUC vs baseline → DEPLOY ``` **Target:** Routing ratio within ±5% of configured split over 1000 requests. Challenger AUC delta clearly positive or negative. --- ### Cost estimate | Resource | Qty | Rate | Estimated cost | |---|---|---|---| | ml.t2.medium endpoint × 2 variants × 2 hours | 2 instances | $0.056/hr | ~$0.22 | | S3 model artifacts (2 × model.tar.gz) | ~20 MB | $0.023/GB-month | <$0.01 | | CloudWatch GetMetricStatistics API | ~10 calls | Free tier | $0.00 | | **Total per session** | | | **~$0.20** | ### Teardown checklist `scripts/teardown.sh` must (in order): - [ ] `aws sagemaker delete-endpoint --endpoint-name ${SAGEMAKER_ENDPOINT_NAME}` - [ ] Wait: `aws sagemaker wait endpoint-deleted --endpoint-name ${SAGEMAKER_ENDPOINT_NAME}` - [ ] `aws sagemaker delete-endpoint-config --endpoint-config-name ${SAGEMAKER_ENDPOINT_CONFIG_NAME}` - [ ] `aws sagemaker delete-model --model-name ${SAGEMAKER_CURRENT_MODEL_NAME}` - [ ] `aws sagemaker delete-model --model-name ${SAGEMAKER_CHALLENGER_MODEL_NAME}` - [ ] `aws s3 rm s3://${S3_BUCKET}/p3-08/ --recursive` - [ ] Print `A/B endpoint and all resources deleted.` ## Source code A reference implementation of A/B Testing — README, guide, source and tests — is at https://github.com/confidentprep14-web/ai-engineering-projects/tree/main/ab-testing.