All in One View

Last updated on 2026-11-02 | 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, control, job orchestration and efficiency.
Describe what S3, a notebook instance, SageMaker jobs and Bedrock each do in this lesson.
Say when a SageMaker job is the tool and when a Bedrock call is.

Most researchers already have somewhere to run ML, a campus HPC cluster, a departmental GPU server, or their own workstation. The cloud is another option, with its own trade-offs in cost, performance, flexibility and control. This episode compares the options, then introduces the four AWS pieces the lesson is built from, S3, a notebook instance, SageMaker jobs and Bedrock.

Cloud vs. the cluster down the hall

Cloud computing means renting compute as needed instead of buying and maintaining it. Whether that is the better choice depends on the workload and on what you already have.

Hardware and cost. A campus cluster or lab server is used for five or more years, so its GPUs are whatever was current when it was bought. Cloud providers offer each new generation soon after release, Blackwell B200 and GB200 today, with Vera Rubin announced for cloud deployment in 2026. The gains between generations are large. NVIDIA’s B200 has more than twice the memory and about seven times the bf16 arithmetic of the A100 on about two and a half times the power, so far more work per watt on paper, with the real speedup depending on the workload (NVIDIA data sheets). The financial trade-off comes down to utilization. Owned hardware costs the same busy or idle, so it is the cheaper option when it stays busy. Rented hardware is billed per second and only while a job runs, so it is the cheaper option for occasional or bursty work, and renting GPUs around the clock soon costs more than owning them.
Control and data. Local hardware gives you control over the software stack, the network and where the data sits, with no quotas and no transfer charges. In the cloud you work within account permissions and service quotas, data moved in or out crosses the network and may be charged, and sensitive data needs a governance check before it leaves campus. Cost needs attention too. An idle rented GPU costs the same as a busy one, so every job in this lesson has a runtime cap, and every resource is stopped when the work is done.
Job orchestration. On a cluster the researcher does the plumbing. You write the batch scripts, load the right modules or build the environment, wait in the queue, and stitch preprocessing, training, evaluation and tuning together as separate submissions, with a scheduler that knows nothing about any of them. A managed service does that part. In SageMaker a job is a script, a container with the framework installed, an instance type and the data’s location, launched from Python. The service starts the machine, runs the script, saves the outputs and shuts the machine down, and a hyperparameter tuning job launches and tracks dozens of training jobs on its own. For a researcher without a sysadmin, that is often the strongest reason to use the cloud, and it is what most of this lesson is about.
Energy efficiency. The comparison depends on the hardware and on how well it is shared. Two things matter most.
- Data center overhead. Not all of the electricity entering a data center reaches its computers. Some goes to cooling and power distribution. Large facilities run by AWS and by national HPC centers report about 15% on top of what the computing hardware draws, against 60 to 80% for typical enterprise and departmental server rooms (Uptime Institute, 2025). Cooling matters more as GPUs draw more power. The newest systems, such as NVIDIA’s 72-GPU GB200 racks, need liquid cooling, which removes heat with less energy than air but needs plumbing most campus data centers do not have, and the large cloud and HPC facilities are the ones that can afford to install it. These figures describe the building, not the GPUs in it.
- Hardware utilization. A GPU does no useful work while idle. A lab’s own GPUs sit unused between experiments. A campus cluster does better, because it pools demand across every group at the institution, and the national pools, such as CHTC‘s OSPool and the National Research Platform, do better still, because they pool it across institutions, often at no direct cost to the researcher. Cloud providers pool demand across a larger customer base and across time zones, so a GPU one group releases is running someone else’s job minutes later. Providers do not publish enough utilization data for a direct comparison. The research pools get much of the same sharing benefit. What they do not have is the providers’ hardware budget, so they adopt each new GPU generation more slowly.

There is no option that is better in every case. Campus and national research computing is often the most economical place to start, above all where it is free to the researcher. The cloud becomes the right tool when a project needs newer hardware, extra capacity for a while, or the orchestration that a managed service does for you and a cluster leaves to you.

S3, the notebook instance, SageMaker and Bedrock

AWS offers many services. This workshop uses four pieces that work 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

The workflow is the diagram. Datasets and results live in S3. We work from a small JupyterLab notebook instance. SageMaker launches the compute-heavy jobs on separate machines. For tasks a model AWS already hosts can do, we call Bedrock instead.

Amazon S3 stores the data and the results. S3 is object storage, organized in buckets. Datasets stay in a bucket instead of on the notebook instance, 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. You create a bucket of your own in the next episode and use it for the rest of the workshop. The larger datasets are already in a shared bucket the instructors staged.

The SageMaker notebook instance is where you work. It is a small machine running JupyterLab, started at the beginning of the workshop and stopped at the end, and every notebook in the lesson runs on it. It is where code is written, scripts are tested on small data, jobs are launched and results are read back. It needs little power, because the expensive work runs elsewhere, so it is an ml.t3.medium at about five cents an hour. The cheap notebook stays up while you work, and the larger instances are paid for only while a job needs them. The Notebooks as Controllers episode creates it, and the two after that connect it to S3 and to a copy of the lesson’s code.

The notebook instance is not the only way to drive AWS. Everything it does can be done from the AWS CLI on your own computer, which the optional The Same Jobs from the Terminal episode shows. We start with the notebook for two reasons. It is one more AWS tool worth knowing, and on the shared account it keeps every launched job visible, in a notebook the instructors can open, which is how a problem gets tracked down fast when a learner hits one.

Amazon SageMaker runs the ML jobs. SageMaker is AWS’s managed machine learning service. You give it a training or inference script, the input data and an instance type. It starts that instance, runs the script on it with the ML frameworks already installed, saves the results and shuts the instance down when the job finishes. What that gives a researcher:

The hardware fits the job. A CPU instance such as ml.m5.large ($0.115 an hour) for classical ML and preprocessing, a GPU instance such as ml.g5.xlarge ($1.41 an hour) for deep learning and fine-tuning. The Instances for ML page has the list. Prices throughout the lesson are us-east-1 on-demand, from the price list of 2026-10-16.
Your own code. Preconfigured containers for scikit-learn, PyTorch, TensorFlow and Hugging Face. A script and a requirements.txt are usually all you supply.
Jobs that manage themselves. No batch scripts and no queue. A job is launched from Python, and a hyperparameter tuning job launches and tracks many training jobs on its own.
A bill you can read. Every job reports its billable seconds, and tags split the bill by project. Training a hundred small models on a dataset under 10 GB costs under $20.

Amazon Bedrock serves hosted models. Sometimes there is no model to train or host. You want to send a model a prompt and get a response. Bedrock gives API access to models AWS already hosts, with no GPU instance to provision. What that gives a researcher:

Nothing to manage. You call the model through an API and pay per token. There is no instance to size, start or shut down.
A catalog of models, and only the catalog. Anthropic, DeepSeek, Amazon and others, for generation and embeddings, picked by ID. A model Bedrock does not carry has to run on SageMaker, which the RAG episodes show.
Cheap at research scale. Embedding a few thousand document chunks and answering a few dozen questions costs cents.
A bill you can read. Bedrock calls carry no tags, so the lesson routes them through tagged inference profiles that put the spend under the right project.

The difference between SageMaker and Bedrock is control against convenience. SageMaker runs your own code on compute you choose. Bedrock runs models AWS hosts, with no infrastructure of yours underneath.

How the episodes use them

Through the workshop, code is written in JupyterLab, datasets and results go to S3, training and other heavy jobs go to SageMaker, and Bedrock is called when a hosted model is enough. None of this needs the notebook instance. Everything here can be done from Python on your own computer or from the AWS command line, and the optional The Same Jobs from the Terminal episode shows how.

Key Points

The cloud offers the current GPU generation on demand, billed per second while a job runs. Campus and national research computing is often cheaper for sustained work, and free to many researchers.
Cost and energy efficiency depend on the hardware generation, how busy it is kept, and the data center’s overhead.
S3 stores the datasets and results. A small SageMaker notebook instance is the workspace.
SageMaker runs training, tuning, processing and inference jobs on instances that exist only for the job.
Bedrock gives API access to hosted models, with no GPU instance to manage.

Content from Data Storage: Setting up S3

Last updated on 2026-11-02 | 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

AWS has several places to keep data, with different uses.

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.

The notebook’s disk versus S3

A SageMaker notebook instance has a disk of its own, the volume you size when you create it, and files there 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. That disk is for the copy of a file you are working on. The data itself lives in S3, where every job can read it, where it survives the instance, and where it costs cents a month.

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)

Store data in S3 and load it onto an instance as needed. This gives:

Separation of storage and compute: data in S3 remains accessible when instances are stopped or deleted.
Easy data sharing: a bucket can be read by team members and by other projects.
Integration with AWS services: SageMaker, Lambda, and other AWS services read directly from and write back to S3.
Scalability: a bucket holds any amount of data; an instance’s disk does not.
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 S3 URI 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 (wattbot2026, badgerscribe, badgercode, braindecoding, poisonedwell, snapshotwi, or sinkorswim for the workshop itself), 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 fails from the job with an error naming the bucket’s region (PermanentRedirect: “the bucket is in this region: …”). Reading data across regions from a notebook does work, but every byte is billed as inter-region transfer ($0.01 per GB between us-east-1 and us-east-2), 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: tags are how the bill is broken down by project, person and purpose. Add these three keys, spelled exactly as shown (they are the account’s approved cost-allocation keys, see KB 72454), and use the same three on every resource you create in this lesson.

Project: your challenge, one of WattBot2026, BadgerScribe, BadgerCode, BrainDecoding, PoisonedWell, SnapshotWI, or SinkOrSwim, the workshop’s own project, if you are not on a challenge yet.
Name: your last and first name run together, e.g. DoeJohn
Purpose: one of five fixed values, here Storage

Tag values are case-sensitive, so use capitalized words with no spaces and the same spelling everywhere. SinkOrSwim and sinkorswim would be two separate rows in Cost Explorer. (The bucket name is lowercase because S3 requires it. Bucket names are not tags.)

The account enforces these tags on every job and notebook instance. A launch with a tag missing, blank or off the lists is refused with AccessDenied ... explicit deny, and the message does not say which tag. Buckets are not checked. A bucket with a wrong tag is created anyway, and its storage cannot be attributed on the bill. The training episode checks the tags in the notebook before each launch and names the problem, and If a Launch Is Refused has the full list of rules and fixes. The rules are not part of SageMaker. The instructor sets them up following the AWS account preparation page.

