Train and deploy an XGBoost model

In this walkthrough, you build a customer churn prediction model using XGBoost in Amazon SageMaker Unified Studio. You:

• Download and explore the IBM Telco Customer Churn dataset

• Prepare tabular data and engineer features

• Train an XGBoost model with a SageMaker AI training job

• Track the experiment and register the model with MLflow

• Deploy the model to a real-time inference endpoint

• Test predictions on held-out data

This walkthrough is intended for data scientists and ML engineers. It assumes you are familiar with Python, SQL, and basic machine learning concepts.

Overview

XGBoost (extreme gradient boosting) is a popular open-source implementation of the gradient-boosted trees algorithm that handles classification, regression, and ranking problems. Using XGBoost with Amazon SageMaker Unified Studio gives you:

• Managed training – Run training as a SageMaker AI training job on a dedicated instance, separate from your notebook.

• Experiment tracking – Log parameters and metrics to a managed MLflow tracking server and register the resulting model.

• Managed hosting – Deploy the trained model to a real-time endpoint with a single call.

Prerequisites

Complete the following before you begin:

• You have an AWS account with permissions to create SageMaker AI resources. For setup instructions, see Set up an AWS account and create an administrator user.

• You have a Amazon SageMaker Unified Studio domain with a user configured. For domain setup, see Quick setup. For user access, see Managing users.

• You have a Amazon SageMaker Unified Studio project created with the All capabilities project profile. For instructions, see Create a project.

• You have an MLflow tracking server created in your project. Copy the tracking server ARN — you supply it in Step 5: Track and register the model with MLflow. For instructions, see Track experiments using MLflow.

You download the dataset in Step 2 — you do not need to bring your own data.

Step 1: Open your notebook and configure your project bucket

This walkthrough supports two notebook environments inside Amazon SageMaker Unified Studio. Launch either one using the following steps.

Option A: SageMaker Notebooks

1. Open your Amazon SageMaker Unified Studio project.

2. In the navigation pane, choose Notebooks.

3. Choose Create Notebook, which provisions a SageMaker AI notebook by default.

Note

The screenshots in this walkthrough show the SageMaker Notebooks experience.

Option B: JupyterLab

1. Open your Amazon SageMaker Unified Studio project.

2. In the navigation pane, under IDEs, choose JupyterLab.

3. Wait for the space to be ready, then create a new notebook with a Python 3 kernel.

4. Install the SageMaker AI SDK and restart the kernel:

   %pip install --upgrade sagemaker

   After the install completes, restart the kernel (Kernel → Restart Kernel) before continuing.

The remaining steps are the same in both environments unless explicitly noted.

Configure your project bucket

Find your project Amazon S3 bucket and prefix. Choose the three-dot menu next to Shared and copy the bucket URI. The URI has the form s3://<bucket>/shared.

Example bucket and prefix values

# Example bucket and prefix values
BUCKET = "amazon-sagemaker-<account-id>-<region>-<suffix>"
PREFIX = "shared"

Important

Use your values any time the code sets a variable called bucket or prefix in the following steps.

Step 2: Download and explore the dataset

This walkthrough uses the publicly available IBM Telco Customer Churn dataset (7,043 customer records).

Download the dataset in your notebook by running the following code in a cell:

Download the dataset

import requests

url = "https://raw.githubusercontent.com/IBM/telco-customer-churn-on-icp4d/master/data/Telco-Customer-Churn.csv"
filename = "Telco-Customer-Churn.csv"

response = requests.get(url)
response.raise_for_status()

with open(filename, 'wb') as f:
    f.write(response.content)

Using JupyterLab instead of SageMaker Notebooks?

Open a terminal in JupyterLab (File → New → Terminal) and download the dataset:

wget https://raw.githubusercontent.com/IBM/telco-customer-churn-on-icp4d/master/data/Telco-Customer-Churn.csv -O Telco-Customer-Churn.csv

In your notebook, load and review the data:

Load and review the data

import pandas as pd

raw = pd.read_csv("Telco-Customer-Churn.csv")
raw.head()

The dataset contains one row per customer with 20 feature columns (including tenure, MonthlyCharges, TotalCharges, Contract, InternetService) and the target column Churn (Yes/No).

Step 3: Prepare the data for XGBoost

Clean the data, engineer features, and place the label in the first column:

Clean data and engineer features

import numpy as np

raw["TotalCharges"] = pd.to_numeric(raw["TotalCharges"], errors="coerce")
raw = raw.dropna(subset=["TotalCharges"])

model_df = pd.DataFrame({
    "churn_label":       (raw["Churn"] == "Yes").astype(int),
    "tenure_months":     raw["tenure"].astype(int),
    "monthly_charges":   raw["MonthlyCharges"].astype(float),
    "total_charges":     raw["TotalCharges"].astype(float),
    "lifetime_value":    raw["MonthlyCharges"].astype(float) * raw["tenure"].astype(int),
    "is_month_to_month": (raw["Contract"] == "Month-to-month").astype(int),
    "is_fiber_optic":    (raw["InternetService"] == "Fiber optic").astype(int),
})

Split into train (70%), validation (15%), and test (15%) sets and upload to Amazon S3:

Split and upload to S3

import boto3, sagemaker

bucket  = BUCKET
prefix  = PREFIX
session = sagemaker.Session(default_bucket=bucket)

shuffled = model_df.sample(frac=1, random_state=1729).reset_index(drop=True)
n = len(shuffled)

