Skip to content

Commit 4042a24

Browse files
SapanaChaudharymufaddal-rohawala
SapanaChaudhary
and
mufaddal-rohawala
authored
documentation: add end-to-end MTRL SOP-Bench example (AgentCore) (#6532)
Add v3-examples/model-customization-examples/serverless/mtrl/, a complete, runnable multi-turn RL example: fine-tune openai-reasoning-gpt-oss-20b on the SOP-Bench aircraft_inspection task with a rollout agent on Amazon Bedrock AgentCore. - 01_build_and_deploy_agent.ipynb: task, tools, reward design (and the reward hack it prevents), AgentCore deploy - 02_train.ipynb: prompt dataset, MultiTurnRLTrainer GRPO run, reading MLflow metrics and traces - 03_evaluate_and_deploy.ipynb: MultiTurnRLEvaluator base vs. fine-tuned, failure diagnosis, deployment options, scoped cleanup - sopbench_agent/: the agent package (entrypoint, agent loop, tools, reward), Dockerfile, and scripts - prerequisites.yaml: IAM roles, S3 bucket, MLflow app (or an existing one), and managed policies for the notebook identity Co-authored-by: Mufaddal Rohawala <89424143+mufaddal-rohawala@users.noreply.github.com>
1 parent 14deb89 commit 4042a24

32 files changed

Lines changed: 4220 additions & 0 deletions

‎v3-examples/model-customization-examples/serverless/mtrl/01_build_and_deploy_agent.ipynb‎

Lines changed: 452 additions & 0 deletions
Large diffs are not rendered by default.

‎v3-examples/model-customization-examples/serverless/mtrl/02_train.ipynb‎

Lines changed: 370 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,370 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"## Multi-Turn RL with SageMaker AI, part 2: train with `MultiTurnRLTrainer`\n",
8+
"\n",
9+
"You turn the task into a prompt dataset, launch a GRPO training job against the agent you deployed\n",
10+
"in part 1, and learn to read a run in MLflow - metrics *and* trajectories.\n",
11+
"\n",
12+
"| # | Notebook | What you do |\n",
13+
"|---|---|---|\n",
14+
"| 1 | `01_build_and_deploy_agent.ipynb` | Explore the task, tools, and reward; deploy the rollout agent to Amazon Bedrock AgentCore |\n",
15+
"| 2 | `02_train.ipynb` | Build the prompt dataset and launch a GRPO training job with `MultiTurnRLTrainer`; read MLflow |\n",
16+
"| 3 | `03_evaluate_and_deploy.ipynb` | Compare base vs. fine-tuned with `MultiTurnRLEvaluator`, diagnose failures, deploy, clean up |\n",
17+
"\n",
18+
"**Before you start:** complete `01_build_and_deploy_agent.ipynb` (prerequisites stack deployed,\n",
19+
"dependencies installed, agent `READY`)."
20+
],
21+
"id": "02-00"
22+
},
23+
{
24+
"cell_type": "markdown",
25+
"metadata": {},
26+
"source": [
27+
"### 1. Configuration"
28+
],
29+
"id": "02-01"
30+
},
31+
{
32+
"cell_type": "code",
33+
"metadata": {},
34+
"execution_count": null,
35+
"outputs": [],
36+
"source": [
37+
"import os\n",
38+
"import sys\n",
39+
"from pathlib import Path\n",
40+
"\n",
41+
"import boto3\n",
42+
"from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer\n",
43+
"\n",
44+
"STACK_NAME = \"mtrl-prerequisites\" # the prerequisites.yaml stack you deployed\n",
45+
"AGENT_NAME = \"sopbench_agentcore\" # the name passed to `agentcore configure --name`\n",
46+
"PREFIX = \"sopbench-mtrl/aircraft\" # S3 prefix for the prompt parquets\n",
47+
"AGENT_DIR = Path(\"sopbench_agent\").resolve()\n",
48+
"\n",
49+
"# Make `!python` / `!agentcore` resolve to this kernel's environment.\n",
50+
"os.environ[\"PATH\"] = f\"{Path(sys.executable).parent}{os.pathsep}{os.environ['PATH']}\"\n",
51+
"\n",
52+
"sess = boto3.session.Session()\n",
53+
"REGION = sess.region_name\n",
54+
"assert REGION, \"Set AWS_REGION (or a default region) to the region you deployed the stack in.\"\n",
55+
"# The training job must run in the AgentCore runtime's region: pin it for every SDK client.\n",
56+
"os.environ[\"AWS_DEFAULT_REGION\"] = os.environ[\"AWS_REGION\"] = REGION\n",
57+
"account = sess.client(\"sts\").get_caller_identity()[\"Account\"]\n",
58+
"\n",
59+
"stack = sess.client(\"cloudformation\").describe_stacks(StackName=STACK_NAME)[\"Stacks\"][0]\n",
60+
"outs = {o[\"OutputKey\"]: o[\"OutputValue\"] for o in stack.get(\"Outputs\", [])}\n",
61+
"BUCKET = outs[\"PromptsBucketName\"]\n",
62+
"ROLE_ARN = outs[\"JobRoleArn\"]\n",
63+
"AGENTCORE_EXEC_ROLE_ARN = outs[\"AgentCoreExecRoleArn\"]\n",
64+
"MLFLOW_APP_ARN = outs[\"MlflowAppArn\"]\n",
65+
"TRAIN_DATASET = f\"s3://{BUCKET}/{PREFIX}/train/aircraft_inspection_train.parquet\"\n",
66+
"VAL_DATASET = f\"s3://{BUCKET}/{PREFIX}/validation/aircraft_inspection_val.parquet\"\n",
67+
"\n",
68+
"\n",
69+
"def find_agent_runtime_arn(name=AGENT_NAME):\n",
70+
" \"\"\"ARN of the READY AgentCore runtime with exactly this name (others may share a prefix).\"\"\"\n",
71+
" for runtime in MultiTurnRLTrainer.list_bedrock_agentcore_runtimes(session=sess):\n",
72+
" if runtime[\"name\"] == name and runtime[\"status\"] == \"READY\":\n",
73+
" return runtime[\"arn\"]\n",
74+
" return None\n",
75+
"\n",
76+
"\n",
77+
"print(f\"account={account} region={REGION} bucket={BUCKET}\")\n",
78+
"AGENT_RUNTIME_ARN = find_agent_runtime_arn()\n",
79+
"S3_OUTPUT_PATH = f\"s3://{BUCKET}/output/\"\n",
80+
"print(\"AGENT_RUNTIME_ARN =\", AGENT_RUNTIME_ARN)\n",
81+
"assert AGENT_RUNTIME_ARN, f\"No READY runtime named {AGENT_NAME} - run 01_build_and_deploy_agent.ipynb.\""
82+
],
83+
"id": "02-02"
84+
},
85+
{
86+
"cell_type": "markdown",
87+
"metadata": {},
88+
"source": [
89+
"### 2. Build and upload the prompt dataset\n",
90+
"\n",
91+
"`scripts/preprocess.py` turns each SOP-Bench task into one parquet row with a single `prompt` column\n",
92+
"- a JSON string with `benchmark`, `instance_id`, `task_input`, and `ground_truth`. The service\n",
93+
"forwards that string verbatim to the agent, so this schema is the contract between the data and the\n",
94+
"agent. It makes a deterministic 80/20 train/validation split (89 / 23 prompts)."
95+
],
96+
"id": "02-03"
97+
},
98+
{
99+
"cell_type": "code",
100+
"metadata": {},
101+
"execution_count": null,
102+
"outputs": [],
103+
"source": [
104+
"!cd {AGENT_DIR} && PYTHONPATH=src python scripts/preprocess.py --s3-prefix s3://{BUCKET}/{PREFIX}"
105+
],
106+
"id": "02-04"
107+
},
108+
{
109+
"cell_type": "markdown",
110+
"metadata": {},
111+
"source": [
112+
"### 3. How a training step works\n",
113+
"\n",
114+
"Each step, the trainer samples a batch of prompts and asks the agent for **`group_size`** rollouts\n",
115+
"per prompt. The reward scores each rollout; **GRPO** computes each rollout's advantage relative to\n",
116+
"its group's mean and updates the policy (a LoRA adapter on the base model).\n",
117+
"\n",
118+
"Because the advantage is *relative within the group*, a group whose rollouts all get the same\n",
119+
"reward contributes **no learning signal** - which is why partial credit and reward variance matter."
120+
],
121+
"id": "02-05"
122+
},
123+
{
124+
"cell_type": "markdown",
125+
"metadata": {},
126+
"source": [
127+
"### 4. Build the trainer\n",
128+
"\n",
129+
"The fine-tuned adapter is registered in a model package group, auto-created as\n",
130+
"`{model}-mtrl-mpg` (checkpoints in `{model}-mtrl-checkpoint-mpg`) and **shared by every MTRL job in\n",
131+
"the account**. To keep your runs separate, pass `output_model_package_group=` and\n",
132+
"`intermediate_checkpoint_model_package_group=` (two different groups)."
133+
],
134+
"id": "02-06"
135+
},
136+
{
137+
"cell_type": "code",
138+
"metadata": {},
139+
"execution_count": null,
140+
"outputs": [],
141+
"source": [
142+
"from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer\n",
143+
"\n",
144+
"MODEL = \"openai-reasoning-gpt-oss-20b\"\n",
145+
"assert MODEL in MultiTurnRLTrainer.list_supported_models(), f\"{MODEL} not supported\"\n",
146+
"\n",
147+
"trainer = MultiTurnRLTrainer(\n",
148+
" model=MODEL,\n",
149+
" agent_env=AGENT_RUNTIME_ARN, # AgentCore runtime ARN (managed path)\n",
150+
" training_dataset=TRAIN_DATASET,\n",
151+
" validation_dataset=VAL_DATASET,\n",
152+
" mlflow_app_arn=MLFLOW_APP_ARN,\n",
153+
" s3_output_path=S3_OUTPUT_PATH,\n",
154+
" role=ROLE_ARN,\n",
155+
" accept_eula=True, # gated model\n",
156+
")\n",
157+
"print(trainer.hyperparameters.get_info()) # authoritative names, types, defaults, and ranges"
158+
],
159+
"id": "02-07"
160+
},
161+
{
162+
"cell_type": "markdown",
163+
"metadata": {},
164+
"source": [
165+
"### 5. Hyperparameters\n",
166+
"\n",
167+
"A ready-to-run recipe. A step takes about 3-4 minutes, so the 10 steps below finish in well under\n",
168+
"an hour; a 3-step test run took about 12 minutes and raised held-out `val/reward/mean` from 2.4 to\n",
169+
"4.7 (out of 5). With 89 prompts and `global_batch_size=32` an epoch is\n",
170+
"only ~3 steps, so `max_epochs` is a high cap and **`max_steps` stops the run**. Lower `max_steps`\n",
171+
"(e.g. 3) to finish faster; the flow is identical.\n",
172+
"\n",
173+
"| Knob | Value | Why |\n",
174+
"|---|---|---|\n",
175+
"| `advantage_method` | `group_based` | GRPO; the library also offers per-turn and other estimators |\n",
176+
"| `group_size` | 4 | rollouts per prompt; larger = better baseline, more rollouts per step |\n",
177+
"| `global_batch_size` | 32 | prompts per step |\n",
178+
"| `sampling_max_tokens` | 4096 | per-rollout token cap; also bounds trajectory length |\n",
179+
"| `learning_rate`, `lora_rank`, `lora_alpha` | 4e-5, 32, 64 | LoRA adapter training |\n",
180+
"| `val_every` | 10 | validation at step 0 and step 10 |\n",
181+
"| `success_threshold` | 5.0 | pass@k counts a rollout as a success when reward >= this; 5.0 (`FINAL_REWARD_SCALE`) = every field correct. The default 1.0 would count ~20%-correct answers. |\n",
182+
"| `rollout_max_concurrency` | 64 | the **only** throttle on AgentCore; lower toward 32 and/or raise `rollout_max_retries` if you see `httpx ConnectTimeout` |\n",
183+
"\n",
184+
"Concurrency trades wall-clock for connection pressure; it does not change the reward the model\n",
185+
"reaches. A longer recipe (`max_steps=30`, `global_batch_size=64`, `sampling_max_tokens=8192`) is in\n",
186+
"`sopbench_agent/scripts/run_rft_training.py`, the non-interactive equivalent of this notebook. It\n",
187+
"took about 2 h 15 min and raised held-out pass@1 from 0.48 to 0.87."
188+
],
189+
"id": "02-08"
190+
},
191+
{
192+
"cell_type": "code",
193+
"metadata": {},
194+
"execution_count": null,
195+
"outputs": [],
196+
"source": [
197+
"from sopbench_agentcore.reward import FINAL_REWARD_SCALE\n",
198+
"\n",
199+
"hp = trainer.hyperparameters\n",
200+
"hp.advantage_method = \"group_based\" # GRPO\n",
201+
"hp.loss_fn = \"ppo\"\n",
202+
"hp.global_batch_size = 32\n",
203+
"hp.group_size = 4\n",
204+
"hp.sampling_max_tokens = 4096\n",
205+
"hp.learning_rate = 4e-5\n",
206+
"hp.lora_rank = 32\n",
207+
"hp.lora_alpha = 64\n",
208+
"hp.max_epochs = 15 # high cap; max_steps governs when the run stops\n",
209+
"hp.max_steps = 10\n",
210+
"hp.val_every = 10\n",
211+
"hp.success_threshold = FINAL_REWARD_SCALE # pass@k success = a fully correct answer (5.0)\n",
212+
"hp.rollout_max_concurrency = 64\n",
213+
"hp.rollout_max_retries = 3"
214+
],
215+
"id": "02-09"
216+
},
217+
{
218+
"cell_type": "markdown",
219+
"metadata": {},
220+
"source": [
221+
"### 6. Launch (non-blocking)"
222+
],
223+
"id": "02-10"
224+
},
225+
{
226+
"cell_type": "code",
227+
"metadata": {},
228+
"execution_count": null,
229+
"outputs": [],
230+
"source": [
231+
"job = trainer.train(wait=False)\n",
232+
"print(\"Job: \", job.job_name)\n",
233+
"print(\"Status:\", job.job_status)\n",
234+
"print(\"MLflow:\", job.get_mlflow_url())"
235+
],
236+
"id": "02-11"
237+
},
238+
{
239+
"cell_type": "code",
240+
"metadata": {},
241+
"execution_count": null,
242+
"outputs": [],
243+
"source": [
244+
"job.refresh()\n",
245+
"print(f\"status={job.job_status} secondary={getattr(job, 'secondary_status', '?')}\")"
246+
],
247+
"id": "02-12"
248+
},
249+
{
250+
"cell_type": "markdown",
251+
"metadata": {},
252+
"source": [
253+
"### 7. Read the run in MLflow\n",
254+
"\n",
255+
"MLflow is your primary observability surface: metrics stream every step, so you diagnose a run\n",
256+
"*while it runs*. Open it from the URL above, or from the SageMaker console (MLflow -> your app ->\n",
257+
"**Open MLflow**) - printed links carry a short-lived token. Your jobs appear in the **Runs** list:\n",
258+
"\n",
259+
"![MLflow Runs list for the openai-reasoning-gpt-oss-20b experiment, one row per training job with status and a duration of roughly 2-3 hours.](images/mlflow-all-runs.png)\n",
260+
"\n",
261+
"**Watch these first:**\n",
262+
"\n",
263+
"| Metric | Read it as |\n",
264+
"|---|---|\n",
265+
"| `rollout/reward/mean`, `rollout/reward/valid_mean` | training reward (valid_mean: over non-zero-advantage groups) - should climb |\n",
266+
"| `analysis/zero_adv_groups` | groups with no variance, so no learning signal - high means a saturated or too-coarse reward |\n",
267+
"| `rollout/tokens/response_max` | if it pins at `sampling_max_tokens`, rollouts truncate mid-answer |\n",
268+
"| `val/reward/mean` | held-out reward every `val_every` steps - should **track** training reward |\n",
269+
"\n",
270+
"Switch to chart mode and MLflow groups related metrics, such as the `rollout/reward` group\n",
271+
"(`max`, `mean`, `min`). Rewards lie in `[-1, 5]`.\n",
272+
"\n",
273+
"Mean climbing and holding is the model learning the output contract and tool sequence; max reaching\n",
274+
"5 means the best rollouts are already perfect; min lifting off the floor (and away from\n",
275+
"the `-1` no-final-answer penalty) means even the worst rollouts improve. The min-max spread is what\n",
276+
"GRPO learns from. If `val/reward/mean` **tracks** training reward, the model is generalizing; if it\n",
277+
"flattens while training reward keeps rising, suspect a reward hack.\n",
278+
"\n",
279+
"In-training validation is forced to greedy `temperature=0`, so its pass@k is degenerate - use the\n",
280+
"standalone evaluator in part 3 for real pass@k."
281+
],
282+
"id": "02-13"
283+
},
284+
{
285+
"cell_type": "markdown",
286+
"metadata": {},
287+
"source": [
288+
"### 8. Read the trajectories\n",
289+
"\n",
290+
"Metrics say the run is healthy; **traces** say what the agent actually did. The **Traces** tab lists\n",
291+
"one row per rollout - add the `reward` column from **Columns** for fast triage:\n",
292+
"\n",
293+
"![MLflow Traces tab: one row per aircraft_inspection rollout with token counts around 22,000-28,000 and a reward column showing mostly 5 with one row at 4.16667.](images/mlflow-traces-all.png)\n",
294+
"\n",
295+
"![MLflow Traces Columns menu with attributes such as Trace ID, Request, Response, Tokens, and Execution time.](images/mlflow-trace-columns.png)\n",
296+
"\n",
297+
"Open a trace: **Summary** shows the exact task inputs and each turn's tool call in order;\n",
298+
"**Details & Timeline** shows every turn's `thinking`, `tool_calls`, and tool result:\n",
299+
"\n",
300+
"![MLflow trace Summary for aircraft a_00262: the prompt JSON followed by turn-1 VerifyAircraftClearance, turn-2 VerifyMechanicalComponents, turn-3 VerifyElectricalSystems, turn-4 CrossCheckSpecifications.](images/mlflow-trace-landing-page.png)\n",
301+
"\n",
302+
"![MLflow trace Details and Timeline: span tree from turn-1 to turn-8 including the three reporting tools, and turn-1's thinking and tool_calls JSON.](images/mlflow-trace-details.png)\n",
303+
"\n",
304+
"Read off a trace: the **tool sequence** (a skipped step or a loop on one tool explains a low\n",
305+
"reward), **where the reward lands** (on the root `trajectory` span and the final turn; intermediate\n",
306+
"turns carry 0), and **per-turn tokens** (which turn is long when `response_max` is high)."
307+
],
308+
"id": "02-14"
309+
},
310+
{
311+
"cell_type": "markdown",
312+
"metadata": {},
313+
"source": [
314+
"### 9. Reattach after a kernel restart\n",
315+
"\n",
316+
"Find your job and reattach - no need to rerun the notebook."
317+
],
318+
"id": "02-15"
319+
},
320+
{
321+
"cell_type": "code",
322+
"metadata": {},
323+
"execution_count": null,
324+
"outputs": [],
325+
"source": [
326+
"from itertools import islice\n",
327+
"\n",
328+
"from sagemaker.train.agent_rft_job import AgentRFTJob\n",
329+
"\n",
330+
"for j in islice(AgentRFTJob.get_all(sort_by=\"CreationTime\", sort_order=\"Descending\"), 20):\n",
331+
" print(f\"{j.job_status:<12} {j.job_name}\")"
332+
],
333+
"id": "02-16"
334+
},
335+
{
336+
"cell_type": "code",
337+
"metadata": {},
338+
"execution_count": null,
339+
"outputs": [],
340+
"source": [
341+
"# job = MultiTurnRLTrainer.attach(job_name=\"<JobName from the list above>\")\n",
342+
"# job.refresh(); print(job.job_status, job.get_mlflow_url())\n",
343+
"# job.stop() # stop a running job"
344+
],
345+
"id": "02-17"
346+
},
347+
{
348+
"cell_type": "markdown",
349+
"metadata": {},
350+
"source": [
351+
"When the job is `Completed`, the LoRA adapter is registered as a SageMaker Model Package\n",
352+
"(`job.output_model_package_arn`). **Next:** `03_evaluate_and_deploy.ipynb`."
353+
],
354+
"id": "02-18"
355+
}
356+
],
357+
"metadata": {
358+
"kernelspec": {
359+
"display_name": "Python 3",
360+
"language": "python",
361+
"name": "python3"
362+
},
363+
"language_info": {
364+
"name": "python",
365+
"pygments_lexer": "ipython3"
366+
}
367+
},
368+
"nbformat": 4,
369+
"nbformat_minor": 5
370+
}

0 commit comments

Comments
 (0)