`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. Copy the key names exactly; a variant spelling does not report.

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/*"
			]
		}
	]
}

The policy, statement by statement:

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-use, which you will pick when you create your notebook instance in the next episode and which every job inherits, and ml-sagemaker-bedrock-use, an equivalent role from an earlier term, 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 role can list the bucket’s name but neither the notebook nor a job can read or write its contents; the first read in the next episode fails with “access denied”.

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 AWS costs $0.09/GB. (AWS waives the first 100 GB a month, but per organization, not per account; a university’s organization uses that elsewhere, so count on the full rate.)
In-region transfer (e.g., S3 to EC2) is free, while cross-region data transfer is charged ($0.01/GB between us-east-1 and us-east-2). 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.01 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 request charges beyond the GETs for training. 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. The download line is about 80% of the total at every size, since storage and transfer both scale with gigabytes: moving a dataset out of AWS once costs about four times what keeping it there for a month does, so keep large data in the region and bring code to it.

Wrong region, 100 GB: 100 reads x 100 GB x $0.01 = $100 of inter-region transfer, on top of the $11.30 above. The one-time download to a laptop was $9; the misplaced bucket costs ten times that, silently, and keeps costing with every job.

Removing unused data (complete after the workshop)

When you are done with the data, remove the files and buckets you no longer need.

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-11-02 | 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 is a transfer out of AWS, $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, capitalized: 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. The Titanic episodes reuse this name as the prefix of their training jobs, which is how you 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: leave the default.
Additional configuration: 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, and the reason is rarely your dataset. The disk only has to hold what the notebook itself opens. A dataset a training job reads stays in S3 and is copied onto the job’s own instance, which has its own disk, set per job, so a 50 GB dataset needs no change here. What cannot go through this notebook at all is anything bigger than its memory, 4 GB on an ml.t3.medium. The Accessing S3 episode says what to do with a file that size, and the Instances for ML page sizes the job instances by dataset. If the disk does fill, stop the instance, choose Edit, and raise the volume size. It can grow, not shrink.
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: leave the default. pip install into the kernels needs no root; root is for system packages, which nothing in the lesson installs.
- 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): leave the defaults.
Git repositories configuration (skip): leave it empty; the lesson repository is cloned from a notebook cell below.

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, one of `WattBot2026`, `BadgerScribe`, `BadgerCode`, `BrainDecoding`, `PoisonedWell`, `SnapshotWI`, or `SinkOrSwim` for the workshop itself
`Name`	your name, e.g. `DoeJohn`
`Purpose`	what this resource is for, `Notebook`

Same spelling as on the bucket. The Project and Purpose lists from the S3 episode apply to every job you launch from this instance, and a value off them is refused at the first launch; the training episode shows what that looks like and how to check first.

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 puts one at the top of every formatted disk for its repair tool, and SageMaker/ is its own disk. Leave it alone. That disk is the volume whose size you set under Additional configuration on the create form (20 GB). 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-11-02 | 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

Each notebook starts by setting up the AWS session it works in:

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

Messages this cell prints

The SageMaker SDK prints a few lines of its own when it is imported and when a session is created, and none of them is an error. sagemaker.config INFO - Not applying SDK defaults from location: ... means it looked for an optional defaults file in two standard places and found none. A line about telemetry says it reports usage to AWS. A warning that this version of the SDK is deprecated means a newer major version exists. The code in this lesson is written for the version the notebook instances ship. All three appear in later episodes too, and can be ignored.

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

There are two ways to get a file from S3 into a notebook: read it straight into memory with get_object, or copy it onto the instance’s disk with download_file. In memory is the usual choice here, since the notebook only inspects data and the jobs read S3 themselves; a copy on disk is for a file you will read many times in one session, or that another program needs as a file.

A) Reading data directly from S3 into memory

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 downloading a copy of titanic_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

download_file writes the object to the instance’s disk, on the 20 GB SageMaker/ volume. There is no transfer charge, because the instance and the bucket are in the same region, and in-region transfer between AWS services is free. The same call from a laptop, with the AWS CLI or with boto3, pulls the data out of AWS at $0.09 per GB. The optional The Same Jobs from the Terminal episode covers working from a laptop and when that charge matters. The file stays on the instance until you delete it.

Example

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…

the instance’s IAM role is ml-sagemaker-use
the bucket policy from the S3 episode names that role

Callout

When the file is bigger than the notebook

Both methods above assume the object fits. The notebook instance has a 20 GB disk and, on an ml.t3.medium, 4 GB of memory, and a model checkpoint or an image archive can be larger than either. A learner in a past class uploaded a multi-gigabyte PyTorch .pth file, found download_file could not fit it on the disk, and found that the stream from get_object could not be handed to torch.load either, because a .pth is a zip archive and torch.load needs to seek inside it, which a stream does not allow. In order of what to try:

Check the size first. show_bucket, below, prints every object’s size, so you know before you copy.
If it fits in memory, read it into memory. get_object gives bytes, and io.BytesIO(bytes) wraps them as a seekable file that torch.load, zipfile and pandas all accept. That is fine up to about a gigabyte on this instance.
If it does not fit the notebook, do not bring it here. Give it to the job that needs it as an input channel, and SageMaker copies it onto the job’s instance, which has the disk and memory the job was sized for. That is how the fine-tuning episode passes 7.5 GB of model weights, which never touch the controller.
A library such as s3fs can present an S3 object as a seekable file, so torch.load works on it directly, but it reads by range requests and is slow for a checkpoint. Use it to read a slice of a large file, not to load a model.

The rule the lesson follows is that the notebook holds what the notebook opens, and everything else goes from S3 to a job.

Check the current size and storage costs of bucket

Bucket size is what storage is billed on, and the bill is a day or two behind, so read the size from the notebook.

The code below calculates your bucket size. The pieces that matter:

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).

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

The same computation is the helper function get_s3_bucket_size(bucket_name). 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)

Then:

PYTHON

helpers.get_s3_bucket_size(bucket_name)

The shared bucket

The S3 episode mentioned a second bucket: the one the instructors staged with the larger datasets, which the fine-tuning and RAG episodes read from. Your role can list it and read from it, and nothing else; an upload to it is refused. Look at what is in it now, so that the names in those episodes are not a surprise. show_bucket prints each top-level folder with how many objects it holds and how big they are:

PYTHON

shared_bucket = "ml-marathon-shared"   # the instructor-staged, read-only bucket; the later episodes have the name filled in
helpers.show_bucket(shared_bucket)

Each folder is one episode’s data: badger-scribe/ holds the scanned pages and the model weights for the fine-tuning episode, wattbot/ the papers and questions for the RAG episodes. The same call with a prefix, helpers.show_bucket(shared_bucket, "wattbot/"), goes one level down.

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.

The rates are S3 Standard in us-east-1 from the AWS price list of 2026-10-16; other regions differ by a few percent and other storage classes by more (AWS S3 Pricing).

PYTHON

# AWS S3 Standard Storage pricing for US East (N. Virginia) region
# Pricing tiers from the AWS price list of 2026-10-16
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}")

The same code is 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}")

Writing output files to S3

Uploading is the reverse call. 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:
    keys = [obj['Key'] for obj in response['Contents']]   # each object's key: its path in the bucket
    print(len(keys), "objects; the first 20:")           # a bucket that has seen a few training jobs holds hundreds of files, so do not print them all
    for key in keys[:20]:
        print(key)
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)   # every key in the bucket, through the paginator
file_list[:20]

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 training episode launches a job that writes its outputs to your bucket. Why does the lesson keep everything in one bucket, organized with prefixes like docs/, instead of creating a bucket per task?

Show me the solution

One way to do it:

PYTHON

size_before = helpers.get_s3_bucket_size(bucket_name)

import time
with open("Notes2.txt", "w") as f:
    f.write("A second note, to see the bucket grow.\n" * 20)
notes2_key = f"docs/Notes2-{int(time.time())}.txt"   # a new key each run; uploading over an existing key of the same size would show no growth
s3.upload_file("Notes2.txt", bucket_name, notes2_key)

size_after = helpers.get_s3_bucket_size(bucket_name)
print(f"before {size_before['size_mb'] * 1024:,.1f} KB, after {size_after['size_mb'] * 1024:,.1f} KB: grew by {(size_after['size_mb'] - size_before['size_mb']) * 1024:.1f} KB")   # the helper returns a dict with size_mb and size_gb

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

Read S3 objects into memory with get_object; copy one to the instance’s disk with download_file only when a program needs it as a file.
Periodically check storage usage and costs to manage S3 budgets.
Write outputs back to your bucket under a prefix, and read the bucket’s size from the notebook to see what accumulates.
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-11-02 | 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 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 goes under version control

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.

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.

Treat the token 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

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

The same token clones a private repository

The lesson repository was public, so git clone in the controllers episode needed no credentials. A private repository, your own or a collaborator’s, needs the token at clone time as well as at push time, and the way to supply it without writing it to disk is to put it in the URL for that one command, from the variable you just set. The repo scope on the token covers it.

PYTHON

# Clone a private repository with the token in the URL for this command only; nothing is written to disk
!git clone https://{username}:{token}@github.com/YOUR_GITHUB_USERNAME/your-private-repo.git /home/ec2-user/SageMaker/your-private-repo

Git stores the URL, token included, in that clone’s .git/config, which is on the instance’s disk. Either strip it after cloning, git remote set-url origin https://github.com/YOUR_GITHUB_USERNAME/your-private-repo.git, and supply the token the same way on each later pull, or revoke the token when you are done, as the cleanup episode’s checklist says. Do not paste a token into a cell as text. The notebook file would then carry it.

Step 2: Configure Git settings

Git records an author on every commit. Set yours in this repository: without --global the setting goes into the repository’s own .git/config, under SageMaker/, which survives a stop. (--global would write to your home directory, which is reset when the instance stops.)

PYTHON

!git config user.name "Your name"             # appears as the author in the commit history
!git config 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, helpers.convert_files() in the lesson’s scripts/helpers.py converts them all. 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.

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, not with 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.

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 and temporary files. (The NumPy archives the training episodes write are the challenge at the end.)

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 adds to the remote instead of being rejected.

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 https://{username}:{token}@github.com/{username}/sagemaker-work.git main --allow-unrelated-histories --no-edit   # the token authenticates the pull too, which a private repository needs

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 four commits in the history: GitHub’s initial commit, your two, and the merge the pull made.

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 pull https://{username}:{token}@github.com/{username}/sagemaker-work.git 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; it converts every .py in the current directory, so keep real modules out of that folder.

Challenge

Challenge: ignore the big files too

The PyTorch episode saves NumPy arrays (.npz) in its working folder, and your own notebooks will do the same in this one. 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 initialized folder needs --allow-unrelated-histories.

Content from Training Models in SageMaker: Intro

Last updated on 2026-11-02 | Edit this page

Estimated time: 65 minutes

Overview

Questions

What are the differences between local training and SageMaker-managed training?
What is a SageMaker training job, and what does it cost?
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.
Launch a training job with the XGBoost estimator: your script, AWS’s container, a rented instance.
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.

Open the notebook from: /ML_with_AWS_SageMaker/notebooks/06_Training-models-in-SageMaker-notebooks.ipynb. When asked for a kernel, pick conda_pytorch; every notebook from here to the end of the lesson uses it. The conda_python3 kernel has no PyTorch, and on current instances no scikit-learn either.

Test on a small sample first, then scale up

The Titanic data trains in a second on the notebook, and nothing about it needs the cloud. What it shows is the pattern. Test the code on a small sample first, on the cheapest machine you have, and once it works, scale up to a more powerful instance when the data or the model calls for one. The notebook instance is that cheap machine here. It is a rental too, at about five cents an hour, but a crash on it costs a few seconds of that, where a crash on a rented GPU costs the minutes the instance spent starting up, at twenty times the rate. A SageMaker training job is that scale-up step. The notebook describes the job and launches it, and SageMaker rents the instance, runs the script on it and releases it when the script ends. The fine-tuning episode follows the same two steps with a 3B model on a GPU, which does need the cloud.

The notebook is not the only way to launch a job. The same launch is a few lines of the AWS CLI from a laptop, which the optional The Same Jobs from the Terminal episode shows, and on your own account that is a fine way to work once you know the pattern. On the shared account, launch from the notebook. Every job then comes from a notebook the administrators can open and read, which is what lets them find and fix a problem quickly if one arises.

This episode walks through both steps once with XGBoost. First the script runs on the notebook, to check that it handles the data and the arguments. Then the same script runs as a training job, with the data supplied from S3, and we look at how the job receives its inputs, where it writes its model, how to estimate its cost before launching and how to read the billable time afterwards.

Initial setup

Set up the project folder

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

Connect to SageMaker and S3

The notebook needs two interfaces to AWS. boto3 reads and writes S3 and inspects resources directly, and the SageMaker SDK provides the higher-level objects that launch training jobs. We also look up the notebook’s IAM role, which the jobs use to read their inputs and write their outputs.

Boto3 API

Boto3 is the AWS SDK for Python: every AWS service, called from Python, with the notebook’s role as the identity.

PYTHON

import boto3                          # AWS SDK for Python: low-level access to every AWS service (S3, EC2, ...)
import pandas as pd
import os
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}")

Check the training script before launching a job

A training job runs your script on a machine SageMaker rents for it, so the script has to work before it is worth renting anything. Run it here first, on the notebook instance, on the small Titanic data. If it crashes here, it costs a few seconds of a five-cent-an-hour machine. If it crashes as a job, you pay for the minutes the job’s instance spent starting up, on a machine that costs more. Every later episode does the same, a few epochs on a subset on the cheapest machine, then the full run.

Challenge

Understanding the XGBoost Training Script

Take a moment to review the ML_with_AWS_SageMaker/scripts/train_xgboost.py script cloned in the controllers episode. 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 is what the SageMaker runs are compared against.
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

For this test we download copies of the Titanic files from S3 onto the notebook instance. The managed job later reads its own input from S3 and never uses these local copies.

Callout

We download the data into the notebook to test the code, as a local pipeline would. In a pipeline that scales, the data lives in S3 and the jobs read it from there.

The download costs nothing beyond a fraction of a cent in requests. The notebook instance and the bucket are in the same region, and transfer between AWS services inside a region is free. Pulling the same files to a laptop with the AWS CLI is a transfer out of AWS at $0.09 per GB, which the optional The Same Jobs from the Terminal episode covers.

Run the following cells to download the training and test CSVs. If the files do not appear in the Jupyter file browser, refresh it. A permissions error usually means the notebook role or the bucket policy does not allow the requested S3 operation.

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_xgboost.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 on the cheapest machine, 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 adds randomness between trees, which reduces overfitting.
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}")

The run takes seconds. The script prints its own training time, and the cell prints the total including data loading; both numbers are the baseline the SageMaker runs below are compared with.

Sanity check: Quick evaluation on test set

Nothing here is SageMaker specific: apply the model to the test set on the notebook instance to check that it trained properly.

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.

Run the script as a SageMaker training job

Now that the script is checked, we submit it to SageMaker. The notebook sends the job configuration, the training code, the instance type, the S3 location of the data and where the output goes. SageMaker starts a separate instance, runs the script in AWS’s XGBoost container, saves the model to S3 and releases the instance when the job ends. The notebook stays free for launching and inspecting jobs. It is not the machine doing the training.

Which instance to start with?

The Titanic data and the XGBoost model are small, so we start with one ml.m5.large CPU instance. The point is to learn how job configuration and billing work, not to make this tiny model train faster. For larger datasets, instance memory and CPU count matter, and the Instances for ML page is the place to start choosing.

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.

Estimators

Every training job runs inside a container image. The XGBoost class we use below is a framework estimator: it looks up AWS’s prebuilt XGBoost image for the region and runs your script inside it. PyTorch, SKLearn, TensorFlow and HuggingFace are the same idea for other frameworks, each with a framework_version that picks the image, and a requirements.txt in source_dir is pip-installed in the container before the script starts, for packages the image lacks. The plain Estimator class takes an image_uri directly and is for a container of your own, or for one of AWS’s built-in algorithms, which train with AWS’s code and no script of yours; the optional challenge at the end of this episode runs XGBoost that way for comparison. The list of prebuilt images says what each one has installed.

Launching the training job

We configure an XGBoost estimator to run the same train_xgboost.py we tested above. The estimator describes the training environment. Its hyperparameters become command-line arguments, and TrainingInput tells SageMaker which data to copy from S3. Before launching we set the cost-tracking tags and estimate how long the instance will be billed. The full code is below, with explanation after it.

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.

So we tag each training job at launch time, with the same three keys as the bucket and the notebook instance.

PYTHON

name = "DoeJohn" # replace with your name
project = "SinkOrSwim"   # your challenge: WattBot2026, BadgerScribe, BadgerCode, BrainDecoding, PoisonedWell or SnapshotWI; SinkOrSwim is the workshop itself
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},
]
helpers.check_launch(job_tags)   # the account's rules, checked here first: AWS's own refusal does not say which rule
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. 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 a few seconds of training plus two or three minutes of startup: about three billable minutes, 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 a key the script has no flag for makes argparse exit with “unrecognized arguments”, which fails the job. (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; the next episode shows what a second instance does
price_per_hour = 0.115 # ml.m5.large, us-east-1 on-demand (2026-10-16), from the Instances for ML page; check it for your region

# The runtime cap, in seconds: SageMaker stops the job and releases the instance when it is reached. Set it a little above the estimate;
# the estimate here is three or four minutes, so 20 minutes leaves room for a slow container pull without letting a hung job bill for hours.
max_run = 20*60 # 20 minutes

# 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

helpers.check_launch(job_tags, instance_type, max_run)   # tags, instance type and cap against the account's rules, before anything is sent

# 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}")

While a job runs, its status is on the console under Training jobs. For a job too long to sit through, see “Launching a job without waiting for it” below.

Callout

If the launch is refused

The account checks every launch against a short list of rules, and refuses one that breaks any of them with an AccessDeniedException ending in with an explicit deny in an identity-based policy. The rules: the three tags Name, Project and Purpose must be present, Name must not be blank, Project must be one of SinkOrSwim, WattBot2026, BadgerScribe, BadgerCode, BrainDecoding, PoisonedWell, SnapshotWI, Purpose must be Storage, Notebook, Training, Tuning or Inference, the instance type must be one the lesson uses, the runtime cap must be 48 hours or less, and endpoints and EC2 instances cannot be created at all. The instance list is single-GPU types only, and that is deliberate: with the 48-hour cap, the most one hung job can bill on the largest allowed type is about $125, where a four-GPU box would be $340 to $1,200.

The error message does not say which rule was broken. AWS does not let an administrator attach an explanation to a refusal. The lesson’s helper module does the explaining instead. helpers.check_launch(job_tags, instance_type, max_run) checks a launch against the same rules before it is sent and says what to change; the tag cell above and every launch in the later episodes call it. helpers.rules() prints the rules. helpers.explain_refusal() is for afterwards: it reads the refused requests of the last hour from CloudTrail, the account’s log of every API call, and says which rule each one broke. All three read the rules from scripts/guardrails.json, the same file the administrator attached to the account, so they cannot disagree with it. A launch that passes check_launch and is still refused is almost always a missing tags=job_tags line on the estimator. The If a Launch Is Refused page has all of this on one page for your own notebooks after the workshop.

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. The minute or so before that, while SageMaker finds an instance, is not billed. (The job’s console page also shows a “Training time”; on an on-demand instance it is the same number, and neither one is your script’s own training time, which only the script’s log has. The two lines the SDK prints when fit() returns, Training seconds and Billable seconds, are the same pair: they differ only on spot instances, which this lesson does not use, so read either one.) 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
    if "BillableTimeInSeconds" not in desc:   # the job has not finished: no bill yet
        print(f"{job_name}: status {desc['TrainingJobStatus']}; billable time and the model path appear once it is Completed")
        return desc
    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 profiler and debugger, written for every training job. Ignore them; disable_profiler=True and debugger_hook_config=False on the estimator turn them off.

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.

The code below:

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}")

Monitoring training

The job’s page in the console has everything the notebook printed and more. Open Training jobs; the list is the whole account’s, so find yours by name, SinkOrSwim-DoeJohn-xgboost-<timestamp>, and click it.

Status: InProgress, Completed, Failed or Stopped.
View logs, under Monitor: the job’s printed output in CloudWatch, kept after the instance is gone. Job Metrics next to it shows CPU, GPU and memory use over the run.
Output: the S3 path of the model artifact.
Stop, under Actions: ends the job and its billing. Check the name first; the list has other people’s jobs in it.

Launching a job without waiting for it

Everything fit() waited for happened on the job’s instance, not this one. Once SageMaker has the request, the job pulls its container, trains, uploads the model to S3 and writes its log to CloudWatch, and none of that needs this notebook. wait=True, the default, only makes the cell block and stream the log here, which is fine for a four-minute job. For a run of hours it means sitting on a blocked cell, or leaving a notebook instance running all night so that a cell can watch a log you could read in CloudWatch.

wait=False returns as soon as the job is accepted. Run the next cell, or stop the notebook instance and go home; the job runs the same either way.

PYTHON

# Launch and return at once: the cell ends in seconds, and the job runs on its own instance
xgboost_estimator.fit({'train': train_input}, wait=False)
job_name = xgboost_estimator.latest_training_job.name
print("launched:", job_name)

While it runs, its status and log are on its page under Training jobs, from any browser. When it finishes, the model is in S3 under output_path, whether or not anything was watching. To read its status and bill from a notebook, the same call as before:

PYTHON

desc = job_cost_report(job_name, price_per_hour)   # prints the status while the job runs; the bill and the model path once it is Completed

If the notebook instance was stopped in between, the kernel’s variables are gone, so set job_name from the printed output or from the Training jobs list before running that cell. xgboost_estimator = XGBoost.attach(job_name, sagemaker_session=session) rebuilds the estimator from the job’s record, waiting for the job if it is still running, if you want the evaluate cells above to work as before.

Interrupting the cell, restarting the kernel and stopping the instance all leave the job running. Only Stop on the job’s page, or the job’s own cap, ends it. A job you launched and forgot bills until one of those happens. Two estimator settings limit what that can cost.

For any run longer than half an hour

A cap you chose. max_run is the most a job can bill, in seconds, and every estimator in this lesson sets it. Put it a little above your estimate, twice the estimate is a fair margin, and never above the account’s 48-hour ceiling, which the account refuses. Leave it out and the SDK sends 24 hours, which on an ml.g5.xlarge is $34 for a job that hung in its first minute.

Checkpoints. The cap ends a hung job, but without checkpoints the work it did is gone with it, and so is the bill. Two estimator arguments fix that: checkpoint_s3_uri= names an S3 prefix, and checkpoint_local_path="/opt/ml/checkpoints" names a folder on the training instance that SageMaker copies to that prefix while the job runs. The script’s part is to save its state into that folder at the end of every epoch and, when it starts, to resume from whatever it finds there. A job that is killed, or that crashes, then costs one epoch instead of the run, and a second launch with the same checkpoint_s3_uri carries on. It is not worth setting up for this one-minute job; the fine-tuning episode’s script does it, and its estimator has the two lines to copy.

Challenge

Launch, lose the kernel, read the bill anyway

Do this once while the job is short, so that it is familiar when the job is long.

The wait=False cell above launched a job; note the name it printed.
From the Jupyter menu choose Kernel, Restart. Every variable is gone, which is what a stopped instance would also do to you, without the minutes of waiting for it to start again.
Rerun the setup cells at the top of the episode (the session, bucket and role), set price_per_hour = 0.115 again, and rerun the cell that defines job_cost_report. Do not rerun the launch cells.
Set job_name to the name you noted and run the job_cost_report cell.

What was the job’s status when you looked, and what did it cost compared with the first job?

Show me the solution

The status is InProgress if you were quick and Completed otherwise. The cost is the same as the first run, since the job did the same work, and it did not care that the kernel it was launched from no longer existed. The only thing that connected the kernel to the job was the name, and if you lose that too, the Training jobs list, filtered by your notebook instance name, has it.

Challenge

Optional: the same job with AWS’s built-in algorithm

Every job so far ran your script in AWS’s XGBoost container. The container also carries AWS’s own XGBoost training code, the built-in algorithm, which you get by using the plain Estimator with the image and no entry_point. Nothing of yours is uploaded, so there is nothing to write; the price is that it does only standard XGBoost. Preprocessing has to happen before the data reaches it, and the data has to arrive as a CSV with no header, every column numeric and the label first. Run the Titanic job that way and compare the bill and the accuracy with the script version.

Show me the solution

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

%%script false --no-raise-error
# ^ remove this line to run the cell; the guard keeps Run All from launching the extra job
from train_xgboost import preprocess_data

# The same titanic_train.csv the script job trained on. The script ran preprocess_data inside the container;
# the built-in algorithm runs no code of ours, so the same preprocessing is applied here instead
X, y = preprocess_data(pd.read_csv(train_filename))   # fill missing values, encode Sex/Embarked, drop text columns

# Same rows, written the way the built-in algorithm requires: 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 under its own prefix so it is not 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)

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, and we are classifying.

PYTHON

%%script false --no-raise-error
# ^ remove this line to run the cell
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

%%script false --no-raise-error
# ^ remove this line to run the cell
desc = job_cost_report(xgboost_estimator_builtin.latest_training_job.name, price_per_hour)

The bill is the same two or three billable minutes, since both jobs spend their time on startup and a fraction of a second on training, and the accuracy matches to within the randomness of subsampling. The built-in route saves writing a script and costs the preprocessing step and every bit of logic of your own. Every later episode uses the script route, because every later job has such logic.

The next episode asks what to do when a job takes too long: a bigger instance, a GPU, or more than one machine, measured on a PyTorch network.

Key Points

Setup: a SageMaker session, the execution role and the bucket name are what every job needs from the notebook.
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.
Estimators: XGBoost(entry_point=...) is AWS’s container running your script, the pattern every later episode uses with PyTorch and HuggingFace.
Training data channels: TrainingInput names the S3 data a job copies onto its instance before the script runs.
Estimate the cost before launching (rate × minutes, startup included), then read back BillableTimeInSeconds with job_cost_report() after every job and compare.
Long runs: launch with wait=False and carry on, or stop the notebook instance; the job runs and finishes without it, and job_cost_report(job_name, ...) reads its bill later. Set max_run a little above the estimate, and checkpoint to S3 on anything longer than half an hour.

Content from Training Models in SageMaker: PyTorch Example

Last updated on 2026-11-02 | 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.

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

The same pattern with a neural network

The previous episode tested an XGBoost script on the notebook and then launched it as a managed training job. This episode follows the same pattern with a small PyTorch neural network on the same Titanic prediction task. Keeping the data and the task fixed means that what changes between runs is only the instance, so CPU and GPU can be compared on their own.

A GPU costs about six times more per hour than the CPU instance, so it has to train the model at least six times faster to be worth it. For a network this small, most of a job’s time is startup, not training. We run the same job on both and read the bills. Then we try two instances instead of one, and see why that does not halve the time.

Setup

We use the same controller notebook, SageMaker role, S3 bucket and Titanic files as in the XGBoost episode. The notebook prepares and uploads the data, then launches the training job on the instance we choose.

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'

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}")

Work in the Titanic folder beside the lesson clone, with the helper module on the path, and record which instance this notebook runs on so that runtimes can be compared.

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

The prediction task is unchanged, and only the training framework is different. So the attention goes to the PyTorch training workflow and to what the choice of instance does to runtime and cost.

Preparing the data

Unlike the XGBoost script, this PyTorch script expects features and labels already processed into NumPy arrays. We prepare those arrays once, save the training and validation sets as .npz files, and upload them to S3. That keeps preprocessing out of the managed job, and every run we compare trains on the same inputs.

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 to the working directory 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 instead of retyping 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 $0.023 per GB per month; a preprocessed copy of a small dataset costs nothing worth counting.

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.")

Check the training script

Before launching a managed job, we run 100 epochs on the controller notebook as a short test. It checks the data paths and the training code and gives a rough time per epoch. It is only a local test. The runs we compare are launched as separate SageMaker jobs.

PYTHON

import torch
import time as t # Measure training time locally

epochs = 100          # a short test run
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}
local_elapsed = t.time() - start_time   # kept for the estimate below
print(f"Local training time: {local_elapsed:.2f} seconds, instance_type = {local_instance}")

The network as a SageMaker training job

Now we submit a 1,000-epoch run as a SageMaker job. Its compute is chosen on the estimator, independently of the notebook that launches it. The first run uses the same ml.m5.large CPU instance as the XGBoost job, so there is a baseline before trying a GPU.

A GPU is not necessarily worth it for this data. The CPU instance costs $0.115 an hour and an ml.g4dn.xlarge GPU instance $0.74 an hour on the Instances for ML page. The GPU would have to save enough time to make up that difference, and the model may be too small to keep it busy. We measure it instead of assuming the more powerful hardware is the better choice.

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"   # your challenge: WattBot2026, BadgerScribe, BadgerCode, BrainDecoding, PoisonedWell or SnapshotWI; SinkOrSwim is the workshop itself
purpose = "Training"        # one of: Storage, Notebook, Training, Tuning, Inference

job_tags = [
    {"Key": "Name", "Value": name},
    {"Key": "Project", "Value": project},
    {"Key": "Purpose", "Value": purpose},
]
helpers.check_launch(job_tags)   # the account's rules, checked here first: AWS's own refusal does not say which rule
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(), a rough figure for how long it will run and what it will cost. There are two ways to make it, and which one applies depends on the size of the model. For a model this small, timing a short test run is more useful than counting the arithmetic, so the measured time comes first and the counted arithmetic, which the larger models later in the workshop need, after it.

Small models: time a short run and scale it. The 100-epoch test gives a time per epoch. Each epoch costs the same as the last, and the script trains on all the rows in one step per epoch, so training time scales with epochs and with rows. Scale the measured time to 1,000 epochs, then add the two or three minutes of container startup and data download, which are billed and are most of the bill for a job this small. The result is not exact, because the notebook and the training instance are different machines, but it says whether the job will cost cents or dollars:

PYTHON

local_seconds_per_epoch = local_elapsed / epochs                 # from the short run above: its 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, from the Instances for ML page

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.7 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. From a few hundred million parameters up, the training arithmetic takes longer than the job’s startup, so the arithmetic is the cost, and it can be counted before any instance exists. The work is counted in floating-point operations, FLOPs, one multiply or one add, and a GPU’s speed in FLOPS, operations per second, about 70 trillion for an A10G. Training one token through a dense network or a transformer costs about 6 FLOPs per parameter, 2 for the forward pass and 4 for the backward one. Multiply by parameters, examples and epochs for the count, and divide by the speed the GPU reaches, about a quarter of its data-sheet peak in practice. It is a scale estimate, not a runtime prediction:

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 the arithmetic comes out as a fraction of a second, which is why timing a short run is the method to use for it:

PYTHON

from train_nn import TitanicNet
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"])          # one example = one passenger, a row of 7 numbers; 569 of them after the validation split
tokens_per_example = 1                                          # the network reads a whole row in one step, so a row is one token
flops = 6 * n_params * n_examples * tokens_per_example * 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 once the arithmetic outweighs the startup, from a few hundred million parameters up, which is where the fine-tuning episode uses it on a 3-billion-parameter model. Its four inputs are the four numbers in the formula, and what each one is depends on the data:

params_billions: every parameter in the model, in billions. TitanicNet has 2,625, which is 0.0000026; a BERT-base text classifier has 110 million, 0.11.
examples: one example is one training item, whatever the model takes in at a time. A passenger’s row here; one tweet, one abstract, one document for a text model; one page image in the fine-tuning episode.
tokens_per_example: how many units the model processes for one item, which is what the arithmetic scales with. A row through TitanicNet is one step, so 1. A text model works in tokens, about 1.3 per English word: a tweet is 40, an abstract 256, a short paper 2,000. An image through a vision transformer is one token per patch, 196 for a 224 × 224 image in 16 × 16 patches.
epochs: passes over the whole set.

The three data numbers multiply to the tokens the model trains on. Here it is on a run that does need the rule: a BERT-sized text classifier trained from scratch on a million abstracts, three passes:

PYTHON

helpers.estimate_training_cost(   # the semicolon at the end keeps Jupyter from also displaying the dictionary the function returns
    params_billions=0.125,      # a BERT-base-sized model, 125 million parameters
    examples=1_000_000,         # a million abstracts
    tokens_per_example=256,     # 190 words per abstract, about 256 tokens
    epochs=3,                   # three full training passes over all abstracts
);

Each line:

tokens: examples × tokens per example × epochs. 768 million tokens is the amount of work; everything below is derived from it.
memory: 16 bytes per parameter for full training (the weight, its 32-bit master copy, its gradient and the optimizer’s two running averages, in the fine-tuning episode’s table) is 2 GB for 125 million parameters, plus 4 GB the helper allows for activations: 6 GB. The cheapest GPU on the lesson’s list with that much memory is the ml.g6.xlarge, one L4 with 24 GB, so that is the instance the rest of the lines price.
compute: 6 FLOPs × 0.125e9 parameters × 768e6 tokens is 5.8e17; the helper multiplies by 1.5 because it assumes gradient checkpointing, which recomputes the forward pass during the backward one to save memory, giving 8.6e17. Divided by a quarter of the L4’s 121 TFLOPS peak, 30 trillion a second, that is 28,500 seconds: 7.9 GPU-hours.
wall-clock: the GPU-hours on one GPU, plus the 10 minutes of startup the helper assumes for a job this size.
cost: the wall-clock hours at the instance’s price, $1.13 an hour, so about $9 a run. The second figure is the same times the number of runs, one here; the hyperparameter-tuning episode is where that multiplier bites.
max_run: twice the wall-clock estimate, in seconds, ready to paste into the estimator. A job that runs past this is killed, so a hung run costs at most $18, not a week.

The function also returns the same numbers as a dictionary, so a notebook can use max_run_seconds and instance_type directly, with no copying from the printout.

What training a language model would cost

Run the same function on two sizes a project might be tempted by: a 1-billion-parameter model and a 7-billion one, each trained from scratch on a million papers, one pass. Nothing is launched; the function only does the arithmetic.

PYTHON

helpers.estimate_training_cost(
    params_billions=1,          # a 1B model, the size of the smallest open-weight language models
    examples=1_000_000,         # a million papers
    tokens_per_example=2000,    # 1,500 words per paper, about 2,000 tokens
    epochs=1,                   # one full training pass over all papers
);

PYTHON

helpers.estimate_training_cost(
    params_billions=7,          # a 7B model, the size most "small" open-weight LLMs are
    examples=1_000_000,         # the same million papers
    tokens_per_example=2000,    # 1,500 words per paper, about 2,000 tokens
    epochs=1,                   # one full training pass over all papers
);

Read the two against the 125M run above. The 1B model still fits on the cheapest GPU, but the arithmetic is 165 GPU-hours: a week of wall-clock on one GPU, $187, and a max_run of 330 hours, which the account refuses, since no job may run longer than 48 hours. A run like that has to be split into checkpointed jobs of a day or two each, or moved to a four-GPU box, an ml.g6e.12xlarge, which finishes in about 14 hours for about $180, and which the class account does not allow. The 7B model no longer fits on one GPU at all: 116 GB of weights, gradients and optimizer state need the four-GPU ml.g6e.12xlarge at $13 an hour, four days of it, $1,270 for one pass over two billion tokens. And two billion tokens is a small corpus; the models on Bedrock were trained on thousands of times more.

That is why nothing in this lesson trains a language model. What the lesson does instead is in the fine-tuning episode, LoRA on a 3-billion-parameter model, which trains tens of millions of adapter parameters on 251 pages for about a dollar, and the RAG episodes, which use a trained model through Bedrock by the token. The function is for the day a project asks for more: run it first, and the number says whether the run is a job, a grant line, or something to do differently.

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, one token per patch for an image through a vision transformer (196 for a 224 × 224 image in 16 × 16 patches); 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, time a short run instead and say what you would measure. Write the estimate down with the max_run you would set, and keep it: the first 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 × 3e8 = 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 short run 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
    if "BillableTimeInSeconds" not in desc:   # the job has not finished: no bill yet
        print(f"{job_name}: status {desc['TrainingJobStatus']}; billable time and the model path appear once it is Completed")
        return desc
    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.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. If the launch is refused with an AccessDeniedException ending in explicit deny, the account’s checks are the cause: a missing tag, an instance type outside the lesson’s list, or a resource kind the lesson never creates; the first training episode’s callout If the launch is refused says what to look at.

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

# The runtime cap, in seconds: SageMaker stops the job and releases the instance when it is reached. Set it a little above the estimate;
# the estimate here is three or four minutes, so 20 minutes leaves room for a slow container pull without letting a hung job bill for hours.
max_run = 20*60 # 20 minutes

# 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.

The same job on a GPU instance

A second job: the same script and the same data, with one change, instance_type, from the ml.m5.large to a GPU instance. A job runs on one instance type; to compare two, you launch two jobs and read both records, which is what the cell at the end of this section does with the earlier job’s desc_cpu and this one’s desc_gpu.

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).

The script does not change

train_nn.py already picks the GPU when there is one and moves the model and the tensors to it, which is why the same file runs on both instance types:

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

Will the GPU job be faster or cheaper?

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); both rates are from the Instances for ML page. 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])

# The two jobs side by side: the ml.m5.large job from earlier (desc_cpu) and this one (desc_gpu); same script, same data, different instance_type
for label, d, rate in [("job on ml.m5.large (CPU)", desc_cpu, prices["ml.m5.large"]), ("job on ml.g4dn.xlarge (GPU)", 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 569 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.6 minutes ($0.005) and the GPU job 3.7 minutes ($0.046): slower, and nine times the price. The GPU job’s wall-clock was 8.6 minutes, because it first sat four and a half minutes in Pending - Training job waiting for capacity; GPU instances are in shorter supply than CPU ones, and the wait is not billed.

A GPU is worth it 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.

More than one machine

When a run is too slow, the moves are in this order, and each one is tried before the next:

A faster single device. The GPU exercise above. On the Instances for ML page that runs from the T4 through the A10G and L4 to the L40S, a fivefold range of speed at a twofold range of price, with no change to the script.
More devices in one box. A multi-GPU instance such as ml.g5.12xlarge or ml.p4d.24xlarge. The GPUs share a link inside the machine, tens to hundreds of GB/s, so splitting a run across them loses little.
More than one instance. The same split across the network between instances, 10 to 400 Gb/s depending on the type, which is ten to a hundred times slower than the link inside a box. This is the last resort, for a run that will not fit one box or would take days on it.

Only the third needs more than a change of instance_type, and on the class account it is the only form available, since the instance list stops at one GPU per instance. The rest of this section shows it on two ml.m5.large CPU instances, where it costs cents, then estimates when it helps and shows the setup for each kind of job.

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. There is nothing to launch to see this; we ran it once while writing the episode, and the record is what the bullets predict: the same 3.5 minutes of wall-clock as the single-instance run, 5.5 billable minutes across the two instances, $0.010, two CloudWatch log streams, algo-1 and algo-2, each printing all 1000 epochs, and two identical saved models. The job worth launching is below.

Data parallelism

To split one training run across two instances instead of running it twice, the script needs three things train_nn.py does 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 three additions 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 single-instance estimator with instance_count=2.

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 say different things: each prints training on 285 of 569 rows (or 284), and only algo-1 prints the epochs and saves the model. That is the proof the work was split, not copied. The run’s final numbers are close to the single-instance run but not identical (our run: validation accuracy 0.818 against 0.804, validation loss 0.96 against 1.05): train_nn.py sets no random seed and train_nn_ddp.py fixes one, so the two runs did not start from the same weights, and a thousand Adam steps on a network that is overfitting by epoch 300 turn any such difference into a visible one. Two runs of train_nn.py differ from each other by as much.

The comparison with the single-instance job, wall-clock and bill:

PYTHON

# Same script, same data, same epochs: one ml.m5.large (desc_cpu) against two with data parallelism (desc_ddp)
for label, d in [("1 x ml.m5.large", desc_cpu), ("2 x ml.m5.large, data parallel", desc_ddp)]:
    n = d["ResourceConfig"]["InstanceCount"]
    wall = (d["TrainingEndTime"] - d["CreationTime"]).total_seconds() / 60
    billable = d["BillableTimeInSeconds"] * n / 60
    print(f"{label:32s} wall-clock {wall:.1f} min, billable {billable:.1f} min on {n} instance(s), ${billable/60*prices['ml.m5.large']:.3f}")

Nothing was gained: the two-instance run took longer and cost twice as much. Each step now ends with a gradient exchange between the two machines, and that exchange has a fixed cost, a network round trip of about a millisecond, plus the gradient’s bytes divided by the link’s bandwidth. What it saves is the forward and backward pass over the 284 rows the other machine took, which for TitanicNet is microseconds. The exchange is a per-step cost that is the same whether the step is tiny or huge; here it is the whole step.

When a second machine helps

helpers.estimate_parallel_speedup puts numbers on that for any model and any box, with the same assumptions as the cost estimator: 6 FLOPs per parameter per token at a quarter of the GPU’s peak, and a gradient of 4 bytes per trainable parameter exchanged once per step over the slowest link in use. It does not overlap the exchange with the computation, which real runs partly do, so its speedup is a floor.

Its inputs differ from the cost estimator’s in one way. The cost estimator takes the whole run, examples × tokens_per_example × epochs, because the bill is the total work. This one takes a single step, tokens_per_step, because the exchange happens once per step and the speedup is the same for every step, so the length of the run does not enter. A step is one batch, so tokens_per_step is the batch size times the tokens per example, the same two numbers as the cost estimator’s with the batch in place of the dataset. The 7B model from the cost section, trained with a batch of 16 documents of 2,000 tokens each, three ways:

PYTHON

# Two single-GPU instances over the network: the class-account form of data parallelism
helpers.estimate_parallel_speedup(
    params_billions=7,
    tokens_per_step=16 * 2000,       # one step: a batch of 16 documents x 2,000 tokens each
    instance_type="ml.g6e.xlarge", instance_count=2,
);

PYTHON

# Four GPUs in one box, linked by PCIe
helpers.estimate_parallel_speedup(
    params_billions=7,
    tokens_per_step=16 * 2000,       # one step: a batch of 16 documents x 2,000 tokens each
    instance_type="ml.g6e.12xlarge",
);

PYTHON

# Eight GPUs in one box, linked by NVLink
helpers.estimate_parallel_speedup(
    params_billions=7,
    tokens_per_step=16 * 2000,       # one step: a batch of 16 documents x 2,000 tokens each
    instance_type="ml.p4d.24xlarge",
);

Read the three together. Two L40S instances over a 20 Gb/s network spend as long exchanging the 28 GB gradient as computing, so two machines finish in the same time as one and bill twice. Four L40S in one box over PCIe get 3.2 of the ideal 4. Eight A100s over NVLink get 7.8 of 8, which is why the big training boxes are built that way.

Data parallelism helps when each step’s computation is large next to the exchange, which means a big model, a big batch and a fast link, in that order of importance. Run the function before renting more than one device, and read the cost line: the speedup is what you buy, the cost ratio is what you pay for it.

Setting it up for each kind of job

A job has one instance type and one count; the count is how many of that type, and every instance gets the same script. What differs is how the work is split, which depends on what is running:

Job	On the estimator	In the script	Data
Your own PyTorch script (this episode)	`distribution={"torch_distributed": {"enabled": True}}` on GPU instances, which starts one process per GPU; `{"mpi": {"enabled": True, "processes_per_host": 1}}` on CPU instances	the three additions in `train_nn_ddp.py`: join the process group, shard the rows, wrap the model in DDP; with a `DataLoader`, a `DistributedSampler` does the sharding	one `.npz` per channel, `FullyReplicated`: every instance downloads it and takes its shard
Hugging Face `Trainer` (the fine-tuning episode)	`distribution={"torch_distributed": {"enabled": True}}`	nothing: `Trainer` detects the processes and does the three things itself; `per_device_train_batch_size` is then per GPU, so the effective batch is that times the device count	as above
A model too large for one GPU’s memory	the same launcher	`fsdp=` in the `TrainingArguments` (or DeepSpeed), which shards the weights, gradients and optimizer state across the GPUs so that a model too big for one fits; each GPU still takes its own batch, so it is data parallelism with sharded state, for memory and not for speed, and it only works inside one box or over a fast network	as above
Built-in XGBoost (the optional challenge in the previous episode)	`instance_count=N`, no `distribution`	nothing: the algorithm trains across instances itself	the CSV split into several objects under one prefix, passed as `TrainingInput(prefix, distribution="ShardedByS3Key")` so each instance downloads only its share
Hyperparameter tuning	`instance_count=1` on the estimator; parallelism is `max_parallel_jobs` on the tuner	nothing	as for a single job

Two things the table does not change. max_run is per job, so a two-instance job with a 2-hour cap bills at most 4 instance-hours. And the account’s rules apply to every instance in the job, so the instance type must be on the list and the tags are the job’s tags.

Key Points

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.
Estimate before you launch: for small models, time a short run and scale it by epochs and rows, plus about three minutes of startup; for large ones, 6 FLOPs per parameter per 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.
GPU or CPU is a measurement: the same script on an ml.g4dn.xlarge was slower and dearer than on an ml.m5.large for this network, because the step is too small to keep a GPU busy; the GPU pays when the model and batch are large.
More than one machine is the last move: a faster single device first, then more GPUs in one box, then more instances. instance_count=2 on its own runs the whole job twice and bills twice; data parallelism needs distribution= on the estimator and, for your own script, three additions: join the process group, shard the rows, wrap the model in DDP.
Parallelism pays when the step’s computation is large next to the gradient exchange: big model, big batch, fast link. helpers.estimate_parallel_speedup gives the speedup and the cost ratio before you rent the devices; on 569 rows of tabular data the measured result was 0.85x the speed at 2.4x the cost.

Content from Hyperparameter Tuning in SageMaker: Neural Network Example

Last updated on 2026-11-03 | 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.
Estimate a tuning job’s cost from max_jobs, max_parallel_jobs and one child job’s minutes with helpers.estimate_tuning_cost, and read back what it cost from its child training jobs.

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

Hyperparameter tuning in SageMaker

The PyTorch episode trained one neural network with a learning rate and an epoch count chosen by hand. Those settings change the result, and choosing them by hand means launching a series of experiments and keeping track of their scores yourself. SageMaker’s hyperparameter tuner does that for us. We say which settings may vary, how a good result is measured, and how many training jobs we are willing to pay for.

Each trial is an ordinary SageMaker training job, launched from the controller notebook on its own instance. The tuner changes the hyperparameters between trials and compares the validation metric they report. That makes the search easier to run, but it does not make it free. Ten trials are ten training jobs on the bill, and what keeps a tuning job cheap is the choice of ranges and the number of models trained. We begin with a single trial to check the configuration before paying for a larger search, the same test-first pattern as the training episodes.

The experiment

We reuse the Titanic network and training data from the PyTorch episode and tune two settings, epochs and learning_rate. One is picked from a short list of values, the other sampled from a continuous range. Two settings are enough to see the mechanics without a new dataset or model. In your own work you may tune more.

There are three pieces to connect. An estimator says how each model is trained, a search space says what may change, and a validation metric tells SageMaker which trials did better. The tuner then bounds the search with the total number of trials, max_jobs, and how many run at once, max_parallel_jobs. Afterwards we read the results, add up the billed compute time, and evaluate the chosen model on the test data.

Set up the project folder

Start in the Titanic project folder, as in the previous episodes:

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

Define the training job

We begin with a PyTorch estimator like the one in the previous episode. It names the training script, the instance, the input paths and the metric to record. When it is handed to the tuner, SageMaker uses it as the template for every trial and substitutes new values for the hyperparameters being searched. The training script does not change.

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: As in every training 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"   # your challenge: WattBot2026, BadgerScribe, BadgerCode, BrainDecoding, PoisonedWell or SnapshotWI; SinkOrSwim is the workshop itself
purpose = "Tuning"        # one of: Storage, Notebook, Training, Tuning, Inference

job_tags = [
    {"Key": "Name", "Value": name},
    {"Key": "Project", "Value": project},
    {"Key": "Purpose", "Value": purpose},
]
helpers.check_launch(job_tags)   # the account's rules, checked here first: AWS's own refusal does not say which rule
print(f"Name: {name}, Project: {project}, Purpose: {purpose}")

Next, we’ll configure the baseline estimator that we plan to do hyperparameter search on. Every price in this episode is from the Instances for ML page, us-east-1 on-demand; the tuning job’s cost is that rate times the sum of the child jobs’ billable minutes.

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), from the Instances for ML page; check it for your region

# The runtime cap, in seconds: SageMaker stops the job and releases the instance when it is reached. Set it a little above the estimate;
# the estimate here is three or four minutes, so 20 minutes leaves room for a slow container pull without letting a hung job bill for hours.
max_run = 20*60 # 20 minutes

# 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
    },
    output_path=f"s3://{bucket_name}/tuning/",   # every child job's model.tar.gz lands here, in your bucket
    sagemaker_session=session,
)

Define the search space

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 search usually needs more jobs to find as good a configuration, and the bill is the number of jobs. 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
}

Configure the tuner

The HyperparameterTuner brings the estimator and the search space together. It also needs an objective. Here it maximizes validation:accuracy, which the training script prints and SageMaker reads out of the log. The strategy chooses which configurations to try, max_jobs caps how many training jobs are launched in total, and max_parallel_jobs says how many run at once. Parallel trials finish sooner, and each one still bills its own instance.

Tuning gains are usually marginal past 20 to 30 configurations, and the search space is where the money goes: a published range for the parameters that matter beats a wide search over all of them.

We first set both limits to one. That single trial checks that the script, the metric extraction, the input files and the permissions all work before we pay for a search.

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
)

Check how the objective metric is recorded

A tuner needs a number from every trial. train_nn.py prints the validation accuracy during training, and the estimator’s metric_definitions regex pulls that value out of the job’s log. The name in the log line, in the metric definition and in the tuner’s objective_metric_name have to agree, or SageMaker has nothing to compare the trials on.

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. Printing every epoch is a measurable share of a step this small. 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.

PY

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.

Launch the first trial

Now we give the tuner the training and validation .npz files the PyTorch episode saved to S3. The two TrainingInput channels map those files to the paths train_nn.py expects. tuner.fit() starts the tuning job. SageMaker launches one training job with one configuration from the search space and records the validation metric it reports. With max_jobs=1 this is a test of the setup, not yet a search.

Estimate first, as before every launch. The training episodes estimated one job. A tuning job is many, and helpers.estimate_tuning_cost does the arithmetic from the same numbers the tuner takes. Its first input is the minutes one child job bills after its startup, which the PyTorch episode’s CPU job measured: about 3 billable minutes, 2 of them startup, so about 1 minute of training, evaluation and model upload. The tuner keeps each instance between child jobs, so startup is paid once per parallel slot and every later job on that slot bills only its own minute. Its inputs are the numbers the tuner takes plus one measurement:

train_minutes_per_job: the billable minutes one child job spends after its startup, on training, evaluation and the model upload. For a small model, time a short run and scale it, as the PyTorch episode did. For a large one, helpers.estimate_training_cost gives it. Here it is about 1 minute.
price_per_hour: the instance’s rate from the Instances for ML page, the same variable the estimator cell above set.
max_jobs: how many child jobs the tuner may launch in total, the same max_jobs you give the tuner.
max_parallel_jobs: how many run at once, the same max_parallel_jobs. This is the number of instances the tuner keeps up, so it is the number of times startup is paid.
startup_minutes: the container pull and data download, 2 by default, from the training episodes’ measured jobs.

It prints the billable minutes with their breakdown, the cost, the wall-clock before the unbilled provisioning, and a max_run for each child job, twice one job’s estimate, to set on the estimator. The whole search, as a dictionary, comes back too.

PYTHON

helpers.estimate_tuning_cost(
    train_minutes_per_job=1,     # a child job's billable minutes after startup, from the PyTorch episode's CPU job
    price_per_hour=price_per_hour,   # ml.m5.large, set with the estimator above
    max_jobs=1,                  # the tuner's max_jobs: one child job, to test the setup
    max_parallel_jobs=1,         # the tuner's max_parallel_jobs: one instance, so startup is paid once
);

One child job, about 3 billable minutes, about $0.01. If the launch is refused with an AccessDeniedException ending in explicit deny, the account’s checks are the cause: a missing tag, an instance type outside the lesson’s list, or a resource kind the lesson never creates; the first training episode’s callout If the launch is refused says what to look at.

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)")

Monitor the tuning job

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)
    if "BillableTimeInSeconds" not in desc:   # the job has not finished: no bill yet
        print(f"{job_name}: status {desc['TrainingJobStatus']}; billable time and the model path appear once it is Completed")
        return desc
    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 and print the total."""
    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:.3f} total")

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. SageMaker’s Bayesian optimization uses the completed trials to choose the next ones, so it spends fewer jobs on unpromising 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. Two at a time is a good default.

Estimate first. The same call with the new limits:

PYTHON

helpers.estimate_tuning_cost(
    train_minutes_per_job=1,     # unchanged: the child jobs are the same job
    price_per_hour=price_per_hour,
    max_jobs=10,                 # ten child jobs
    max_parallel_jobs=2,         # two at a time, so startup is paid twice and the other eight jobs reuse warm instances
);

About 14 billable minutes, $0.03, and about 7 minutes of wall-clock plus provisioning. Startup is paid twice, once per slot, not ten times. Write the numbers down and compare them with the bill below.

PYTHON

import time as t

# 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, the estimate’s 14. Read the per-job lines and the reason the estimate counts startup per slot 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.

More jobs in parallel

Bayesian optimization chooses each job’s values from the results so far. With max_parallel_jobs well below max_jobs, later jobs learn from earlier ones; with all ten at once, nothing has finished when the values are chosen, and the search is a random one. The All at once challenge below measures what that does to the result, the wall-clock and the bill.

Extracting and evaluating the best model after tuning

Tuning takes about 10 minutes. After SageMaker completes the hyperparameter tuning job, the results, including the trained models for each configuration, are stored in an S3 bucket. The steps to locate and evaluate the best model on test data:

Understanding the folder structure:
- Each child job uploads its model to the estimator’s output_path, under its own job name: s3://{bucket}/tuning/{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.

The code:

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 = {k: v for k, v in best_job_desc["HyperParameters"].items() if k in hyperparameter_ranges}   # the record also carries the SDK's own sagemaker_* entries
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 the PyTorch 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.

Preprocessing lives in a function so that the test set is transformed exactly as the training set was. 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 PyTorch 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}")

Challenge

All at once

Rerun the 10-job search with max_parallel_jobs=10 (equal to max_jobs). Before you launch it, run helpers.estimate_tuning_cost with the new limit and 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. The estimator says 30 billable minutes, ten slots’ startup plus ten jobs, against 14 for two slots. 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).

PY

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: 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.
Estimate the whole search before launching it: helpers.estimate_tuning_cost takes the tuner’s own limits and one child job’s billable minutes; startup is paid once per parallel slot, which is why two at a time is cheaper than ten at once.
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-11-04 | Edit this page

Estimated time: 180 minutes

Overview

Questions

When is fine-tuning worth it?
What does a fine-tuning run cost, and how do I estimate that before launching it?
How do I fine-tune a model as a SageMaker training job?
How do I tell whether the fine-tune helped?

Objectives

Explain what drives the cost of a fine-tuning run (model size × tokens × epochs × number of attempts), and what to do before launching one at any size: set a time cap, run a small slice first, write the estimate down, and checkpoint.
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 a model?

In the earlier training episodes, we built relatively small models from scratch. That works when we have a well-defined problem and enough training data, but it isn’t practical for tasks that require the broad capabilities of a modern language or vision model with billions of weights. Training one of those models from scratch takes thousands of GPUs running for months, tens of millions of GPU-hours, and enormous datasets. The PyTorch episode put numbers on that with helpers.estimate_training_cost. One pass of a 7B language model over a million papers came to about $1,270 on a four-GPU box, and the models behind ChatGPT are far larger and trained on far more.

Fortunately, we don’t have to start from scratch. Pretrained models already know how to interpret images, understand language, and generate text. Fine-tuning lets us build on those capabilities by training a model further on examples from our own task.

Consider historical handwriting recognition. A general-purpose vision-language model, such as the open-weight Qwen3.8-27B or a commercial one like Gemini, may be able to transcribe ordinary printed text and modern handwriting, yet struggle with nineteenth-century documents, unfamiliar letterforms, or the conventions used by a particular archive. Giving it better instructions can help, but there are limits to what prompting can accomplish when the model has never learned to recognize the writing itself.

Fine-tuning offers another approach. Show the model examples of the documents alongside their correct transcriptions, and train it to produce better results on similar material.

This can be surprisingly effective. CHURRO is a 3-billion-parameter vision-language model developed for historical text recognition. Its researchers fine-tuned an existing Qwen model on roughly 100,000 historical document pages and reported better transcription accuracy than Gemini 2.5 Pro at substantially lower inference cost (paper). Rather than using an increasingly large general-purpose model, they trained a smaller one to become particularly good at the task they cared about. The same pattern shows up in other fields:

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

That’s the approach we’ll explore in this episode. We’ll take CHURRO, which already specializes in historical documents, and adapt it further to a collection of handwritten German letters from UW-Madison’s Badger Scribe challenge.

How fine-tuning works

Fine-tuning uses the same basic training process we’ve already seen. The model makes predictions on training examples, we measure how far those predictions are from the correct answers, and gradient descent adjusts its parameters to reduce the error.

The difference is where we begin. Rather than initializing the model with random weights, we start with weights learned during an earlier, much larger training process. Our examples don’t need to teach the model everything about language or images. They only need to help it adapt what it already knows.

We could update every parameter in the model, which is called full fine-tuning. That’s how CHURRO itself was created. But even a relatively small model with 3 billion parameters is expensive to train this way, because training requires much more GPU memory than simply running the model.

LoRA adds new small weights beside the frozen model

The alternative to full fine-tuning is parameter-efficient fine-tuning, or PEFT, a family of methods that leave the pretrained weights frozen and train a small number of new parameters instead. Low-rank adaptation, or LoRA, is the most widely used form of PEFT. It keeps the pretrained weights frozen and adds small sets of trainable parameters to selected layers. Rather than directly changing a large weight matrix, LoRA learns a much smaller adjustment to its output.

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.

Nothing in the original model is trained. The trained parameters are the new matrices, and they are a small fraction of the model’s size, under one percent in this episode. We still run the whole model during training, but we only calculate and store weight updates for the new matrices. That saves enough memory to fine-tune CHURRO on a single GPU with 24 GB of memory; the next section says how much.

LoRA also produces a relatively small adapter that we can save separately from the original model. To use the fine-tuned model later, we load the pretrained weights together with the adapter.

This makes LoRA particularly useful for research groups. We can take an existing open-weight model, adapt it to a specialized dataset, evaluate whether the adaptation helped, and keep the resulting model without the expense of full fine-tuning.

Training, full fine-tuning and LoRA on one scale

The GPU-hours put the three on one scale. Training from scratch is at the top, the full fine-tune that made CHURRO in the middle, and this episode’s LoRA run at the bottom.

Horizontal bar chart on a logarithmic GPU-hours axis from 0.1 to 10 million. Training from scratch, every weight from random: 460,000 GPU-hours. Full fine-tuning, every weight from a pretrained model: 800 GPU-hours. LoRA fine-tuning, an adapter under one percent of the weights: 0.5 GPU-hours. Each bar is labeled with the GPUs, the time and the electricity. Energy is GPU electricity at the rated draw, a lower bound.

Why do the three differ so much? Training from scratch and full fine-tuning do the same work per step. Every parameter is updated, which means keeping four things in GPU memory for each one, the weight itself, its gradient, a 32-bit copy of the weight for the update to land in, and the optimizer’s two running averages. That is 16 bytes per parameter, so a 3B model needs about 48 GB before activations and other overhead, which means several GPUs. What separates those two is how much data they go through, trillions of tokens against a hundred thousand pages. LoRA cuts the work in a different place. It updates only the adapter and leaves the model’s own weights frozen, and a frozen weight costs 2 bytes, not 16, with nothing else to store for it. That is why the same model fits one GPU, and for our task it is all we need.

Our task, transcribing historical German letters

For this exercise, we’ll use the Badger Scribe challenge, which asks participants to produce faithful transcriptions of scanned historical documents. The competition includes several kinds of archival material, but we’ll concentrate on letters from 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, transcribed by the Institute’s Kurrent Transcription Group.

These letters were written in Kurrent, a historical form of German handwriting whose letterforms differ substantially from modern cursive. Even a model trained on historical documents can struggle with it, making this a useful test of whether additional fine-tuning helps.

We have 295 labeled pages from this collection. We’ll reserve approximately 15% for evaluation and train on the remaining 251, using the corresponding human-produced transcriptions as our targets.

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

Our goal is straightforward. Does fine-tuning CHURRO on these examples improve its ability to transcribe previously unseen pages?

We’ll measure that using character error rate (CER), the same metric used by the competition. A lower CER means the transcription contains fewer character-level mistakes.

What we’ll do, and what it will cost

We’ll run the training on AWS SageMaker, as a managed training job, with no interactive GPU notebook left running. SageMaker will provision the GPU, run our training script, save the model adapter and evaluation results, and shut down the instance when the job finishes.

The entire run, including evaluation before and after fine-tuning, takes about 33 minutes on an ml.g5.xlarge instance and costs approximately $0.75, based on our measured October 2026 run. That small cost is deliberate. Fine-tuning can become expensive when we move to larger models, longer datasets, or repeated experiments, so we’ll learn how to estimate memory and runtime requirements before launching the job. We’ll also set a time limit to prevent an unexpected training run from continuing indefinitely.

By the end, we’ll have fine-tuned a vision-language model, compared its predictions against the original model, and checked how closely the actual bill matched our estimate.

Three places to run this fine-tune

This episode runs the fine-tune as a SageMaker job. The same model, data and scripts also run on Google Colab’s free GPU and on UW-Madison’s Center for High Throughput Computing, as two standalone versions in the lesson repository’s finetune_extras/ folder. Pick by what you have and what you want to learn.

	AWS SageMaker	Google Colab	CHTC
Who can use it	Anyone with an AWS account and a card or credits	Anyone with a Google account	UW-Madison researchers and students, with a faculty sponsor
Cost of this run	About `$0.75`	Free	Free
GPU	A10G, 24 GB, bf16	T4, 16 GB, no bf16, so the frozen weights go to 4-bit	L40S, A100, H100 or H200, 48 to 141 GB, bf16; NVLink on the A100, H100 and H200 servers
Time for this run	33 minutes, starts in a minute	About an hour and a half, starts at once if a GPU is free	About 20 minutes, after a queue wait of minutes to hours
Largest LoRA fine-tune	No hard limit, only cost. 7B on the lesson’s `$1.41`/h card, 70B at `$13`/h	about 7B, in 4-bit	about 50B in bf16 on an H200, a 70B in 4-bit on one 80 GB card, larger across a server’s cards
Largest full fine-tune	No hard limit, only cost. 3B at `$7`/h, 70B at `$127`/h	none, a 3B needs 52 GB	a 7B on one H200, a 70B across the H100 or H200 servers’ cards
What you learn beyond the fine-tune	Estimating and reading a bill, tags, job guard rails	Fitting a model into a small GPU	Describing a job for a scheduler, containers, staging data

For fine-tuning, CHTC is the first choice for UW-Madison researchers. It is free, its GPUs are large, and the queue is the only price. Colab is for practicing LoRA with nothing to set up. The cloud has the newest and fastest GPUs, and a fine-tune on a large dataset sometimes needs them, but that is where the bills in this episode’s cost table come from, so it is the choice when the other two cannot do the job, not the default.

None of the three runtime limits is a wall. The training script writes a checkpoint at the end of every epoch and resumes from it, so a run longer than Colab’s session, CHTC’s job lane or an AWS max_run is several runs that each pick up where the last one stopped. Each version says where its checkpoints go.

Callout

What the class account allows

Do not fully fine-tune anything in this workshop, 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; a full fine-tune of a 7B model needs more memory than any instance the account allows. The larger runs appear in this episode only as estimates, so that you can recognize one before you launch it.

Step 0: Estimate the cost before you launch anything

Before launching a fine-tune, write down what it should cost. If the bill comes in at three times your number, something is wrong, and you want to find that out on the first run, not the fiftieth. Two questions decide the cost of any training run. Will the model fit on the GPU, and how long will it run? Both can be answered on paper, to within a factor of two, before anything is rented.

How much GPU memory it needs

Every parameter in the model takes up GPU memory, and how much depends on what you are doing with it. Running the model needs only the weight. A full fine-tune needs the weight, its gradient, a 32-bit copy of the weight, and the two running averages Adam keeps for it, the four numbers per weight from the How fine-tuning works section above. A LoRA fine-tune keeps the weights frozen in 16-bit and adds a small adapter that carries its own gradients and optimizer state.

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

What you are doing with the model	Bytes per parameter	What has to be in memory for each parameter
Inference, FP8	1	the weight, quantized
Inference, bf16	2	the weight
LoRA fine-tune	~2.5	the frozen 16-bit weight, plus a share of the adapter with its gradients and optimizer state
Full fine-tune, mixed precision	16	the 16-bit weight (2), its gradient (2), a 32-bit copy of the weight (4), Adam’s two running averages (8)

Multiply the bytes per parameter by the parameter count, add about 4 GB for activations (more for long sequences or big batches), and you have the GPU memory the job needs. The second table does that for four model sizes and names the smallest SageMaker box that holds each one. The inference column uses the precision people run each size at. Models up to about 8B are served in bf16, which is how the Processing Jobs episode runs Qwen3-8B. Above that the weights are quantized, FP8 for the 27B as the GPU-notebook episode does, 4-bit for a 70B on one card. Training cannot be quantized that way, so the two fine-tuning columns are 16-bit throughout:

Model	Inference, as run in practice	LoRA fine-tune	Full fine-tune
3B	bf16, 6 GB: `ml.g5.xlarge`, 1 × 24 GB, `$1.41`/h	12 GB: `ml.g5.xlarge`	52 GB: `ml.g5.12xlarge`, 4 × 24 GB, `$7.09`/h
7B	bf16, 14 GB: `ml.g5.xlarge`	22 GB: `ml.g5.xlarge`, with nothing to spare	116 GB: `ml.g6e.12xlarge`, 4 × 48 GB, `$13.12`/h
27B	FP8, 27 GB: `ml.g6e.xlarge`, 1 × 48 GB, `$2.61`/h	72 GB: `ml.g5.12xlarge`, 4 × 24 GB, `$7.09`/h	436 GB: `ml.p5.48xlarge`, 8 × 80 GB, `$63.30`/h
70B	4-bit, 35 GB: `ml.g6e.xlarge`; FP8, 70 GB: `ml.g6e.12xlarge`, 4 × 48 GB, `$13.12`/h	179 GB: `ml.g6e.12xlarge`	1,124 GB: two `ml.p5.48xlarge`, `$127`/h

A 7B model that runs on one 24 GB GPU needs 116 GB to fully fine-tune, a four-GPU box at $13 an hour with the model sharded across the GPUs. Training cannot run on quantized weights the way inference can, because each step’s update is too small for 8-bit or 4-bit numbers to hold. LoRA avoids that jump by freezing the model and training a few million adapter parameters, which is why this episode fits on one 24 GB card.

How long will it run?

How long a job runs is how much work it does divided by how fast the GPU does it. The work is counted in floating-point operations, FLOPs, one multiply or one add on one number. The GPU’s speed is a rate, floating-point operations per second, written FLOPS with a capital S, and a trillion a second is a TFLOPS, which is how a data sheet quotes a GPU.

The work per token is easy to count, because every parameter does the same small amount of it. Running the model forward is one multiply and one add per parameter, 2 FLOPs. Passing the error back through the model to find what each parameter should have done is another 2. Updating the parameters is another 2. A full fine-tune does all three, 6 FLOPs per parameter per token. LoRA leaves the base weights alone, so it skips the update and does 4.

The three matrix products per token; a full fine-tune does all three, LoRA skips the weight update

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

The data sheet quotes a GPU’s peak speed, about 70 TFLOPS for an A10G, 312 for an A100, 989 for an H100. A training job reaches about a quarter of that, because the GPU spends much of its time waiting for data to arrive from memory. Cost is GPU-hours times the hourly rate, plus the startup minutes every job spends pulling its container and data, which are billed too.

For the class run, 251 pages at about 1,600 tokens each for 2 passes is 803,000 tokens. LoRA on 3.75B parameters is 4 × 3.75 billion × 803,000, about 1.2 × 10¹⁶ FLOPs. A quarter of an A10G’s 70 TFLOPS is 17 trillion a second, 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 same as the estimate.

The same arithmetic as a function

helpers.estimate_finetune_cost does both calculations for any model and any data, picks the cheapest instance that fits (or uses the one you name), and prints the memory needed, the GPU-hours, the cost per run and per project, and a max_run to set. It is the same function the PyTorch episode used as estimate_training_cost, which is this one with every weight trained from scratch, so the assumptions are the same there and here. Bytes per parameter from the first table, 6 or 4 FLOPs per parameter per token, times 1.5 for gradient checkpointing, a quarter of the GPU’s peak, and 10 minutes of startup per job. Its inputs are the quantities in the two rules, with the class run as the example:

params_billions: every parameter the model holds, including parts you are not training. CHURRO is a 3B language model plus a vision tower, 3.75B in all.
examples: one example is one training item, whatever the model is given one at a time. Here it is one page image with its transcription, and there are 251 of them. For a text model it would be one document; for a speech model, one recording.
tokens_per_example: how many tokens the model sees for one of those items. Tokens, not words, because the model never sees words. English text is about 1.3 tokens per word, more for German and for old spellings. An image on a Qwen-VL style model becomes one token per 28 × 28 pixel block after resizing, 14-pixel patches merged two by two, about 1,300 for a megapixel. For a page here that is about 1,300 for the image plus about 300 for the transcription, 1,600; the Setup section counts it with the model’s own processor, so it is not a guess.
epochs: passes over the data. The three numbers multiply to the tokens the model sees, 251 × 1,600 × 2, about 800,000. If you know the total tokens and not the breakdown, put the total in examples and 1 in tokens_per_example; the arithmetic is the same.
method: "lora" or "full", which sets the bytes per parameter and the FLOPs per token.
instance_type: one of the boxes in the price list; leave it out and the cheapest one that fits is chosen.
runs: how many times you expect to run it. A first run is never the last; five to ten is realistic for a project before any sweep.

PYTHON

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

step0_estimate = helpers.estimate_finetune_cost(
    params_billions=3.75,        # CHURRO-3B: the 3B language model plus its vision tower
    examples=251,                # one example = one page image with its transcription
    tokens_per_example=1600,     # ~1,300 image tokens + ~300 text tokens; counted in Setup below
    epochs=2,                    # passes over the 251 pages
    method="lora",
    instance_type="ml.g5.xlarge",
    runs=1,                      # one run today; the printout's last cost line is this times runs
)

The function says about 27 minutes and $0.64, 17 minutes of training plus 10 of startup. The hand calculation said 12 minutes of training, and the measured run took 12. The estimator is high because it adds half again for gradient checkpointing, which our script uses to save memory by recomputing the forward pass during the backward one. A third either way is normal for this rule. The job you launch in Step 1 bills about 32 minutes, because it also scores holdout pages before and after training, about 14 minutes the estimator does not count. Keep step0_estimate. Step 2 puts it next to the job’s own bill.

The same call for a 7B model on 5,000 documents of about 600 words (800 tokens) each, three passes. Two arguments change besides the size. runs goes from 1 to 8, because a project reruns a fine-tune as it fixes data and tries settings, and the printout’s last line prices all eight. instance_type is left out, so the function picks the cheapest instance the model fits on.

PYTHON

helpers.estimate_finetune_cost(
    params_billions=7,           # a 7B model, LoRA on top
    examples=5000,               # 5,000 documents
    tokens_per_example=800,      # about 600 words each
    epochs=3,                    # three passes over the documents
    method="lora",
    runs=8,                      # the eight attempts a project takes; the last cost line is for all of them
);                               # no instance_type: the cheapest that fits is chosen

Batch size is not an input, and does not need to be. A bigger batch keeps the GPU busier, which is the utilization argument (0.25 by default; a well-fed GPU reaches 0.35 or so), and it needs more memory for activations, which is activations_gb.

Which instance to rent

The memory rule names the smallest box that holds a job. These are the single-GPU and multi-GPU SageMaker training instances from the price list; the full table, with CPU instances and notes on each GPU, 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`