train      = shuffled.iloc[: int(0.70 * n)]
validation = shuffled.iloc[int(0.70 * n) : int(0.85 * n)]
test       = shuffled.iloc[int(0.85 * n) :]

train.to_csv("/tmp/train.csv", header=False, index=False)
validation.to_csv("/tmp/validation.csv", header=False, index=False)
test.to_csv("/tmp/test.csv", index=False)

for name in ["train", "validation", "test"]:
    boto3.Session().resource("s3").Bucket(bucket).Object(
        f"{prefix}/datasets/{name}/{name}.csv"
    ).upload_file(f"/tmp/{name}.csv")

print("Splits written to", f"s3://{bucket}/{prefix}/datasets/")

After the cell finishes, the files appear in your notebook file browser.

Step 4: Train the XGBoost model

Configure and launch the training job:

Configure and launch training job

import sagemaker
from sagemaker.estimator import Estimator
from sagemaker.inputs import TrainingInput
from sagemaker import image_uris, get_execution_role

bucket  = BUCKET
prefix  = PREFIX
session = sagemaker.Session(default_bucket=bucket)
region  = session.boto_region_name
role    = get_execution_role()

image_uri = image_uris.retrieve(framework="xgboost", region=region, version="1.7-1")

xgb = Estimator(
    image_uri=image_uri,
    role=role,
    instance_count=1,
    instance_type="ml.m5.xlarge",
    output_path=f"s3://{bucket}/{prefix}/model",
    sagemaker_session=session,
)

xgb.set_hyperparameters(
    objective="binary:logistic",
    eval_metric="auc",
    num_round=200,
    max_depth=5,
    eta=0.2,
    subsample=0.8,
)

train_in = TrainingInput(f"s3://{bucket}/{prefix}/datasets/train/", content_type="csv")
val_in   = TrainingInput(f"s3://{bucket}/{prefix}/datasets/validation/", content_type="csv")

xgb.fit({"train": train_in, "validation": val_in})

Note

Training takes approximately 5 minutes. Monitor progress under AI/ML → Training jobs in Amazon SageMaker Unified Studio.

Step 5: Track and register the model with MLflow

Log hyperparameters and register the trained model:

Log and register model with MLflow

import os, mlflow

mlflow.set_tracking_uri("<your-mlflow-tracking-server-arn>")
mlflow.set_experiment("customer-churn-xgboost")

with mlflow.start_run(run_name="xgboost-churn") as run:
    mlflow.log_params(xgb.hyperparameters())
    model_uri = f"s3://{bucket}/{prefix}/model/{xgb.latest_training_job.name}/output/model.tar.gz"
    mlflow.register_model(model_uri=model_uri, name="customer-churn-xgboost")

To view the experiment and registered model, choose MLflow in the navigation pane, open the MLflow Tracking Servers tab, and choose Open MLflow. In the MLflow UI, choose Models in the navigation pane to see the registered customer-churn-xgboost model.

Step 6: Deploy the model to a real-time endpoint

Deploy the model to a real-time endpoint:

Deploy to endpoint

from sagemaker.serializers import CSVSerializer

predictor = xgb.deploy(
    initial_instance_count=1,
    instance_type="ml.m5.large",
    serializer=CSVSerializer(),
)

Important

The endpoint runs on a dedicated ml.m5.large instance and incurs costs as long as it is in service. Delete the endpoint when you are finished testing (see Step 8: Clean up).

Monitor deployment status under the Inference endpoints tab in the navigation pane. The endpoint is ready when status shows InService.

Step 7: Test predictions

Score the test set:

Score the test set

test = pd.read_csv(f"s3://{bucket}/{prefix}/datasets/test/test.csv")
features = test.drop(columns=["churn_label"])

scores = predictor.predict(features.to_numpy()).decode("utf-8")
test["churn_probability"] = [float(s) for s in scores.strip().split()]

test[["churn_label", "churn_probability"]].head(10)

Values closer to 1.0 indicate higher churn likelihood.

Troubleshooting

The following table lists common issues and their resolutions.

Issue	Resolution
AttributeError: module 'sagemaker' has no attribute 'Session'	Run %pip install --upgrade sagemaker, restart the kernel, then re-run.
S3UploadFailedError: AccessDenied when uploading training data	In Amazon SageMaker Unified Studio, your IAM role has s3:PutObject permission only on your project path (for example, s3://bucket/dzd-<domain-id>/<project-id>/dev/*). Include the full project prefix when constructing Amazon S3 paths.
ResourceLimitExceeded on training job	Check Service Quotas → Amazon SageMaker AI for the instance type, or switch to one with available quota.
Poor validation AUC	Increase num_round, tune max_depth and eta, or add more features.
Endpoint stuck in Creating	Check CloudWatch logs for the endpoint. Verify instance availability in your Region.
Older SageMaker AI SDK causes import or attribute errors after upgrade	Run %pip install --upgrade --force-reinstall sagemaker and restart the kernel.

For more information, see the SageMaker AI troubleshooting guide.

Step 8: Clean up

To avoid ongoing charges, delete the resources that you created in this walkthrough.

1. Delete the endpoint:

   predictor.delete_endpoint()

2. Delete the Amazon S3 datasets and model artifacts at s3://{BUCKET}/{PREFIX}/.

3. If you created a domain for this walkthrough, delete it. For instructions, see Deleting a domain.