All in One View

Last updated on 2026-10-20 | Edit this page

Estimated time: 15 minutes

Overview

Questions

When is cloud compute the right tool for a researcher who already has a cluster or a workstation?
What do SageMaker and Bedrock add on top of renting a machine?

Objectives

Compare cloud compute with campus HPC and lab hardware on hardware generation, cost and control.
Introduce S3, SageMaker and Bedrock.
Preview the two parts of the lesson.

This episode compares cloud compute with the compute researchers usually already have, then introduces the three AWS services the lesson uses: S3, SageMaker and Bedrock.

Cloud vs. the cluster down the hall

Most researchers already have somewhere to run ML, such as a campus HPC cluster, a departmental GPU server, or a lab workstation. How the cloud compares:

Buy once and maintain, or rent the latest. A cluster or lab server is bought once, maintained by someone, and used for five or more years with whatever GPUs were current at purchase. Cloud providers offer each new generation soon after release (Blackwell B200 and GB200 today, Vera Rubin announced for cloud deployment in 2026), billed per second and only while a job runs. The largest systems (72-GPU GB200 racks) are liquid-cooled; most campus data centers are air-cooled and cannot host them without a facility upgrade. Owned hardware costs the same busy or idle, so for GPUs that stay busy owning is cheaper, and for bursty research work renting is.
Control and data. On-prem gives full control over the software stack, no quotas, no egress charges, and data that never leaves campus. In the cloud you work within quotas and account permissions, and data moved in or out crosses the network. An idle rented GPU costs the same as a busy one, so every job in this lesson has a runtime cap.
Efficiency comes from sharing. Two things decide the energy per unit of useful work: how much of the facility’s power reaches the hardware, and how much of the time the hardware is busy. On the first, power usage effectiveness (PUE, total facility power over IT power, 1.0 is perfect) tracks facility scale: AWS reports a global average of 1.14 for 2025, a 2025 survey of US research HPC puts TACC at about 1.15 and MGHPCC near 1.14, while the Uptime Institute’s 2025 survey puts enterprise and colocation data centers at 1.58-1.80, where a departmental server room also sits. On the second, a GPU is only efficient while it is busy. A lab’s own GPUs idle whenever the lab is not training; a campus cluster pools demand across a university and keeps its hardware busier; a cloud provider pools demand across the world and across time zones, so a GPU that one group releases is running someone else’s job minutes later. Providers do not publish GPU utilization, so this part is reasoning rather than data, but a single institution cannot smooth demand the way a global pool can. (PUE figures are self-reported and measure the building, not the GPUs in it.)

SageMaker and Bedrock

Amazon SageMaker is AWS’s managed service for machine learning. You write the training or inference script and pick a machine type; SageMaker runs the script on that machine as a job, with the ML frameworks already installed, and shuts the machine down when the job finishes. In practice that means:

High-performance compute only when needed: SageMaker lets you develop interactively in lightweight, inexpensive notebook environments (or your own laptop) and then launch training, tuning, or inference jobs on more powerful instance types only when necessary. You pay for expensive compute only while a job runs.
Support for custom scripts: Most training and inference scripts can be run using pre-configured estimators or containers that come with popular ML frameworks such as scikit-learn, PyTorch, TensorFlow, and Hugging Face already installed. In many cases, you can simply include a requirements.txt file to add any additional dependencies you need. When you need more control, SageMaker also supports fully custom Docker containers, so you can bring your own code, dependencies, and environments for training, tuning, and inference.
Flexible compute options: SageMaker lets you easily select instance types tailored to your project needs. For exploratory analysis, use a lightweight CPU (e.g., ml.m5.large). For compute-intensive tasks, such as training deep learning models, you can switch to GPU instances for faster processing. We’ll cover instances more in-depth throughout the lesson (and how to select them), but here’s a preview of the different types:
- CPU instances (e.g., ml.m5.large — $0.115/hour): Suitable for general ML workloads, feature engineering, and inference tasks.
- Memory-optimized instances (e.g., ml.r5.2xlarge — approx. $0.65/hour): Best for handling large datasets in memory.
- GPU instances (e.g., ml.g5.xlarge — $1.41/hour): Optimized for compute-intensive tasks like deep learning training and fine-tuning, offering accelerated processing.
- For more details, check out the supplemental “Instances for ML” page. We’ll discuss this topic more throughout the lesson. (Prices here and throughout the lesson are us-east-1 on-demand list prices from the AWS price list of 2026-10-16; other regions are the same or within a few percent.)
Parallel training and tuning: training jobs can run across several instances for large models and datasets, and hyperparameter tuning jobs run many training jobs in parallel with a few lines of configuration.
Ease of orchestration / Simplified ML pipelines: Traditional high-performance computing (HPC) or high-throughput computing (HTC) environments often require researchers to break ML workflows into separate batch jobs, manually orchestrating each step (e.g., submitting preprocessing, training, cross-validation, and evaluation as distinct tasks and stitching the results together later). In SageMaker, a training or tuning job is launched from a notebook with the script, the data location and the instance type, and the orchestration is handled for you.
Cost tracking: every job reports its billable seconds, and every resource can be tagged so that the bill can be broken down by project. For many common research workloads SageMaker is cheap. For example, training roughly 100 small to medium-sized models (e.g., logistic regression, random forests, or lightweight deep learning models with a few million parameters) on a small dataset (under 10GB) can cost under $20, making it accessible for many research projects.

Amazon Bedrock is the other AWS service this lesson uses. Instead of running your script on a machine, Bedrock serves models that AWS hosts through an API. In practice that means:

Nothing to start or stop: you call the API and pay per token; when you stop calling it, you stop paying. There is no instance to size, launch, or forget.
A catalog of hosted models: DeepSeek, Claude, Amazon’s own models and others, for generation and for embeddings, each selectable by model ID with no download or GPU.
Cheap at research scale: embedding a few thousand document chunks and answering a few dozen questions costs cents, which is why the RAG part of the lesson starts here.
Only the catalog: if the model you need is not on it, Bedrock cannot serve it. That is the one case where the lesson moves RAG onto SageMaker.
Cost tracking by inference profile: Bedrock calls carry no tags by default; the lesson routes them through tagged application inference profiles so they show up on the bill like everything else.

Amazon S3 is where data lives. Buckets hold files of any size, every SageMaker job reads its inputs from a bucket and writes its outputs back to one, and storage costs about two cents per GB per month. In this lesson you create one bucket of your own in the next episode and use it for the rest of the workshop; the larger datasets are in a shared bucket the instructors staged.

How the pieces fit together

Diagram of the lesson's AWS components: S3 buckets at the top, a notebook instance on the left launching SageMaker training, tuning and processing jobs in the center, and Bedrock on the right called as an API

Everything in the lesson is one of these four boxes. S3 holds your data and every job’s outputs. The notebook instance is where you work; it is small because it only launches things. SageMaker jobs are where compute happens, on an instance that exists only for the job. Bedrock is an API: no instance at all. Everything the notebook does here can also be done from the AWS command line or from Python on your own machine; the optional last episode shows the CLI equivalents.

Lesson outline

The rest of the lesson is hands-on, in two parts:

SageMaker: training, tuning and fine-tuning. Set up an S3 bucket and a small notebook instance, launch training jobs (XGBoost and PyTorch) and a hyperparameter tuning job from it, then fine-tune a vision-language model for handwriting transcription as a training job.
Retrieval-augmented generation (RAG). Build a RAG pipeline over research papers on Amazon Bedrock, then run the same pipeline as SageMaker Processing jobs with an open-weight model.

Every job is tagged, capped, estimated before launch and checked against the bill after.

Key Points

The cloud offers the current GPU generation on demand and bills per second while a job runs; on-prem hardware offers full control and a fixed cost but a fixed generation and a queue.
Renting is cheaper for bursty research work; owning is cheaper for GPUs that stay busy.
An idle rented GPU costs the same as a busy one, so every job in this lesson is launched with a runtime cap.
SageMaker runs your script in a preconfigured framework container on an instance that exists only for the job; Bedrock serves AWS-hosted models per token with nothing to start or stop.

Content from Data Storage: Setting up S3

Last updated on 2026-10-21 | Edit this page

Estimated time: 20 minutes

Overview

Questions

How can I store and manage data effectively in AWS for SageMaker workflows?
What are the best practices for using S3 versus EC2 storage for machine learning projects?

Objectives

Explain data storage options in AWS for machine learning projects.
Describe the advantages of S3 for large datasets and multi-user workflows.
Create an S3 bucket in the class region with a bucket policy and cost-allocation tags, and upload the Titanic data to it.
Estimate S3 storage, transfer and request costs for a small dataset.

Storing data on AWS

Machine learning and AI projects rely on data, making efficient storage and management essential. AWS provides several options for storing data, each with different use cases and trade-offs.

Consult your institution’s IT before handling sensitive data in AWS

When using AWS for research, ensure that no restricted or sensitive data is uploaded to S3 or any other AWS service unless explicitly approved by your institution’s IT or cloud security team. For projects involving sensitive or regulated data (e.g., HIPAA, FERPA, or proprietary research data), consult your institution’s cloud security or compliance team to explore approved solutions. This may include encryption, restricted-access storage, or dedicated secure environments. If unsure about data classification, review your institution’s data security policies before uploading.

Options for storage: EC2 Instance or S3

When working with SageMaker and other AWS services, you have options for data storage, primarily EC2 instances or S3.

What is an EC2 instance?

An Amazon EC2 (Elastic Compute Cloud) instance is a virtual server: a machine with a chosen amount of CPU, memory, disk and optionally GPU that you rent by the second. Elastic is the point of the name: you can start one, resize it to a different type, add more, and shut them down whenever you want, instead of owning a fixed machine. EC2 instances come in hundreds of types, from small web servers to 8-GPU nodes. Every SageMaker notebook instance and every SageMaker job runs on an EC2 instance underneath; SageMaker’s ml.t3.medium, ml.g5.xlarge and so on are EC2 instance types with an ml. prefix and SageMaker’s software and billing on top.

When to store data directly on EC2

Using an EC2 instance for data storage can be useful for temporary or small datasets, especially during processing within a Jupyter notebook. However, this storage is not persistent; if the instance is deleted, the data is erased. (On a SageMaker notebook instance, files under /home/ec2-user/SageMaker survive a Stop and are lost when the instance is Deleted; pip installs and files anywhere else on the instance are lost on every Stop.) Therefore, EC2 is ideal for one-off experiments or intermediate steps in data processing.

Callout

Limitations of EC2 storage

Scalability: EC2 storage is limited to the instance’s disk capacity, so it may not be ideal for very large datasets.
Cost: EC2 storage can be more costly for long-term use compared to S3.
Data Persistence: EC2 data may be lost if the instance is stopped or terminated, unless using Elastic Block Store (EBS) for persistent storage. On a SageMaker notebook instance, only /home/ec2-user/SageMaker survives a Stop, and nothing survives a Delete.

What is an S3 bucket?

Storing data in an S3 bucket is generally preferred for machine learning workflows on AWS, especially when using SageMaker. An S3 bucket is a container in Amazon S3 (Simple Storage Service) where you can store, organize, and manage data files. Buckets act as the top-level directory within S3 and can hold a virtually unlimited number of files and folders, making them ideal for storing large datasets, backups, logs, or any files needed for your project. You access objects in a bucket via a unique S3 URI (e.g., s3://your-bucket-name/your-file.csv), which you can use to reference data across various AWS services like EC2 and SageMaker.

Callout

Benefits of using S3 (recommended for SageMaker and ML workflows)

For flexibility, scalability, and cost efficiency, store data in S3 and load it into EC2 as needed. This setup allows:

Separation of storage and compute: The most essential advantage. Data in S3 remains accessible even when EC2 instances are stopped or terminated, reducing costs and improving workflow flexibility.
Easy data sharing: Datasets in S3 are easier to share with team members or across projects compared to EC2 storage.
Integration with AWS services: SageMaker, Lambda, and other AWS services can read directly from and write back to S3, streamlining ML workflows.
Scalability: S3 handles large datasets efficiently, enabling storage beyond the limits of an EC2 instance’s disk space.
Cost efficiency: S3 storage is generally lower cost than expanding EC2 disk volumes, and you only pay for the storage you use.
Data persistence: Unlike EC2 storage, which is lost when an instance is deleted, S3 ensures long-term data availability.

Recommended approach: S3 buckets

In order to upload our titanic dataset to an S3 bucket on AWS, we’ll follow the below summary procedure (details follow):

Log in to AWS Console and navigate to S3.
Create a new bucket or use an existing one.
Upload your dataset files.
Use the object URL to reference your data in future experiments.

Detailed procedure

Log in to the AWS Console through the sign-in link your instructor emailed you. It looks like https://<org>.awsapps.com/start/#/console?account_id=...&role_name=... and opens the console directly in the workshop account and role after you sign in with your university login. Working on your own account, use the regular console sign-in instead.

2. Navigate to S3

Type “S3” in the search bar
Recommended: Select the star icon to save S3 as a bookmark in your AWS toolbar
Select S3 - Scalable Storage in the Cloud

3. Create a new bucket

Click Create bucket. The form has the sections below; most defaults are right. Change only what is listed.

General configuration

AWS Region: the region your instructor named (in the workshop account, US East (Ohio), us-east-2). This must be the same region your notebook instance and jobs run in.
Bucket type: General purpose (the default).
Bucket namespace: Global namespace. The console recommends the account Regional namespace, but that appends your account ID, region and -an to every bucket name, which would turn sinkorswim-doejohn into sinkorswim-doejohn-123456789012-us-east-2-an everywhere in the lesson. Global keeps the name as you type it.
Bucket name: challenge-lastnamefirstname, all lowercase, e.g. sinkorswim-doejohn: the challenge you are working on (sinkorswim, wattbot2026, badgerscribe, badgercode, …), a hyphen, then your last and first name run together. S3 requires lowercase letters, digits and hyphens only, 3 to 63 characters, starting and ending with a letter or digit. Buckets are the one place AWS forces lowercase; for the notebook instance and the tags we use the same words with capitals, SinkOrSwim and DoeJohn. This is the only bucket you will create in the workshop. Every later episode writes its outputs (trained models, embeddings, predictions) into this same bucket.

Callout

Wrong region costs money and breaks jobs

A SageMaker training or processing job reads only from buckets in its own region; a bucket created elsewhere shows up later as “no such bucket” or “access denied” from the job. Reading data across regions from a notebook does work, but every byte is billed as inter-region transfer (about $0.01-0.02 per GB each way), and a model that reads its training set every epoch pays that on every pass. Check the region before you click Create; it cannot be changed afterwards.

Object Ownership: leave ACLs disabled (the default). Access is managed through the bucket policy you add in the next step and through IAM roles, which is what AWS recommends.

Block Public Access settings for this bucket: leave Block all public access on (the default). Nothing in this lesson needs public access; if you ever do, use a policy or a signed URL instead of opening the bucket.

Bucket Versioning: leave Disable (the default). Versioning keeps every overwritten copy and raises storage costs; it is useful for datasets that change over time, not for this workshop.

Tags (NOT OPTIONAL): tags are how the bill is broken down by project, person and purpose; this is the one section you must fill in. Add the following three keys, spelled exactly (they are the account’s approved cost-allocation keys, see KB 72454). You will put the same three keys on every resource you create in this lesson.

Project: your challenge name (SinkOrSwim, WattBot2026, BadgerScribe, BadgerCode, …), e.g. SinkOrSwim
Name: your last and first name run together, e.g. DoeJohn
Purpose: one of five fixed values, here Storage

AWS allows tag values of up to 256 characters made of letters, digits, spaces and + - = . _ : / @, and they are case-sensitive. We use capitalised words with no spaces, and the same spelling on every resource: SinkOrSwim and sinkorswim would be separate rows in Cost Explorer.

Purpose takes one of five values, so that the bill can be summed by kind of work rather than by individual job. Which job it was is already in the job’s name.

`Purpose`	Used for
`Storage`	your bucket
`Notebook`	the controller notebook instance
`Training`	training jobs, including the fine-tune
`Tuning`	hyperparameter tuning jobs
`Inference`	Bedrock calls, RAG jobs, batch transcription

Screenshot showing required tags for an S3 bucket

Callout

Why these three keys, and why they must be spelled exactly

Tags only show up in cost reports if the tag key has been activated as a cost allocation tag at the account level. On UW–Madison’s AWS accounts that activation is done centrally by the Public Cloud team for a fixed list of approved keys, documented at AWS - Tagging with an AWS Account (KB 72454). Name, Project, and Purpose are on that list; ModelSize, Team, project (lowercase) or any other spelling are not, and spend tagged with them is invisible in Cost Explorer. The KB page says it plainly: cut and paste, because any deviation stops the tag from reporting.

Two more things to know now, because they come up in every later episode:

Tags never propagate. Tagging this bucket tags the bucket only. The notebook instance, every training job, processing job and Bedrock inference profile you create later needs its own copy of the same three tags.
Outside UW, you activate the keys yourself, once, under Billing and Cost Management → Cost allocation tags. Until you do, the tags exist on the resources but not on the bill.

Default encryption: leave SSE-S3 (the default). Objects are encrypted at rest with keys S3 manages; KMS keys only matter when you need to control the keys yourself.

Advanced settings: leave closed.

Click Create bucket at the bottom.

4. Edit bucket policy

Once the bucket is created, the console opens the new bucket’s page (a green “Successfully created bucket” banner at the top, an empty Objects list below). If you land on the bucket list instead, click your bucket’s name. We’ll have to edit our bucket’s policy to allow ourselves proper access to any files stored there (e.g., read from bucket, write to bucket). To set these permissions:

Click the Permissions tab (next to Objects, Metadata, Properties).
Scroll down to Bucket policy and click Edit. Paste the following policy, editing the bucket name “sinkorswim-doejohn” to reflect your bucket’s name and replacing <account-id> with the 12-digit AWS account ID shown at the top right of the console (your instructor will also give it to you)

JSON

{
	"Version": "2012-10-17",
	"Statement": [
		{
			"Effect": "Allow",
			"Principal": {
			    "AWS": [
			        "arn:aws:iam::<account-id>:role/ml-sagemaker-use",
			        "arn:aws:iam::<account-id>:role/ml-sagemaker-bedrock-use"
		        ]
			},
			"Action": [
				"s3:GetObject",
				"s3:PutObject",
				"s3:DeleteObject",
				"s3:ListBucket",
				"s3:ListMultipartUploadParts"
			],
			"Resource": [
				"arn:aws:s3:::sinkorswim-doejohn",
				"arn:aws:s3:::sinkorswim-doejohn/*"
			]
		}
	]
}

What the policy says, line by line:

Principal is who gets access. SageMaker never acts as you; it acts as an IAM role attached to your notebook instance and to every job you launch (the execution role). In the workshop account there are two: ml-sagemaker-bedrock-use, which you will pick when you create your notebook instance in the next episode and which every job inherits, and the older SageMaker-only ml-sagemaker-use, listed so that a notebook created with it still reads the bucket. The <account-id> is the 12-digit number shown next to the account name at the top right of the console.
Action is what they may do: read objects, write objects, delete objects, list the bucket’s contents, and list the parts of a large upload in progress.
Resource is where: the bucket itself (for listing) and everything in it (/*, for the object actions).

Without this, the notebook can see the bucket but a job’s TrainingInput fails with “access denied” when it tries to list and download its data.

Callout

Doing this on your own account

Everything above transfers; only the role name changes.

Find your execution role. When you create a notebook instance (next episode) you either pick an existing role or let the console create one, named like AmazonSageMaker-ExecutionRole-20261020T120000. Its ARN is shown on the instance’s page under Permissions and encryption, and looks like arn:aws:iam::<account-id>:role/service-role/AmazonSageMaker-ExecutionRole-....
Put that ARN in the policy as the only entry under Principal, and your bucket name under Resource.
Or skip the bucket policy by putting sagemaker in the bucket name: the console-created role carries the AmazonSageMakerFullAccess policy, which already grants read and write on any bucket whose name contains sagemaker. A bucket named sinkorswim-doejohn is not covered, which is why the bucket policy is needed at all.

Either way, access comes from one of two places: the role’s own IAM policy (what the role may do anywhere) or the bucket’s policy (what this bucket allows). In the same account, allowing in either one is enough. The SageMaker execution role documentation covers the role side. UW-Madison researchers can bring questions about roles, policies and accounts to the Public Cloud team’s weekly office hours (Thursdays at 2 pm on Zoom) or email cloud-services@cio.wisc.edu.

Two buckets appear in this workshop, and this is the only one you create. Your bucket holds the data you bring (the Titanic files, uploaded below) and everything the later episodes produce. The larger datasets for the fine-tuning and RAG episodes (hundreds of scanned letters, a corpus of PDFs) live in a shared, read-only bucket that the instructors stage before class: one copy for everyone, instead of a hundred copies of the same files. You read from it and never upload to it; in the workshop account it is ml-marathon-shared, and the episodes that use it have the name in their first cell.

5. Upload files to the bucket

If you haven’t downloaded these files yet (part of workshop setup), download the data for this workshop: data.zip
- Extract the zip folder contents (Right-click -> Extract all on Windows; Double-click on mac)
- Save the two data files (train and test) to a location where they can easily be accessed. E.g., …
  - ~/Downloads/data/titanic/titanic_train.csv
  - ~/Downloads/data/titanic/titanic_test.csv
Navigate to the Objects tab of your bucket, then Upload.
Add Files (titanic_train.csv, titanic_test.csv) and click Upload to complete. These two files are what the next episodes read: the Accessing S3 episode loads them into a notebook straight from this bucket, and the training and tuning episodes point their jobs at the same two objects.

6. Take note of S3 URI for your data

After uploading a file to S3, click on the file to locate its Object URI (e.g., s3://sinkorswim-doejohn/titanic_train.csv). The Uniform Resource Identifier (URI) is a unique address that specifies the location of the file within S3. This URI is essential for referencing data in AWS services like SageMaker, where it will be used to load data for processing and model training.

S3 bucket costs

S3 bucket storage incurs costs based on data storage, data transfer, and request counts.

Storage costs

Storage is charged per GB per month. Typical: Storing 10 GB costs approximately $0.23/month in S3 Standard ($0.023/GB-month in us-east-1).
Pricing Tiers: S3 offers multiple storage classes (Standard, Intelligent-Tiering, Glacier, etc.), with different costs based on access frequency and retrieval times. Standard S3 fits most purposes.

Data transfer costs

Uploading data to S3 is free.
Downloading data (out of S3) incurs charges (~$0.09/GB). Be sure to take note of this fee, as it can add up fast for large datasets.
In-region transfer (e.g., S3 to EC2) is free, while cross-region data transfer is charged (~$0.02/GB). This is one reason your bucket and your SageMaker resources must be in the same region.
Deleting objects is free.

Request costs

Every operation on a bucket is an API request, and each one is billed by count. Two kinds matter: GET ($0.0004 per 1,000), charged every time an object is read, and PUT/LIST ($0.005 per 1,000), charged every time an object is written or a prefix is listed. One request is one object, whatever its size: reading a 5 GB file is one GET; reading 5,000 one-kilobyte files is 5,000 GETs.

A GET happens whenever code opens an object: pd.read_csv("s3://..."), s3.download_file(...), a training job pulling its input channel (one GET per object in the prefix), a PyTorch Dataset that fetches images from S3 one by one. The last case is the one that bites in ML: a dataset of 100,000 images read once per epoch is 100,000 GETs per epoch, and a 50-epoch run is 5 million requests, about $2. Still small, but no longer zero, and a hyperparameter sweep multiplies it by the number of trials. The usual fix is the one SageMaker already uses for input channels: copy the data to the instance once (File mode) and read it from local disk during training.

For this lesson, requests round to nothing: two CSVs read a few dozen times. Over a research project, the costs rank in a stable order: storage is the one that accumulates (every dataset, model and output you forget to delete keeps billing monthly), transfer out is the one that spikes (one bulk download of results to a laptop), and requests matter only with many small files read repeatedly.

To calculate specific costs based on your needs, storage class, and region, refer to AWS’s S3 Pricing Information.

Challenge

Challenge: Estimating Storage Costs

1. Estimate the total cost of storing 1 GB in S3 for one month assuming:

Storage duration: 1 month
Storage region: us-east-1
Storage class: S3 Standard
Data will be retrieved 100 times for model training and tuning by SageMaker jobs in the same region (GET requests)
Data will be downloaded once to your laptop and then deleted after the project concludes

Hints

S3 storage cost: $0.023 per GB per month (us-east-1)
Data transfer cost (download to your laptop): $0.09 per GB (us-east-1 out to internet)
In-region transfer (reads by SageMaker jobs): free
GET requests cost: $0.0004 per 1,000 requests (each model training will incur one GET request)
Deleting objects: free
Check the AWS S3 Pricing page for more details.

2. Repeat the above calculation using the following dataset sizes: 10 GB, 100 GB, 1 TB (1024 GB)

3. Wrong region. Suppose the 100 GB bucket had been created in us-east-1 but the SageMaker jobs run in us-east-2. The 100 training reads are now cross-region transfers at $0.02 per GB each. What does that add, and how does it compare to the download charge?

Show me the solution

Using the S3 Standard rate in us-east-1:

1 GB:
- Storage: 1 GB * $0.023 = $0.023
- Download: 1 GB * $0.09 = $0.09
- GET Requests: 100 requests * $0.0004 per 1,000 = $0.00004
- Deletion: $0
- Total Cost: $0.11304
10 GB:
- Storage: 10 GB * $0.023 = $0.23
- Download: 10 GB * $0.09 = $0.90
- GET Requests: 100 requests * $0.0004 per 1,000 = $0.00004
- Deletion: $0
- Total Cost: $1.13004
100 GB:
- Storage: 100 GB * $0.023 = $2.30
- Download: 100 GB * $0.09 = $9.00
- GET Requests: 100 requests * $0.0004 per 1,000 = $0.00004
- Deletion: $0
- Total Cost: $11.30004
1 TB (1024 GB):
- Storage: 1024 GB * $0.023 = $23.55
- Download: 1024 GB * $0.09 = $92.16
- GET Requests: 100 requests * $0.0004 per 1,000 = $0.00004
- Deletion: $0
- Total Cost: $115.71004

These costs assume no additional request charges beyond those for storage, the download, and GET requests for training.

Two things to notice. The GET line is the same four-thousandths of a cent at every size: request costs depend on how many objects you touch, not how big they are. And the download line grows from 80% of the total at 1 GB to almost all of it at 1 TB: for a dataset of any size, moving it out of AWS costs several times more than keeping it there for a month, so plan to keep large data in the region and bring code to it, not the other way round.

Wrong region, 100 GB: 100 reads x 100 GB x $0.02 = $200 of inter-region transfer, on top of the $11.30 above. The one-time download to a laptop was $9; the misplaced bucket costs twenty times that, silently, and keeps costing with every job. This is why the region is the first setting on the create-bucket form and why the warning above tells you to check it before clicking Create.

Removing unused data (complete after the workshop)

After you are done using your data, it’s important to practice good resource stewardship and remove the unneeded files/buckets.

Option 1: Delete data only (if you plan to reuse bucket for other datasets)

Go to S3, navigate to the bucket.
Select files to delete, then Actions > Delete.
An empty bucket costs nothing: S3 bills for stored bytes and requests, and an empty bucket has neither. Keeping it is fine; the lesson’s end-of-workshop cleanup deletes it only because the shared account is swept between terms.

Option 2: Delete the S3 bucket entirely (you no longer need the bucket or data)

Select the bucket, click Actions > Delete.
Type the bucket name to confirm deletion.

Please complete option 2 following this workshop. On your own account, either option stops the bill: the charges come from the objects, not the bucket.

Key Points

Use S3 for scalable, cost-effective, and flexible storage.
EC2 storage is fairly uncommon, but may be suitable for small, temporary datasets; only /home/ec2-user/SageMaker survives a Stop, and nothing survives a Delete.
Create the bucket in the same region as your SageMaker notebook, and tag it with the same three keys (Name, Project, Purpose) you will put on every other resource.
Track your S3 storage costs, data transfer, and requests to manage expenses.
Regularly delete unused data or buckets to avoid ongoing costs.

Content from Notebooks as Controllers

Last updated on 2026-10-21 | Edit this page

Estimated time: 20 minutes

Overview

Questions

How do you set up and use SageMaker notebooks for machine learning tasks?
How can you manage compute resources efficiently using SageMaker’s controller notebook approach?

Objectives

Explain why the lesson works from a SageMaker notebook instance, and what the controller pattern is.
Create a tagged notebook instance on a small instance type, open the lesson notebooks in it, and stop it when done.

Setting up our notebook environment

In this episode we set up the SageMaker notebook instance that every later episode runs from.

What a SageMaker notebook instance is

A SageMaker notebook instance is a virtual machine that AWS runs for you with JupyterLab already installed, a set of Python kernels that come with the AWS SDKs (boto3 and the SageMaker Python SDK), and an IAM role attached that decides what the code running on it may do in your account. You create it from the console and open JupyterLab in a browser tab. Compute is billed per second while it is running and not while it is stopped. The instance’s disk is separate: the files you keep under /home/ec2-user/SageMaker/ survive a stop, and that storage volume keeps billing at about $0.14 per GB-month (about $2.80 a month for the 20 GB this lesson asks for) until the instance is deleted.

Using the notebook as a controller

The notebook instance is a controller: a small, cheap machine where you write code and use the SageMaker Python SDK to launch larger instances for model training, batch processing, and hyperparameter tuning. Those instances start, run your script, write their results to S3, and shut themselves down, so you pay for powerful hardware only for the minutes a job needs it.

Diagram from the overview episode: S3 buckets at the top, the notebook instance on the left launching SageMaker training, tuning and processing jobs, and Bedrock on the right

The notebook is the box on the left of the diagram from the overview episode. Everything else in the picture is launched or called from it, and every later episode uses this same instance.

Why we work from a notebook instance in this lesson

The SageMaker SDK runs anywhere Python does, so the lesson could be run from a laptop. We use a notebook instance for four reasons:

No credential setup. The IAM role attached to the instance authorizes every call made from it, so there is nothing to install or log in to. Working from a laptop is not hard, aws configure sso and aws sso login take a few minutes, but in a room of a hundred people that is a hundred setups, plus SSO sessions expiring mid-episode and the occasional person pasting an access key instead. The role sidesteps all of it.
Everyone has the same environment. The kernels, the SDK version and the paths are identical for every learner, so a cell that works for the instructor works for you.
Instructors can see what happened. Every notebook instance in the shared account shows up in the console, and an instructor can open yours and see the notebook that launched a job, with its outputs. Jobs launched from a laptop are just as visible in the job list, but the code that launched them is on your machine, so when something breaks or a charge looks wrong there is less to go on.
No egress charges. AWS bills data transfer when it leaves AWS, not when it moves within a region. The notebook instance sits in the same region as your bucket, so reading or downloading files from S3 to the instance is free. Pulling the same files to a laptop costs $0.09 per GB wherever the laptop is: “region” means the AWS data center your resources live in, and a machine outside AWS is never in one.

We’ll follow these steps to create our first “SageMaker notebook instance”.

1. Navigate to SageMaker

In the AWS Console, search for SageMaker.
Recommended: Select the star icon next to Amazon SageMaker AI to save SageMaker as a bookmark in your AWS toolbar
Select Amazon SageMaker AI

2. Create a new notebook instance

In the SageMaker left-side menu, click on Notebook instances, then click Create notebook instance.
Notebook name: the same words as your bucket, capitalised: Challenge-LastnameFirstname, e.g. SinkOrSwim-DoeJohn. One instance serves the whole workshop; the individual .ipynb notebooks inside it are where the per-episode work lives. SageMaker allows letters, digits and hyphens (no spaces or underscores, up to 63 characters); capitals are fine here and in tags, only bucket names must be lowercase. This name is also reused as the prefix of every training job you launch from the notebook, which is how you will find your jobs in a shared account.
Instance type: sets the CPU, memory and price of the virtual machine. Our notebook is a controller, so pick a small one: ml.t3.medium (2 vCPU, 4 GB RAM, $0.05/hour). The jobs we launch from it will run on more powerful instances; the supplemental Instances for ML page compares them.
Platform identifier: This is an internal AWS setting related to the environment version and underlying platform. You can leave this as the default.
Additional configuration (leave collapsed): holds the Volume size in GB of the instance’s persistent disk, default 5 GB. Expand it and set 20 GB: the RAG episodes unpack about a gigabyte of PDFs and page images, the fine-tuning episode keeps its data here, and 5 GB fills up on the second run. The disk is the one thing a stopped instance keeps billing for, $0.14 per GB-month, so 20 GB is $2.80 a month until the instance is deleted; do not go larger than that without a reason.
Permissions and encryption:
- IAM role: Select ml-sagemaker-use. It is pre-configured for this workshop with access to SageMaker and, for the RAG episodes, Bedrock, with some restrictions to prevent overuse or misuse of resources; ml-sagemaker-bedrock-use, if the list offers it, is the same thing. Outside of the workshop, you create or select a role that includes the AmazonSageMakerFullAccess policy and, for Bedrock, the policy in the Account preparation checklist instructor page.
- Root access: Determines whether the user can run administrative commands within the notebook instance. You should Enable root access to allow installing additional packages if/when needed.
- Encryption key (skip): While we won’t use this feature for the workshop, it is possible to specify a KMS key for encrypting data at rest if needed.
Network (skip): Networking settings are optional. Configure them if you’re working within a specific VPC or need network customization.
Git repositories configuration (skip): You don’t need to complete this configuration. Instead, we’ll run a clone command from our notebook later to get our repo setup. This approach is a common strategy (allowing some flexiblity in which repo you use for the notebook).

Tags (NOT OPTIONAL): the same three keys you put on your bucket, with the same values except Purpose. Tags never propagate from one resource to another, so the instance needs its own copy; the previous episode’s callout on why these keys and no others applies here too.

Key	Value
`Project`	your challenge name (`SinkOrSwim`, `WattBot2026`, `BadgerScribe`, `BadgerCode`, …), e.g. `SinkOrSwim`
`Name`	your name, e.g. `DoeJohn`
`Purpose`	what this resource is for, `Notebook`

Screenshot showing the three required tags on a SageMaker notebook instance

Click Create notebook instance.

Load pre-filled Jupyter notebooks

The new instance appears in the list with status Pending. It takes three to five minutes to change to InService, and the Open JupyterLab link does not appear until it does; refreshing the page does not speed it up. Starting a stopped instance later takes about as long. Once it is InService, click Open JupyterLab next to it. In JupyterLab, create a new notebook with the conda_python3 kernel (we aren’t training anything yet, so the plain Python environment is enough).

In the new notebook, run:

SH

!pwd
!ls

The first line prints /home/ec2-user/SageMaker: new notebooks start in that folder. The second lists it: the notebook you just created, and a folder called lost+found. That folder is not yours and holds nothing. Linux creates one at the top of every formatted disk, as the place its repair tool would put file fragments after a crash; it appears here because SageMaker/ is its own disk, and it is the sign that you are looking at that disk and not the instance’s system drive. Leave it alone. SageMaker/ is a separate disk, the volume whose size sits under Additional configuration on the create form (the 20 GB you set). The rest of the instance lives on the system disk and is reset when the instance stops; SageMaker/ is the only place that survives a stop. JupyterLab’s file browser opens there for that reason, and it is where everything in this lesson goes.

Now clone the lesson repository into it:

SH

!git clone https://github.com/carpentries-incubator/ML_with_AWS_SageMaker.git

Then open ML_with_AWS_SageMaker/notebooks/04_Accessing-S3-via-SageMaker-notebooks.ipynb from the file browser to begin the first notebook.

At the end of class: stop the instance

A notebook instance bills for every second it is InService. At the end of class, go to Notebook instances, select yours and choose Stop. Your files under /home/ec2-user/SageMaker (including the repos you clone) survive a Stop and are there when you start it again; packages you pip install are not, so expect to reinstall them. When the workshop is over, Delete the instance: a stopped instance still bills for its storage volume ($0.14 per GB-month), and the cleanup episode walks through this.

Key Points

Use a minimal SageMaker notebook instance as a controller to manage larger, resource-intensive tasks.
The notebook does not do the heavy computation; later episodes use it to launch training, tuning and processing jobs on separate instances.
Every resource gets the same three tags; they never propagate, so each new resource needs its own.
Stop the notebook instance at the end of class, and delete it when the workshop is over.

Content from Accessing and Managing Data in S3 with SageMaker Notebooks

Last updated on 2026-10-21 | Edit this page

Estimated time: 30 minutes

Overview

Questions

How can I load data from S3 into a SageMaker notebook?
How do I monitor storage usage and costs for my S3 bucket?
What steps are involved in pushing new data back to S3 from a notebook?

Objectives

Read data directly from an S3 bucket into memory in a SageMaker notebook.
Check storage usage and estimate costs for data in an S3 bucket.
Upload new files from the SageMaker environment back to the S3 bucket.

Open the notebook from: /ML_with_AWS_SageMaker/notebooks/04_Accessing-S3-via-SageMaker-notebooks.ipynb. Select the conda_python3 kernel (the same one you used in the last episode).

Set up AWS environment

To begin each notebook, it’s important to set up an AWS environment that will allow seamless access to the necessary cloud resources. Here’s what we’ll do to get started:

Define the Role: We’ll use get_execution_role() to retrieve the IAM role associated with the SageMaker instance. This role specifies the permissions needed for interacting with AWS services like S3, which allows SageMaker to securely read from and write to storage buckets.
Initialize the SageMaker Session: Next, we’ll create a sagemaker.Session() object, which will help manage and track the resources and operations we use in SageMaker, such as training jobs and model artifacts. The session acts as a bridge between the SageMaker SDK commands in our notebook and AWS services.
Set Up an S3 Client using boto3: Using boto3, we’ll initialize an S3 client for accessing S3 buckets directly. Boto3 is the official AWS SDK for Python, allowing developers to interact programmatically with AWS services like S3, EC2, and Lambda.

Starting with these initializations prepares our notebook environment to efficiently interact with AWS resources for model development, data management, and deployment.

Start with the role and the session, and print both. The role is the identity every call from this notebook uses. Check the region line against the bucket you made in the S3 episode: the two must match, or the training jobs in later episodes will fail to find their data and anything you do read will be billed as cross-region transfer. If they differ, the fix is a new bucket in this region, not a new notebook instance.

Work from a project folder beside the lesson clone, so downloads and outputs stay out of the repository. The Titanic episodes all use this one; the fine-tuning and RAG episodes make their own.

PYTHON

!mkdir -p /home/ec2-user/SageMaker/titanic   # a project folder beside the lesson clone, so downloads and outputs stay out of the repository
%cd /home/ec2-user/SageMaker/titanic

PYTHON

import sagemaker                      # SageMaker Python SDK: higher-level API for training jobs, tuning, deployment
from sagemaker import get_execution_role

role = sagemaker.get_execution_role() # the IAM role attached to this notebook instance; every call below acts as this role
session = sagemaker.Session()         # tracks the region, default bucket and SageMaker client for this notebook
print(f"Execution Role: {role}")
print(f"AWS Region: {session.boto_region_name}")  # The region of the notebook instance; your bucket must be in this same region

Callout

Warnings this cell prints

sagemaker.config INFO - Not applying SDK defaults from location: ...: an info line, not an error. The SDK looked for an optional defaults file in two standard places and found none.

SageMakerV2DeprecationWarning: You are using the SageMaker Python SDK v2 ...: the kernels ship SDK v2 and this lesson is written for it. v3 is a rewrite with a different API, so do not install it here. Ignore the warning.

Next, an S3 client, and a listing of the buckets this role can see.

PYTHON

import boto3                          # AWS SDK for Python: low-level access to every AWS service (S3, EC2, ...)

s3 = boto3.client('s3')               # client for S3 API calls (list, upload, download)
bucket_names = [bucket["Name"] for bucket in s3.list_buckets()["Buckets"]]
print(f"Available S3 Buckets: {bucket_names}")  # Every bucket in the account this role can list, including other learners' buckets in a shared account

Reading data from S3

You can either (A) read data from S3 into memory or (B) download a copy of your S3 data into your notebook instance. Since we are using SageMaker notebooks as controllers—rather than performing training or tuning directly in the notebook—the best practice is to read data directly from S3 whenever possible. However, there are cases where downloading a local copy may be useful. We’ll show you both strategies.

A) Reading data directly from S3 into memory

This is the recommended approach for most workflows. By keeping data in S3 and reading it into memory when needed, we avoid local storage constraints and ensure that our data remains accessible for SageMaker training and tuning jobs.

Pros:

Scalability: Data remains in S3, allowing multiple training/tuning jobs to access it without duplication.
Efficiency: No need to manage local copies or manually clean up storage.
Cost-effective: Avoids unnecessary instance storage usage.

Cons:

Network dependency: Requires internet access to S3.
Potential latency: Reading large datasets repeatedly from S3 may introduce small delays. This approach works best if you only need to load data once or infrequently.

Example: Reading data from S3 into memory

Our data is stored in the S3 bucket you created in the previous episode, named ‘challenge-lastnamefirstname’ (e.g., sinkorswim-doejohn). We can use the following code to read data directly from S3 into memory in the Jupyter notebook environment, without actually downloading a copy of train.csv as a local file.

PYTHON

import pandas as pd
# Define the S3 bucket and object key
bucket_name = 'sinkorswim-doejohn'  # replace with your S3 bucket name

# Read the train data from S3
key = 'titanic_train.csv'  # the object key: the file's path inside the bucket
response = s3.get_object(Bucket=bucket_name, Key=key)  # fetch the object; its bytes arrive as a stream in response['Body']
train_data = pd.read_csv(response['Body'])

# Read the test data from S3
key = 'titanic_test.csv'
response = s3.get_object(Bucket=bucket_name, Key=key)
test_data = pd.read_csv(response['Body'])

# check shape
print(train_data.shape)
print(test_data.shape)

# Inspect the first few rows of the DataFrame
train_data.head()

B) Download copy into notebook environment

In some cases, downloading a local copy of the dataset may be useful, such as when performing repeated reads in an interactive notebook session. This download costs nothing: it moves data from S3 to a SageMaker notebook instance in the same region, which AWS does not bill. The same download_file run on your laptop would be a transfer out of AWS at $0.09 per GB.

Pros:

Faster access for repeated operations: Avoids repeated S3 requests.
Works offline: Useful if running in an environment with limited network access.

Cons:

Consumes instance storage: Notebook instances have limited space.
Requires manual cleanup: Downloaded files remain until deleted.

Example

PYTHON

!pwd

PYTHON

# Define the S3 bucket and file location
key = "titanic_train.csv"  # Path to your file in the S3 bucket
local_file_path = "/home/ec2-user/SageMaker/titanic/titanic_train.csv"  # Local path to save the file

# Copy the object to the instance's disk (unlike get_object above, which only read it into memory)
s3.download_file(bucket_name, key, local_file_path)
!ls "/home/ec2-user/SageMaker/titanic/"

Note: You may need to hit refresh on the file explorer panel to the left to see this file. If you get any permission issues…

check that you have selected the appropriate policy for this notebook
check that your bucket has the appropriate policy permissions

Check the current size and storage costs of bucket

It’s a good idea to periodically check how much storage you have used in your bucket. You can do this from a Jupyter notebook in SageMaker by using the Boto3 library, which is the AWS SDK for Python. This will allow you to calculate the total size of objects within a specified bucket.

The code below will calculate your bucket size for you. Here is a breakdown of the important pieces in the next code section:

Paginator: Since S3 buckets can contain many objects, we use a paginator to handle large listings.
Size calculation: We sum the Size attribute of each object in the bucket.
Unit conversion: The size is given in bytes, so dividing by 1024 ** 2 converts it to megabytes (MB).

Note: If your bucket has very large objects or you want to check specific folders within a bucket, you may want to refine this code to only fetch certain objects or folders.

PYTHON

# Initialize the total size counter (bytes)
total_size_bytes = 0

# Use a paginator to handle large bucket listings (one list call returns at most 1000 objects)
# This ensures that even if the bucket contains many objects, we can retrieve all of them
paginator = s3.get_paginator("list_objects_v2")

# Iterate through all pages of object listings
for page in paginator.paginate(Bucket=bucket_name):
    # 'Contents' contains the list of objects in the current page, if available
    for obj in page.get("Contents", []):  
        total_size_bytes += obj["Size"]  # Add each object's size to the total

# Convert the total size to gigabytes for cost estimation
total_size_gb = total_size_bytes / (1024 ** 3)

# Convert the total size to megabytes for easier readability
total_size_mb = total_size_bytes / (1024 ** 2)

# Print the total size in MB
print(f"Total size of bucket '{bucket_name}': {total_size_mb:.2f} MB")

# Print the total size in GB
#print(f"Total size of bucket '{bucket_name}': {total_size_gb:.2f} GB")

Using the lesson’s helper functions

We have added code to calculate bucket size to a helper function called get_s3_bucket_size(bucket_name) for your convenience. There are also some other helper functions in scripts/helpers.py of the lesson repository, which you cloned to /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/ in the previous episode.

Tell Python where the lesson’s scripts/ folder is:

PYTHON

import sys
# Put the lesson's scripts/ folder first on Python's import path
sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")
import helpers  # the lesson's helper functions (bucket size, cost estimate, object listing)

Our helpers.py file can be found in ML_with_AWS_SageMaker/scripts/helpers.py. With this file downloaded, you can call this function via…

PYTHON

helpers.get_s3_bucket_size(bucket_name)

Check storage costs of bucket

To estimate the storage cost of your Amazon S3 bucket directly from a Jupyter notebook in SageMaker, you can use the following approach. This method calculates the total size of the bucket and estimates the monthly storage cost based on AWS S3 pricing.

Note: AWS S3 pricing varies by region and storage class. The example below uses the S3 Standard storage class pricing for the US East (N. Virginia) region as of November 1, 2024. Please verify the current pricing for your specific region and storage class on the AWS S3 Pricing page.

PYTHON

# AWS S3 Standard Storage pricing for US East (N. Virginia) region
# Pricing tiers as of November 1, 2024
first_50_tb_price_per_gb = 0.023  # per GB for the first 50 TB
next_450_tb_price_per_gb = 0.022  # per GB for the next 450 TB
over_500_tb_price_per_gb = 0.021  # per GB for storage over 500 TB

# Calculate the cost based on the size
if total_size_gb <= 50 * 1024:
    # Total size is within the first 50 TB
    cost = total_size_gb * first_50_tb_price_per_gb
elif total_size_gb <= 500 * 1024:
    # Total size is within the next 450 TB
    cost = (50 * 1024 * first_50_tb_price_per_gb) + \
           ((total_size_gb - 50 * 1024) * next_450_tb_price_per_gb)
else:
    # Total size is over 500 TB
    cost = (50 * 1024 * first_50_tb_price_per_gb) + \
           (450 * 1024 * next_450_tb_price_per_gb) + \
           ((total_size_gb - 500 * 1024) * over_500_tb_price_per_gb)

print(f"Estimated monthly storage cost: ${cost:.5f}")
print(f"Estimated annual storage cost: ${cost*12:.5f}")

For your convenience, we have also added this code to a helper function.

PYTHON

monthly_cost, storage_size_gb = helpers.calculate_s3_storage_cost(bucket_name)
print(f"Estimated monthly cost ({storage_size_gb:.4f} GB): ${monthly_cost:.5f}")
print(f"Estimated annual cost ({storage_size_gb:.4f} GB): ${monthly_cost*12:.5f}")

Important Considerations:

Pricing Tiers: AWS S3 pricing is tiered. The first 50 TB per month is priced at $0.023 per GB, the next 450 TB at $0.022 per GB, and storage over 500 TB at $0.021 per GB. Ensure you apply the correct pricing tier based on your total storage size.
Region and Storage Class: Pricing varies by AWS region and storage class. The example above uses the S3 Standard storage class pricing for the US East (N. Virginia) region. Adjust the pricing variables if your bucket is in a different region or uses a different storage class.
Additional Costs: This estimation covers storage costs only. AWS S3 may have additional charges for requests, data retrievals, and data transfers. For a comprehensive cost analysis, consider these factors as well.

For detailed and up-to-date information on AWS S3 pricing, please refer to the AWS S3 Pricing page.

Writing output files to S3

As your analysis generates new files or demands additional documentation, you can upload files to your bucket as demonstrated below. For this demo, we write a small Notes.txt file on the instance and check that it exists:

PYTHON

notes_file_path = "/home/ec2-user/SageMaker/titanic/Notes.txt"
with open(notes_file_path, "w") as f:
    f.write("Workshop notes.\n")

!ls -l "/home/ec2-user/SageMaker/titanic/Notes.txt"

Then upload it. Here it goes under a docs/ prefix; drop the docs/ to put it in the bucket’s root instead.

PYTHON

notes_s3_key = "docs/Notes.txt"   # the object key: a "docs/" prefix plus the file name
# local file -> bucket + key; S3 has no real folders, the "docs/" prefix exists as soon as one key uses it
s3.upload_file(notes_file_path, bucket_name, notes_s3_key)
print(f"Uploaded {notes_file_path} to s3://{bucket_name}/{notes_s3_key}")

After uploading, we can view the objects/files available on our bucket using…

PYTHON

# List and print all objects in the bucket
response = s3.list_objects_v2(Bucket=bucket_name)  # one page (up to 1000 keys); the paginator above handles more

# Check if there are objects in the bucket
if 'Contents' in response:
    for obj in response['Contents']:
        print(obj['Key'])  # Print the object's key (its path in the bucket)
else:
    print("The bucket is empty or does not exist.")

Alternatively, we can substitute this for a helper function call as well.

PYTHON

file_list = helpers.list_S3_objects(bucket_name)
file_list

Challenge

Challenge: before and after

Using helpers.get_s3_bucket_size(bucket_name), record the size of your bucket. Then create a second file (say docs/Notes2.txt with a few lines of text), upload it under the docs/ prefix, and measure the bucket size again. By how much did it grow, and what does that cost per month at $0.023 per GB?
The next episode launches a training job that writes its outputs to your bucket. Why does the lesson keep everything in one bucket, organised with prefixes like docs/, rather than creating a bucket per task?

Show me the solution

One way to do it:

PYTHON

size_before = helpers.get_s3_bucket_size(bucket_name)

with open("Notes2.txt", "w") as f:
    f.write("A second note, to see the bucket grow.\n" * 20)
s3.upload_file("Notes2.txt", bucket_name, "docs/Notes2.txt")

size_after = helpers.get_s3_bucket_size(bucket_name)
print(size_before, size_after)

The bucket grows by the size of the file (a few hundred bytes here), and at $0.023 per GB-month that is well under a thousandth of a cent per month. The same two lines, run before and after a training job, tell you what the job wrote and whether something is accumulating that you should clean up.

A prefix is just part of the object key, so docs/Notes.txt, xgboost/output/... and wattbot/chunks.jsonl can all live in one bucket and still be listed, downloaded or deleted as groups. One bucket means one policy, one set of tags, one region to get right, and one thing to empty and delete at the end; a bucket per task would repeat the setup each time and lengthen the cleanup list.

Key Points

Load data from S3 into memory for efficient storage and processing.
Periodically check storage usage and costs to manage S3 budgets.
Use SageMaker to upload analysis results and maintain an organized workflow.
The lesson repository lives at /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/; every later episode adds its scripts/ folder to sys.path and imports helpers.

Content from (Optional) Pushing Your Notebooks to GitHub from SageMaker

Last updated on 2026-10-21 | Edit this page

Estimated time: 35 minutes

Overview

Questions

How can I put the code I write on a SageMaker notebook instance under version control on GitHub?
How do I authenticate to GitHub from a notebook instance without leaving a password or key behind?
How can I convert notebooks to .py files and keep .ipynb files out of version control?

Objectives

Create a GitHub repository for your own SageMaker work and connect a folder on the notebook instance to it.
Configure Git on a notebook instance to push over HTTPS with a GitHub Personal Access Token (PAT), entered with getpass.
Convert .ipynb files to .py files for cleaner version control.

Open prefilled .ipynb notebook

Open the notebook from: /ML_with_AWS_SageMaker/notebooks/05_Interacting-with-code-repo.ipynb. Select the conda_python3 kernel.

This episode is optional self-study: it is not needed to run the rest of the lesson, but it is the workflow you will want once you start writing code of your own in SageMaker. It also doubles as a short git refresher. You need a GitHub account for it; nothing else in the lesson does.

Step 0: What we are versioning, and where

So far every notebook you have run came from the lesson repository, which you cloned read-only in the notebooks-as-controllers episode. You cannot push to that repository. This episode sets up the usual situation instead: your own work in your own repository. We will create an empty repository on GitHub, make a matching folder on the notebook instance, copy a notebook into it, and push.

On GitHub, click New repository (the + in the top-right corner). Name it sagemaker-work (any name is fine; the cells below use this one). Leave it Public or make it Private, as you prefer. Tick Add a README file so the repository is not empty, which gives Step 6 a pull with something in it. Click Create repository.
Back in JupyterLab, make the folder on the instance and turn it into a git repository. We put it under /home/ec2-user/SageMaker/, the one directory that survives stopping the instance.

PYTHON

%cd /home/ec2-user/SageMaker/
!mkdir -p sagemaker-work
%cd sagemaker-work
!git init -b main   # start an empty git repository here, with "main" as its first branch

Connect it to the repository you just created, replacing YOUR_GITHUB_USERNAME:

PYTHON

# A remote is a copy of the repository elsewhere; "origin" is the conventional name for the main one (here, GitHub)
!git remote add origin https://github.com/YOUR_GITHUB_USERNAME/sagemaker-work.git
!git remote -v   # list the remotes and their URLs, to check the address

Step 1: A GitHub personal access token (PAT)

A notebook instance is a machine that other people may be able to reach (an administrator, or whoever uses it after you) and that you may delete next month. Neither a saved password nor an SSH key belongs on it. GitHub’s answer is a Personal Access Token: a credential you create with only the permissions you need and an expiry date, use over HTTPS, and revoke in one click. We will enter it with getpass each session so that it never appears in a cell or a saved notebook.

Important Note: A PAT is a password with a scope. Give it the minimum permissions, set an expiry, and never paste it into code.

Generate the token on GitHub

Go to Settings by clicking on your profile picture in the upper-right corner of GitHub.
Click Developer settings at the very bottom of the left sidebar.
Select Personal access tokens, then Tokens (classic).
Click Generate new token (classic).
Give it a descriptive name (e.g., “SageMaker notebook instance”) and an expiry (30 days is plenty for a workshop).
Select the minimum permissions needed:
- For a public repository: only public_repo.
- For a private repository: repo.
Click Generate token and copy it immediately: you cannot see it again once you leave the page.

Caution: Treat your PAT like a password. Store it in a password manager, not in a file on the instance, and revoke it (same page, Delete) when the workshop is over; the cleanup episode’s checklist reminds you.

Enter it with `getpass`

input() and getpass.getpass() prompt you in the notebook; the token is held in a Python variable for this kernel session and is not written into the notebook file.

PYTHON

import getpass   # prompts without echoing what you type, so the token never shows on screen or in the saved notebook

username = input("GitHub Username: ")
token = getpass.getpass("GitHub Personal Access Token (PAT): ")  # kept in this variable for the kernel session only

Note: when you re-run the whole notebook later, this cell prompts again, because nothing is saved.

Step 2: Configure Git settings

Git records an author on every commit. Set yours once; --global stores it in your home directory on the instance, so it persists across restarts of this instance (not across instances).

PYTHON

!git config --global user.name "Your name"             # appears as the author in the commit history
!git config --global user.email your_email@example.edu  # use the email on your GitHub account so commits link to your profile

Step 3: Convert `.ipynb` to `.py`

We want our notebook under version control, but a .ipynb file is JSON with outputs and metadata baked in, which makes diffs unreadable and repositories large. The usual practice is to version a plain .py copy of the code and ignore the notebooks.

Benefits of converting to `.py` before committing

Cleaner version control: readable diffs, reviewable pull requests, sane merges.
Script compatibility: a .py file runs from the command line and imports as a module.
Smaller repository: no outputs, no images, no cell metadata.

Install Jupytext, which converts in both directions.

PYTHON

!pip install jupytext   # converts between .ipynb and .py in both directions

Copy the notebook from the Accessing S3 episode into our repository folder (we are still in /home/ec2-user/SageMaker/sagemaker-work) and convert it there. Copying keeps the lesson clone untouched and puts the .py file where git will see it.

PYTHON

# Copy the notebook into the repository folder ("." is the current directory); the lesson clone stays untouched
!cp /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/notebooks/04_Accessing-S3-via-SageMaker-notebooks.ipynb .
!jupytext --to py 04_Accessing-S3-via-SageMaker-notebooks.ipynb   # writes a .py next to it: code and markdown, no outputs

[jupytext] Reading 04_Accessing-S3-via-SageMaker-notebooks.ipynb in format ipynb
[jupytext] Writing 04_Accessing-S3-via-SageMaker-notebooks.py

With several notebooks, convert them all in a loop:

PYTHON

import subprocess
import os

notebooks = [f for f in os.listdir() if f.endswith('.ipynb')]
for notebook in notebooks:
    output_file = notebook.replace('.ipynb', '.py')
    subprocess.run(["jupytext", "--to", "py", notebook, "--output", output_file])  # the same shell command, run from Python
    print(f"Converted {notebook} to {output_file}")

The same loop is helpers.convert_files() in the lesson’s scripts/helpers.py. It converts every .ipynb in the current directory, so run it from inside the folder you mean to convert:

PYTHON

import sys
# Put the lesson's scripts/ folder on Python's import path, as in the Accessing S3 episode
sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")
import helpers
helpers.convert_files(direction="notebook_to_python")

Step 4: Add and commit the `.py` file

Check the status of the repository. Make sure you are in the repository folder first.

PYTHON

%cd /home/ec2-user/SageMaker/sagemaker-work
!git status   # shows the untracked files (the .ipynb and the .py); nothing is staged yet

Add and commit. We add the .py by name rather than git add ., so the .ipynb copy is never tracked (Step 5 ignores it for good).

PYTHON

!git add 04_Accessing-S3-via-SageMaker-notebooks.py   # stage: mark this file to go into the next commit
!git commit -m "Add Accessing-S3 notebook as a script"   # a commit is a snapshot of the staged files; the message says what changed, specifically

Check the status again.

PYTHON

!git status   # the .py is committed; only the untracked .ipynb remains, which Step 5 ignores

Step 5: Ignore notebooks and clutter with `.gitignore`

A .gitignore file lists patterns git should never track. Ignoring *.ipynb across the repository means a stray git add . cannot commit a notebook by accident.

Confirm you are in the repository folder.

PYTHON

!pwd

Create the file (hidden in the file browser because it starts with a dot; ls -a shows it) and add the notebook pattern.

PYTHON

with open(".gitignore", "a") as gitignore:   # "a" appends, so re-running the cell keeps earlier rules
    gitignore.write("# Jupyter notebooks: we version the .py copies instead\n*.ipynb\n")  # one pattern per line; * matches any name
!cat .gitignore

While we are here, ignore the usual clutter too: caches, temporary files, and the NumPy archives the training episodes write next to notebooks.

PYTHON

with open(".gitignore", "a") as gitignore:
    gitignore.write("\n# Cache and temp files\n__pycache__/\n*.tmp\n*.log\n")
!cat .gitignore

Add and commit .gitignore so the rule travels with the repository.

PYTHON

!git add .gitignore   # .gitignore is itself tracked, so everyone who clones gets the same rules
!git commit -m "Ignore notebooks, caches and temp files"

Step 6: Pull, then push

Before pushing, pull: the remote already has the README that GitHub created, and in a shared repository it may have your collaborators’ commits. Pulling first means your push is a simple addition rather than a rejected one.

Which pull strategy

Git has to decide how to combine your local commits with remote ones. The three options:

Merge (pull.rebase false): the remote changes come in as a merge commit. Simplest, shows where you integrated. The default and the right choice for most collaborative work.
Rebase (pull.rebase true): your commits are replayed on top of the remote ones, giving a linear history. Cleaner, but conflicts have to be resolved commit by commit.
Fast-forward only (pull.ff only): refuses to pull unless nothing diverged. Useful as a safety check, not as a daily setting.

We use merge. Because our folder and the GitHub repository started separately (one from git init, one from Create repository), git also needs --allow-unrelated-histories for this very first pull; never again after that.

PYTHON

!git config pull.rebase false   # pull strategy: merge (remote commits come in through a merge commit)
# pull = fetch origin's main and merge it into ours; the flag is needed only this once, as the two histories share no commit
!git pull origin main --allow-unrelated-histories --no-edit

If git reports a merge conflict, each conflicting file has to be resolved one way or the other and then staged and committed. You will not get one here (the README and your script are different files), but this is what it looks like:

PYTHON

# Example only.
# Keep YOUR version of a file:
# !git checkout --ours 04_Accessing-S3-via-SageMaker-notebooks.py
# Keep the REMOTE version of a file:
# !git checkout --theirs README.md
# Then stage and commit the resolution:
# !git add 04_Accessing-S3-via-SageMaker-notebooks.py README.md
# !git commit -m "Resolve merge conflicts"

Push with the PAT

The token goes into the URL for this one command, built from the variables getpass filled. It is not stored by git and does not appear in the notebook.

PYTHON

github_url = f"github.com/{username}/sagemaker-work.git"
# IPython substitutes the Python variables in {braces}; user:token@ in an HTTPS URL authenticates this one push and is not stored
!git push https://{username}:{token}@{github_url} main

Open https://github.com/YOUR_GITHUB_USERNAME/sagemaker-work in a browser: the README, the .py file and .gitignore should all be there, with your two commits in the history.

Step 7: Pulling `.py` files and converting back to notebooks

Suppose a collaborator edits the .py file on GitHub, or you edit it from another machine. Pull as usual:

PYTHON

!git config pull.rebase false
!git pull origin main   # a routine pull: no extra flag now that the histories are joined

Then convert the script back to a notebook with Jupytext. The .ipynb it writes is ignored by git (Step 5), so the round trip leaves the repository clean.

PYTHON

# --to notebook converts the other way; the .ipynb it writes matches the *.ipynb rule in .gitignore
!jupytext --to notebook 04_Accessing-S3-via-SageMaker-notebooks.py --output 04_Accessing-S3-via-SageMaker-notebooks.ipynb

To convert every .py in the folder back to notebooks, helpers.convert_files(direction="python_to_notebook") does the loop; remember it converts every .py in the current directory, so keep real modules out of that folder.

Challenge

Challenge: ignore the big files too

The training episodes save NumPy arrays (.npz) next to your notebooks. Those are binary and can be large, which makes them a poor fit for git. Add *.npz to your repository’s .gitignore, commit it, push with your PAT, and verify the change on GitHub.

Show me the solution

Still inside /home/ec2-user/SageMaker/sagemaker-work:

PYTHON

with open(".gitignore", "a") as gitignore:
    gitignore.write("\n# NumPy archives saved by the training episodes\n*.npz\n")
!cat .gitignore
!git add .gitignore
!git commit -m "Ignore .npz arrays"
!git push https://{username}:{token}@github.com/{username}/sagemaker-work.git main

Then open https://github.com/YOUR_GITHUB_USERNAME/sagemaker-work/blob/main/.gitignore in a browser: the *.npz line should be there, and the commit list should show “Ignore .npz arrays” as the latest commit. If the push is rejected with “authentication failed”, the PAT has expired or lacks the public_repo (or repo) scope; generate a new one and re-run the getpass cell.

Key Points

Your own work belongs in your own repository: create one on GitHub, git init a folder under /home/ec2-user/SageMaker/, and connect the two with git remote add.
Authenticate over HTTPS with a scoped, expiring Personal Access Token entered through getpass; nothing is saved on the instance, and the token is revoked when you are done.
Version .py copies of notebooks (Jupytext converts both ways) and ignore *.ipynb, caches and data archives in .gitignore.
Pull before you push; the first pull into a freshly initialised folder needs --allow-unrelated-histories.

Content from Training Models in SageMaker: Intro

Last updated on 2026-10-21 | Edit this page

Estimated time: 65 minutes

Overview

Questions

What are the differences between local training and SageMaker-managed training?
How do Estimator classes in SageMaker streamline the training process for various frameworks?
How do you estimate what a training job will cost before launching it, and check what it cost afterwards?

Objectives

Understand the difference between training locally in a SageMaker notebook and using SageMaker’s managed infrastructure.
Learn to configure and use SageMaker’s Estimator classes for different frameworks (e.g., XGBoost, PyTorch, SKLearn).
Tell apart the two ways to run XGBoost on SageMaker, your script in AWS’s container versus AWS’s built-in algorithm in the same container, and know when each fits.
Conduct training with data stored in S3 and monitor training job status using the SageMaker console.
Estimate the cost of a job before launching it, tag it so it shows up on the bill, and read back its billable time afterwards.

Initial setup

1. Open prefilled .ipynb notebook

Open the notebook from: /ML_with_AWS_SageMaker/notebooks/06_Training-models-in-SageMaker-notebooks.ipynb. When asked for a kernel, pick conda_pytorch; we will use the same kernel for the next two episodes. The conda_python3 kernel has no PyTorch, and on current instances no scikit-learn either.

2. Move to the project folder and import the helpers

Work from the Titanic project folder the S3 episode created, and make the lesson repository’s scripts/ importable. The training scripts for this episode (train_xgboost.py) and the next (train_nn.py) are in that folder.

PYTHON

# The Titanic project folder from the S3 episode, beside the lesson clone: downloads and outputs land here, not in the repository
!mkdir -p /home/ec2-user/SageMaker/titanic
%cd /home/ec2-user/SageMaker/titanic
import sys
# Put the lesson's scripts/ folder on Python's import path
sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")
import helpers  # lesson helper functions; train_xgboost.py lives in the same folder

3. Initialize SageMaker environment

This code initializes the AWS SageMaker environment by defining the SageMaker role and S3 client. It also specifies the S3 bucket and key for accessing the Titanic training dataset stored in an S3 bucket.

Boto3 API

Boto3 is the official AWS SDK for Python, allowing developers to interact programmatically with AWS services like S3, EC2, and Lambda. It provides both high-level and low-level APIs, making it easy to manage AWS resources and automate tasks. With built-in support for paginators, waiters, and session management, Boto3 simplifies working with AWS credentials, regions, and IAM permissions. It’s ideal for automating cloud operations and integrating AWS services into Python applications.

PYTHON

import boto3                          # AWS SDK for Python: low-level access to every AWS service (S3, EC2, ...)
import pandas as pd
import sagemaker                      # SageMaker Python SDK: higher-level API for training jobs, tuning, deployment
from sagemaker import get_execution_role

# Initialize the SageMaker role (will reflect notebook instance's policy)
role = sagemaker.get_execution_role()  # the IAM role attached to this notebook instance; every call below acts as this role
print(f'role = {role}')

# Initialize an S3 client to interact with Amazon S3, allowing operations like uploading, downloading, and managing objects and buckets.
s3 = boto3.client('s3')

# Define the S3 bucket that we will load from
bucket_name = 'sinkorswim-doejohn'  # replace with your S3 bucket name

# Define train/test filenames
train_filename = 'titanic_train.csv'  # object keys inside the bucket (no prefix: they sit at the bucket root)
test_filename = 'titanic_test.csv'

Create a SageMaker session to manage interactions with Amazon SageMaker, such as training jobs, model deployments, and data input/output. The region comes from the notebook instance itself, so the same notebook works in any account.

PYTHON

region = boto3.Session().region_name  # inherited from the notebook instance; do not hard-code a region
session = sagemaker.Session()         # tracks the region, default bucket and SageMaker client for this notebook
sagemaker_client = session.sagemaker_client  # the boto3 SageMaker client inside the session; we'll use it to describe jobs
print(f"region = {region}")

Testing train.py on this notebook’s instance

In this next section, we will learn how to take a model training script that was written/designed to run locally, and deploy it to more powerful instances (or many instances) using SageMaker. This is helpful for machine learning jobs that require extra power, GPUs, or benefit from parallelization. However, before we try exploiting this extra power, it is essential that we test our code thoroughly! We don’t want to waste unnecessary compute cycles and resources on jobs that produce bugs rather than insights.

General guidelines for testing ML pipelines before scaling

Run tests locally first (if feasible) to avoid unnecessary AWS charges. Here, we assume that local tests are not feasible due to limited local resources. Instead, we use our SageMaker instance to test our script on a minimally sized EC2 instance.
Use a small dataset subset (e.g., 1-5% of data) to catch issues early and speed up tests.
Start with a small/cheap instance before committing to larger resources. Visit the Instances for ML page for guidance.
Log everything to track training times, errors, and key metrics.
Verify correctness first before optimizing hyperparameters or scaling.

Discussion

What tests should we do before scaling?

Before scaling to multiple or more powerful instances (e.g., training on larger/multiple datasets in parallel or tuning hyperparameters in parallel), it’s important to run a few quick sanity checks to catch potential issues early. In your group, discuss:

Which checks do you think are most critical before scaling up?
What potential issues might we miss if we skip this step?

Show me the solution

Which checks do you think are most critical before scaling up?

Data loads correctly – Ensure the dataset loads without errors, expected columns exist, and missing values are handled properly.
Overfitting check – Train on a small dataset (e.g., 100 rows). If it doesn’t overfit, there may be a data or model setup issue.
Loss behavior check – Verify that training loss decreases over time and doesn’t diverge.
Training time estimate – Run on a small subset to estimate how long full training will take.
Memory estimate - Estimate the memory needs of the algorithm/model you’re using, and understand how this scales with input size.
Save & reload test – Ensure the trained model can be saved, reloaded, and used for inference without errors.

What potential issues might we miss if we skip the above checks?

Silent data issues – Missing values, unexpected distributions, or incorrect labels could degrade model performance.
Code bugs at scale – Small logic errors might not break on small tests but could fail with larger datasets.
Inefficient training runs – Without estimating runtime, jobs may take far longer than expected, wasting AWS resources.
Memory or compute failures – Large datasets might exceed instance memory limits, causing crashes or slowdowns.
Model performance issues – If a model doesn’t overfit a small dataset, there may be problems with features, training logic, or hyperparameters.

Callout

Know Your Data Before Modeling

The sanity checks above focus on validating the code, but a model is only as good as the data it’s trained on. A deeper look at feature distributions, correlations, and potential biases is critical before scaling up. We won’t cover that here, but it’s essential to keep in mind for any ML/AI practitioner.

Challenge

Understanding the XGBoost Training Script

Take a moment to review the ML_with_AWS_SageMaker/scripts/train_xgboost.py script we just cloned into our notebook. This script handles preprocessing, training, and saving an XGBoost model, while also adapting to both local and SageMaker-managed environments.

Try answering the following questions:

Data Preprocessing: What transformations are applied to the dataset before training?
Training Function: What does the train_model() function do? Why do we print the training time?
Command-Line Arguments: What is the purpose of argparse in this script? How would you modify the script if you wanted to change the number of training rounds?
Handling Local vs. SageMaker Runs: How does the script determine whether it is running in a SageMaker training job or locally (within this notebook’s instance)?
Training and Saving the Model: What format is the dataset converted to before training, and why? How is the trained model saved, and where will it be stored?

After reviewing, discuss any questions or observations with your group.

Show me the solution

Data Preprocessing: The script fills missing values (Age with median, Embarked with mode), converts categorical variables (Sex and Embarked) to numerical values, and removes columns that don’t contribute to prediction (Name, Ticket, Cabin).
Training Function: The train_model() function takes the training dataset (dtrain), applies XGBoost training with the specified hyperparameters, and prints the training time. Printing training time helps compare different runs and ensures that scaling decisions are based on performance metrics.
Command-Line Arguments: argparse allows passing parameters like max_depth, eta, num_round, etc., at runtime without modifying the script. To change the number of training rounds, you would update the --num_round argument when running the script: python train_xgboost.py --num_round 200
Handling Local vs. SageMaker Runs: The script uses os.environ.get("SM_CHANNEL_TRAIN", ".") and os.environ.get("SM_MODEL_DIR", ".") to detect whether it’s running in SageMaker. SM_CHANNEL_TRAIN is the directory where SageMaker stores input training data, and SM_MODEL_DIR is the directory where trained models should be saved. If these environment variables are not set (e.g., running locally), the script defaults to "." (current directory).
Training and Saving the Model: The dataset is converted into XGBoost’s DMatrix format, which is optimized for memory and computation efficiency. The trained model is saved with XGBoost’s own save_model() as xgboost-model.json, stored either in the SageMaker SM_MODEL_DIR (if running in SageMaker) or in the local directory. The JSON format can be loaded by any later XGBoost version, which matters here because the SageMaker container and the notebook kernel do not run the same version.

Download data into notebook environment

It can be convenient to have a copy of the data (i.e., one that you store in your notebook’s instance) to allow us to test our code before scaling things up.

Callout

While we demonstrate how to download data into the notebook environment for testing our code (previously setup for local ML pipelines), keep in mind that S3 is the preferred location for dataset storage in a scalable ML pipeline.

Run the next code chunk to download data from S3 to notebook environment. You may need to hit refresh on the file explorer panel to the left to see this file. If you get any permission issues…

check that you have selected the appropriate policy for this notebook
check that your bucket has the appropriate policy permissions

PYTHON

# Define the S3 bucket and file location
file_key = f"{train_filename}"  # Path to your file in the S3 bucket
local_file_path = f"./{train_filename}"  # Local path to save the file

# Download the file using the s3 client variable we initialized earlier
# This writes a copy to the instance's disk, unlike get_object, which reads the object into memory
s3.download_file(bucket_name, file_key, local_file_path)
print("File downloaded:", local_file_path)

We can do the same for the test set.

PYTHON

# Define the S3 bucket and file location
file_key = f"{test_filename}"  # Path to your file in the S3 bucket
local_file_path = f"./{test_filename}"  # Local path to save the file

# Download the file
s3.download_file(bucket_name, file_key, local_file_path)
print("File downloaded:", local_file_path)

Logging runtime & instance info

To compare our local runtime with future experiments, we’ll need to know what instance was used, as this will greatly impact runtime in many cases. We can extract the instance name for this notebook using…

PYTHON

# Replace with your notebook instance name.
# This does NOT refer to specific ipynb files, but to the SageMaker notebook instance.
notebook_instance_name = 'SinkOrSwim-DoeJohn'

# Describe the notebook instance (sagemaker_client was created from our session above)
response = sagemaker_client.describe_notebook_instance(NotebookInstanceName=notebook_instance_name)

# Display the status and instance type
print(f"Notebook Instance '{notebook_instance_name}' status: {response['NotebookInstanceStatus']}")
local_instance = response['InstanceType']  # e.g. ml.t3.medium; recorded so runtimes can be compared across instances
print(f"Instance Type: {local_instance}")

Helper: `get_notebook_instance_info()`

You can also use the get_notebook_instance_info() function found in the lesson’s scripts/helpers.py to retrieve this info for your own project.

PYTHON

helpers.get_notebook_instance_info(notebook_instance_name, region)

Test train.py on this notebook’s instance (or when possible, on your own machine) before doing anything more complicated (e.g., hyperparameter tuning on multiple instances)

PYTHON

!pip install xgboost scikit-learn   # neither notebook kernel ships scikit-learn; the script imports both

Local test

Run the training script here in the notebook first, on the small data, with the hyperparameters passed in as variables. A local run finds bugs for free, and the same script then runs unchanged as a SageMaker job.

PYTHON

import time as t # we'll use the time package to measure runtime

start_time = t.time()

# Define your parameters. These python vars wil be passed as input args to our train_xgboost.py script using %run

max_depth = 3 # Sets the maximum depth of each tree in the model to 3. Limiting tree depth helps control model complexity and can reduce overfitting, especially on small datasets.
eta = 0.1 #  Sets the learning rate to 0.1, which scales the contribution of each tree to the final model. A smaller learning rate often requires more rounds to converge but can lead to better performance.
subsample = 0.8 # Specifies that 80% of the training data will be randomly sampled to build each tree. Subsampling can help with model robustness by preventing overfitting and increasing variance.
colsample_bytree = 0.8 # Specifies that 80% of the features will be randomly sampled for each tree, enhancing the model's ability to generalize by reducing feature reliance.
num_round = 100 # Sets the number of boosting rounds (trees) to 100. More rounds typically allow for a more refined model, but too many rounds can lead to overfitting.
train_file = 'titanic_train.csv' #  Points to the location of the training data

# %run executes the script in this kernel; IPython substitutes the Python variables in {braces} into the arguments
%run /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/train_xgboost.py --max_depth {max_depth} --eta {eta} --subsample {subsample} --colsample_bytree {colsample_bytree} --num_round {num_round} --train {train_file}

# Measure and print the time taken
print(f"Total local runtime: {t.time() - start_time:.2f} seconds, instance_type = {local_instance}")

Training on this relatively small dataset should take less than a minute, but as we scale up with larger datasets and more complex models in SageMaker, tracking both training time and total runtime becomes essential for efficient debugging and resource management.

Note: Our code above includes print statements to monitor dataset size, training time, and total runtime, which provides insights into resource usage for model development. We recommend incorporating similar logging to track not only training time but also total runtime, which includes additional steps like data loading, evaluation, and saving results. Tracking both can help you pinpoint bottlenecks and optimize your workflow as projects grow in size and complexity, especially when scaling with SageMaker’s distributed resources.

Sanity check: Quick evaluation on test set

This next section isn’t SageMaker specific, but it does serve as a good sanity check to ensure our model is training properly. Here’s how you would apply the outputted model to your test set using your local notebook instance.

Score the model you just trained locally on the test set, reusing the script’s own preprocess_data so the test rows are transformed exactly as the training rows were.

PYTHON

import xgboost as xgb
import pandas as pd
import numpy as np
from sklearn.metrics import accuracy_score
from train_xgboost import preprocess_data  # importable because scripts/ is on sys.path

# Load the test data
test_data = pd.read_csv('./titanic_test.csv')

# Preprocess the test data using the imported preprocess_data function
X_test, y_test = preprocess_data(test_data)

# Convert the test features to DMatrix for XGBoost
dtest = xgb.DMatrix(X_test)

# Load the trained model from the saved file
model = xgb.Booster()
model.load_model('./xgboost-model.json')  # the file the local %run wrote to the current directory

# Make predictions on the test set
preds = model.predict(dtest)
predictions = np.round(preds)  # Round predictions to 0 or 1 for binary classification

# Calculate and print the accuracy of the model on the test data
accuracy = accuracy_score(y_test, predictions)
print(f"Test Set Accuracy: {accuracy:.4f}")

A reasonably high test set accuracy suggests our code/model is working correctly.

Training via SageMaker: your script inside AWS’s XGBoost container

Unlike “local” training (using this notebook), this next approach leverages SageMaker’s managed infrastructure to handle resources, parallelism, and scalability. By specifying instance parameters, such as instance_count and instance_type, you can control the resources allocated for training.

Which instance to start with?

In this example, we start with one ml.m5.large instance, which is suitable for small- to medium-sized datasets and simpler models. Using a single instance is often cost-effective and sufficient for initial testing, allowing for straightforward scaling up to more powerful instance types or multiple instances if training takes too long. See here for further guidance on selecting an appropriate instance for your data/model: Instances for ML page

Challenge

Challenge: pick an instance for a bigger job

Open the supplemental Instances for ML page. A colleague has a 20 GB tabular dataset (a CSV of sensor readings) and wants to train gradient-boosted trees on it, like we are about to do with the Titanic data. Which instance would you pick for the training job, and why?

Show me the solution

The table’s rule of thumb is RAM of at least 1.5× the dataset, so 20 GB of data wants about 32 GB of memory: ml.m5.2xlarge (8 vCPU, 32 GiB, about $0.46/h) is the natural fit for a tree-ensemble training job, or ml.m5.4xlarge (64 GiB) if feature engineering will add many columns. No GPU: gradient-boosted trees on tabular data are CPU work, so a g4dn or g5 would cost more and run no faster.

Two ways to run XGBoost on SageMaker, and why the names mislead

Every training job runs inside a container image, and for XGBoost that image comes from AWS either way. The question is whose training code runs in it:

	Your script	AWS’s built-in algorithm
Estimator	`XGBoost(entry_point="train_xgboost.py", ...)`	`Estimator(image_uri=<AWS XGBoost image>, ...)`, no `entry_point`
Who wrote the training code	you: the same `train_xgboost.py` you just ran with `%run`	AWS; nothing of yours is uploaded
What `hyperparameters` are	flags for your script’s `argparse`	XGBoost’s own parameter names
Data	whatever your script reads; it preprocesses the raw CSV itself	a headerless, all-numeric CSV with the label first, prepared beforehand
When it fits	preprocessing, custom logic, anything the algorithm alone cannot do	standard tabular XGBoost with nothing to write

The class names point the wrong way. XGBoost is a framework estimator: a convenience that looks up AWS’s XGBoost image and runs your script in it; PyTorch, SKLearn and HuggingFace are the same idea for other frameworks. The plain Estimator takes an image directly, and when that image is one of AWS’s built-in algorithms the image’s own code does the training. We do the script version first, because it is the one that generalises to every later episode, and the built-in algorithm at the end of this one.

Overview of Estimator classes in SageMaker

To launch this training “job”, we’ll use the XGBoost Estimator. In SageMaker, Estimator classes streamline the configuration and training of models on managed instances. Each Estimator can work with custom scripts and be enhanced with additional dependencies by specifying a requirements.txt file, which is automatically installed at the start of training. Here’s a breakdown of some commonly used Estimator classes in SageMaker:

1. `Estimator` (Base Class)

Purpose: Runs whatever image you name: one of AWS’s built-in algorithms, or a container of your own.
Configuration: Requires an image_uri; with a built-in algorithm there is no entry_point, the image’s own code trains.
Dependencies: You can use requirements.txt to install Python packages or configure a custom Docker container with pre-baked dependencies.
Ideal Use Cases: Custom algorithms or models that need tailored environments not covered by built-in containers.

2. `XGBoost` Estimator

Purpose: AWS’s XGBoost container running your script; the class saves you looking up the image.
Configuration:
- entry_point: Path to a custom script, useful for additional preprocessing or unique training workflows.
- framework_version: Select XGBoost version, e.g., "1.7-1".
- dependencies: Specify additional packages through requirements.txt to enhance preprocessing capabilities or incorporate auxiliary libraries.
Ideal Use Cases: Tabular data modeling using gradient-boosted trees; cases requiring custom preprocessing or tuning logic.

3. `PyTorch` Estimator

Purpose: Configures training jobs with PyTorch for deep learning tasks.
Configuration:
- entry_point: Training script with model architecture and training loop.
- instance_type: e.g., ml.g5.xlarge for GPU acceleration.
- framework_version and py_version: Define specific versions (e.g., "2.3" and "py311").
- dependencies: Install any required packages via requirements.txt to support advanced data processing, data augmentation, or custom layer implementations.
Ideal Use Cases: Deep learning models, particularly complex networks requiring GPUs and custom layers.

4. `SKLearn` Estimator

Purpose: Supports scikit-learn workflows for data preprocessing and classical machine learning.
Configuration:
- entry_point: Python script to handle feature engineering, preprocessing, or training.
- framework_version: Version of scikit-learn, e.g., "1.2-1".
- dependencies: Use requirements.txt to install any additional Python packages required by the training script.
Ideal Use Cases: Classical ML workflows, extensive preprocessing, or cases where additional libraries (e.g., pandas, numpy) are essential.

5. `TensorFlow` Estimator

Purpose: Designed for training and deploying TensorFlow models.
Configuration:
- entry_point: Script for model definition and training process.
- instance_type: Select based on dataset size and computational needs.
- dependencies: Additional dependencies can be listed in requirements.txt to install TensorFlow add-ons, custom layers, or preprocessing libraries.
Ideal Use Cases: NLP, computer vision, and transfer learning applications in TensorFlow.

6. `HuggingFace` Estimator

Purpose: Provides managed containers for training and fine-tuning models using the Hugging Face transformers library.
Configuration:
- entry_point: Custom script for training (e.g., train.py).
- transformers_version, pytorch_version, py_version: Define framework versions.
- dependencies: Optional requirements.txt for extra libraries.
Ideal Use Cases: NLP, vision, or multimodal tasks using pretrained Transformer models.

Callout

Configuring custom environments with `requirements.txt`

For all these Estimators, adding a requirements.txt file as a dependencies argument ensures that additional packages are installed before training begins. This approach allows the use of specific libraries that may be critical for custom preprocessing, feature engineering, or model modifications. Here’s how to include it (illustration only; we don’t run this):

from sagemaker.sklearn.estimator import SKLearn
sklearn_estimator = SKLearn(
    base_job_name=f"{notebook_instance_name}-xgboost",   # job name = who + which model; SageMaker appends a timestamp
    entry_point="train_script.py",
    role=role,
    instance_count=1,
    instance_type="ml.m5.large",
    max_run=max_run,
    tags=job_tags,
    output_path=f"s3://{bucket_name}/output",
    framework_version="1.2-1",
    dependencies=['requirements.txt'],  # Adding custom dependencies
    hyperparameters={
        "max_depth": 5,
        "eta": 0.1,
        "subsample": 0.8,
        "num_round": 100
    }
)

This setup simplifies training, allowing you to maintain custom environments directly within SageMaker’s managed containers, without needing to build and manage your own Docker images. The AWS SageMaker Documentation provides lists of pre-built container images for each framework and their standard libraries, including details on pre-installed packages.

Deploying to other instances

For this deployment, we configure the XGBoost estimator with a custom training script, train_xgboost.py, and define hyperparameters directly within the SageMaker setup. Here’s the full code, with some additional explanation following the code.

Cost tracking

When you launch a SageMaker training job from a notebook, SageMaker creates new managed resources (EC2 instances, attached storage, logs) on your behalf. These resources do not automatically inherit the notebook instance’s tags.

To avoid this, we explicitly tag each training job at launch time. This ensures that compute usage is traceable to a project, a purpose, and a human-readable name, even after the job has completed.

PYTHON

name = "DoeJohn" # replace with your name
project = "SinkOrSwim" # replace with your challenge name
purpose = "Training"        # one of: Storage, Notebook, Training, Tuning, Inference

# Tags: key/value labels attached to the job's resources, so its cost shows up under your name and project on the bill
job_tags = [
    {"Key": "Name", "Value": name},
    {"Key": "Project", "Value": project},
    {"Key": "Purpose", "Value": purpose},
]
print(f"Name: {name}, Project: {project}, Purpose: {purpose}")

Estimate before you launch

Estimate the cost of every job before you run it: instance rate x expected minutes, plus two or three minutes of startup. Two things happen before your script runs, and only one of them is billed. First SageMaker finds an instance for you: that is provisioning, it takes a minute or so, and it is free. Then the instance pulls the training container and downloads your data: that is startup, and the meter is running. In this lesson, startup always means the billed part. An ml.m5.large costs about $0.115/hour (us-east-1 on-demand; see the Instances for ML page). Our local test finished in under a minute, so the estimate is roughly 1 minute of training + 2-3 minutes of startup = 3-4 billable minutes, or under a cent.

One thing to watch in the estimator below is hyperparameters. It is how values get from this notebook into the script. SageMaker turns the dictionary into command-line flags, --train titanic_train.csv --max_depth 3 --eta 0.1 ..., and appends them to the command that starts train_xgboost.py inside the container, exactly as %run passed them a moment ago; the script’s argparse block reads them. So the keys must be the script’s argument names, and anything the script does not define a flag for is ignored. (The job’s log prints the full command, under “Invoking script with the following command”, which is the place to look when a value does not seem to arrive.)

PYTHON

from sagemaker.inputs import TrainingInput
from sagemaker.xgboost.estimator import XGBoost   # framework estimator: AWS's XGBoost container running your own script

# Define instance type/count we'll use for training
instance_type="ml.m5.large"   # the managed machine SageMaker launches for the job, separate from this notebook's instance
instance_count=1 # always start with 1 (see "When training takes too long" at the end of this episode)
price_per_hour = 0.115 # ml.m5.large, us-east-1 on-demand (2026-10-16); check the instances page for your region

# Define max runtime in seconds to ensure you don't use more compute time than expected. Use a generous threshold (2x expected time but < 2 days) so that work isn't interrupted without any gains.
max_run = 2*60*60 # 2 hours; SageMaker stops the job and releases the instance when this is reached

# Define S3 paths for input and output
train_s3_path = f's3://{bucket_name}/{train_filename}'  # an S3 URI: bucket plus key; SageMaker takes input in this form

# we'll store all results in a subfolder called xgboost on our bucket. This folder will automatically be created if it doesn't exist already.
output_folder = 'xgboost'
output_path = f's3://{bucket_name}/{output_folder}/'   # S3 prefix (a "folder") where SageMaker uploads the job's outputs

# The XGBoost framework estimator: AWS's XGBoost container, running our script
xgboost_estimator = XGBoost(
    base_job_name=f"{notebook_instance_name}-xgboost",   # prefix of the job name: who + which model; SageMaker appends a timestamp to make it unique
    max_run=max_run, # in seconds; always include (max 48 hours)
    entry_point='train_xgboost.py',      # the script SageMaker runs inside the container on the training instance
    source_dir='/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts',  # Directory where your script is located (uploaded to the job with it)
    role=role,                           # the job assumes this role to read its input and write its output in S3
    tags=job_tags,                       # the cost-tracking tags defined above
    instance_count=instance_count,
    instance_type=instance_type,
    output_path=output_path,
    sagemaker_session=session,
    framework_version="1.7-1",           # XGBoost container version
    hyperparameters={                    # passed to the script as command-line arguments: --train ..., --max_depth ...
        'train': train_file,
        'max_depth': max_depth,
        'eta': eta,
        'subsample': subsample,
        'colsample_bytree': colsample_bytree,
        'num_round': num_round
    }
)

# Define input data
# TrainingInput: where the data is and its format; SageMaker copies it onto the training instance before the script runs
train_input = TrainingInput(train_s3_path, content_type='csv')

# Measure and start training time
start = t.time()
# fit() launches the job and blocks until it finishes (wait=False would return at once); the 'train' channel becomes SM_CHANNEL_TRAIN
xgboost_estimator.fit({'train': train_input})
end = t.time()

print(f"Runtime for training on SageMaker: {end - start:.2f} seconds, instance_type: {instance_type}, instance_count: {instance_count}")

When running longer training jobs, you can check on their status periodically from the AWS SageMaker Console (where we originally launched our Notebook instance) on left side panel under “Training”.

Hyperparameters

The hyperparameters section in this code defines the input arguments of train_xgboost.py. The first is the name of the training input file, and the others are hyperparameters for the XGBoost model, such as max_depth, eta, subsample, colsample_bytree, and num_round.

TrainingInput

Additionally, we define a TrainingInput object containing the training data’s S3 path, to pass to .fit({'train': train_input}). SageMaker uses TrainingInput to download your dataset from S3 to a temporary location on the training instance. This location is mounted and managed by SageMaker and can be accessed by the training job if/when needed.

What did that cost?

Timeline of a SageMaker training job showing the unbilled wait for an instance and the billed phases from container pull to model upload describe_training_job returns BillableTimeInSeconds: the seconds the instance was yours, from container pull to model upload, per instance. It also returns TrainingTimeInSeconds, which on on-demand instances is the same number (it only differs on spot instances), so do not expect it to isolate your script. The minute or so before that, while SageMaker finds an instance, is not billed. SageMaker bills per second, with no rounding, so cost is billable seconds x instance count x hourly rate. We’ll wrap this in a small helper and call it after every job in this lesson.

PYTHON

def job_cost_report(job_name: str, price_per_hour: float):
    desc = session.sagemaker_client.describe_training_job(TrainingJobName=job_name)  # the job's full record: status, timings, resources, output
    n = desc["ResourceConfig"]["InstanceCount"]
    billable = desc["BillableTimeInSeconds"] * n   # BillableTimeInSeconds: container pull to model upload, charged per instance
    wall = (desc["TrainingEndTime"] - desc["CreationTime"]).total_seconds()   # from the launch request to the end, including unbilled provisioning
    print(f"{job_name}: status {desc['TrainingJobStatus']}, wall-clock {wall/60:.1f} min, billable {billable/60:.1f} min on {n} instance(s), cost ${billable/3600*price_per_hour:.3f} at ${price_per_hour}/h")
    print(f"    model: {desc['ModelArtifacts']['S3ModelArtifacts']}")   # the model.tar.gz this job uploaded to output_path
    return desc

desc = job_cost_report(xgboost_estimator.latest_training_job.name, price_per_hour)  # latest_training_job: the job fit() just ran

The model: line is where the trained model went. Everything else the job wrote sits next to it under the same prefix; list it straight from the notebook:

PYTHON

# Everything the job wrote, straight from S3: the model tarball plus the profiler's small files
job_prefix = f"{output_folder}/{xgboost_estimator.latest_training_job.name}/"
for obj in s3.list_objects_v2(Bucket=bucket_name, Prefix=job_prefix).get("Contents", []):
    print(f"{obj['Size'] / 1024:8.1f} KB  s3://{bucket_name}/{obj['Key']}")

Compare this with the estimate you made above. The job’s log (scroll up in the fit() output) shows the script itself trained in a fraction of a second, yet the bill is about two minutes. The rest is container pull, data download and model upload: the fixed overhead of every job, the same whether you train for one minute or one hour. The extra minute between calling fit() and the first billed second is SageMaker finding an instance, which is free.

What landed in S3

The listing above is what you would see in the S3 console; open it once to see the same thing the console’s way.

In your bucket, under xgboost/<job-name>/ (the prefix from output_path plus the job name):

output/model.tar.gz: the model artifacts, a tarball of whatever the script saved to SM_MODEL_DIR (here xgboost-model.json). This is the file we download below.
profiler-output/ and debug-output/: small files from SageMaker’s built-in profiler, written for every training job. Ignore them; they can be turned off with disable_profiler=True on the estimator.

The job’s log is not in S3. It is in CloudWatch, the AWS service that stores the printed output of every job: everything the script prints goes to a CloudWatch log stream named after the job, and it stays there after the instance is gone. To read it: open Training jobs, click the job’s name, scroll to Monitor and click View logs, then click the one stream listed. What you see is the same text fit() streamed into the notebook, which is useful the day a job you launched last night failed while the kernel was off.

A second bucket also appeared, sagemaker-<region>-<account-id>. The SageMaker SDK creates this default bucket the first time it needs somewhere to stage files: before the job started, fit() zipped source_dir and uploaded it there as <job-name>/source/sourcedir.tar.gz for the training instance to download. It holds kilobytes per job and nothing you need, but it is shared by everyone in the account and carries no tags, so it is one more thing on the cleanup list. If you want everything in your own bucket, create the session as sagemaker.Session(default_bucket=bucket_name).

Extracting trained model from S3 for final evaluation

To evaluate the model on a test set after training, we’ll go through these steps:

Download the trained model from S3.
Load and preprocess the test dataset.
Evaluate the model on the test data.

Here’s how you can implement this in your SageMaker notebook. The following code will:

Download the model.tar.gz file containing the trained model from S3.
Load the titanic_test.csv data we downloaded earlier and preprocess it as needed.
Use the XGBoost model to make predictions on the test set and then compute accuracy or other metrics on the results.

If additional metrics or custom evaluation steps are needed, you can add them in place of or alongside accuracy.

PYTHON

# Model results are saved in auto-generated folders. Use xgboost_estimator.latest_training_job.name to get the folder name
# Key of the model artifact: <output prefix>/<job name>/output/model.tar.gz, a tarball of whatever the script saved to SM_MODEL_DIR
model_s3_path = f'{output_folder}/{xgboost_estimator.latest_training_job.name}/output/model.tar.gz'
print(model_s3_path)
local_model_path = 'model.tar.gz'

# Download the trained model from S3
s3.download_file(bucket_name, model_s3_path, local_model_path)

# Extract the model file
import tarfile
with tarfile.open(local_model_path) as tar:
    tar.extractall()

Now the same evaluation on the model the job trained: load it from the unpacked archive and score the test set the same way as before.

PYTHON

import xgboost as xgb
import pandas as pd
import numpy as np
from sklearn.metrics import accuracy_score
from train_xgboost import preprocess_data

# Load the test data
test_data = pd.read_csv('./titanic_test.csv')

# Preprocess the test data using the imported preprocess_data function
X_test, y_test = preprocess_data(test_data)

# Convert the test features to DMatrix for XGBoost
dtest = xgb.DMatrix(X_test)

# Load the trained model from the saved file (the container's XGBoost 1.7 wrote it; the kernel's newer XGBoost reads it)
model = xgb.Booster()
model.load_model('./xgboost-model.json')

# Make predictions on the test set
preds = model.predict(dtest)
predictions = np.round(preds)  # Round predictions to 0 or 1 for binary classification

# Calculate and print the accuracy of the model on the test data
accuracy = accuracy_score(y_test, predictions)
print(f"Test Set Accuracy: {accuracy:.4f}")

Now the other route from the table at the top of this section: the same AWS XGBoost container, but with AWS’s own training code instead of our script. Nothing of ours is uploaded; we hand the algorithm data in the shape it expects and XGBoost’s own parameters. Below, we’ll discuss both the code and pros and cons of the image-based setup compared to the custom script approach.

The built-in algorithm: AWS’s training code in the same container

When the job is standard XGBoost on a table, there is nothing to write: the built-in algorithm trains, evaluates and saves the model itself. The price is that it does only that. Preprocessing has to happen before the data reaches it, and the data has to arrive in its format.

Comparison: your script vs. the built-in algorithm

Feature	Your script (`XGBoost` with `entry_point`)	Built-in algorithm (`Estimator` with `image_uri`, no script)
Flexibility	Allows for custom preprocessing, data transformation, or advanced parameterization. Requires a Python script and custom dependencies can be added through `requirements.txt`.	Limited to XGBoost’s built-in functionality, no custom preprocessing steps without additional customization.
Simplicity	Requires setting up a script with `entry_point` and managing dependencies. Ideal for specific needs but requires configuration.	Streamlined for fast deployment without custom code. Simple setup and no need for custom scripts.
Performance	Similar performance, though potential for overhead with additional preprocessing.	Optimized for typical XGBoost tasks with faster startup. May offer marginally faster time-to-first-train.
Use Cases	Ideal for complex workflows requiring unique preprocessing steps or when adding specific libraries or functionalities.	Best for quick experiments, standard workflows, or initial testing on datasets without complex preprocessing.

When to use each approach:

Your script: whenever the job needs preprocessing, feature engineering, or any logic of your own around the training call. This is the route every later episode uses.
Built-in algorithm: standard XGBoost on a table that is already in its format, with nothing to write.

Below is the same Titanic job on the built-in algorithm.

Preparing data for the built-in algorithm

The built-in algorithm will not run our preprocess_data() for us, and it expects a CSV with no header row, every column numeric, and the label in the first column. The raw titanic_train.csv (header, Name and Sex strings, label in the second column) fails at data load. So we run the same preprocessing we used in train_xgboost.py here in the notebook, write the result in the expected layout, and upload it to its own prefix in our bucket.

PYTHON

from train_xgboost import preprocess_data

# Same preprocessing our script applies: fill missing values, encode Sex/Embarked, drop text columns
X, y = preprocess_data(pd.read_csv(train_filename))

# Built-in XGBoost wants: no header, all numeric, label first
builtin_filename = "train_builtin.csv"
pd.concat([y, X], axis=1).to_csv(builtin_filename, header=False, index=False)

# Upload to a dedicated prefix so it doesn't get confused with the raw CSV
builtin_s3_key = f"xgboost/builtin/{builtin_filename}"
s3.upload_file(builtin_filename, bucket_name, builtin_s3_key)
builtin_s3_path = f"s3://{bucket_name}/{builtin_s3_key}"
print("Uploaded:", builtin_s3_path)

Setting up the data path

With our script, the data path was just another flag the script read. The built-in algorithm has no flags: the data arrives as a TrainingInput channel it knows to look for. Here there is no script of ours in the container, so the dictionary goes to AWS’s own training code, which reads it as XGBoost’s parameters: the keys must be XGBoost’s parameter names (max_depth, eta, num_round), not flags we invented, and there is no train key because the data arrives as a channel. We also add objective, because the built-in algorithm’s default objective is regression rather than binary classification.

PYTHON

from sagemaker.estimator import Estimator # when using images, we use the general Estimator class: a container image, no script of ours

# Define instance type/count we'll use for training
instance_type="ml.m5.large"
instance_count=1 # always start with 1

# Use Estimator directly for built-in container without specifying entry_point
xgboost_estimator_builtin = Estimator(
    base_job_name=f"{notebook_instance_name}-xgboost",   # job name = who + which model; SageMaker appends a timestamp
    max_run=max_run, # in seconds; always include (max 48 hours)
    # image_uris.retrieve looks up AWS's prebuilt XGBoost container for this region; its built-in training code replaces entry_point
    image_uri=sagemaker.image_uris.retrieve("xgboost", region, version="1.7-1"),
    role=role,
    tags=job_tags,
    instance_count=instance_count,
    instance_type=instance_type,
    output_path=output_path,
    sagemaker_session=session,
    hyperparameters={                    # go straight to the XGBoost algorithm, not to a script of ours
        'objective': 'binary:logistic',  # the built-in default is regression; we are classifying
        'max_depth': max_depth,
        'eta': eta,
        'subsample': subsample,
        'colsample_bytree': colsample_bytree,
        'num_round': num_round
    }
)

# Define input data: the headerless, label-first CSV we just uploaded
train_input_builtin = TrainingInput(builtin_s3_path, content_type="text/csv")  # content_type tells the built-in algorithm how to parse the data

# Measure and start training time
start = t.time()
xgboost_estimator_builtin.fit({'train': train_input_builtin})
end = t.time()

print(f"Runtime for training on SageMaker: {end - start:.2f} seconds, instance_type: {instance_type}, instance_count: {instance_count}")

Then the cost report:

PYTHON

desc = job_cost_report(xgboost_estimator_builtin.latest_training_job.name, price_per_hour)

Monitoring training

To view and monitor your SageMaker training job, follow these steps in the AWS Management Console. Since training jobs may be visible to multiple users in your account, it’s essential to confirm that you’re interacting with your own job before making any changes.

Navigate to the SageMaker Console
- Go to the AWS Management Console and open the SageMaker service (can search for it)
View training jobs
- In the left-hand navigation menu, select Training jobs (or open this link in the class region). You’ll see a list of recent training jobs, which may include jobs from other users in the account.
Verify your training Job
- Identify your job by looking for the specific name format (e.g., SinkOrSwim-DoeJohn-YYYY-MM-DD-HH-MM-SS-XXX) generated when you launched the job. Click on its name to access detailed information. Cross-check the job details, such as the Instance Type and Input data configuration, with the parameters you set in your script.
Monitor the job status
- Once you’ve verified the correct job, click on its name to access detailed information:
  - Status: Confirms whether the job is InProgress, Completed, or Failed.
  - Logs: the View logs link opens the job’s output in CloudWatch, the AWS logging service, where it is kept after the job ends; Job Metrics shows CPU, GPU and memory use over the run.
  - Output Data: Shows the S3 location with the trained model artifacts.
Stopping a training job
- Before stopping a job, ensure you’ve selected the correct one by verifying job details as outlined above.
- If you’re certain it’s your job, go to Training jobs in the SageMaker Console, select the job, and choose Stop from the Actions menu. Confirm your selection, as this action will halt the job and release any associated resources.
- Important: Avoid stopping jobs you don’t own, as this could disrupt other users’ work and may have unintended consequences.

Following these steps helps ensure you only interact with and modify jobs you own, reducing the risk of impacting other users’ training processes.

When training takes too long

When training time becomes excessive, two main options can improve efficiency in SageMaker:

Option 1: Upgrading to a more powerful instance
Option 2: Using multiple instances for distributed training

Generally, Option 1 is the preferred approach and should be explored first.

Option 1: Upgrade to a more powerful instance (preferred starting point)

Upgrading to a more capable instance, particularly one with GPU capabilities, is often the simplest and most cost-effective way to speed up training. Check the Instances for ML page for guidance.

When to use a single instance upgrade:
- Dataset size – The dataset is small to moderate (e.g., <10 GB), fitting comfortably within memory.
- Model complexity – XGBoost models are typically small enough to fit in memory.
- Training time – If training completes in a few hours but could be faster, upgrading may help.

Upgrading a single instance is usually the most efficient option. It avoids the communication overhead of multi-instance setups and works well for small to medium datasets.

Option 2: Use multiple instances for distributed training

If upgrading a single instance doesn’t sufficiently reduce training time, distributed training across multiple instances may be a viable alternative. This only helps when your script and data are set up to split the work.

When to consider multiple instances

Using multiple instances makes sense when:
- Dataset size – The dataset is large and doesn’t fit comfortably in memory.
- Expected training time – A single instance takes too long (e.g., >10 hours).
- Need for faster training – Parallelization can speed up training but introduces communication overhead.

If scaling to multiple instances, monitoring training time and efficiency is critical. In many cases, a single, more powerful instance may be more cost-effective than multiple smaller ones.

In the next episode we train a PyTorch network, try Option 1 (a GPU instance), and then run a two-instance job to see what instance_count=2 does and does not do.

Key Points

Environment initialization: Setting up a SageMaker session, defining roles, and specifying the S3 bucket are essential for managing data and running jobs in SageMaker.
Local vs. managed training: Always test your code locally (on a smaller scale) before scaling things up. This avoids wasting resources on buggy code that doesn’t produce reliable results.
Estimator classes: SageMaker provides framework-specific Estimator classes (e.g., XGBoost, PyTorch, SKLearn) to streamline training setups, each suited to different model types and workflows.
Your script vs. the built-in algorithm: both run in AWS’s XGBoost container. XGBoost(entry_point=...) runs your script in it and is the pattern every later episode uses; Estimator(image_uri=...) with no script runs AWS’s own training code, which needs the data pre-shaped and takes XGBoost’s own parameters.
Training data channels: Using TrainingInput ensures SageMaker manages data efficiently.
Cost habit: Estimate cost before launching (rate x minutes + ~3 minutes startup), then read back BillableTimeInSeconds with job_cost_report() after every job and compare.
Scaling: Upgrade the single instance first; consider multiple instances only when data no longer fits in memory or a run takes more than ~10 hours.

Content from Training Models in SageMaker: PyTorch Example

Last updated on 2026-10-21 | Edit this page

Estimated time: 50 minutes

Overview

Questions

When should you consider using a GPU instance for training neural networks in SageMaker, and what are the benefits and limitations?
How does SageMaker handle data parallelism and model parallelism, and which is suitable for typical neural network training?

Objectives

Preprocess the Titanic dataset for efficient training using PyTorch.
Save and upload training and validation data in .npz format to S3.
Understand the trade-offs between CPU and GPU training for smaller datasets.
Train a PyTorch model as a SageMaker training job and compare CPU and GPU instance types on time and cost.
Differentiate between data parallelism and model parallelism, and determine when each is appropriate in SageMaker.

Initial setup: open prefilled .ipynb notebook

Open the notebook from: /ML_with_AWS_SageMaker/notebooks/07_Training-models-in-SageMaker-notebooks-part2.ipynb. Select the conda_pytorch kernel.

Setup notebook as controller

Once your notebook is open, we can setup our SageMaker controller as usual:

PYTHON

import boto3
import pandas as pd
import sagemaker
from sagemaker import get_execution_role

# Initialize the SageMaker role (will reflect notebook instance's policy)
role = sagemaker.get_execution_role()
print(f'role = {role}')

# Initialize an S3 client to interact with Amazon S3, allowing operations like uploading, downloading, and managing objects and buckets.
s3 = boto3.client('s3')

# Define the S3 bucket that we will load from
bucket_name = 'sinkorswim-doejohn'  # replace with your S3 bucket name

# Define train/test filenames
train_filename = 'titanic_train.csv'
test_filename = 'titanic_test.csv'

Create a SageMaker session to manage interactions with Amazon SageMaker, such as training jobs, model deployments, and data input/output. As before, the region comes from the notebook instance.

PYTHON

region = boto3.Session().region_name  # inherited from the notebook instance; do not hard-code a region
session = sagemaker.Session()
sagemaker_client = session.sagemaker_client  # low-level boto3 client; describe_training_job() below uses it
print(f"region = {region}")

We should also record our local instance information to report this information during testing. First, let’s make sure we’re starting in the same location to access helper functions

PYTHON

# The Titanic project folder from the S3 episode, beside the lesson clone: downloads and outputs land here, not in the repository
!mkdir -p /home/ec2-user/SageMaker/titanic
%cd /home/ec2-user/SageMaker/titanic
import sys
sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")
import helpers

PYTHON

notebook_instance_name = 'SinkOrSwim-DoeJohn' # replace with your notebook instance name

local_instance_info = helpers.get_notebook_instance_info(notebook_instance_name, region)
local_instance = local_instance_info['InstanceType']
local_instance

If you are starting from a fresh notebook instance, download the Titanic CSVs from S3 again.

PYTHON

s3.download_file(bucket_name, train_filename, train_filename)
s3.download_file(bucket_name, test_filename, test_filename)

Training a neural network with SageMaker

Let’s see how to do a similar experiment, but this time using PyTorch neural networks. We will again demonstrate how to test our custom model train script (train_nn.py) before deploying to SageMaker, and discuss some strategies (e.g., using a GPU) for improving train time when needed.

Preparing the data (compressed npz files)

When deploying a PyTorch model on SageMaker, it’s helpful to prepare the input data in a format that’s directly accessible and compatible with PyTorch’s data handling methods. The next code cell will prep our npz files from the existing csv versions.

Callout

Why are we using this file format?

Optimized data loading:
The .npz format stores arrays in a compressed, binary format, making it efficient for both storage and loading. PyTorch can easily handle .npz files, especially in batch processing, without requiring complex data transformations during training.
Batch compatibility:
When training neural networks in PyTorch, it’s common to load data in batches. By storing data in an .npz file, we can quickly load the entire dataset or specific parts (e.g., X_train, y_train) into memory and feed it to the PyTorch DataLoader, enabling efficient batched data loading.
Reduced I/O overhead in SageMaker:
Storing data in .npz files minimizes the I/O operations during training, reducing time spent on data handling. This is especially beneficial in cloud environments like SageMaker, where efficient data handling directly impacts training costs and performance.
Consistency and compatibility:
Using .npz files allows us to ensure consistency between training and validation datasets. Each file (train_data.npz and val_data.npz) stores the arrays in a standardized way that can be easily accessed by keys (X_train, y_train, X_val, y_val). This structure is compatible with PyTorch’s Dataset class, making it straightforward to design custom datasets for training.
Support for multiple data types:
.npz files support storage of multiple arrays within a single file. This is helpful for organizing features and labels without additional code. Here, the train_data.npz file contains both X_train and y_train, keeping everything related to training data in one place. Similarly, val_data.npz organizes validation features and labels, simplifying file management.

In summary, saving the data in .npz files ensures a smooth workflow from data loading to model training in PyTorch, leveraging SageMaker’s infrastructure for a more efficient, structured training process.

Open ML_with_AWS_SageMaker/scripts/train_nn.py in the file browser and read it before going on; it is short. Three things to find: TitanicNet, a three-layer network with 7 inputs and one output; the training loop in train(), which prints validation:accuracy = ... once per hundred epochs (the hyperparameter-tuning episode’s metric regex reads that line); and the last lines, which save the model to /opt/ml/model inside a job and next to the script otherwise, the same trick as train_xgboost.py. Questions now are cheaper than questions once it is running on a billed instance.

We reuse the preprocess_data() function from the lesson’s scripts/train_nn.py rather than re-typing the preprocessing here: the same function will later be applied to the test set, so train and test are transformed identically. It returns the fitted StandardScaler, and we save it with joblib so the next episode can load it.

PYTHON

!pip install scikit-learn   # the notebook kernels do not ship it; train_nn.py imports it

Load the Titanic CSV, run it through the script’s preprocess_data, split off a validation set, and save the arrays and the scaler. The training script reads these files; the scaler is kept so test data can later be transformed the same way.

PYTHON

import numpy as np
import joblib
from sklearn.model_selection import train_test_split
from train_nn import preprocess_data

# Load and preprocess the Titanic dataset: encode Sex/Embarked, fill missing Age/Fare, standardize features
df = pd.read_csv(train_filename)
X, y, scaler = preprocess_data(df)   # scaler is fit on the training data only

# Split the data
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

# Save the preprocessed data (and the fitted scaler) to our local jupyter environment
np.savez('train_data.npz', X_train=X_train, y_train=y_train)  # .npz = several numpy arrays in one compressed file, read back by key
np.savez('val_data.npz', X_val=X_val, y_val=y_val)
joblib.dump(scaler, 'scaler.joblib')  # keep the fitted scaler: the test set must be transformed with these same means/stds, never refit
print(X_train.shape, X_val.shape)

Next, we will upload our compressed files (and the scaler) to our S3 bucket. Storage is fairly cheap on AWS (around $0.023 per GB per month), but be mindful of uploading too much data. It may be convenient to store a preprocessed version of the data at times, but try not to store too many versions that aren’t being actively used.

PYTHON

train_file = "train_data.npz"  # Local file path in your notebook environment
val_file = "val_data.npz"  # Local file path in your notebook environment
scaler_file = "scaler.joblib"

# Upload the training and validation files (and the scaler) to the bucket root
for f in [train_file, val_file, scaler_file]:
    s3.upload_file(f, bucket_name, f)

print("Files successfully uploaded to S3.")

Testing on notebook instance

You should always test code thoroughly before scaling up and using more resources. Here, we will test our script using a small number of epochs (100) — just to verify our setup is correct.

PYTHON

import torch
import time as t # Measure training time locally

epochs = 100          # small smoke test
learning_rate = 0.001

start_time = t.time()
%run /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/train_nn.py --train train_data.npz --val val_data.npz --epochs {epochs} --learning_rate {learning_rate}
print(f"Local training time: {t.time() - start_time:.2f} seconds, instance_type = {local_instance}")

Deploying PyTorch neural network via SageMaker

Now that we have tested things locally, we can try to train with a larger number of epochs (1000). We can do this easily by invoking the PyTorch estimator. Our notebook is currently configured to use ml.m5.large. We will use the same instance type with the below code (using our notebook as a controller).

Should we use a GPU?: Since this dataset is fairly small, we don’t necessarily need a GPU for training. Considering costs, the m5.large is $0.115/hour, while the cheapest GPU instance (ml.g4dn.xlarge) is $0.74/hour. However, for larger datasets (> 1 GB) and models, we may want to consider a GPU if training time becomes cumbersome (see Instances for ML). We’ll test that in the exercise below.

Cost tracking

As in the XGBoost episode, every training job needs its own tags: a job does not inherit the notebook instance’s tags, and an untagged job is anonymous on the bill. Define them once here and pass tags=job_tags to every estimator below.

PYTHON

name = "DoeJohn" # replace with your name
project = "SinkOrSwim" # replace with your challenge name
purpose = "Training"        # one of: Storage, Notebook, Training, Tuning, Inference

job_tags = [
    {"Key": "Name", "Value": name},
    {"Key": "Project", "Value": project},
    {"Key": "Purpose", "Value": purpose},
]
print(f"Name: {name}, Project: {project}, Purpose: {purpose}")

Estimating a training job before you launch it

Every job in this lesson gets an estimate written down before fit(), and the habit matters more than the numbers here, which are cents. There are two ways to make the estimate, and which one applies depends on the size of the model.

Small models: time a smoke test and scale it. For a network like TitanicNet, the arithmetic is nothing; the time is Python overhead per epoch plus the job’s fixed startup. So measure the 100-epoch run you just did, scale to the epochs and data you plan, and add the startup:

PYTHON

local_seconds_per_epoch = (t.time() - start_time) / epochs     # from the smoke test above: wall-clock over its 100 epochs
planned_epochs = 1000
startup_minutes = 3                                            # container pull + data download, billed; the same for every job
rate = 0.115                                                   # $/hour, ml.m5.large

train_minutes = local_seconds_per_epoch * planned_epochs / 60
estimate = (train_minutes + startup_minutes) / 60 * rate
print(f"{local_seconds_per_epoch * 1000:.1f} ms per epoch locally -> {train_minutes:.1f} min of training for {planned_epochs} epochs")
print(f"estimate: {train_minutes + startup_minutes:.0f} min billed, about ${estimate:.3f} on ml.m5.large")

Expect a couple of minutes of training plus two or three of startup: about a cent. (The ml.t3.medium notebook is slower than the ml.m5.large the job runs on, so the training part of the estimate is high. Our run billed 2.8 minutes in total, $0.005.) If the planned run has ten times the rows, multiply the training minutes by ten; the startup does not change.

Large models: count the arithmetic. Once a model has millions of parameters, compute is the cost, and it can be counted before any instance exists. Training one example through a dense network or a transformer costs about 6 floating-point operations per parameter (2 forward, 4 backward); multiply by parameters, examples and epochs, and divide by what the GPU does per second, about a quarter of its data-sheet peak in practice:

FLOPs       = 6 × parameters × examples × tokens per example × epochs
GPU-seconds = FLOPs ÷ (GPU peak FLOPS × 0.25)
cost        = GPU-hours × instance price per hour + startup

On TitanicNet this says the training is nothing, which is the point of the two regimes:

PYTHON

n_params = sum(p.numel() for p in TitanicNet().parameters())   # the network's weights and biases: how big the model is
n_examples = len(np.load("train_data.npz")["X_train"])
flops = 6 * n_params * n_examples * planned_epochs
print(f"{n_params:,} parameters x {n_examples} examples x {planned_epochs} epochs = {flops:.1e} FLOPs")
print(f"on an A10G at a quarter of 70 TFLOPS: {flops / (70e12 * 0.25):.4f} seconds of arithmetic; the bill is the startup")

The same rule, in helpers.estimate_training_cost, picks an instance and prints the time, cost and max_run for any model size. It only earns its keep from about 1e15 FLOPs up, a few million parameters over a few million examples, which is where the fine-tuning episode uses it on a 3-billion-parameter model. Here it is on a run that does need it: a 125-million-parameter text model, a million training examples of 256 tokens, three passes:

PYTHON

helpers.estimate_training_cost(params_billions=0.125, examples=1_000_000, tokens_per_example=256, epochs=3)

About eight GPU-hours and nine dollars on the cheapest instance that fits, eight hours of wall-clock on one GPU; a job you want capped and tagged before it starts.

Challenge

Estimate your own first run

Take the model you expect to train for your own project or challenge, or the biggest one you have trained before, and estimate its first SageMaker run before you ever launch one. You need four numbers:

Parameters. For a model you have in PyTorch, sum(p.numel() for p in model.parameters()); for a model card, it is in the name (-7B, -125M).
Examples and tokens per example. Rows for tabular data, sequence length for text, pixels ÷ 784 for an image through a vision transformer; for a convolutional network use the FLOPs-per-image figure from its paper instead of the 6 × parameters rule.
Epochs. Small data usually takes three to five passes, large corpora one to three; the number is a budget, not a given.
Instance. From the Instances for ML page, or let the helper pick the cheapest that fits.

If the FLOPs come out below 1e15, use the smoke-test method instead and say what you would measure. Write the estimate down with the max_run you would set, and keep it: the first real bill goes next to it.

Show me the solution

There is no single answer; the solution is the written estimate. One worked example for scale: a 7B language model fully fine-tuned on 50,000 documents of 2,000 tokens for 3 epochs is 6 × 7e9 × 1e8 = 1.3e19 FLOPs, 45 GPU-hours on A100s, about $140 a run; the fine-tuning episode works that one through. A 10-million-parameter network on 100,000 rows for 20 epochs is 1.2e14 FLOPs: below the line, so time a smoke test and scale it, and expect the startup minutes to dominate.

We’ll also bring back the job_cost_report() helper from the previous episode, plus the hourly rates of the instances we’re about to use.

PYTHON

def job_cost_report(job_name: str, price_per_hour: float):
    desc = session.sagemaker_client.describe_training_job(TrainingJobName=job_name)  # the job record the console shows
    n = desc["ResourceConfig"]["InstanceCount"]
    billable = desc["BillableTimeInSeconds"] * n   # container pull to model upload, charged per instance
    wall = (desc["TrainingEndTime"] - desc["CreationTime"]).total_seconds()   # from the launch request to the end, including unbilled provisioning
    print(f"{job_name}: status {desc['TrainingJobStatus']}, wall-clock {wall/60:.1f} min, billable {billable/60:.1f} min on {n} instance(s), cost ${billable/3600*price_per_hour:.3f} at ${price_per_hour}/h")
    print(f"    model: {desc['ModelArtifacts']['S3ModelArtifacts']}")   # the model.tar.gz this job uploaded to output_path
    return desc

# us-east-1 on-demand rates (2026-10-16); check the instances page for your region
prices = {"ml.m5.large": 0.115, "ml.m5.xlarge": 0.23, "ml.g4dn.xlarge": 0.74}

Set up the estimator for a longer run: the epochs, the instance, where the output goes, and a two-hour max_run cap.

PYTHON

from sagemaker.pytorch import PyTorch
from sagemaker.inputs import TrainingInput

epochs = 1000          # more epochs now that the script is known to work
instance_count = 1
instance_type = "ml.m5.large"
output_path = f's3://{bucket_name}/output_nn/' # this folder will auto-generate if it doesn't exist already

# Define max runtime in seconds to ensure you don't use more compute time than expected. Use a generous threshold (2x expected time but < 2 days) so that work isn't interrupted without any gains.
max_run = 2*60*60 # 2 hours

# Define the PyTorch estimator and pass hyperparameters as arguments
pytorch_estimator = PyTorch(
    base_job_name=f"{notebook_instance_name}-nn",   # who + which model; SageMaker appends a timestamp to make each job name unique
    max_run=max_run, # in seconds; always include (max 48 hours)
    entry_point="/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/train_nn.py",  # copied into the container and run there
    role=role,
    tags=job_tags,  # per job; nothing is inherited from the notebook instance
    instance_type=instance_type, # with this small dataset, we don't necessarily need a GPU for fast training. 
    instance_count=instance_count,
    # framework_version + py_version pick the prebuilt PyTorch training image the job runs in
    framework_version="2.3",  # if fit() fails with an image error, check available versions with: sagemaker.image_uris.retrieve("pytorch", region, version="2.3", py_version="py311", image_scope="training", instance_type="ml.m5.large")
    py_version="py311",
    output_path=output_path,
    sagemaker_session=session,
    hyperparameters={  # passed to train_nn.py as command-line flags: --train ... --epochs 1000 ...
        "train": "/opt/ml/input/data/train/train_data.npz",  # SageMaker will mount this path
        "val": "/opt/ml/input/data/val/val_data.npz",        # SageMaker will mount this path
        "epochs": epochs,
        "learning_rate": learning_rate
    }
)

# Define input paths: each TrainingInput is a channel; its S3 object is copied to /opt/ml/input/data/<channel>/ in the container
# content_type is a label only; our script reads the file itself
train_input = TrainingInput(f"s3://{bucket_name}/train_data.npz", content_type="application/x-npz")
val_input = TrainingInput(f"s3://{bucket_name}/val_data.npz", content_type="application/x-npz")

# Start the training job and time it
start = t.time()
pytorch_estimator.fit({"train": train_input, "val": val_input})
end = t.time()

print(f"Runtime for training on SageMaker: {end - start:.2f} seconds, instance_type: {instance_type}, instance_count: {instance_count}")

PYTHON

desc_cpu = job_cost_report(pytorch_estimator.latest_training_job.name, prices[instance_type])

For a job this small, almost all of the billable minutes are startup.

Deploying PyTorch neural network via SageMaker with a GPU instance

In this section, we’ll implement the same procedure as above, but using a GPU-enabled instance for potentially faster training. While GPU instances are more expensive, they can be cost-effective for larger datasets or more complex models that require significant computational power.

Selecting a GPU Instance

For a small dataset like ours, we don’t strictly need a GPU, but for larger datasets or more complex models, a GPU can reduce training time. Here, we’ll select an ml.g4dn.xlarge instance, which provides a single GPU and costs approximately $0.74/hour (check Instances for ML page for detailed pricing).

Code modifications for GPU use

Using a GPU requires minor changes in your training script (train_nn.py). Specifically, you’ll need to: 1. Check for GPU availability in PyTorch. 2. Move the model and tensors to the GPU device if available.

Enabling PyTorch to use GPU in `train_nn.py`

The following code snippet enables GPU support in train_nn.py:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TitanicNet().to(device)
X_train, y_train = X_train.to(device), y_train.to(device)

Challenge

CPU or GPU: predict, then measure

We’re about to run the identical job on an ml.g4dn.xlarge (one NVIDIA T4 GPU, about $0.74/hour) instead of the ml.m5.large ($0.115/hour). Before you run the cell:

Predict which instance finishes first (total wall-clock, as fit() reports it).
Predict which costs less for this job.

Then run the GPU job, call job_cost_report() for it, and compute cost = rate x billable minutes for both runs. Were you right? Explain the result.

PYTHON

instance_type = "ml.g4dn.xlarge"  # one NVIDIA T4 GPU; train_nn.py moves to it on its own via torch.cuda.is_available()
instance_count = 1

pytorch_estimator_gpu = PyTorch(
    base_job_name=f"{notebook_instance_name}-nn",
    max_run=max_run, # in seconds; always include (max 48 hours)
    entry_point="/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/train_nn.py",
    role=role,
    tags=job_tags,
    instance_type=instance_type,
    instance_count=instance_count,
    framework_version="2.3",
    py_version="py311",
    output_path=output_path,
    sagemaker_session=session,
    hyperparameters={
        "train": "/opt/ml/input/data/train/train_data.npz",
        "val": "/opt/ml/input/data/val/val_data.npz",
        "epochs": epochs,
        "learning_rate": learning_rate
    }
)

start = t.time()
pytorch_estimator_gpu.fit({"train": train_input, "val": val_input})
end = t.time()
print(f"Runtime for training on SageMaker: {end - start:.2f} seconds, instance_type: {instance_type}, instance_count: {instance_count}")

desc_gpu = job_cost_report(pytorch_estimator_gpu.latest_training_job.name, prices[instance_type])

for label, d, rate in [("CPU ml.m5.large", desc_cpu, prices["ml.m5.large"]), ("GPU ml.g4dn.xlarge", desc_gpu, prices["ml.g4dn.xlarge"])]:
    billable_min = d["BillableTimeInSeconds"] / 60
    print(f"{label}: billable {billable_min:.1f} min x ${rate}/h = ${billable_min/60*rate:.3f}")

Show me the solution

Typically the CPU run finishes about as fast or faster, and it is several times cheaper. Both jobs spend roughly the same few minutes on startup (pulling the container, downloading data; the unbilled wait for an instance comes before that), and the GPU container is larger, so it often takes a little longer to start. The actual training is a few seconds on either instance: the network has about 2,600 parameters and the whole training set is 713 rows, so each full-batch step is tiny and the T4 sits almost idle, and the one-time CUDA initialization adds a few seconds of its own. Startup dominates both bills, and the GPU bill is about 6x higher per minute (0.74 / 0.115), so the GPU costs several times more for no speedup. In our run the CPU job billed 2.8 minutes ($0.005) and the GPU job 3.6 minutes ($0.045): slower, and nine times the price.

A GPU pays off when the per-step work is large enough to keep it busy: bigger models, bigger batches, or datasets of hundreds of MB and up. If you ever do need more GPU than a g4dn.xlarge, the next sensible step is an ml.g5.xlarge (~$1.41/hour), not a multi-GPU box.

Callout

GPUs can be slow for small datasets/models

This performance discrepancy might be due to the following factors:

Small dataset/model size: When datasets and models are small, the overhead of transferring data between the CPU and GPU, as well as managing the GPU, can slow things down. For very small models and datasets, CPUs are often faster since there’s minimal data to process.
GPU initialization overhead: Every time a training job starts on a GPU, there’s a small overhead for initializing CUDA libraries. For short jobs, this setup time can make the GPU appear slower overall.
Batch size: GPUs perform best with larger batch sizes since they can process many data points in parallel. If the batch size is too small, the GPU is underutilized, leading to suboptimal performance. You may want to try increasing the batch size to see if this reduces training time.
Instance type: Some GPU instances, like the ml.g4dn series, have less computational power than the larger g5 series. They’re better suited for inference or lightweight tasks rather than intense training, so a more powerful instance (e.g., ml.g5.xlarge) could help for larger tasks.

If training time continues to be critical, sticking with a CPU instance may be the best approach for smaller datasets. For larger, more complex models and datasets, the GPU’s advantages should become more apparent.

Distributed Training for Neural Networks in SageMaker

The previous episode gave the rule of thumb: upgrade the single instance first, and consider multiple instances only when the dataset no longer fits in one instance’s memory (tens of GB or more) or a single-instance run takes longer than about 10 hours. The GPU exercise was Option 1. Let’s now look at Option 2, starting with what instance_count=2 does.

What `instance_count=2` does on its own

Two panels comparing instance_count=2 without distribution, where each instance trains on the full data, with data parallelism, where batches are split and gradients synced

Setting instance_count=2 on an estimator makes SageMaker start two instances and run the full training script on each of them. Nothing is split:

TrainingInput defaults to distribution="FullyReplicated", so every instance downloads the complete dataset.
train_nn.py has no distributed code, so each instance trains its own independent copy of the model on all of the data for all 1000 epochs, and each saves a model.
SageMaker bills BillableTimeInSeconds x InstanceCount.

So without a distribution strategy, two instances means the same job twice, in parallel, at double the cost. Estimate first: about the same 3-4 minutes of wall-clock as the single-instance run, but billed for two instances, so about $0.01 on ml.m5.large (our run: 3.5 minutes of wall-clock, 5.5 billable minutes across the two instances, $0.010).

PYTHON

instance_count = 2   # demonstration: NOT expected to be faster, and expected to cost double
instance_type = "ml.m5.large"

pytorch_estimator_2x = PyTorch(
    base_job_name=f"{notebook_instance_name}-nn",
    max_run=max_run, # in seconds; always include (max 48 hours)
    entry_point="/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/train_nn.py",
    role=role,
    tags=job_tags,
    instance_type=instance_type,
    instance_count=instance_count,  # no distribution= argument, so each instance runs the whole script on the whole dataset
    framework_version="2.3",
    py_version="py311",
    output_path=output_path,
    sagemaker_session=session,
    hyperparameters={
        "train": "/opt/ml/input/data/train/train_data.npz",
        "val": "/opt/ml/input/data/val/val_data.npz",
        "epochs": epochs,
        "learning_rate": learning_rate
    }
)

start = t.time()
pytorch_estimator_2x.fit({"train": train_input, "val": val_input})
end = t.time()
print(f"Runtime for training on SageMaker: {end - start:.2f} seconds, instance_type: {instance_type}, instance_count: {instance_count}")

desc_2x = job_cost_report(pytorch_estimator_2x.latest_training_job.name, prices[instance_type])
print(f"Single-instance job billed {desc_cpu['BillableTimeInSeconds']/60:.1f} min; two-instance job billed {desc_2x['BillableTimeInSeconds']*2/60:.1f} min")

Look at the job’s log: open Training jobs, click the job’s name, scroll to Monitor, click View logs. This time CloudWatch lists two streams, algo-1 and algo-2, and each one prints all 1000 epochs. The wall-clock time is about the same as the single-instance run (slightly longer, if anything, because two instances have to be provisioned), and the billable minutes are double.

Data parallelism, done properly

To split one training run across two instances instead of running it twice, the script needs three things the demo above did not have, and the estimator needs one. scripts/train_nn_ddp.py is train_nn.py with those three added; open the two side by side, the differences are marked 1, 2 and 3 in the comments.

Join the process group. The launcher starts one copy of the script on each instance. The copies are numbered from zero, and a copy’s number is its rank: with two instances, one copy is rank 0 and the other rank 1, and the count of copies, 2, is the world size. Each copy learns its own rank and the world size from environment variables the launcher sets, and setup_distributed reads them and calls init_process_group, after which the copies can talk to each other. The rank is the only thing that differs between the copies, and everything below that has to differ between instances keys off it. With no launcher the function returns rank 0 of 1 and the script trains like train_nn.py.
Shard the data. Each copy keeps the rows whose index, counting from zero, leaves its rank as the remainder when divided by the world size: rank 0 takes rows 0, 2, 4, … and rank 1 takes rows 1, 3, 5, …, which is X_train[rank::world_size]. Half the training set each, no row seen twice, and the copies agree on the split without talking, because each one knows only its rank.
Wrap the model in DistributedDataParallel. After every loss.backward(), DDP averages the gradients across the processes before optimizer.step(), so both copies of the model take the same step and stay identical. That averaged gradient over two half-shards is the gradient over the whole set, which is why the result is one training run and not two. Only the copy with rank 0 prints the metrics and saves the model, since every copy holds the same weights and one file is enough; rank 0 is the conventional choice because every world has one.

On the estimator, distribution={"mpi": {"enabled": True, "processes_per_host": 1}} is what starts the processes: SageMaker launches the script on every instance through MPI, which sets the rank and world-size variables the script reads and lets the instances find each other by name. (torch_distributed, the torchrun launcher, is the same idea for GPU instances; the script handles either.) Everything else is the two-instance estimator from the demo.

PYTHON

pytorch_estimator_ddp = PyTorch(
    base_job_name=f"{notebook_instance_name}-nn-ddp",
    max_run=max_run,
    entry_point="/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/train_nn_ddp.py",   # the three additions, marked in the script
    role=role,
    tags=job_tags,
    instance_type="ml.m5.large",
    instance_count=2,
    distribution={"mpi": {"enabled": True, "processes_per_host": 1}},   # start one process per instance and connect them; without this, two copies again
    framework_version="2.3",
    py_version="py311",
    output_path=output_path,
    sagemaker_session=session,
    hyperparameters={
        "train": "/opt/ml/input/data/train/train_data.npz",
        "val": "/opt/ml/input/data/val/val_data.npz",
        "epochs": epochs,
        "learning_rate": learning_rate,
    },
)

start = t.time()
pytorch_estimator_ddp.fit({"train": train_input, "val": val_input})
print(f"Data-parallel run: {(t.time() - start) / 60:.1f} min wall-clock on 2 x ml.m5.large")
desc_ddp = job_cost_report(pytorch_estimator_ddp.latest_training_job.name, prices["ml.m5.large"])

Read the log before the bill. The two CloudWatch streams now say different things: each prints training on 446 of 891 rows (or 445), and only algo-1 prints the epochs and saves the model. The final validation loss matches the single-instance run to the third decimal, because the averaged gradient is the same gradient; that match is the proof that the work was split rather than copied. The wall-clock is not shorter. Each step now costs a round of gradient exchange between two machines, which takes longer than the 445 rows it saves, and the bill is still two instances’ worth. That is the rule of thumb again, measured: data parallelism pays when the step is big enough that the exchange is a small part of it, which on tabular data of 891 rows it never is.

Model parallelism is the other kind, for a model too large for one device’s memory, where the layers themselves are spread across devices (smdistributed.modelparallel, or PyTorch’s FSDP). There is no example here because no model in this lesson needs it: the fine-tuning episode’s memory arithmetic shows what does, a model whose weights plus their gradients and optimizer state exceed one GPU, and the usual answer at this scale is a LoRA fine-tune, which shrinks the problem instead of spreading it.

For the built-in XGBoost image from the previous episode the recipe is simpler, because the algorithm already knows how to train across instances: split the training CSV into several S3 objects under one prefix and pass TrainingInput(prefix, content_type="text/csv", distribution="ShardedByS3Key"), so each instance receives a different subset of the files. Even then, a dataset of 891 rows will only show overhead.

Either way, the rule of thumb stands: a bigger single instance first, and multiple instances only once data no longer fits in memory or a run exceeds ~10 hours.

Key Points

Efficient data handling: The .npz format is optimized for efficient loading, reducing I/O overhead and enabling batch compatibility for PyTorch’s DataLoader.
Startup dominates small jobs: Every SageMaker training job waits a minute or so for an instance (not billed), then spends one to three minutes on container pull and data download (billed) before the script runs; for a job whose script finishes in seconds, that overhead is the whole bill.
GPU training: While beneficial for larger models or datasets, GPUs may introduce overhead for smaller tasks; selecting the right instance type is critical for cost-efficiency.
Data parallelism vs. model parallelism: Data parallelism splits data across instances and synchronizes model weights, suitable for typical neural network tasks. Model parallelism, which divides model layers, is ideal for very large models that exceed memory capacity.
instance_count alone is not distributed training: Without distribution= on the estimator and distributed code in the script, instance_count=2 runs the full script on the full data on each instance and bills twice. Data parallelism is three additions to the script (join the process group, shard the rows, wrap the model in DistributedDataParallel) and one argument on the estimator; train_nn_ddp.py has them, and reaches the single-instance run’s loss on half the rows per instance.
SageMaker configuration: By adjusting instance counts and types, SageMaker supports scalable training setups. Starting with CPU training and scaling as needed with GPUs or distributed setups allows for performance optimization.
Testing locally first: Before deploying large-scale training in SageMaker, test locally with a smaller setup to ensure code correctness and efficient resource usage.
Estimate before you launch: for small models, time a smoke test and scale it by epochs and rows, plus about three minutes of startup; for large ones, 6 FLOPs per parameter per example-token, divided by a quarter of the GPU’s peak. helpers.estimate_training_cost does the second; write the number down and compare it to the bill.

Content from Hyperparameter Tuning in SageMaker: Neural Network Example

Last updated on 2026-10-21 | Edit this page

Estimated time: 60 minutes

Overview

Questions

How can we efficiently manage hyperparameter tuning in SageMaker?
How can we parallelize tuning jobs to optimize time without increasing costs?

Objectives

Set up and run a hyperparameter tuning job in SageMaker.
Define ContinuousParameter and CategoricalParameter for targeted tuning.
Log and capture objective metrics for evaluating tuning success.
Optimize tuning setup to balance cost and efficiency, including parallelization.
Read back what a whole tuning job cost from its child training jobs.

Initial setup: open prefilled .ipynb notebook

Open the notebook from: /ML_with_AWS_SageMaker/notebooks/08_Hyperparameter-tuning.ipynb. Select the conda_pytorch kernel.

Hyperparameter tuning in SageMaker

To conduct efficient hyperparameter tuning with neural networks (or any model) in SageMaker, we’ll leverage SageMaker’s hyperparameter tuning jobs while carefully managing parameter ranges and model count. Here’s an overview of the process, with a focus on both efficiency and cost-effectiveness.

Key steps

The overall process involves these seven below steps.

Setup estimator
Define parameter ranges
Set up HyperparameterTuner object
Prepare training script to log metrics
Set data paths and launch tuner.fit()
Monitor tuning job from SageMaker console
Extract best model for final evaluation

Code example for SageMaker hyperparameter tuning with neural networks

We’ll walk through each step in detail by tuning a neural network. Specifically, we will test out different values for our epochs and learning_rate parameters. We are sticking to just two hyperparameters for demonstration purposes, but you may wish to explore additional parameters in your own work.

This setup provides:

Explicit control over epochs using CategoricalParameter, allowing targeted testing of specific values.
Efficient sampling for learning_rate using ContinuousParameter, covering a defined range for a balanced approach.
Cost control by setting moderate max_jobs and max_parallel_jobs.

By managing these settings and configuring metrics properly, you can achieve a balanced and efficient approach to hyperparameter tuning for neural networks.

0. Directory setup

Just to make sure we are all on the same directory starting point, let’s cd to our instance’s root directory

PYTHON

# The Titanic project folder from the S3 episode, beside the lesson clone: downloads and outputs land here, not in the repository
!mkdir -p /home/ec2-user/SageMaker/titanic
%cd /home/ec2-user/SageMaker/titanic
import sys
sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")
import helpers

1. Setup estimator

To kick off our hyperparameter tuning for a neural network model, we’ll start by defining the SageMaker Estimator. The estimator setup here is very similar to our previous episode, where we used it to configure and train a model directly. However, this time, rather than directly running a training job with the estimator, we’ll be using it as the foundation for hyperparameter tuning.

In SageMaker, the estimator serves as a blueprint for each tuning job, specifying the training script, instance type, and key parameters like data paths and metrics. Once defined, this estimator will be passed to a Hyperparameter Tuner that manages the creation of multiple training jobs with various hyperparameter combinations, helping us find an optimal configuration.

Here’s the setup for our PyTorch estimator, which includes specifying the entry script for training (train_nn.py) and defining hyperparameters that will remain fixed across tuning jobs. The hyperparameters we’re setting up to tune include epochs and learning_rate, with a few specific values or ranges defined.

We’ll use our notebook instance name again to label the training jobs launched in this episode

PYTHON

import boto3

notebook_instance_name = 'SinkOrSwim-DoeJohn' # adjust to your notebook name. we'll use this variable to name the training jobs launched by the tuner.

# the following code just verifies you have the right name for your notebook instance
region = boto3.Session().region_name  # inherited from the notebook instance; do not hard-code a region
sagemaker_client = boto3.client('sagemaker', region_name=region) # Initialize SageMaker client
response = sagemaker_client.describe_notebook_instance(NotebookInstanceName=notebook_instance_name) # Describe the notebook instance

# Display the status and instance type
print(f"Notebook Instance '{notebook_instance_name}' status: {response['NotebookInstanceStatus']}")
local_instance = response['InstanceType']
print(f"Instance Type: {local_instance}")

Cost tracking: Just like in the previous episode, we explicitly tag each training job at launch time. This ensures that compute usage is traceable to a project, a purpose, and a human-readable name, even after the job has completed. Tags on the HyperparameterTuner are copied to every child training job it launches, so we only need to set them once, on the tuner (we also set them on the estimator so a plain .fit() on it is tagged too).

PYTHON

name = "DoeJohn" # replace with your name
project = "SinkOrSwim" # replace with your challenge name
purpose = "Tuning"        # one of: Storage, Notebook, Training, Tuning, Inference

job_tags = [
    {"Key": "Name", "Value": name},
    {"Key": "Project", "Value": project},
    {"Key": "Purpose", "Value": purpose},
]
print(f"Name: {name}, Project: {project}, Purpose: {purpose}")

Next, we’ll configure the baseline estimator that we plan to do hyperparameter search on.

PYTHON

import sagemaker
# range types for the search: ContinuousParameter (float range), IntegerParameter (whole-number range), CategoricalParameter (explicit list)
from sagemaker.tuner import HyperparameterTuner, IntegerParameter, ContinuousParameter, CategoricalParameter
from sagemaker.pytorch import PyTorch
from sagemaker.inputs import TrainingInput
from sagemaker import get_execution_role

# Initialize role, bucket, and SageMaker session variables
role = get_execution_role()
bucket_name = 'sinkorswim-doejohn'  # replace with your S3 bucket name
session = sagemaker.Session()
s3 = boto3.client('s3')
print("Role:", role)
print("Region:", session.boto_region_name)

# Define instance type/count we'll use for training
instance_type="ml.m5.large"
instance_count=1 # always start with 1; tuning parallelism comes from max_parallel_jobs, not instance_count
price_per_hour = 0.115 # ml.m5.large, us-east-1 on-demand (2026-10-16); check the instances page for your region

# Define max runtime in seconds to ensure you don't use more compute time than expected. Use a generous threshold (2x expected time but < 2 days) so that work isn't interrupted without any gains.
max_run = 2*60*60 # 2 hours

# Define the PyTorch estimator with entry script and environment details
pytorch_estimator = PyTorch(
    base_job_name=f"{notebook_instance_name}-nn",   # who + which model; the tuner appends its own suffix to each child job
    max_run=max_run, # in seconds; always include (max 48 hours)
    entry_point="/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/train_nn.py",  # Your script for training
    role=role,
    tags=job_tags,  # tags a plain .fit() on this estimator; the tuner below sets its own
    instance_count=instance_count,
    instance_type=instance_type,
    framework_version="2.3",  # if a job fails with an image error, check with: sagemaker.image_uris.retrieve("pytorch", region, version="2.3", py_version="py311", image_scope="training", instance_type="ml.m5.large")
    py_version="py311",
    # the Regex pulls the number out of each "validation:accuracy = 0.8123" line the script prints; SageMaker records it as the metric
    metric_definitions=[{"Name": "validation:accuracy", "Regex": "validation:accuracy = ([0-9\\.]+)"}],
    hyperparameters={
        "train": "/opt/ml/input/data/train/train_data.npz",  # SageMaker will mount this path
        "val": "/opt/ml/input/data/val/val_data.npz",        # SageMaker will mount this path
        "epochs": 100, # Placeholder initial value. Will be overridden by the tuning values tested
        "learning_rate": 0.001 # Placeholder initial value. Will be overridden by the tuning values tested
    },
    sagemaker_session=session,
)

2. Define hyperparameter ranges

In SageMaker, you must explicitly define ranges for any hyperparameters you want to tune. SageMaker supports ContinuousParameter, IntegerParameter and CategoricalParameter types:

ContinuousParameter allows SageMaker to dynamically sample numeric values within a specified range, making it ideal for broad, exploratory tuning. The total number of values tested can be controlled through the upcoming max_jobs parameter, which defines how many different combinations SageMaker will evaluate.
IntegerParameter does the same for whole-number settings such as a batch size or layer width.
CategoricalParameter specifies exact values for SageMaker to test, which is useful when you want the model to only try a specific set of options.

By default, SageMaker uses Bayesian optimization, adjusting future selections based on previous results to efficiently find optimal values. You can also set the strategy to “Random” for uniform sampling across the range, which is effective in larger or more exploratory search spaces. Random sampling may end up costing much more in time and resources, however. Generally, we recommend sticking with the default setting.

PYTHON

# Hyperparameter tuning ranges
hyperparameter_ranges = {
    "epochs": CategoricalParameter([100, 1000, 10000]),       # exact values to try, nothing in between; adjust as needed
    "learning_rate": ContinuousParameter(0.001, 0.1),  # any float in this range; the Bayesian strategy picks where to sample
}

Hyperparameter considerations in neural nets

When tuning hyperparameters in neural networks, it’s essential to prioritize parameters that directly impact model performance while remaining mindful of diminishing returns and potential overfitting. Below are some core hyperparameters to consider and general strategies for tuning:

Learning Rate: Often the most impactful parameter, learning rate controls the speed of convergence. A smaller learning rate can lead to more stable, though slower, convergence, while a larger rate may speed things up but risks overshooting optimal values. Testing ranges like 0.0001 to 0.1 with a ContinuousParameter is common practice, and Bayesian search can help home in on ideal values.
Batch Size: Larger batch sizes often yield faster training times and may improve stability, but this can risk bypassing useful local minima, especially with small datasets. Smaller batch sizes can capture more nuanced updates but are computationally expensive. Ranges from 16 to 256 are worth exploring for most use cases, although, for large datasets or high-capacity models, even larger values may be beneficial.
Number of Epochs: While larger epochs allow the model to learn from data longer, increasing epochs doesn’t always equate to better performance and can lead to overfitting. Exploring CategoricalParameter([50, 100, 500, 1000]) can help balance performance without excessive training costs.
Layer Width and Depth: Increasing the width or depth (i.e., number of neurons and layers) can improve model capacity but also risks overfitting if the dataset is small or lacks variability. Testing a range of layer sizes or depths (e.g., 32, 64, 128 neurons per layer) can reveal whether additional complexity yields benefits. Notably, understanding double descent is essential here, as larger networks may initially seem to overfit before unexpectedly improving in the second descent—a phenomenon worth investigating in high-capacity networks.
Regularization Parameters: Regularization techniques, such as dropout rates or weight decay, can help control overfitting by limiting model capacity. For example, dropout rates from 0.1 to 0.5 or weight decay values of 0.0001 to 0.01 often balance underfitting and overfitting effectively. Higher regularization might inhibit learning, especially if the model is relatively small.
Early Stopping: While not a traditional hyperparameter, setting up early stopping based on the validation performance can prevent overfitting without the need to exhaustively test for epoch limits. By allowing the model to halt once performance plateaus or worsens, early stopping can improve efficiency in hyperparameter tuning.
Special Phenomena - Grokking and Double Descent: For certain complex datasets or when tuning particularly large models, keep an eye on phenomena like grokking (sudden shifts from poor to excellent generalization) and double descent (an unexpected second drop in error after initial overfitting). These behaviors are more likely to appear in models with high capacity and prolonged training periods, potentially requiring longer epochs or lower learning rates to observe.

In summary, hyperparameter tuning is a balancing act between expanding model capacity and mitigating overfitting. Prioritize parameters that have shown past efficacy in similar problems, and limit the search to a manageable range—often 20–30 model configurations are sufficient to observe gains. This approach keeps resource consumption practical while achieving meaningful improvements in model performance.

3. Set up HyperparameterTuner object

In step 3, we set up the HyperparameterTuner, which controls the tuning process by specifying the…

estimator: Here, we connect the previously defined pytorch_estimator to tuner, ensuring that the tuning job will run with our PyTorch model configuration.
objective: The metric_definitions and objective_metric_name indicate which metric SageMaker should monitor to find the best-performing model; in this case, we’re tracking “validation:accuracy” and aiming to maximize it. We’ll show you how to setup your training script to report this information in the next step.
hyperparameter_ranges: Defined above.
strategy: SageMaker uses a Bayesian strategy by default, which iteratively improves parameter selection based on past performance to find an optimal model more efficiently. Although it’s possible to adjust to a “Random” strategy, Bayesian optimization generally provides better results, so it’s recommended to keep this setting.
max_jobs and max_parallel_jobs: Finally, we set max_jobs to control the total number of configurations SageMaker will explore and max_parallel_jobs to limit the number of jobs that run simultaneously, balancing resource usage and tuning speed. Since SageMaker tests different hyperparameter values dynamically, it’s important to limit total parallel jobs to <= 4.

Resource-Conscious Approach: To control costs and energy-needs, choose efficient instance types and limit the search to impactful parameters at a sensible range, keeping resource consumption in check. Hyperparameter tuning often does yield better performing models, but these gains can be marginal after exhausting a reasonable search window of 20-30 model configurations. As a researcher, it’s also important to do some digging on past work to see which parameters may be worthwhile to explore. Make sure you understand what each parameter is doing before you adjust them.

Always start with max_jobs = 1 and max_parallel_jobs=1. Before running the full search, let’s test our setup by setting max_jobs = 1. This will test just one possible hyperparameter configuration. This critical step helps ensure our code is functional before attempting to scale up.

PYTHON

# Tuner configuration
tuner = HyperparameterTuner(
    base_tuning_job_name=f"{name}-nn-tune"[:20],   # tuning job names are capped at 32 characters and the SDK appends a 12-character timestamp, so the base gets at most 20; child jobs are named <tuning job>-001-<hash>
    estimator=pytorch_estimator,  # the blueprint: every child training job is this estimator with different hyperparameters
    # same regex as the estimator; the value it captures is what the tuner compares across child jobs
    metric_definitions=[{"Name": "validation:accuracy", "Regex": "validation:accuracy = ([0-9\\.]+)"}],
    objective_metric_name="validation:accuracy",  # Ensure this matches the metric name exactly (the "Name" above)
    objective_type="Maximize",                   # Specify if maximizing or minimizing the metric
    hyperparameter_ranges=hyperparameter_ranges,  # the search space defined above
    strategy="Bayesian",  # Default setting (recommend sticking with this!); can adjust to "Random" for uniform sampling across the range
    max_jobs=1,                # Always start with 1 job for debugging purposes. Adjust based on exploration needs (keep below 30 to be kind to environment). 
    max_parallel_jobs=1,         # Always start with 1 for debugging purposes. Recommended to keep this value <= 4 since SageMaker tests values dynamically.
    tags=job_tags,               # propagated to every training job the tuner launches
)

4. Prepare training script to log metrics

To prepare train_nn.py for hyperparameter tuning, we added code to log validation metrics in a format that SageMaker recognizes for tracking. In the training loop, we added a print statement for the validation accuracy in a specific format that SageMaker can capture.

Note: It’s best to use an if statement to only print out metrics periodically (e.g., every 100 epochs), so that your print time does not eat up too much of your training time. It may be a little counter-intuitive that printing can slow things down so dramatically, but it truly does become a significant factor if you’re doing it every epoch. On the flipside of this, you don’t want to print metrics so infrequently that you lose resolution in the monitored validation accuracy. Choose a number between 100-1000 epochs or divide your total epoch count by ~25 to yield a reasonable range.

if (epoch + 1) % 100 == 0 or epoch == epochs - 1:
    print(f"validation:accuracy = {val_accuracy:.4f}", flush=True)  # Log for SageMaker metric tracking

Paired with this, our metric_definitions above uses a regular expression "validation:accuracy = ([0-9\\.]+)" to extract the val_accuracy value from each log line. This regex specifically looks for validation:accuracy =, followed by a floating-point number, which corresponds to the format of our log statement in train_nn.py.

5. Set data paths and launch tuner.fit()

In step 5, we define the input data paths for the training job and launch the hyperparameter tuning process. Using TrainingInput, we specify the S3 paths to our train_data.npz and val_data.npz files (uploaded in the previous episode). This setup ensures that SageMaker correctly accesses our training and validation data during each job in the tuning process. We then call tuner.fit and pass a dictionary mapping each data channel (“train” and “val”) to its respective path. This command initiates the tuning job, triggering SageMaker to begin sampling hyperparameter combinations, training the model, and evaluating performance based on our defined objective metric. Once the job is launched, SageMaker handles the tuning process automatically, running the specified number of configurations and keeping track of the best model parameters based on validation accuracy.

Estimate first: one child job is the same ~1 minute of training plus ~3 minutes of startup we saw in the previous episode, so about $0.01.

PYTHON

# Define the input paths
train_input = TrainingInput(f"s3://{bucket_name}/train_data.npz", content_type="application/x-npz")
val_input = TrainingInput(f"s3://{bucket_name}/val_data.npz", content_type="application/x-npz")

import time as t

# Launch the hyperparameter tuning job (blocks until the single job has finished)
start_time = t.time()
tuner.fit({"train": train_input, "val": val_input})
print(f"Hyperparameter tuning job finished in {t.time() - start_time:.1f} seconds (wall-clock, including instance startup)")

6. Monitor tuning job from SageMaker console

While the above cell runs, we can check on the progress by visiting the SageMaker Console and finding the “Training” tab located on the left panel. Click “Hyperparameter tuning jobs” to view running jobs.

What did that cost?

The money is spent by the training jobs the tuner launches, and that is where SageMaker records billable time. We bring back job_cost_report() from the previous episodes and run it over every child job of the tuner. The list of child jobs comes from tuner.analytics().dataframe() (one row per training job, with its hyperparameters and final objective value).

PYTHON

def job_cost_report(job_name: str, price_per_hour: float):
    desc = session.sagemaker_client.describe_training_job(TrainingJobName=job_name)
    n = desc["ResourceConfig"]["InstanceCount"]
    billable = desc["BillableTimeInSeconds"] * n
    wall = (desc["TrainingEndTime"] - desc["CreationTime"]).total_seconds()   # launch request to end, including unbilled provisioning
    print(f"{job_name}: status {desc['TrainingJobStatus']}, wall-clock {wall/60:.1f} min, billable {billable/60:.1f} min on {n} instance(s), cost ${billable/3600*price_per_hour:.3f} at ${price_per_hour}/h")
    print(f"    model: {desc['ModelArtifacts']['S3ModelArtifacts']}")   # the model.tar.gz this job uploaded to output_path
    return desc

def tuning_cost_report(tuner, price_per_hour):
    """Call job_cost_report on every training job the tuner launched; return the total cost in dollars."""
    jobs_df = tuner.analytics().dataframe()  # one row per child training job: its hyperparameters, FinalObjectiveValue and TrainingJobName
    total_billable = 0
    for job_name in jobs_df["TrainingJobName"]:
        desc = job_cost_report(job_name, price_per_hour)
        total_billable += desc["BillableTimeInSeconds"] * desc["ResourceConfig"]["InstanceCount"]
    total_cost = total_billable / 3600 * price_per_hour
    print(f"Tuning job {tuner.latest_tuning_job.name}: {len(jobs_df)} training jobs, {total_billable/60:.1f} billable min, ${total_cost:.2f} total")
    return total_cost

tuning_cost_report(tuner, price_per_hour)

Scaling up our approach

If all goes well, we can scale up the experiment with the below code. In this configuration, we’re scaling up the search by allowing SageMaker to test more hyperparameter configurations (max_jobs=10) while setting max_parallel_jobs=2 to manage parallelization efficiently. With two jobs running at once, SageMaker will be able to explore potential improvements more quickly than in a fully sequential setup, while still dynamically selecting promising values as it learns from completed jobs. This balance leverages SageMaker’s Bayesian optimization, which uses completed trials to inform subsequent ones, helping to avoid redundant testing of less promising parameter combinations. Setting max_parallel_jobs higher than 2-4 reduces tuning efficiency, as SageMaker’s ability to learn from completed jobs decreases when too many jobs run simultaneously.

With this approach, SageMaker is better able to refine configurations without overloading resources or risking inefficient exploration, making max_parallel_jobs=2 a solid default for most use cases.

Estimate first: 10 child jobs at roughly 3 billable minutes each is about 30 billable minutes on ml.m5.large, so about $0.06 in total. Wall-clock is five rounds of two jobs, so expect 15-20 minutes. Write both numbers down; the bill below comes in well under them, for a reason worth knowing.

PYTHON

import time as t # always a good idea to keep a runtime of your experiments 

# Configuration variables
max_jobs = 10
max_parallel_jobs = 2

# Define the Tuner configuration
tuner = HyperparameterTuner(
    base_tuning_job_name=f"{name}-nn-tune"[:20],   # tuning job names are capped at 32 characters and the SDK appends a 12-character timestamp, so the base gets at most 20; child jobs are named <tuning job>-001-<hash>
    estimator=pytorch_estimator,
    metric_definitions=[{"Name": "validation:accuracy", "Regex": "validation:accuracy = ([0-9\\.]+)"}],
    objective_metric_name="validation:accuracy",  # Ensure this matches the metric name exactly
    objective_type="Maximize",                   # Specify if maximizing or minimizing the metric
    hyperparameter_ranges=hyperparameter_ranges,
    max_jobs=max_jobs,
    max_parallel_jobs=max_parallel_jobs,
    tags=job_tags,
)

# Track start time
start_time = t.time()

# Launch the hyperparameter tuning job
tuner.fit({"train": train_input, "val": val_input})

# Calculate runtime
runtime = t.time() - start_time

# Print confirmation with runtime and configuration details
print(f"Tuning runtime: {runtime:.2f} seconds, Instance Type: {instance_type}, Max Jobs: {max_jobs}, Max Parallel Jobs: {max_parallel_jobs}")

Monitoring tuning

Initial Jobs: SageMaker starts by running only max_parallel_jobs (2 in this case) as the initial batch. As each job completes, new jobs from the remaining pool are triggered until max_jobs (10) is reached.
Job Completion: Once the first few jobs complete, SageMaker will continue to launch the remaining jobs up to the maximum of 10, but no more than two at a time.

Once it has finished, the cost report for all 10 child jobs:

PYTHON

tuning_cost_report(tuner, price_per_hour)

Our run finished in 9.3 minutes of wall-clock and billed 14.1 minutes in total, $0.03: half the estimate. Read the per-job lines and the reason is visible. The first two child jobs billed about 3 minutes each, like every job so far. The other eight billed about 1 minute each. A tuning job keeps its instances between child jobs, so jobs 3 through 10 ran on the two instances that were already up, with the container pulled and the data downloaded. The startup overhead that dominates a single small job was paid twice, not ten times.

Can/should we run more instances in parallel?

Setting max_parallel_jobs to 10 (equal to max_jobs) will indeed launch all 10 jobs in parallel. This approach won’t affect the total cost (since cost is based on the number of total jobs, not how many run concurrently), but it can impact the final results and resource usage pattern due to SageMaker’s ability to dynamically select hyperparameter values to test to maximize efficiency and improve model performance. This adaptability is especially useful for neural networks, which often have a large hyperparameter space with complex interactions. Here’s how SageMaker’s approach impacts typical neural network training:

1. Adaptive Search Strategies

SageMaker offers Bayesian optimization for hyperparameter tuning. Instead of purely random sampling, it learns from previous jobs to choose the next set of hyperparameters more likely to improve the objective metric.
For neural networks, this strategy can help converge on better-performing configurations faster by favoring promising areas of the hyperparameter space and discarding poor ones.

2. Effect of `max_parallel_jobs` on adaptive tuning

When using Bayesian optimization, a lower max_parallel_jobs (e.g., 2–4) can allow SageMaker to iteratively adjust and improve its choices. Each batch of jobs informs the subsequent batch, which may yield better results over time.
Conversely, if all jobs are run in parallel (e.g., max_parallel_jobs=10), SageMaker can’t learn and adapt within a single batch, making this setup more like a traditional grid or random search. This approach is still valid, especially for small search spaces, but it doesn’t leverage the full potential of adaptive tuning.

3. Practical impact on neural network training

For simpler models or smaller parameter ranges, running jobs in parallel with a higher max_parallel_jobs works well and quickly completes the search.
For more complex neural networks or large hyperparameter spaces, an adaptive strategy with a smaller max_parallel_jobs may yield a better model with fewer total jobs by fine-tuning hyperparameters over multiple iterations.

Summary

For fast, straightforward tuning: Set max_parallel_jobs closer to max_jobs for simultaneous testing.
For adaptive, refined tuning: Use a smaller max_parallel_jobs (like 2–4) to let SageMaker leverage adaptive tuning for optimal configurations.

This balance between exploration and exploitation is particularly impactful in neural network tuning, where training costs can be high and parameters interact in complex ways.

Extracting and evaluating the best model after tuning

Tuning takes about 10 minutes to complete — not bad for 10 models! After SageMaker completes the hyperparameter tuning job, the results, including the trained models for each configuration, are stored in an S3 bucket. Here’s a breakdown of the steps to locate and evaluate the best model on test data.

Understanding the folder structure:
- SageMaker saves each tuning job’s results in the specified S3 bucket under a unique prefix.
- For the best model, SageMaker stores the model artifact in the format s3://{bucket}/{job_name}/output/model.tar.gz. Each model is compressed as a .tar.gz file containing the saved model parameters.
Retrieve and load the best model:
- Using the tuner.best_training_job() method, you can get the name of the best-performing job.
- From there, retrieve the S3 URI of the best model artifact, download it locally, and extract the files for use.
Prepare test data for final assessment of model generalizability
- Using the scaler that was fit on the training data.
Evaluate the model on test data:
- Once extracted, load the saved model weights and evaluate the model on your test dataset to get the final performance metrics.

Here’s the code to implement these steps:

View best model details and storage info

We can easily view the best hyperparameters from the tuning procedure…

PYTHON

# 1. Get the best training job from the completed tuning job
best_job_name = tuner.best_training_job()  # the child job with the highest validation:accuracy
print("Best training job name:", best_job_name)

# 2. Use describe_training_job to retrieve full details, including hyperparameters...
best_job_desc = session.sagemaker_client.describe_training_job(TrainingJobName=best_job_name)
best_hyperparameters = best_job_desc["HyperParameters"]
print("Best hyperparameters:", best_hyperparameters)

# ...  and model URI (location on S3)
best_model_s3_uri = best_job_desc['ModelArtifacts']['S3ModelArtifacts']  # the model.tar.gz that job uploaded when it finished
print(f"Best model artifact S3 URI: {best_model_s3_uri}")

Retrieve and load best model

Download the winning job’s model.tar.gz and unpack it; SageMaker packs whatever the script saved to /opt/ml/model into that archive.

PYTHON

import tarfile

# Download and extract the model artifact: SageMaker packs whatever the script saved to /opt/ml/model into model.tar.gz
local_model_path = "best_model.tar.gz"
model_bucket, model_key = best_model_s3_uri.split('/')[2], '/'.join(best_model_s3_uri.split('/')[3:])  # s3://bucket/key -> bucket, key
s3.download_file(model_bucket, model_key, local_model_path)

# Extract the model files from the tar.gz archive
with tarfile.open(local_model_path, 'r:gz') as tar:
    tar.extractall()
print("Best model downloaded and extracted.")

Prepare test set as test_data.npz

In our previous episode, we converted our train dataset into train/validate subsets, and saved them out as .npz files for efficient processing. We’ll need to preprocess our test data the same way to evaluate it on our model.

Note: It’s always a good idea to keep preprocessing code as a function so you can apply the same exact procedure across datasets with ease. We’ll import our preprocessing function from train_nn.py, and pass it the scaler that was fit on the training data (saved as scaler.joblib in the previous episode). Calling preprocess_data(test_df) without the scaler would fit a new one on the test set and transform the features differently from what the model saw in training.

PYTHON

!pip install scikit-learn   # train_nn.py imports it; not in the notebook kernels

Fetch the test CSV and the training-set scaler, transform the test rows the same way the training rows were, and save the arrays.

PYTHON

import pandas as pd
import numpy as np
import joblib
from train_nn import preprocess_data

# Fetch the test CSV and the training-set scaler from S3 (both uploaded in earlier episodes)
s3.download_file(bucket_name, "titanic_test.csv", "titanic_test.csv")
s3.download_file(bucket_name, "scaler.joblib", "scaler.joblib")
scaler = joblib.load("scaler.joblib")  # the StandardScaler fit on the training data in the previous episode

# Preprocess test data with the same scaler used for training
test_df = pd.read_csv("titanic_test.csv")
X_test, y_test, _ = preprocess_data(test_df, scaler=scaler)  # transform only; passing the scaler stops it refitting on test data

# Save processed data for testing
np.savez('test_data.npz', X_test=X_test, y_test=y_test)

Evaluate the model on test data

Rebuild the network from the script’s class, load the saved weights, and measure accuracy on the test set.

PYTHON

from train_nn import TitanicNet
from train_nn import calculate_accuracy
import torch

# Load the model (saved as 'nn_model.pth' by train_nn.py)
model = TitanicNet()
model.load_state_dict(torch.load("nn_model.pth", map_location="cpu"))
model.eval()

# Load test data
test_data = np.load("test_data.npz")
X_test = torch.tensor(test_data['X_test'], dtype=torch.float32)
y_test = torch.tensor(test_data['y_test'], dtype=torch.float32).unsqueeze(1)

# Evaluate the model on the test set
with torch.no_grad():
    predictions = model(X_test)
    accuracy = calculate_accuracy(predictions, y_test)
    print(f"Test Accuracy: {accuracy:.4f}")

Conclusions

In about 10 minutes, we produced a model that is 77-80% accurate on the test set. However, this performance does come at a cost (albeit manageable if you’ve stuck with our advice thus far). This next section will help you assess the total compute time that was used by your tuning job.

Each child job is billed like any training job; the tuning job itself costs nothing. Here’s a breakdown of what’s happening:

Billable Time (BillableTimeInSeconds): the seconds each child job’s instance was yours, from container pull through the script to model upload. The wait for an instance before that is not billed. (TrainingTimeInSeconds is the same number on on-demand instances; it only differs on spot.)
Instance Count: billable seconds are multiplied by the number of instances the job used.
Per-Second Billing: SageMaker bills per second, with no rounding to minutes or hours. For short jobs like ours, container pull and data download are most of the bill, which is why ten short jobs cost more than one job that trained ten times as long.

For convenience, we have added this as a function in helpers.py

PYTHON

helpers.calculate_tuning_job_time(tuner, region)  # same per-child-job sum as tuning_cost_report, packaged in helpers.py

Challenge

All at once

Rerun the 10-job search with max_parallel_jobs=10 (equal to max_jobs). Before you launch it, predict what changes in (a) total cost, (b) wall-clock time, and (c) the quality of the search. Then run it and check (a) and (b) with tuning_cost_report() and the runtime print.

PYTHON

tuner_parallel = HyperparameterTuner(
    base_tuning_job_name=f"{name}-nn-tune"[:20],   # tuning job names are capped at 32 characters and the SDK appends a 12-character timestamp, so the base gets at most 20; child jobs are named <tuning job>-001-<hash>
    estimator=pytorch_estimator,
    metric_definitions=[{"Name": "validation:accuracy", "Regex": "validation:accuracy = ([0-9\\.]+)"}],
    objective_metric_name="validation:accuracy",
    objective_type="Maximize",
    hyperparameter_ranges=hyperparameter_ranges,
    max_jobs=10,
    max_parallel_jobs=10,
    tags=job_tags,
)

start_time = t.time()
tuner_parallel.fit({"train": train_input, "val": val_input})
print(f"Tuning runtime: {t.time() - start_time:.2f} seconds with max_parallel_jobs=10")
tuning_cost_report(tuner_parallel, price_per_hour)
print("Best objective, sequential-ish run:", tuner.analytics().dataframe()["FinalObjectiveValue"].max())
print("Best objective, all-parallel run:  ", tuner_parallel.analytics().dataframe()["FinalObjectiveValue"].max())

Show me the solution

Cost: higher, not the same. Our all-parallel run billed 27.5 minutes ($0.05) against 14.1 minutes ($0.03) for the two-at-a-time run. Ten instances each pull the container and download the data, so the startup overhead is paid ten times; with two at a time it was paid twice and the other eight child jobs reused warm instances. (If your account has a quota on concurrent ml.m5.large training instances, the request may be rejected or throttled, which is the practical limit on this approach.)
Wall-clock: roughly one job’s duration. Our run took 4.3 minutes against 9.3, so about twice as fast, not five times: the two-at-a-time run’s later jobs were short because they skipped startup.
Search quality: with all 10 launched at once, the Bayesian strategy has no completed results to learn from, so the 10 configurations are effectively a random sample of the space. On a 2-parameter problem this small you will usually see a similar best objective either way; on a larger search space with more jobs, the sequential-ish run tends to find a better configuration for the same number of jobs.

Challenge

Tune the batch size too

Suppose you want to add batch_size to the search as IntegerParameter(16, 256). What has to change in train_nn.py for that to do anything, and what has to change in the notebook?

Show me the solution

The tuner can only vary what the script reads. Right now train_nn.py has no --batch_size argument and trains on the full dataset every epoch (one model(X_train) call per epoch), so a batch_size hyperparameter would be passed on the command line, rejected by argparse, and the job would fail. Three changes are needed:

In train_nn.py: add parser.add_argument('--batch_size', type=int, default=64), pass it to train(), and replace the full-batch step with a mini-batch loop, for example by wrapping the tensors in a TensorDataset and iterating over a DataLoader(dataset, batch_size=batch_size, shuffle=True) inside each epoch (with optimizer.zero_grad(), forward, loss, backward() and step() per batch).

parser.add_argument('--batch_size', type=int, default=64)
...
loader = DataLoader(TensorDataset(X_train, y_train), batch_size=batch_size, shuffle=True)
for epoch in range(epochs):
    model.train()
    for xb, yb in loader:
        optimizer.zero_grad()
        loss = criterion(model(xb), yb)
        loss.backward()
        optimizer.step()

In the estimator: add a placeholder "batch_size": 64 to hyperparameters, so the argument is always passed.
In the ranges: add "batch_size": IntegerParameter(16, 256) to hyperparameter_ranges.

Two side effects worth anticipating: an epoch now takes several optimizer steps instead of one, so the same epochs values train for longer (consider shrinking the epochs choices), and the search space has grown, so 10 jobs will cover it less well; 20-30 would be a reasonable budget, which at these prices is still well under a dollar.

Key Points

Tuner = estimator + ranges + objective: A HyperparameterTuner wraps an ordinary estimator; the script only has to print the objective metric in a format the metric_definitions regex can capture.
Start with one job: max_jobs=1, max_parallel_jobs=1 first, to catch script and configuration errors before paying for a search.
Parallelism buys time, not quality: Cost is per child training job; max_parallel_jobs changes wall-clock time and, with the Bayesian strategy, how much each job can learn from the ones before it. 2-4 parallel jobs is a good default.
Tags propagate from the tuner: Tags on the tuner are copied to every child training job, the one case in SageMaker where tags are inherited.
Read the bill from the child jobs: SageMaker bills per second of BillableTimeInSeconds x instance count on each training job; container pull and data download are inside that number and dominate short jobs.
Transform the test set with the training scaler: Load the saved scaler and pass it to preprocess_data(test_df, scaler=scaler); never refit on test data.

Content from Fine-tuning a Vision-Language Model with Training Jobs

Last updated on 2026-10-21 | Edit this page

Estimated time: 180 minutes

Overview

Questions

When is fine-tuning worth it?
Why does fine-tuning cost anywhere from a dollar to tens of thousands, and how do I estimate where my run falls before I launch it?
How do I run a GPU fine-tune as a SageMaker training job so the meter stops the moment it finishes?
How do I measure whether the fine-tune helped?

Objectives

Explain what drives the cost of a fine-tuning run (model size × tokens × epochs × number of attempts) and why the same four habits apply at $1 and at $10,000: a time cap, a small first run, and an estimate written down before launch.
Launch a LoRA fine-tune of a 3B vision-language model (CHURRO-3B) as a HuggingFace training job on a single 24 GB GPU, with tags, a unique job name and a max_run cap.
Read a job’s billable time and compare it to the estimate.
Compare the base model and the fine-tuned model on held-out pages using the competition’s own character error rate (CER).
Know how to scale the same job up to the full dataset (homework) and to multi-GPU instances for larger models.

Open the notebook from: /ML_with_AWS_SageMaker/notebooks/09_Finetuning-BadgerScribe.ipynb. Select the conda_pytorch kernel.

Why fine-tune an LLM

The training episodes so far trained small models from scratch. This episode fine-tunes a large language model with LoRA, the cheap form of fine-tuning explained below: a vision-language model with billions of parameters, pretrained on more data and compute than any lab could repeat. That is what makes it a different kind of job, in what it can do and in what it costs.

General-purpose models are trained on the average of the internet, and research problems are rarely average. A 19th-century German hand, a lab’s own annotation conventions, species names, a regional dialect, the layout of a specific archive: the largest commercial model has seen little of it, and no prompt can make up for that.

Done on the right data, a small fine-tuned model beats a much larger general one. The model we LoRA fine-tune in this episode is itself the example: CHURRO, a 3-billion-parameter vision-language model, was fully fine-tuned from Qwen2.5-VL-3B on about 100,000 pages of historical documents, and its authors report that it exceeds the accuracy of Gemini 2.5 Pro on historical text recognition at 15.5x lower cost (paper). A model that size also runs on one GPU you control, which matters when the data cannot leave your institution, when results have to be reproducible with fixed weights, and when the job is a million pages rather than a demo. CHURRO is not a one-off: the same pattern shows up wherever a research task has its own data.

Dot plot of benchmark scores for three task-specific models against a much larger general model on the same benchmark; the specialist is ahead in each row by 6.5, 2.2 and 1.2 points

Sources: CHURRO, EMNLP 2025; Goedel-Prover-V2, ICLR 2026; MedGemma technical report, 2025. Scores are as reported by each model’s authors.

The trick

Fine-tuning is ordinary training that starts from a trained model instead of random weights. You show the model your examples, it computes how wrong its output was, and gradient descent nudges the weights toward your data. Because the model already knows how to read, write and look at images, it needs hundreds of examples to learn your task where training from scratch would need millions. Everything it learned before stays in the weights; you are adjusting them, not replacing them.

Full fine-tuning is expensive

Training every weight means storing four numbers per weight instead of one. There is the weight itself, and its gradient, which says which way to nudge it this step. Then there is what the optimizer remembers: Adam, the optimizer everyone uses, keeps for every weight a running average of its recent gradients and a running average of their squares, so that it can give each weight its own step size. And the weight has to be kept in 32-bit as well as 16-bit, because each step’s nudge is far smaller than a 16-bit number can hold; applied to a 16-bit weight it would round away to nothing. Two bytes for the working weight, two for the gradient, four for the 32-bit weight and eight for Adam’s two averages: 16 bytes per parameter, against 2 for inference. A 3B model that runs in 6 GB needs 48 GB to fully fine-tune; a 7B model needs 112 GB, which is already a multi-GPU instance at $25 an hour or more, and the bill scales with every epoch and every retry.

Do not fully fine-tune anything in this workshop (that is, update all the weights), and do not LoRA fine-tune a model larger than 7B on the class account. A 7B LoRA needs about 22 GB, which is the most the 24 GB GPU the class uses will hold; anything larger needs the 48 GB instance, and above about 13B no instance the account allows holds the model at all. On a few hundred examples a 7B LoRA run is still a few dollars, so retries are cheap. A full fine-tune of a 7B model needs 112 GB, more than any instance the class account allows, and would be a few hundred dollars per run on one that did; a 70B model is $6,000 to $12,000 per run, and a hyperparameter sweep multiplies that by ten. Those runs exist in this episode only as estimates, so that you can recognise one before you launch it.

LoRA

LoRA (low-rank adaptation) freezes the pretrained weights and trains a few million extra parameters in small matrices bolted onto the attention and MLP layers. For each large weight matrix W (d × d), it adds two thin ones, B (d × r) and A (r × d), with a rank r of 8 to 64; the layer computes W·x + B·A·x, and only A and B get gradients.

Left: a d by d weight matrix, all trained, 16 bytes per weight. Right: the same matrix frozen in bf16 plus two thin trained matrices B (d by r) and A (r by d); output is W x plus B A x.

The frozen weights can stay in 16-bit, since nothing updates them, and the gradients and optimizer state only exist for the adapter. The adapter is about 100 MB, trains on a single 24 GB GPU for a 3B model, and recovers most of the gain of a full fine-tune for a task like this: the model still learns your hands and your conventions, it just learns them through the adapter. This episode uses LoRA.

Scale, and how easy it is to get wrong

What makes fine-tuning an LLM different is scale. The training episodes ran for minutes on a $0.12-an-hour CPU; a full fine-tune runs for hours or days on a GPU, a LoRA fine-tune for minutes to hours, and every setting multiplies the bill: model size, epochs, image resolution, instance type, and whether anyone is watching when the job should have ended. Any one of them off by 10x is a 10x bill, and the first run is never the last. Fine-tuning is one of the most common surprise costs we see hurt people, and it is the reason this episode spends as long on estimating and capping a run as on running it. We measure the gain ourselves: the episode ends by scoring the base model and the LoRA fine-tuned model on the same held-out pages.

The task: Badger Scribe

Badger Scribe is a Kaggle challenge built from UW Digital Collections material: 19th-century surveyors’ field notebooks, craftsmen’s account books, and German immigrant letters in Kurrent script. Given a page image, produce a verbatim transcription. The score is character error rate (CER): edits per reference character, lower is better.

This lesson uses the German letters, the Anna Seifert correspondence held by the Max Kade Institute: 191 letters and postcards written between 1932 and 1942 by one woman to her son in Kurrent (the old German handwriting taught in schools until the 1940s, whose letterforms look nothing like today’s cursive: an e is two short strokes, an s can take three shapes, and most Germans alive today cannot read it), transcribed by the Institute’s Kurrent Transcription Group. It is the part of the competition data where every label was produced by a person, and it is most of the hidden test set.

split	pages	documents	label source
train	295	4 letter bundles	all human-transcribed, ~620 characters per page
test	141	5 other letter bundles	labels held by Kaggle

Small is the point. 295 labelled pages is a realistic amount of data for a library or a lab to have, and it is enough to move a model that was pre-trained on the right kind of material. Kurrent is also a genuinely different script, not a style of the one the model already reads, so this is a case where the model has to learn something new.

We LoRA fine-tune CHURRO-3B. It already reads old handwriting across 46 language clusters; we are teaching it these hands and these transcription conventions.

Step 0: Estimate the cost before you launch anything

Before launching a fine-tune, write down what it should cost. If the bill is 3x your number, something is wrong, and you want to know on the first run, not the fiftieth.

What fine-tuning costs, from one dollar to ten thousand

The cost of a run is GPU-hours × the hourly rate, and GPU-hours scale with model size × tokens seen × epochs, multiplied by how many runs you do before one works. Our class run sits at the very bottom of that curve. The rows below use current SageMaker on-demand prices in us-east-1 (AWS price list, October 2026):

Run	GPUs	Instances rented (price per instance-hour)	Hours	GPU-hours	Cost = instance-hours × price
This episode: 3B LoRA, all 251 Kurrent training pages, 2 passes, 0.7M tokens seen	1 × A10G	1 × `ml.g5.xlarge` (`$1.41`)	0.53	0.53	`$0.75`
Homework: the same pages, 4 passes, 40 pages scored, 1.4M tokens	1 × A10G	1 × `ml.g5.xlarge` (`$1.41`)	1.1	1.1	~`$1.50`
Badger Scribe teams: 3B LoRA, the whole train set (619 pages, all three categories), 2 passes, 1.8M tokens	1 × A10G	1 × `ml.g5.xlarge` (`$1.41`)	0.85	0.85	~`$1.20`
7B full fine-tune, 50k documents, 300M tokens (the challenge below)	8 × A100	1 × `ml.p4d.24xlarge` (`$25.25`)	5.5	44	~`$140`
70B LoRA, one day	8 × A100	1 × `ml.p4d.24xlarge` (`$25.25`)	24	192	~`$600`
CHURRO-3B’s own training: full fine-tune, 100k pages	32 × H100	4 × `ml.p5.48xlarge` (`$63.30`)	25	800	~`$6,300`
70B full fine-tune, two days	32 × H100	4 × `ml.p5.48xlarge` (`$63.30`)	48	1,536	~`$12,000`
Any row above × a 10-run hyperparameter sweep					× 10

Tokens seen is examples × tokens per example × passes over the data, and it is what the time rule below multiplies by model size. The number of passes is a choice, and the class run’s two passes are a budget choice for a 40-minute job, not a recommendation: large instruction-tuning corpora are usually passed over one to three times, but a few hundred examples usually take three to five. Each extra pass adds one more unit of the same row, so a ten-pass run over the whole train set is about $10, and the sweep row still applies on top. Step 2 shows how the per-epoch curve decides it.

The class row is measured (2026-10-21: 31.7 billed minutes, $0.75); the homework and teams rows are that run’s seconds per page applied to more passes and pages. The teams row is the competition’s full training set; the shared bucket holds only the German letters, so a team trains on everything from the competition’s own data download.

GPUs come in fixed boxes: an ml.g5.xlarge is one A10G, an ml.p4d.24xlarge is eight A100s, an ml.p5.48xlarge is eight H100s, and you rent whole boxes. So “32 H100s for 25 hours” is four ml.p5.48xlarge for 25 hours: 100 instance-hours at $63.30 is $6,330. GPU-hours (GPUs × hours) is what the time estimate below produces; divide by the GPUs per box to get instance-hours, and multiply by the box’s price to get the bill.

Horizontal bar chart on a logarithmic dollar axis: this episode 0.95, homework 1.90, the full Badger Scribe train set 2.25, 7B full fine-tune 140, 70B LoRA 600, CHURRO-3B 6,300, 70B full fine-tune 12,000, with the ten-run sweep shown as a lighter extension of each bar

Two rows to note. The CHURRO row is the model you are about to LoRA fine-tune: Stanford spent on the order of 800 H100-hours making it, and you get the benefit for a dollar because you are adapting it, not retraining it. Starting from a model that already knows most of the task is the largest cost saving available, and it is decided before any GPU is involved. The sweep row: a first run is rarely the last, so a “$140 full fine-tune” is usually a $700 to $1,400 project. That is the number that surprises research groups. (The 7B row is worked out in the challenge below; the 70B rows are rule-of-thumb magnitudes; the CHURRO figures come from a published review of the paper.)

The habits in this episode are cheap practice at a dollar a run. At $140 a run they are the difference between a project and a budget incident; at $6,000 a run, one forgotten max_run or one wrong instance type is a grant line.

Back of the envelope: will it fit, and how long will it take?

Two rules of thumb answer most of the question before you touch a console. They are good to within a factor of two, which is what you need to tell a $5 run from a $500 one.

1. Memory: will it fit on the GPU? Count bytes per parameter for what you are doing, multiply by the parameter count, then add a few GB for activations (more for long sequences and big batches).

Stacked bar chart of bytes per parameter for inference versus LoRA versus full fine-tuning

What you are doing	Bytes per parameter	Why	3B	7B	27B	70B
Inference, 4-bit weights	0.5	weights only	1.5 GB	3.5 GB	14 GB	35 GB
Inference, FP8	1	weights only	3 GB	7 GB	27 GB	70 GB
Inference, bf16	2	weights only	6 GB	14 GB	54 GB	140 GB
LoRA fine-tune	~2.5	frozen bf16 weights + a small adapter with its gradients and optimizer + activations	8 GB	18 GB	68 GB	175 GB
Full fine-tune, mixed precision	16	16-bit weights 2 + gradients 2 + 32-bit copy of the weights 4 + Adam’s two running averages 8	48 GB	112 GB	430 GB	1.1 TB

Note the jump from the inference rows to the training rows. Training cannot run on quantized weights the way inference can: each step’s update is too small for 8-bit or 4-bit numbers to hold, so every trained weight keeps a 32-bit copy plus Adam’s two averages, 16 bytes against 2. That is why a 7B model you can run on a 24 GB GPU needs 112 GB plus activations to fully fine-tune, which means several GPUs with the model sharded across them, at $25 to $63 an hour. LoRA sidesteps it by freezing the model and training a few million adapter parameters, which is why this episode fits on one 24 GB card.

2. Time: how many GPU-hours? Training a token costs three matrix products per layer, each one multiply and one add per weight. Full fine-tuning does all three; LoRA skips the weight update for the frozen base.

Three matrix products per token: forward (2 FLOPs per weight), error passed back to the inputs (2), and the weight update (2); full fine-tuning does all three for 6, LoRA skips the update for 4

tokens      = examples × tokens per example × passes over the data
FLOPs       = 6 × parameters × tokens      (4 × for LoRA)
GPU-seconds = FLOPs ÷ (GPU peak FLOPS × 0.25)      a real job reaches about a quarter of the data-sheet peak

Peak FLOPS is on the GPU’s data sheet: about 70 trillion per second for an A10G, 312 for an A100, 989 for an H100. Cost is GPU-hours × the hourly rate, plus about 10 minutes of startup per job.

The class run: 251 pages × 1,600 tokens × 2 passes = 803k tokens. LoRA on 3.75B parameters: 4 × 3.75e9 × 8.03e5 ≈ 1.2e16 FLOPs. A quarter of an A10G’s 70 TFLOPS is 17 TFLOPS, so about 700 seconds, call it 12 minutes. Measured, the same instance trains at 1.45 seconds per page-pass, so 502 page-passes take about 12 minutes: the rule lands on the number.

helpers.estimate_finetune_cost does the same arithmetic for any model size, example count, tokens per example and passes. It picks the cheapest instance that fits, or the one you name, and prints memory, GPU-hours, cost per run and per project, and the max_run to set:

PYTHON

import sys
sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")
import helpers

helpers.estimate_finetune_cost(
    params_billions=3.75,        # CHURRO-3B: 3B language model + vision tower
    examples=251, tokens_per_example=1600, epochs=2,   # tokens, not words: ~1,300 for the image + ~300 of text; counted in Setup below
    method="lora",
    instance_type="ml.g5.xlarge",
)

And on the kind of run a marathon team might be weighing, with the attempts a real project takes:

PYTHON

helpers.estimate_finetune_cost(
    params_billions=7, examples=5000, tokens_per_example=800, epochs=3,
    method="lora", runs=8,       # instance_type=None: the cheapest that fits
)

tokens_per_example is the input you have to think about, and it is tokens, not words: the model never sees words. Text is about 1.3 tokens per English word, more for German and for old spellings; an image on a Qwen-VL style model is one token per 28 × 28 pixel block after resizing, about 1,300 for a megapixel; for audio, EEG windows or anything else, the model card says what one example becomes. Do not guess it when you can count it: the Setup section below counts the tokens in one page with the model’s own processor. The estimate is good to a factor of two, and it is the number to write down before the launch.

More GPUs do not make it cheaper. GPU-hours are the cost; spreading them across eight GPUs changes the wall-clock time, not the bill, and only helps if the job is large enough to keep all eight busy. You add GPUs when the memory row above says one GPU cannot hold the job, not to save money.

Challenge

Estimate a run you have not launched

A lab wants to fully fine-tune a 7B model on 50,000 documents of ~2,000 tokens for 3 epochs. They are choosing between ml.p4d.24xlarge (8 × A100 40 GB, $25.25/h, ~312 TFLOPS bf16 per GPU) and renting a single ml.g5.xlarge; the supplemental Instances for ML page has both, with what each box holds and how its price covers all its GPUs. Using the two rules: does it fit on the g5? How many GPU-hours on the p4d, and what does one run cost? What does a realistic project cost?

Show me the solution

It does not fit on the g5: full fine-tuning a 7B model needs ~112 GB plus activations, and the g5 has 24 GB. On the p4d the 8 × 40 GB = 320 GB holds it with the model sharded.

Tokens: 50,000 × 2,000 × 3 = 3 × 10^8. FLOPs: 6 × 7 × 10^9 × 3 × 10^8 ≈ 1.26 × 10^19. A quarter of an A100’s 312 TFLOPS is 78 TFLOPS: 1.26 × 10^19 / 7.8 × 10^13 ≈ 1.6 × 10^5 GPU-seconds ≈ 45 GPU-hours, which is about 6 hours of wall-clock on 8 GPUs, about $140 for one run (the “7B full fine-tune” row in the table above). The helper gives the same answer:

PYTHON

helpers.estimate_finetune_cost(params_billions=7, examples=50_000, tokens_per_example=2000, epochs=3,
                               method="full", instance_type="ml.p4d.24xlarge",
                               gradient_checkpointing=False, runs=10)

A realistic project is not one run. The first run finds a data bug, the second a bad learning rate, then two or three are compared, then the one that worked is rerun on the final data: five to ten runs, so $700 to $1,400, before anyone touches a hyperparameter sweep. That is the number to put in the proposal, and it is why the four habits matter more as the model gets bigger.

The instance

The model we LoRA fine-tune fits on one 24 GB GPU. Single-GPU and multi-GPU SageMaker training instances, same price list (the full table, with CPU instances and the GPU-by-GPU notes, is on the supplemental Instances for ML page):

Instance	GPU	GPU memory	`$`/hour
`ml.g4dn.xlarge`	1 × T4	16 GB	`$0.74`
`ml.g6.xlarge`	1 × L4	24 GB	`$1.13`
`ml.g5.xlarge`	1 × A10G	24 GB	`$1.41`
`ml.g6e.xlarge`	1 × L40S	48 GB	`$2.61`
`ml.g5.12xlarge`	4 × A10G	96 GB	`$7.09`
`ml.p4d.24xlarge`	8 × A100	320 GB	`$25.25`
`ml.p5.48xlarge`	8 × H100	640 GB	`$63.30`

Multi-GPU instances are the right tool when the model, its optimizer state and a batch do not fit on one GPU: a full fine-tune of a 7B model, or anything at 70B. They are the wrong tool when the job does fit on one GPU, because the price goes up 5x to 45x and a job too small to split gets no faster. A 3B LoRA run fits with room to spare, so ml.g5.xlarge is the workhorse for this episode. A T4 is too small and too slow for a 3B vision model with page-sized images; ml.g6.xlarge is a fine substitute if g5 capacity is short; ml.g6e.xlarge is roughly 2–3x faster for less than 2x the price, so it becomes cheaper per run once training outweighs the fixed startup (Exercise 2).

Scaling the same job to more GPUs

Nothing in this episode is specific to small runs. The training script uses Hugging Face Trainer, which handles multi-GPU data parallelism when SageMaker launches it under torchrun. To run the same script on a multi-GPU instance, change the estimator, not the script:

instance_type="ml.p4d.24xlarge",                        # 8 x A100, or ml.g5.12xlarge for 4 x A10G
distribution={"torch_distributed": {"enabled": True}},  # one process per GPU, each on a different slice of the data
max_run=6 * 60 * 60,                                     # 6 hours, in seconds: from the Step 0 estimate, with margin

Each GPU holds a full copy of the model, so the effective batch becomes batch_size x grad_accum x number of GPUs, and the learning rate usually needs a modest increase. Two lines in the script should be wrapped in if trainer.is_world_process_zero(): so that evaluation and saving run once rather than on every process. When the model, its optimizer state and a batch do not fit on one GPU (a full fine-tune of 7B or more; see the bytes-per-parameter table), data parallelism is not enough: pass fsdp="full_shard" in TrainingArguments (or a DeepSpeed ZeRO config) so the model is sharded across the GPUs, and use instance_count=2 or more for multi-node runs.

The cost controls do not change with size, but the stakes do. An ml.p4d.24xlarge bills $25.25 an hour and an ml.p5.48xlarge $63.30 an hour from the moment the job starts, and a max_run of six hours is a $150 to $380 ceiling on a single launch. Estimate GPU-hours from the Step 0 rules, set max_run to about twice that estimate, tag the job, and run a slice of the data on one GPU first so that a bug costs a dollar rather than a few hundred. Do not launch a multi-GPU job on the class account.

The run

A training job bills from the moment its instance is ready until the job ends, including the minutes spent pulling the container and downloading data and model weights. In class we train on all 251 German-letter training pages for 2 passes and score 15 held-out pages before and after, all inside one job:

Phase	Time	Cost
Waiting for the instance (not billed)	~1 min	`$0`
Startup: container pull, ~130 MB of images, 7.5 GB of model weights, model load	~5 min	`$0.12`
Score 15 holdout pages with the base model, ~21 s each	~5 min	`$0.12`
Training: 251 pages × 2 passes, ~1.5 s per page-pass with gradient checkpointing	~12 min	`$0.28`
Score 8 of the holdout pages at the end of each pass, ~17 s each, for the curve	~5 min	`$0.11`
Score the same 15 pages with the LoRA fine-tuned model, ~17 s each	~4 min	`$0.10`
Save and upload the adapter	~1 min	`$0.02`
One run	~33 min wall-clock, ~32 min billed	`$0.75`

Every row is measured, from this exact job run on 2026-10-21: 32.5 minutes from launch to finish, 31.7 billed, $0.75. Scoring a page takes about 17 to 21 seconds, not the 10 a first guess might assume, so the two 15-page evaluations and the curve points are 14 of the 32 minutes: nearly half the bill goes to measuring, not training.

Training time scales linearly with pages × epochs, and roughly linearly with image resolution (every 28 × 28 pixel block of the page becomes one token in the model’s sequence). Doubling max_pixels doubles the time; doubling epochs doubles it again. Those two settings, left at careless values, are why the estimate must be written down before the first run and compared to the real BillableTimeInSeconds after it.

We set max_run to about twice the estimate. If a run goes past that, SageMaker kills it and we lose at most a dollar, not a weekend.

The whole episode

What	Instance	Time	Cost
Controller notebook (no GPU needed)	`ml.t3.medium`	3 h	`$0.15`
LoRA fine-tune on all 251 pages, 2 passes, with before/after scoring and a per-pass curve	`ml.g5.xlarge`	~33 min	`$0.75`
Exercise 1: same job at 2 megapixels	`ml.g5.xlarge`	~60 min	`$1.40`
Total			~`$2.60`

Compare the classic mistake: doing all of this interactively on a ml.g5.xlarge notebook instance. The class itself costs about the same, but a notebook instance keeps billing after you close the tab. Forgotten over a weekend: $90. Forgotten for a month: $1,029. A training job cannot be forgotten, because it ends.

Callout

The four habits

Cap it. Every job gets max_run. Every notebook GPU instance gets stopped when you stand up.
Start small. The first launch of any training script uses a slice of the data and a short cap. The in-class run below is that small first run; the homework scales it up only after it has been seen to work.
Estimate it, then check. Write the number down. After the run, read BillableTimeInSeconds from the job description and compare.
Checkpoint it. Save progress to S3 as the job runs, so that a killed, interrupted or crashed run costs one epoch, not the whole bill.

Setup

Before class (instructor checklist)

Students never download the data themselves; four things have to be in place before the session. Each is written out in the instructor pages.

Stage the data in the shared read-only bucket from the lesson’s data release (badger-scribe-data-v1, ~130 MB): a few minutes from CloudShell, a laptop or a notebook instance; commands and the read-only bucket policy in Shared data bucket: staging before each term (the same bucket holds the WattBot data for the RAG episodes). The bucket must survive the end-of-workshop cleanup, or be recreated from the release next term.
Stage the base model in the same bucket (badger-scribe/model/churro-3B/, ~7.5 GB, once, from a laptop or notebook instance, not CloudShell). Every job reads the weights from S3 as a model channel instead of 100 students pulling 7.5 GB each from Hugging Face in the same minute.
Quotas: ml.g5.xlarge for training job usage at least the class size (request 100). Default is often 0 to 4 and increases take days; file three to four weeks ahead. Full list in Account preparation checklist.
Keep the bucket name ml-marathon-shared; the episodes have it filled in. If you stage under another name, tell students to change shared_bucket in the first code cell.

Fallback instance if g5 capacity is short: ml.g6.xlarge (L4, 24 GB, $1.13/h).

This episode runs from a CPU notebook instance (ml.t3.medium or ml.m5.xlarge). The notebook is only a controller: it stages data, launches jobs, and reads results. The GPU lives in the job.

There is no data upload in this episode. The competition data is published as a release on the lesson repository (three zip files: labels, train images, test images), and your instructor has staged it in a shared read-only bucket:

s3://<shared-bucket>/badger-scribe/train.csv, test.csv, metadata.csv, metric.py, sample_submission.csv, transcription_conventions.md
s3://<shared-bucket>/badger-scribe/images/<page_id>.jpg: train and test pages together, about 130 MB
s3://<shared-bucket>/badger-scribe/model/churro-3B/: a copy of the CHURRO-3B weights (~7.5 GB), so that jobs read them from S3 instead of every student pulling them from Hugging Face at once

The training job reads the images and the model straight from that bucket; the only files your notebook downloads are the small CSVs. Outputs go to your own bucket from the S3 episode, under badger-scribe/.

The first cell is the same setup as every episode: a SageMaker session, the region and the execution role, and an S3 client. Then the names this episode needs: the shared bucket the jobs read from, your own bucket they write to, your username for job names, and the two S3 prefixes that become the job’s train and model input channels.

PYTHON

import os
import json
import time

import boto3
import pandas as pd
import sagemaker
from sagemaker import get_execution_role

import helpers   # already on sys.path from the Step 0 cell; imported again so Setup runs on its own after a kernel restart

session = sagemaker.Session()
region = session.boto_region_name
role = get_execution_role()
s3_client = boto3.client("s3", region_name=region)
print("Region:", region)
print("Role:", role)

Now the names, after moving into a project folder of our own under the instance root. Everything this episode downloads or stages (data/, code/) goes there, beside the lesson clone rather than inside it. Two buckets: the shared one the jobs read from, and your own, where everything they write lands. Your username goes into job names, which SageMaker requires to be unique in the account. The two S3 prefixes at the end are what the estimator mounts into the job as its train and model input channels.

PYTHON

!mkdir -p /home/ec2-user/SageMaker/badger-scribe   # a project folder beside the lesson clone, so data and job code stay out of the repository
%cd /home/ec2-user/SageMaker/badger-scribe

shared_bucket = "ml-marathon-shared"   # the instructor-staged, read-only bucket with the data and model weights; change only on your own account
bucket_name = "sinkorswim-doejohn"     # <-- EDIT ME: your own bucket from the S3 episode; everything the jobs write goes here
username = "doejohn"                   # <-- EDIT ME: lowercase letters/digits only; used in job names

data_prefix = "badger-scribe"
train_s3_uri = f"s3://{shared_bucket}/{data_prefix}/"                  # training channel: csvs + images/
model_s3_uri = f"s3://{shared_bucket}/{data_prefix}/model/churro-3B/"  # model channel: staged CHURRO-3B weights

local_data_dir = "./data"
os.makedirs(local_data_dir, exist_ok=True)

Download the four small files (the two page lists, the metadata and the competition’s scoring script) and look at how the training pages split by category and by who transcribed them.

PYTHON

for fname in ["train.csv", "test.csv", "metadata.csv", "metric.py"]:
    s3_client.download_file(shared_bucket, f"{data_prefix}/{fname}", os.path.join(local_data_dir, fname))

train_df = pd.read_csv(os.path.join(local_data_dir, "train.csv"), keep_default_na=False)   # keep_default_na=False: an empty cell stays "" instead of becoming NaN, so every transcription is a string
test_df = pd.read_csv(os.path.join(local_data_dir, "test.csv"), keep_default_na=False)

print("train pages:", len(train_df), " test pages:", len(test_df))
print(train_df.groupby(["category", "label_source"]).size())

Not all 295 labelled pages are trained on. The script first sets 15% of them aside, by category and with a fixed seed, so the same 44 pages are held out for every run and every learner; those are the pages every CER in this episode is measured on. The 251 that remain are what the class run trains on, for two passes; the homework trains on the same 251 for four. The holdout is scored three ways: 15 pages before and after training, and 8 of those 15 at the end of every pass. Here are those numbers for this run, from the same settings the training job gets in Step 1:

PYTHON

holdout_frac, eval_pages, curve_pages = 0.15, 15, 8   # the Step 1 hyperparameters that decide the split

n_holdout = round(len(train_df) * holdout_frac)
n_pool = len(train_df) - n_holdout
print(f"{len(train_df)} labelled pages")
print(f"  held out for scoring:   {n_holdout}  ({holdout_frac:.0%}, fixed seed, never trained on)")
print(f"  available for training: {n_pool}")
print(f"  trained on:             all {n_pool}, 2 passes in class, 4 in the homework")
print(f"  scored before/after:    {eval_pages} of the {n_holdout} holdout pages; {curve_pages} of those after every pass")
print(f"test pages (no labels, scored only by Kaggle): {len(test_df)}")

Look at one page and its transcription. Note what “verbatim” means here: the writer’s spelling and punctuation, umlauts and ß kept as written, Kurrent letterforms mapped to their Latin equivalents, and # for a character the transcriber could not read.

PYTHON

row = train_df[train_df.category == "kade_letters"].iloc[3]
print(row.page_id, "|", row.category, "|", row.label_source)
print(row.text[:600])

PYTHON

from IPython.display import Image, display

page_path = os.path.join(local_data_dir, f"{row.page_id}.jpg")
s3_client.download_file(shared_bucket, f"{data_prefix}/images/{row.page_id}.jpg", page_path)
display(Image(filename=page_path, width=500))

How many tokens is a page?

The cost estimate in Step 0 used 1,600 tokens per page. That number is tokens, not words: the model never sees words. The image becomes one token per 28 × 28 pixel block once it is resized to max_pixels, and the transcription becomes pieces of words, more pieces for German and for old spellings than for modern English. Count it with the model’s own processor rather than guess it; the processor’s files are a few megabytes, so the 7.5 GB of weights stay in S3.

PYTHON

%pip install -q "transformers>=4.56"   # the processor that turns an image + text into the model's tokens; same version as the training job

Fetch the processor’s config files from the shared bucket and load the processor with the same pixel limits the training script uses. Then run one page through it twice: once with the transcription, to count everything, and once with only the prompt, to find where the transcription’s tokens start. The count splits three ways: the image, the instruction prompt the script sends with every page (the transcription rules from build_prompt plus the chat-template markup around them, the same on every page), and the transcription itself.

PYTHON

from transformers import AutoProcessor
import badger_common as bc

proc_dir = os.path.join(local_data_dir, "churro-processor")
os.makedirs(proc_dir, exist_ok=True)
for fname in ["config.json", "preprocessor_config.json", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
              "added_tokens.json", "vocab.json", "merges.txt", "chat_template.json", "video_preprocessor_config.json"]:
    s3_client.download_file(shared_bucket, f"{data_prefix}/model/churro-3B/{fname}", os.path.join(proc_dir, fname))
processor = AutoProcessor.from_pretrained(proc_dir, min_pixels=bc.DEFAULT_MIN_PIXELS, max_pixels=bc.DEFAULT_MAX_PIXELS)

Now the page you displayed above, prepared exactly as the training script prepares it: the image resized to the pixel cap, the instruction prompt built for its category, and the transcription attached as the assistant’s answer. Running the prompt alone through the processor a second time tells us how many tokens come before the transcription starts.

PYTHON

image = bc.load_page_image(page_path, bc.DEFAULT_MAX_PIXELS)   # resized exactly as the training script resizes it
prompt = bc.build_prompt(row.category)
full = processor.apply_chat_template(bc.full_messages(prompt, row.text), tokenize=False)
enc = processor(text=[full], images=[image], return_tensors="pt")
n_total = enc["input_ids"].shape[1]
n_prompt = processor(text=[processor.apply_chat_template(bc.user_messages(prompt), tokenize=False, add_generation_prompt=True)],
                     images=[image], return_tensors="pt")["input_ids"].shape[1]
n_image = int((enc["image_grid_thw"].prod(dim=1) // 4).sum())   # 14-px patches merged 2x2: one token per 28x28 block

print(f"page {row.page_id}: {image.size[0]} x {image.size[1]} px after resizing -> {n_image} image tokens")
print(f"instruction prompt (same text on every page): {n_prompt - n_image} tokens")
print(f"transcription: {len(row.text.split())} words, {len(row.text)} characters -> {n_total - n_prompt} tokens")
print(f"tokens the model sees for this page: {n_total}")

Two things to notice in the output. The image is close to the 1,280-token cap whatever the page’s original size, because every scan is larger than a megapixel and gets shrunk to it. And the transcription comes out near two tokens per word, not the 1.3 of the English rule of thumb: German, 1930s spelling and the writer’s abbreviations all split into more pieces. That is why the guess in Step 0 is a guess, and why the next cell measures the average instead.

One page is a point, not an average. Text length varies a lot from page to page; image tokens hardly vary, because almost every scan is larger than max_pixels and gets resized to the same cap. So count text tokens for every training page with the tokenizer alone, which needs no images and takes seconds, and image tokens over a sample of pages.

PYTHON

text_tokens = train_df.text.map(lambda t: len(processor.tokenizer(t)["input_ids"]))   # every page, tokenizer only

sample = train_df.sample(20, random_state=0)
image_tokens = []
for pid in sample.page_id:
    path = os.path.join(local_data_dir, f"{pid}.jpg")
    if not os.path.exists(path):
        s3_client.download_file(shared_bucket, f"{data_prefix}/images/{pid}.jpg", path)
    grid = processor.image_processor(images=[bc.load_page_image(path, bc.DEFAULT_MAX_PIXELS)], return_tensors="pt")["image_grid_thw"]
    image_tokens.append(int((grid.prod(dim=1) // 4).sum()))

Put the three parts together. The instruction prompt is the same text on every page, so its count from the single page above is reused as a constant.

PYTHON

n_instructions = n_prompt - n_image                      # the instruction prompt: the same text, so the same count, on every page
avg_text, avg_image = text_tokens.mean(), sum(image_tokens) / len(image_tokens)
avg_tokens_per_page = avg_image + n_instructions + avg_text
print(f"text tokens per page, all {len(train_df)} training pages: mean {avg_text:.0f}, min {text_tokens.min()}, max {text_tokens.max()}")
print(f"image tokens per page, {len(image_tokens)} sampled pages: mean {avg_image:.0f}")
print(f"average tokens per page: {avg_image:.0f} image + {n_instructions} instruction prompt + {avg_text:.0f} transcription = {avg_tokens_per_page:.0f}")

That average is the tokens_per_example the estimate should have used. Re-run the estimator with it and compare with the 1,600 guess:

PYTHON

helpers.estimate_finetune_cost(params_billions=3.75, examples=251, tokens_per_example=int(avg_tokens_per_page), epochs=2,
                               method="lora", instance_type="ml.g5.xlarge")

Multiply by pages and passes and that is the “tokens seen” in the cost table. Halve max_pixels and the image count halves with it; that is what Exercise 1 changes.

The scripts

Jobs run in containers, so the code they execute has to be shipped with the job. Copy the scripts from the lesson repository into a code/ directory:

PYTHON

import shutil

os.makedirs("code", exist_ok=True)
for f in ["badger_common.py", "finetune_churro_lora.py", "churro_transcribe_batch.py"]:
    shutil.copy(f"/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/{f}", f"code/{f}")

with open("code/requirements.txt", "w") as f:  # pip-installed inside the container before the script starts
    f.write("\n".join(["peft>=0.14", "accelerate", "pillow", "pandas"]))

print(os.listdir("code"))

badger_common.py: the prompt, image loading, and a copy of the competition’s CER. Shared by the other two so the model is trained on exactly the prompt it is later asked.
finetune_churro_lora.py: the training job entry point. Holds out 15% of pages per category, filters to the categories you ask for, scores a sample of holdout pages with the base model, trains a LoRA adapter, scores the same pages again, and saves everything.
churro_transcribe_batch.py: used in the homework. Transcribes any list of pages with the base model or with the adapter merged in, and writes predictions.csv in the submission format.

Open code/finetune_churro_lora.py and read it before launching it. The pieces worth understanding:

# Only the language model's projection layers get LoRA weights.
# The vision tower ("visual") stays frozen.
lora = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05, task_type="CAUSAL_LM",
    target_modules=r"^(?!.*visual).*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)$",
)

# The loss is computed on the transcription only. Prompt and image tokens
# are masked to -100 so the model is not graded on repeating the question.
labels[i, :n_prompt] = -100

And the one setting that controls cost more than any other:

PYTHON

# Every 28x28 block of the page becomes one token. 1280 blocks ~ 1 megapixel.
DEFAULT_MAX_PIXELS = 1280 * 28 * 28

PYTHON

name = "DoeJohn"         # replace with your name
project = "SinkOrSwim"   # replace with your challenge name

def make_job_tags(purpose: str):   # one of: Storage, Notebook, Training, Tuning, Inference
    return [
        {"Key": "Name", "Value": name},
        {"Key": "Project", "Value": project},
        {"Key": "Purpose", "Value": purpose},
    ]

Step 1: Launch the LoRA fine-tune

One job does everything: scores 15 held-out German-letter pages with the base model, trains on the 251 German-letter training pages for 2 passes, scoring 8 of the holdout pages at the end of each pass, scores the same 15 pages again, and saves the adapter. The hyperparameters below are what the Step 0 estimate was built from.

First the guard rails, each one a named variable so that the printout below shows exactly what the job is allowed to cost and where it saves its progress.

PYTHON

instance_type = "ml.g5.xlarge"            # one A10G GPU, 24 GB; the GPU lives in the job, not in this notebook
g5_price = 1.41                           # $/hour, us-east-1 on-demand
estimated_minutes = 32                    # the Step 0 estimate
max_run_minutes = 65                      # about twice the estimate; SageMaker kills the job at this cap, whatever it is doing
run_id = time.strftime("%m%d-%H%M")                                   # launch time, so every run gets a new name: job names must be unique in the account
training_job_name = f"badger-{username}-churro-lora-2ep-{run_id}"   # who, which model, which recipe, when; "lora" in the name because it is a LoRA fine-tune
checkpoint_s3_uri = f"s3://{bucket_name}/{data_prefix}/checkpoints/{training_job_name}/"   # the script saves the adapter here at the end of every epoch

print(f"job name:     {training_job_name}")
print(f"instance:     {instance_type} at ${g5_price}/h")
print(f"estimate:     {estimated_minutes} min  ->  about ${g5_price * estimated_minutes / 60:.2f}")
print(f"max_run cap:  {max_run_minutes} min  ->  worst case ${g5_price * max_run_minutes / 60:.2f}, then SageMaker stops it")
print(f"checkpoints:  {checkpoint_s3_uri}")
print(f"tags:         {make_job_tags('Training')}")

Callout

Four guard rails on every fine-tuning job

max_run is the only thing standing between a bug and your whole budget. A training loop that diverges, hangs on a bad batch or waits on a dead data loader bills the same as one that works. The cap turns “whatever it takes” into a number you chose: here $1.88 worst case. Set it to about twice the estimate, never leave it at the default of 24 hours, and on a $25-an-hour instance treat it as the price of the launch.
A unique job_name prevents the double launch. Running the launch cell twice fails with ResourceInUse instead of starting a second job. Without it, SageMaker generates a fresh name every time and happily runs two copies.
Checkpoints to S3 bound the wasted work. The script saves the adapter and optimizer state to /opt/ml/checkpoints at the end of every epoch, SageMaker syncs that folder to checkpoint_s3_uri while the job runs, and when a job is launched with the same checkpoint_s3_uri the script resumes from the latest checkpoint. A job killed by max_run or a crash costs at most one epoch, not the whole run. On a ten-hour run, this is the difference between losing ten minutes and losing $150.
Tags make the bill readable. Every job carries Name, Project and Purpose; without them a GPU job is anonymous spend that nobody can attribute.

Next, what the script will be told to do. These hyperparameters reach the script as command-line flags (--epochs 3 --lr 0.0002 ...), and they are what the Step 0 estimate was built from. Three of them set the cost: the number of pages, the number of passes, and max_pixels. The rest shape the training: a LoRA learning rate, the adapter rank, an effective batch of four pages through gradient accumulation, and which holdout pages get scored when.

PYTHON

hyperparameters = {
    # no model_id: the script reads the base weights from the "model" channel mounted at fit()
    "epochs": 2,                      # two passes over the 251 pages; the per-pass curve below says what each one bought
    "lr": 2e-4,                       # LoRA takes a higher learning rate than a full fine-tune (~1e-5)
    "lora_r": 16,                     # adapter rank: how many trainable parameters LoRA adds per layer
    "batch_size": 1,
    "grad_accum": 4,                  # effective batch of 4 pages -> ~125 optimizer steps over 2 passes
    "max_pixels": 1280 * 28 * 28,     # ~1 megapixel per page; tokens per page, and so cost, scale with it
    "eval_pages": 15,                 # holdout pages scored before (eval_before=1) and after training
    "eval_before": 1,
    "curve_pages": 8,                 # of those 15, scored again at the end of every epoch: the CER curve
}

Then the estimator: which script and folder to ship, which prebuilt container to run it in (the three version numbers pick a Hugging Face image with transformers, PyTorch and Python), which instance, the guard rails from above, and where the output and checkpoints go.

PYTHON

from sagemaker.huggingface import HuggingFace

model_output_path = f"s3://{bucket_name}/{data_prefix}/models/"   # the job's model.tar.gz lands under here, in your own bucket

estimator = HuggingFace(
    entry_point="finetune_churro_lora.py",
    source_dir="code/",                   # the whole directory ships with the job: script, badger_common.py, requirements.txt
    role=role,
    tags=make_job_tags("Training"),
    instance_type=instance_type,
    instance_count=1,
    volume_size=50,                       # GB of disk for the container, weights and images
    max_run=max_run_minutes * 60,         # the cap, in seconds: SageMaker stops the job here no matter what
    checkpoint_s3_uri=checkpoint_s3_uri,  # SageMaker syncs checkpoint_local_path to here while the job runs
    checkpoint_local_path="/opt/ml/checkpoints",   # where the script writes a checkpoint at the end of each epoch
    # the three versions select the prebuilt Hugging Face training image (transformers + PyTorch + Python)
    transformers_version="4.56",
    pytorch_version="2.8",
    py_version="py312",
    output_path=model_output_path,
    sagemaker_session=session,
    hyperparameters=hyperparameters,
)

Now launch. Two details in this one line matter:

job_name is the one defined with the guard rails, with the launch time in it: SageMaker requires training job names to be unique in the account, and a name that is fixed fails with ResourceInUse the second time the cell runs. Rerunning the guard-rail cell gives a new name and a new checkpoint prefix. To resume a job that was killed, keep the old checkpoint_s3_uri and launch under the new name.
wait=True (the default) blocks the notebook until the job ends and streams its log here. For a 25-minute job that is the right choice: you cannot accidentally queue a second launch while you wait, and everything worth doing in the meantime happens on the job’s page under Training jobs, not in this notebook.

The second input channel, model, is the copy of CHURRO-3B your instructor staged in the shared bucket: SageMaker mounts it at /opt/ml/input/data/model and the script loads from there, so a hundred jobs starting in the same minute copy 7.5 GB from S3 instead of each pulling it from Hugging Face.

PYTHON

# two channels: "train" is mounted at /opt/ml/input/data/train, "model" at /opt/ml/input/data/model inside the container
# the job is created with the tags given to the estimator (tags=make_job_tags("Training")); job_cost_report reads them back below
# wait=True blocks this cell and streams the log here; the job keeps running (and billing) even if you interrupt the cell
start = time.time()
estimator.fit({"train": train_s3_uri, "model": model_s3_uri}, job_name=training_job_name, wait=True)
print(f"Wall-clock from launch to finish: {(time.time() - start) / 60:.1f} min (the billed part is reported below)")

Callout

Interrupting the cell does not stop the job

The square “interrupt kernel” button stops the notebook from watching the job. The GPU instance keeps running and keeps billing until the script finishes or max_run expires. To stop a job, use the console (SageMaker AI → Training → Training jobs → select → Stop) or:

PYTHON

# session.sagemaker_client.stop_training_job(TrainingJobName=training_job_name)

Billing stops within a minute or two. If you launched with the wrong instance, the wrong data or the wrong number of epochs, stop it immediately; the cents already spent are not a reason to let it run.

While it runs

Open a second browser tab on the SageMaker console: SageMaker AI → Training → Training jobs. Find your job and watch its status move through Starting, Downloading, Training, Uploading, Completed. Click the job’s name. Its page shows the instance type, the max_run you set, the tags, and under Monitor a View logs button that opens CloudWatch, the AWS service that stores every job’s printed output; click the one stream listed, named after the job. It is the same text the notebook cell is streaming, but it stays available after the kernel is gone, which is what you will use when a job you launched yesterday fails overnight.

Things to find in the log while you wait:

pages: 295 total -> 251 train / 44 holdout: the split.
trainable params: ...: what fraction of the 3B parameters LoRA trains.
[eval before] 15 holdout pages in ...s (... s/page) -> macro CER ...: the base model’s score, and the per-page time that the homework inference jobs will cost.
The Trainer’s loss lines every 5 steps, and [eval epoch1] ... at the end of each epoch: the points of the curve you plot in Step 2.

Step 2: Read the results and the bill

When the cell returns, ask SageMaker what it charged for, and check that the job carries your tags; a job without them is spend nobody can attribute. The billable time is the number to compare against the estimate every time.

PYTHON

def job_cost_report(job_name: str, price_per_hour: float):
    desc = session.sagemaker_client.describe_training_job(TrainingJobName=job_name)
    billable = desc["BillableTimeInSeconds"]           # the bill: container pull, 7.5 GB model copy, the script, model upload
    wall = (desc["TrainingEndTime"] - desc["CreationTime"]).total_seconds()        # launch request to end
    provisioning = (desc["TrainingStartTime"] - desc["CreationTime"]).total_seconds()   # waiting for the instance: not billed
    print(f"job:            {job_name}")
    print(f"status:         {desc['TrainingJobStatus']}")
    print(f"wall-clock:     {wall / 60:.1f} min from launch to finish, of which {provisioning / 60:.1f} min provisioning (not billed)")
    print(f"billable time:  {billable / 60:.1f} min  (what you pay for; the script's own timing is in train_summary.json)")
    print(f"cost:           ${billable / 3600 * price_per_hour:.3f} at ${price_per_hour}/h")
    print(f"model:          {desc['ModelArtifacts']['S3ModelArtifacts']}")   # the model.tar.gz this job uploaded to output_path
    tags = session.sagemaker_client.list_tags(ResourceArn=desc["TrainingJobArn"])["Tags"]   # tags live on the job, not in describe_training_job
    print(f"tags:           {', '.join(t['Key'] + '=' + t['Value'] for t in tags) or 'NONE: this spend is unattributed'}")
    return desc

desc = job_cost_report(training_job_name, g5_price)
adapter_s3_uri = desc["ModelArtifacts"]["S3ModelArtifacts"]   # the model: line above, kept for the download below

Compare the billable minutes with train_minutes in the summary you download next: the difference is the fixed overhead of every job, container pull, data download and the 7.5 GB model copy from S3. It is the same whether you train for one minute or one hour, which is why a small first run is cheap and why batching work into fewer, longer jobs is cheaper than many tiny ones.

The whole output of the job is a model.tar.gz of about 100 MB: the LoRA adapter, the processor configuration, the holdout list, the evaluation files, the loss and CER curves (train_log.json) and a summary. The 3B base model is never copied anywhere.

PYTHON

import tarfile

run_dir = os.path.join(local_data_dir, training_job_name)
os.makedirs(run_dir, exist_ok=True)
bucket, key = adapter_s3_uri.replace("s3://", "").split("/", 1)   # s3://bucket/key -> bucket, key
s3_client.download_file(bucket, key, os.path.join(run_dir, "model.tar.gz"))
with tarfile.open(os.path.join(run_dir, "model.tar.gz")) as tar:   # everything the script saved to /opt/ml/model
    tar.extractall(run_dir)
print(sorted(os.listdir(run_dir)))

summary = json.load(open(os.path.join(run_dir, "train_summary.json")))   # train_pages, train_minutes, seconds_per_page, CER before/after
print(json.dumps(summary, indent=2))

Before vs. after

Both evaluations ran on the same 15 pages, so the difference is the LoRA fine-tune and nothing else.

PYTHON

before = pd.read_csv(os.path.join(run_dir, "eval_before.csv"), keep_default_na=False)   # empty predictions stay "" rather than NaN
after = pd.read_csv(os.path.join(run_dir, "eval_after.csv"), keep_default_na=False)

# cer = character error rate per page: edits / characters in the reference, lower is better
comparison = before[["page_id", "cer"]].merge(after[["page_id", "cer"]], on="page_id", suffixes=("_base", "_lora"))
comparison["gain"] = comparison.cer_base - comparison.cer_lora   # positive = the LoRA fine-tune helped on that page
print(comparison.sort_values("gain", ascending=False).round(3).to_string(index=False))
print(f"\nmean CER  base: {comparison.cer_base.mean():.3f}   LoRA: {comparison.cer_lora.mean():.3f}")

Look at the page that improved most. The LoRA fine-tuned model should have picked up the conventions (umlauts and ß as written, no modernised spelling, Kurrent s and e forms read correctly) as well as this particular writer’s hand.

PYTHON

best = comparison.sort_values("gain", ascending=False).iloc[0].page_id
ref = before.set_index("page_id").loc[best, "text_ref"]
print("REFERENCE\n", ref[:500])
print("\nBASE\n", before.set_index("page_id").loc[best, "text_pred"][:500])
print("\nTUNED\n", after.set_index("page_id").loc[best, "text_pred"][:500])

How many passes were worth paying for

Before and after says whether the LoRA fine-tune helped. The curve says whether you bought the right amount of it. The job recorded the training loss every 5 steps and scored 8 of the holdout pages at the end of every epoch, with the base model as epoch 0.

PYTHON

import matplotlib.pyplot as plt

log = json.load(open(os.path.join(run_dir, "train_log.json")))
loss = pd.DataFrame(log["loss"])
curve = pd.DataFrame(log["holdout_cer_by_epoch"])

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 3.8))
ax1.plot(loss.epoch, loss.loss, marker=".")
ax1.set(xlabel="epoch", ylabel="training loss", title="Training loss")
ax2.plot(curve.epoch, curve.macro_cer, marker="o")
ax2.set(xlabel="epoch (0 = base model)", ylabel="holdout CER", title=f"Holdout CER on {log['curve_pages']} pages")
ax2.set_xticks(curve.epoch)
plt.tight_layout(); plt.show()
print(curve.to_string(index=False))

Read the right-hand plot, not the left. Training loss almost always keeps falling; it is measured on pages the model is being fitted to. The holdout CER is the number the competition scores, and it has three shapes:

Still falling at the last epoch: another pass would probably have helped, and the time rule says what it costs (one more unit of the training row). That is the usual result on a few hundred pages, and the reason the table’s “passes” column is a choice and not a rule.
Flat: you have what this data can give; more passes buy nothing. Stop here next time and save the money.
Rising while the loss still falls: the model is memorising the training pages. Fewer passes, or more data.

A test run on 120 pages for 3 passes gave 0.246 (base), 0.086, 0.069, 0.061: a big drop on the first pass, then smaller ones, still falling at the third. The third pass bought 0.008 of CER for five minutes of GPU, which is why the class run stops at two and the homework goes to four: that is the trade the curve lets you see.

Eight pages is a noisy estimate (one hard page moves it), so read the trend, not the third decimal.

Callout

The evaluation sample is too small, on purpose

This job holds out 44 pages and scores 15 of them, 8 per curve point, to keep the run near half an hour. Those numbers are chosen for the clock, not for the statistics: one misread page moves an 8-page CER by several points. In your own work, keep most of the data for training, but score enough held-out examples that the number is not noise. Twenty is a floor. With a few hundred examples, 50 is a reasonable holdout; with thousands, 100 or more. Scoring is linear in pages (here about 18 s each), so score the curve on a smaller fixed sample each pass and the whole holdout only at the end.

Callout

How many epochs, in practice

The number of passes is fixed before launch because it is the cost cap: max_run is set from it, and the learning-rate schedule (cosine here) is laid out over it, decaying to zero at the planned end. So the common practice is not to guess a magic number but to pick a budget, score the holdout every epoch inside it, and keep the best epoch:

Small data, a few hundred to a few thousand examples: plan 3 to 5 passes. One pass rarely settles; ten usually memorises.
Large instruction-tuning sets, tens of thousands and up: 1 to 3 passes. A second pass over that much data often helps less than it costs.
Keep the best checkpoint, not the last. The script saves one per epoch; a training run that scores the holdout each epoch can reload the best (load_best_model_at_end=True in Hugging Face’s Trainer, with save_strategy and eval_strategy both set to epoch).
Early stopping (stop when the holdout has not improved for N epochs, EarlyStoppingCallback in Trainer) is the automated form. It is used less with LLMs than with small models, because each evaluation costs minutes of GPU and because stopping early leaves the cosine schedule unfinished; most teams plan the budget and pick the best epoch afterwards instead.

Whatever you do, the maximum is decided before launch. That is what makes the cost estimate a number and not a hope.

Challenge

Was the estimate right?

Compare the billable time from job_cost_report with the ~32 minutes estimated in Step 0. Which phase was off, and by how much? (train_minutes and seconds_per_page in train_summary.json tell you.)
Using your measured seconds per page, estimate the homework run: the same 251 pages × 4 passes, 40 evaluation pages before and after, four curve points, and the same startup overhead. Write the number down; you will compare it to the real bill after the homework.

Show me the solution

Expect about 32 billed minutes. Scoring is the line most likely to surprise: at 17 to 21 seconds a page, the two 15-page evaluations and the two curve points are about 14 minutes, nearly half the bill. Training is close to 1.5 s per page-pass; if you measured 3, the 502 page-passes take 25 minutes instead of 12.
At 1.5 s per page-pass: 1,004 page-passes is 25 minutes; two 40-page evaluations at ~19 s are 25; four curve points are 9; startup and upload 6: about 65 minutes and $1.50. Scale every line by the ratio you measured.

Callout

Why the holdout number flatters you

The script holds out 15% of pages, so pages from the same letter bundle appear in both the training and holdout halves. The competition’s hidden test set is split by document: the five test bundles are letters the model has never seen. All the letters are from one writer, which softens the gap here, but her hand, pen and paper change over a decade, so expect the leaderboard CER to be somewhat worse than the holdout CER. A document-level holdout is the fix; it is left as an exercise because with only four bundles in train, holding one out means giving up a quarter of the data.

Step 3: What did this episode cost?

Add it up from the job descriptions, not from memory.

PYTHON

total = 0.0
# every job this episode launched shares the badger-<username> prefix, so one listing finds them all, earlier runs under your name included
for job in session.sagemaker_client.list_training_jobs(NameContains=f"badger-{username}", MaxResults=50)["TrainingJobSummaries"]:
    d = session.sagemaker_client.describe_training_job(TrainingJobName=job["TrainingJobName"])
    cost = d.get("BillableTimeInSeconds", 0) / 3600 * g5_price
    total += cost
    print(f"{job['TrainingJobName']:<40} {d['TrainingJobStatus']:<10} {d.get('BillableTimeInSeconds', 0) / 60:5.1f} min  ${cost:.2f}")
print(f"\nGPU spend so far: ${total:.2f}  (plus the CPU notebook at $0.05/h)")

Then compare that number to the estimate in Step 0. If you are within 50%, your estimating is good enough to be trusted with a bigger run. If you are not, find out which line was wrong before launching anything else.

Homework: a four-pass LoRA run and (optionally) a Kaggle submission

Everything below is the same pattern at a larger scale: about $1.50 for the LoRA fine-tune and $0.80 for inference, roughly two hours of wall-clock time that you do not need to sit through. Launch, go do something else, come back.

The cells that launch jobs or read their output are commented out, so that Run All in class does not start two hours of GPU work. When you do the homework, uncomment them one at a time.

The four-pass LoRA run

Same estimator and the same 251 pages, four passes instead of two, 40 evaluation pages instead of 15, and a 130-minute cap. Write your estimate down first (Step 2’s challenge). Launch with wait=False here only because you will not be watching: it returns immediately and you poll later. The explicit job name still protects you from a double launch.

The same guard rails first: a new job name, its own checkpoint folder, and the cap written down with its worst case.

PYTHON

homework_job_name = f"badger-{username}-churro-lora-4ep-{time.strftime('%m%d-%H%M')}"   # a new name per launch, as above
homework_checkpoint_s3_uri = f"s3://{bucket_name}/{data_prefix}/checkpoints/{homework_job_name}/"
print(f"max_run cap: 130 min -> worst case ${g5_price * 130 / 60:.2f}; checkpoints -> {homework_checkpoint_s3_uri}")

Two settings change: four passes, and 40 evaluation pages, so the before/after numbers rest on more than 15 pages. grad_accum 8 doubles the effective batch, which suits a longer run.

PYTHON

homework_hyperparameters = {
    # no model_id: read from the "model" channel, as in class
    "epochs": 4,
    "lr": 2e-4,
    "lora_r": 16,
    "batch_size": 1,
    "grad_accum": 8,
    "max_pixels": 1280 * 28 * 28,
    "holdout_frac": 0.15,             # share of pages per category held out of training
    "eval_pages": 40,
    "eval_before": 1,
}

Same estimator as in class, except for the job name, the checkpoint folder and the 130-minute cap.

PYTHON

homework_estimator = HuggingFace(
    entry_point="finetune_churro_lora.py",
    source_dir="code/",
    role=role,
    tags=make_job_tags("Training"),
    instance_type=instance_type,
    instance_count=1,
    volume_size=50,
    max_run=130 * 60,                     # 130 minutes, in seconds: twice the estimate
    checkpoint_s3_uri=homework_checkpoint_s3_uri,
    checkpoint_local_path="/opt/ml/checkpoints",
    transformers_version="4.56",
    pytorch_version="2.8",
    py_version="py312",
    output_path=model_output_path,
    sagemaker_session=session,
    hyperparameters=homework_hyperparameters,
)

Launch and come back later; the next cell polls for it.

PYTHON

# # wait=False: returns at once; poll the job with describe_training_job in the next cell
# homework_estimator.fit({"train": train_s3_uri, "model": model_s3_uri}, job_name=homework_job_name, wait=False)
# print("launched:", homework_job_name)

PYTHON

# # Come back later and run this until it says Completed. Everything below needs the adapter this run produces.
# desc = session.sagemaker_client.describe_training_job(TrainingJobName=homework_job_name)
# print(desc["TrainingJobStatus"], "-", desc.get("SecondaryStatus"))
# homework_adapter_s3_uri = None
# if desc["TrainingJobStatus"] == "Completed":
#     desc = job_cost_report(homework_job_name, g5_price)
#     homework_adapter_s3_uri = desc["ModelArtifacts"]["S3ModelArtifacts"]

Base vs. LoRA fine-tuned over the holdout and the test set

The evaluation inside the training job covers a sample. To score the whole holdout and produce predictions for the 141 test pages, run the transcription script as a Processing job, once without the adapter and once with it. Each run transcribes about 185 pages in roughly 15 minutes. Processing jobs are introduced in the RAG part of the lesson; do this homework after that part, or treat the launch cell below as a preview: it is the same pattern as a training job, with a script, an instance type, tags and a runtime cap.

First, build the list of pages and upload it together with the metadata.

PYTHON

holdout_df = pd.read_csv(os.path.join(local_data_dir, training_job_name, "holdout_page_ids.csv"))

pages_df = pd.concat([
    train_df[train_df.page_id.isin(holdout_df.page_id)][["page_id", "doc_id", "category", "text"]],
    test_df[["page_id", "doc_id", "category"]].assign(text=""),
], ignore_index=True)
print(len(pages_df), "pages to transcribe")

The holdout rows keep their transcriptions so they can be scored; the test rows have none. Upload the list and the metadata to your bucket, and note the two S3 locations the jobs will read.

PYTHON

infer_prefix = f"{data_prefix}/inference"
pages_df.to_csv(os.path.join(local_data_dir, "pages.csv"), index=False)
for fname in ["pages.csv", "metadata.csv"]:
    s3_client.upload_file(os.path.join(local_data_dir, fname), bucket_name, f"{infer_prefix}/{fname}")

images_s3_uri = f"s3://{shared_bucket}/{data_prefix}/images/"
infer_data_s3_uri = f"s3://{bucket_name}/{infer_prefix}/"

The script this job runs is code/churro_transcribe_batch.py, the third file copied in Setup. Open it: the docstring at the top lists its inputs and outputs, main loads the model, merges the adapter into the base weights if one is given (PeftModel.from_pretrained(...).merge_and_unload()), transcribes the pages in batches, writes predictions.csv, and scores only the rows that have a reference text. The same script with and without --adapter_tar is the whole comparison.

Then one function that launches a transcription job, in four parts:

inputs: S3 prefixes copied into the container before the script starts: the images and the staged base model (both from the shared bucket; a Processing job has no model channel, so the model directory is passed to the script as --model_id), your CSVs, and, optionally, the adapter;
arguments: the script’s command-line flags, the processor’s equivalent of hyperparameters;
the processor: role, tags, instance, container versions and the runtime cap, as on the estimator;
run: which script, which inputs, and where the output is copied when the job ends.

PYTHON

from sagemaker.huggingface import HuggingFaceProcessor
from sagemaker.processing import ProcessingInput, ProcessingOutput

def run_transcription_job(label: str, adapter_uri: str = None):
    output_uri = f"s3://{bucket_name}/{data_prefix}/predictions/{label}/"
    inputs = [  # each ProcessingInput copies an S3 prefix to a container path before the script starts
        ProcessingInput(source=images_s3_uri, destination="/opt/ml/processing/input/images"),
        ProcessingInput(source=model_s3_uri, destination="/opt/ml/processing/input/model"),
        ProcessingInput(source=infer_data_s3_uri, destination="/opt/ml/processing/input/data"),
    ]
    # arguments = command-line flags for the script (a processor has no hyperparameters=)
    arguments = ["--model_id", "/opt/ml/processing/input/model", "--pages_csv", "pages.csv",
                 "--batch_size", "4", "--max_pixels", str(1280 * 28 * 28)]
    if adapter_uri:
        inputs.append(ProcessingInput(source=adapter_uri, destination="/opt/ml/processing/input/adapter"))
        arguments += ["--adapter_tar", "/opt/ml/processing/input/adapter/model.tar.gz"]

    processor = HuggingFaceProcessor(
        role=role,
        tags=make_job_tags("Inference"),
        instance_type=instance_type,
        instance_count=1,
        volume_size_in_gb=50,             # same role as volume_size on the estimator
        transformers_version="4.56",      # same image family as the training job
        pytorch_version="2.8",
        py_version="py312",
        sagemaker_session=session,
        max_runtime_in_seconds=45 * 60,   # the processor's name for max_run: 45 minutes, in seconds; the estimate is 15 minutes
    )
    processor.run(
        code="churro_transcribe_batch.py",
        source_dir="code/",
        inputs=inputs,
        # ProcessingOutput: whatever the script writes under source is copied to destination when the job ends
        outputs=[ProcessingOutput(output_name="predictions", source="/opt/ml/processing/output", destination=output_uri)],
        arguments=arguments,
        job_name=f"badger-{username}-churro-transcribe-{label}-{time.strftime('%m%d-%H%M')}",   # unique per launch, like the training jobs
        wait=False, logs=False,           # launch and return, so the two jobs run at the same time
    )
    return output_uri

Launch it twice: once with the base model alone, once with the homework adapter merged in. Not before the four-pass run has finished: the first line checks.

PYTHON

# assert homework_adapter_s3_uri, "the four-pass run has not finished; rerun the poll cell above until it says Completed"
# base_output_uri = run_transcription_job("base")
# lora_output_uri = run_transcription_job("lora", adapter_uri=homework_adapter_s3_uri)

Both jobs run in parallel on their own GPU, which doubles the instantaneous spend but halves the wall-clock time; the total cost is the same. Watch them under SageMaker AI → Data preparation → Processing jobs, and when both are Completed, score the holdout rows with metric.py, the exact file that scores the leaderboard.

PYTHON

# import subprocess
#
# def fetch_predictions(output_uri: str, label: str) -> pd.DataFrame:
#     bucket, key = output_uri.replace("s3://", "").split("/", 1)
#     local = os.path.join(local_data_dir, f"predictions_{label}.csv")
#     s3_client.download_file(bucket, key + "predictions.csv", local)
#     return pd.read_csv(local, keep_default_na=False)
#
# preds = {label: fetch_predictions(uri, label) for label, uri in [("base", base_output_uri), ("lora", lora_output_uri)]}

Write the holdout answers to a file, keep only the holdout rows of each prediction file, and run metric.py on each pair the way Kaggle runs it on a submission.

PYTHON

# holdout_solution = train_df[train_df.page_id.isin(holdout_df.page_id)][["page_id", "text", "category"]]
# holdout_solution.to_csv(os.path.join(local_data_dir, "holdout_solution.csv"), index=False)
#
# for label, df in preds.items():
#     df[df.page_id.isin(holdout_solution.page_id)].to_csv(os.path.join(local_data_dir, f"holdout_pred_{label}.csv"), index=False)
#     print(f"\n=== {label} ===")
#     print(subprocess.run(
#         ["python", os.path.join(local_data_dir, "metric.py"),
#          "--solution", os.path.join(local_data_dir, "holdout_solution.csv"),
#          "--submission", os.path.join(local_data_dir, f"holdout_pred_{label}.csv")],
#         capture_output=True, text=True).stdout)

Submit

The lora prediction file already contains every test page. Keep only those rows, in the order test.csv lists them. If you want a leaderboard score, this file is what you upload on the competition’s Submit Prediction page (two submissions per day; this is the only step that needs a Kaggle account, and it is optional).

PYTHON

# submission = test_df[["page_id"]].merge(preds["lora"], on="page_id", how="left").fillna({"text": ""})
# assert len(submission) == len(test_df) and submission.text.str.len().gt(0).all(), "every test page needs a prediction"
#
# submission_path = os.path.join(local_data_dir, "submission_churro_lora.csv")
# submission.to_csv(submission_path, index=False)
# s3_client.upload_file(submission_path, bucket_name, f"{data_prefix}/submissions/submission_churro_lora.csv")
# print("wrote", submission_path, len(submission), "rows")

Finally, rerun Step 3 and compare the total with the estimate you wrote down before the four-pass run.

Cleanup

The training and processing jobs have already ended; there is nothing to stop. Check Training jobs and Processing jobs anyway: nothing should say InProgress.
Delete what you no longer need from s3://<your-bucket>/badger-scribe/ (the models/ and predictions/ prefixes). The adapter is small, but prediction outputs accumulate across runs.
Stop your notebook instance. It is a CPU instance at five cents an hour, which is exactly the kind of charge that goes unnoticed for a semester.

Bonus exercises

Not part of the class session. Exercise 1 relaunches the class job with one change and takes about an hour of GPU time; Exercise 2 is arithmetic only.

Challenge

Exercise 1: Image resolution

Relaunch the same job with "max_pixels": 2560 * 28 * 28 (~2 megapixels, twice the 1280 * 28 * 28 of the class run) and the job name badger-{username}-churro-lora-2ep-2mp. Before it starts, write down your prediction for its billable time, and raise max_run to match: the 65-minute cap was set for the 1-megapixel job. Then compare train_minutes and seconds_per_page in the two summaries, and the before/after CER. Did the extra pixels help?

Show me the solution

Time per page roughly doubles: the sequence is twice as long and attention cost grows faster than linearly. On these scans the strokes are legible well below 1 megapixel, so the extra pixels mostly buy tokens that carry no information, and CER usually moves very little. Resolution is the first thing to lower when a fine-tune is over budget, and the first to raise if the model is misreading small characters such as the ½ fractions.

Challenge

Exercise 2: Cheaper per run, or just faster?

ml.g6e.xlarge (one L40S, 48 GB) costs $2.61/h, 1.85x the price of ml.g5.xlarge, and runs this workload roughly 2.5–3x faster. It also has twice the GPU memory, so batch_size can be 2 or 4 instead of 1, which uses the GPU more efficiently.

Using the Step 0 “The run” table, estimate the cost of the class job on ml.g6e.xlarge. Then do the same for the homework run (25 min of training, 34 min of scoring, 6 min of startup and upload on the g5). Which phase does not get faster, and when does the faster instance save money?

Show me the solution

Startup and upload (~6 min, $0.26 at $2.61/h) do not speed up: the container pull and the S3 copies are the same on any instance. For the class job, training drops from ~12 to ~4 minutes and the ~14 minutes of scoring to ~5, so the job is about 15 minutes and ~$0.65: a little less than the $0.75 g5 run. At this size the two are nearly a wash, because the fixed overhead is a third of the faster job.

The homework run on the g5 is 25 + 34 + 6 = 65 minutes, about $1.50. On the g6e at 2.5–3x: ~9 minutes of training, ~12 of scoring, the same 6 of overhead, about 27 minutes and ~$1.20. The longer the GPU work, the more the faster card saves: the g5 bill grows linearly with pages and passes while the g6e bill grows at a third of the slope. The g6e only loses on a job so short that the fixed overhead dominates, a 5-minute smoke test, say.

The hourly rate is the wrong thing to compare; the cost per run is what matters, and the fixed overhead is a larger share of it on faster hardware. The catch is quotas: the account needs a service quota for ml.g6e.xlarge training jobs, and those often default to zero.

Key Points

Fine-tuning cost is GPU-hours × hourly rate, and GPU-hours scale with model size × tokens × epochs × attempts. The same recipe runs from under a dollar (3B, LoRA, a few hundred pages) to tens of thousands (70B, full fine-tune, days on multi-GPU nodes). Starting from a model that already knows most of the task is the biggest saving available.
Training jobs, not GPU notebooks, are the safe way to run it: they end, max_run guarantees it, a unique job_name prevents a second copy launching by accident, and checkpoints synced to S3 mean a killed job costs an epoch rather than the run. Interrupting a notebook cell does not stop a job.
Estimate before launching (startup overhead + pages × epochs × seconds per page, at the instance’s rate), then compare to BillableTimeInSeconds.
Image resolution (max_pixels) and epochs change cost by multiples; the instance size changes it by 5–45x, which is right when the model does not fit on one GPU and waste when it does.
Score the holdout at the end of every epoch: the curve, not a rule of thumb, says how many passes were worth paying for. Measure the gain on the same held-out pages before and after, with the task’s own metric, and remember that a page-level holdout is optimistic when the real test is split by document.

Content from Overview of RAG Workflows on AWS

Last updated on 2026-10-21 | Edit this page

Estimated time: 15 minutes

Retrieval-Augmented Generation (RAG) on AWS

Retrieval-Augmented Generation (RAG) is a pattern where you retrieve relevant context from your data and then generate an answer using that context. Unlike model training, a standard RAG workflow does not fine‑tune or train a model — it combines retrieval + inference only.

This episode introduces the major ways to build RAG systems on AWS, explains why Amazon Bedrock should be your default, and prepares us for later episodes where we experiment with each approach.

Overview

Questions

What is Retrieval‑Augmented Generation (RAG)?
What are the main architectural options for running RAG on AWS?
Why is Bedrock the recommended default, and when should you reach for SageMaker instead?
How do the billing models differ, and what does it cost if you forget to shut something down?

Objectives

Understand that RAG does not require training or fine‑tuning a model.
Recognize the major architectural patterns for RAG systems on AWS.
Explain why pay‑per‑token Bedrock inference is the safe default for most research RAG work.
Identify the specific situations where self‑managed GPU compute (Processing Jobs, notebook GPUs, or endpoints) is actually worth it.
Understand the core cost and control trade‑offs that drive which approach to use.

What is RAG?

Diagram of the RAG pipeline with the embedder and generator highlighted as the two model calls

RAG combines two steps:

Retrieve: Search your document store (vector DB or FAISS index) to find relevant text.
Generate: Provide those retrieved chunks to a large language model (LLM) to answer a question.

No model weights are updated. No backprop. No training job.
RAG is an inference‑only pattern that layers retrieval logic around an LLM.

This matters for how you should think about compute. The only things a RAG pipeline needs from a model are two stateless calls:

“Turn this text into an embedding vector.”
“Given this prompt (question + retrieved chunks), generate an answer.”

Neither of those requires you to own a GPU. They are exactly the kind of calls that a managed inference API is built to serve.

Start Here: Bedrock Is the Default

For standard RAG inference, use Amazon Bedrock. Bedrock exposes embedding and generation models (Amazon Titan, Anthropic Claude, DeepSeek, Meta Llama, and others) as pay‑per‑token APIs. Your retrieval code — chunking, indexing, similarity search, reranking, evaluation — stays entirely under your control and can run on a cheap CPU notebook or your laptop. Only the two model calls above are sent to Bedrock.

The reason this is the default comes down to how you are billed:

	Bedrock (managed API)	SageMaker GPU instance / endpoint
Billing unit	Per token processed	Per hour the instance exists
Cost while idle	`$0` — nothing is running between calls	Full hourly rate, 24/7, whether or not you use it
What you must remember to do	Know the per-token price before a big batch; route calls through a tagged profile	Stop or delete the instance every time you finish
Model updates	Managed by AWS / the provider	Your problem (containers, CUDA, VRAM)
Model choice	Curated catalog (incl. proprietary models)	Any open‑weight model you can load

A GPU instance is a rental that keeps the meter running until you turn it off. Forgetting to do so is the single most common way research teams burn through cloud budgets. Bedrock has no equivalent failure mode: if you stop calling it, you stop paying.

What forgetting actually costs

To make this concrete, here is what a GPU notebook instance costs if it is left running after you close your laptop (on‑demand prices at time of writing; see the instances page):

Instance	GPU	`$`/hour	Forgotten overnight (12 h)	Forgotten over a weekend (64 h)	Forgotten for a month (730 h)
ml.g4dn.xlarge	1 × T4 (16 GB)	`$0.74`	~`$9`	~`$47`	~`$540`
ml.g5.xlarge	1 × A10G (24 GB)	`$1.41`	~`$17`	~`$90`	~`$1,030`
ml.g5.2xlarge	1 × A10G (24 GB)	`$1.52`	~`$18`	~`$97`	~`$1,110`
ml.g6e.xlarge	1 × L40S (48 GB)	`$2.61`	~`$31`	~`$167`	~`$1,905`
ml.p4d.24xlarge	8 × A100 (40 GB)	`$25.25`	~`$303`	~`$1,616`	~`$18,430`

Compare that to running the entire WattBot pipeline from this lesson on Bedrock. Using DeepSeek V3 (check the Bedrock pricing page for current per‑token prices) and Titan Text Embeddings V2 ($0.02 per million tokens):

Step	Rough volume	Cost
Embed ~9,000 chunks (~3M tokens)	3M input tokens	~`$0.06`
Answer and explain the 245 training questions (two calls each, ~1.5K‑token prompts) with DeepSeek V3	650K in / 40K out	~`$0.50`
Total		about `$0.55`

A single forgotten weekend on the cheapest GPU instance costs more than fifty full runs of the pipeline on Bedrock; a forgotten weekend on the ml.g6e.xlarge the optional GPU-notebook episode uses costs more than a hundred. That asymmetry is why Bedrock is the default, not just an alternative.

Callout

Bedrock trade-offs to know about

Bedrock is not free of downsides; it’s just that none of them are “you forgot and it kept charging you.”

Less control over architecture. You can only call models in the Bedrock catalog. If your research depends on a specific open‑weight model, a custom fine‑tune, or hacking on the model internals, you’ll need to host it yourself.
Per‑token cost can exceed GPU cost for very large offline batches. If you are embedding tens of millions of chunks with a small open model, a short‑lived GPU job may be cheaper per unit of work. Do the arithmetic before assuming either way.
Model access and region availability. Some models must be enabled in the console first, and the catalog varies by region. Plan for this before a workshop or deadline.
Quotas. On‑demand throughput is rate‑limited per model and per account. For large batch workloads, use the Bedrock batch inference API or request a quota increase rather than hammering the on‑demand endpoint.
Data handling. Bedrock does not use your prompts to train models and keeps data in‑region, but you should still check your institution’s policy before sending sensitive documents to any hosted API.
Cost tracking needs one extra step. On‑demand Bedrock calls carry no tags. To attribute spend to a project you route calls through a tagged application inference profile (see “Tagging” below and the Bedrock episode). It is a one‑time setup, but skipping it leaves the usage anonymous on the bill.

Approaches to Running RAG on AWS

With that framing, here are the main patterns, ordered from “reach for this first” to “only when you really need it.”

Comparison table of four ways to run RAG on AWS: Bedrock API (default, pay per token, nothing to shut down), Processing job (batch with open models, bills per second while the job runs, terminates itself), GPU notebook (learning and debugging only, bills per hour until stopped, about $90 if forgotten over a weekend) and Inference endpoint (steady traffic only, bills 24/7, about $1,000 if idle for a month).

1. Amazon Bedrock — managed embedding and generation APIs (default)

Call embedding and generation models via API from your RAG pipeline. You still own and manage retrieval logic (chunking, indexing, reranking, evaluation), but you outsource the heavy model lifecycle work entirely. You pay per token, nothing runs between calls, and there is nothing to shut down. Bedrock also gives RAG systems access to proprietary models (e.g., Claude) that would otherwise need to be purchased and integrated separately.

Use it for: essentially any RAG system where a model in the Bedrock catalog is acceptable — prototyping, research tools, course projects, and production chatbots alike.

2. SageMaker Processing Jobs — short-lived batch GPU compute

For large corpora, or when you need a specific open‑weight model that Bedrock doesn’t offer, treat parts of the RAG pipeline — especially embedding and batch generation — as offline batch jobs rather than a live model. You run a short‑lived Hugging Face Processing job that spins up a GPU instance, loads your model, processes all the chunked text in one shot, saves the results to S3, and then terminates itself. Because the GPU only exists while the job runs, you get self‑managed models without the forgot‑to‑shut‑it‑down risk. This is the safest way to use a GPU on SageMaker.

The trade‑off is latency: starting a Processing job takes several minutes, so this pattern is for “compute once, use many times” workloads, not per‑query retrieval. Launching a job per user question would be far too slow.

Use it for: embedding very large corpora with an open model, periodic batch regeneration, or research that requires a model you can’t get through Bedrock.

3. A single GPU-backed notebook instance — for learning and debugging

For small‑ to medium‑sized models (< 70 B), you can pick a GPU instance (e.g., ml.g5.xlarge), load your embedding and generation models directly in the notebook, and run RAG end‑to‑end there. The architecture is the simplest to understand, and it’s genuinely the best way to learn what happens inside a RAG pipeline, because every step is a line of Python you can inspect.

But it is the most expensive way to run RAG and the easiest one to get wrong. The GPU bills by the hour from the moment the instance starts until you stop it, including the long stretches while you download PDFs, chunk text, read results, or go to lunch. It is also the route most likely to be forgotten, because a notebook instance looks like a workspace rather than a running job. If you use this approach:

Treat it as a learning exercise, not a deployment pattern.
Stop the instance the moment you are done, and verify the status on the Notebook instances page.
Set up an AWS Budget alert and consider a lifecycle configuration that auto‑stops idle notebooks.
For larger models the hourly rate (e.g., $25/hour for a p4d.24xlarge) becomes a limiting factor very quickly.

Use it for: walking through RAG internals step by step, debugging a model locally before packaging it as a job, or short experiments where you will be at the keyboard the whole time.

4. Long-lived SageMaker inference endpoints — self-hosted interactive serving

For applications that need low‑latency, interactive RAG (APIs, chatbots, dashboards) and that must use a model you host yourself, you can deploy your own embedding and generation models as SageMaker inference endpoints and call them from your retrieval service. This gives you full control over the model, scaling policies, and autoscaling.

It is also the most expensive option when traffic is low or bursty, because you are paying for capacity that stays online even when nobody is querying the system — the same always‑on problem as a notebook GPU, but now intentionally and indefinitely. Before choosing this route, ask whether Bedrock already serves an acceptable model: it provides exactly this interactive, low‑latency experience with none of the idle cost. Endpoints are the right answer only when the model itself is the reason Bedrock won’t work.

Use it for: production systems with steady traffic that require a custom or open‑weight model. We do not build one hands‑on in this lesson.

Tagging: Every Billable Resource Needs Its Own Tags

Cost tracking in this workshop relies on three tags — Name, Project, Purpose — so that Cost Explorer can break the shared bill down by challenge. The part that trips people up is that tags never propagate. Tagging your notebook instance tags the notebook instance and nothing else. Every other resource you create is a separate line item that is anonymous unless you tag it yourself at creation time.

Resource	What it bills	How it gets tagged
Notebook instance	Hourly, while running	Tags you set in the console when creating it
S3 bucket	Storage + requests	Tags you set in the console when creating it
Training / tuning job	Hourly, while the job runs	`tags=` argument on the estimator or tuner, every job
Processing job	Hourly, while the job runs	`tags=` argument on the processor, every job
Inference endpoint	Hourly, while deployed	`tags=` argument on `deploy()`
Bedrock on‑demand call	Per token	Cannot be tagged directly. Create an application inference profile with tags and pass its ARN as `modelId` instead of the model ID

Three things follow from this:

Each job gets its own tags. A SageMaker job launched from a tagged notebook does not inherit the notebook’s tags. Pass tags=job_tags on every estimator, tuner, and processor, exactly as the training, tuning, and Processing‑job episodes do. Vary Purpose per job so the bill tells you what the money went to.
Bedrock is tagged through a profile, not a call. invoke_model and converse have no tags parameter. An application inference profile is a free, model‑specific wrapper that carries tags; usage routed through it is attributed to those tags. One profile per base model (so one for the embedding model and one for the generation model). Note that Bedrock’s tag format is lowercase key/value, while SageMaker’s is Key/Value; the tag names stay the same.
Only the three approved keys count. Name, Project and Purpose, spelled exactly; any other key is ignored by the bill, for the reason the S3 episode gives.

When Do You Use Which Approach?

Your situation	Recommended route
You need to embed and query a document collection and a catalog model (Titan, Claude, Llama, …) is fine	Bedrock
You want an interactive chatbot or API over your documents	Bedrock
You need a specific open‑weight model and the work is batch (embedding a corpus, bulk generation)	Processing Jobs
You are learning how RAG works, or debugging a model step by step	Notebook GPU — and stop it when you leave
You need interactive serving of a model Bedrock doesn’t offer, with steady traffic	Inference endpoint

If you are unsure, pick Bedrock. The worst case is a slightly higher per‑token bill; the worst case with the other routes is a four‑figure surprise on next month’s invoice.

Callout

How the hands-on episodes are ordered

The next two episodes walk through RAG on Bedrock first, then on Processing Jobs. The Bedrock episode is the one to start from: it builds the whole WattBot pipeline, including chunking the corpus, from a cheap CPU notebook with nothing left running afterward. The Processing Jobs episode reruns the same pipeline on a self-hosted open-weight model that Bedrock does not carry, as jobs that end on their own, so you can see what Bedrock abstracts away and what it costs to do yourself. An optional self-study episode after them runs the same self-hosted pipeline interactively on a GPU notebook instance. Those first two are the defaults. The GPU-notebook episode is for the case where you need to poke at a self-hosted model interactively, and only if you will remember to stop the instance. All three RAG episodes share one module of pipeline code (scripts/wattbot_rag.py in the lesson repository).

When you build your own RAG system afterward, start from the Bedrock episode.

Key Points

RAG is an inference‑only workflow: no training or fine‑tuning required, so you do not need to own a GPU to run one.
Amazon Bedrock is the default for RAG inference: you pay per token, nothing runs between calls, and there is nothing to forget to shut down.
Self‑managed GPU compute (Processing Jobs, notebook GPUs, inference endpoints) is justified when you need a model Bedrock doesn’t offer or when very large batch volumes make per‑token pricing uncompetitive — not as a starting point.
Among the GPU options, Processing Jobs are the safest because the instance terminates itself; notebook GPUs and endpoints bill by the hour until you stop them.
A forgotten GPU instance costs tens of dollars overnight and hundreds to thousands over a month; a full WattBot run on Bedrock costs well under a dollar.
Tags never propagate: tag every job at launch, and route Bedrock calls through a tagged application inference profile, or the spend is untraceable.
The hands‑on episodes start with Bedrock, then show the same pipeline on Processing Jobs; the notebook-GPU version is optional self-study.

Content from RAG with Bedrock

Last updated on 2026-10-21 | Edit this page

Estimated time: 50 minutes

Overview

Questions

How do we build a complete RAG pipeline on AWS without provisioning a GPU?
How do we turn a collection of PDFs into retrievable, citable chunks stored in S3?
How do we call Bedrock-hosted embedding and generation models so that the spend is tagged to our project?

Objectives

Read the WattBot metadata, questions and PDF corpus from the shared bucket, and cache what you derive from them (chunks.jsonl) in your own bucket for the later episodes.
Convert PDFs into page-level records and overlapping text chunks.
Create tagged application inference profiles so that every Bedrock call is attributed to your project.
Embed chunks with Amazon Titan Text Embeddings and retrieve relevant context with cosine similarity.
Generate WattBot-format answers and explanations with a Bedrock-hosted model (DeepSeek V3) and score them against train_QA.csv.
Build a second index from page images with a multimodal embedding model and measure what it changes for the table and figure questions.

Open the notebook from: /ML_with_AWS_SageMaker/notebooks/11_RAG-Bedrock.ipynb. Select the conda_pytorch kernel. This is the first hands-on RAG episode, and it starts where the overview told you to start: Amazon Bedrock. We build a complete RAG pipeline for the WattBot 2026 challenge from a small, CPU-only SageMaker notebook. Nothing in this episode needs a GPU: the core model work — both text embeddings and answer generation — runs on Bedrock-hosted models that bill per token and leave nothing running between calls. We’ll use:

an Amazon Titan Text Embeddings V2 model to turn WattBot chunks into vectors, and
a DeepSeek V3 model hosted on Bedrock to generate answers and explanations.

We also do the corpus preparation here: unpack the WattBot PDF corpus, split it into pages and overlapping chunks, and cache the chunks on S3. The Processing Jobs episode (and the optional GPU-notebook episode) reuse those chunks and run the same retrieval, evaluation, and scoring logic on self-hosted models, so you can compare cost, latency, and complexity directly.

Why Bedrock for WattBot?

When you self-host models (as the Processing Jobs episode does next), you are responsible for picking a model, managing versions, and making sure your instance has enough VRAM. That’s fine for experiments, but it can get painful once multiple teams or challenges want to reuse the same pipeline.

Running your embedding + generation steps on Amazon Bedrock gives you a few nice properties:

Managed, up-to-date models. You can use high‑quality models from Anthropic, Amazon, DeepSeek, and others without worrying about container images or CUDA versions.
Pay for what you use (in tokens). Instead of paying for a GPU instance that might sit idle, you pay per token (input + output) when you call the model. For some workloads this is cheaper; for large offline batches with smaller models, a dedicated GPU can still win.
Easier sharing and governance. It’s easier to standardize on a small set of Bedrock models across courses, hackathons, or labs than to manage many separate GPU instances.

In this notebook, both the embedding and answer/explanation steps run on Bedrock-hosted models. The embeddings are saved as embeddings_bedrock.npy so you can compare them against the self-hosted run in the Processing Jobs episode.

Models used in this episode

We’ll work with Amazon Bedrock–hosted foundation models for both embedding and generation:

Amazon Titan Text Embeddings V2 (amazon.titan-embed-text-v2:0)
- General‑purpose text‑embedding model for semantic search, retrieval, clustering, and classification.
- Supports configurable embedding dimensions (256, 512, or 1,024) and has presets tuned for retrieval or binary indexing.
- AWS does not publish the exact number of parameters for Titan models; you can treat it as a modern transformer specialized for embeddings rather than free‑form text generation.
DeepSeek V3 (deepseek.v3-v1:0 via Bedrock)
- An open‑weight mixture‑of‑experts model (671B parameters in total, about 37B active per token) that Bedrock serves on‑demand, so there is no GPU to manage.
- Particularly useful when you want many calls (e.g., one per question) and care about low latency and lower per‑token pricing compared to flagship models (check the Bedrock pricing page for current rates).
- The same code works for any chat model in the catalog. Bedrock has one chat call, converse, that takes the same input for every provider (a system prompt and a list of user and assistant messages) and returns the reply in the same shape, so switching base_generation_model_id to a Claude, Llama, or Mistral model ID changes nothing else.
(Optional) Multimodal models for tables and figures
- Bedrock also exposes multimodal models that can reason over images, charts, and document layouts (for example, Claude models with vision, or Amazon Titan Multimodal Embeddings). These are a good fit if much of your evidence lives in figures, tables, or scanned PDFs.
- To use them from Bedrock you send both text and image content in a single request:
  - Pre‑process PDFs by rendering pages (or cropping individual tables/figures) to images using a tool like pdf2image or a headless browser.
  - Base64‑encode those images and include them as image parts alongside text in the model request.
  - For multimodal embeddings, you call a Titan multimodal embedding model with an inputImage (and optionally inputText) payload to obtain a single vector that mixes visual and textual information.
- This notebook stays with text‑only embeddings + generation to keep the workflow simple, but the same RAG pattern extends naturally to multimodal models once you add an image‑extraction step to your preprocessing pipeline.

For the full list of available models (Claude variants, Amazon models, and partner models), open the Model catalog in the Amazon Bedrock console. Each entry has a model card with capabilities, typical use cases, and pricing, for exploring alternatives for your own RAG system.

Callout

Choosing the two models

The catalog tells you what exists and what it costs; it does not tell you what is good. Three places do, and all three are free to read:

Embedding models: the MTEB leaderboard scores embedding models on retrieval benchmarks. Filter to the retrieval task and your language; the top of the table changes every few months, and Titan V2 is a sound middle-of-the-pack choice that Bedrock serves without setup.
Generation models: LMArena ranks chat models by blind human preference, and Artificial Analysis plots quality against price per million tokens, which is the chart that matters for RAG: the answer is written from retrieved text, so a mid-priced model often reads as well as a flagship at a tenth of the cost.
Your own questions: the benchmarks rank models on other people’s data. train_QA.csv has labelled questions for this corpus, and the evaluation at the end of this episode scores a model on them and prices the run from its token counts. Changing base_generation_model_id and rerunning is the whole experiment. A model that wins on a leaderboard and loses on your questions has told you something the leaderboard cannot.

The habit is the same as for instances: pick the cheapest candidate that scores well enough, measure it on your data, and move up only when the measurement says to.

Notebook + dataset setup

For this episode you need a SageMaker notebook instance, but not a GPU one. A small CPU instance such as ml.t3.medium is enough, because the models run on Bedrock rather than in the notebook. The notebook’s execution role must allow Bedrock access, which a default SageMaker role does not. On the class account both roles the console offers carry it. On your own account, attach the Bedrock policy from the Account preparation checklist instructor page to your role before starting; the first Bedrock call fails with an access-denied error naming the role until it is there.

See Instances for ML for further guidance on instance types.

Two buckets, no uploads

This episode reads from two buckets and you create neither of them here:

The shared bucket, ml-marathon-shared, already set in the first cell below, holds the WattBot inputs under wattbot/: metadata.csv (the index of papers), train_QA.csv (labelled questions and answers), corpus.zip (the PDFs) and pages.jsonl (the text of every PDF page, extracted once). It is read-only and there is one copy for the whole class.
Your bucket from the S3 episode (challenge-lastnamefirstname) is where this episode writes what it derives, under a wattbot/ prefix: the chunked text (wattbot/chunks.jsonl), the embeddings (wattbot/embeddings/) and the predictions (wattbot/solutions/). Its tags cover its storage cost; your notebook instance carries the tags you gave it when you created it, and in this episode those two resources plus the per-token Bedrock calls (tagged through an inference profile, see “Cost tracking” below) are the only things being billed.

Callout

Working through the lesson on your own?

If there is no instructor-staged bucket, stage the three input files in your own bucket under the same keys and point both variables at it: metadata.csv and train_QA.csv are in data/data.zip in the lesson repository, corpus.zip is built once by scripts/build_wattbot_corpus.py, and instructors/shared-data-bucket.md has the upload commands. Then set shared_bucket = bucket_name in the first cell below and everything else works unchanged.

Step 1 – Configure the notebook

Three things are new here compared with the training episodes, so the setup is split into four cells.

The session, region and S3 client are as before, plus the execution role, printed because the Bedrock calls below run as it.

PYTHON

import os
import json
import time
from typing import Dict, Any, List
from pathlib import Path

import boto3
from botocore.config import Config
from botocore.exceptions import ClientError
import pandas as pd
import numpy as np
import sagemaker

session = sagemaker.Session()
region = session.boto_region_name
s3 = boto3.client("s3", region_name=region)
role = sagemaker.get_execution_role()
print("Role:", role)   # the execution role: on the class account it carries the Bedrock permissions; on your own, it must

The models. RAG needs two: an embedding model, which turns a piece of text into a vector so that similar passages land near each other, and a generation model, the chat model that writes the answer from the passages it is given. Both are catalog models on Bedrock, named by a model ID; the model catalog in the console lists everything available in the region with its ID and price, and the supported models page in the documentation is the same list as a table. Prices are per million tokens on the Bedrock pricing page: in October 2026, DeepSeek V3 on demand is $0.62 in and $1.85 out, and Titan Text Embeddings V2 is $0.02 in, which is why the whole WattBot run below costs cents. We write the IDs down now but will not call them by these IDs directly: the cost-tracking step below wraps each one in a tagged inference profile first, because that is the only way Bedrock spend gets attributed to a project.

PYTHON

# Base model IDs. Any chat model in the Bedrock catalog works as the generator (a Claude or Llama ID, say),
# because converse, Bedrock's chat call, takes the same messages and returns the same shape for every provider.
base_generation_model_id = "deepseek.v3-v1:0"             # writes the answers; $0.62 in / $1.85 out per million tokens (Oct 2026)
base_embedding_model_id = "amazon.titan-embed-text-v2:0"  # turns text into vectors for retrieval; $0.02 per million tokens

The files. The inputs live in the shared bucket: the paper list, the questions with their answers, and the zipped corpus of PDFs. Everything this episode derives from them goes to your own bucket, under the same wattbot/ prefix: the chunked text, the chunk embeddings, and the answers the pipeline produces.

PYTHON

!mkdir -p /home/ec2-user/SageMaker/wattbot   # a project folder beside the lesson clone, so data and outputs stay out of the repository
%cd /home/ec2-user/SageMaker/wattbot

shared_bucket = "ml-marathon-shared"   # the instructor-staged, read-only bucket with the data and model weights; change only on your own account
bucket_name = "sinkorswim-doejohn"     # <-- EDIT ME: your own bucket from the S3 episode

# inputs, in the shared bucket
metadata_key = "wattbot/metadata.csv"      # one row per paper: id, title, url
train_key = "wattbot/train_QA.csv"         # the questions, with reference answers for scoring
corpus_key = "wattbot/corpus.zip"          # the PDFs, one per paper

# derived, written to YOUR bucket
chunks_key = "wattbot/chunks.jsonl"                                # the corpus cut into passages (end of the chunking step)
embeddings_key = "wattbot/embeddings/embeddings_bedrock.npy"       # one vector per chunk
solutions_key = "wattbot/solutions/wattbot_solutions_bedrock.csv"  # the pipeline's answers, in the competition's format

local_data_dir = "./data"              # working copies on this notebook instance
os.makedirs(local_data_dir, exist_ok=True)

Two Bedrock clients. Bedrock has two APIs. bedrock-runtime is the one that calls models (converse, the chat call, for answers; invoke_model, the raw call with a provider-specific body, for embeddings); it is what every answer and every embedding in this episode goes through. bedrock is the control plane: listing models, creating the inference profiles that carry our tags. The runtime client gets a retry setting because on-demand Bedrock throughput is rate-limited per model: a loop of a few thousand embedding calls will occasionally be throttled, and the adaptive mode backs off and retries instead of raising.

PYTHON

bedrock_runtime = boto3.client(
    "bedrock-runtime",
    region_name=region,
    config=Config(retries={"max_attempts": 10, "mode": "adaptive"}),   # back off and retry when throttled
)
bedrock = boto3.client("bedrock", region_name=region)                  # control plane: list models, create inference profiles

Step 2 – Get the shared WattBot helpers

The retrieval, prompting, answer parsing and scoring logic of the WattBot pipeline is the same whether the models run on Bedrock, in a Processing job or on a notebook GPU. It lives once in the lesson repository, in scripts/wattbot_rag.py, and every RAG episode imports it; the only code that differs per episode is the two model calls. Clone the repository and copy the module next to this notebook (the Processing Jobs episode ships the same file inside its jobs).

PYTHON

!cp /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/wattbot_rag.py /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/Score.py .

PYTHON

from wattbot_rag import (
    ANSWER_SYSTEM_PROMPT,
    build_docid_to_url,
    compute_wattbot_score,
    retrieve_top_k,
    run_single_qa,
)

Open wattbot_rag.py in the Jupyter file browser as you go: each helper is short, and the docstring at the top says which episode uses what.

Step 3 – Load the metadata and training questions

Pull the two CSV files from the shared bucket into the notebook and inspect them. Two helpers first: download_from_s3 copies one object from the shared bucket (or any bucket you name) into the local data folder, and smart_read_csv reads a CSV by trying a few encodings in turn, because author names and titles in metadata.csv are not always valid UTF-8.

PYTHON

import zipfile


def download_from_s3(key: str, local_name: str, bucket: str = shared_bucket) -> str:
    """Download a file from S3 (the shared bucket by default) to local_data_dir."""
    local_path = os.path.join(local_data_dir, local_name)
    print(f"Downloading s3://{bucket}/{key} -> {local_path}")
    s3.download_file(bucket, key, local_path)
    return local_path


def smart_read_csv(path: str) -> pd.DataFrame:
    """Read a CSV, trying utf-8 and then the common Windows and Latin encodings."""
    encodings = ["utf-8", "latin1", "ISO-8859-1", "cp1252"]
    last_error = None
    for enc in encodings:
        try:
            return pd.read_csv(path, encoding=enc)
        except Exception as e:
            last_error = e
    if last_error is not None:
        raise last_error
    raise RuntimeError(f"Unable to read CSV at {path}")

Now the two files: train_QA.csv is the labelled questions and answers we score against, metadata.csv is one row per paper with its id, title and URL.

PYTHON

train_qa_path = download_from_s3(train_key, "train_QA.csv")      # labelled questions + answers: what we score against
metadata_path = download_from_s3(metadata_key, "metadata.csv")   # one row per paper (id, title, url): the corpus index

train_df = smart_read_csv(train_qa_path)
metadata_df = smart_read_csv(metadata_path)

print("train_QA.csv columns:", train_df.columns.tolist())
print("metadata.csv columns:", metadata_df.columns.tolist())
print("\nNumber of training QAs:", len(train_df))
print("Number of metadata rows:", len(metadata_df))

train_df.head(15)

Step 4 – Get the PDF corpus

The corpus is the set of PDFs listed in metadata.csv, one per paper, named <id>.pdf. Downloading them from the publishers’ sites is slow and unreliable when fifty notebooks do it at once, so the instructors did it once with scripts/build_wattbot_corpus.py and staged the result as corpus.zip in the shared bucket. We download that one file, unpack it, and delete the zip: the instance’s disk is 20 GB at most, and the zip plus the unpacked PDFs are close to 700 MB.

PYTHON

PDF_DIR = os.path.join(local_data_dir, "pdfs")
os.makedirs(PDF_DIR, exist_ok=True)

corpus_path = download_from_s3(corpus_key, "corpus.zip")
with zipfile.ZipFile(corpus_path, "r") as zf:
    zf.extractall(PDF_DIR)
os.remove(corpus_path)   # the zip is 270 MB and the unpacked PDFs are 400 MB; the notebook disk is small

pdf_files = sorted(f for f in os.listdir(PDF_DIR) if f.lower().endswith(".pdf"))
print(f"{len(pdf_files)} PDFs in {PDF_DIR}; metadata lists {len(metadata_df)} papers")

Callout

Where the corpus came from

metadata.csv has a url column pointing at each paper. scripts/build_wattbot_corpus.py in the lesson repository walks that column, downloads every PDF, zips them, and then extracts the text of every page into pages.jsonl; it reports any URL that failed so the instructor can decide what to do about it. If you are working on your own, run it once and upload both files to your bucket under wattbot/ (with shared_bucket = bucket_name, as in the callout above). That is the whole difference between “the data” and “a copy of the data”: the corpus is built once and read many times.

Step 5 – Turn PDFs into page-level “documents”

The next unit of work is a list of page-level records, one per PDF page. The instructor’s staging script already built it for this corpus (pages.jsonl in the shared bucket), so below we define the function that does the extraction, then load the staged result rather than rerun it. Each record stores:

text: page text (as extracted by pypdf).
doc_id: short ID from metadata.csv (e.g., strubell2019).
title: title of the document.
url: original PDF URL.
page_num: zero-based page index.
page_label: label used inside the PDF (often 1-based).

Later, we will chunk these pages into smaller overlapping segments for embedding.

Why we page-chunk first

We split the PDF into pages before chunking because pages give us a stable, easy-to-interpret unit.
This helps with:

Keeping metadata (doc ID, URL, page labels) tied to the text.
Debugging retrieval — it’s much easier to understand what the model saw if we know which page(s) were used.
Cleaning text before making smaller overlapping chunks.
Flexibility later — once pages are structured, we can try different chunk sizes or strategies without re-extracting the PDF.

In short: pages first → then chunks keeps the workflow cleaner and easier to reason about.

PYTHON

%pip install -q pypdf

Open each PDF listed in the metadata with pypdf and turn every page into one record: its text, the document id, title and URL, and the page number. The function is the reference for what a page record is; on a CPU notebook it takes several minutes over this corpus, so the cell after it first looks for pages.jsonl, the same records written once when the corpus was staged.

PYTHON

from pypdf import PdfReader

def pdfs_to_page_docs(metadata: pd.DataFrame, pdf_dir: str) -> List[Dict[str, Any]]:
    """Load each PDF into a list of page-level dictionaries.

    Each dict has keys: text, doc_id, title, url, page_num, page_label, total_pages.
    """
    page_docs: List[Dict[str, Any]] = []

    for _, row in metadata.iterrows():
        doc_id = str(row["id"]).strip()
        title = str(row.get("title", "")).strip()
        url = str(row.get("url", "")).strip()

        pdf_path = os.path.join(pdf_dir, f"{doc_id}.pdf")
        if not os.path.exists(pdf_path):
            print(f"Missing PDF for {doc_id}, skipping.")
            continue

        try:
            reader = PdfReader(pdf_path)
        except Exception as e:
            print(f"Failed to read {pdf_path}: {e}")
            continue

        total_pages = len(reader.pages)
        for i, page in enumerate(reader.pages):
            try:
                text = page.extract_text() or ""
            except Exception as e:
                print(f"Failed to extract text from {doc_id} page {i}: {e}")
                text = ""

            text = text.strip()
            if not text:
                # Still keep the page so we know it exists, but mark it as empty
                text = "[[EMPTY PAGE TEXT – see original PDF for tables/figures]]"

            page_docs.append(
                {
                    "text": text,
                    "doc_id": doc_id,
                    "title": title,
                    "url": url,
                    "page_num": i,
                    "page_label": str(i + 1),
                    "total_pages": total_pages,
                }
            )

    return page_docs


start = time.time()
try:
    pages_path = download_from_s3("wattbot/pages.jsonl", "pages.jsonl")   # staged next to corpus.zip by the corpus build script
    page_docs = [json.loads(line) for line in open(pages_path)]
    print(f"Loaded {len(page_docs)} pre-extracted page records in {time.time() - start:.1f} s")
except ClientError:
    page_docs = pdfs_to_page_docs(metadata_df, PDF_DIR)
    print(f"Extracted {len(page_docs)} page-level records from {len(metadata_df)} PDFs in {(time.time() - start) / 60:.1f} min")
page_docs[0] if page_docs else None

A paper whose PDF the corpus build could not download has no pages here, and a question that cites it can only score zero, which would measure the download rather than the pipeline. drop_questions_without_papers removes those questions from the evaluation set and says which papers were missing; the unanswerable questions, which cite nothing, stay.

PYTHON

from wattbot_rag import drop_questions_without_papers

train_df = drop_questions_without_papers(train_df, {p["doc_id"] for p in page_docs})
print(f"{len(train_df)} questions can be answered from this corpus")

Step 6 – Text chunking

RAG systems typically break documents into chunks so that:

Each chunk is long enough to carry meaningful context.
No chunk is so long that it blows up the embedding/LLM context window.

For this workshop we will implement a simple sliding-window chunker that operates on characters:

chunk_size_chars: maximum characters per chunk (e.g., 1,000–1,500).
chunk_overlap_chars: overlap between consecutive chunks (e.g., 200).

In your own work, you may wish to plug in more sophisticated semantic chunking methods (e.g., splitting on headings, section titles, or sentence boundaries). For now, we’ll keep the implementation explicit and easy to debug.

PYTHON

def split_text_into_chunks(
    text: str,
    chunk_size_chars: int = 1200,     # max characters per chunk (~300 tokens)
    chunk_overlap_chars: int = 200,   # consecutive chunks share this much, so a sentence cut at the edge survives
) -> List[str]:
    """Split `text` into overlapping character-based chunks.

    This is a simple baseline; more advanced versions might:
    - split on sentence boundaries, or
    - merge short paragraphs and respect section headings.
    """
    text = text.strip()
    if not text:
        return []

    chunks: List[str] = []
    start = 0
    text_len = len(text)

    while start < text_len:
        end = min(start + chunk_size_chars, text_len)
        chunk = text[start:end]
        chunks.append(chunk)
        if end == text_len:
            break
        # Move the window forward, keeping some overlap
        start = end - chunk_overlap_chars

    return chunks


def make_chunked_docs(
    page_docs: List[Dict[str, Any]],
    chunk_size_chars: int = 1200,
    chunk_overlap_chars: int = 200,
) -> List[Dict[str, Any]]:
    """Turn page-level records into smaller overlapping text chunks.

    Each chunk keeps a pointer back to its document and page metadata.
    """
    chunked: List[Dict[str, Any]] = []
    for page in page_docs:
        page_text = page["text"]
        chunks = split_text_into_chunks(
            page_text,
            chunk_size_chars=chunk_size_chars,
            chunk_overlap_chars=chunk_overlap_chars,
        )
        for idx, chunk_text in enumerate(chunks):
            chunked.append(
                {
                    "text": chunk_text,
                    "doc_id": page["doc_id"],
                    "title": page["title"],
                    "url": page["url"],
                    "page_num": page["page_num"],
                    "page_label": page["page_label"],
                    "total_pages": page["total_pages"],
                    "chunk_idx_in_page": idx,   # position within the page; the rest is copied from the page record
                }
            )
    return chunked

Chunking is cheap CPU work, but there is no reason to repeat it on every run. We save the chunks locally as data/chunks.jsonl (one JSON record per line) and reload the file if it already exists.

PYTHON

chunks_jsonl_path = os.path.join(local_data_dir, "chunks.jsonl")

def save_chunked_docs_jsonl(path, chunks):
    with open(path, "w", encoding="utf-8") as f:
        for rec in chunks:
            json.dump(rec, f, ensure_ascii=False)
            f.write("\n")


def load_chunked_docs_jsonl(path):
    with open(path, "r", encoding="utf-8") as f:
        return [json.loads(line) for line in f]

# -------------------------------------------------------------------
# Cached chunking logic
# -------------------------------------------------------------------
if os.path.exists(chunks_jsonl_path):
    print(f"Found existing chunk file: {chunks_jsonl_path}")
    chunked_docs = load_chunked_docs_jsonl(chunks_jsonl_path)
    print("Loaded chunked docs:", len(chunked_docs))
else:
    print("No chunk file found. Running chunking step...")
    chunked_docs = make_chunked_docs(page_docs)
    save_chunked_docs_jsonl(chunks_jsonl_path, chunked_docs)
    print(f"Saved chunked docs to {chunks_jsonl_path}")

# Show first chunk
print("Raw pages:", len(page_docs))
print("Chunked docs:", len(chunked_docs))
chunked_docs[0] if chunked_docs else None

PYTHON

# Upload to YOUR bucket so the Processing Jobs episode (and future runs) can reuse the
# chunks without repeating the unpacking and chunking steps.
print(f"Uploading chunked docs to s3://{bucket_name}/{chunks_key} ...")
s3.upload_file(chunks_jsonl_path, bucket_name, chunks_key)
print("Upload complete.")

Step 7 – Cost tracking: Bedrock calls carry no tags unless you route them through a profile

Every SageMaker job in this lesson is tagged at launch, because tags never propagate from the notebook. Bedrock is stricter still: an on-demand invoke_model or converse call has no tags parameter at all. If you call deepseek.v3-v1:0 or amazon.titan-embed-text-v2:0 by its model ID, the usage lands in the account’s Bedrock bill as anonymous per-token charges. The notebook’s tags, the bucket’s tags, and the tags you set on anything else are irrelevant to it. On a shared account that means nobody can tell which team spent what.

The fix is not a tag on the call; it is a different thing to call. An application inference profile is a small resource you create in your own account: an alias that points at one base model and carries your three tags. It has its own ARN, and that ARN goes wherever a model ID goes. Nothing else changes: same model, same price per token, same answers. The only difference is in the bill, where usage made through the alias shows up with the alias’s tags.

Before and after, in the generation call:

	`modelId` you pass	Who it is billed to
Raw model ID	`deepseek.v3-v1:0`	nobody: anonymous DeepSeek spend on the account
Your profile	`arn:aws:bedrock:us-east-2:<account-id>:application-inference-profile/adewnu8wqxpy`	`Name=DoeJohn, Project=SinkOrSwim, Purpose=Inference`

So “cost tracking on Bedrock” is: create the alias once per model (one for DeepSeek, one for Titan, because an alias points at exactly one model), keep its ARN in a variable, and use that variable as modelId from then on. The profile is free; you pay only for tokens, as before.

One call per model does all of it. tagged_profile_arn in wattbot_rag.py looks for a profile named after your project and user name, creates it with the three tags if it is not there, and returns its ARN. Open the file if you want to see the forty lines; nothing in them is worth memorising. What is worth remembering is the sentence above: one tagged profile per model, its ARN in place of the model ID.

PYTHON

from wattbot_rag import tagged_profile_arn

name = "DoeJohn"         # replace with your name
project = "SinkOrSwim"   # replace with your challenge name

# From here on, these ARNs are what we pass as modelId. The base IDs are never sent to the runtime again.
bedrock_model_id = tagged_profile_arn(bedrock, base_generation_model_id, "wattbot-gen", name, project)
embedding_model_id_bedrock = tagged_profile_arn(bedrock, base_embedding_model_id, "wattbot-embed", name, project)

Callout

The same thing from a terminal

Outside a notebook it is one CLI command per model, which is all the helper wraps:

BASH

aws bedrock create-inference-profile --inference-profile-name SinkOrSwim-DoeJohn-wattbot-gen \
  --model-source copyFrom=arn:aws:bedrock:us-east-2::foundation-model/deepseek.v3-v1:0 \
  --tags key=Name,value=DoeJohn key=Project,value=SinkOrSwim key=Purpose,value=Inference

It prints the profile ARN; pass that as modelId. The helper adds only the reuse check, so rerunning does not create duplicates.

Callout

Checking the profiles

You can confirm the routing in the Bedrock console under Inference and Assessment → Cross-region inference → Application inference profiles, or with bedrock.list_inference_profiles(typeEquals="APPLICATION").
The keys must be the three approved ones, Name, Project and Purpose, spelled exactly; any other key is ignored by the bill (the S3 episode says why).
The profile itself has no hourly or storage charge. You only pay per token, exactly as before.

Challenge

Who is paying for these tokens?

List the application inference profiles in the account and find your two. For each, print its name, the model it points at, and its tags (bedrock.list_tags_for_resource takes the profile ARN). Then answer: if a teammate skipped this step and called deepseek.v3-v1:0 by its raw model ID for the rest of the episode, what would Cost Explorer show when the organizers filter the month’s Bedrock spend by Project?

Show me the solution

PYTHON

profiles = bedrock.list_inference_profiles(typeEquals="APPLICATION")["inferenceProfileSummaries"]
print(f"{len(profiles)} application inference profiles in this account/region")
for p in profiles:
    if p["inferenceProfileArn"] in {bedrock_model_id, embedding_model_id_bedrock}:
        tags = bedrock.list_tags_for_resource(resourceARN=p["inferenceProfileArn"])["tags"]
        print(p["inferenceProfileName"], "->", [m["modelArn"].split("/")[-1] for m in p["models"]])
        print("   tags:", {t["key"]: t["value"] for t in tags})

Filtering by Project shows only the usage that went through a tagged profile. The teammate’s calls are still on the bill, under the Bedrock service line for deepseek.v3-v1:0, but with no Project, Name or Purpose value: in the grouped view they fall into the “No tag key: Project” bucket, together with every other untagged call in the account. There is no way to attribute them after the fact; tags are applied when the usage is metered, not later. That is why the base model ID is never passed to the runtime in this episode.

Step 8 – Embed the chunks with Titan

Everything the RAG loop needs is now in memory: the chunks, the questions, and a way to resolve a document id to its URL for the ref_url column.

PYTHON

print(f"Chunks: {len(chunked_docs)}")
print(f"Train QAs: {len(train_df)}")
print(f"Metadata rows: {len(metadata_df)}")

docid_to_url = build_docid_to_url(metadata_df)   # doc_id -> url, for the ref_url column of the submission
print(f"docid_to_url has {len(docid_to_url)} entries.")

Titan Text Embeddings V2 takes one string per invoke_model call and returns a 1,024-dimensional normalized vector. embedding_model_id_bedrock is the ARN of our tagged profile for it, so every call is attributed to our Project/Name/Purpose tags. Three cells: where the result is saved, the call that embeds one string, and the loop that embeds every chunk once and caches the matrix. First the paths.

PYTHON

data_dir = Path(local_data_dir)
emb_save_path = data_dir / "embeddings_bedrock.npy"             # the finished matrix, one row per chunk
emb_partial_path = data_dir / "embeddings_bedrock.partial.npy"  # progress file, so an interrupted run resumes

The call itself. invoke_model takes a provider-specific JSON body; Titan’s is {"inputText": ...}, and its response carries the vector under embedding. embed_query is the same call wrapped for retrieval time, when we embed one question at a time.

PYTHON

def bedrock_embed_text(text: str, model_id: str = embedding_model_id_bedrock) -> List[float]:
    """Call the Titan embedding model (through our tagged profile) for one string."""
    response = bedrock_runtime.invoke_model(
        modelId=model_id,
        contentType="application/json",
        accept="application/json",
        body=json.dumps({"inputText": text}),   # Titan's own request shape; add "dimensions": 256 or 512 to shrink the 1,024 default
    )
    response_body = json.loads(response["body"].read())
    embedding = response_body.get("embedding")   # a list of floats, one per dimension
    if embedding is None:
        raise ValueError(f"No 'embedding' found in response: {response_body}")
    return embedding


def embed_query(question: str) -> np.ndarray:
    """The retrieval-time half of the pipeline: the same model, one question at a time."""
    return np.asarray(bedrock_embed_text(question), dtype="float32")

Embedding the corpus, about 9,000 chunks, is ten minutes and a few cents, so do not do it twice. The cell after next looks for the finished matrix locally, then in the shared bucket if use_instructor_embeddings is True, then in your bucket, and only then computes it, saving progress every 500 chunks so an interrupted run picks up where it stopped.

Callout

In a class, use the instructor’s copy

Bedrock allows this account 300,000 Titan embedding tokens per minute, and the quota is not adjustable. This corpus is about 3 million tokens, so one notebook alone needs ten minutes at the full quota; a room of them shares it, so everyone’s loop slows to a crawl of retries. In class, leave use_instructor_embeddings in the next cell at True: the cache cell then downloads the matrix your instructor computed with the same chunks and the same model from the shared bucket and skips the calls. Working on your own, set it to False and let the loop run; that is the step to practise.

PYTHON

use_instructor_embeddings = True                            # True in class: download the instructor's matrix. False: embed the 9,000 chunks yourself, ten minutes at the quota
shared_embeddings_key = "wattbot/embeddings_bedrock.npy"   # where the instructor staged it in the shared bucket

One small helper asks S3 whether an object exists, and one checks that a cached matrix belongs to these chunks: a matrix from an earlier corpus has a different number of rows, and using it would pair every question with the wrong passages.

PYTHON

def _s3_object_exists(bucket: str, key: str) -> bool:
    try:
        s3.head_object(Bucket=bucket, Key=key)
        return True
    except ClientError:
        return False


def load_embeddings_if_matching(path) -> Any:
    """The matrix at `path` if it has one row per chunk, else None (a matrix from a different corpus)."""
    matrix = np.load(path)
    if matrix.shape[0] == len(chunked_docs):
        return matrix
    print(f"{path} has {matrix.shape[0]} rows but there are {len(chunked_docs)} chunks: a matrix from a different corpus, ignoring it")
    return None


# Cache order: local file, the instructor's copy (in class), your bucket, then recompute (and upload for next time)
chunk_embeddings = None
if emb_save_path.exists():
    chunk_embeddings = load_embeddings_if_matching(emb_save_path)
if chunk_embeddings is None and use_instructor_embeddings and _s3_object_exists(shared_bucket, shared_embeddings_key):
    s3.download_file(shared_bucket, shared_embeddings_key, str(emb_save_path))
    print(f"Downloaded the instructor's embeddings from s3://{shared_bucket}/{shared_embeddings_key}")
    chunk_embeddings = load_embeddings_if_matching(emb_save_path)
if chunk_embeddings is None and _s3_object_exists(bucket_name, embeddings_key):
    s3.download_file(bucket_name, embeddings_key, str(emb_save_path))
    print(f"Downloaded your cached embeddings from s3://{bucket_name}/{embeddings_key}")
    chunk_embeddings = load_embeddings_if_matching(emb_save_path)
if chunk_embeddings is None:
    # Resume from a partial save if a previous run was interrupted
    all_embeddings = list(np.load(emb_partial_path)) if emb_partial_path.exists() else []
    start = len(all_embeddings)
    print(f"Computing embeddings via Bedrock from chunk {start} of {len(chunked_docs)}...")
    t0 = time.time()

    for idx in range(start, len(chunked_docs)):
        all_embeddings.append(bedrock_embed_text(chunked_docs[idx]["text"]))
        if (idx + 1) % 500 == 0 or idx + 1 == len(chunked_docs):
            np.save(emb_partial_path, np.array(all_embeddings, dtype="float32"))
            print(f"Embedded {idx + 1} / {len(chunked_docs)} chunks (progress saved)")

    print(f"Embedded {len(chunked_docs) - start} chunks in {time.time() - t0:.1f} seconds")
    chunk_embeddings = np.array(all_embeddings, dtype="float32")   # the embedding matrix: one row per chunk, one column per dimension
    np.save(emb_save_path, chunk_embeddings)
    emb_partial_path.unlink(missing_ok=True)
    s3.upload_file(str(emb_save_path), bucket_name, embeddings_key)
    print(f"Saved embeddings to {emb_save_path} and s3://{bucket_name}/{embeddings_key}")

assert chunk_embeddings.shape[0] == len(chunked_docs), "embeddings and chunks are out of step; delete the cached files and rerun"
print("Embeddings shape:", chunk_embeddings.shape)

Step 9 – Retrieve: cosine similarity over the embedding matrix

Retrieval is pure NumPy: embed the question with the same model, take the cosine similarity against every chunk, and keep the top k. retrieve_top_k from wattbot_rag.py does that (open the file to see it: it is a normalize-and-matrix-multiply). No vector database is needed for a few thousand chunks; the optional GPU-notebook episode discusses when one becomes worth it.

PYTHON

first_row = train_df.iloc[0]
print("Sample question:", first_row["question"])

# Cosine similarity of the question vector against every row of the matrix; keep the k best
test_retrieved = retrieve_top_k(embed_query(first_row["question"]), chunk_embeddings, chunked_docs, k=3)
for r in test_retrieved:   # score is the cosine similarity; 1.0 would be an identical vector
    snippet = r["text"].replace("\n", " ")[:160]
    print(f"- score={r['score']:.3f} | doc_id={r['doc_id']} | page={r['page_label']} | {snippet}...")

Step 10 – Generate: DeepSeek V3 through the `converse` API

Bedrock has two runtime APIs. invoke_model takes a provider-specific request body (which is why the Titan call above builds a {"inputText": ...} body by hand). converse is provider-neutral: a system prompt, a list of messages, and an inferenceConfig, in the same shape for DeepSeek, Claude, Llama and the rest. We use converse for generation so the model is a one-line swap. The response carries the generated text and a usage block with the input and output token counts, which is what Bedrock bills on; we keep every usage block so we can price the run at the end.

PYTHON

bedrock_usage = []   # one {"inputTokens": ..., "outputTokens": ...} per call

def call_bedrock_chat(
    system_prompt: str,
    user_prompt: str,
    max_new_tokens: int = 512,
    model_id: str = bedrock_model_id,   # the tagged profile ARN, never the raw model ID
) -> str:
    """Send one system + user turn to a Bedrock chat model and return the generated text."""
    try:
        response = bedrock_runtime.converse(
            modelId=model_id,
            # system: instructions that apply to the whole exchange
            system=[{"text": system_prompt}],
            # messages: the conversation so far, here a single user turn
            messages=[{"role": "user", "content": [{"text": user_prompt}]}],
            # inferenceConfig: generation settings; temperature 0 makes the output repeatable
            inferenceConfig={"maxTokens": max_new_tokens, "temperature": 0},
        )
    except ClientError as e:
        print(f"ERROR calling Bedrock model {model_id}: {e}")
        raise

    bedrock_usage.append(response["usage"])   # token counts for this call, kept so we can price the run later
    # Some models return extra content blocks (e.g. reasoning); keep only the text ones.
    text_blocks = [block["text"] for block in response["output"]["message"]["content"] if "text" in block]
    return "\n".join(text_blocks).strip()


print(call_bedrock_chat("You are a concise assistant.", "What is 2 + 2? Answer in one short sentence.", max_new_tokens=32))
print("usage of that call:", bedrock_usage[-1])

That function has the signature the shared pipeline expects, generate(system_prompt, user_prompt, max_new_tokens) -> str, and embed_query has the other one, embed_query(question) -> vector. Those two callables are the only difference between this episode and the self-hosted ones.

What the pipeline asks the model

run_single_qa in wattbot_rag.py makes two calls per question. The first asks for a JSON object in the WattBot submission format, from the retrieved context only; the second, made only when the first produced an answer, asks for a one-to-three-sentence explanation of how the cited material supports it. The system prompt of the first call is the contract that makes the output machine-readable:

PYTHON

print(ANSWER_SYSTEM_PROMPT)

A question is marked unanswerable (is_blank everywhere, as the competition requires) when the model says so or when the best retrieval score is below retrieval_threshold.

Step 11 – Run the WattBot evaluation with Bedrock

Now we can loop over the questions, run retrieval + Bedrock generation, and write a wattbot_solutions_bedrock.csv file (locally and to your bucket).

train_QA.csv has 245 questions, and at two calls each that is about twenty minutes. For class we score a stratified sample instead: sample_questions_by_type draws ten questions from each evidence type (Quote, Table, Figure, Math, unanswerable, cross-paper, reconcile), about 55 questions after overlaps, so every row of the breakdown in the next step has something in it. The first 55 rows of the file would be nearly all Quote questions. The seed fixes the draw, and the Processing Jobs episode uses the same one, so the two runs are compared on the same questions. Set per_type=0 to run the whole file.

The Processing Jobs episode reuses this same loop with self-hosted models; the only difference there is that the embedding, answer and explanation phases call local models instead of Bedrock-hosted ones.

PYTHON

from wattbot_rag import sample_questions_by_type

eval_df = sample_questions_by_type(train_df, per_type=10, seed=0)   # per_type=0 for all 245
print(f"Scoring {len(eval_df)} of {len(train_df)} questions")

results = []
run_start = time.time()

for _, row in eval_df.iterrows():
    print("#" * 96)
    print(f"QUESTION: {row['question']}")

    out = run_single_qa(
        row=row,                             # one train_QA.csv row: id, question, expected answer_unit
        embed_query=embed_query,             # the two model callables: the only episode-specific inputs
        generate=call_bedrock_chat,
        chunk_embeddings=chunk_embeddings,   # the index: the matrix plus the chunks it was built from
        chunked_docs=chunked_docs,
        docid_to_url=docid_to_url,           # for the ref_url column
        top_k=8,                             # chunks handed to the model as context
        retrieval_threshold=0.25,            # best cosine score below this => unanswerable (is_blank)
        verbose=True,
    )

    print(f"ANSWER: {out['answer']}")
    print(f"ref_ids: {out['ref_id']}")
    print(f"EXPLANATION: {out['explanation']}")
    results.append(out)

results_df = pd.DataFrame(results)
print(f"Answered {len(results_df)} questions in {(time.time() - run_start) / 60:.1f} min")

output_dir = "outputs"
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, "wattbot_solutions_bedrock.csv")
results_df.to_csv(output_path, index=False)
s3.upload_file(output_path, bucket_name, solutions_key)
print(f"Wrote predictions to {output_path} and s3://{bucket_name}/{solutions_key}")

results_df.head()

Step 12 – Score against the training answers

WattBot scores each prediction with three components: 0.75 for answer_value (numeric within ±0.1% relative tolerance, categorical match after normalisation, and zero if you answer a question the corpus cannot answer), 0.20 for the F1 between your ref_id set and the ground-truth citations, and 0.05 for refusing cleanly, with every evidence field blank, when the answer is is_blank. The number comes from Score.py, the competition’s own metric file, which compute_wattbot_score in wattbot_rag.py calls for the headline score so that what you see here is what the leaderboard would report. It also prints the share of correct answers by evidence type, the seven flag columns of train_QA.csv: Quote (the answer is a sentence in a paper), Table, Figure, Math (a calculation from numbers in the text), is_NA (the corpus cannot answer it, and the right response is is_blank), CrossPaper (numbers from two papers have to be combined) and Reconcile (two figures that sound alike but measure different things, such as water withdrawn against water consumed, or a popular claim against a measured value, and the answer is the ratio or the distinction between them; quoting either figure alone is the wrong answer the question is built to catch). A question can carry several. Read that table before changing anything: figure and table questions failing while quote questions pass points at text extraction, which loses figures and flattens tables, not at the model. The wrapper adds a per-row score, joins on id, and prints the worst-scoring rows so you can see why the pipeline lost points. train_QA.csv has no separate “unanswerable” column, so gt_is_na_col="is_blank" tells the scorer to derive that flag from answer_value == "is_blank" on both sides.

PYTHON

scored_df = compute_wattbot_score(
    train_qa_path=train_qa_path,    # ground truth
    preds_path=output_path,         # our predictions, joined to it on id
    n_examples=20,
)

Why were the wrong ones wrong?

The score tells you which questions failed, not why. Two different failures look the same in that table. Either retrieval never surfaced the right paper, so the model had nothing to work with, or the right passage was in front of it and it still declined or misread it. The cell below separates them for the lowest-scoring rows: it joins the model’s own explanation and the passages it cited back onto the score, then reruns retrieval for the question (one embedding call) and reports whether the ground-truth paper was among the top-k chunks and what the best cosine score was.

PYTHON

from wattbot_rag import retrieve_top_k, split_ref_ids

worst = scored_df.sort_values("wattbot_score").head(8)
by_id = results_df.set_index("id")   # the model's answer, explanation and cited passages, keyed by question id

for _, row in worst.iterrows():
    qid = row["id"]
    gt_refs = set(split_ref_ids(row["ref_id_gt"]))   # train_QA.csv stores the list as a string
    retrieved = retrieve_top_k(embed_query(row["question_gt"]), chunk_embeddings, chunked_docs, k=8)   # the merge suffixed both sides' columns
    retrieved_docs = [r["doc_id"] for r in retrieved]
    print("=" * 96)
    print(f"{qid}: {row['question_gt']}")
    print(f"  expected: {row['answer_value_gt']}  from {sorted(gt_refs)}")
    print(f"  predicted: {row['answer_value_pred']}  citing {row['ref_id_pred']}")
    print(f"  retrieval: top cosine {retrieved[0]['score']:.3f}; top-8 papers {retrieved_docs}")
    print(f"  ground-truth paper retrieved: {'yes' if gt_refs & set(retrieved_docs) else 'NO'}")
    if by_id.loc[qid, "explanation"] != "is_blank":
        print(f"  model's explanation: {by_id.loc[qid, 'explanation']}")
        print(f"  passages it cited: {str(by_id.loc[qid, 'supporting_materials'])[:400]}")

Read the two lines at the bottom of each block together. “Ground-truth paper retrieved: NO” is a retrieval problem: a different chunk size, a reranker, or a better embedding model is the fix, and no generator will help. “Yes” with is_blank is the model declining on evidence it had, which the system prompt and the retrieval threshold control. “Yes” with a wrong value is a reading error, and the cited passage usually shows what it latched onto: a nearby number, a different unit, a different model in the same table. And sometimes the explanation shows the model read correctly and the label is the problem: a question asking what inference “can account for” labelled 90 where three papers in the corpus give 60, 70 and 80 to 90 percent, so a range is arguably the better answer, or a fact that appears in several papers with the label naming only one of them. Labels are written by people. When the explanation convinces you, the fix is to the label, and the competition organisers want to hear about it. Each of those is a different lever, and the score alone would have you pulling the wrong one.

Challenge

What did that run cost?

Every converse response carried a usage block, and call_bedrock_chat kept them all in bedrock_usage. Sum the input and output tokens over the evaluation run and price them at the current on-demand rates for DeepSeek V3 on the Bedrock pricing page (fill the two prices in below). Add the embedding step: ~9,000 chunks of 1,200 characters at Titan Text Embeddings V2’s $0.02 per million input tokens. How does the total compare with the overview’s estimate for “the whole WattBot pipeline on Bedrock”, and with the ml.g5.xlarge GPU the Processing Jobs episode rents for the same questions, $1.41 an hour?

Show me the solution

The two prices are from the Bedrock pricing page for DeepSeek on-demand in October 2026; update them for the model and region you used. The rest sums the usage blocks and prints the cost.

PYTHON

price_in_per_m = 0.62    # $ per million input tokens: DeepSeek v3.2 on-demand in us-east-2, October 2026; check the pricing page for your model and region
price_out_per_m = 1.85   # $ per million output tokens, same source

in_tokens = sum(u["inputTokens"] for u in bedrock_usage)
out_tokens = sum(u["outputTokens"] for u in bedrock_usage)
gen_cost = in_tokens / 1e6 * price_in_per_m + out_tokens / 1e6 * price_out_per_m

# Titan bills per input token; ~4 characters per token is a fair estimate for English prose
embed_tokens = sum(len(ch["text"]) for ch in chunked_docs) / 4
embed_cost = embed_tokens / 1e6 * 0.02

print(f"{len(bedrock_usage)} generation calls: {in_tokens:,} input + {out_tokens:,} output tokens -> ${gen_cost:.3f}")
print(f"~{embed_tokens:,.0f} embedding tokens -> ${embed_cost:.3f}")
print(f"Total ~${gen_cost + embed_cost:.2f}; the Processing Jobs episode rents an ml.g5.xlarge at $1.41 an hour for the same questions")

With 53 questions, two calls each and ~1,500-token prompts, the run was 140K input and 9K output tokens: ten cents at DeepSeek V3’s prices, plus five cents of embeddings for the 2.4 million tokens of corpus. All 245 questions is about half a dollar, which is the overview’s estimate for the whole pipeline. The Processing Jobs episode answers the same 53 questions on a rented ml.g5.xlarge and bills about $0.60 of GPU time for them, four times this run, most of it loading a model that Bedrock already has loaded.

Challenge

Swap the models

Everything above is two model IDs and two callables. The model catalog has dozens of chat models and several embedding models, and “Choosing the two models” at the top of the episode said the leaderboards cannot tell you which one wins on your questions. Find out.

A different generator. Pick another chat model from the catalog (for example us.amazon.nova-2-lite-v1:0, us.meta.llama4-maverick-17b-instruct-v1:0, or a Mistral or Claude model if your account has access), wrap it in a tagged profile as in Step 7, and rerun the Step 11 loop on the same 55-question sample with call_bedrock_chat pointed at it. Score it with compute_wattbot_score and price it from its usage blocks. Which evidence types moved, and what did a point of score cost?
A different embedder. Retrieval is the other half. The Ohio catalog has one other embedding model, Cohere Embed v4 (cohere.embed-v4:0), and its request body is different from Titan’s: a list of texts and an input_type that is search_document for the corpus and search_query for the question. Embed the chunks with it under a new file name so the Titan matrix is kept (ten minutes at the quota, so not in class), rerun Step 11 with DeepSeek unchanged, and score.

Keep the rest fixed. One change per experiment, or you will not know what moved the score.

Show me the solution

Both solution cells start with %%script false --no-raise-error, which makes Jupyter skip the cell, so a Run All of the episode does not spend ten minutes here; delete that line to run one.

Generator. call_bedrock_chat already takes a model_id, so the swap is a profile and a lambda:

PYTHON

%%script false --no-raise-error
# Challenge solution: delete the line above to run this cell. It is skipped on Run All so the episode's own cells are not delayed by it.
alt_generation_model_id = "us.amazon.nova-2-lite-v1:0"   # any chat model in the catalog with a converse API
alt_model_arn = tagged_profile_arn(bedrock, alt_generation_model_id, "wattbot-gen-alt", name, project)

usage_before = len(bedrock_usage)
alt_results = []
for _, row in eval_df.iterrows():
    alt_results.append(run_single_qa(
        row=row, embed_query=embed_query,
        generate=lambda s, u, max_new_tokens=512: call_bedrock_chat(s, u, max_new_tokens, model_id=alt_model_arn),
        chunk_embeddings=chunk_embeddings, chunked_docs=chunked_docs, docid_to_url=docid_to_url,
        top_k=8, retrieval_threshold=0.25,
    ))
alt_path = os.path.join(output_dir, "wattbot_solutions_bedrock_alt.csv")
pd.DataFrame(alt_results).to_csv(alt_path, index=False)

alt_usage = bedrock_usage[usage_before:]   # only this run's calls
print(f"{sum(u['inputTokens'] for u in alt_usage):,} input + {sum(u['outputTokens'] for u in alt_usage):,} output tokens; price them from the pricing page")
compute_wattbot_score(train_qa_path=train_qa_path, preds_path=alt_path, n_examples=5)

Read the per-type lines against the DeepSeek run. A cheaper model that holds its own on Quote questions and loses on Math and CrossPaper is the usual pattern; whether that trade is worth it depends on which questions you have.

Embedder. First list the embedding models the region’s catalog has, because model IDs come and go and an ID that is not in the list fails at the profile step with “The provided model identifier is invalid”:

PYTHON

for m in bedrock.list_foundation_models(byOutputModality="EMBEDDING")["modelSummaries"]:
    print(f"{m['modelId']:<45} {m['providerName']}")

Cohere’s request takes a list of texts and says whether they are documents or a query; its response is a list of vectors, nested under the number type asked for. One wrinkle: the catalog lists cohere.embed-v4:0 for Ohio, but creating a profile on it fails with “does not support On Demand inference”, because in this region the model is served only through the US cross-region profile, us.cohere.embed-v4:0. tagged_profile_arn catches that error and switches to the us. ID itself, which is the same thing that happens for Llama 4 and several other models. Everything else is Step 8 again: a tagged profile, one call per chunk, a cached matrix, and a query function that uses the same model.

PYTHON

%%script false --no-raise-error
# Challenge solution: delete the line above to run this cell. Ten minutes of embedding, so not in class.
cohere_model_id = "cohere.embed-v4:0"   # listed for Ohio but served only through the US cross-region profile; tagged_profile_arn switches to us.cohere.embed-v4:0 on its own
cohere_arn = tagged_profile_arn(bedrock, cohere_model_id, "wattbot-embed-cohere", name, project)

def cohere_embed(text: str, input_type: str = "search_document") -> np.ndarray:
    """One Cohere Embed v4 vector; input_type is search_document for the corpus, search_query for a question."""
    body = {"texts": [text[:2000]], "input_type": input_type, "embedding_types": ["float"]}
    response = bedrock_runtime.invoke_model(modelId=cohere_arn, contentType="application/json", accept="application/json", body=json.dumps(body))
    return np.asarray(json.loads(response["body"].read())["embeddings"]["float"][0], dtype="float32")

cohere_emb_path = data_dir / "embeddings_cohere.npy"   # a new file: the Titan matrix stays
if cohere_emb_path.exists():
    cohere_embeddings = np.load(cohere_emb_path)
else:
    t0 = time.time()
    cohere_embeddings = np.vstack([cohere_embed(ch["text"]) for ch in chunked_docs])
    np.save(cohere_emb_path, cohere_embeddings)
    print(f"Embedded {len(chunked_docs)} chunks in {(time.time() - t0) / 60:.1f} min")

def embed_query_cohere(question: str) -> np.ndarray:
    return cohere_embed(question, input_type="search_query")

Then the Step 11 loop with embed_query=embed_query_cohere and chunk_embeddings=cohere_embeddings, DeepSeek unchanged, and the scorer. The hit-rate cell from Step 13 is the sharper instrument here: compute it for the new matrix before scoring, because an embedder that finds the paper less often cannot be rescued by any generator. A cheaper variant of the same experiment keeps Titan and adds "dimensions": 256 to its request body in Step 8, which asks whether a four-times smaller index retrieves as well.

Step 13 – Multimodal embeddings for the table and figure questions

Look back at the evidence-type table in Step 12. The Table and Figure rows are usually the worst, and the reason is upstream of the model. pypdf gives us a page’s text: a table comes out as a run of numbers with its column headers somewhere else on the page, and a figure comes out as its caption and nothing more. A question whose answer is a bar in a plot has nothing in the chunks to find, and a question whose answer is a cell in a table is being read from a flattened copy.

A multimodal embedding model is one model that maps text and images into the same vector space, so a text question can be compared with an image directly: the picture of a table lands near a sentence about its contents. Bedrock has two: Amazon Nova Multimodal Embeddings (text, images, documents, video and audio, released 2025) and the older Titan Multimodal Embeddings G1 (text and images). Both are per-call, like Titan Text, and billed per image plus per text token; the Bedrock pricing page has the rates, on the order of a hundredth of a cent per image. One catch: in October 2026 both are offered in us-east-1 only (check the model catalog for that region). A Bedrock call is an HTTPS request and the region is a parameter, so from Ohio that is a second client pointed at Virginia; it also means the images leave the region, which is the thing to check against your institution’s data policy before doing this with documents that are not public papers.

This step cuts the tables and figures out of the PDFs as images, builds a second index from them next to the text one, measures whether it finds the right paper for table and figure questions more often, and then hands the images to a model that can look at them. Each cell is short because the plumbing is in wattbot_rag.py.

Extract the tables and figures

Why cut them out at all, rather than hand the embedding model a whole page, or the PDF itself (Nova accepts documents as input)? For the same reason Step 6 cut the text into chunks instead of embedding each paper whole. An embedding is a fixed-length vector, a thousand numbers, whatever goes in; the more that goes in, the less of each thing survives in it. A page holding two tables, a figure and three paragraphs gets one vector that is a blur of all six, and a question about one of the tables has to match that blur. A crop of the table gets a vector that is about the table. The same holds at generation time: the vision model reads a crop at full resolution for a few hundred tokens, where a page costs a thousand tokens and shows the table at a quarter of the size. We tried whole pages first on this corpus: they helped figure questions a little and tables not at all. The crops, as the callout at the end shows, help both.

A PDF has no table or figure objects to pull out. It stores text runs, raster images, and drawing commands (lines, fills, curves), and a plot is a few thousand of the last. PyMuPDF ships three detectors that turn those back into boxes: find_tables finds tables from their rulings or their text alignment, get_image_info lists the embedded raster images, and cluster_drawings groups the drawing commands of a vector figure into one bounding box. extract_figures_and_tables in wattbot_rag.py runs the three on every page, drops page-sized background fills, rules and logos, pairs each box with the nearest caption line starting “Table N” or “Figure N” (tables are captioned above, figures below) and widens the crop to include that caption, so the image carries the identifier a question will use. A caption whose box no detector found falls back to the whole page, so nothing the text mentions is missed. Open the function; it is sixty lines and every rule in it is a line you can change. The detectors are the slow part, a second or two per page on this instance, so the cell runs them only on the pages whose text carries a caption, under half of the corpus; in class, with use_instructor_embeddings set, it downloads the crops your instructor made instead.

PYTHON

%pip install -q pymupdf

PYTHON

from wattbot_rag import extract_figures_and_tables, find_caption_pages, embed_multimodal, image_blocks, hit_at_k

crop_dir = data_dir / "crops"
crops_path = data_dir / "crops.jsonl"
shared_crops_key = "wattbot/crops.zip"   # the instructor's crops, staged like the embeddings

if not crops_path.exists() and use_instructor_embeddings and _s3_object_exists(shared_bucket, shared_crops_key):
    s3.download_file(shared_bucket, shared_crops_key, str(data_dir / "crops.zip"))
    with zipfile.ZipFile(data_dir / "crops.zip") as zf:
        zf.extractall(data_dir)
    (data_dir / "crops.zip").unlink()
    print("Downloaded the instructor's crops")

if crops_path.exists():
    crops = [json.loads(line) for line in open(crops_path)]
else:
    caption_pages: Dict[str, Dict[int, set]] = {}   # per paper, the pages whose text carries a caption, and which kinds: the detectors run only there
    for p in find_caption_pages(page_docs):
        caption_pages.setdefault(p["doc_id"], {})[p["page_num"]] = {"table" if c.startswith("Table") else "figure" for c in p["captions"]}
    t0 = time.time()
    crops = []
    for i, (doc_id, page_nums) in enumerate(sorted(caption_pages.items())):
        crops += extract_figures_and_tables(os.path.join(PDF_DIR, f"{doc_id}.pdf"), doc_id, crop_dir, pages=page_nums)
        if (i + 1) % 10 == 0:
            print(f"{i + 1} papers, {len(crops)} crops so far ({time.time() - t0:.0f} s)")
    with open(crops_path, "w") as f:
        f.writelines(json.dumps(c) + "\n" for c in crops)
    print(f"Extracted {len(crops)} crops from {len(caption_pages)} papers in {(time.time() - t0) / 60:.1f} min")

n_tables = sum(c["kind"] == "table" for c in crops)
n_pages = sum(c["whole_page"] for c in crops)
print(f"{n_tables} tables and {len(crops) - n_tables} figures; {n_pages} of them are whole-page fallbacks")
crops[:3]

Look at a few before trusting them: a crop that is a logo, a page rule, or half a plot tells you which threshold to change.

PYTHON

from IPython.display import Image, display

for c in [c for c in crops if c["kind"] == "table"][:2] + [c for c in crops if c["kind"] == "figure"][:2]:
    print(f"{c['doc_id']} page {c['page_label']}: {c['caption']}")
    display(Image(filename=c["png_path"], width=500))

Embed the crops

The embedding model lives in us-east-1, so this is a second runtime client pointed there. One difference from the text models: Bedrock rejects this model’s request body when it arrives through an application inference profile, so these calls use the model ID directly and are the one Bedrock spend in this lesson that carries no project tags. embed_multimodal builds the model’s request body (Nova’s says whether a vector is for the index or for a query, and that the image is a document rather than a photo) and returns a unit vector. Set multimodal_model_id to amazon.titan-embed-image-v1 to try the older model; the function handles both.

PYTHON

multimodal_region = "us-east-1"
multimodal_model_id = "amazon.nova-2-multimodal-embeddings-v1:0"   # called by ID: this model does not accept calls through an inference profile

bedrock_runtime_mm = boto3.client("bedrock-runtime", region_name=multimodal_region,
                                  config=Config(retries={"max_attempts": 10, "mode": "adaptive"}))

v = embed_multimodal(bedrock_runtime_mm, multimodal_model_id, image_png=Path(crops[0]["png_path"]).read_bytes())
print(f"{crops[0]['doc_id']} {crops[0]['caption'][:40]!r} -> vector of {v.shape[0]} dimensions")

One call per crop, cached like the text matrix: locally, then in your bucket, then computed. The loop runs six calls at a time through a thread pool and prints every fifty; a thousand images takes about two minutes. The question side uses the same model with purpose="query".

PYTHON

crop_emb_path = data_dir / "embeddings_crops_multimodal.npy"
crop_emb_key = "wattbot/embeddings/embeddings_crops_multimodal.npy"

if crop_emb_path.exists():
    crop_embeddings = np.load(crop_emb_path)
elif use_instructor_embeddings and _s3_object_exists(shared_bucket, "wattbot/embeddings_crops_multimodal.npy"):
    s3.download_file(shared_bucket, "wattbot/embeddings_crops_multimodal.npy", str(crop_emb_path))
    crop_embeddings = np.load(crop_emb_path)
    print("Downloaded the instructor's crop embeddings")
elif _s3_object_exists(bucket_name, crop_emb_key):
    s3.download_file(bucket_name, crop_emb_key, str(crop_emb_path))
    crop_embeddings = np.load(crop_emb_path)
else:
    from concurrent.futures import ThreadPoolExecutor

    def embed_crop(c):
        return embed_multimodal(bedrock_runtime_mm, multimodal_model_id, image_png=Path(c["png_path"]).read_bytes())

    t0 = time.time()
    rows = []
    with ThreadPoolExecutor(max_workers=6) as pool:   # one image per call: six in flight at a time
        for i, vec in enumerate(pool.map(embed_crop, crops)):
            rows.append(vec)
            if (i + 1) % 50 == 0:
                print(f"Embedded {i + 1} / {len(crops)} crops ({time.time() - t0:.0f} s, {(i + 1) / (time.time() - t0):.1f} per second)")
    crop_embeddings = np.vstack(rows)
    np.save(crop_emb_path, crop_embeddings)
    s3.upload_file(str(crop_emb_path), bucket_name, crop_emb_key)
    print(f"Embedded {len(crops)} crops in {(time.time() - t0) / 60:.1f} min")

assert crop_embeddings.shape[0] == len(crops), "crop embeddings and crops are out of step; delete the cached file and rerun"
print("Crop embeddings shape:", crop_embeddings.shape)

def embed_query_mm(question: str) -> np.ndarray:
    return embed_multimodal(bedrock_runtime_mm, multimodal_model_id, text=question, purpose="query")

Does the image index find the right table or figure?

Retrieval is the same retrieve_top_k as Step 9, over the crop matrix instead of the chunk matrix. Two measures, both hit rates, higher is better, 1.00 meaning every question had it. The first is coarse: is a ground-truth paper among the top results? train_QA.csv names the paper for every question. The second is the one that matters for this step: is the table or figure the answer sits in among the crops the model would be shown? Most table and figure questions name it in supporting_materials (“Table 3”, “Fig. 5, upper panel”), and cited_captions pulls that identifier out; crop_matches checks a crop’s caption against it. For the text index, the crops it “shows” are those on the pages its top chunks came from, which is what the third vision run below attaches. Quote-only questions are the control for the paper measure, where the image index should have no advantage.

PYTHON

from wattbot_rag import cited_captions, crop_matches

k = 8
groups = {"Table": train_df[train_df["Table"] == 1], "Figure": train_df[train_df["Figure"] == 1],
          "Quote only": train_df[(train_df["Quote"] == 1) & (train_df["Table"] == 0) & (train_df["Figure"] == 0)].head(40)}
crops_by_page: Dict[tuple, list] = {}
for c in crops:
    crops_by_page.setdefault((c["doc_id"], c["page_num"]), []).append(c)

def crops_behind(chunks):   # the crops on the pages the retrieved chunks came from
    return [c for ch in chunks for c in crops_by_page.get((ch["doc_id"], ch["page_num"]), [])]

print("Hit rate (higher is better). Paper: a ground-truth paper is in the top 8 results. Table/figure: the one the answer")
print("sits in is among the crops shown (the image index's top 3; for the text index, the crops on its top chunks' pages)\n")
print(f"{'questions':<12}{'n':>4}{'paper, text':>14}{'paper, image':>14}{'crop, text':>14}{'crop, image':>14}")
for label, qs in groups.items():
    gt = [split_ref_ids(r) for r in qs["ref_id"]]
    text_chunks = [retrieve_top_k(embed_query(q), chunk_embeddings, chunked_docs, k=k) for q in qs["question"]]
    image_crops = [retrieve_top_k(embed_query_mm(q), crop_embeddings, crops, k=3) for q in qs["question"]]
    paper_text = hit_at_k([[c["doc_id"] for c in r] for r in text_chunks], gt)
    paper_image = hit_at_k([[c["doc_id"] for c in r] for r in image_crops], gt)
    # the crop measure, over the questions that name their table or figure
    targets = [(refs, cited_captions(sm)) for refs, sm in zip(gt, qs["supporting_materials"])]
    crop_text, crop_image, n_named = [], [], 0
    for (refs, labels), chunks_i, crops_i in zip(targets, text_chunks, image_crops):
        if not labels:
            continue
        n_named += 1
        crop_text.append(any(crop_matches(c, d, l) for c in crops_behind(chunks_i) for d in refs for l in labels))
        crop_image.append(any(crop_matches(c, d, l) for c in crops_i for d in refs for l in labels))
    crop_cols = f"{np.mean(crop_text):>14.2f}{np.mean(crop_image):>14.2f}  (of {n_named} that name one)" if n_named else f"{'':>14}{'':>14}"
    print(f"{label:<12}{len(qs):>4}{paper_text:>14.2f}{paper_image:>14.2f}{crop_cols}")

Read the paper columns first: both indexes find the paper for nearly every table and figure question, so at that level there is nothing to win, and the image index is behind. The crop columns are the real comparison. They say whether the text route, chunks and the crops on their pages, or the image route, the crops themselves, more often puts the exact table or figure in front of the model. If the image column is higher there while the paper column is lower, the image index is doing a different job than the text index, finding the object rather than the document, and the next cells show whether the model can use it. If both crop columns are low, the fix is upstream, in the extraction: a table the detector missed became a whole-page fallback, and a page is a weak key.

Let the model see the table or figure

Finding the table is half of it; DeepSeek V3 reads text only, so a figure retrieved as an image is no use to it. converse accepts image blocks in a user message for models that can see, and the swap is the model ID: Amazon Nova 2 Lite reads images and documents and is in the catalog for Ohio through the cross-region profile below (check the model catalog; us.amazon.nova-lite-v1:0 is the cheaper, weaker alternative). call_vision_chat is call_bedrock_chat with an images argument: for each attached crop it adds a text label naming the paper and the caption, then the image, so the model can cite what it reads.

PYTHON

base_vision_model_id = "us.amazon.nova-2-lite-v1:0"
vision_model_id = tagged_profile_arn(bedrock, base_vision_model_id, "wattbot-vision", name, project)
vision_usage = []

def call_vision_chat(system_prompt: str, user_prompt: str, max_new_tokens: int = 512, images=()) -> str:
    """One system + user turn, with zero or more (label, png_bytes) images attached to the user turn."""
    content = [{"text": user_prompt}]
    if images:
        content.append({"text": "The images below are tables and figures from the same papers; read values from them where the text above is incomplete."})
        content.extend(image_blocks(images))
    response = bedrock_runtime.converse(
        modelId=vision_model_id,
        system=[{"text": system_prompt}],
        messages=[{"role": "user", "content": content}],
        inferenceConfig={"maxTokens": max_new_tokens, "temperature": 0},
    )
    vision_usage.append(response["usage"])
    return "\n".join(b["text"] for b in response["output"]["message"]["content"] if "text" in b).strip()

def crop_label(c) -> str:
    return f"{c['doc_id']}, page {c['page_label']}, {c['caption'][:60]}" if c["caption"] else f"{c['doc_id']}, page {c['page_label']}"

print(call_vision_chat("You are a concise assistant.", "What is in the attached image? One sentence.", 64,
                       images=[(crop_label(crops[0]), Path(crops[0]["png_path"]).read_bytes())]))

Three runs over the Table and Figure questions, all through the vision model, so that the only thing changing between them is what the model is shown. The first is the control: text chunks only. The second attaches the top three crops from the image index. The third attaches the crops from the pages the top text chunks came from: if the hit-rate table says the text index is the better finder, the question is whether finding with text and reading with the image is the better split. Each run is the same loop as Step 11, so every question costs an answer call and an explanation call; the runs with images take about five minutes each, because every attached crop is hundreds of input tokens. In class, with use_instructor_embeddings set, the cell downloads the three result files from your instructor’s run of it instead and the next cell scores them; running the loops yourself is for after class.

PYTHON

tf_df = train_df[(train_df["Table"] == 1) | (train_df["Figure"] == 1)]
print(f"{len(tf_df)} table and figure questions; each run makes an answer call and an explanation call per question, so about {6 * len(tf_df)} calls for three runs")

def crops_for(question: str, images_from: str):
    """The crops to attach: none, the image index's top three, or those on the pages behind the top text chunks."""
    if images_from == "image_index":
        return retrieve_top_k(embed_query_mm(question), crop_embeddings, crops, k=3)
    if images_from == "text_index":
        found = []
        for c in crops_behind(retrieve_top_k(embed_query(question), chunk_embeddings, chunked_docs, k=8)):
            if c not in found:
                found.append(c)
        return found[:3]
    return []

def run_tf(images_from: str) -> pd.DataFrame:
    rows = []
    for i, (_, row) in enumerate(tf_df.iterrows(), start=1):
        images = [(crop_label(c), Path(c["png_path"]).read_bytes()) for c in crops_for(row["question"], images_from)]
        generate = lambda s, u, max_new_tokens=512, images=images: call_vision_chat(s, u, max_new_tokens, images=images)
        rows.append(run_single_qa(row=row, embed_query=embed_query, generate=generate, chunk_embeddings=chunk_embeddings,
                                  chunked_docs=chunked_docs, docid_to_url=docid_to_url, top_k=8, retrieval_threshold=0.25))
        if i % 10 == 0:
            print(f"  {i} / {len(tf_df)} questions ({time.time() - t0:.0f} s)")
    return pd.DataFrame(rows)

vision_runs = {"text chunks only": "none", "chunks + crops from the image index": "image_index", "chunks + crops behind the text chunks": "text_index"}
vision_files = {images_from: f"wattbot_solutions_vision_{images_from}.csv" for images_from in vision_runs.values()}

t0 = time.time()
for label, images_from in vision_runs.items():
    local_path = os.path.join(output_dir, vision_files[images_from])
    shared_key = f"wattbot/solutions/{vision_files[images_from]}"
    if use_instructor_embeddings and _s3_object_exists(shared_bucket, shared_key):   # in class: the instructor's run of this cell, scored in the next
        s3.download_file(shared_bucket, shared_key, local_path)
        print(f"{label}: downloaded the instructor's results")
        continue
    print(f"{label}: running")
    run_tf(images_from).to_csv(local_path, index=False)
    s3.upload_file(local_path, bucket_name, shared_key)   # kept in your bucket too, next to the Step 11 predictions
    print(f"{label}: done at {(time.time() - t0) / 60:.1f} min")

Score all three the same way as Step 12. The scorer compares the rows it finds, so the breakdown is over these questions only; the Table and Figure lines are the ones to read, across the three runs and against the DeepSeek run.

PYTHON

for label, images_from in vision_runs.items():
    print("=" * 96)
    print(f"Nova 2 Lite, {label}")
    compute_wattbot_score(train_qa_path=train_qa_path, preds_path=os.path.join(output_dir, f"wattbot_solutions_vision_{images_from}.csv"), n_examples=3)

Callout

What to expect, and what it cost

Measured once on this corpus, 50 table and figure questions, Nova 2 Lite throughout. Retrieval at the paper level was a tie that the text index won: 0.94 against 0.89 on table questions, 0.94 against 0.88 on figures, and 0.82 against 0.65 on quote-only questions, where a picture has nothing to offer. Then reading. Text chunks alone scored 0.64, with 57% of table answers and 59% of figure answers correct. Chunks plus the image index’s top three crops scored 0.75, with 69% of tables and 71% of figures correct and citation F1 up from 0.85 to 0.93. Chunks plus the crops on the pages behind the top text chunks, the route the paper-level hit rate seemed to favour, scored 0.67. Two readings. Showing the model the table or figure is worth about ten points on the questions that depend on one, tables included, which whole pages at 100 dpi had not managed. And finding the paper is not finding the table: the image index chose better crops than the text route did, because a question about a table matches the picture of that table better than it matches the prose on the same page. Fifty questions is a small sample, so read the direction, not the decimals. For the bill, every crop is a few hundred input tokens at Nova 2 Lite’s rate on the pricing page, so a run with images costs two to three times the control and takes three to four times as long: 1.6, 4.0 and 3.0 minutes for the three runs here. vision_usage holds all three runs’ usage blocks, and the embedding of the 2,213 crops is the per-image rate times that count, a few cents.

Wrap‑up: Bedrock versus GPU‑based runs

You now have a complete WattBot evaluation running on Bedrock, with per‑token billing and nothing left running. The Processing Jobs episode reruns the same evaluation on a self‑hosted model, and the optional GPU-notebook episode does it once more interactively, so by the end you will have three versions to compare:

Bedrock (this episode) using hosted embedding and generation models with per‑token billing.
SageMaker Processing job running an open‑source model in batch with on-demand compute.
Notebook GPU instance (optional) using the same locally loaded open‑source model.

When deciding between these options in practice:

Use Bedrock or other hosted APIs when:
- You want to try the latest frontier models quickly.
- You only need to run a modest number of questions, or you are still prototyping.
- You prefer a simple, token‑based cost model and don’t want to manage GPU capacity.
Use self‑hosted models on GPU instances when:
- You need a specific open-weight model that the catalog does not carry.
- You expect to run large batches repeatedly (e.g., many thousands of questions).
- You want tight control over which architectures/checkpoints you run or fine‑tune.
- You already have institutional access to cost‑effective on‑prem or cloud GPUs.

The core RAG evaluation logic stays identical across all three episodes, which is the main takeaway: once you have a clean retrieval + normalization pipeline (like WattBot’s, in wattbot_rag.py), swapping out the generator is mostly a matter of re‑implementing embed_query and generate for each compute option you care about.

Concluding remarks: Bedrock models are one piece of the RAG puzzle

In this episode we used Bedrock-hosted models for both embedding and generation. Larger, higher-quality models can definitely help a ton — especially on messy real-world questions — but it’s important to remember that they are still just one component in your RAG system.

Bigger or newer models do not magically fix weak retrieval. If your chunks are poorly aligned with the questions, a very strong LLM will still struggle.
Most of the long‑term accuracy gains in RAG systems come from the plumbing around the LLMs, including:
- smarter / semantic chunking strategies
- good metadata and filtering
- reranking or multi‑stage retrieval
- domain‑specific heuristics and post‑processing
Cost and latency live in tension with quality. Larger models (or higher token budgets) often improve answers, but at the cost of more inference time and higher per‑request spend. Bedrock makes it easier to experiment with that tradeoff by switching models without rewriting your pipeline.

As you adapt this notebook to your own projects, treat the LLM choice as one tunable component in a larger system. Iterating on chunking, indexing, and retrieval policies will almost always give you more headroom than swapping between already-good models.

Key Points

A complete RAG pipeline on AWS needs no GPU: a CPU notebook plus Bedrock-hosted embedding and generation models is enough.
Corpus preparation (unpacking the staged PDFs, page extraction, chunking) is cheap CPU work; cache chunks.jsonl on S3 so later episodes and runs can reuse it.
Bedrock calls carry no tags; route them through a tagged application inference profile or the spend is anonymous on the bill.
converse is the provider-neutral way to call a chat model on Bedrock; its usage block is what you are billed on. DeepSeek V3 is the generator here, and any catalog model is a one-line swap.
Bedrock bills per token and leaves nothing running between calls, which is why it is the default route for RAG inference.
The model is one tunable component; most accuracy gains come from chunking, retrieval, and post-processing.

Content from RAG with Processing Jobs

Last updated on 2026-10-21 | Edit this page

Estimated time: 50 minutes

Overview

Questions

When does a RAG pipeline need a self-hosted, open-weight model instead of Bedrock, and what is the safest way to rent the GPU for it?
How do we run the embedding and generation steps of the WattBot pipeline as SageMaker Processing jobs that end on their own?
How do we pass data in and out of a job through S3, and tag the job so its cost is traceable?

Objectives

Explain why Bedrock’s catalog is not enough when you need a specific open-weight model, and why a Processing job is the default way to self-host one for batch work.
Launch a GPU Processing job that embeds the cached WattBot chunks with a self-hosted embedding model and writes the vectors to S3.
Launch a second Processing job that runs Qwen3-8B (an open-weight model in a size Bedrock does not carry) with transformers and runs the full RAG loop over the training questions.
Read a job’s inputs, outputs, arguments and runtime cap, and relate its duration to what it cost.
Score the batch predictions with the WattBot metric and compare against the Bedrock run.

Open the notebook from: /ML_with_AWS_SageMaker/notebooks/12_RAG-Processing-jobs.ipynb. Select the conda_pytorch kernel. ## When Bedrock does not have your model

The Bedrock episode is the default for RAG: per-token billing, nothing left running, two API calls. Bedrock’s catalog is curated, though. It carries Amazon’s models, Anthropic, Meta, Mistral, DeepSeek, and a few others, and within a family only some sizes, usually a generation behind: for Qwen, the serverless lineup is still the Qwen3 series. Hugging Face hosts well over a million open-weight checkpoints. The one you want (the newest release, a specific size, a fine-tune for scientific text, a model a paper just published) is usually not on Bedrock, and often never will be.

When that happens you host the model yourself on a GPU you rent, and the question becomes how. There are three ways to rent a GPU on SageMaker:

	Bills	Stops when	Best for
GPU notebook instance	by the hour from start to stop	you stop it	interactive debugging, prototyping at the keyboard
Processing or training job	by the second while the job runs	the script ends, or `max_runtime_in_seconds`	batch work: embed a corpus, answer a list of questions, train
Inference endpoint	by the hour while deployed	you delete it	serving a model to an application with steady traffic

For research RAG, almost all the work is batch: embed everything once, answer a list of questions, score them. That is a job, and a job is the one option that cannot be forgotten. This episode runs the WattBot pipeline that way. The notebook itself is a small CPU instance and never touches the model; it stages data, launches two GPU jobs, and reads the results back.

The generator is Qwen3-8B, an 8-billion-parameter open-weight model (Apache 2.0). Bedrock serves the 32B and 235B members of that family but not this size, which is the usual situation: the size you want is the one missing. The job runs it with transformers, the same library the fine-tuning episode trained with: from_pretrained downloads the weights onto the GPU, and generate writes the answer to a chat prompt, the self-hosted counterpart of what Bedrock does behind converse. Inference needs no gradients, so the weights are the memory cost: 8B parameters at two bytes each in bf16 is about 16 GB, which fits the 24 GB GPU of an ml.g5.xlarge with room for the KV cache. (The fine-tuning episode has the full memory arithmetic; the short version is that inference needs only the weights and training needs eight times more per parameter.)

Why bother serving a model yourself when Bedrock is a call away? Because open-weight models that fit one GPU are now close to the frontier on the kinds of tasks research code runs, on data that never leaves your account. The newest of this family, Qwen3.8-27B, fits one GPU and lands within a few points of a frontier commercial model:

Paired horizontal bars comparing a 27B open-weight model on one GPU with a frontier commercial model; within a few points on agent tasks and science questions, ahead on chart reasoning

Qwen figures are from the Qwen3.8-27B model card and were not reproduced here; Opus 4.8 figures are from its published evaluations.

Callout

What this episode costs

Job	Instance	`$`/h	Expected duration	Cost
Embed ~9,000 chunks	`ml.g5.xlarge` (A10G, 24 GB)	`$1.41`	~10 min incl. startup	~`$0.20`
RAG over a 53-question sample with Qwen3-8B	`ml.g5.xlarge` (A10G, 24 GB)	`$1.41`	~20 min: 1 provisioning, ~4 downloading and loading 16 GB of weights, ~13 generating	~`$0.40`
Controller notebook	`ml.t3.medium`	`$0.05`	2 h	`$0.10`

About $0.70 for the episode, and the GPUs are off the moment each script ends. The RAG job answers the same stratified sample of questions as the Bedrock episode (ten per evidence type, 53 questions once the two papers the corpus lacks are dropped); all 220 is about an hour of generation, so run that at home. Both jobs carry max_runtime_in_seconds, so a hung job costs at most the cap. Compare the same pipeline on a GPU notebook instance (the optional GPU-notebook episode at the end of the lesson): the meter runs until someone stops it, and a forgotten ml.g5.xlarge is $34 a day.

The trade-off is latency. A job takes minutes to start, so this pattern is for batch questions, not a chatbot. For interactive use, Bedrock answers in seconds; for interactive use of a model Bedrock does not have, an inference endpoint does, at an hourly rate you have to remember to turn off.

In this episode we will:

Reuse the chunks cached in your bucket by the Bedrock episode and the inputs in the shared bucket.
Run two short-lived GPU Processing jobs: one computes embeddings for all chunks; the second runs Qwen3-8B with transformers and the full RAG loop (retrieval + generation) over all training questions.
Score the result and compare it with the Bedrock run.

Setup

The scripts

Jobs run in containers, so the code they execute has to be shipped with the job: a Processing job cannot import functions defined in this notebook. The lesson repository has the three files we need as standalone scripts. Clone it:

scripts/embedding_inference.py: generic embedding script (JSONL with a text field in, embeddings.npy out).
scripts/hf_embedder.py: the embedding model wrapper both jobs use, on plain transformers (mean pooling and normalisation, thirty lines), so neither job has to install anything for it.
scripts/wattbot_rag_batch.py: the WattBot RAG loop as a batch job: loads the inputs, runs the generator, writes wattbot_solutions.csv.
scripts/wattbot_rag.py: the shared pipeline logic (retrieval, prompts, answer parsing, run_single_qa, the scorer) that the Bedrock episode imported too. The batch script imports it, so it ships with the job; the notebook imports it for scoring.

The cell below copies them out of the lesson clone into the project folder, one directory per job: code_embed/ for the embedding job and code_rag/ for the RAG job. The split follows from how a job gets its code. run(source_dir=...) uploads that whole directory into the container, the same way the training estimator’s source_dir did in the fine-tuning episode, so each directory holds exactly the files its job imports. The copies of wattbot_rag.py and Score.py next to the notebook are for the scoring step at the end, which runs in the notebook, not in a job.

PYTHON

!mkdir -p /home/ec2-user/SageMaker/wattbot   # the same project folder the Bedrock episode used, beside the lesson clone
%cd /home/ec2-user/SageMaker/wattbot

# One source directory per job: SageMaker uploads the whole directory into the container
!mkdir -p code_embed code_rag
# The embedding job's script
!cp /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/embedding_inference.py /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/hf_embedder.py code_embed/
# The RAG job's script plus the shared module it imports (the container cannot see this notebook)
!cp /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/wattbot_rag_batch.py /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/wattbot_rag.py /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/hf_embedder.py code_rag/
# The notebook's own copies, for scoring at the end: the pipeline module and the competition's metric file
!cp /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/wattbot_rag.py /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/Score.py .
!ls code_embed code_rag

Next, set up the AWS SDK, SageMaker session, and bucket names.

PYTHON

import os
import json

import boto3
import time
import numpy as np
import pandas as pd
import sagemaker
from sagemaker import get_execution_role
from sagemaker.huggingface import HuggingFaceProcessor         # launches a job in the prebuilt Hugging Face container
from sagemaker.processing import ProcessingInput, ProcessingOutput   # S3 <-> container path mappings for a job

session = sagemaker.Session()
region = session.boto_region_name
role = get_execution_role()   # the IAM role the jobs run under to read and write S3

shared_bucket = "ml-marathon-shared"   # the instructor-staged, read-only bucket with the data and model weights; change only on your own account
bucket_name = "sinkorswim-doejohn"     # <-- EDIT ME: your own bucket (chunks in, outputs out)
s3_client = boto3.client("s3", region_name=region)

local_data_dir = "./data"
os.makedirs(local_data_dir, exist_ok=True)

print("Region:", region)
print("Role:", role)
print("Bucket:", bucket_name)

Step 1 – Check the inputs

The jobs read three files from S3: train_QA.csv and metadata.csv from the shared bucket, and chunks.jsonl from your bucket, where the Bedrock episode cached it. The PDFs themselves are not needed again: chunking was done once, and the jobs work from the chunks. We download the three files to the notebook only to confirm they exist and to look at them; the jobs fetch their own copies from S3.

First the keys and the full S3 URIs. The URIs are the form a ProcessingInput takes, so they are built here once and reused when the jobs are launched.

PYTHON

train_qa_key = "wattbot/train_QA.csv"
metadata_key = "wattbot/metadata.csv"
chunks_s3_key = "wattbot/chunks.jsonl"

# Full S3 URIs: this is the form a ProcessingInput takes
train_qa_s3 = f"s3://{shared_bucket}/{train_qa_key}"
metadata_s3 = f"s3://{shared_bucket}/{metadata_key}"
chunks_s3_uri = f"s3://{bucket_name}/{chunks_s3_key}"

def smart_read_csv(path: str) -> pd.DataFrame:
    """Read a CSV as utf-8, falling back to latin-1 for the odd accented author name."""
    try:
        return pd.read_csv(path)
    except UnicodeDecodeError:
        return pd.read_csv(path, encoding="latin-1")

Then download the three files, skipping any already on disk, and load them so we can check the counts before spending GPU time on them.

PYTHON

train_qa_path = os.path.join(local_data_dir, "train_QA.csv")
metadata_path = os.path.join(local_data_dir, "metadata.csv")
local_chunks_path = os.path.join(local_data_dir, "chunks.jsonl")

if not os.path.exists(train_qa_path):
    s3_client.download_file(shared_bucket, train_qa_key, train_qa_path)
if not os.path.exists(metadata_path):
    s3_client.download_file(shared_bucket, metadata_key, metadata_path)
if not os.path.exists(local_chunks_path):
    s3_client.download_file(bucket_name, chunks_s3_key, local_chunks_path)   # written by the Bedrock episode

train_df = smart_read_csv(train_qa_path)
metadata_df = smart_read_csv(metadata_path)
with open(local_chunks_path, "r", encoding="utf-8") as f:
    chunked_docs = [json.loads(line) for line in f]

print("Train QAs:", len(train_df))
print("Metadata rows:", len(metadata_df))
print("Chunks:", len(chunked_docs), "from", chunks_s3_uri)
train_df.head(3)

If the last download fails with NoSuchKey, go back to the Bedrock episode and run its chunking step; this episode starts from its output.

Step 2 – Processing Job 1: embed all chunks on a GPU

Now we launch a short-lived Hugging Face Processing job that:

Downloads chunks.jsonl from S3.
Loads thenlper/gte-large from Hugging Face with the container’s own transformers.
Encodes each chunk into an embedding vector.
Saves the full matrix as embeddings.npy back to S3.

We use the same embedding_inference.py script across projects; here it expects a JSONL file with a text field.

Nothing to install

Neither job installs anything. The Hugging Face container ships transformers and PyTorch, which is all the embedding wrapper and the generator need, so both jobs start as soon as the instance is up. A requirements.txt next to a script would be pip-installed before it runs; neither job here has one.

Cost tracking: tag every job

Tags do not propagate in AWS. The tags you put on your notebook instance cover the notebook’s hourly cost and nothing else. Each Processing job below launches its own GPU instance, which is a separate billable resource with no tags unless we pass them at launch, exactly as we did for training and tuning jobs earlier in the workshop. An untagged job shows up in Cost Explorer as anonymous SageMaker spend that nobody can attribute to a team.

Define the tags once and pass tags=job_tags to every HuggingFaceProcessor you create. Use a distinct Purpose per job so you can tell the embedding job from the batch RAG job on the bill.

PYTHON

name = "DoeJohn"         # replace with your name
project = "SinkOrSwim"   # replace with your challenge name

def make_job_tags(purpose: str):   # one of: Storage, Notebook, Training, Tuning, Inference
    return [
        {"Key": "Name", "Value": name},
        {"Key": "Project", "Value": project},
        {"Key": "Purpose", "Value": purpose},   # differs per job so the bill tells them apart
    ]

Launch the embedding job

Open scripts/embedding_inference.py first; it is under fifty lines. It reads a JSONL file of chunks from the job’s input directory, embeds them with whatever Hugging Face model --model_id names, through the thirty-line wrapper in hf_embedder.py, normalises the vectors so a dot product is a cosine similarity, and saves one embeddings.npy to the output directory. Nothing in it knows about WattBot, which is the point of a Processing job script: inputs in, outputs out, no state.

A HuggingFaceProcessor is the Processing-job counterpart of the estimator: the same role, tags, instance type and container versions, plus a runtime cap. Its run call names the script, mounts the inputs, and says where the output goes. Once it is launched, the job appears on the Processing jobs page of the SageMaker console (under Data preparation in the left menu), the Processing counterpart of the Training jobs page you used earlier.

In class the instructor has run this job already and staged its output. The switch below copies that matrix into your bucket, at the key the job would have written, and the launch cell then skips the job: the second job, which you do launch, has the same shape and is where the class spends its GPU minutes. On your own, set it to False and run the job; it is eight minutes.

PYTHON

def _s3_object_exists(bucket: str, key: str) -> bool:
    try:
        s3_client.head_object(Bucket=bucket, Key=key)
        return True
    except s3_client.exceptions.ClientError:
        return False

embedding_model_id = "thenlper/gte-large"
emb_output_prefix = "wattbot/embeddings"
emb_output_path = f"s3://{bucket_name}/{emb_output_prefix}/"

use_instructor_embeddings = True                       # True in class: copy the instructor's gte-large matrix and skip the job. False: run the job
shared_gte_key = "wattbot/embeddings_gte_large.npy"    # where the instructor staged it in the shared bucket

run_embedding_job = not (use_instructor_embeddings and _s3_object_exists(shared_bucket, shared_gte_key))
if not run_embedding_job:
    s3_client.copy({"Bucket": shared_bucket, "Key": shared_gte_key}, bucket_name, f"{emb_output_prefix}/embeddings.npy")
    print(f"Copied the instructor's embeddings to s3://{bucket_name}/{emb_output_prefix}/embeddings.npy; the embedding job is skipped")

PYTHON


embedding_processor = HuggingFaceProcessor(
    base_job_name="WattBot-embed-gte-large",   # SageMaker appends a timestamp to make each job name unique
    role=role,
    tags=make_job_tags("Inference"),   # tags are per job; nothing is inherited from the notebook
    instance_type="ml.g5.xlarge",      # the GPU the job rents (A10G, 24 GB): plenty for a 335M-parameter embedder
    instance_count=1,
    transformers_version="4.56",       # these three pick the prebuilt container image
    pytorch_version="2.8",
    py_version="py312",
    sagemaker_session=session,
    max_runtime_in_seconds=2 * 60 * 60,   # 2 hours, in seconds (2 h x 60 min x 60 s); the estimate is ~8 minutes. SageMaker kills the job here, so a hang costs at most this
)

if run_embedding_job:   # False in class when the instructor's matrix was copied in above
    # run() uploads the code, starts the instance, streams the job's log here and returns when the script ends (wait=True is the default)
    start = time.time()
    embedding_processor.run(
        code="embedding_inference.py",     # the script the container runs...
        source_dir="code_embed/",          # ...uploaded with everything beside it (hf_embedder.py)
        inputs=[
            ProcessingInput(
                source=chunks_s3_uri,                       # S3 object (or prefix) copied in before the script starts
                destination="/opt/ml/processing/input",    # where it appears inside the container
            )
        ],
        outputs=[
            ProcessingOutput(
                output_name="embeddings",
                source="/opt/ml/processing/output",        # whatever the script writes here...
                destination=emb_output_path,               # ...is copied to this S3 prefix when the script ends
            )
        ],
        arguments=[                        # command-line flags, as the script's argparse expects them
            "--model_id", embedding_model_id,
            "--input_filename", "chunks.jsonl",
            "--text_key", "text",
            "--input_dir", "/opt/ml/processing/input",
            "--output_dir", "/opt/ml/processing/output",
        ],
    )

    print(f"Embedding job complete: {(time.time() - start) / 60:.1f} min wall-clock, including instance startup")

Check on the running job in the AWS Console

To view the job running from the AWS Console, you can visit SageMaker AI, and then find the “Data Preparation” dropdown menu on the left side panel. Click that to find “Processing jobs”, or open this link in the class region (change us-east-2 in it for another region).

It takes about 10 minutes in total for the job to complete: a couple of provisioning the instance and downloading the embedding model, eight of embedding. This is the downside of launching jobs, but the good news is that we only need to launch one embedding job for our RAG pipeline. This strategy also ensures we’re only paying for the GPU while the job runs.

Sanity-check the embeddings locally

We can download embeddings.npy back into the notebook and inspect its shape to confirm the job ran successfully.

PYTHON

local_embeddings_path = os.path.join(local_data_dir, "embeddings.npy")
embeddings_key = f"{emb_output_prefix}/embeddings.npy"   # the ProcessingOutput destination + the filename the script wrote

s3_client.download_file(bucket_name, embeddings_key, local_embeddings_path)
chunk_embeddings = np.load(local_embeddings_path)

print("Embeddings shape:", chunk_embeddings.shape, "for", len(chunked_docs), "chunks")

Step 3 – Processing Job 2: full WattBot RAG over all questions

For the second job, we pass four inputs:

chunks.jsonl – serialized chunks (your bucket)
embeddings.npy – precomputed chunk embeddings (your bucket, written by Job 1)
train_QA.csv – training questions (shared bucket)
metadata.csv – to resolve ref_id → URL (shared bucket)

The --per_type 10 argument draws the same stratified sample the Bedrock episode scored, ten questions per evidence type with the same seed, so the two runs are compared on the same questions; leave it out to answer them all. Open scripts/wattbot_rag_batch.py and read main before launching it. It is short, because the pipeline lives in scripts/wattbot_rag.py, the same module the Bedrock episode imported:

retrieve_top_k: cosine similarity over the embedding matrix
run_single_qa: retrieval, the answer call, the hybrid unanswerable rule (retrieval threshold OR the model’s own is_blank), the explanation call
compute_wattbot_score: the metric we use at the end

What the batch script adds is the two model callables: embed_query wraps the same gte-large model and pooling the embedding job used (hf_embedder.py) (the question has to be embedded with the same model as the chunks), and generate_batch wraps transformers.generate on the generator (Qwen3-8B here) with Qwen’s thinking mode off, taking several prompts at once. The loop itself is run_batch_qa, the batched twin of the run_single_qa the Bedrock episode used: the same retrieval, prompts and unanswerable rule, with the model calls grouped eight at a time because a GPU decoding one prompt is idle most of the time. The job writes wattbot_solutions.csv in the WattBot submission format.

Callout

Why not vLLM?

vLLM is the open-source inference engine most people serve open-weight models with: you give it a model name, it loads the weights onto the GPU and answers chat requests, batching many at once so that throughput scales. That is what you want behind an API with users. This job has about 100 prompts and no users waiting, and transformers.generate, already in the container, can decode several prompts at once too: run_batch_qa in wattbot_rag.py hands them to the model eight at a time (--batch_size), which is where most of vLLM’s speed on a batch like this comes from, with nothing to install. vLLM’s install is several gigabytes of billed GPU time before the first answer. The GPU-notebook episode uses vLLM, where its 8-bit weights are needed to fit a larger model.

PYTHON

emb_output_s3 = f"s3://{bucket_name}/{emb_output_prefix}/embeddings.npy"

print("chunks:    ", chunks_s3_uri)
print("embeddings:", emb_output_s3)
print("train_QA:  ", train_qa_s3)
print("metadata:  ", metadata_s3)

If the launch below fails at once with ResourceLimitExceeded naming an instance type, nothing in the notebook is wrong: the account’s quota for that type is too low, and the increase is requested in Service Quotas (days, for GPU types). On your own account, request ml.g5.xlarge for processing job usage before you start.

What launching a job means

This is the first GPU job you launch yourself in this episode (the embedding job above was staged for you in class), so here is the whole mechanism in plain terms. You describe a machine and point at a script. SageMaker rents the machine, copies the files you named from S3 onto its disk, runs your script on it, copies whatever the script wrote in its output folder back to S3, and shuts the machine down. The notebook does nothing but wait and show you the script’s printed output. Nothing from this notebook’s memory goes along: not chunked_docs, not train_df, not any function defined above. The script sees only the files you copied in and the arguments you gave it, which is why every input is an S3 location and every setting is a command-line flag. That is the same thing estimator.fit did in the fine-tuning episode, with one difference: a training job hands back a trained model, a Processing job hands back whatever files the script wrote, here a CSV of answers.

The description is two cells. The first builds a HuggingFaceProcessor: which machine, which container image, how long it may run, whose tags it carries. Two of its settings matter more than on the embedding job. volume_size_in_gb=60 attaches a 60 GB volume to the instance, mounted under /opt/ml/processing; the container’s own root disk is small, so the 16 GB of model weights must not land there. env={"HF_HOME": ...} points Hugging Face’s download cache at that volume, which is where the weights go.

PYTHON

rag_output_prefix = "wattbot/solutions"                       # where this job's output lands in your bucket
rag_output_path = f"s3://{bucket_name}/{rag_output_prefix}/"

generation_model_id = "Qwen/Qwen3-8B"   # the Hugging Face name the job downloads; ~16 GB of bf16 weights

rag_processor = HuggingFaceProcessor(
    base_job_name="WattBot-rag-batch",   # SageMaker appends a timestamp; this is the name you look for on the Processing jobs page
    role=role,                           # the IAM role the job runs under, to read your bucket and the shared one
    tags=make_job_tags("Inference"),     # separate job, separate tags: nothing is inherited from the embedding job
    instance_type="ml.g5.xlarge",        # 1 x A10G, 24 GB: holds the bf16 8B model plus the KV cache
    instance_count=1,
    volume_size_in_gb=60,                # disk attached to the instance, mounted at /opt/ml/processing; the weights go here
    env={"HF_HOME": "/opt/ml/processing/hf_cache"},   # environment variable in the container: Hugging Face downloads to this path
    transformers_version="4.56",         # these three pick the prebuilt container image, the same one the embedding job used
    pytorch_version="2.8",
    py_version="py312",
    sagemaker_session=session,
    max_runtime_in_seconds=2 * 60 * 60,  # the cap: SageMaker kills the job after 2 hours; the estimate is 20 minutes
)

The processor is only a description; nothing runs until run. Its arguments are the three things a job needs: the code (code and source_dir), the data (inputs, each an S3 location and the container path it is copied to before the script starts), and where the results go (outputs, a container path copied to S3 when the script ends). The arguments list is the script’s command line, in the form its argparse expects.

PYTHON

start = time.time()
rag_processor.run(
    code="wattbot_rag_batch.py",       # the script the container runs...
    source_dir="code_rag/",            # ...uploaded with everything beside it: wattbot_rag.py and hf_embedder.py, which it imports
    inputs=[                           # each input: an S3 object, and the container path it is copied to before the script starts
        ProcessingInput(
            source=chunks_s3_uri,                              # the chunks, from your bucket
            destination="/opt/ml/processing/input/chunks",     # four inputs, each in its own subdirectory
        ),
        ProcessingInput(
            source=emb_output_s3,                              # the gte-large matrix the embedding job wrote
            destination="/opt/ml/processing/input/embeddings",
        ),
        ProcessingInput(
            source=train_qa_s3,                                # the questions, from the shared bucket
            destination="/opt/ml/processing/input/train",
        ),
        ProcessingInput(
            source=metadata_s3,                                # the paper index, to fill the ref_url column
            destination="/opt/ml/processing/input/metadata",
        ),
    ],
    outputs=[
        ProcessingOutput(
            output_name="solutions",
            source="/opt/ml/processing/output",                # whatever the script writes here...
            destination=rag_output_path,                       # ...is copied to this S3 prefix when the script ends
        )
    ],
    arguments=[                        # the script's command line: python wattbot_rag_batch.py --input_dir ... --top_k 8 ...
        "--input_dir", "/opt/ml/processing/input",             # where the inputs above were mounted
        "--output_dir", "/opt/ml/processing/output",           # where the script must write for the output copy to pick it up
        "--embedding_model_id", embedding_model_id,            # the same model that embedded the chunks, to embed each question
        "--qwen_model_id", generation_model_id,                # the generator
        "--top_k", "8",                                        # chunks handed to the model per question
        "--retrieval_threshold", "0.25",                       # best cosine below this => unanswerable
        "--per_type", "10",                                    # the same stratified sample the Bedrock episode scored (seed 0); drop it for all the questions
        "--batch_size", "8",                                   # prompts decoded together; 1 would be one at a time, several times slower
    ],
)

print(f"RAG batch job complete: {(time.time() - start) / 60:.1f} min wall-clock, including instance startup and the weight download")

While it runs, read what the container is printing. Open the Processing jobs page, click the job’s name, scroll to Monitor and click View logs. That opens CloudWatch, the AWS service that stores every job’s printed output, at a list with one entry named after your job; click it. The text is the container’s stdout, kept after the job ends, updated while it runs (refresh). You will see, in order: the weight download (16 GB in shards), Loading model with transformers, a line per batch of answers and of explanations, then one line per question. The whole job is the 20 minutes in the cost callout; the generation is the second half of it.

Step 4 – Download predictions and evaluate

Finally, we download wattbot_solutions.csv from S3, inspect a few rows, and compute the WattBot score against train_QA.csv with the same compute_wattbot_score the Bedrock episode used.

PYTHON

solutions_key = f"{rag_output_prefix}/wattbot_solutions.csv"
local_solutions_path = os.path.join(local_data_dir, "wattbot_solutions.csv")

s3_client.download_file(bucket_name, solutions_key, local_solutions_path)
solutions_df = pd.read_csv(local_solutions_path)
solutions_df.head()

Score the predictions with the competition’s own Score.py, through the same wrapper the Bedrock episode used, so the two runs compare directly.

PYTHON

from wattbot_rag import compute_wattbot_score

results_df = compute_wattbot_score(
    train_qa_path=train_qa_path,
    preds_path=local_solutions_path,
    n_examples=10,
)

Compare the overall score with the Bedrock run. The retrieval logic, prompts and scoring are identical, so any difference is the two models: gte-large versus Titan for retrieval, Qwen3-8B versus DeepSeek V3 for generation. In one run on the same 53 questions, Qwen3-8B scored 0.57 against DeepSeek V3’s 0.66: its citations were better (F1 0.70 against 0.57) and its answers worse (51% against 66% correct), with the cross-paper and calculation questions losing the most. An 8B model against a 671B one is not a fair fight, and the gap is the price of a model you host yourself on one small GPU.

Challenge

What did the two jobs cost?

Ask SageMaker what it recorded for each job: session.sagemaker_client.describe_processing_job(ProcessingJobName=...) returns CreationTime, ProcessingStartTime and ProcessingEndTime. The job names are embedding_processor.latest_job.job_name and rag_processor.latest_job.job_name. For each job, compute the minutes between creation and start (provisioning, before the instance is billed) and between start and end (the billed window), and price the billed window at $1.41/h, the ml.g5.xlarge rate both jobs used. Then, from the RAG job’s CloudWatch log timestamps, estimate how much of its billed window was setup (the weight download and model load) and how much was generation. How does the total compare with the cost callout, and with the Bedrock run?

Show me the solution

describe_processing_job returns three timestamps; the gap between the first two is unbilled provisioning, between the last two the billed run.

PYTHON

def job_cost(job_name: str, price_per_hour: float):
    d = session.sagemaker_client.describe_processing_job(ProcessingJobName=job_name)   # the record SageMaker keeps for the job
    # CreationTime -> ProcessingStartTime: provisioning the instance, not billed
    provisioning = (d["ProcessingStartTime"] - d["CreationTime"]).total_seconds() / 60
    # ProcessingStartTime -> ProcessingEndTime: the script running, billed by the second
    billed = (d["ProcessingEndTime"] - d["ProcessingStartTime"]).total_seconds() / 60
    cost = billed / 60 * price_per_hour
    print(f"{job_name}: {d['ProcessingJobStatus']}")
    print(f"  wall-clock:                {provisioning + billed:.1f} min from launch to end")
    print(f"  provisioning (not billed): {provisioning:.1f} min")
    print(f"  billed window:             {billed:.1f} min at ${price_per_hour}/h = ${cost:.2f}")
    print(f"  output:                    {d['ProcessingOutputConfig']['Outputs'][0]['S3Output']['S3Uri']}")   # where the job's files went
    return cost

total = job_cost(embedding_processor.latest_job.job_name, 1.41) if run_embedding_job else 0.0   # no job to price when the matrix was copied in
total += job_cost(rag_processor.latest_job.job_name, 1.41)
print(f"GPU total for the episode: ${total:.2f}")

Expect the embedding job to bill about 8 minutes (about $0.20) and the RAG job about 17 minutes (about $0.40), so about $0.60 of GPU time, the callout’s figure. In the RAG job’s log, the first 4 minutes of the billed window are the download and the model load; the rest is the questions, two generation calls each. Setup is a third of what you paid for, which is the argument for keeping an HF_HOME cache on S3 (skip the download) once a job is run routinely. The Bedrock run of the same questions cost cents, and the self-hosted route is only justified here because the model is not in the catalog.

Challenge

The runtime cap

Both processors set max_runtime_in_seconds=2 * 60 * 60. Predict what would happen if the RAG job were launched with max_runtime_in_seconds=600 instead: what does the job status become, what is in the output prefix, and what did it cost? Then compute the most a hung job can bill under each job’s actual cap, and compare it with the cost of an ml.g6e.xlarge notebook instance left running over a weekend (64 hours).

Show me the solution

With a 600-second cap the RAG job is killed by SageMaker ten minutes after ProcessingStartTime, while it is still downloading or loading the weights. The status becomes Stopped (the failure reason reads MaxRuntimeExceeded), the output prefix is empty because the script never reached the point of writing wattbot_solutions.csv, and the bill is the full 600 seconds: 600 / 3600 × $1.41 ≈ $0.24 for nothing. A cap has to leave room for the setup that precedes the work, which is why ours is six times the 20-minute estimate rather than ten percent over it.

Under the caps we set, the worst case for a job that hangs (a stuck download, a model that never finishes loading) is 2 h × $1.41 = $2.82 for either job. An ml.g5.xlarge notebook instance forgotten over a weekend is 64 h × $1.41 = $90, and nothing caps it. That is the difference between a job and an instance: the job’s worst case is a number you chose before you launched it.

Key Points

Bedrock’s catalog is curated and lags releases; when you need a specific open-weight model you rent a GPU and serve it yourself.
Of the ways to rent a GPU, a Processing job is the one that cannot be forgotten: it bills while the script runs and ends on its own, with max_runtime_in_seconds as the backstop. Use it for batch work; use Bedrock or an endpoint when you need answers in seconds.
Inference needs only the weights: Qwen3-8B in bf16 is ~16 GB and fits one 24 GB GPU.
A job’s inputs and outputs are S3 prefixes mounted into the container; the notebook only stages files and reads results, so it stays a cheap CPU instance.
The pipeline logic (wattbot_rag.py) ships with the job unchanged from the Bedrock episode; only the two model callables differ.
Tags do not propagate from the notebook: every job gets its own tags=.

Content from (Optional) RAG with a Notebook GPU

Last updated on 2026-10-21 | Edit this page

Estimated time: 45 minutes

Overview

Questions

How can we run a basic Retrieval-Augmented Generation (RAG) pipeline entirely from a single GPU-backed SageMaker notebook?
How do we go from cached text chunks to a searchable embedding space using a self-hosted model?
How can we generate WattBot-style answers (including citations and evidence) that follow the competition’s scoring conventions?

Objectives

Verify that our SageMaker notebook instance has a working GPU and compatible Python environment.
Reuse the WattBot metadata and questions from the shared bucket, and the chunks.jsonl the Bedrock episode cached in your own bucket.
Implement a simple, explicit “from scratch” embedding pipeline with a self-hosted embedding model, without relying on FAISS or production vector DBs.
Build a small retrieval helper that finds the most relevant chunks for a question using cosine similarity in embedding space.
Serve a 27B open-weight generator (Qwen3.8-27B in FP8, via vLLM) on the notebook’s GPU and wire it to the retriever to produce WattBot-format answers (including answer, ref_id, ref_url, and supporting_materials).
Explain why inference uses quantized weights rather than bf16, and what a self-hosted GPU costs per hour, per forgotten weekend, and per forgotten month.
Add a second LLM pass that generates short explanations and marks whether the evidence comes from text, figures, tables, or a combination.

Callout

Optional, self-study episode

This episode is not taught in the workshop. The Bedrock and Processing Jobs episodes are the default ways to run RAG; this one is for the case where you need to poke at a self-hosted model interactively, and only if you will remember to stop the instance. Do not run this notebook on the controller instance from the other episodes. Create a new GPU notebook instance for it (Setup below). On a CPU instance the vLLM install and the 27 GB model load will exhaust memory and freeze your JupyterLab. It runs the same pipeline as the Processing Jobs episode on a GPU notebook instance, which is the right tool for prototyping alone at a keyboard and the wrong one for a room of a hundred people with meters running. Work through it on your own when you want to see model loading, GPU memory and generation as plain Python. Read the cost callout first.

When Bedrock does not have your model

In the Bedrock episode we built the full WattBot pipeline on Amazon Bedrock, the recommended default for RAG inference because it bills per token and leaves nothing running between calls. Bedrock’s catalog is curated, though. It has Amazon’s own models, Anthropic, Meta, Mistral, DeepSeek, and a handful of others, and even within a family it carries a few sizes and lags the releases: for Qwen, the serverless lineup is still the Qwen3 generation (Qwen3-32B, Qwen3-235B and the Qwen3-Coder models). Hugging Face hosts well over a million open-weight checkpoints. The one you want (the newest release, a specific size, a specific fine-tune, a model tuned on scientific text) is usually not on Bedrock, and often never will be.

When that happens you host the model yourself on a GPU you rent. This episode does that in its simplest form: a single GPU notebook instance, with Qwen3.8-27B as the generator. It is a 27-billion-parameter open-weight (Apache 2.0) model released in August 2026 that is not offered on Bedrock, large enough to follow the WattBot answer format reliably, and small enough to fit on one 48 GB GPU once its weights are stored in 8-bit floating point. Everything Bedrock abstracts away (loading an embedding model, loading a generator, fitting both in GPU memory) becomes visible as plain Python. Treat it as a learning exercise and a prototyping workbench, not a deployment pattern.

Callout

This instance bills `$2.61` an hour until you stop it, and nothing stops it for you

A GPU notebook instance is a rental with the meter running. It bills by the hour from the moment it starts until the moment you stop it, whether it is generating answers, sitting at an idle prompt, or showing a closed browser tab on a laptop that went home for the weekend. Nothing in this episode turns it off for you. The instance we use, ml.g6e.xlarge (one L40S GPU, 48 GB), costs $2.61 per hour on-demand in us-east-1 (AWS price list, October 2026). That is:

Left running for	Cost
This episode (~2 hours)	~`$5`
Overnight (12 h)	~`$31`
A weekend (64 h)	~`$167`
A week	~`$438`
A month	~`$1,906`

The same evaluation on Bedrock cost well under a dollar. One forgotten weekend on this instance costs a thousand times more than the work it was rented for. A forgotten GPU notebook is the most common way research groups overspend cloud credits.

Two further charges people do not expect:

A stopped instance still bills for its storage. We give this notebook a 100 GB volume for model weights. SageMaker charges $0.14 per GB-month for notebook storage, so a stopped-but-not-deleted instance costs about $14 a month forever. When you are done with the episode, delete the instance, not just stop it, unless you will use it again within days.
Idle is billed the same as busy. There is no discount for a GPU that is loaded and waiting. Loading a 27B model takes a few minutes, which tempts people to leave the instance up “so I don’t have to reload”. That is the $1,906.

The habits that keep this safe:

Stop the instance when you stand up, every time, from the Notebook instances page (select it, Stop). Then look at the status column until it says Stopped.
Set a reminder on your phone or calendar before you start the instance, for two hours from now.
Prefer jobs for anything beyond exploration: a Processing job (the Processing Jobs episode) or a training job ends on its own, and max_runtime_in_seconds (max_run on a training job) guarantees it.

Self-hosting on a notebook GPU is nonetheless the right starting point for understanding the alternatives, because it is:

Easy to understand on the AWS side (pick a GPU instance, pick a model, go)
Ideal for learning what retrieval and generation do, step by step
The fastest loop for experimenting with prompts and models that Bedrock does not carry

The Processing Jobs episode runs this same pipeline as short-lived jobs, so the GPU only bills while it is embedding or generating; that is the pattern to use for real work. Later patterns (inference endpoints, batch transform) build on the same idea.

Overview: WattBot RAG on a single notebook GPU

In this episode we rerun the WattBot 2026 RAG pipeline from the Bedrock episode, this time entirely from a single GPU-backed SageMaker notebook with self-hosted models. The generator is Qwen3.8-27B, the newer and larger relative of the Qwen3-8B the Processing Jobs episode ran: an open-weight model that fits one 48 GB GPU in FP8 and sits within a few points of a frontier commercial model on agent tasks and science questions (see that episode’s chart), which is why it is worth the trouble of hosting.

In this episode we will:

Reuse the WattBot data you prepared in the Bedrock episode.
- Pull train_QA.csv and metadata.csv from the shared bucket and the cached chunks.jsonl from your own.
Implement the model steps explicitly in code.
- Embed chunks with an open embedding model (thenlper/gte-large) running on the notebook’s GPU.
- Implement cosine-similarity search over the embedding matrix.
Serve an open-weight generator that Bedrock does not offer.
- Load Qwen3.8-27B with FP8 weights through vLLM on a single ml.g6e.xlarge GPU.
- Construct WattBot-style answers that we can compare against train_QA.csv.
Add an explanation pass.
- Use an LLM to look at the retrieved evidence, the answer, and citations.
- Generate a short explanation and label the evidence type: [text], [table], [figure] or a combination such as [text+table].

Notebook + dataset setup

For this episode you need a notebook instance with a GPU that can hold the FP8 27B model plus the embedding model plus working memory for generation. Create a new notebook instance (same steps as the controller-notebook episode, same three tags) with:

Instance type: ml.g6e.xlarge (1 × L40S, 48 GB GPU memory, $2.61/h). The FP8 weights are ~27 GB; the rest is for the embedding model and the KV cache that generation needs for long prompts. The L40S runs FP8 natively. A 24 GB instance such as ml.g5.xlarge cannot hold this model; a 16 GB ml.g4dn.xlarge cannot hold anything in this size class.
Volume size: 100 GB (the default of 5 GB will not fit the weights or the vLLM install). This volume bills at $0.14 per GB-month even while the instance is stopped.
Tags: the same Project, Name, Purpose keys as every other resource; Purpose is Notebook.

Once it is InService, open JupyterLab, clone the lesson repository as in the controller-notebook episode (this is a new instance, so it has no copy yet), and open /ML_with_AWS_SageMaker/notebooks/13_RAG-GPU-notebook.ipynb. Select the conda_pytorch kernel.

See Instances for ML for further guidance. This instance bills by the hour from the moment it starts, whether or not it is doing anything: stop it when you leave, and delete it when the episode is over.

We reuse the same two buckets as the Bedrock episode: the shared bucket for metadata.csv and train_QA.csv, and your own bucket for the wattbot/chunks.jsonl you cached there. If you skipped that episode, go back and run its chunking step first. Your notebook instance carries the tags you gave it when you created it and your bucket carries its own tags, so in this episode those two resources are the only things being billed.

Verify GPU and basic environment

PYTHON

# nvidia-smi lists the GPU, its memory and the driver; it fails when the instance has no GPU
!nvidia-smi || echo "No GPU detected – please switch to a GPU-backed instance (ml.g6e.xlarge) before running this notebook."

PYTHON

# also verify you have selected the conda_pytorch kernel
import torch
print("torch cuda available:", torch.cuda.is_available())   # False means PyTorch cannot use the GPU even if nvidia-smi sees it
print("num gpus:", torch.cuda.device_count())

Install vLLM first

We install the inference server before loading anything else, because the install may replace PyTorch and then needs a kernel restart; doing it now means nothing is lost by restarting.

Callout

Installing vLLM takes a few minutes, and the meter is running

vLLM is the open-source inference engine that loads an open-weight model onto the GPU and answers chat requests from Python, the self-hosted counterpart of Bedrock’s runtime (the Processing Jobs episode introduces it). pip install vllm downloads several gigabytes (it brings its own CUDA kernels and may replace the kernel’s PyTorch build) and takes 5–10 minutes. Qwen3.8 needs vLLM 0.17 or newer. The 27 GB of model weights take another 4–6 minutes to download the first time. At $2.61 an hour that is about 60 cents of waiting, which is fine; what is not fine is doing it twice because the instance was deleted and recreated. Keep the Hugging Face cache on the 100 GB volume (the HF_HOME line in the serving cell below does that) so the weights survive a stop/start.

If the install changes PyTorch, restart the kernel (Kernel → Restart) before continuing.

PYTHON

# Several GB; it may replace torch, so restart the kernel (Kernel -> Restart) if the log says it did
%pip install -q "vllm>=0.17" "peft>=0.18"   # Qwen3.8 needs vLLM 0.17+, which brings its own PyTorch; peft upgraded to match the transformers it installs

Import data from bucket into notebook

The same setup as the Bedrock episode: session, region, role, an S3 client, the two bucket names and a local data folder, plus a small download helper.

PYTHON

import os
import json
import time

import boto3
import pandas as pd
import numpy as np

import sagemaker
from sagemaker import get_execution_role


# Initialize SageMaker + AWS basics
session = sagemaker.Session()
region = session.boto_region_name
role = get_execution_role()
s3_client = boto3.client("s3", region_name=region)

print("Region:", region)
print("Role:", role)

PYTHON

def download_s3_object(bucket: str, key: str, local_path: str) -> None:
    os.makedirs(os.path.dirname(local_path), exist_ok=True)
    print(f"Downloading s3://{bucket}/{key} -> {local_path}")
    s3_client.download_file(bucket, key, local_path)

PYTHON

!mkdir -p /home/ec2-user/SageMaker/wattbot   # a project folder beside the lesson clone (this is a new instance, so it starts empty)
%cd /home/ec2-user/SageMaker/wattbot

shared_bucket = "ml-marathon-shared"   # the instructor-staged, read-only bucket with the data and model weights; change only on your own account
bucket_name = "sinkorswim-doejohn"     # <-- EDIT ME: your own bucket (chunks live here)

# Local working directory in the notebook instance
local_data_dir = "./data"

session_start = time.time()   # the meter on this GPU instance started when it booted; this marks where the episode's work began

print("Local data dir:", local_data_dir)

Get the shared WattBot helpers

The retrieval, prompts, answer parsing, per-question orchestration and the scorer are the same code the Bedrock and Processing Jobs episodes used: scripts/wattbot_rag.py in the lesson repository. Copy the module next to this notebook. What this episode adds is the two model callables, embed_query and generate, running on the notebook’s own GPU.

PYTHON

!cp /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/wattbot_rag.py /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/hf_embedder.py /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/Score.py .

Import the pieces of the shared RAG module we reuse: the prompts, the retrieval and scoring functions, and the per-question runner.

PYTHON

from wattbot_rag import (
    ANSWER_SYSTEM_PROMPT,
    build_docid_to_url,
    compute_wattbot_score,
    embed_query_with,
    explanation_system_prompt,
    retrieve_top_k,
    run_single_qa,
)

Download the two WattBot input files from the shared bucket.

PYTHON

# Download metadata.csv and train_QA.csv from the shared bucket
metadata_key = "wattbot/metadata.csv"
train_qa_key = "wattbot/train_QA.csv"

metadata_path = os.path.join(local_data_dir, "metadata.csv")
train_qa_path = os.path.join(local_data_dir, "train_QA.csv")

download_s3_object(shared_bucket, metadata_key, metadata_path)
download_s3_object(shared_bucket, train_qa_key, train_qa_path)

Step 1 – Model libraries and safe CSV loading

PYTHON

import torch
from hf_embedder import HFEmbedder   # thirty lines on plain transformers: encode(list of strings) -> unit-length matrix; the Processing jobs use the same file

A CSV reader that tries a few encodings, because some metadata files are not valid UTF-8.

PYTHON

def smart_read_csv(path: str) -> pd.DataFrame:
    """Try several encodings when reading a CSV file.

    Some CSVs (especially those with special characters in author names or titles)
    may not be valid UTF-8. This helper rotates through common encodings and raises
    the last error only if all fail.
    """
    encodings = ["utf-8", "latin1", "ISO-8859-1", "cp1252"]
    last_error = None
    for enc in encodings:
        try:
            return pd.read_csv(path, encoding=enc)
        except Exception as e:
            last_error = e
    if last_error is not None:
        raise last_error
    raise RuntimeError(f"Unable to read CSV at {path}")


train_df = smart_read_csv(train_qa_path)
metadata_df = smart_read_csv(metadata_path)

print("train_QA.csv columns:", train_df.columns.tolist())
print("metadata.csv columns:", metadata_df.columns.tolist())
print("\nNumber of training QAs:", len(train_df))
print("Number of metadata rows:", len(metadata_df))

train_df.head(15)

Step 2 – Load the cached chunks from S3

In the Bedrock episode we turned every PDF into page-level records, split those pages into overlapping chunks, and uploaded the result to S3 as chunks.jsonl. Chunking is CPU work, so there is no reason to repeat it while a GPU is billing. We download the cached file and load it.

PYTHON

chunks_s3_key = "wattbot/chunks.jsonl"          # in YOUR bucket, written by the Bedrock episode
chunks_jsonl_path = os.path.join(local_data_dir, "chunks.jsonl")

def load_chunked_docs_jsonl(path):
    with open(path, "r", encoding="utf-8") as f:
        return [json.loads(line) for line in f]

if not os.path.exists(chunks_jsonl_path):
    download_s3_object(bucket_name, chunks_s3_key, chunks_jsonl_path)

chunked_docs = load_chunked_docs_jsonl(chunks_jsonl_path)
print("Chunked docs:", len(chunked_docs))
chunked_docs[0] if chunked_docs else None

Each chunk keeps a pointer back to its document and page metadata (doc_id, title, url, page_num, page_label, chunk_idx_in_page), which we use later to cite sources.

Step 3 – Build an embedding matrix

Now we embed each chunk into a vector using an open embedding model. For WattBot, a strong and relatively efficient choice is:

`thenlper/gte-large` (Recommended baseline embedder)

Size / parameters: ~335M parameters, roughly 1.3–1.4 GB in BF16/FP16 when loaded on GPU. Fits cleanly on T4 (16 GB), L4, A10G, A10, A100, and all g5.* instances. Offers noticeably better retrieval quality than smaller 100M–150M models without requiring high-end GPU memory. Runs comfortably on g4dn.xlarge, g5.xlarge, or g5.2xlarge during workshops. Lets participants see meaningful improvements from chunking and retrieval methods without excessive compute cost.
Intended use: General-purpose retrieval and semantic search across academic PDFs, sustainability reports, and mixed-domain long-form documents. Stronger semantic coherence than gte-base or MiniLM, but still lightweight enough for workshop hardware.
Throughput expectations:
- CPU only: workable for small corpora (<2k chunks) but slow for anything larger.
- GPU (T4, L4, A10G, A100) with batch sizes around 64–128:
  - 20k–40k chunks/min on L4 or A10G
  - 10k–15k chunks/min on T4
  - 50k+ chunks/min on A100

We will:

Load the embedding model on GPU if available.
Encode all chunks in batches.
Store the resulting matrix as a torch.Tensor or numpy.ndarray along with the original chunked_docs list.

Later, we’ll implement a small retrieval helper that does cosine-similarity search over this matrix—no additional indexing library required.

PYTHON

# We'll use a stronger embedding model now that we have a GPU.
# This model has ~335M parameters and benefits from GPU acceleration,
# but is still reasonable to run on a single 24 GB GPU.
embedding_model_id = "thenlper/gte-large"

print("CUDA available for embeddings:", torch.cuda.is_available())

# Single shared embedder object that we can pass around; it uses the GPU when one is available
embedder = HFEmbedder(embedding_model_id)

Embed every chunk in batches on the GPU and stack the results into one matrix, one row per chunk.

PYTHON

def embed_texts(embedder, docs, batch_size: int = 32) -> np.ndarray:
    """Embed all chunk texts into a dense matrix of shape (N, D)."""
    texts = [d["text"] for d in docs]
    all_embeddings = []
    start = time.time()
    for i in range(0, len(texts), batch_size):
        batch = texts[i : i + batch_size]
        emb = embedder.encode(
            batch,
            convert_to_numpy=True,
            show_progress_bar=False,
            normalize_embeddings=True,   # unit-length vectors, so a dot product is the cosine similarity
        )
        all_embeddings.append(emb)
    embeddings = np.vstack(all_embeddings) if all_embeddings else np.zeros((0, 768))
    print(f"Computed embeddings for {len(texts)} chunks in {time.time() - start:.1f}s")
    return embeddings

PYTHON


chunk_embeddings = embed_texts(embedder, chunked_docs)
chunk_embeddings.shape

4. Build a simple retrieval step (cosine similarity)

We are not using a heavy vector database in this episode.

Instead, we:

Embed each chunk with thenlper/gte-large (done above).
Embed each question.
Compute cosine similarity between the question embedding and all chunk embeddings.
Take the top–k most similar chunks as our retrieved context.

This keeps the retrieval logic completely transparent for teaching, while still matching the spirit of production systems that use FAISS, Chroma, Weaviate, etc.

When might FAISS or a vector database be worth exploring?

For small–to–medium experiments (a few thousand to maybe tens of thousands of chunks), this “plain NumPy + cosine similarity” approach is usually enough. You might consider FAISS or a full vector DB when:

Your corpus gets big
Once you’re in the hundreds of thousands to millions of chunks, brute-force similarity search can become slow and memory-hungry. FAISS and friends provide approximate nearest neighbor search that scales much better.
You need low-latency, repeated queries
If many users (or a web app) will hit your RAG system concurrently, you’ll want:
- fast indexing,
- efficient caching, and
- sub-second query latency.
  Vector DBs are designed for this use case.
You need rich filtering or metadata search
Vector DBs often support:
- filtering by metadata (e.g., paper = "chung2025", year > 2021),
- combining keyword + vector search (“hybrid search”),
- role-based access control and multi-tenant setups.
You want to share an index across services
If multiple notebooks, microservices, or teams need to reuse the same embedding index, a shared FAISS index or hosted vector DB is much easier to manage than passing around .npy files.
You need GPU-accelerated or distributed search
FAISS can use GPUs and sharding to speed up search on very large embedding collections. This is overkill for our teaching demo (and the Wattbot project in general), but very relevant for production-scale systems.

In this episode we deliberately stick with a simple in-memory index so the retrieval step is easy to inspect and debug. In your own projects you can swap out the retrieval layer for FAISS or a vector DB without changing the overall RAG architecture: the model still sees “top–k retrieved chunks” as context.

retrieve_top_k from wattbot_rag.py is that search: normalize the query and the matrix, one matrix multiply for the cosine similarities, argsort for the top k. Open the file to see it; it is a dozen lines. embed_query_with wraps the sentence-transformer as the embed_query(question) callable the rest of the pipeline expects.

PYTHON

embed_query = embed_query_with(embedder)   # embed_query(question) -> vector, with the same model that embedded the chunks

# Quick sanity check for `retrieve_top_k` on the first training question
first_row = train_df.iloc[0]
test_question = first_row["question"]
print("Sample question:", test_question)

test_retrieved = retrieve_top_k(
    query_embedding=embed_query(test_question),
    chunk_embeddings=chunk_embeddings,
    chunked_docs=chunked_docs,
    k=3,
)

print(f"Top {len(test_retrieved)} retrieved chunks:")
for r in test_retrieved:
    snippet = r["text"].replace("\n", " ")
    if len(snippet) > 160:
        snippet = snippet[:160] + "..."
    print(f"- score={r['score']:.3f} | doc_id={r['doc_id']} | page={r['page_label']} | snippet={snippet}")

5. Serve the generator: Qwen3.8-27B in FP8, via vLLM

For this episode we use Qwen3.8-27B, an open-weight model that Bedrock does not carry, in its FP8 release (Qwen/Qwen3.8-27B-FP8), served with vLLM. It is a dense 27B model with a hybrid attention design (most layers use a linear-attention variant with a fixed-size state, so the KV cache stays small even for long prompts), a native 262K context, and a “thinking” mode that is on by default and that we turn off, because WattBot wants a JSON answer, not a reasoning trace.

Why FP8 weights and not bf16? Inference is not training. No gradients, no optimizer state, nothing is updated, so the only memory that matters is the weights themselves plus the KV cache for the prompt being processed. A 27B model in bf16 is about 54 GB of weights and does not fit a 48 GB GPU; the same weights in 8-bit floating point are about 27 GB, which fits with room for the embedding model and generation. FP8 is near-lossless for a model this size, and it is faster, because decoding is limited by how fast the GPU can read weights from memory. For inference on rented GPUs, reduced-precision weights are the norm (FP8 where the GPU supports it, 4-bit integer where it does not or where memory is tighter); bf16 is what you use when you have memory to spare and a model small enough not to need it.

Why vLLM and not transformers.pipeline? vLLM is an inference server: it has the kernels for FP8 and for this model’s hybrid attention, batches requests, and manages the KV cache so long prompts do not run out of memory. It is what you would deploy behind an API, and its offline LLM class is a two-line replacement for a transformers.pipeline.

We will:

Call the model once to propose an answer and supporting evidence.
Call it a second time with a smaller prompt to generate a short explanation (<= 100 characters).

PYTHON

import os
# Keep model weights on the notebook's 100 GB volume rather than the small root disk.
os.environ["HF_HOME"] = "/home/ec2-user/SageMaker/hf_cache"

from vllm import LLM, SamplingParams

qwen_model_id = "Qwen/Qwen3.8-27B-FP8"   # pre-quantized FP8 checkpoint; vLLM reads that from its config, so no quantization= argument is needed

load_start = time.time()
llm = LLM(
    model=qwen_model_id,                         # Hugging Face id; the weights download to HF_HOME on first use
    max_model_len=8192,                          # longest prompt + answer we need; the model allows 262K, we do not
    gpu_memory_utilization=0.85,                 # leave ~7 GB for the embedding model and PyTorch overhead
    limit_mm_per_prompt={"image": 0, "video": 0},  # text only: skip reserving memory for the vision encoder
)
print(f"Model loaded in {(time.time() - load_start) / 60:.1f} min (download on first use, then weights into GPU memory)")
print("vLLM engine ready:", qwen_model_id)

PYTHON

!nvidia-smi --query-gpu=memory.used,memory.total --format=csv

The number above is what you are paying for: a 48 GB GPU, most of it now holding model weights and a reserved KV cache. It stays allocated, and billed, until you stop the instance.

PYTHON

def call_qwen_chat(system_prompt: str, user_prompt: str, max_new_tokens: int = 384) -> str:
    """Send a chat prompt to the vLLM engine and return only the generated text."""
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]
    params = SamplingParams(temperature=0.0, max_tokens=max_new_tokens)
    outputs = llm.chat(
        messages,
        sampling_params=params,
        chat_template_kwargs={"enable_thinking": False},   # answer directly; no reasoning trace
        use_tqdm=False,
    )
    return outputs[0].outputs[0].text.strip()

print("Generator model and helper loaded.")

One throwaway question, to confirm the generator answers before spending time on the real ones.

PYTHON


# Quick sanity check for `call_qwen_chat`
test_system_prompt = "You are a concise assistant who answers simple questions clearly."
test_user_prompt = "What is 2 + 2? Answer in one short sentence."

test_response = call_qwen_chat(
    system_prompt=test_system_prompt,
    user_prompt=test_user_prompt,
    max_new_tokens=32,
)
print(f"Generator ({qwen_model_id}) test response: {test_response}")

6. The prompts

call_qwen_chat has the signature the shared pipeline expects, generate(system_prompt, user_prompt, max_new_tokens) -> str, and embed_query has the other one. run_single_qa in wattbot_rag.py makes two calls per question with them:

The first asks for a JSON object with answer, answer_value (normalized scalar, range or category), answer_unit, ref_id (list of document ids), is_blank and supporting_materials. The retrieved chunks go in the user turn, labelled [doc_id, page N] so the model can cite them, along with the expected answer_unit from train_QA.csv as a hint, so the number comes back in the unit the ground truth uses.
The second, made only when the first produced an answer, asks for a short explanation. With with_evidence_type=True it also asks the model to name the main type of evidence (text, table, figure, or a combination) in the strict format evidence_type: explanation, and the code turns that into [evidence_type] explanation for the CSV, so the tag is parsed rather than trusted.

PYTHON

print(ANSWER_SYSTEM_PROMPT)
print("-" * 80)
print(explanation_system_prompt(with_evidence_type=True))

7. Run over the full WattBot training set

For every row of train_QA.csv, run_single_qa:

Embeds the question and retrieves the top-k chunks.
Asks Qwen for the answer JSON and parses it (a malformed reply counts as unanswerable).
Decides answerability with a hybrid rule: the question is is_blank if the model said so or the top retrieval cosine score is below retrieval_threshold. Weak retrieval means the model is about to answer from memory instead of the papers, which is what RAG exists to prevent. A lower threshold marks fewer questions unanswerable, a higher one more; change it, rerun, and rescore to see the trade-off.
Copies answer_unit from the row (never guessed), maps each ref_id to its URL via metadata.csv, and makes the explanation call.
Returns one submission-format row; we collect them into train_solutions_qwen.csv.

PYTHON

docid_to_url = build_docid_to_url(metadata_df)

all_results = []
max_N = np.inf   # set to e.g. 5 for a smoke test
run_start = time.time()

for idx, row in train_df.iterrows():
    if idx >= max_N:
        break
    print("#" * 96)
    print(f"QUESTION: {row['question']}")

    res = run_single_qa(
        row=row,
        embed_query=embed_query,             # the embedding model on this GPU
        generate=call_qwen_chat,             # the vLLM engine on this GPU; the Bedrock episode passed its converse wrapper here
        chunk_embeddings=chunk_embeddings,
        chunked_docs=chunked_docs,
        docid_to_url=docid_to_url,
        top_k=8,
        retrieval_threshold=0.25,
        with_evidence_type=True,             # ask the explanation call to tag [text], [table], [figure]
        verbose=True,
    )
    print(f"ANSWER: {res['answer']}")
    print(f"ref_ids: {res['ref_id']}")
    print(f"EXPLANATION: {res['explanation']}")
    all_results.append(res)

solutions_df = pd.DataFrame(all_results)
print(f"Answered {len(solutions_df)} questions in {(time.time() - run_start) / 60:.1f} min on the notebook GPU")
solutions_path = os.path.join(local_data_dir, "train_solutions_qwen.csv")
solutions_df.to_csv(solutions_path, index=False)

print(f"Saved solutions for {len(solutions_df)} questions to: {solutions_path}")
solutions_df.head(20)

Compare answers to ground truth

WattBot scores each prediction with the competition’s own Score.py: 0.75 for answer_value (numeric within ±0.1% relative tolerance, categorical match after normalisation, zero if you answer a question the corpus cannot answer), 0.20 for citation F1 against the ground-truth ref_id set, and 0.05 for a clean refusal when the answer is is_blank. compute_wattbot_score from wattbot_rag.py calls it for the headline number, adds a per-row score, and prints the worst rows with the model’s explanation.

PYTHON

results_df = compute_wattbot_score(
    train_qa_path=train_qa_path,
    preds_path=solutions_path,
    n_examples=10,
)

Change retrieval_threshold or top_k in the loop above, rerun it, and rescore: with the model loaded, a full pass over the training set takes a few minutes, which is what a GPU notebook is for. Then stop the instance.

Recap and next steps

In this episode, we:

Pulled the WattBot corpus chunks that the Bedrock episode cached on S3 into our notebook environment, and reused the pipeline code (wattbot_rag.py) from the Bedrock and Processing Jobs episodes unchanged.
Created vector embeddings for each chunk and used similarity search to retrieve relevant context for a given question.
Used an LLM to generate answers from retrieved context and wrote results out to a CSV for later scoring and analysis.
Handled unanswerable questions with an is_blank flag so the system can explicitly say “I don’t know” when the evidence isn’t there.

This is just a first pass at a RAG pipeline: it works, but there’s a lot of headroom to improve both accuracy and robustness. Some natural next steps:

Increase the size/quality of models used for embedding and generation: Try stronger embedding models (e.g., larger sentence-transformers or domain-tuned embeddings) and more capable LLMs for answer generation, especially if you have GPU budget.
Add a reranking step: Instead of sending the top-k raw nearest neighbors directly to the LLM, use a cross-encoder or reranker model to re-score those candidates and send only the best ones.
Handle figures and tables more carefully: Many key numbers live in tables, figure captions, or plots. Consider:
- OCR / table-parsing tools (e.g., pytesseract, table extractors, PDF parsers).
- Multimodal models that can embed or interpret figures and diagrams, not just text.
- Separate chunking strategies for captions, tables, and main text.
Enrich chunks with metadata: Attach metadata like section headings (e.g., Methods, Results), paper ID, year, or paragraph type. You can:
- Filter or boost chunks by metadata at retrieval time.
- Use metadata in the prompt so the LLM knows where evidence is coming from.
Look for LLMs tuned for scientific literature: Experiment with models that are explicitly trained or finetuned on scientific text (e.g., arXiv / PubMed) so they:
- Parse equations and technical language more reliably.
- Are less likely to hallucinate when reading dense scientific prose.

As you iterate, the goal is to treat this notebook as a baseline RAG “workbench”: you can swap in better models, smarter retrieval strategies, and richer document preprocessing without changing the overall pipeline structure.

What did this episode cost?

There is no job record to ask here; the bill is simply the instance’s hours. This counts from the setup cell, so the boot, the install and every pause since are in it, and it keeps counting until you stop the instance.

PYTHON

g6e_price = 2.61   # $/hour, ml.g6e.xlarge on-demand
hours = (time.time() - session_start) / 3600
print(f"this notebook has been running {hours * 60:.0f} min since the setup cell: ${hours * g6e_price:.2f} so far, and counting")
print(f"the same pipeline as Processing jobs billed about $0.60 (that episode); on Bedrock, about fifteen cents")

The Processing Jobs episode does the same work, with the same wattbot_rag.py, as short-lived SageMaker jobs, so the GPU only bills while it is embedding or generating. Stop this instance now, and delete it if you are done with the episode.

Key Points

Why self-host at all: Bedrock’s catalog is curated and lags releases; most open-weight models (including Qwen3.8-27B) are not on it. When the model you need is not offered, you rent a GPU and serve it yourself.
Inference uses reduced-precision weights: no gradients means the weights are the memory cost. FP8 puts a 27B model on one 48 GB GPU (ml.g6e.xlarge); bf16 would need 54 GB and a multi-GPU instance.
The meter never stops on its own: ml.g6e.xlarge is $2.61/h, $167 per forgotten weekend, $1,906 per forgotten month, plus $14/month for the storage volume even when stopped. Stop it when you stand up; delete it when the episode is over.
Local-first RAG: For teaching (and small corpora), we avoid an external vector database and instead perform cosine similarity search over in-memory embeddings.
Ground-truth units: The answer_unit column is always copied directly from train_QA.csv, never guessed by the LLM.
Two-stage LLM use: One call focuses on answering and citing; a second, lighter call produces a short explanation tagged with an evidence type.
WattBot conventions: We respect the Kaggle competition format, using is_blank for unanswerable questions and for missing fields.
Scalability path: The same logic can later be swapped to FAISS/Chroma and larger models, while preserving the interface used here.

Content from (Optional) The Same Jobs from the Terminal

Last updated on 2026-10-20 | Edit this page

Estimated time: 30 minutes

Overview

Questions

What does the AWS command-line interface (CLI) let me do that the lesson’s notebooks did, and what should I not do with it?
How do I see, and stop, everything that is costing money from a terminal?
What request does the SageMaker SDK send when I call estimator.fit()?

Objectives

Authenticate the CLI with short-lived Identity Center credentials, never with access keys.
List, inspect, and stop training jobs, processing jobs, tuning jobs and notebook instances from the terminal.
Copy data to and from S3 in bulk with aws s3 sync.
Launch one training job from its JSON description and recognise it as the request the SDK built in the training episodes.
State which resources the lesson never creates from the CLI, and why the account refuses them.

Optional, self-study episode

Everything in this lesson went through the SageMaker Python SDK in a notebook. The same account can be used from a terminal with the AWS CLI: to check on a job from another machine, to stop something when no notebook is open, to copy results out of S3 in bulk, or to script a launch. This episode shows the CLI equivalents of what you already did, in four groups: look, stop, copy, launch.

Callout

Read this before the first command

The CLI, the SDK and the console are the same API. Nothing here is more powerful than a notebook cell. The difference is that a terminal shows nothing unless you ask: a notebook instance left running is visible in the console the next day, while an endpoint created from a terminal on a Friday is not noticed until the bill arrives. A past participant of this workshop spent a large share of a year’s credits that way. Three rules apply on the shared account and are worth keeping on your own:

Authenticate with aws sso login only. Never aws configure with an access key ID and secret. Those are long-lived, and they end up in shell histories and committed files.
Launch only jobs (training, processing, tuning), with the same three tags and the same MaxRuntimeInSeconds the notebooks always set. Jobs end on their own. Endpoints, EC2 instances and GPU notebook instances do not, and the lesson never creates them from the CLI, or at all.
An AccessDeniedException on the commands below is intended. The workshop permission set denies endpoint and EC2 creation, instance types outside the lesson’s list, and any job launched without the tags. Do not look for a workaround.

Install AWS CLI v2 if you do not have it, then register the account’s Identity Center once. It asks for the start URL and region of the sign-in page you used in the browser (for the workshop account they are in your sign-in email), opens a browser to approve the device, lets you pick the account and permission set, and asks for a default region (us-east-2 for the workshop account) and a profile name.

BASH

aws configure sso

From then on, every session starts with a login, and a check that you are who you think you are, in the account you think you are in:

BASH

aws sso login --profile workshop
export AWS_PROFILE=workshop          # PowerShell: $env:AWS_PROFILE="workshop"
aws sts get-caller-identity

The output names the account ID and an AWSReservedSSO_... role. The credentials expire after a few hours; aws sso login again when a command returns ExpiredToken. Nothing that would still work next month is written anywhere.

If you are on the notebook instance’s own terminal (JupyterLab, File → New → Terminal), skip this step: the instance’s role is already the identity, which is the reason the lesson works from one.

2. Look: what is running and what did it cost

The SageMaker SDK calls you used map one-to-one onto aws sagemaker subcommands. --query is a JMESPath filter that keeps the output to what you asked about.

BASH

# Everything that is running right now, by type. Each list should be short, or empty.
aws sagemaker list-training-jobs   --status-equals InProgress --query 'TrainingJobSummaries[].TrainingJobName'
aws sagemaker list-processing-jobs --status-equals InProgress --query 'ProcessingJobSummaries[].ProcessingJobName'
aws sagemaker list-hyper-parameter-tuning-jobs --status-equals InProgress --query 'HyperParameterTuningJobSummaries[].HyperParameterTuningJobName'
aws sagemaker list-notebook-instances --query 'NotebookInstances[].[NotebookInstanceName,NotebookInstanceStatus,InstanceType]' --output table
aws sagemaker list-endpoints       # must print an empty list on the workshop account

describe-training-job is the call behind job_cost_report() in the training episodes. The same number is there: billable time, which includes the container pull and data download you paid for.

BASH

JOB=<a training job name from the list above>
aws sagemaker describe-training-job --training-job-name $JOB \
  --query '{status:TrainingJobStatus, instance:ResourceConfig.InstanceType, count:ResourceConfig.InstanceCount, billable_s:BillableTimeInSeconds, tags:Tags}'

Cost is billable seconds × instance count × the hourly rate, exactly as in the notebooks. Tags are returned by a separate call, which is also how you find the untagged job that will show up as “unallocated” on the bill:

BASH

ARN=$(aws sagemaker describe-training-job --training-job-name $JOB --query TrainingJobArn --output text)
aws sagemaker list-tags --resource-arn $ARN --query 'Tags[].[Key,Value]' --output table

3. Stop: the commands worth memorising

These are the most useful CLI commands to know. Each one stops a charge.

BASH

aws sagemaker stop-training-job   --training-job-name   <name>
aws sagemaker stop-processing-job --processing-job-name <name>
aws sagemaker stop-hyper-parameter-tuning-job --hyper-parameter-tuning-job-name <name>   # stops every child job
aws sagemaker stop-notebook-instance --notebook-instance-name <name>

A stopped job keeps its logs and whatever it had already written to S3; a stopped notebook instance keeps /home/ec2-user/SageMaker/. Stopping is never the expensive mistake. Deleting an instance when you are done with it is the step that stops the storage charge too:

BASH

aws sagemaker delete-notebook-instance --notebook-instance-name <name>   # only after it reports Stopped

4. Copy: S3 in bulk

The S3 episode uploaded through the console and the SDK. For many files, or for pulling a job’s outputs down to a laptop, the CLI is the better tool:

BASH

aws s3 ls s3://<your-bucket>/ --recursive --human-readable --summarize     # every object, and the bucket's total size
aws s3 sync s3://<your-bucket>/xgboost/ ./xgboost-outputs/                 # download a prefix; only new or changed files move
aws s3 sync ./my-data/ s3://<your-bucket>/my-data/                         # upload a folder the same way
aws s3 rm s3://<your-bucket>/scratch/ --recursive                          # delete a prefix; this one asks no questions

Downloads to a laptop are a transfer out of AWS ($0.09 per GB). For anything larger than a few hundred megabytes, keep the data in the region and let a job read it, which is the lesson’s pattern.

5. Launch: one training job, as JSON

Every estimator.fit() in the training episodes assembled a JSON request and sent it to CreateTrainingJob. Here is that request. Fill in the five placeholders from your notebook: the image URI (run print(xgboost_estimator.training_image_uri()) in the training episode’s notebook), your role ARN (print(role)), your bucket, your tag values, and a job name, which must be unique in the account.

BASH

cat > train-xgboost.json <<'JSON'
{
  "TrainingJobName": "<notebook-instance-name>-cli-xgboost-1",
  "AlgorithmSpecification": {
    "TrainingImage": "<xgboost image URI for your region>",
    "TrainingInputMode": "File"
  },
  "RoleArn": "<execution role ARN>",
  "HyperParameters": { "objective": "binary:logistic", "num_round": "100", "max_depth": "5", "eta": "0.1" },
  "InputDataConfig": [
    {
      "ChannelName": "train",
      "ContentType": "text/csv",
      "DataSource": { "S3DataSource": {
        "S3DataType": "S3Prefix",
        "S3Uri": "s3://<your-bucket>/xgboost/builtin/train_builtin.csv",
        "S3DataDistributionType": "FullyReplicated" } }
    }
  ],
  "OutputDataConfig": { "S3OutputPath": "s3://<your-bucket>/xgboost/cli-output/" },
  "ResourceConfig": { "InstanceType": "ml.m5.large", "InstanceCount": 1, "VolumeSizeInGB": 10 },
  "StoppingCondition": { "MaxRuntimeInSeconds": 1200 },
  "Tags": [
    { "Key": "Project", "Value": "<ChallengeName>" },
    { "Key": "Name",    "Value": "<LastnameFirstname>" },
    { "Key": "Purpose", "Value": "Training" }
  ]
}
JSON
aws sagemaker create-training-job --cli-input-json file://train-xgboost.json

Three fields to notice. StoppingCondition.MaxRuntimeInSeconds is the max_run you set on every estimator; the SDK would have sent 48 hours if you had not. Tags is the tags= argument, and without it the request is refused on the workshop account. ResourceConfig sets the instance type and count, which determine the cost; nothing in a JSON file prints a price, so check it as carefully as you would in a notebook.

The CLI returns the moment the request is accepted; nothing waits. To follow the job the way wait=True did:

BASH

aws sagemaker wait training-job-completed-or-stopped --training-job-name <the job name>   # blocks until it ends
aws sagemaker describe-training-job --training-job-name <the job name> --query '{status:TrainingJobStatus, billable_s:BillableTimeInSeconds, failure:FailureReason}'

Estimate first, compare after, as always: this is the same ~4 minutes billable on ml.m5.large, about a cent.

Callout

What the lesson never does from the CLI

aws sagemaker create-endpoint / create-endpoint-config: an endpoint is a server that runs until you delete it, with no runtime cap. The lesson serves models with processing jobs and Bedrock so that nothing of this kind exists in the account.
aws sagemaker create-notebook-instance with a GPU type: a GPU notebook is $60 a day while you are asleep. The optional GPU-notebook episode creates one through the console, where it stays visible, and deletes it the same hour.
aws ec2 run-instances: bare virtual machines have no tags required, no runtime cap and no SageMaker job list to appear in.

On the workshop account all three are denied by the permission set, as is any instance type outside the lesson’s list and any job without the three tags; the instructor page Account preparation checklist has the policy. On your own account, the same denies take about fifteen minutes to set up.

Challenge

Challenge: the Friday-evening check

You are leaving for the weekend and have a terminal, not a browser. Write the commands you would run, in order, to make sure the account is costing nothing until Monday, and the one command to run if each check finds something. Then run them against your own resources.

Show me the solution

BASH

aws sso login --profile workshop && export AWS_PROFILE=workshop
aws sagemaker list-training-jobs   --status-equals InProgress --query 'TrainingJobSummaries[].TrainingJobName'
aws sagemaker list-processing-jobs --status-equals InProgress --query 'ProcessingJobSummaries[].ProcessingJobName'
aws sagemaker list-hyper-parameter-tuning-jobs --status-equals InProgress --query 'HyperParameterTuningJobSummaries[].HyperParameterTuningJobName'
aws sagemaker list-notebook-instances --query 'NotebookInstances[?NotebookInstanceStatus==`InService`].NotebookInstanceName'
aws sagemaker list-endpoints --query 'Endpoints[].EndpointName'

For anything the first four print: the matching stop-... command from section 3. A job you recognise as yours and expect to finish within its MaxRuntimeInSeconds can be left; a job you do not recognise is a question for its owner, and on a shared account the tags (list-tags) tell you who that is. For an endpoint: aws sagemaker delete-endpoint --endpoint-name <name>, then find out how it was created, because the permission set should have refused it. Buckets cost almost nothing over a weekend and need no action.

Key Points

The CLI, the SDK and the console are one API; the CLI shows you nothing you do not ask for, so its first use is to look and to stop, not to launch.
Authenticate with aws sso login and short-lived credentials; never store an access key.
list-*, describe-* and stop-* for training, processing and tuning jobs and notebook instances cover almost every reason to open a terminal.
A create-training-job JSON is what estimator.fit() sent: the same tags, the same MaxRuntimeInSeconds, the same instance choice, with no price printed next to it.
Endpoints, EC2 instances and GPU notebooks are not jobs and do not end; the workshop account denies creating them, and your own account should too.

Content from Resource Management and Monitoring

Last updated on 2026-10-21 | Edit this page

Estimated time: 35 minutes

Overview

Questions

How can I monitor and manage AWS resources to avoid unnecessary costs?
What steps are necessary to clean up SageMaker and S3 resources after the workshop?
What best practices can help with efficient resource utilization?

Objectives

Understand how to shut down SageMaker notebook instances to minimize costs.
Clean up S3 storage and stop running jobs.
Explore basic resource management strategies and tools for AWS.

Initial setup

Open the notebook from: /ML_with_AWS_SageMaker/notebooks/15_Resource-management-cleanup.ipynb. Any kernel works; conda_python3 is fine. This episode is mostly console steps, but the cells below ask AWS what is still running or stored under your name, so you can check the list against reality before and after you clean up.

PYTHON

import boto3
import sagemaker
import sys

sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")
import helpers

session = sagemaker.Session()
sm = session.sagemaker_client
region = session.boto_region_name

name = "DoeJohn"                              # replace with your name, as used in job names and tags
project = "SinkOrSwim"                        # replace with your challenge name, as used in the Bedrock profile names
bucket_name = "sinkorswim-doejohn"            # replace with your bucket

What is still running under your name?

Everything this lesson launched carries your name in its job name or its tags, so one pass over the job lists shows what is still billing. A completed job holds nothing; an InProgress one bills until it ends or hits its cap.

PYTHON

print("Notebook instances (bill while InService; volume bills while Stopped):")
for nb in sm.list_notebook_instances()["NotebookInstances"]:
    print(f"  {nb['NotebookInstanceName']:<40} {nb['NotebookInstanceStatus']:<12} {nb['InstanceType']}")

print("\nJobs still InProgress with your name in them:")
running = 0
for kind, lister, key, name_key in [
    ("training",   sm.list_training_jobs,              "TrainingJobSummaries",            "TrainingJobName"),
    ("tuning",     sm.list_hyper_parameter_tuning_jobs, "HyperParameterTuningJobSummaries", "HyperParameterTuningJobName"),
    ("processing", sm.list_processing_jobs,            "ProcessingJobSummaries",          "ProcessingJobName"),
]:
    for job in lister(StatusEquals="InProgress", NameContains=name, MaxResults=50)[key]:
        print(f"  {kind:<11} {job[name_key]}")
        running += 1
print("  none" if running == 0 else f"  {running} job(s) still billing: stop them from the console if you no longer need them")

Then storage, which bills whether or not anything runs:

PYTHON

helpers.get_s3_bucket_size(bucket_name)

And the Bedrock inference profiles, which cost nothing to keep but are yours to remove (the instructors’ cleanup does not touch them):

PYTHON

bedrock = boto3.client("bedrock", region_name=region)
mine = f"{project}-{name}-"   # the Bedrock episode named every profile <project>-<name>-<label>, so this prefix is yours and nobody else's
profiles = [p for p in bedrock.list_inference_profiles(typeEquals="APPLICATION")["inferenceProfileSummaries"]
            if p["inferenceProfileName"].startswith(mine)]
for p in profiles:
    print(p["inferenceProfileName"], p["inferenceProfileArn"])
print("none" if not profiles else f"{len(profiles)} profile(s) with your prefix; the next cell deletes them")

Deleting them is one call per profile. The cell deletes only the profiles the cell above listed, the ones whose name starts with your project and your name, so nobody else’s are touched even though the account’s policy would allow it; read the list before you set the flag. The notebook instance needs the Bedrock role from the RAG episodes for this call; on the old role it fails with an access-denied error, and the Bedrock console does the same job by hand (Inference and Assessment → Cross-region inference → Application inference profiles → select → Delete).

PYTHON

confirm_delete = False   # set to True after reading the list above

if confirm_delete:
    for p in profiles:
        bedrock.delete_inference_profile(inferenceProfileIdentifier=p["inferenceProfileArn"])
        print("deleted", p["inferenceProfileName"])
    profiles = []
else:
    print("nothing deleted; set confirm_delete = True and rerun to delete the profiles listed above")

Now the console steps. Rerun the cells above when you are done; every list should be empty or Stopped.

Shutting down notebook instances

Notebook instances in SageMaker are billed per second, so it’s essential to stop or delete them when they are no longer needed. Earlier in the Notebooks as controllers episode, we discussed using lower-cost instance types like ml.t3.medium (approximately $0.05/hour) for controlling workflows. While this makes open notebooks less costly than larger instances, it’s still a good habit to stop or delete notebooks to avoid unnecessary spending, especially if left idle for long periods.

Navigate to SageMaker in the AWS Console.
In the left-hand menu, click Notebook instances.
Locate your notebook instance and select it.
Choose Stop to shut it down temporarily or Delete to permanently remove it. > Tip: If you plan to reuse the notebook later, stopping it is sufficient. Deleting is recommended if you are finished with the workshop.

A stopped instance stops the compute charge but keeps its storage volume, which bills at about $0.14 per GB-month. The default 5 GB is cents; the 100 GB volume the optional GPU-notebook episode asks for is $14 a month. If you enlarged the volume to hold model weights, delete the instance once the episode is done.

Cleaning up S3 storage

While S3 storage is relatively inexpensive, cleaning up unused buckets and files helps keep costs minimal and your workspace organized.

Navigate to the S3 Console.
Locate the bucket(s) you created for this workshop.
Open the bucket and select any objects (files) you no longer need.
Click Delete to remove the selected objects.
To delete an entire bucket:
- Empty the bucket by selecting Empty bucket under Bucket actions.
- Delete the bucket by clicking Delete bucket.

Reminder: Earlier in the workshop, we set up tags for S3 buckets. Use these tags to filter and identify workshop-related buckets, ensuring that only unnecessary resources are deleted.

Monitoring and stopping active jobs

SageMaker charges for training, tuning and Processing jobs while they run, so make sure to terminate unused jobs. Every job in this lesson was launched with a runtime cap (max_run on training and tuning jobs, max_runtime_in_seconds on Processing jobs), so a forgotten job ends on its own.

In the SageMaker Console, go to Training jobs, Hyperparameter tuning jobs, or Processing jobs.
Identify any active jobs that you no longer need.
Select the jobs and click Stop. > Tip: Review the job logs to ensure you’ve saved the results before stopping a job.

Billing and cost monitoring

Managing your AWS expenses is vital to staying within budget. Follow these steps to monitor and control costs:

Set up billing alerts:
- Go to the AWS Billing Dashboard.
- Navigate to Budgets and create a budget alert to track your spending.
Review usage and costs:
- Use the AWS Cost Explorer in the Billing Dashboard to view detailed expenses by service, such as SageMaker and S3.
Use tags for cost tracking:
- Refer to the tags you set up earlier in the workshop. Remember that tags never propagate between resources: your notebook instance, each S3 bucket, each training/tuning/Processing job, and each Bedrock application inference profile carries its own tags, and anything you launched without tags is untraceable after the fact.
- Only the three approved keys count on the bill, Name, Project and Purpose, spelled exactly (the S3 episode says why); a tag you just added can take up to 24 hours to appear in Cost Explorer.
- Bedrock application inference profiles cost nothing to keep, but delete the ones you created once a project is finished; the cells at the top of this episode do it for the ones carrying your name.

End-of-workshop checklist

Work through this list before you sign off for the last time.

Notebook instances: Stop, then Delete (SageMaker → Notebook instances). Includes the optional GPU notebook instance if you created one.
S3 bucket: empty and delete your challenge-lastnamefirstname bucket (S3 → select bucket → Empty, then Delete). Download anything you want to keep first, once: egress is $0.09/GB. Leave the shared data bucket alone; it is the instructors’.
Default SDK bucket sagemaker-<region>-<account-id>: created automatically the first time a job staged its code. In the shared account the instructors delete it; on your own account, empty and delete it with the rest.
Training, tuning and Processing jobs: confirm none are still running (SageMaker → each job list, filter by status InProgress). Completed jobs hold no billable resource; their artifacts are in your bucket and go with it.
Bedrock application inference profiles: delete the ones you created, with the two cells at the top of this episode (they list and delete only the profiles named with your project and your name) or from the Bedrock console.
GitHub personal access token: revoke it on GitHub (Settings → Developer settings → Personal access tokens) if you made one for the optional PAT episode.

Every stop and delete in this list has a one-line CLI equivalent; the optional episode The Same Jobs from the Terminal collects them, for the day you need to stop something from a machine with no browser.

In the shared workshop account, the instructors sweep the account with aws-nuke after the workshop. That sweep deletes notebook instances and the learner buckets only, so if you forget those, they will still go. It does not touch Bedrock application inference profiles, so those are yours to remove.

Best practices for resource management

Efficient resource management can save significant costs and improve your workflows. Below are some best practices:

Automate resource cleanup: Use AWS SDK or CLI scripts to automatically shut down instances and clean up S3 buckets when not in use. Learn more about automation with AWS CLI.
Schedule resource usage: Schedule instance start and stop times using AWS Lambda. Learn how to schedule tasks with AWS Lambda.
Test workflows locally first: Before scaling up experiments in SageMaker, test them locally to minimize cloud usage and costs. Learn about SageMaker local mode.
Use cost tracking tools: Explore AWS’s cost allocation tags and budget tracking features. Learn more about cost management in AWS.

By following these practices and leveraging the additional resources provided, you can optimize your use of AWS while keeping costs under control.

Challenge

Challenge: your own cleanup list

Without looking at the checklist above, list every billable resource you created in this workshop, and for each one write down where it is deleted (which console page, or which command). Then compare with the checklist. Which items would have kept billing if you had only stopped your notebook instance and walked away?

Show me the solution

A complete list for someone who did every episode, including the optional ones:

Resource	Created in	Bills while it exists?	Where it is deleted
S3 bucket `challenge-lastnamefirstname` and everything in it (the Titanic files and `.npz` arrays at the root, and the `docs/`, `xgboost/`, `output_nn/`, `wattbot/` and `badger-scribe/` prefixes)	S3 episode	Yes: storage per GB-month	S3 console → Empty bucket → Delete bucket
The SDK’s default bucket `sagemaker-<region>-<account-id>` (code tarballs staged by every `fit()` and `run()`)	Training episode, automatically	Yes, a little: storage per GB-month, shared by the whole account	Instructors delete it at the end of term; on your own account, empty and delete it like any other bucket
Controller notebook instance (`ml.t3.medium`)	Notebooks as controllers	Yes while InService; its volume bills while Stopped	SageMaker → Notebook instances → Stop → Delete
Training jobs (XGBoost, PyTorch), tuning jobs, Processing jobs	Training, tuning, RAG and fine-tuning episodes	Only while running; all capped by `max_run` / `max_runtime_in_seconds`	SageMaker → job lists → Stop if InProgress; nothing to delete afterwards
Bedrock application inference profiles	RAG on Bedrock	No (the calls billed per token when made)	The delete cell at the top of this episode, or the Bedrock console
GPU notebook instance (`ml.g6e.xlarge`, 100 GB volume)	Optional GPU-notebook episode	Yes: `$2.61`/h InService, `$14`/month for the volume when Stopped	SageMaker → Notebook instances → Delete
GitHub personal access token	Optional PAT episode	No, but it is a credential	GitHub → Settings → Developer settings → revoke

Stopping the controller and walking away leaves the bucket billing for storage, a stopped GPU instance’s volume billing at $0.14/GB-month, and any Processing job that had not hit its max_run yet. Model artifacts in your bucket are the only thing worth keeping; download them once, then delete the bucket.

Key Points

Always stop or delete notebook instances when not in use to avoid charges.
Stop any training, tuning or Processing job you no longer need.
Bedrock application inference profiles are not swept by the instructors’ aws-nuke run; remove them yourself.
Regularly clean up unused S3 buckets and objects to save on storage costs.
Monitor your expenses through the AWS Billing Dashboard and set up alerts.
Use tags (set up earlier in the workshop) to track and monitor costs by resource.
Following best practices for AWS resource management can significantly reduce costs and improve efficiency.

All in One View

Overview

Questions

Objectives

Cloud vs. the cluster down the hall

SageMaker and Bedrock

How the pieces fit together

Lesson outline

Overview

Questions

Objectives

Storing data on AWS

Consult your institution’s IT before handling sensitive data in AWS

Options for storage: EC2 Instance or S3

What is an EC2 instance?

When to store data directly on EC2

Limitations of EC2 storage

What is an S3 bucket?

Benefits of using S3 (recommended for SageMaker and ML workflows)

Recommended approach: S3 buckets

Detailed procedure

1. Sign in to the AWS Management Console

2. Navigate to S3

3. Create a new bucket

Wrong region costs money and breaks jobs

Why these three keys, and why they must be spelled exactly

4. Edit bucket policy

JSON

Doing this on your own account

5. Upload files to the bucket

6. Take note of S3 URI for your data

S3 bucket costs

Storage costs

Data transfer costs

Request costs

Challenge: Estimating Storage Costs

Show me the solution

Removing unused data (complete after the workshop)

Overview

Questions

Objectives

Setting up our notebook environment

What a SageMaker notebook instance is

Using the notebook as a controller

Why we work from a notebook instance in this lesson

1. Navigate to SageMaker

2. Create a new notebook instance

Load pre-filled Jupyter notebooks

SH

SH

At the end of class: stop the instance

Overview

Questions

Objectives

Set up AWS environment

PYTHON

PYTHON

Warnings this cell prints

PYTHON

Reading data from S3

A) Reading data directly from S3 into memory

Example: Reading data from S3 into memory

PYTHON

B) Download copy into notebook environment

Example

PYTHON

PYTHON

Check the current size and storage costs of bucket

PYTHON

Using the lesson’s helper functions

PYTHON

PYTHON

Check storage costs of bucket

PYTHON

PYTHON

Writing output files to S3

PYTHON

PYTHON

PYTHON

PYTHON

Enter it with `getpass`

Step 3: Convert `.ipynb` to `.py`

Benefits of converting to `.py` before committing

Step 4: Add and commit the `.py` file

Step 5: Ignore notebooks and clutter with `.gitignore`

Step 7: Pulling `.py` files and converting back to notebooks