A 3B LoRA run fits a 24 GB GPU 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 to 3x faster for less than twice the price, so it becomes cheaper per run once training outweighs the fixed startup.

Multi-GPU instances are for jobs that do not fit on one GPU, a full fine-tune of a 7B model or anything at 70B. They do not make a job cheaper. GPU-hours are the cost, and spreading them across eight GPUs shortens the wall-clock time and leaves the bill alone, and only shortens it when the job is big enough to keep all eight busy. For a job that fits one GPU the price goes up 5x to 45x and the job gets no faster. You add GPUs when the memory table says one GPU cannot hold the job, not to save money.

When a job does need several GPUs, the training script’s Hugging Face Trainer runs data-parallel across every GPU of the instance once the estimator sets distribution={"torch_distributed": {"enabled": True}}; the PyTorch episode’s table of setups per kind of job has the details, including sharding the model with fsdp= when it does not fit one GPU. LoRA splits across GPUs better than a full fine-tune does, because what the GPUs exchange after every step is the gradient, and in a LoRA run that is the adapter alone, under 1% of the model. helpers.estimate_parallel_speedup takes method="lora" for this. As in the PyTorch episode, it takes one step instead of the whole run, tokens_per_step, the batch size times the tokens per example, because the exchange it prices happens once per step. The class run on a four-GPU box, as an estimate only:

