# Train and deploy an XGBoost model
<a name="sagemaker-xgboost-recipe"></a>

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
<a name="xgboost-overview"></a>

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
<a name="xgboost-prerequisites"></a>

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](https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html).
+ You have a Amazon SageMaker Unified Studio domain with a user configured. For domain setup, see [Quick setup](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-unified-studio-quick-setup.html). For user access, see [Managing users](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-unified-studio-manage-users.html).
+ You have a Amazon SageMaker Unified Studio project created with the **All capabilities** project profile. For instructions, see [Create a project](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-unified-studio-create-project.html).
+ 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](#xgboost-step5). For instructions, see [Track experiments using MLflow](https://docs.aws.amazon.com/sagemaker/latest/dg/mlflow.html).

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
<a name="xgboost-step1"></a>

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

### Option A: SageMaker Notebooks
<a name="xgboost-step1-option-a"></a>

1. Open your Amazon SageMaker Unified Studio project.

1. In the navigation pane, choose **Notebooks**.

1. 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
<a name="xgboost-step1-option-b"></a>

1. Open your Amazon SageMaker Unified Studio project.

1. In the navigation pane, under **IDEs**, choose **JupyterLab**.

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

1. 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
<a name="xgboost-configure-bucket"></a>

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
<a name="xgboost-code-bucket-prefix"></a>

```
# 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
<a name="xgboost-step2"></a>

This walkthrough uses the publicly available [IBM Telco Customer Churn dataset](https://github.com/IBM/telco-customer-churn-on-icp4d) (7,043 customer records).

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

### Download the dataset
<a name="xgboost-code-download-dataset"></a>

```
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
<a name="xgboost-code-load-review"></a>

```
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
<a name="xgboost-step3"></a>

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

### Clean data and engineer features
<a name="xgboost-code-clean-engineer"></a>

```
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
<a name="xgboost-code-split-upload"></a>

```
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
<a name="xgboost-step4"></a>

Configure and launch the training job:

### Configure and launch training job
<a name="xgboost-code-training-job"></a>

```
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
<a name="xgboost-step5"></a>

Log hyperparameters and register the trained model:

### Log and register model with MLflow
<a name="xgboost-code-mlflow-register"></a>

```
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
<a name="xgboost-step6"></a>

Deploy the model to a real-time endpoint:

### Deploy to endpoint
<a name="xgboost-code-deploy-endpoint"></a>

```
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](#xgboost-step8)).

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

## Step 7: Test predictions
<a name="xgboost-step7"></a>

Score the test set:

### Score the test set
<a name="xgboost-code-score-test"></a>

```
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
<a name="xgboost-troubleshooting"></a>

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](https://docs.aws.amazon.com/sagemaker/latest/dg/troubleshooting.html).

## Step 8: Clean up
<a name="xgboost-step8"></a>

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

1. Delete the endpoint:

   ```
   predictor.delete_endpoint()
   ```

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

1. If you created a domain for this walkthrough, delete it. For instructions, see [Deleting a domain](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-unified-studio-delete-domain.html).