MTRL: Add end-to-end MTRL SOP-Bench example (AgentCore)

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

@@ -0,0 +1,370 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Multi-Turn RL with SageMaker AI, part 2: train with `MultiTurnRLTrainer`\n",
+    "\n",
+    "You turn the task into a prompt dataset, launch a GRPO training job against the agent you deployed\n",
+    "in part 1, and learn to read a run in MLflow - metrics *and* trajectories.\n",
+    "\n",
+    "| # | Notebook | What you do |\n",
+    "|---|---|---|\n",
+    "| 1 | `01_build_and_deploy_agent.ipynb` | Explore the task, tools, and reward; deploy the rollout agent to Amazon Bedrock AgentCore |\n",
+    "| 2 | `02_train.ipynb` | Build the prompt dataset and launch a GRPO training job with `MultiTurnRLTrainer`; read MLflow |\n",
+    "| 3 | `03_evaluate_and_deploy.ipynb` | Compare base vs. fine-tuned with `MultiTurnRLEvaluator`, diagnose failures, deploy, clean up |\n",
+    "\n",
+    "**Before you start:** complete `01_build_and_deploy_agent.ipynb` (prerequisites stack deployed,\n",
+    "dependencies installed, agent `READY`)."
+   ],
+   "id": "02-00"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 1. Configuration"
+   ],
+   "id": "02-01"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import boto3\n",
+    "from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer\n",
+    "\n",
+    "STACK_NAME = \"mtrl-prerequisites\"  # the prerequisites.yaml stack you deployed\n",
+    "AGENT_NAME = \"sopbench_agentcore\"  # the name passed to `agentcore configure --name`\n",
+    "PREFIX = \"sopbench-mtrl/aircraft\"  # S3 prefix for the prompt parquets\n",
+    "AGENT_DIR = Path(\"sopbench_agent\").resolve()\n",
+    "\n",
+    "# Make `!python` / `!agentcore` resolve to this kernel's environment.\n",
+    "os.environ[\"PATH\"] = f\"{Path(sys.executable).parent}{os.pathsep}{os.environ['PATH']}\"\n",
+    "\n",
+    "sess = boto3.session.Session()\n",
+    "REGION = sess.region_name\n",
+    "assert REGION, \"Set AWS_REGION (or a default region) to the region you deployed the stack in.\"\n",
+    "# The training job must run in the AgentCore runtime's region: pin it for every SDK client.\n",
+    "os.environ[\"AWS_DEFAULT_REGION\"] = os.environ[\"AWS_REGION\"] = REGION\n",
+    "account = sess.client(\"sts\").get_caller_identity()[\"Account\"]\n",
+    "\n",
+    "stack = sess.client(\"cloudformation\").describe_stacks(StackName=STACK_NAME)[\"Stacks\"][0]\n",
+    "outs = {o[\"OutputKey\"]: o[\"OutputValue\"] for o in stack.get(\"Outputs\", [])}\n",
+    "BUCKET = outs[\"PromptsBucketName\"]\n",
+    "ROLE_ARN = outs[\"JobRoleArn\"]\n",
+    "AGENTCORE_EXEC_ROLE_ARN = outs[\"AgentCoreExecRoleArn\"]\n",
+    "MLFLOW_APP_ARN = outs[\"MlflowAppArn\"]\n",
+    "TRAIN_DATASET = f\"s3://{BUCKET}/{PREFIX}/train/aircraft_inspection_train.parquet\"\n",
+    "VAL_DATASET = f\"s3://{BUCKET}/{PREFIX}/validation/aircraft_inspection_val.parquet\"\n",
+    "\n",
+    "\n",
+    "def find_agent_runtime_arn(name=AGENT_NAME):\n",
+    "    \"\"\"ARN of the READY AgentCore runtime with exactly this name (others may share a prefix).\"\"\"\n",
+    "    for runtime in MultiTurnRLTrainer.list_bedrock_agentcore_runtimes(session=sess):\n",
+    "        if runtime[\"name\"] == name and runtime[\"status\"] == \"READY\":\n",
+    "            return runtime[\"arn\"]\n",
+    "    return None\n",
+    "\n",
+    "\n",
+    "print(f\"account={account} region={REGION} bucket={BUCKET}\")\n",
+    "AGENT_RUNTIME_ARN = find_agent_runtime_arn()\n",
+    "S3_OUTPUT_PATH = f\"s3://{BUCKET}/output/\"\n",
+    "print(\"AGENT_RUNTIME_ARN =\", AGENT_RUNTIME_ARN)\n",
+    "assert AGENT_RUNTIME_ARN, f\"No READY runtime named {AGENT_NAME} - run 01_build_and_deploy_agent.ipynb.\""
+   ],
+   "id": "02-02"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 2. Build and upload the prompt dataset\n",
+    "\n",
+    "`scripts/preprocess.py` turns each SOP-Bench task into one parquet row with a single `prompt` column\n",
+    "- a JSON string with `benchmark`, `instance_id`, `task_input`, and `ground_truth`. The service\n",
+    "forwards that string verbatim to the agent, so this schema is the contract between the data and the\n",
+    "agent. It makes a deterministic 80/20 train/validation split (89 / 23 prompts)."
+   ],
+   "id": "02-03"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!cd {AGENT_DIR} && PYTHONPATH=src python scripts/preprocess.py --s3-prefix s3://{BUCKET}/{PREFIX}"
+   ],
+   "id": "02-04"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 3. How a training step works\n",
+    "\n",
+    "Each step, the trainer samples a batch of prompts and asks the agent for **`group_size`** rollouts\n",
+    "per prompt. The reward scores each rollout; **GRPO** computes each rollout's advantage relative to\n",
+    "its group's mean and updates the policy (a LoRA adapter on the base model).\n",
+    "\n",
+    "Because the advantage is *relative within the group*, a group whose rollouts all get the same\n",
+    "reward contributes **no learning signal** - which is why partial credit and reward variance matter."
+   ],
+   "id": "02-05"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 4. Build the trainer\n",
+    "\n",
+    "The fine-tuned adapter is registered in a model package group, auto-created as\n",
+    "`{model}-mtrl-mpg` (checkpoints in `{model}-mtrl-checkpoint-mpg`) and **shared by every MTRL job in\n",
+    "the account**. To keep your runs separate, pass `output_model_package_group=` and\n",
+    "`intermediate_checkpoint_model_package_group=` (two different groups)."
+   ],
+   "id": "02-06"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer\n",
+    "\n",
+    "MODEL = \"openai-reasoning-gpt-oss-20b\"\n",
+    "assert MODEL in MultiTurnRLTrainer.list_supported_models(), f\"{MODEL} not supported\"\n",
+    "\n",
+    "trainer = MultiTurnRLTrainer(\n",
+    "    model=MODEL,\n",
+    "    agent_env=AGENT_RUNTIME_ARN,          # AgentCore runtime ARN (managed path)\n",
+    "    training_dataset=TRAIN_DATASET,\n",
+    "    validation_dataset=VAL_DATASET,\n",
+    "    mlflow_app_arn=MLFLOW_APP_ARN,\n",
+    "    s3_output_path=S3_OUTPUT_PATH,\n",
+    "    role=ROLE_ARN,\n",
+    "    accept_eula=True,                     # gated model\n",
+    ")\n",
+    "print(trainer.hyperparameters.get_info())  # authoritative names, types, defaults, and ranges"
+   ],
+   "id": "02-07"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 5. Hyperparameters\n",
+    "\n",
+    "A ready-to-run recipe. A step takes about 3-4 minutes, so the 10 steps below finish in well under\n",
+    "an hour; a 3-step test run took about 12 minutes and raised held-out `val/reward/mean` from 2.4 to\n",
+    "4.7 (out of 5). With 89 prompts and `global_batch_size=32` an epoch is\n",
+    "only ~3 steps, so `max_epochs` is a high cap and **`max_steps` stops the run**. Lower `max_steps`\n",
+    "(e.g. 3) to finish faster; the flow is identical.\n",
+    "\n",
+    "| Knob | Value | Why |\n",
+    "|---|---|---|\n",
+    "| `advantage_method` | `group_based` | GRPO; the library also offers per-turn and other estimators |\n",
+    "| `group_size` | 4 | rollouts per prompt; larger = better baseline, more rollouts per step |\n",
+    "| `global_batch_size` | 32 | prompts per step |\n",
+    "| `sampling_max_tokens` | 4096 | per-rollout token cap; also bounds trajectory length |\n",
+    "| `learning_rate`, `lora_rank`, `lora_alpha` | 4e-5, 32, 64 | LoRA adapter training |\n",
+    "| `val_every` | 10 | validation at step 0 and step 10 |\n",
+    "| `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",
+    "| `rollout_max_concurrency` | 64 | the **only** throttle on AgentCore; lower toward 32 and/or raise `rollout_max_retries` if you see `httpx ConnectTimeout` |\n",
+    "\n",
+    "Concurrency trades wall-clock for connection pressure; it does not change the reward the model\n",
+    "reaches. A longer recipe (`max_steps=30`, `global_batch_size=64`, `sampling_max_tokens=8192`) is in\n",
+    "`sopbench_agent/scripts/run_rft_training.py`, the non-interactive equivalent of this notebook. It\n",
+    "took about 2 h 15 min and raised held-out pass@1 from 0.48 to 0.87."
+   ],
+   "id": "02-08"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from sopbench_agentcore.reward import FINAL_REWARD_SCALE\n",
+    "\n",
+    "hp = trainer.hyperparameters\n",
+    "hp.advantage_method = \"group_based\"   # GRPO\n",
+    "hp.loss_fn = \"ppo\"\n",
+    "hp.global_batch_size = 32\n",
+    "hp.group_size = 4\n",
+    "hp.sampling_max_tokens = 4096\n",
+    "hp.learning_rate = 4e-5\n",
+    "hp.lora_rank = 32\n",
+    "hp.lora_alpha = 64\n",
+    "hp.max_epochs = 15                    # high cap; max_steps governs when the run stops\n",
+    "hp.max_steps = 10\n",
+    "hp.val_every = 10\n",
+    "hp.success_threshold = FINAL_REWARD_SCALE  # pass@k success = a fully correct answer (5.0)\n",
+    "hp.rollout_max_concurrency = 64\n",
+    "hp.rollout_max_retries = 3"
+   ],
+   "id": "02-09"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 6. Launch (non-blocking)"
+   ],
+   "id": "02-10"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "job = trainer.train(wait=False)\n",
+    "print(\"Job:   \", job.job_name)\n",
+    "print(\"Status:\", job.job_status)\n",
+    "print(\"MLflow:\", job.get_mlflow_url())"
+   ],
+   "id": "02-11"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "job.refresh()\n",
+    "print(f\"status={job.job_status} secondary={getattr(job, 'secondary_status', '?')}\")"
+   ],
+   "id": "02-12"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 7. Read the run in MLflow\n",
+    "\n",
+    "MLflow is your primary observability surface: metrics stream every step, so you diagnose a run\n",
+    "*while it runs*. Open it from the URL above, or from the SageMaker console (MLflow -> your app ->\n",
+    "**Open MLflow**) - printed links carry a short-lived token. Your jobs appear in the **Runs** list:\n",
+    "\n",
+    "![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",
+    "\n",
+    "**Watch these first:**\n",
+    "\n",
+    "| Metric | Read it as |\n",
+    "|---|---|\n",
+    "| `rollout/reward/mean`, `rollout/reward/valid_mean` | training reward (valid_mean: over non-zero-advantage groups) - should climb |\n",
+    "| `analysis/zero_adv_groups` | groups with no variance, so no learning signal - high means a saturated or too-coarse reward |\n",
+    "| `rollout/tokens/response_max` | if it pins at `sampling_max_tokens`, rollouts truncate mid-answer |\n",
+    "| `val/reward/mean` | held-out reward every `val_every` steps - should **track** training reward |\n",
+    "\n",
+    "Switch to chart mode and MLflow groups related metrics, such as the `rollout/reward` group\n",
+    "(`max`, `mean`, `min`). Rewards lie in `[-1, 5]`.\n",
+    "\n",
+    "Mean climbing and holding is the model learning the output contract and tool sequence; max reaching\n",
+    "5 means the best rollouts are already perfect; min lifting off the floor (and away from\n",
+    "the `-1` no-final-answer penalty) means even the worst rollouts improve. The min-max spread is what\n",
+    "GRPO learns from. If `val/reward/mean` **tracks** training reward, the model is generalizing; if it\n",
+    "flattens while training reward keeps rising, suspect a reward hack.\n",
+    "\n",
+    "In-training validation is forced to greedy `temperature=0`, so its pass@k is degenerate - use the\n",
+    "standalone evaluator in part 3 for real pass@k."
+   ],
+   "id": "02-13"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 8. Read the trajectories\n",
+    "\n",
+    "Metrics say the run is healthy; **traces** say what the agent actually did. The **Traces** tab lists\n",
+    "one row per rollout - add the `reward` column from **Columns** for fast triage:\n",
+    "\n",
+    "![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",
+    "\n",
+    "![MLflow Traces Columns menu with attributes such as Trace ID, Request, Response, Tokens, and Execution time.](images/mlflow-trace-columns.png)\n",
+    "\n",
+    "Open a trace: **Summary** shows the exact task inputs and each turn's tool call in order;\n",
+    "**Details & Timeline** shows every turn's `thinking`, `tool_calls`, and tool result:\n",
+    "\n",
+    "![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",
+    "\n",
+    "![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",
+    "\n",
+    "Read off a trace: the **tool sequence** (a skipped step or a loop on one tool explains a low\n",
+    "reward), **where the reward lands** (on the root `trajectory` span and the final turn; intermediate\n",
+    "turns carry 0), and **per-turn tokens** (which turn is long when `response_max` is high)."
+   ],
+   "id": "02-14"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 9. Reattach after a kernel restart\n",
+    "\n",
+    "Find your job and reattach - no need to rerun the notebook."
+   ],
+   "id": "02-15"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from itertools import islice\n",
+    "\n",
+    "from sagemaker.train.agent_rft_job import AgentRFTJob\n",
+    "\n",
+    "for j in islice(AgentRFTJob.get_all(sort_by=\"CreationTime\", sort_order=\"Descending\"), 20):\n",
+    "    print(f\"{j.job_status:<12} {j.job_name}\")"
+   ],
+   "id": "02-16"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# job = MultiTurnRLTrainer.attach(job_name=\"<JobName from the list above>\")\n",
+    "# job.refresh(); print(job.job_status, job.get_mlflow_url())\n",
+    "# job.stop()   # stop a running job"
+   ],
+   "id": "02-17"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "When the job is `Completed`, the LoRA adapter is registered as a SageMaker Model Package\n",
+    "(`job.output_model_package_arn`). **Next:** `03_evaluate_and_deploy.ipynb`."
+   ],
+   "id": "02-18"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "pygments_lexer": "ipython3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}