PYTHON

# The class run on a 4 × A10G instance: the gradient exchanged is the adapter's, not the model's.
# An estimate, not a launch: a 12-minute job does not need four GPUs, and the box costs 5x the single one.
helpers.estimate_parallel_speedup(
    params_billions=3.75,
    tokens_per_step=4 * 1600,        # one step, not the run: the job trains 1 page per GPU, so 4 pages on 4 GPUs, x 1,600 tokens a page
    instance_type="ml.g5.12xlarge",
    method="lora",
);

Near 4x on 4 GPUs at about the same total cost, which is the case for a multi-GPU box when a LoRA run on a million pages has to finish by a deadline. It is not a reason to rent one for a run that finishes in twelve minutes, and the class account does not allow multi-GPU instances at all.

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

With the two rules, any run can be priced before it exists. The table does that for runs from this episode’s to the ones that make a model like CHURRO, with SageMaker on-demand prices in us-east-1 from October 2026. GPUs come in fixed boxes, so “32 H100s for 25 hours” is four ml.p5.48xlarge for 25 hours, 100 instance-hours at $63.30:

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.8M 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.6M 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, 2M tokens	1 × A10G	1 × `ml.g5.xlarge` (`$1.41`)	0.85	0.85	~`$1.20`
7B full fine-tune, 50,000 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, 100,000 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, repeated in a 10-run hyperparameter sweep					× 10

The class row is measured. The homework and team rows take that run’s seconds per page and apply them to more passes and more pages. The 7B row is worked out in the challenge below, the 70B rows are order-of-magnitude figures, and the CHURRO row comes from a published review of the paper.

Horizontal bar chart on a logarithmic dollar axis: this episode 0.75, homework 1.50, the full Badger Scribe train set 1.20, 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

The CHURRO bar is the model you are about to 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 your task is the biggest saving available, and it is decided before any GPU is involved. The sweep extension is the one that surprises research groups. A first run is rarely the last, so a $140 full fine-tune is usually a $700 to $1,400 project.

The number of passes is a choice. Two passes for the class run is a budget decision for a 33-minute job, not advice. Large instruction-tuning corpora usually get one to three passes, and a few hundred examples usually get three to five. Each extra pass adds the same amount again, so ten passes over the whole train set is about $5. Step 2 shows how to read the per-pass curve and decide.

Challenge

Estimate a run you have not launched

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

Show me the solution

It does not fit on the g5, because a full fine-tune of a 7B model needs about 112 GB plus activations, and the g5 has 24 GB. On the p4d, 8 × 40 GB = 320 GB holds it with the model sharded across the GPUs.

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, 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 estimator 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 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. That is five to ten runs, $700 to $1,400, before anyone touches a hyperparameter sweep. That is the number for the proposal. The cap, the small first run, the written estimate and the checkpoints are the same at this size; the bill for skipping them is not.

What the class run costs, phase by phase

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 job as run on 2026-10-21, 32.5 minutes from launch to finish, 31.7 billed, $0.75. Scoring a page takes 17 to 21 seconds, 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 grows linearly with pages × epochs, and roughly linearly with image resolution, since 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 gets written down before the first run and compared to the BillableTimeInSeconds the job reports 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, with the controller notebook and the exercise:

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.30`

Compare the classic mistake, doing all of this interactively on an 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 ends automatically.

Callout

Before any fine-tuning run

Set max_run on every job, and stop any notebook GPU instance when you stand up.
Make the first launch of a training script a small one, a slice of the data and a short cap. The in-class run below is that small run; the homework scales it up after it has been seen to work.
Write the cost estimate down before launching. Afterwards, read BillableTimeInSeconds from the job description and compare.
Save checkpoints to S3 as the job runs, so a killed 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 on the Instances for ML page).

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; your notebook downloads only the small CSVs, a few page images to look at, and the processor files. 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 sys
sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")   # repeated from Step 0 so that Setup runs on its own after a kernel restart
import helpers

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 and not 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)

What the jobs will read, before downloading any of it: the page lists and metadata at the top level, the scanned pages under images/, and the staged CHURRO-3B weights under model/. The weights are 7.5 GB and never come to this notebook. They go to each job as its model channel, and SageMaker copies them onto the job’s instance, which is the pattern the Accessing S3 episode’s callout gives for any file bigger than the controller.

PYTHON

helpers.show_bucket(shared_bucket, f"{data_prefix}/")   # the shared bucket's folder for this episode: files, and each sub-folder's count and size

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 labeled 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)} labeled 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. The 28 is the model’s own number. CHURRO’s Qwen vision encoder cuts the image into 14 × 14 pixel patches and merges each 2 × 2 group of patches into one token for the language model, so one token covers 28 × 28 pixels, and both numbers are in the processor’s config as patch_size and merge_size. The same method works for any model. Run one example through the model’s own processor and read the length of the input_ids it returns, which is the count whatever the model does inside. What differs is how the image tokens arise, and the model card and the processor config say which. Qwen’s models tile the image at one token per 28 × 28 pixels, so a bigger page is more tokens. Gemma 3 and PaliGemma resize every image to one fixed size and spend 256 tokens on it however large the page. LLaVA spends 576. The transcription becomes pieces of words, more pieces for German and for old spellings than for modern English. The count depends on the model, so it has to come from this model’s own processor: every model has its own tokenizer vocabulary, and every vision model its own patch size and resizing rules, so the same page is a different number of tokens for a different model. CHURRO’s processor is the code the training job will run on every page, so a count from it is the job’s count, not an estimate. Its 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 the way 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 the way 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())   # Qwen-specific breakdown: 14-px patches merged 2x2, one token per 28x28 block. n_total above is the count for any model

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}")

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. Step 0 of this episode ran helpers.estimate_finetune_cost with a guess of 1,600 tokens per page; run it again with the measured average and compare:

Challenge

Count the page for another model

Pick another open vision-language model and count what the same page costs there. llava-hf/llava-1.5-7b-hf is not gated and its processor downloads in seconds. Before running anything, read its model card and predict two things: how many image tokens a page will take, and whether max_pixels will change that. Then load its processor, build the same message with its chat template, the image, the instruction and the transcription, and read the input_ids length. Run the text without the image to separate the two.

Show me the solution

PYTHON

from transformers import AutoProcessor
llava = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt}]},
            {"role": "assistant", "content": [{"type": "text", "text": row.text}]}]
text = llava.apply_chat_template(messages, tokenize=False)
n_llava = llava(text=[text], images=[image], return_tensors="pt")["input_ids"].shape[1]
n_llava_text = len(llava.tokenizer(text.replace("<image>", ""))["input_ids"])
print(f"LLaVA 1.5: {n_llava} tokens for the page, {n_llava - n_llava_text} of them for the image; CHURRO: {n_total}")

LLaVA 1.5 resizes every image to 336 × 336 pixels and cuts it into 24 × 24 patches of 14 pixels, 576 image tokens however large the page, so max_pixels changes nothing for it. The page comes to about 900 tokens, 576 for the image and the rest for the instruction and the German text, which its Llama tokenizer splits into a different number of pieces than Qwen’s. Fewer tokens is a cheaper step. It is not a better model for this task. A 336-pixel image of a letter page is too small to read the handwriting, which is why CHURRO spends 1,280 tokens on a megapixel. The token count is cost. Accuracy has to be measured.

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");

The measured count is a tenth below the guess, and the estimate moves with it: 26 minutes and $0.60 against 27 and $0.64. Step 1 keeps the 27-minute figure; a tenth either way is well inside what the rule can tell, which is about a factor of two. 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.

Now the same run as a full fine-tune. The arithmetic goes up by half, 6 FLOPs per parameter per token instead of 4, but that is not what changes the bill:

PYTHON

helpers.estimate_finetune_cost(params_billions=3.75, examples=251, tokens_per_example=int(avg_tokens_per_page), epochs=2,
                               method="full");   # no instance_type: the cheapest that holds it

Memory is what changes. The weights, their gradients and the optimizer’s state come to 64 GB against 13, so the run no longer fits a 24 GB GPU, and the cheapest instance that holds it is the four-GPU ml.g5.12xlarge at $7.09 an hour, five times the price. The four GPUs share the work, so the job is shorter, and the run comes to about $1.87 against $0.60. On 251 pages that is a dollar; the same jump at 7B is the Step 0 table’s $140 a run, and the two warnings the function prints are the account’s rules: no full fine-tune, no multi-GPU instance.

Two more sizes on the same 251 pages, each both ways. Nothing is launched; read the memory line and the instance it forces.

PYTHON

helpers.estimate_finetune_cost(params_billions=7, examples=251, tokens_per_example=int(avg_tokens_per_page), epochs=2, method="lora");

PYTHON

helpers.estimate_finetune_cost(params_billions=7, examples=251, tokens_per_example=int(avg_tokens_per_page), epochs=2, method="full");

PYTHON

helpers.estimate_finetune_cost(params_billions=70, examples=251, tokens_per_example=int(avg_tokens_per_page), epochs=2, method="lora");

PYTHON

# instance_type named on purpose: without it the function finds nothing on the list that holds the job and raises
helpers.estimate_finetune_cost(params_billions=70, examples=251, tokens_per_example=int(avg_tokens_per_page), epochs=2, method="full",
                               instance_type="ml.p5.48xlarge");   # the largest box on the list: 8 x H100, 640 GB

A 7B LoRA still fits one GPU, 22 GB on a 24 GB L4, for about $0.50; the full fine-tune of the same model needs 116 GB, a four-GPU box at $13 an hour, and the account refuses it. At 70B the LoRA needs 179 GB, so even the adapter run is a four-GPU job the account refuses, and the full fine-tune needs 1,124 GB: nothing on the list holds it, and the function says what the largest box would do, start, bill ten minutes of startup at $63 an hour, and die out of memory. That run is two or more such boxes with the model sharded across them, which is the $12,000 row in the cost table once the data is a real corpus and not 251 pages. On 251 pages the arithmetic is minutes everywhere, so these bills are mostly each box’s price for ten minutes of startup; what the four runs show is the box a model forces you onto, which is the first thing to find out about a fine-tune, before the data or the epochs.

A code folder for the job

The three scripts the job needs are already on this instance, in the lesson clone. The job cannot use them from there: it runs in a container on another machine, and the only code it gets is the one directory named by source_dir= on the estimator, uploaded whole. That directory should hold what the job imports and a requirements.txt listing what the container has to install, and nothing else. So make a small code/ folder here with copies of the three files, instead of pointing source_dir at the lesson’s scripts/ directory, which would upload every script in the lesson and need a requirements file written into the repository:

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"]))   # peft is Hugging Face's library of PEFT methods; it is where get_peft_model and LoraConfig come from

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 the same 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.

Challenge

Read the script before you pay for it

Open code/finetune_churro_lora.py in the Jupyter file browser. It is about 400 lines and you are about to launch it on a GPU. Find the answers in the file:

Which layers get LoRA weights, and which part of the model gets none? What would the regular expression in target_modules also match without its (?!.*visual) part?
Roughly how many parameters will train? Work it out from r=16 and the shapes of the language model’s projections, which the cell before the hyperparameters read from config.json (36 layers; q_proj and o_proj are 2048 × 2048, k_proj and v_proj are 2048 × 256, and gate_proj, up_proj and down_proj are 2048 × 11008), then compare with the trainable params: line in the job’s log once it runs.
Which tokens does the loss count? Find every place labels are set to -100.
The script stops itself in two situations before training a single step. Which, and why are they there?

Show me the solution

Only the language model’s q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj and down_proj. Every module in the vision tower has visual in its name, and the negative lookahead excludes them. Without it, the vision blocks’ own gate_proj, up_proj and down_proj would match too, and the model would be relearning how to see from 251 pages.
A LoRA adapter on a matrix of shape in × out adds r × (in + out) parameters. Per layer: 16 × 4096 × 2 for q_proj and o_proj, 16 × 2304 × 2 for k_proj and v_proj, 16 × 13056 × 3 for the MLP, about 831,000; times 36 layers is about 30 million, under 1% of the 3.75 billion. The log line trainable params: ... || all params: ... || trainable%: ... should agree. If it said 100%, nothing was frozen and the run is a full fine-tune at a full fine-tune’s memory and price.
The transcription only. labels[enc["attention_mask"] == 0] = -100 drops the padding, labels[i, :n_prompt] = -100 drops the prompt and the image tokens, and -100 is the value PyTorch’s cross-entropy ignores. The model is graded on what it writes, not on repeating the question.
Right after the adapter is attached: it raises if any trainable tensor is in the vision tower, or if more than 5% of the model is trainable. Either means the job is not the LoRA run that was estimated, so it stops before the GPU-hours are spent; a job that raises there has cost its startup minutes and nothing more.

PYTHON

name = "DoeJohn"         # replace with your name
project = "SinkOrSwim"   # your challenge: WattBot2026, BadgerScribe, BadgerCode, BrainDecoding, PoisonedWell or SnapshotWI; SinkOrSwim is the workshop itself

def make_job_tags(purpose: str):   # one of: Storage, Notebook, Training, Tuning, Inference
    return helpers.check_launch([   # checked against the account's rules first, with a message that names the problem
        {"Key": "Name", "Value": name},
        {"Key": "Project", "Value": project},
        {"Key": "Purpose", "Value": purpose},
    ])

make_job_tags("Training");   # fails here, with the fix, and not at the launch; the semicolon keeps Jupyter from also displaying the list

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 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, from the Instances for ML page
estimated_minutes = 27                    # 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.53 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 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 one epoch and losing the whole run: $14 on the g5, $250 on a p4d.
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 2 --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.

One of them, lora_r, the adapter rank, decides how many parameters train, and that number is worth knowing before the launch for two reasons. It is what the script’s safeguard checks, since a run that trains 5% of the model instead of 1% is a different, more expensive job, and it is the trainable params: line the job prints, so working it out now means you can tell whether the job is doing what you asked.

The arithmetic is short. LoRA attaches to seven weight matrices in every layer of the language model, the four attention projections (q_proj, k_proj, v_proj, o_proj) and the three in the MLP (gate_proj, up_proj, down_proj). On a matrix of shape in × out, an adapter of rank r adds r × (in + out) parameters, the two thin matrices from the LoRA figure. Add that up over the seven matrices, multiply by the number of layers, and you have the count. The widths come from the model’s config.json. For any model on Hugging Face that file is one click away, under Files and versions on the model page (CHURRO-3B, for this one), and the same page’s safetensors viewer lists every weight tensor with its shape. The model card’s prose is a quicker glance but not always complete, so read the file when a number matters. Here the file is staged next to the weights in the shared bucket.

PYTHON

s3_client.download_file(shared_bucket, f"{data_prefix}/model/churro-3B/config.json", "config.json")   # the model's architecture file
with open("config.json") as f:
    cfg = json.load(f)
text = cfg.get("text_config", cfg)            # newer configs keep the language model's sizes under text_config, older ones at the top level

layers = text["num_hidden_layers"]            # how many times the seven matrices repeat
hidden = text["hidden_size"]                  # the width of the model's main vector
inter = text["intermediate_size"]             # the MLP's wider middle layer
kv = hidden * text["num_key_value_heads"] // text["num_attention_heads"]   # k_proj and v_proj are narrower: several attention heads share one key and value

projections = {                               # (in, out) of each matrix LoRA attaches to, in one layer
    "q_proj": (hidden, hidden), "k_proj": (hidden, kv), "v_proj": (hidden, kv), "o_proj": (hidden, hidden),
    "gate_proj": (hidden, inter), "up_proj": (hidden, inter), "down_proj": (inter, hidden),
}
print(f"{layers} layers; " + ", ".join(f"{name} {i} x {o}" for name, (i, o) in projections.items()))

total_params = 3.75e9                         # the whole model, vision tower included
for lora_r in (8, 16, 32):
    per_layer = sum(lora_r * (i + o) for i, o in projections.values())   # r x (in + out) for each of the seven matrices
    trainable = per_layer * layers
    print(f"rank {lora_r:>2}: {trainable / 1e6:5.1f}M trainable parameters, {100 * trainable / total_params:.2f}% of the model")

Rank 16, the lesson’s choice, trains about 30 million parameters, under 1% of the model, and doubling the rank doubles it. The adapter’s gradients and optimizer state, 16 bytes for each of those 30 million, come to about half a gigabyte, which is why a LoRA run fits beside the frozen 7.5 GB of weights on a 24 GB card. The job’s log confirms the count in its trainable params: line, and the script refuses to train if more than 5% of the model turns out to be trainable.

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: about 30 million trainable parameters, from the cell above; 32 would double that
    "batch_size": 1,
    "grad_accum": 4,                  # effective batch of 4 pages -> ~125 optimizer steps over 2 passes
    "max_pixels": 1280 * 28 * 28,     # the page is resized to at most this many pixels, about a megapixel; every 28 x 28 block is one token, so about 1,280 image tokens a page, as Setup measured. The setting that moves cost most
    "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
}

What one training step does. The Trainer takes one page at a time, since batch_size is 1. It runs the page image, the instruction prompt and the human transcription through the model, and measures the loss on the transcription tokens only, since the collator masks the prompt and image tokens out of it. The gradient of that loss flows back through the frozen model into the adapter matrices, the only parameters with gradients. After grad_accum pages, four here, the summed gradient makes one AdamW update to the adapter, so a pass over 251 pages is about 63 updates and two passes about 126. The learning rate ramps up over the first 3% of updates and then falls along a cosine curve to zero at the last one. Every update costs a forward and a backward pass over about 1,500 tokens.

Every setting, and what it moves. The dictionary sets the ones the lesson chose. The others keep the script’s defaults, which the script’s --help lists.

Setting	Value here	What it does	What changing it costs or buys
`epochs`	2	Passes over the training pages. Training time scales with it, and the holdout curve at the end of each pass says whether the next one helped.	A third pass is 50% more training time. The homework runs four.
`lr`	2e-4	The size of each update to the adapter. LoRA uses about ten times a full fine-tune’s rate (~1e-5), because the adapter starts at zero and has few parameters.	Too high and the loss climbs or jumps around. Too low and two passes do not get far.
`lora_r`	16	The adapter rank, the inner width of the two thin matrices. The trainable parameter count scales with it, about 30 million at 16.	32 doubles the trainable parameters and the adapter’s memory, and on a few hundred pages tends to overfit sooner.
`lora_alpha`	32 (default)	Scales the adapter’s output by alpha ÷ r, 2 here. The script’s default, not set in the run.	Keep alpha at twice r when changing r, or the effective learning rate of the adapter changes with it.
`lora_dropout`	0.05 (default)	Dropout on the adapter’s input, a little regularization against a few hundred pages.	Rarely worth changing.
`batch_size`	1	Pages per forward pass per GPU. A page is about 1,500 tokens with its image, so one page at a time is what fits beside the model on a 24 GB card.	2 needs roughly twice the activation memory and does not fit the lesson’s card.
`grad_accum`	4	Pages whose gradients are summed before one update. The effective batch is `batch_size` × `grad_accum`.	8 halves the number of updates per pass, each a little smoother. The script’s default is 8; the lesson uses 4 so that 251 pages give enough updates in two passes.
`max_pixels`	1,003,520 (1280 × 28 × 28)	The page is resized to at most this many pixels, and every 28 × 28 block is one token, so about 1,280 image tokens a page.	Halving it halves the image tokens and the training time, and loses detail on small handwriting. Doubling it doubles both.
`eval_pages`	15	Holdout pages transcribed and scored, before training when `eval_before` is 1 and again after.	Scoring is 20 to 70 seconds a page, so 15 pages is 5 to 15 minutes each time. 40 pages, the homework’s choice, gives a steadier number at three times the scoring time.
`eval_before`	1	Score the base model on the same pages before training, so the before-and-after comparison is on identical pages.	0 skips it and saves the time. Then there is nothing to compare against.
`curve_pages`	8	Of the 15 holdout pages, the ones scored again at the end of every pass, which draws the holdout CER curve.	0 skips the curve. More pages make the curve steadier and each pass longer.
`holdout_frac`	0.15 (default)	The fraction of pages, within each category, held out of training with a fixed seed. 44 of 295 here.	Changing it changes which pages are scored, so runs with different values are not comparable.
`seed`	42 (default)	Fixes the holdout split, the evaluation sample and the training order, so every run scores the same pages.
`max_new_tokens`	1024 (default)	The most tokens the model may generate when transcribing a page during scoring.	A page is about 200 text tokens. The cap matters only when the model rambles, which the base model does on pages it cannot read.

One check of all of it before the estimator exists, as check_launch checked the tags. check_finetune_launch applies the account’s rules to the tags, the instance type and the cap, and then this lesson’s own rules for a fine-tune: LoRA and not a full fine-tune, a model of at most 7B, a cap between the estimate and four times it, checkpoints on any run that may pass an hour, a fixed job name, and sane values for the passes and the pixel budget. It raises with every problem named, or prints what the launch is allowed to cost.

PYTHON

helpers.check_finetune_launch(
    tags=make_job_tags("Training"), instance_type=instance_type, max_run=max_run_minutes * 60,
    estimated_minutes=estimated_minutes, price_per_hour=g5_price,
    params_billions=3.75, method="lora",
    job_name=training_job_name, checkpoint_s3_uri=checkpoint_s3_uri, hyperparameters=hyperparameters,
);

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.
If the launch is refused with an AccessDeniedException ending in explicit deny, the account’s checks are the cause: a missing tag, an instance type outside the lesson’s list, or a resource kind the lesson never creates; the first training episode’s callout If the launch is refused says what to look at.
wait=True (the default) blocks the notebook until the job ends and streams its log here. For a 33-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, open the Training jobs page in the console (SageMaker AI → Training → Training jobs), select it and click Stop, or run:

PYTHON

# Only if you need to stop the job: remove the # from the next line and run the cell
# 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 no 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: ...: how many parameters LoRA trains, about 30 million in the adapter, against the 3.75B it leaves frozen. The safeguards: line after it means the script’s two checks passed.
[eval before] 15 holdout pages in ...s (... s/page) -> macro CER ...: the base model’s score, and the time per page one at a time; the homework’s inference jobs batch pages and are faster per page.
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

BillableTimeInSeconds is the job’s own meter and it is final the moment the job ends; Cost Explorer shows the same charge a day or two later, so the comparison is done here, from the job, and not from the billing console.

PYTHON

billed_min = desc["BillableTimeInSeconds"] / 60
estimated_min = step0_estimate["wall_clock_hours"] * 60
print(f"Step 0 estimate: {estimated_min:.0f} min, ${step0_estimate['cost_per_run']:.2f}   (training plus 10 min of startup)")
print(f"Billed:          {billed_min:.0f} min, ${billed_min / 60 * g5_price:.2f}")
print(f"Ratio billed / estimate: {billed_min / estimated_min:.2f}; the scoring passes, which the estimator does not model, are about 14 of the billed minutes")

Within a factor of two is the estimator’s claim, and the bill should land inside it. If it does not, the summary file downloaded next says which phase ran long.

Compare the billable minutes with train_minutes in the summary you download next. In our run the gap was 15 minutes: nine of them the two scorings of 15 holdout pages before and after training, and the rest the fixed overhead of every job, container pull, data download and the 7.5 GB model copy from S3. That overhead 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 "", not 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}")

Our run went from a mean CER of 0.267 to 0.065 on the 15 pages: a quarter of the characters wrong before, one in fifteen after, for 75 cents of GPU. 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 memorizing 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 memorizes.
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.

Challenge

Was the estimate right?

Compare the billable time from job_cost_report with the 27 minutes estimated in Step 0 and the 32 measured in its phase table. 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 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 not within 50%, find out which line was wrong before launching anything else.

Homework: a four-pass LoRA run and (optionally) predictions on the test pages

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.

Callout

The homework cells are switched off

Every cell below that launches a job or needs a job’s output starts with the line %%script false --no-raise-error, which makes Jupyter skip the cell, so Run All in class does not start two hours of GPU work. When you do the homework, delete that first line from a cell and run it. Do the cells in order: each one needs the one before it to have finished.

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 because you will not be watching: the cell returns immediately, and a later cell checks on the job. 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,
}

The same check as in class, with the homework’s cap and name:

PYTHON

helpers.check_finetune_launch(
    tags=make_job_tags("Training"), instance_type="ml.g5.xlarge", max_run=130 * 60,
    estimated_minutes=65, price_per_hour=g5_price, params_billions=3.75, method="lora",
    job_name=homework_job_name, checkpoint_s3_uri=homework_checkpoint_s3_uri, hyperparameters=homework_hyperparameters,
);

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,
)

With wait=False, the launch cell returns as soon as SageMaker accepts the job. The job runs on its own instance for about an hour, so stop this notebook instance if you like and come back. The cell after it asks SageMaker for the job’s status: run it when you are back, and once the status is Completed it reads the bill and the location of the adapter the job saved.

PYTHON

%%script false --no-raise-error
# homework cell, switched off by the line above: delete that line to run it
# wait=False: the cell returns as soon as the job is accepted; the job runs regardless of this notebook
homework_estimator.fit({"train": train_s3_uri, "model": model_s3_uri}, job_name=homework_job_name, wait=False)
print("launched:", homework_job_name)

PYTHON

%%script false --no-raise-error
# homework cell, switched off by the line above: delete that line to run it
# Run this when you are back; it prints the job's status. Everything below needs the adapter, so wait for Completed.
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=60 * 60,   # the processor's name for max_run: 60 minutes, in seconds; the estimate is 15 minutes batched, 55 if it ran one page at a time
    )
    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

%%script false --no-raise-error
# homework cell, switched off by the line above: delete that line to run it
assert homework_adapter_s3_uri, "the four-pass run has not finished; rerun the status 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 on the Processing jobs page (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

%%script false --no-raise-error
# homework cell, switched off by the line above: delete that line to run it
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

%%script false --no-raise-error
# homework cell, switched off by the line above: delete that line to run it
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)

Write the test-page predictions

The lora prediction file already contains every one of the 141 test pages. Keep only those rows, in the order test.csv lists them, and you have the fine-tuned model’s transcriptions of pages nobody has labeled.

That file is not a Kaggle submission. The lesson’s data is the German-letter part of Badger Scribe. The competition’s training set has 619 pages in three categories, the 295 letters plus 288 pages of surveyors’ field notes and 36 of 1830s account books, and its test set has 208 pages, the 141 letters plus 31 and 36 of the others. A submission needs a row for all 208. To get a leaderboard score, download the competition data from Kaggle, run the same fine-tune on its train.csv and images, which the script takes unchanged since its --categories default is all of them, run the transcriber over its test.csv, and build the file from its sample_submission.csv. This step is optional, and only it needs a Kaggle account.

PYTHON

%%script false --no-raise-error
# homework cell, switched off by the line above: delete that line to run it
test_preds = test_df[["page_id"]].merge(preds["lora"], on="page_id", how="left").fillna({"text": ""})
assert len(test_preds) == len(test_df) and test_preds.text.str.len().gt(0).all(), "every test page needs a prediction"

test_preds_path = os.path.join(local_data_dir, "test_predictions_churro_lora.csv")
test_preds.to_csv(test_preds_path, index=False)
s3_client.upload_file(test_preds_path, bucket_name, f"{data_prefix}/predictions/test_predictions_churro_lora.csv")
print("wrote", test_preds_path, len(test_preds), "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, a 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-{run_id}, with a fresh run_id. 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 on the Instances for ML page, 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 phase 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 about two-thirds of the slope. The g6e only loses on a job so short that the fixed overhead dominates, a 5-minute test run, say.

Compare the cost per run, not the hourly rate: 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, not 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 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-11-03 | Edit this page

Estimated time: 15 minutes

Overview

Questions

What is Retrieval-Augmented Generation (RAG)?
What are the four ways to run inference on AWS, and which does a RAG pipeline need?
Why is Bedrock the default, and when is SageMaker the right choice instead?
How do the billing models differ, and what does it cost to forget to shut something down?

Objectives

Explain that RAG trains nothing, so running one does not need a GPU of your own.
Name the four ways to run inference on AWS and what each one bills for.
Say why pay-per-token Bedrock is the default for most research RAG work.
Say when self-managed GPU compute (Processing jobs, notebook GPUs, endpoints) is worth it.
Weigh the cost and control trade-offs that decide which to use.

Retrieval-Augmented Generation (RAG) on AWS

Retrieval-Augmented Generation (RAG) is a pattern where you retrieve relevant passages from your own documents and then generate an answer from them. The model is used, not changed. No weights are updated and no training job runs, so the documents can change every week and nothing has to be retrained. Using a trained model to get an output from it, a prediction, an embedding or an answer, is called inference, and everything in this part of the lesson is inference. The example throughout is WattBot, a question-answering challenge over a corpus of research papers on AI and energy.

There are four ways to run inference on AWS. Call a model through Bedrock’s managed API, run it in a Processing job that ends on its own, load it on a GPU notebook instance, or deploy it as an inference endpoint. This episode compares the four, says why Bedrock is the one to start with, and sets up the episodes that follow, which build the same RAG pipeline in three of these ways and weigh the trade-offs of each.

What is RAG?

Diagram of the RAG pipeline with the embedder and generator highlighted as the two model calls

A language model cannot reliably answer questions about documents it has never seen. RAG connects the model to a collection of material you control. You break the documents into passages, turn each passage into a vector with an embedding model, and keep the vectors so you can search them later. When someone asks a question, you embed the question with the same model and find the passages whose vectors are most similar. Those passages go to a generation model together with the question, and it writes an answer grounded in what was retrieved, citing the passages that support it. When the evidence is not there, a well-built pipeline says so instead of guessing.

So there are two model calls, and neither is training. The embedding model is called once per chunk when the corpus is indexed and once per question. The generation model is called once per question. Both are inference with no state between calls, so neither needs a GPU of your own. The rest of the pipeline, preparing documents, searching vectors, assembling prompts and scoring answers, is ordinary Python. That code stays the same while the place the models run changes, which is what the next three episodes show.

Bedrock is the default for inference

For most RAG pipelines, and for inference in general, Amazon Bedrock is the fastest and most cost-effective starting point on AWS. Bedrock serves embedding and generation models (Amazon Titan, Anthropic Claude, DeepSeek, Meta Llama and others) through pay-per-token APIs. Your retrieval code, chunking, indexing, similarity search, evaluation, stays yours and runs on a cheap CPU notebook or a laptop. Only the two model calls go to Bedrock.

The difference is in how each one bills.

	Bedrock (managed API)	SageMaker GPU instance / endpoint
Billing unit	Per token processed	Per second while the instance runs, at an hourly rate
Cost while idle	`$0`: nothing runs between calls	The full hourly rate, around the clock, used or not
What you have to remember	The per-token price before a big batch; route calls through a tagged profile	Stop or delete the instance, every time
Model updates	AWS and the provider	You: containers, CUDA, GPU memory
Model choice	A curated catalog, proprietary models included	Any open-weight model you can load

Self-hosting on SageMaker gives you control over the weights, the libraries and the hardware, which matters for a specific open-weight model or your own fine-tune. The trade is that a rented GPU bills while the instance runs, whether or not it is generating answers, and research work is mostly occasional experiments, not a steady stream of requests. A GPU instance is a rental whose meter runs until you turn it off, and forgetting to is a common way a research team spends its cloud budget on nothing. Bedrock has no such failure. Stop calling it and you stop paying.

What forgetting costs

A GPU notebook instance left running after you close your laptop (on-demand prices at the time of writing; the instances page has the current ones):

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,906`
ml.p4d.24xlarge	8 × A100 (40 GB)	`$25.25`	~`$303`	~`$1,616`	~`$18,430`

Against that, the whole WattBot pipeline of this lesson on Bedrock, with DeepSeek V3 for generation (the Bedrock pricing page has current per-token prices) and Titan Text Embeddings V2 ($0.02 per million tokens):

Step	Rough volume	Cost
Embed ~9,000 chunks (~2.4M tokens)	2.4M input tokens	~`$0.05`
Answer and explain the 220 answerable training questions (two calls each, ~1.5K-token prompts) with DeepSeek V3	650K in / 40K out	~`$0.50`
Total		about `$0.55`

One 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.

Callout

Where Bedrock falls short

Model choice. You can only call models in the catalog. If the work depends on a specific open-weight model, your own fine-tune, or changes to the model’s internals, you host it yourself.
Very large batches. Embedding tens of millions of chunks with a small open model can be cheaper per unit on a short-lived GPU job than per token. Do the arithmetic before assuming either way.
Access and regions. Some models have to be enabled in the console first, and the catalog differs by region. Check before a workshop or a deadline.
Quotas. On-demand throughput is rate-limited per model and per account. For a large batch, use the batch inference API or request a quota increase instead of retrying against the limit.
Data handling. Bedrock does not train on your prompts and keeps data in the region, or within the US when a us. cross-region profile is used; your institution’s policy on sending documents to a hosted API still applies.
Cost tracking takes one extra step. On-demand calls carry no tags. To attribute spend to a project you route calls through a tagged application inference profile (the Tagging section below, and the Bedrock episode). Set up once; skip it and the usage is anonymous on the bill.

Four ways to run inference on AWS

The four differ less in what the model computes than in how long the hardware stays rented. Here they are in the order to try them.

1. Amazon Bedrock, the default

Call the embedding and generation models from your pipeline. You keep the retrieval logic; AWS keeps the models. You pay per token, nothing runs between calls, and there is nothing to shut down. Bedrock is also where proprietary models such as Claude are available without a separate contract.

Use it for: any RAG system where a catalog model will do, from prototypes and course projects to research tools and production chatbots.

2. SageMaker Processing jobs

For a large corpus, or for an open-weight model Bedrock does not carry, run the embedding and the batch generation as jobs instead of against a live model. A Hugging Face Processing job starts a GPU instance, loads the model, works through all the chunks or questions, writes the results to S3 and ends. The GPU exists only while the job runs, so you get a self-managed model without the risk of forgetting to stop it. This is the safest way to use a GPU on SageMaker.

The cost is latency. A job takes several minutes to start, so this is for work done once and used many times, not for answering one question at a time.

Use it for: embedding a large corpus with an open model, periodic batch regeneration, or research that needs a model Bedrock does not offer.

3. A GPU notebook instance

For models up to about 70B you can choose a GPU instance (an ml.g5.xlarge, say), load the embedding and generation models in the notebook, and run RAG end to end there. It is the simplest setup to follow, and the best way to learn what happens inside a RAG pipeline, because every step is a line of Python you can inspect.

It is also the most expensive way to run RAG and the easiest to get wrong. The GPU bills by the hour from the moment the instance starts until you stop it, through every stretch spent downloading PDFs, chunking text, reading results or at lunch. It is the route most often forgotten, because a notebook instance looks like a workspace, not a running job. If you use it:

Treat it as a learning exercise, not a way to deploy anything.
Stop the instance the moment you are done, and check its status on the Notebook instances page.
Set an AWS Budget alert, and consider a lifecycle configuration that stops idle notebooks.
For larger models the hourly rate ($25 an hour for a p4d.24xlarge; the Instances for ML page has the list) becomes the limit quickly.

Use it for: walking through RAG step by step, debugging a model before packaging it as a job, or short experiments where you stay at the keyboard.

4. SageMaker inference endpoints

For an application that needs interactive, low-latency answers (an API, a chatbot, a dashboard) from a model you host yourself, deploy the embedding and generation models as SageMaker inference endpoints and call them from your retrieval service. You control the model, the scaling and the hardware.

It is also the most expensive option when traffic is low or bursty, because you pay for capacity that stays up when nobody is asking, the same always-on problem as a notebook GPU, now by design and without end. Before choosing it, ask whether Bedrock already serves a model that will do; it gives the same interactive latency with no idle cost. Endpoints are the answer only when the model itself is why Bedrock will not work.

Use it for: production systems with steady traffic that need a custom or open-weight model. The lesson does not build one.

Bedrock charges per call, a Processing job charges for one finite run, and a notebook or an endpoint charges for as long as it is left running. Those lifecycles, not the RAG algorithm, decide which to use.

Every billable resource needs its own tags

Cost tracking in this workshop relies on three tags, Name, Project and Purpose, so that Cost Explorer can split the shared bill by challenge. Tags never propagate. Tagging your notebook instance tags the notebook instance and nothing else. Every other resource you create is a separate line on the bill, anonymous unless you tag it when you create it.

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	Per second while the job runs, at an hourly rate	`tags=` argument on the estimator or tuner, every job
Processing job	Per second while the job runs, at an hourly rate	`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:

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, as the training, tuning and Processing-job episodes do. Vary Purpose per job so the bill says 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 wrapper around one model 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. Bedrock’s tag format is lowercase key/value where SageMaker’s is Key/Value; the tag names are the same.
Only the three approved keys count. Name, Project and Purpose, with that spelling and capitalization; the bill ignores any other key, for the reason the S3 episode gives.

Which to use

Your situation	Route
You need to embed and query a document collection and a catalog model (Titan, Claude, Llama, …) will do	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, stopped when you leave
You need interactive serving of a model Bedrock does not offer, with steady traffic	Inference endpoint

If in doubt, Bedrock. Its worst case is a somewhat higher per-token bill; the worst case on the other routes is a four-figure invoice next month.

Callout

The order of the hands-on episodes

The next two episodes build RAG on Bedrock first, then on Processing jobs. The Bedrock episode is the one to start from. It builds the whole WattBot pipeline, chunking the corpus included, from a cheap CPU notebook with nothing left running afterward. The Processing-jobs episode reruns the same pipeline on an open-weight model Bedrock does not carry, as jobs that end on their own, which shows what Bedrock does for you and what it costs to do yourself. An optional self-study episode runs the same self-hosted pipeline interactively on a GPU notebook instance, for the case where you need to work on a self-hosted model by hand, and only if you will remember to stop the instance. All three 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 inference only. Nothing is trained or fine-tuned, so running one does not need a GPU of your own.
Bedrock is the default for 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 for a model Bedrock does not offer, or for batch volumes so large that per-token pricing loses; it is not a starting point.
Of the GPU options, Processing jobs are the safest, because the instance ends 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 under a dollar.
Tags never propagate. Tag every job at launch and route Bedrock calls through a tagged application inference profile, or the spend cannot be traced.
The hands-on episodes start with Bedrock, then the same pipeline on Processing jobs; the notebook-GPU version is optional self-study.

Content from RAG with Bedrock

Last updated on 2026-11-02 | 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 uses Amazon Bedrock, the default the overview recommended. We build the whole RAG pipeline for the WattBot 2026 challenge from a small CPU-only SageMaker notebook. Nothing here needs a GPU, because both the text embeddings and the answer generation run on Bedrock-hosted models, which bill per token and leave nothing running between calls. The two models are Amazon Titan Text Embeddings V2, which turns WattBot chunks into vectors, and DeepSeek V3, which writes the answers and explanations.

This episode also prepares the corpus. It unpacks the WattBot PDFs, splits them into pages and overlapping chunks, and caches the chunks in S3. The Processing-jobs episode and the optional GPU-notebook episode reuse those chunks and run the same retrieval, evaluation and scoring on self-hosted models, so cost, latency and effort can be compared directly.

Why Bedrock for WattBot?

With a self-hosted model, as in the Processing-jobs episode, you choose the model, keep its version, and make sure the instance has the GPU memory for it. That is fine for an experiment and a burden once several teams want the same pipeline. Running the embedding and generation steps on Bedrock means:

Managed models. Models from Anthropic, Amazon, DeepSeek and others, with no container images or CUDA versions of your own.
Paying per token. Instead of a GPU instance that may sit idle, you pay for the input and output tokens of each call. For most workloads that is cheaper; for a large offline batch with a small model, a GPU job can still win.
One set of models for a group. A lab or a course can agree on a few Bedrock models more easily than on a fleet of GPU instances.

In this notebook both the embedding step and the answer and explanation steps run on Bedrock. The embeddings are saved as embeddings_bedrock.npy, so they can be compared with the self-hosted run in the Processing-jobs episode.

Models used in this episode

Both models are Bedrock-hosted foundation models:

Amazon Titan Text Embeddings V2 (amazon.titan-embed-text-v2:0): a general text-embedding model for retrieval, clustering and classification. It returns vectors of 256, 512 or 1,024 dimensions, with presets for retrieval. AWS does not publish its parameter count.
DeepSeek V3 (deepseek.v3-v1:0): an open-weight mixture-of-experts model, 671B parameters in total and about 37B active per token, that Bedrock serves on demand, so there is no GPU to manage. Its per-token price is well below the flagship models’ (the Bedrock pricing page has current rates), which matters when there is one call per question. The same code works for any chat model in the catalog, because Bedrock’s one chat call, converse, takes the same input for every provider, a system prompt and a list of messages, and returns the reply in the same shape, so changing base_generation_model_id to a Claude, Llama or Mistral ID changes nothing else.
Multimodal models for tables and figures: Bedrock also serves models that embed or read images. Step 13 uses one, Amazon Nova Multimodal Embeddings, on the table and figure questions.

The model catalog in the Bedrock console lists every model available in the region (Claude variants, Amazon models and partner models), each with a model card and a link to its pricing.

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, the chart to read for RAG, since 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 labeled 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, not 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 setup cells below, holds the WattBot inputs under wattbot/: metadata.csv (the index of papers), train_QA.csv (labeled 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 four 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 setup 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 sys
sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")   # the lesson's helper module, cloned in the controllers episode
import helpers
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; prices are on the pricing page, 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, because 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 in the shared bucket are 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, which will hold 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)

The inputs, as they sit in the shared bucket. The three files named above are at the top of wattbot/; the rest of the folder is what the instructor staged for the caching steps later in the episode.

PYTHON

helpers.show_bucket(shared_bucket, "wattbot/")   # every file under the prefix, with its size

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, for listing models and 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, so 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. Copy the module from the lesson clone to the folder 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,
    split_ref_ids,
    run_single_qa,
)

Open wattbot_rag.py in the Jupyter file browser as you go. Each helper is short.

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 come 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 labeled 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")      # labeled 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, because the instance’s volume is 20 GB, 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 the cell below loads the staged result and falls back to extracting only if it is missing. 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 chunks keep their page

Every chunk carries the page it came from, so the document id, URL and page label stay attached to the text, a retrieved chunk can be traced to the page the model saw, the text can be cleaned at the page level before chunking, and chunk sizes can be changed later without extracting the PDFs again.

PYTHON

%pip install -q pypdf

pdfs_to_page_docs in wattbot_rag.py opens each PDF listed in the metadata with pypdf and turns every page into one record holding its text, the document id, title and URL, and the page number. It is forty lines of plain Python with nothing AWS in it; open the file to read it. On a CPU notebook it takes several minutes over this corpus, so the cell first looks for pages.jsonl, the same records written once when the corpus was staged, and runs the extraction only if that file is not there.

PYTHON

from wattbot_rag import pdfs_to_page_docs   # the extraction: open each PDF with pypdf, one record per page; in scripts/wattbot_rag.py

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, not 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

Pages are split into chunks because a page is the wrong size for retrieval: too long, and one vector has to stand for several topics at once; too short, and a chunk carries no context. The chunker here is a sliding window over characters:

chunk_size_chars: the most characters in a chunk (1,000 to 1,500 is usual).
chunk_overlap_chars: how many characters consecutive chunks share (200 here), so a sentence cut by a boundary is whole in one of them.

In your own work, splitting on headings, sections or sentence boundaries, called semantic chunking, is the usual next step.

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 through a tagged 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 comes to three things. Create the alias once per model (one for DeepSeek, one for Titan, because an alias points at a single 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. 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"   # your challenge: WattBot2026, BadgerScribe, BadgerCode, BrainDecoding, PoisonedWell or SnapshotWI; SinkOrSwim is the workshop itself
print(f"Tags on the profiles, and so on every Bedrock call: Name={name}, Project={project}, Purpose=Inference")   # check them first: a profile tagged DoeJohn bills as DoeJohn

# 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

POWERSHELL

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 with bedrock.list_inference_profiles(typeEquals="APPLICATION"); the console’s inference-profile pages list only the system-defined profiles.
The keys must be the three approved ones, Name, Project and Purpose, with that spelling and capitalization; 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, 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 this. 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, so 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. The next three cells set where the result is saved, define the call that embeds one string, and run 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. This corpus is about 2.4 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 practice.

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, since 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 with 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, a normalize and a 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 with 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. It takes 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), 53 questions after overlaps and the two missing papers, so every row of the breakdown in the next step has something in it. The first 53 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 220 remaining questions
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 normalization, 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, which are 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.

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


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, where 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, as with a question asking what inference “can account for” labeled 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 organizers 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 on-demand, 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.50 of GPU time for them, three 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 53-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 also lists 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, so 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, and the picture of a table lands near a sentence about its contents. Bedrock has two from Amazon, 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 Nova Multimodal Embeddings is offered in us-east-1 and not in Ohio (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, instead of handing 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 about eighty 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 is that 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, not 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; the 2,213 crops take about four 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, not 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, where 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; our runs took 1.6, 4.0 and 3.0 minutes, the image runs longer 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. At the paper level the text index wins. A ground-truth paper is in its top 8 for 94% of table questions and 94% of figure questions, against 77% and 71% for the image index, and for 82% against 55% of quote-only questions, where a picture has nothing to offer. At the crop level the two routes are even. The crop the answer sits in is among those shown for 55% of table questions by the text route and 52% by the image index, and for 57% against 50% of figure questions. 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 favor, 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 the crop-level hit rates say where the image index’s lead does not come from. It finds the answer’s own crop no more often than the text route does. The difference is in what else the model is shown, three crops chosen by their similarity to the question against every crop on the pages behind the top chunks. 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

You now have a complete WattBot evaluation on Bedrock, billed per token, with nothing left running. The Processing-jobs episode reruns it on a self-hosted model, and the optional GPU-notebook episode once more by hand, so by the end there are three versions of the same pipeline to compare, Bedrock, a Processing job and a notebook GPU. The code they share is the retrieval, prompting and scoring in wattbot_rag.py; what each one implements is the two callables, embed_query and generate.

Use Bedrock or another hosted API when a catalog model will do, when the number of questions is modest or you are still prototyping, and when you would rather pay per token than manage a GPU. Use a self-hosted model on a GPU when the work needs a model the catalog does not carry, when you will run large batches repeatedly, when you need control over the exact checkpoint, or when you already have GPUs you are not paying for by the hour.

The model is one part of a RAG system, and not the part where most of the accuracy is won. A bigger or newer model does not fix weak retrieval. If the chunks do not match the questions, a strong model still has nothing to read. The gains over time come from the work around the model, chunking, metadata and filtering, reranking, post-processing, and the score breakdown by evidence type in Step 12 says which of them to work on. A bigger model or a longer prompt raises the bill per question as well as the score, and Bedrock lets you measure that trade by changing one model ID.

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-11-02 | Edit this page

Estimated time: 50 minutes

Overview

Questions

When does a RAG pipeline need a self-hosted, open-weight model instead of Bedrock?
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 built the whole WattBot pipeline without managing a GPU, and that stays the simplest choice whenever Bedrock carries a model that will do. Research often depends on a particular open-weight checkpoint, though. A newly released model, a smaller size, a fine-tune for scientific text, a model a paper just published. Bedrock’s catalog is curated. 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, and the one you want is usually not on Bedrock, and often never will be.

In that case you run the model on hardware you rent through SageMaker, and the next decision is how long to keep that hardware. SageMaker rents a GPU three ways.

	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 releases the instance on its own, so it cannot be forgotten. This episode runs the WattBot pipeline as two jobs. The first turns the cached chunks into embeddings with thenlper/gte-large. The second loads the generator, retrieves evidence for each question and writes answers in WattBot’s submission format. The notebook stays a small CPU instance that never touches a model. It stages data, launches the 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, where 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. An 8B model is small enough to run unquantized, 8B parameters at two bytes each in bf16, about 16 GB, which fits the 24 GB GPU of an ml.g5.xlarge with room for the KV cache. Larger models are quantized for inference, to one byte per parameter or less, which is how the GPU-notebook episode fits a 27B model on one card. (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.)

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

Rates are us-east-1 on-demand from the Instances for ML page.

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`	~15 min: 1 provisioning, ~8 downloading and loading 16 GB of weights, ~6 generating	~`$0.35`
Controller notebook	`ml.t3.medium`	`$0.05`	2 h	`$0.10`

About $0.65 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 that follows), where 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 the same 53-question sample as the Bedrock episode.
Score the result and compare it with the Bedrock run. The retrieval logic, prompts and scorer are unchanged, so the comparison is of the models, not of the pipeline.

Setup

A code folder for each job

A Processing job runs in a container on a separate machine. It cannot see the Python variables or files that exist only in this notebook, and the only code it gets is the one directory named by source_dir=, uploaded whole, the same way the training estimator’s source_dir worked in the fine-tuning episode. The four scripts the jobs need are already on this instance, in the lesson clone, but pointing source_dir at the lesson’s whole scripts/ directory would upload every script in the workshop. So each job gets a small folder holding just the files it imports, copied out of the clone. The four files:

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 normalization, 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 into the project folder, code_embed/ for the embedding job and code_rag/ for the RAG job. 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

Now connect the notebook to SageMaker and S3. The shared bucket holds the challenge data. Your own bucket holds the chunks the Bedrock episode cached and will receive the job outputs.

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, because 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, as with the 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"   # your challenge: WattBot2026, BadgerScribe, BadgerCode, BrainDecoding, PoisonedWell or SnapshotWI; SinkOrSwim is the workshop itself

import sys
sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")
import helpers   # check_launch: the account's rules, with a message that names the problem

def make_job_tags(purpose: str):   # one of: Storage, Notebook, Training, Tuning, Inference
    return helpers.check_launch([   # checked against the account's rules first, with a message that names the problem
        {"Key": "Name", "Value": name},
        {"Key": "Project", "Value": project},
        {"Key": "Purpose", "Value": purpose},   # differs per job so the bill tells them apart
    ])

print("Tags on every job this episode launches:", make_job_tags("Inference"))   # check them before launching: a job tagged DoeJohn is on the bill as DoeJohn

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, normalizes 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. Inputs in, outputs out, no state.

A HuggingFaceProcessor is the Processing-job counterpart of the estimator, with 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

In the console the job is under SageMaker AI → Data preparation → Processing jobs, or at this link for 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. That startup is the cost of a job, paid once. The embedding job runs once per corpus, and the GPU bills only while it 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: the WattBot RAG loop as a job

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, with the same retrieval, prompts and unanswerable rule, with the model calls grouped four at a time because a GPU decoding one prompt is idle most of the time. The group is a trade. More prompts at once is faster, and each one adds its own cache of attention keys and values to the 16 GB the weights already hold, so a group of long prompts can run the 24 GB card out of memory. The script halves the group and retries when that happens, so a long batch costs time, not the job. 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 four 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 is refused with an AccessDeniedException ending in explicit deny, the account’s checks are the cause: a missing tag, an instance type outside the lesson’s list, or a resource kind the lesson never creates; the first training episode’s callout If the launch is refused says what to look at. 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 this is what happens when you launch one. 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, saying 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= sets two environment variables in the container. HF_HOME points Hugging Face’s download cache at that volume, which is where the weights go, and PYTORCH_CUDA_ALLOC_CONF tells PyTorch to reuse freed GPU memory in place instead of leaving it fragmented, which matters when a batch of prompts is close to what the card holds.

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 variables in the container: Hugging Face downloads to this path...
         "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"},   # ...and PyTorch reuses freed GPU memory instead of fragmenting it, which is what a batch of long prompts runs out of
    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", "4",                                   # prompts decoded together; 1 would be one at a time, several times slower. 8 ran out of GPU memory on long prompts
    ],
)

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, Loading model with transformers, the weight download (16 GB in shards), a line per batch of answers and of explanations, then one line per question. The whole job is the 15 minutes in the cost callout; the generation is the last six of them.

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.54 against DeepSeek V3’s 0.69. Its citations were better (F1 0.70 against 0.60) and its answers worse (47% against 70% correct), with the cross-paper and calculation questions losing the most, a fifth of each right against three fifths. 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 14 minutes (about $0.32), so about $0.50 of GPU time, the callout’s GPU figure. In the RAG job’s log, the first 8 minutes of the billed window are the download and the model load; the last 6 are the questions, two generation calls each. Setup is more than half 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

What a shorter runtime cap would have done

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 the cap each job set, 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, not 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.g6e.xlarge notebook instance forgotten over a weekend is 64 h × $2.61 = $167, and nothing caps 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 is small enough to run unquantized, ~16 GB in bf16 on one 24 GB GPU, with the rest for the prompts’ attention cache; how many prompts decode together is bounded by that remainder. Larger models are quantized to fit.
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-11-03 | 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.
Build the embedding matrix and the retrieval step with a self-hosted embedding model, in plain NumPy.
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 instead of 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 work on a self-hosted model by hand, 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.

Why run RAG on a GPU notebook?

The Bedrock and Processing Jobs episodes are the better patterns for most research work. Bedrock serves the models for you, and a Processing job rents a GPU only long enough to finish a batch. A notebook GPU gives you something different, an interactive place to load a model, watch GPU memory, try retrieval strategies and change prompts without repackaging a job after every experiment.

That matters when you need a model Bedrock does not offer and want to work on it by hand. This episode runs the WattBot pipeline with Qwen3.8-27B as the generator, a 27-billion-parameter open-weight (Apache 2.0) model released in August 2026 that Bedrock does not carry, loaded in FP8 through vLLM beside a local embedding model. Everything a managed model API hides, loading an embedding model, loading a generator, fitting both in GPU memory, is a cell you can run and change.

A 27B model is large enough that the choice of GPU and weight precision matters. At about one byte per weight, FP8 puts the weights near 27 GB, and a 48 GB L40S leaves room for the embedding model and the attention cache that generation needs. The 24 GB GPU the earlier episodes used cannot hold this setup. That makes the notebook a good place to learn about model serving, and an expensive place to leave idle.

Treat this as an interactive experiment, not a deployment pattern. Create the separate GPU notebook described below. Do not run it on the CPU controller from the other episodes.

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 a 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.

What keeps 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.

What we’ll build

We reuse the WattBot 2026 document chunks and evaluation questions prepared in the Bedrock episode. The pipeline keeps its purpose. What changes is that its two managed model calls become two models running on this GPU, thenlper/gte-large for embeddings and Qwen3.8-27B for answers. Qwen3.8-27B is the newer and larger relative of the Qwen3-8B the Processing Jobs episode ran, 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.

First we load the cached chunks and embed them into a searchable matrix. Then we implement cosine-similarity retrieval so a question brings back the most relevant passages. Finally we load Qwen through vLLM, pass those passages into the same WattBot prompts, and generate answers, citations and short explanations that label the evidence type, [text], [table], [figure] or a combination such as [text+table]. Scoring those answers against the competition’s references tells us how the self-hosted pipeline performs, not only whether it runs.

The workflow is kept close to the earlier episodes on purpose. With the documents, prompts and scoring the same, the differences in model serving are the only thing left to compare.

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 on the Instances for ML page). 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; the next cell stops the notebook"

PYTHON

# also verify you have selected the conda_pytorch kernel
import os
os.environ["HF_HOME"] = "/home/ec2-user/SageMaker/hf_cache"   # before any Hugging Face import: every model download goes to the volume, not the small root disk
import torch

# Stop here, and so stop a Run All, unless this instance has a GPU big enough for the 27B model:
# the install below is several gigabytes and the weights are 27 GB, neither worth downloading onto the wrong machine
if not torch.cuda.is_available():
    raise RuntimeError("No GPU: this notebook needs an ml.g6e.xlarge notebook instance. Stop this instance and create one as described above.")
gpu_memory_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
if gpu_memory_gb < 40:
    raise RuntimeError(f"GPU has {gpu_memory_gb:.0f} GB; the FP8 27B model needs a 48 GB GPU (ml.g6e.xlarge). Stop this instance and create one as described above.")
print(f"GPU: {torch.cuda.get_device_name(0)}, {gpu_memory_gb:.0f} GB")

Install vLLM first

The install may replace PyTorch and then needs a kernel restart, so it comes before anything is loaded.

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"   # Qwen3.8 needs vLLM 0.17+, which brings its own PyTorch

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 with an open embedding model:

`thenlper/gte-large`

The embedder the Processing Jobs episode used has 335 million parameters, 0.7 GB in bf16, 1,024-dimensional vectors, and about 1,100 chunks a minute on an A10G, so the 9,000 chunks here take a few minutes on this GPU.

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, with no indexing library.

PYTHON

# We'll use a stronger embedding model now that we have a GPU.
# This model has ~335M parameters and benefits from GPU acceleration,
# and fits beside the 27B generator on this 48 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, 1024))
    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)

Retrieval here is cosine similarity between the question’s vector and every chunk’s vector, with the top k chunks as context, and nothing more. For nine thousand chunks that is one matrix multiply. A vector database (FAISS, Chroma, Weaviate) is the same step built for millions of chunks or many concurrent users, with approximate search, metadata filters and a shared index; swapping it in changes nothing else in the pipeline.

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 embedder 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, with 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, not the small root disk.

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 full set.

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, labeled [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, not 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)
from wattbot_rag import drop_questions_without_papers
train_df = drop_questions_without_papers(train_df, {d["doc_id"] for d in chunked_docs})   # as in the other RAG episodes: questions citing papers the corpus lacks are not scored

all_results = []
max_N = np.inf   # set to e.g. 5 for a short 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 normalization, 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 pass over all 245 questions takes tens of minutes, and a max_N of 50 a few. 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 a first pass at a RAG pipeline. Next steps, in rough order of what they usually gain:

Stronger models: a larger or domain-tuned embedding model, or a more capable generator, if the GPU budget allows.
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.
Models tuned on scientific text (arXiv, PubMed), which read equations and technical language more reliably and invent less when the prose is dense.

Models, retrieval and preprocessing can each be swapped without changing the rest of the pipeline.

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, from the Instances for ML page
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 on the 53-question sample billed about $0.50 as Processing jobs; 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, an 80 GB GPU or 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.
Retrieval is a matrix multiply: cosine similarity over the in-memory embedding matrix is enough for a corpus this size; a vector database is the same step for a corpus of millions.
Units come from the data: answer_unit is copied from train_QA.csv, never guessed by the model.
Two calls per question: one answers and cites; a second, shorter one writes the explanation and names the evidence type.
The competition’s format: is_blank marks unanswerable questions and missing fields, as the scorer expects.

Content from (Optional) The Same Jobs from the Terminal

Last updated on 2026-11-04 | 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.
Match each SDK call and helper the episodes used to the CLI command that does the same thing.
List and inspect training jobs, processing jobs, tuning jobs and notebook instances from the terminal, stop your own by name, and read a job’s bill the way job_cost_report() did.
Copy data to and from S3 in bulk with aws s3 sync.
Relaunch the training episode’s XGBoost job from its JSON description, with the same script, hyperparameters, cap and tags.
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 on a SageMaker notebook instance. That instance is a cost and a chore. It bills while it is on, five cents an hour for the ml.t3.medium, a dollar a day, and someone has to remember to stop it. It is also a second place to work, a Jupyter session in a browser tab, when your code, your editor and your files are on your laptop.

None of that is required. The notebook instance was a convenience for teaching and for governance in a shared account. SageMaker is an API, and the AWS CLI calls the same API from a terminal. A job launched from your laptop needs no instance at all. The request goes from your terminal to SageMaker, the job runs on a machine rented for it and ends, and the model lands in S3. For a project that launches a job a day and otherwise works locally, that removes the notebook instance’s whole cost and the habit of stopping it. The terminal is also the tool for checking on a job from another machine, stopping something when no notebook is open, and copying results out of S3 in bulk.

This episode does five things from a laptop terminal. It signs in to the account with Identity Center, lists what is running and reads what a job cost, stops your own jobs and your own notebook instance by name, copies data to and from S3, and relaunches the training episode’s XGBoost job from a JSON file.

Two things differ from the notebook. The first is identity. On the instance, its role signed every call. On a laptop you sign in yourself, which is section 1. The second is checking. The lesson’s notebooks printed a cost estimate and ran a launch check before every job. A JSON request does neither on its own, so section 5 runs the lesson’s checker on the file before the create call. It applies the account’s rules and prints the most the request can bill.

On a shared workshop account, such as the ML Marathon’s, keep launching from the notebook instance. It is required there, because every learner’s launches and the code behind them then sit on an instance the administrators can see, which is how they help when something goes wrong. Launching from the terminal is for your own account, and for practicing it in this episode. Checking on jobs and stopping them from the terminal is fine on any account.

What the notebooks did, and the command that does the same

The CLI, the SageMaker SDK, boto3 and the console all talk to the same AWS API. The SDK’s estimators built requests for it, the boto3 calls in the cleanup and RAG episodes sent requests to it directly, and the console shows its results. So every call the episodes made has a CLI command with the same name and the same fields, and a boto3 method name such as describe_training_job is the CLI command describe-training-job with the underscores swapped for hyphens. This table is the map. The sections below work through the rows that matter most.

Episode	In the notebook	From the terminal
Controllers, Training	`sagemaker.Session()`, `get_execution_role()`	`aws sts get-caller-identity`, and the role ARN is the `RoleArn` field of any job you launched
Training	`helpers.get_notebook_instance_info(name, region)`	`aws sagemaker describe-notebook-instance --notebook-instance-name <name>`
Accessing S3	`s3.upload_file`, `s3.download_file`, `helpers.show_bucket(bucket)`	`aws s3 cp`, `aws s3 sync`, `aws s3 ls --recursive`
Training, PyTorch	`xgboost_estimator.fit({'train': train_input})`	`aws sagemaker create-training-job --cli-input-json file://train-xgboost.json` (section 5)
Training	`fit(wait=True)`, the default	`aws sagemaker wait training-job-completed-or-stopped --training-job-name <name>`
Training	`job_cost_report(job_name, price)`	`aws sagemaker describe-training-job --training-job-name <name>` (section 2)
Hyperparameter tuning	`tuner.fit(...)`	`aws sagemaker create-hyper-parameter-tuning-job --cli-input-json file://...`
Hyperparameter tuning	`tuning_cost_report(tuner, price)`, `tuner.best_training_job()`	`aws sagemaker list-training-jobs-for-hyper-parameter-tuning-job`, `describe-hyper-parameter-tuning-job --query BestTrainingJob.TrainingJobName`
PyTorch, Hyperparameter tuning, Fine-tuning	`helpers.estimate_training_cost(...)`, `helpers.estimate_tuning_cost(...)`, `helpers.estimate_finetune_cost(...)`	the same functions from a laptop (section 5)
Fine-tuning	`helpers.check_launch(...)`, `helpers.check_finetune_launch(...)`	`python scripts/check_request.py <request>.json --max-cost <dollars>` (section 5)
RAG with Processing jobs	`rag_processor.run(...)`, `describe_processing_job(...)`	`aws sagemaker create-processing-job --cli-input-json file://...`, `aws sagemaker describe-processing-job --processing-job-name <name>`
RAG with Bedrock	`bedrock.list_inference_profiles(typeEquals="APPLICATION")`, `bedrock_runtime.converse(...)`	`aws bedrock list-inference-profiles --type-equals APPLICATION`, `aws bedrock-runtime converse --model-id <profile ARN> ...`
Cleanup	`sm.list_notebook_instances()`, `bedrock.delete_inference_profile(...)`	`aws sagemaker list-notebook-instances`, `aws bedrock delete-inference-profile --inference-profile-identifier <ARN>` (section 3)

The SDK’s job classes (XGBoost, PyTorch, HuggingFace, HuggingFaceProcessor, HyperparameterTuner) are builders for the JSON these commands take. describe-* on any job you launched from a notebook prints that JSON back, which is how section 5 gets its file.

What to replace in the commands

Anything in angle brackets, <like this>, is a value of yours. Replace it, brackets included. This table lists every placeholder in the episode and how to find its value.

Placeholder	What it is	Where to find it
`<notebook-instance-name>`, then `$NB`	The name you gave your notebook instance in the controllers episode	The first column of `aws sagemaker list-notebook-instances` in section 2, where it is set once as `$NB`
`<your-bucket>`, then `$BUCKET`	The bucket you created in the S3 episode	`aws s3 ls` prints every bucket you can see; set once as `$BUCKET` in section 2
`<a training job name>`	Any training job you ran	The list command in section 2, or the name the training notebook printed when `fit()` ran
`<your xgboost job name>`	The XGBoost job from the training episode, named `<notebook-instance-name>-xgboost-<timestamp>`	`aws sagemaker list-training-jobs --name-contains xgboost --query 'TrainingJobSummaries[].TrainingJobName'`
`<name>` in section 3	A job or notebook instance to stop	The lists in section 2
`<image>`, `<role>`, `<code>`	The container, the role and the uploaded code of the XGBoost job	Printed by the `describe-training-job` command at the start of section 5
`<LastnameFirstname>`, `<ChallengeName>`	Your `Name` and `Project` tag values	The same values you put on your bucket and your notebook instance
`<ARN>`	A Bedrock inference profile’s ARN	The second column of `aws bedrock list-inference-profiles` in section 3

Callout

Read this before the first command

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, and 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 automatically. Endpoints, EC2 instances and GPU notebook instances do not, and the lesson never creates them from the CLI, or at all.
A refused create command is the account’s rules, not a fault. The workshop account carries a guard-rail policy that denies endpoint and EC2 creation, instance types outside the lesson’s list, caps over 48 hours, and any job whose three tags are missing, blank or off the Project and Purpose lists. The administrator attaches that policy, to the notebook role and to the permission set your laptop signs in with. You attach nothing. If a create command is refused, that policy is why, and AWS says only AccessDenied. The checker in section 5 applies the same rules to a request before it is sent, so you see the reason instead. (Instructors: the policy and where it is attached are on the Account preparation checklist page.)

This episode runs from your own laptop, in a terminal, so it works when no notebook is open. CloudShell and the SageMaker notebook instance already have the AWS CLI installed, and on the instance it uses the instance’s role, so there is no sign-in. On a laptop you install the CLI and sign in yourself.

Install AWS CLI v2 from the installer for your operating system. Then open a terminal, Terminal on macOS, Git Bash or PowerShell on Windows, and check that it is on your path:

BASH

aws --version

Register the account’s Identity Center, once. Run the command below in the same terminal and answer its prompts in this order. A profile is a named set of these answers, saved in ~/.aws/config on your laptop, so that the CLI knows which account, role and region a command is for. One laptop can hold several, one per account you use.

BASH

aws configure sso

Prompt	Answer
SSO session name	`workshop`
SSO start URL	Your sign-in link from the instructor’s email, up to and including `/start`, with nothing after it (for UW-Madison it begins with `https://uw-madison` and ends in `awsapps.com/start`)
SSO region	The region Identity Center runs in, from the same email (`us-east-2` for UW-Madison). It is not always the account’s working region
SSO registration scopes	Enter, for the default
A browser tab opens	Sign in with your university login if asked. On the page headed Allow botocore-client-workshop to access your data?, choose Allow access
Account	The workshop account, whose ID is in the sign-in email. If your login has only one, the CLI picks it for you
Role	The permission set your login carries (for UW-Madison, the one with `sagemaker` in its name). If there is only one, the CLI picks it
Default client region	`us-east-2` for the workshop account
Default output format	Enter, which is `json`
Profile name	`workshop`, in place of the long default it offers, since the rest of the episode uses that name

The wizard ends by printing aws sts get-caller-identity --profile workshop. That is the identity check from the next block.

The role here is the permission set, not the notebook role. The permission set is the identity you sign in with, the same one the console uses. The notebook role, ml-sagemaker-use, is what a job runs as, and it stays in the job’s RoleArn in section 5.

The commands in this episode are the same in Bash (macOS, Linux, Git Bash on Windows) and in PowerShell. Where the shell’s own syntax differs, variables, line continuation, writing a file, the two are shown as tabs.

From then on, every session starts with three commands. The login fetches credentials for the profile, which takes you through the browser once more. Setting AWS_PROFILE makes every later command in that terminal use the profile without a --profile flag on each one. The identity check confirms 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
aws sts get-caller-identity

POWERSHELL

aws sso login --profile workshop
$env:AWS_PROFILE = "workshop"
aws sts get-caller-identity

The login ends with Successfully logged into Start URL: .... The identity check prints the account ID and an Arn that names an AWSReservedSSO_... role, the permission set, and ends with your email. Check the region too, since SageMaker resources are regional and a job in another region does not appear in any list:

BASH

aws configure get region     # us-east-2 for the workshop account; if it is wrong, run aws configure sso again

The credentials expire after a few hours. Run aws sso login again when a command returns ExpiredToken. Nothing that would still work next month is written anywhere.

If you are in the notebook instance’s own terminal instead (JupyterLab, File → New → Terminal), skip the install and the sign-in. The instance’s role is already the identity, and every command below works there as written.

The sections below are the AWS commands, so that you see what the SDK was doing. Most of the lesson’s helpers wrapped one of these calls and are not needed here. The ones that do arithmetic, the cost estimators and the launch check, come back in section 5, from the lesson’s scripts/helpers.py, plain Python that runs on a laptop as it did on the notebook instance. Clone the lesson once and install boto3, which the module imports:

BASH

git clone -q https://github.com/carpentries-incubator/ML_with_AWS_SageMaker.git   # once, anywhere on the laptop
pip install boto3

2. What is running, and what a job cost

First the lists. Each should be short, or empty. --query is a JMESPath filter that keeps the output to what you asked about, and the cleanup episode’s first cell ran the same four list_* calls through boto3.

BASH

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

Your notebook instance is in the fourth list, and aws s3 ls prints your bucket. Those two names come up in most commands below, so set them as shell variables once per terminal session, after the login, and the commands use $NB and $BUCKET from here on:

BASH

aws s3 ls
NB='<notebook-instance-name>'
BUCKET='<your-bucket>'

POWERSHELL

aws s3 ls
$NB = "<notebook-instance-name>"
$BUCKET = "<your-bucket>"

Then one job’s bill. job_cost_report() in the training episodes called describe_training_job and read four fields from it, the status, the instance count, BillableTimeInSeconds, and the model’s S3 path, and computed the wall-clock from CreationTime and TrainingEndTime. This is the same call with the same fields. Pick a job from the list below, newest first with its status, and put its full name in place of <a training job name>. On the shared account the list is everyone’s jobs, so do not take the newest. Take the one named after your notebook instance:

BASH

aws sagemaker list-training-jobs --sort-by CreationTime --sort-order Descending --query 'TrainingJobSummaries[].[TrainingJobName,TrainingJobStatus]' --output table
# only your own, since the lesson names every job after its notebook instance: filter on that name
# aws sagemaker list-training-jobs --name-contains $NB --sort-by CreationTime --sort-order Descending --query 'TrainingJobSummaries[].[TrainingJobName,TrainingJobStatus]' --output table
# the same list in name order, which groups each learner's jobs together
# aws sagemaker list-training-jobs --sort-by Name --sort-order Ascending --query 'TrainingJobSummaries[].[TrainingJobName,TrainingJobStatus]' --output table

BASH

JOB='<a training job name>'
aws sagemaker describe-training-job --training-job-name $JOB \
  --query '{status:TrainingJobStatus, instance:ResourceConfig.InstanceType, count:ResourceConfig.InstanceCount, billable_s:BillableTimeInSeconds, created:CreationTime, ended:TrainingEndTime, model:ModelArtifacts.S3ModelArtifacts}'

POWERSHELL

$JOB = "<a training job name>"
aws sagemaker describe-training-job --training-job-name $JOB `
  --query '{status:TrainingJobStatus, instance:ResourceConfig.InstanceType, count:ResourceConfig.InstanceCount, billable_s:BillableTimeInSeconds, created:CreationTime, ended:TrainingEndTime, model:ModelArtifacts.S3ModelArtifacts}'

Cost is billable_s × count ÷ 3600 × the hourly rate from the Instances for ML page, the arithmetic job_cost_report() did. The helper printed it. Here you do the multiplication. billable_s can be missing for a job that failed. Read the status and FailureReason before reading anything into its cost.

The other job kinds have the same shape. The Processing Jobs episode’s job_cost() read CreationTime, ProcessingStartTime and ProcessingEndTime from describe_processing_job, and the tuning episode’s tuning_cost_report() looped describe_training_job over the tuner’s child jobs:

BASH

aws sagemaker describe-processing-job --processing-job-name <name> --query '{status:ProcessingJobStatus, created:CreationTime, started:ProcessingStartTime, ended:ProcessingEndTime, instance:ProcessingResources.ClusterConfig.InstanceType}'
aws sagemaker list-training-jobs-for-hyper-parameter-tuning-job --hyper-parameter-tuning-job-name <name> --query 'TrainingJobSummaries[].[TrainingJobName,TrainingJobStatus,FinalHyperParameterTuningJobObjectiveMetric.Value]' --output table
aws sagemaker describe-hyper-parameter-tuning-job --hyper-parameter-tuning-job-name <name> --query 'BestTrainingJob.TrainingJobName'   # tuner.best_training_job()

Tags are returned by a separate call. The Name tag is how you confirm a job is yours before you stop it in the next section, and a job with no tags is the one that shows on the bill under “No tag key”:

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

POWERSHELL

$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. Stopping jobs and notebook instances

Each of these stops a charge. They are the calls behind the Stop buttons the cleanup episode walks through in the console. Every one takes a name, and on a shared account the name must be yours. The lists in section 2 show everyone’s jobs and instances. Before any stop command, confirm the Name tag with the list-tags call from section 2, and stop only what you started. Someone else’s job is a question for them, or for the instructor.

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/ and keeps billing for that volume. Deleting the instance is the step that stops the storage charge, and the cleanup episode’s last two items, the bucket and the Bedrock inference profiles, have their CLI forms too:

BASH

aws sagemaker delete-notebook-instance --notebook-instance-name <name>   # only after it reports Stopped
aws s3 rb s3://$BUCKET --force                                             # empties the bucket, then deletes it; the cleanup episode's delete_objects loop plus delete_bucket
aws bedrock list-inference-profiles --type-equals APPLICATION --query 'inferenceProfileSummaries[].[inferenceProfileName,inferenceProfileArn]' --output table
aws bedrock delete-inference-profile --inference-profile-identifier <ARN>  # one per profile, as the cleanup episode's loop did

4. Copying to and from S3

The Accessing S3 episode moved single files with s3.upload_file and s3.download_file and listed the bucket with helpers.show_bucket. For many files, or for pulling a job’s outputs down to a laptop, the CLI is the better tool, and these are the same keys those episodes used:

BASH

aws s3 ls s3://$BUCKET/ --recursive --human-readable --summarize             # helpers.show_bucket, for the whole bucket
aws s3 cp titanic_train.csv s3://$BUCKET/titanic_train.csv                     # s3.upload_file(train_filename, bucket_name, train_filename)
aws s3 cp s3://$BUCKET/xgboost/<job-name>/output/model.tar.gz model.tar.gz    # the download in the training episode's evaluate step
aws s3 sync s3://$BUCKET/xgboost/ ./xgboost-outputs/                           # every job's outputs under the prefix; only new or changed files move
aws s3 rm s3://$BUCKET/scratch/ --recursive                                    # delete a prefix; this one asks no questions

The last line, sync --delete and the rb --force in section 3 remove data for good and ask nothing first. Check the bucket name twice before any of them, and on the shared account run them only on your own bucket.

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. Launching the training episode’s job from a JSON file

xgboost_estimator.fit({'train': train_input}) in the training episode assembled a JSON request and sent it to CreateTrainingJob. SageMaker kept that request, and describe-training-job on the job you ran prints it back. Three fields come from there, because they are specific to your account and region, and the rest of the file below is the estimator’s arguments under their API names.

First find the job’s name. The training episode named it after your notebook instance, so it is the one with xgboost in it:

BASH

aws sagemaker list-training-jobs --name-contains "$NB-xgboost" --query 'TrainingJobSummaries[].TrainingJobName'

Put that name in place of <your xgboost job name> below.

BASH

JOB='<your xgboost job name>'
aws sagemaker describe-training-job --training-job-name $JOB \
  --query '{image:AlgorithmSpecification.TrainingImage, role:RoleArn, code:HyperParameters.sagemaker_submit_directory}'

POWERSHELL

$JOB = "<your xgboost job name>"
aws sagemaker describe-training-job --training-job-name $JOB `
  --query '{image:AlgorithmSpecification.TrainingImage, role:RoleArn, code:HyperParameters.sagemaker_submit_directory}'

The describe output is a description, not a request. It carries fields AWS fills in after the fact, so it cannot be saved and sent back as it is, which is why the file below is written out from the estimator’s arguments and takes only three values from the description.

image is the XGBoost 1.7-1 container that framework_version="1.7-1" picked. role is the role the notebook printed. code is where fit() uploaded source_dir, the sourcedir.tar.gz in the SDK’s default bucket that the training episode pointed out under What landed in S3. Reusing it means this job runs the same train_xgboost.py the notebook’s job ran. Save the file below as train-xgboost.json, in the folder your terminal is in, with any text editor. The file is the same whichever shell you use, and a file cannot read shell variables, so the two names are written out in it. Put the three values in place of <image>, <role> and <code>, your bucket in place of <your-bucket>, your notebook instance’s name in TrainingJobName, and your tag values in place of <LastnameFirstname> and <ChallengeName>. The code value prints with quotation marks inside it; put only the s3://... path in place of <code>, since the file supplies the quotes. The -cli-1 suffix makes the job name unique; use -cli-2 for a second run, since SageMaker refuses a name it has seen.

JSON

{
  "TrainingJobName": "<notebook-instance-name>-xgboost-cli-1",
  "AlgorithmSpecification": {
    "TrainingImage": "<image>",
    "TrainingInputMode": "File"
  },
  "RoleArn": "<role>",
  "HyperParameters": {
    "train": "\"titanic_train.csv\"",
    "max_depth": "3",
    "eta": "0.1",
    "subsample": "0.8",
    "colsample_bytree": "0.8",
    "num_round": "100",
    "sagemaker_program": "\"train_xgboost.py\"",
    "sagemaker_submit_directory": "\"<code>\"",
    "sagemaker_region": "\"us-east-2\"",
    "sagemaker_container_log_level": "20"
  },
  "InputDataConfig": [
    {
      "ChannelName": "train",
      "ContentType": "csv",
      "DataSource": { "S3DataSource": {
        "S3DataType": "S3Prefix",
        "S3Uri": "s3://<your-bucket>/titanic_train.csv",
        "S3DataDistributionType": "FullyReplicated" } }
    }
  ],
  "OutputDataConfig": { "S3OutputPath": "s3://<your-bucket>/xgboost/" },
  "ResourceConfig": { "InstanceType": "ml.m5.large", "InstanceCount": 1, "VolumeSizeInGB": 30 },
  "StoppingCondition": { "MaxRuntimeInSeconds": 1200 },
  "Tags": [
    { "Key": "Name",    "Value": "<LastnameFirstname>" },
    { "Key": "Project", "Value": "<ChallengeName>" },
    { "Key": "Purpose", "Value": "Training" }
  ]
}

Read it against the estimator cell in the training episode. Every field is one of its arguments:

In the estimator	In the JSON
`base_job_name=f"{notebook_instance_name}-xgboost"` plus the SDK’s timestamp	`TrainingJobName`, which you write out in full
`framework_version="1.7-1"`	`AlgorithmSpecification.TrainingImage`
`role=role`	`RoleArn`
`hyperparameters={'train': train_file, 'max_depth': max_depth, ...}`	`HyperParameters`, the same six keys and values, as strings
`entry_point='train_xgboost.py'`, `source_dir=.../scripts`	`sagemaker_program` and `sagemaker_submit_directory`, which the SDK adds to the hyperparameters so the container knows what to run
`TrainingInput(train_s3_path, content_type='csv')` and `fit({'train': ...})`	`InputDataConfig`, one channel named `train`
`output_path=f's3://{bucket_name}/xgboost/'`	`OutputDataConfig.S3OutputPath`
`instance_type="ml.m5.large"`, `instance_count=1`	`ResourceConfig`, with the SDK’s default 30 GB volume
`max_run=20*60`	`StoppingCondition.MaxRuntimeInSeconds`, 1200
`tags=job_tags`	`Tags`, the same three keys

The string values with escaped quotes, "\"titanic_train.csv\"", are how the SDK encodes a text hyperparameter, and the container decodes them before they reach argparse in the script, so --train titanic_train.csv arrives as it did from the notebook. Numbers are plain strings.

Before sending it, estimate it, as the notebooks did before every launch. Nothing in a JSON file prints a price, so the hand method is the instance’s hourly rate from the Instances for ML page times the minutes you expect, and MaxRuntimeInSeconds at about twice that. This XGBoost job is startup and little else, two to three minutes on ml.m5.large, under a cent. For a job with real compute the lesson’s estimators do the arithmetic, with the same arguments the notebooks gave them. They are called from python -c, with the lesson’s scripts folder put on the path first. These are the PyTorch episode’s estimate, the tuning episode’s and the fine-tuning episode’s, and each prints the time, the cost and the max_run to put in the request:

BASH

python -c "import sys; sys.path.insert(0, 'ML_with_AWS_SageMaker/scripts'); import helpers; helpers.estimate_training_cost(params_billions=0.125, examples=1_000_000, tokens_per_example=256, epochs=3, instance_type='ml.g6.xlarge')"
python -c "import sys; sys.path.insert(0, 'ML_with_AWS_SageMaker/scripts'); import helpers; helpers.estimate_tuning_cost(train_minutes_per_job=1, price_per_hour=0.115, max_jobs=8)"
python -c "import sys; sys.path.insert(0, 'ML_with_AWS_SageMaker/scripts'); import helpers; helpers.estimate_finetune_cost(params_billions=3.75, examples=251, tokens_per_example=1600, epochs=2, instance_type='ml.g5.xlarge')"

Then check the request. scripts/check_request.py reads the JSON and does two things. It applies the account’s rules from scripts/guardrails.json, as helpers.check_launch did before every launch in the notebooks, so a Project or Purpose off the lists, an instance type off the list or a cap over 48 hours is named here, where AWS would refuse it silently. Then it prints the most the request can bill: every instance it asks for, running to its cap, at the instance’s hourly price, and for a tuning job every child job it allows. --max-cost is a dollar amount you choose, the most you are willing to spend on this job. When the worst case is above it, the checker exits with an error, and the launch command below is written so that an error stops it. Here it is 1, one dollar, because the job’s worst case is four cents:

BASH

python ML_with_AWS_SageMaker/scripts/check_request.py train-xgboost.json --max-cost 1   # refuse if this job could bill more than $1

OUTPUT

train-xgboost.json: tags, instance type and cap pass the account's rules
worst case if every job runs to its cap: 1 x 0.33 h x $0.115/h = $0.04

The same command takes a Processing job’s or a tuning job’s request, and the tuning job’s worst case multiplies by its MaxNumberOfTrainingJobs, which is where a tuning request surprises people.

The check catches configuration and cost mistakes before a request reaches AWS. It does not grant permission to launch the job or guarantee that AWS accepts it. The account’s own policy decides that, and on the shared workshop account the rule is the one above, launch from the notebook instance.

Then send it. This one line is estimator.fit(), and the && in front of it means the create call runs only when the check passed:

BASH

python ML_with_AWS_SageMaker/scripts/check_request.py train-xgboost.json --max-cost 1 && aws sagemaker create-training-job --cli-input-json file://train-xgboost.json

The CLI returns the moment the request is accepted. Nothing waits, which is fit(wait=False). To follow the job the way the default wait=True did, then read its bill as in section 2:

BASH

aws sagemaker wait training-job-completed-or-stopped --training-job-name "$NB-xgboost-cli-1"   # blocks until it ends
aws sagemaker describe-training-job --training-job-name "$NB-xgboost-cli-1" --query '{status:TrainingJobStatus, billable_s:BillableTimeInSeconds, failure:FailureReason, model:ModelArtifacts.S3ModelArtifacts}'

Expect the same two to three billable minutes on ml.m5.large as the notebook’s job, under a cent, and the same xgboost-model.json inside the model.tar.gz it names.

The PyTorch job, a Processing job and a tuning job follow the same route. describe-training-job, describe-processing-job and describe-hyper-parameter-tuning-job on the ones you ran print the request, and create-training-job, create-processing-job and create-hyper-parameter-tuning-job take it back with --cli-input-json. A PyTorch job’s HyperParameters carries epochs and learning_rate and two input channels. A Processing job’s AppSpecification.ContainerArguments is the arguments= list. A tuning job’s file has the estimator’s request nested under TrainingJobDefinition and the ranges under HyperParameterTuningJobConfig.ParameterRanges.

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 session.
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 guard-rail policy the administrator attached, as is any instance type outside the lesson’s list, any cap over 48 hours and any job whose tags are missing or off the lists. On your own account, the policy is on the instructor page Account preparation checklist. Attaching it to your own user or role takes an afternoon and is worth it.

Challenge

Challenge: what is still running on Friday evening

You are leaving for the weekend and have a terminal, not a browser. Write the commands you would run, in order, to make sure nothing of yours is costing money until Monday, and the one command to run if a check finds something. Then run them.

Show me the solution

BASH

aws sso login --profile workshop && export AWS_PROFILE=workshop

POWERSHELL

aws sso login --profile workshop; $env:AWS_PROFILE = "workshop"

BASH

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 of yours the first four print, confirm it with list-tags and run the matching stop-... command from section 3. A job you expect to finish within its MaxRuntimeInSeconds can be left. A job that is not yours is its owner’s, and the tags say who that is. For an endpoint, aws sagemaker delete-endpoint --endpoint-name <name>, then find out how it was created, because on the workshop account the guard-rail policy refuses it. Buckets cost almost nothing over a weekend and need no action.

Key Points

The CLI, the SDK and the console are one API. Every call the notebooks made has a CLI command with the same name and fields, and describe-* on a job you ran prints the request the SDK built.
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.
describe-training-job is job_cost_report(). Cost is billable seconds × instance count × the hourly rate, and you do the multiplication.
A create-training-job JSON is what estimator.fit() sent, the same script, hyperparameters, channel, cap and tags, with no price printed next to it.
Endpoints, EC2 instances and GPU notebooks are not jobs and do not end. The workshop account’s notebook role denies creating them, and your own account should too.

Content from Resource Management and Monitoring

Last updated on 2026-11-02 | 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.

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.

Initial setup

PYTHON

import boto3
import sagemaker
import sys

sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")   # the lesson's helper module, cloned in the controllers episode
import helpers

session = sagemaker.Session()
sm = session.sagemaker_client          # the low-level SageMaker client: list_*, describe_*, stop_* calls
region = session.boto_region_name      # the region the notebook runs in (us-east-2 on the workshop account)

name = "DoeJohn"                              # replace with your name, as used in job names and tags
project = "SinkOrSwim"   # the Project value used in your Bedrock profile names
bucket_name = "sinkorswim-doejohn"            # replace with your bucket

What is still running under your name?

One pass over the three job lists shows every job still running in the account; yours are the ones with your notebook name in the training jobs and the WattBot- Processing jobs you launched. A completed job holds nothing; an InProgress one bills until it ends or hits its cap.

PYTHON

# Every notebook instance in the account (the list is shared, so other learners' instances appear too), with its status.
# InService bills the instance's hourly rate; Stopped bills only the volume.
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']}")

# The three kinds of job, each with its own list call. The call returns a dictionary; the list of jobs is under a key
# whose name differs per kind ("TrainingJobSummaries", ...), and so does the key holding the job's name, so the loop
# carries both. StatusEquals keeps only running jobs; NameContains keeps the ones with your name in the job name.
print("\nJobs still InProgress in the account:")
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", MaxResults=50)[key]:   # every running job in the account; the Processing jobs are not named after you, so no name filter
        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)   # sums the size of every object in the bucket; S3 bills per GB-month whether or not anything runs

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)   # the Bedrock control-plane client: profiles are created and deleted here, not on the runtime client
mine = f"{project}-{name}-"   # the Bedrock episode named every profile <project>-<name>-<label>, so this prefix is yours and nobody else's
# list_inference_profiles returns every application profile in the account; keep the ones whose name starts with your prefix
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"])   # the ARN is what the delete call below takes
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.

PYTHON

confirm_delete = False   # set to True after reading the list above

if confirm_delete:
    for p in profiles:   # only the profiles the cell above listed: yours, by prefix
        bedrock.delete_inference_profile(inferenceProfileIdentifier=p["inferenceProfileArn"])   # one call per profile; there is no undo
        print("deleted", p["inferenceProfileName"])
    profiles = []        # so a rerun of this cell does not try to delete them twice
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 bill per second at an hourly rate while they run, so 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. A small instance left running still bills a dollar a day.

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 20 GB the controllers episode set is $2.80 a month; 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

S3 storage is cheap, and a bucket left behind bills every month until it is deleted.

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, so that only workshop resources are deleted.

Monitoring and stopping active jobs

SageMaker charges for training, tuning and Processing jobs while they run, so stop jobs you no longer need. 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

To see what you are spending:

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. To turn a line on the bill back into hours, the rates are on the Instances for ML page.
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 stopping 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.

Rules the account enforces, and how to set them on your own

Several things you could have done wrong in this workshop were refused before they cost anything: a job launched without the three tags, with a blank name, with a Project that is not on the account’s list, with an instance type outside the lesson’s list, a Processing job with no runtime cap, and any endpoint or EC2 instance at all. None of that is SageMaker’s own behavior. It is a set of rules the account administrator attached to the role your notebook runs as, written in the same IAM policy language that gave the role its permissions, and each one says: refuse this request if it looks like this. An AWS account with no such rules lets a tagged job and an untagged one through alike, and lets a notebook instance run until someone remembers it.

If you administer an account of your own, a lab’s or a project’s, the same rules are an afternoon’s work, and they hold across the console, the SDK and the CLI, because all three are the same API. The lesson’s Account preparation checklist instructor page has the policy as JSON and the one command that attaches it to a role; its Overnight safeguards section has the three measures that go with it. In order of what they prevent for what they cost:

Rule	What it stops	How it is done
Every job carries the three cost tags, with real values	spend nobody can attribute	an IAM deny on `CreateTrainingJob` and the like when `aws:RequestTag/Name` is missing or blank, or `Project` and `Purpose` are missing or not on their lists
Only instance types you have priced	the `$63`-an-hour box chosen by accident	an IAM deny when `sagemaker:InstanceTypes` is not on your list
A runtime cap on every job, at most 48 hours	a hung job billing for a week	an IAM deny when `sagemaker:MaxRuntimeInSeconds` is missing or too large
No endpoints, no EC2	the two resource kinds with no runtime cap at all	an IAM deny on `CreateEndpoint` and `ec2:RunInstances`
Notebook instances stop at 11 pm	the forgotten GPU notebook, `$63` a night	a ten-line scheduled Lambda function
A budget per person and a brake on the total	a quiet overspend	AWS Budgets on the `Name` tag, with an action that attaches a deny policy at the ceiling

A deny applies to everyone who uses the role, including you, so your own launches need real tag values and listed instance types too. And the error a refused request gets, AccessDeniedException ... with an explicit deny in an identity-based policy, does not say which rule matched, so write the rules down where the people using the account will find them. A deny cannot carry a custom message either: AWS says explicit deny in an identity-based policy and nothing more, which is why this lesson keeps a copy of the policy next to its helper code and checks every launch against it in the notebook first, with helpers.check_launch, so the message you get names the rule. The If a Launch Is Refused page is the same list in one place, for the day a job of your own is refused.

What you can do without an administrator:

Automate the cleanup you do by hand. The inventory cells at the top of this episode are the SDK calls; a script that runs them and stops what it finds is the Lambda in the table above.
Test on a slice first. Every expensive run in this lesson was preceded by a small one on the same code: a few epochs, a subset of pages, a handful of questions.
Estimate, then compare. Write down what a run should cost before launching it, read BillableTimeInSeconds afterwards, and find out why when the two disagree, as the fine-tuning episode does.
Tag everything, yourself, every time. Cost allocation tags are the only way a shared bill becomes individual bills, and they never propagate.

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_runtime_in_seconds 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.
The refusals you met in this workshop (untagged jobs, wrong instance types, endpoints) are IAM deny rules an administrator attached to the role, not SageMaker defaults; on an account you administer, the instructor pages have the policy to copy.
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.

All in One View

Overview

Questions

Objectives

Cloud vs. the cluster down the hall

S3, the notebook instance, SageMaker and Bedrock

How the episodes use them

Overview

Questions

Objectives

Storing data on AWS

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

The notebook’s disk versus S3

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

Messages 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

When the file is bigger than the notebook

Check the current size and storage costs of bucket

PYTHON

Using the lesson’s helper functions

PYTHON

PYTHON

The shared bucket

PYTHON

Check storage costs of bucket

PYTHON

PYTHON

Writing output files to S3

PYTHON

PYTHON

PYTHON

PYTHON

Challenge: before and after

Show me the solution

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

Helper: `get_notebook_instance_info()`