diff --git a/v3-examples/model-customization-examples/serverless/mtrl/01_build_and_deploy_agent.ipynb b/v3-examples/model-customization-examples/serverless/mtrl/01_build_and_deploy_agent.ipynb new file mode 100644 index 0000000000..ced7fa7695 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/01_build_and_deploy_agent.ipynb @@ -0,0 +1,452 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multi-Turn RL with SageMaker AI, part 1: build and deploy the rollout agent\n", + "\n", + "Production LLM agents rarely answer in one shot: they call tools, read results, and adjust course\n", + "over many turns. **Multi-turn reinforcement learning (MTRL)** on Amazon SageMaker AI optimizes the\n", + "whole trajectory, not just the final answer.\n", + "\n", + "This three-notebook series is a complete, working MTRL example. You fine-tune\n", + "`openai-reasoning-gpt-oss-20b` to follow a Standard Operating Procedure (SOP) on the\n", + "[SOP-Bench](https://github.com/amazon-science/SOP-Bench) `aircraft_inspection` task, with a\n", + "rollout agent hosted on **Amazon Bedrock AgentCore**. Every piece - environment, tools, reward,\n", + "agent - is real, runnable code in [`sopbench_agent/`](sopbench_agent/) that you can adapt to your\n", + "own task.\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", + "**In this notebook** you will:\n", + "1. Deploy the prerequisites and install dependencies.\n", + "2. Explore the task, its tools, and a real task instance.\n", + "3. Understand how the agent participates in the training loop and connects to the MTRL service.\n", + "4. Check that the reward is trustworthy - and see the reward hack it guards against.\n", + "5. Deploy the agent to AgentCore.\n", + "\n", + "> **Note:** This is sample code, for non-production usage. Work with your security and legal teams\n", + "> to meet your organizational security, regulatory, and compliance requirements before deployment." + ], + "id": "01-00" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "![SageMaker AI MTRL service overview: your agent RL environment (on AgentCore, EKS, EC2, Fargate, or custom) talks to the SageMaker AI MTRL service through the Job Runtime APIs (Sample, SampleWithResponseStream, CompleteRollout, UpdateReward). The MTRL service runs training and evaluation, driven by the MultiTurnRLTrainer, MultiTurnRLEvaluator, and ModelBuilder SDKs, and emits a Model Package, checkpoints, CloudWatch logs, and MLflow metrics.](images/service-overview.jpg)\n", + "\n", + "SageMaker AI runs the training loop and, each step, asks your **agent** for a group of rollouts.\n", + "You supply three things: the **environment** (the task and its tools), the **reward function**, and\n", + "the **agent** that ties them together. The agent calls the policy model served by the training job\n", + "(`Sample` / `SampleWithResponseStream`), executes tools, and reports the reward\n", + "(`CompleteRollout` / `UpdateReward`) - the `@sagemaker_rft_handler` decorator does the reporting\n", + "for you." + ], + "id": "01-01" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 1. Prerequisites\n", + "\n", + "**Account and region.** Use a sandbox/dev account in a region where SageMaker AI multi-turn RL and\n", + "Amazon Bedrock AgentCore are available (this example was tested in `us-east-1`). Set `AWS_REGION`\n", + "(or your default region) to that region before you start. The training job role is intentionally\n", + "broad *within* its account and region (see the comments in `prerequisites.yaml`).\n", + "\n", + "**Deploy the prerequisites stack** once per account. [`prerequisites.yaml`](prerequisites.yaml)\n", + "creates the MTRL training job role, the AgentCore agent execution role, a private KMS-encrypted S3\n", + "bucket for prompts and output, a managed MLflow app, and two managed policies:\n", + "\n", + "```bash\n", + "aws cloudformation deploy --stack-name mtrl-prerequisites \\\n", + " --template-file prerequisites.yaml --capabilities CAPABILITY_NAMED_IAM\n", + "# Accounts allow 3 MLflow apps by default. To reuse one instead of creating one, add:\n", + "# --parameter-overrides ExistingMlflowAppArn=\n", + "```\n", + "\n", + "**Permissions for whoever runs these notebooks** (for example your SageMaker Studio execution\n", + "role): attach `AmazonSageMakerFullAccess`, `BedrockAgentCoreFullAccess`, and the two managed\n", + "policies the stack outputs as `NotebookPolicyArn` and `AgentCoreToolkitPolicyArn`.\n", + "\n", + "**Python environment.** Python 3.10+ (3.11 recommended). Docker is *not* required:\n", + "`agentcore deploy` builds the image on AWS CodeBuild. Install the dependencies, then **restart the\n", + "kernel**." + ], + "id": "01-02" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "%pip install -q -r sopbench_agent/requirements-dev.txt\n", + "%pip install -q --no-deps -e sopbench_agent" + ], + "id": "01-03" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2. Configuration\n", + "\n", + "Every account-specific value is read from the prerequisites stack outputs - nothing to paste." + ], + "id": "01-04" + }, + { + "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}\")" + ], + "id": "01-05" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3. The environment: SOP-Bench `aircraft_inspection`\n", + "\n", + "The task models a pre-flight airworthiness inspection. The agent receives an SOP - a multi-step\n", + "protocol for verifying an aircraft's mechanical systems, electrical systems, and component\n", + "authenticity - plus input parameters for one aircraft (tail number, component serial numbers,\n", + "battery status, weights, ...). It must **follow the SOP step by step, calling tools to gather\n", + "information**, then emit a structured report with six decision fields.\n", + "\n", + "Why it is a good RL environment:\n", + "\n", + "* **Multi-turn by nature** - 3-7 tool calls per instance; not solvable in one shot.\n", + "* **Many decision fields** - six output fields give natural partial credit and enough variance\n", + " across rollouts for GRPO to learn from.\n", + "* **Conditional logic** - incident and mismatch reporting only apply when something fails.\n", + "* **Cheap and deterministic** - tools are in-memory lookups, so a rollout takes seconds and is\n", + " reproducible. A good RL environment is cheap, reproducible, and representative of production -\n", + " built from the production tools, not pointed at production.\n", + "\n", + "Load one real instance from the installed `amazon-sop-bench` package:" + ], + "id": "01-06" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "import json\n", + "from sopbench_agentcore.samples import load_instance\n", + "\n", + "instance = load_instance(\"aircraft_inspection\", task_id=110)\n", + "print(instance[\"instance_id\"])\n", + "print(\"task_input :\", json.dumps(json.loads(instance[\"task_input\"]), indent=2))\n", + "print(\"ground_truth:\", instance[\"ground_truth\"])" + ], + "id": "01-07" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This is what the trainer sends the agent as `payload[\"prompt\"]` (as a JSON string). The agent only\n", + "*uses* `task_input`; `ground_truth` is held out and only read by the reward after the rollout.\n", + "`nan` marks a field that does not apply to this aircraft - the model must still emit it as `\"NA\"`.\n", + "\n", + "The seven tools the agent can call, straight from the benchmark:" + ], + "id": "01-08" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "from amazon_sop_bench.benchmarks import load_benchmark\n", + "\n", + "benchmark = load_benchmark(\"aircraft_inspection\")\n", + "for spec in benchmark.tools.get_tool_specs():\n", + " s = spec[\"toolSpec\"]\n", + " print(f\"{s['name']:<28} {s['description'][:90]}\")" + ], + "id": "01-09" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4. How the agent works\n", + "\n", + "The agent is a small package in [`sopbench_agent/src/sopbench_agentcore/`](sopbench_agent/src/sopbench_agentcore/):\n", + "\n", + "| File | Role |\n", + "|---|---|\n", + "| `main.py` | AgentCore entrypoint; parses the prompt, runs the loop, scores, returns `{\"reward\": ...}` |\n", + "| `openai_agent.py` | The multi-turn tool-use loop over an OpenAI-compatible client |\n", + "| `tools.py` | Adapts SOP-Bench tools into function-calling specs and an executor |\n", + "| `prompts.py` | System prompt and the output contract |\n", + "| `reward.py` | Scores the final answer field by field (dense, task success, and the training reward) |\n", + "| `config.py` | Environment-driven settings (turn cap, sampling fallbacks, debug flags) |\n", + "\n", + "The entrypoint stacks two decorators - one per platform:\n", + "\n", + "```python\n", + "from bedrock_agentcore import BedrockAgentCoreApp\n", + "from sagemaker.train.rft import sagemaker_rft_handler\n", + "\n", + "app = BedrockAgentCoreApp()\n", + "\n", + "@app.entrypoint # AgentCore: route the HTTP request to this function\n", + "@sagemaker_rft_handler # MTRL: report the returned \"reward\" back to the trainer\n", + "def invoke(payload: dict) -> dict:\n", + " return run_rollout(payload)\n", + "```\n", + "\n", + "**The loop.** Seed the conversation with the SOP (system) and the task (user), then repeat: call\n", + "the policy model; if it returns tool calls, execute them and append the results; if it returns\n", + "plain text, that is the final answer. Stop when the model stops calling tools or hits `max_turns`.\n", + "A bad tool call is fed back to the model as an error string so it can self-correct. The shipped\n", + "loop also resamples empty responses (and fails the rollout if they persist, so the trainer retries it),\n", + "recovers an answer GPT-OSS writes only in its reasoning channel on a turn without tool calls, and treats a `400`/`500` from the\n", + "endpoint as the end of the rollout (scoring what was captured).\n", + "\n", + "**Connecting to the MTRL service.** The agent never hosts a model: the training job serves the\n", + "policy model and injects its endpoint per rollout in `payload[\"metadata\"][\"endpoint\"]` (sampling\n", + "params in `payload[\"inferenceParams\"]`). The agent authenticates with a short-lived token - no\n", + "static keys:\n", + "\n", + "```python\n", + "from sagemaker.core.token_generator import generate_token\n", + "from sagemaker.train.rft import get_inference_headers\n", + "\n", + "client = OpenAI(base_url=endpoint + \"/v1\",\n", + " api_key=generate_token(region=region),\n", + " default_headers=get_inference_headers())\n", + "```\n", + "\n", + "Each `client.chat.completions.create(...)` is a `Sample` call on the SageMaker AI Job Runtime\n", + "(`job-runtime.sagemaker.{region}.api.aws`). The contract is fixed - **accept a prompt, run a\n", + "multi-turn rollout, return a reward** - so the trainer only ever needs the agent's runtime ARN." + ], + "id": "01-10" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 5. The reward - and the reward hack it prevents\n", + "\n", + "`reward.py` scores the rollout's **last tagged answer** against the ground truth, field by field,\n", + "with SOP-Bench's own parser and value comparison, so scores line up with the published benchmark:\n", + "\n", + "* **`dense`** - fraction of fields correct (0-1): partial credit gives a usable gradient early on.\n", + "* **TSR** (task success rate) - 1 only if *every* field is correct: the metric you report.\n", + "* The **training reward** is `dense x 5`, so it ranges from **0** (all wrong) to **5** (perfect);\n", + " a rollout that never emits a tagged answer scores **-1**. Only the last answer counts, so\n", + " repeating an answer earns nothing extra.\n", + "\n", + "**The scar this guards against.** An earlier agent learned to drop the required answer wrapper,\n", + "and a lenient reward parser still accepted it: training reward went **up** while the evaluator\n", + "scored it **down**. Here an untagged answer is simply no answer (-1), and `prompts.py` requires\n", + "**every** expected field (`\"NA\"` when not applicable).\n", + "\n", + "**Matching rules.** Values compare case-insensitively, null-like values (`\"NA\"`, `null`) are\n", + "equivalent, and - as in SOP-Bench - a value that contains the expected one also matches, so\n", + "`\"failed (low battery)\"` counts as `\"failed\"`. That containment rule would also let one answer that\n", + "lists every option (`\"success fail retest\"`) collect credit on every task, so the reward adds one\n", + "guard: a value that matches **more than one** of the field's allowed values never matches.\n", + "Containment is otherwise kept for benchmark parity, which also means a negation such as\n", + "`\"unsuccessful\"` still matches `\"success\"`.\n", + "\n", + "**Tag names matter for GPT-OSS.** The prompt asks for the answer as\n", + "`{...}` - not SOP-Bench's native\n", + "``. gpt-oss confuses a tag containing the word \"final\" with its Harmony `final`\n", + "channel and writes the answer into the reasoning channel, so the endpoint returns an empty message.\n", + "Renaming the tags cut empty responses from 8.7% to 0.2%. `reward.py` maps them back before scoring,\n", + "and still accepts SOP-Bench's own formats.\n", + "\n", + "Score four canned answers against the instance above (`scripts/inspect_reward.py`; no model or AWS\n", + "calls):" + ], + "id": "01-11" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "!cd {AGENT_DIR} && PYTHONPATH=src python scripts/inspect_reward.py" + ], + "id": "01-12" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "* **Case 2** loses credit and fails TSR: every expected field must be emitted.\n", + "* **Case 3** loses that one field: a value that matches several allowed values never counts.\n", + "* **Case 4** has the right values but no answer tags, so the rollout has not answered: **-1**.\n", + "\n", + "Run the offline tests of the agent loop, the reward, and the entrypoint (fake model client, no\n", + "endpoint or AWS calls):" + ], + "id": "01-13" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "!cd {AGENT_DIR} && PYTHONPATH=src python scripts/test_local.py" + ], + "id": "01-14" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 6. Deploy the agent to AgentCore\n", + "\n", + "`agentcore configure` writes a local `.bedrock_agentcore.yaml` pointing at the execution role from\n", + "the prerequisites stack (it has `AmazonSageMakerJobRuntimeAccess`, which the agent needs to call\n", + "`Sample`, `CompleteRollout`, and `UpdateReward`). `agentcore deploy` then builds the arm64 image on\n", + "CodeBuild, pushes it to ECR, and creates the runtime - a few minutes. No local container engine\n", + "is needed; if `configure` stalls at a `Password:` prompt, a locally installed engine (such as Finch)\n", + "is asking for `sudo` while the CLI probes for it - run the command from a terminal instead." + ], + "id": "01-15" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "!cd {AGENT_DIR} && AGENTCORE_SUPPRESS_RECOMMENDATION=1 agentcore configure --non-interactive \\\n", + " --entrypoint src/sopbench_agentcore/main.py --name {AGENT_NAME} \\\n", + " --requirements-file requirements.txt --region {REGION} --deployment-type container \\\n", + " --protocol HTTP --disable-memory --execution-role {AGENTCORE_EXEC_ROLE_ARN}" + ], + "id": "01-16" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "!cd {AGENT_DIR} && AGENT_NAME={AGENT_NAME} ./scripts/deploy.sh" + ], + "id": "01-17" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "import time\n", + "\n", + "for _ in range(60):\n", + " AGENT_RUNTIME_ARN = find_agent_runtime_arn()\n", + " if AGENT_RUNTIME_ARN:\n", + " break\n", + " time.sleep(20)\n", + "print(\"AGENT_RUNTIME_ARN =\", AGENT_RUNTIME_ARN)\n", + "assert AGENT_RUNTIME_ARN, f\"No READY runtime named {AGENT_NAME}; check the deploy output above.\"" + ], + "id": "01-18" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Runtime logs land in CloudWatch under `/aws/bedrock-agentcore/runtimes/-DEFAULT`. To\n", + "debug a rollout, redeploy with `DEBUG=1 ./scripts/deploy.sh` for verbose per-turn logs (large log\n", + "volume - not for full training runs).\n", + "\n", + "#### AgentCore notes\n", + "\n", + "* **arm64-only runtime.** Don't pin `--platform=linux/arm64` in the Dockerfile; CodeBuild builds\n", + " arm64, and `agentcore deploy --local` builds native amd64 for local testing.\n", + "* **2048 MB image cap (non-adjustable).** The image stays under it by using the modular\n", + " `sagemaker-train` + `sagemaker-core` packages and a multi-stage Dockerfile.\n", + "* **Region.** The training job must run in the runtime's region.\n", + "* **Concurrency.** AgentCore has no server-side cap, so the trainer's `rollout_max_concurrency` is\n", + " the *only* throttle on concurrent inference connections (more in notebook 2).\n", + "* **Other hosting.** The same contract works for agents on Amazon EKS, Amazon EC2, AWS Fargate, or\n", + " your own infrastructure, connected through a Lambda forwarder - see\n", + " [Preparing your agent](https://docs.aws.amazon.com/sagemaker/latest/dg/model-customize-mtrl-agent.html).\n", + "\n", + "**Next:** `02_train.ipynb` builds the prompt dataset and launches training against this runtime." + ], + "id": "01-19" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/v3-examples/model-customization-examples/serverless/mtrl/02_train.ipynb b/v3-examples/model-customization-examples/serverless/mtrl/02_train.ipynb new file mode 100644 index 0000000000..8216f27e49 --- /dev/null +++ b/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=\"\")\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 +} diff --git a/v3-examples/model-customization-examples/serverless/mtrl/03_evaluate_and_deploy.ipynb b/v3-examples/model-customization-examples/serverless/mtrl/03_evaluate_and_deploy.ipynb new file mode 100644 index 0000000000..aefc8f7dfb --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/03_evaluate_and_deploy.ipynb @@ -0,0 +1,406 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multi-Turn RL with SageMaker AI, part 3: evaluate, deploy, and clean up\n", + "\n", + "Training reward going up is necessary but not sufficient. You score the base and fine-tuned models\n", + "side by side on held-out prompts with `MultiTurnRLEvaluator`, read what the model actually learned,\n", + "diagnose what still fails, see the deployment options, and tear everything down.\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:** a `Completed` training job from `02_train.ipynb`, and the agent still `READY`." + ], + "id": "03-00" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 1. Configuration" + ], + "id": "03-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", + "\n", + "AGENT_RUNTIME_ARN = find_agent_runtime_arn()\n", + "S3_OUTPUT_PATH = f\"s3://{BUCKET}/eval-output/\"\n", + "assert AGENT_RUNTIME_ARN, f\"No READY runtime named {AGENT_NAME} - evaluation runs rollouts against it.\"" + ], + "id": "03-02" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2. Attach to your training job" + ], + "id": "03-03" + }, + { + "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": "03-04" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "JOB_NAME = \"\" # TODO: paste from the list above\n", + "job = MultiTurnRLTrainer.attach(job_name=JOB_NAME)\n", + "job.refresh()\n", + "print(\"status:\", job.job_status, \"| model package:\", job.output_model_package_arn)" + ], + "id": "03-05" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3. Evaluate base vs. fine-tuned in one run\n", + "\n", + "Pass the completed job as `model`: the evaluator resolves the fine-tuned model package, the agent,\n", + "and the base model from it. `evaluate_base_model=True` scores **both** in one pipeline - the honest\n", + "before/after. It runs as its own job (`JobCategory=AgentRFTEvaluation`) with\n", + "`num_prompts x eval_group_size` full rollouts per model - 23 x 8 x 2 models = 368 here - so expect\n", + "several minutes.\n", + "\n", + "For real pass@k you need sampling temperature > 0 **and** `eval_group_size >= max(pass_k_values)`.\n", + "pass@k counts a rollout as a success when its reward is at least `success_threshold`. The\n", + "default is 1.0, which on this reward would count an answer with only ~20% of fields right, so set\n", + "it to 5.0: a success is then a fully correct answer, the same as TSR.\n", + "The evaluator runs *your* agent and reward on held-out data: it measures generalization, not\n", + "independence from your reward - a wrong reward is wrong on held-out data too." + ], + "id": "03-06" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "from sagemaker.train.evaluate import MultiTurnRLEvaluator\n", + "\n", + "from sopbench_agentcore.reward import FINAL_REWARD_SCALE\n", + "\n", + "comparison = MultiTurnRLEvaluator(\n", + " model=job,\n", + " dataset=VAL_DATASET,\n", + " s3_output_path=S3_OUTPUT_PATH + \"comparison/\",\n", + " mlflow_resource_arn=MLFLOW_APP_ARN,\n", + " role=ROLE_ARN,\n", + " accept_eula=True,\n", + " evaluate_base_model=True, # base + fine-tuned, side by side\n", + ")\n", + "hp = comparison.hyperparameters\n", + "print(hp.get_info())\n", + "hp.temperature = 0.7 # >0 so samples differ\n", + "hp.eval_group_size = 8 # must be >= max(pass_k_values)\n", + "hp.pass_k_values = [1, 2, 4, 8]\n", + "hp.success_threshold = FINAL_REWARD_SCALE # pass@k success = every field correct (5.0)\n", + "\n", + "execution = comparison.evaluate() # non-blocking\n", + "print(\"evaluation started:\", execution.arn)" + ], + "id": "03-07" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "execution.wait() # blocks while the rollouts run" + ], + "id": "03-08" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4. Read the results\n", + "\n", + "Each model is logged as its own MLflow run (`base-model-eval`, `fine-tuned-model-eval`) under the\n", + "experiment `mtrl-eval-openai-reasoning-gpt-oss-20b`. Compare them side by side:\n", + "\n", + "| Metric | What it tells you |\n", + "|---|---|\n", + "| `eval/reward/pass_at_1` | **primary success metric** - probability one sample is fully correct (with `success_threshold=5.0`) |\n", + "| `eval/reward/pass_at_k` | probability at least one of *k* samples is fully correct |\n", + "| `eval/reward/mean` | average reward, in `[-1, 5]` |\n", + "| `eval/reward/zero_frac` | fraction of rollouts scoring exactly 0 - an answer with every field wrong (no answer scores -1) |\n", + "| `eval/turns/mean` | high values suggest looping |\n", + "| `eval/tokens/response_max` | pinned at the cap means truncation |\n", + "\n", + "### 5. What the model actually learned\n", + "\n", + "An example from an earlier run, on the **serial-number-mismatch** failure case (installed serial differs from the\n", + "expected one, weight out of tolerance, low battery):\n", + "\n", + "* **Step 0, no answer (reward -1).** The base model calls the right tools in the right order and\n", + " even detects every problem in its reasoning - but **never emits the final answer block**.\n", + "* **Step 10, reward 5.** Same archetype; now it closes the loop with a complete report:\n", + "\n", + "```json\n", + "{\"aircraft_ready\": \"FALSE\", \"mechanical_inspection_result\": \"fail\",\n", + " \"electrical_inspection_result\": \"retest\", \"component_incident_response\": \"failed\",\n", + " \"component_mismatch_response\": \"failed\", \"cross_check_reporting_response\": \"failed\"}\n", + "```\n", + "\n", + "What changed: it **closes the loop** (emits the report at all - the single biggest lift), it\n", + "**propagates failures** instead of defaulting to `\"NA\"`, and its reasoning is **more directed**. A\n", + "rising reward curve is usually a **format/behavior story before a knowledge story** - the base model\n", + "often knows the answer and loses the reward on the contract. Reproduce it: in the Traces tab, sort\n", + "by `reward` and compare an early low-reward trace with a late high-reward one on the same archetype.\n", + "\n", + "### 6. Diagnose failures\n", + "\n", + "Read failing rollouts, not just the aggregate:\n", + "\n", + "* **`pass_at_1 = 0` + high `eval/turns/mean`** - looping without solving; check tool usage.\n", + "* **`pass_at_1 = 0` + low `eval/tokens/response_mean`** - empty or malformed answers; usually the\n", + " output contract (a missing field, or no answer tags at all).\n", + "* **Training reward far above the held-out score** - the reward parser accepts what the evaluator\n", + " rejects: the reward hack from part 1.\n", + "* **`eval/reward/mean` near -1** - the model is not answering at all (no answer tags); check the\n", + " prompt contract and `max_turns`.\n", + "* **Many failed rollouts** (training logs, CloudWatch) - persistent empty responses, endpoint\n", + " errors before an answer, or malformed prompts fail the rollout; the trainer retries each up to\n", + " `rollout_max_retries`. Check the agent is `READY` and the dataset format.\n", + "\n", + "Fixes usually live in the reward or the agent prompt, not in training longer." + ], + "id": "03-09" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 7. Deploy (reference)\n", + "\n", + "The job's output is a **LoRA adapter** registered as a Model Package. Two paths - shown for\n", + "reference and not run here, because a deployed model bills until you delete it. Both take the\n", + "trainer or job directly and resolve the model package, base model, and serving container for you.\n", + "\n", + "**SageMaker AI real-time endpoint** (needs the full `sagemaker` package - `pip install sagemaker` -\n", + "and a GPU instance; a quota increase may apply):\n", + "\n", + "```python\n", + "import uuid\n", + "\n", + "from sagemaker.serve import ModelBuilder\n", + "\n", + "model_builder = ModelBuilder(model=job, instance_type=\"ml.g6e.48xlarge\")\n", + "model_builder.accept_eula = True # gated base model\n", + "model_builder.build()\n", + "endpoint = model_builder.deploy(\n", + " endpoint_name=f\"mtrl-endpoint-{uuid.uuid4().hex[:8]}\",\n", + " instance_type=\"ml.g6e.48xlarge\",\n", + " initial_instance_count=1,\n", + ")\n", + "```\n", + "\n", + "**Amazon Bedrock Custom Model Import** (managed inference, no endpoint to run):\n", + "\n", + "```python\n", + "from sagemaker.serve.bedrock_model_builder import BedrockModelBuilder\n", + "\n", + "bedrock_builder = BedrockModelBuilder(model=job)\n", + "bedrock_builder.deploy(\n", + " imported_model_name=\"mtrl-finetuned-model\",\n", + " role_arn=\"\",\n", + ")\n", + "```\n", + "\n", + "See [Import a customized model into Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-import-model.html).\n", + "Either way, re-run the evaluator against the deployed model to confirm the numbers hold." + ], + "id": "03-10" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 8. Best practices and next steps\n", + "\n", + "* **Environment:** cheap, reproducible, representative - you will run it thousands of times.\n", + "* **Evaluation first:** decide how you judge success independently of the training reward before\n", + " you train, so you catch reward hacking early.\n", + "* **Reward:** reward the outcome, give partial credit, align the parser with the evaluator.\n", + "* **Turn budget:** cap `max_turns` / `sampling_max_tokens` so trajectories don't truncate.\n", + "* **Monitoring:** reward mean, zero-advantage groups, response length, and whether validation tracks\n", + " training.\n", + "\n", + "**Densify further.** This agent returns one scalar reward per rollout. MTRL also accepts a **list of\n", + "per-turn rewards**, paired with `advantage_method = \"group_based_per_turn\"`, so credit lands on the\n", + "turn that earned it. To try it, change `compute_reward` to return one reward per sampled turn (note\n", + "that in-loop resamples of empty responses are sampled turns too). Add per-turn shaping only when you\n", + "can say what a \"good turn\" is - over-shaping invites new reward hacks.\n", + "\n", + "**Bring your own task:** swap `tools.py`, `prompts.py`, `reward.py`, and `scripts/preprocess.py` for\n", + "your domain; the trainer, evaluator, and agent structure stay the same. See also\n", + "[Best practices for multi-turn RL in Amazon SageMaker AI](https://aws.amazon.com/blogs/machine-learning/best-practices-for-multi-turn-reinforcement-learning-in-amazon-sagemaker-ai/)." + ], + "id": "03-11" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 9. Clean up\n", + "\n", + "Set `CONFIRM_CLEANUP = True` and run the cells to stop this example's training job and evaluation\n", + "(if still running), destroy the agent runtime (and its ECR repository), empty the bucket, and delete the\n", + "prerequisites stack. Also delete any endpoint you deployed in section 7.\n", + "\n", + "> `agentcore destroy` also deletes the execution role you passed to `agentcore configure` - here the\n", + "> stack's `AmazonBedrockAgentCoreSDKRuntime-...` role - and CloudFormation does not notice. To deploy\n", + "> the agent again later, delete and re-create the prerequisites stack first.\n", + "\n", + "> **Permissions.** Besides the policies listed in part 1, cleanup needs `ecr:DeleteRepository`,\n", + "> `s3:ListBucketVersions` / `s3:DeleteObjectVersion` on the bucket (the stack's notebook policy\n", + "> grants these), and - for the stack deletion - CloudFormation and IAM delete permissions. Run the\n", + "> last cell with the identity that created the stack if your notebook role lacks those." + ], + "id": "03-12" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "CONFIRM_CLEANUP = False\n", + "\n", + "if CONFIRM_CLEANUP:\n", + " # Stop only what this example started - never other jobs in the account.\n", + " if \"execution\" in globals():\n", + " execution.refresh()\n", + " if execution.status.overall_status in (\"Starting\", \"Executing\"):\n", + " execution.stop()\n", + " print(\"stopped evaluation\", execution.arn)\n", + " job.refresh()\n", + " if job.job_status == \"InProgress\":\n", + " job.stop()\n", + " print(\"stopped training job\", job.job_name)" + ], + "id": "03-13" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "if CONFIRM_CLEANUP:\n", + " !cd {AGENT_DIR} && agentcore destroy -a {AGENT_NAME} --delete-ecr-repo --force" + ], + "id": "03-14" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "if CONFIRM_CLEANUP:\n", + " sess.resource(\"s3\").Bucket(BUCKET).object_versions.delete() # versioned bucket\n", + " cfn = sess.client(\"cloudformation\")\n", + " cfn.delete_stack(StackName=STACK_NAME)\n", + " cfn.get_waiter(\"stack_delete_complete\").wait(StackName=STACK_NAME)\n", + " print(\"deleted\", STACK_NAME)" + ], + "id": "03-15" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Left in place on purpose: the fine-tuned model package version (delete it in the SageMaker console\n", + "under Model registry if you don't need it), the `SagemakerEvaluation-MTRLEvaluation` pipeline (shared\n", + "by all MTRL evaluations in the account), MLflow runs, and the CloudWatch log groups under\n", + "`/aws/bedrock-agentcore/runtimes/` and `/aws/codebuild/bedrock-agentcore-*`." + ], + "id": "03-16" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/v3-examples/model-customization-examples/serverless/mtrl/README.md b/v3-examples/model-customization-examples/serverless/mtrl/README.md new file mode 100644 index 0000000000..14db6923f0 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/README.md @@ -0,0 +1,84 @@ +# Multi-turn RL with SageMaker AI: an end-to-end SOP-Bench example + +A complete, working multi-turn reinforcement learning (MTRL) example. It fine-tunes +`openai-reasoning-gpt-oss-20b` to follow a multi-step Standard Operating Procedure with tool calls +on the [SOP-Bench](https://github.com/amazon-science/SOP-Bench) `aircraft_inspection` task. The +rollout agent is hosted on Amazon Bedrock AgentCore. Unlike the template-style +`mtrl_finetuning_example_notebook_v3_prod.ipynb`, every piece here (environment, tools, reward, +agent, data) is real code that you can run as-is and then adapt to your own task. + +| # | Notebook | What you do | +|---|---|---| +| 1 | [`01_build_and_deploy_agent.ipynb`](01_build_and_deploy_agent.ipynb) | Explore the task, tools, and reward (including the reward hack it prevents), then deploy the rollout agent to AgentCore | +| 2 | [`02_train.ipynb`](02_train.ipynb) | Build the prompt dataset, launch a GRPO job with `MultiTurnRLTrainer`, and read the metrics and trajectories in MLflow | +| 3 | [`03_evaluate_and_deploy.ipynb`](03_evaluate_and_deploy.ipynb) | Compare base and fine-tuned models with `MultiTurnRLEvaluator`, diagnose failures, deploy, and clean up | + +## Quick start + +1. **Use a sandbox/dev account** in a region with SageMaker AI multi-turn RL and Amazon Bedrock + AgentCore (tested in `us-east-1`; set `AWS_REGION`). Deploy the prerequisites stack once. It creates + the job role, the agent execution role, an S3 bucket, an MLflow app, and two managed policies: + ```bash + aws cloudformation deploy --stack-name mtrl-prerequisites \ + --template-file prerequisites.yaml --capabilities CAPABILITY_NAMED_IAM + ``` + Accounts allow 3 MLflow apps by default. To reuse an existing app, add + `--parameter-overrides ExistingMlflowAppArn=`. +2. **Grant permissions to the identity that runs the notebooks.** Attach `AmazonSageMakerFullAccess`, + `BedrockAgentCoreFullAccess`, and the stack outputs `NotebookPolicyArn` and + `AgentCoreToolkitPolicyArn`. +3. **Run the notebooks in order** from this folder, using Python 3.10+ (3.11 recommended). + Notebook 1 installs the dependencies. Docker is not needed, because the agent image builds on + AWS CodeBuild. + +The default run is 10 training steps, which takes well under an hour (about 3-4 minutes per step). +Lower `max_steps` in notebook 2 for a quicker pass. The longer 30-step recipe in +`sopbench_agent/scripts/run_rft_training.py` takes about 2 h 15 min. + +## What's in this folder + +``` +prerequisites.yaml CloudFormation: IAM roles, S3 bucket, MLflow app, notebook policies +sopbench_agent/ the rollout agent (an installable package + AgentCore Dockerfile) + src/sopbench_agentcore/ main.py (entrypoint), openai_agent.py (agent loop), tools.py, + prompts.py, reward.py, config.py, samples.py + scripts/ deploy.sh, preprocess.py, inspect_reward.py, test_local.py, + run_rft_training.py (non-interactive training), run_manual_eval.py +images/ figures used by the notebooks +``` + +## Adapting it to your task + +Keep the trainer, the evaluator, and the agent structure. Replace these task-specific parts: +* `tools.py`: the environment and its tools; +* `prompts.py`: the output contract; +* `reward.py`: return `{"reward": }` (or `{"status": "error"}` to fail the rollout); +* `scripts/preprocess.py`: one `prompt` column, which the service forwards verbatim to the agent. + +This example is built for SOP-Bench's multi-field `aircraft_inspection` benchmark. Moving to another +benchmark means updating `SUPPORTED_BENCHMARKS` and `BENCHMARK_NAME` (`config.py`, `Dockerfile`), +`EXCLUDED_FIELDS` (`samples.py`), the example answer in `prompts.py`, and the dataset paths in +`scripts/preprocess.py`, the notebooks, and `scripts/run_rft_training.py`. + +Agents hosted on Amazon EKS, Amazon EC2, AWS Fargate, or your own infrastructure use the same +contract through a Lambda forwarder. See +[Preparing your agent](https://docs.aws.amazon.com/sagemaker/latest/dg/model-customize-mtrl-agent.html). + +## Cost and cleanup + +You pay for the following until you delete them: +* the training and evaluation jobs; +* AgentCore invocations; +* the CodeBuild image build; +* S3 storage; +* the MLflow app; +* any endpoint you deploy. + +Section 9 of notebook 3 tears everything down. + +## Third-party data + +The agent installs `amazon-sop-bench` from +[github.com/amazon-science/SOP-Bench](https://github.com/amazon-science/SOP-Bench), which is +licensed under CC-BY-NC-4.0, and uses its benchmark data at runtime. No SOP-Bench data is stored in +this folder. diff --git a/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-all-runs.png b/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-all-runs.png new file mode 100644 index 0000000000..c5b004bc36 Binary files /dev/null and b/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-all-runs.png differ diff --git a/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-trace-columns.png b/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-trace-columns.png new file mode 100644 index 0000000000..56bd685642 Binary files /dev/null and b/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-trace-columns.png differ diff --git a/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-trace-details.png b/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-trace-details.png new file mode 100644 index 0000000000..9efd2825e5 Binary files /dev/null and b/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-trace-details.png differ diff --git a/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-trace-landing-page.png b/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-trace-landing-page.png new file mode 100644 index 0000000000..40fe582f05 Binary files /dev/null and b/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-trace-landing-page.png differ diff --git a/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-traces-all.png b/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-traces-all.png new file mode 100644 index 0000000000..c46988cead Binary files /dev/null and b/v3-examples/model-customization-examples/serverless/mtrl/images/mlflow-traces-all.png differ diff --git a/v3-examples/model-customization-examples/serverless/mtrl/images/service-overview.jpg b/v3-examples/model-customization-examples/serverless/mtrl/images/service-overview.jpg new file mode 100644 index 0000000000..03a0ae877f Binary files /dev/null and b/v3-examples/model-customization-examples/serverless/mtrl/images/service-overview.jpg differ diff --git a/v3-examples/model-customization-examples/serverless/mtrl/prerequisites.yaml b/v3-examples/model-customization-examples/serverless/mtrl/prerequisites.yaml new file mode 100644 index 0000000000..c04fb4e4b1 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/prerequisites.yaml @@ -0,0 +1,466 @@ +AWSTemplateFormatVersion: '2010-09-09' +Description: > + Prerequisites for the SageMaker AI multi-turn RL (MTRL) SOP-Bench example. Creates the + MTRL training job role, the Amazon Bedrock AgentCore agent execution role, a private S3 + bucket for prompts and job output, a managed MLflow app, and two managed policies to attach + to the identity that runs the notebooks. Deploy in a sandbox/dev account, in a region with + SageMaker AI multi-turn RL and Amazon Bedrock AgentCore. + +Parameters: + ExistingMlflowAppArn: + Type: String + Default: "" + Description: > + Optional. ARN of an existing SageMaker MLflow app to log to. Leave empty to create one + (accounts have a default quota of 3 MLflow apps). + +Conditions: + CreateMlflowApp: !Equals [!Ref ExistingMlflowAppArn, ""] + +Resources: + + # Training prompts + job output. Private, encrypted, versioned. + PromptsBucket: + Type: AWS::S3::Bucket + # checkov:skip=CKV_AWS_18: Sample-scoped bucket (private, KMS-encrypted, versioned, all public access blocked) for training prompts and job output. Server access logging would require a second log bucket; enable it if your account policy requires it. + Properties: + BucketName: !Sub "sm-mtrl-${AWS::AccountId}-${AWS::Region}" + PublicAccessBlockConfiguration: + BlockPublicAcls: true + BlockPublicPolicy: true + IgnorePublicAcls: true + RestrictPublicBuckets: true + BucketEncryption: + ServerSideEncryptionConfiguration: + - ServerSideEncryptionByDefault: + SSEAlgorithm: aws:kms + VersioningConfiguration: + Status: Enabled + + # SageMaker AI assumes this role during an AgentRFT job. At run time the job needs + # to: invoke the AgentCore runtime, read prompts + write output/checkpoints to S3, + # DECRYPT the KMS-encrypted prompts bucket, access the gated model package, write + # MLflow traces, and emit CloudWatch logs. This is exactly the AWS-published, + # purpose-built AmazonSageMakerJobFullAccess policy (S3/KMS/model-package/MLflow/ + # AgentCore/logs, all same-account-conditioned). It replaces the older + # AmazonSageMakerFullAccess + hand-rolled inline, which lacked KMS (encrypted-bucket + # decrypt) and model-package access -- both would fail a real job at run time. + # Ref: docs.aws.amazon.com/aws-managed-policy/latest/reference/AmazonSageMakerJobFullAccess.html + SageMakerFineTuningJobRole: + Type: AWS::IAM::Role + # checkov:skip=CKV_AWS_111: TrainingWorkloadSupplemental grants only ec2 Describe*/network-interface, ecr GetAuthorizationToken/image-pull, and cloudwatch:PutMetricData -- standard SageMaker training actions without resource-level support, which must be Resource:"*". Narrowing them makes the MTRL SDK's pre-flight role validation fail. + # checkov:skip=CKV_AWS_109: The only IAM action granted is iam:PassRole (EvaluationPipelinePassRole), scoped to passing THIS role only and conditioned to the SageMaker services (iam:PassedToService), so the evaluation pipeline can run its evaluation job under this role. No permissions-management actions are granted. + # checkov:skip=CKV_AWS_63: EvaluationPipelineSageMakerFull uses sagemaker:* bounded to arn:aws:sagemaker:::* because the MultiTurnRLEvaluator pipeline calls a broad server-side set of SageMaker APIs (CreateAction/AddTags/CreateJob/...) not covered by AmazonSageMakerJobFullAccess. Deploy in a dedicated sandbox account; narrow it for shared accounts. + Properties: + RoleName: !Sub "SageMakerFineTuningJobRole-${AWS::Region}" + AssumeRolePolicyDocument: + Version: '2012-10-17' + Statement: + # The SageMaker AI MTRL (AgentRFT) service assumes this job role as principal + # 'job.sagemaker.amazonaws.com' at CreateJob time (verified: without it, + # CreateJob fails "Could not assume role ... allows principal + # 'job.sagemaker.amazonaws.com'"). We keep the classic 'sagemaker.amazonaws.com' + # too so any non-AgentRFT SageMaker path that assumes the role still works. + - Effect: Allow + Principal: + Service: + - job.sagemaker.amazonaws.com + - sagemaker.amazonaws.com + Action: sts:AssumeRole + ManagedPolicyArns: + - arn:aws:iam::aws:policy/AmazonSageMakerJobFullAccess + # Supplement to JobFullAccess. The MTRL SDK's IamRoleResolver validates this + # role via context-free iam:SimulatePrincipalPolicy against the generic + # 'training' action set. JobFullAccess condition-gates its ec2 statements on + # aws:ResourceAccount and omits ecr pull + cloudwatch:PutMetricData entirely, + # so the context-free simulation reports all 14 as missing -> RoleValidationError. + # These must be granted UNCONDITIONALLY: any condition needs request context the + # validator does not supply, which re-triggers the false-missing (verified via + # simulate-principal-policy). This is the AWS-standard SageMaker training action + # set (ec2 VPC networking, ecr image pull, cloudwatch metrics); ec2:Describe* and + # ecr:GetAuthorizationToken are *-only by definition. + Policies: + - PolicyName: TrainingWorkloadSupplemental + PolicyDocument: + Version: '2012-10-17' + Statement: + - Sid: Ec2VpcNetworking + Effect: Allow + Action: + - ec2:CreateNetworkInterface + - ec2:CreateNetworkInterfacePermission + - ec2:DeleteNetworkInterface + - ec2:DeleteNetworkInterfacePermission + - ec2:DescribeNetworkInterfaces + - ec2:DescribeVpcs + - ec2:DescribeSubnets + - ec2:DescribeSecurityGroups + - ec2:DescribeDhcpOptions + Resource: "*" + - Sid: EcrImagePull + Effect: Allow + Action: + - ecr:GetAuthorizationToken + - ecr:BatchCheckLayerAvailability + - ecr:BatchGetImage + - ecr:GetDownloadUrlForLayer + Resource: "*" + - Sid: CloudWatchMetrics + Effect: Allow + Action: cloudwatch:PutMetricData + Resource: "*" + # MultiTurnRLEvaluator runs as a SageMaker Pipeline (JobCategory=AgentRFTEvaluation) + # under THIS role. That pipeline creates its own AgentRFTEvaluation job and records + # lineage (Actions/Artifacts/Contexts + associations), calling a broad set of + # SageMaker APIs (CreateAction, AddTags, CreateJob, ...) that the training-scoped + # AmazonSageMakerJobFullAccess does not grant -- each surfaces as a separate + # "not authorized to perform: sagemaker:" one run at a time. Rather than + # enumerate the full server-side call set, grant sagemaker:* scoped to this + # account+region only. Deploy this stack in a dedicated sandbox account. + - Sid: EvaluationPipelineSageMakerFull + Effect: Allow + Action: sagemaker:* + Resource: !Sub "arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:*" + # The eval pipeline's CreateJob launches the AgentRFTEvaluation job to RUN UNDER + # this same role, which requires iam:PassRole on it -- otherwise CreateJob fails + # "not authorized to perform: iam:PassRole". Scoped to passing THIS role only, and + # conditioned to SageMaker services so the role can't be handed to anything else. + - Sid: EvaluationPipelinePassRole + Effect: Allow + Action: iam:PassRole + Resource: !Sub "arn:aws:iam::${AWS::AccountId}:role/SageMakerFineTuningJobRole-${AWS::Region}" + Condition: + StringEquals: + iam:PassedToService: + - sagemaker.amazonaws.com + - job.sagemaker.amazonaws.com + + # The AgentCore-hosted agent runs under this role: it calls the RFT Runtime + # for inference + reward reporting (managed policy) and pulls its image / logs. + # Point `agentcore configure --execution-role` at this role's ARN. + AgentCoreExecRole: + Type: AWS::IAM::Role + Properties: + RoleName: !Sub "AmazonBedrockAgentCoreSDKRuntime-${AWS::Region}-sopbench" + AssumeRolePolicyDocument: + Version: '2012-10-17' + Statement: + - Effect: Allow + Principal: + Service: bedrock-agentcore.amazonaws.com + Action: sts:AssumeRole + Condition: + StringEquals: + aws:SourceAccount: !Ref "AWS::AccountId" + ArnLike: + aws:SourceArn: !Sub "arn:aws:bedrock-agentcore:${AWS::Region}:${AWS::AccountId}:*" + ManagedPolicyArns: + - arn:aws:iam::aws:policy/AmazonSageMakerJobRuntimeAccess + Policies: + - PolicyName: AgentRuntimeBaseline + PolicyDocument: + Version: '2012-10-17' + Statement: + - Sid: EcrPull + Effect: Allow + Action: + - ecr:BatchCheckLayerAvailability + - ecr:GetDownloadUrlForLayer + - ecr:BatchGetImage + Resource: !Sub "arn:aws:ecr:${AWS::Region}:${AWS::AccountId}:repository/*" + - Sid: EcrAuthToken + Effect: Allow + Action: ecr:GetAuthorizationToken + Resource: "*" + - Sid: Logs + Effect: Allow + Action: + - logs:CreateLogGroup + - logs:CreateLogStream + - logs:PutLogEvents + Resource: !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/bedrock-agentcore/*" + + # MLflow assumes this role to read/write the tracking artifact store in S3. + # Trust is service-scoped to SageMaker; permissions are scoped to the mlflow/ + # prefix of the prompts bucket only. + MlflowServiceRole: + Type: AWS::IAM::Role + Condition: CreateMlflowApp + Properties: + RoleName: !Sub "SageMakerMlflowRole-${AWS::Region}-sopbench" + AssumeRolePolicyDocument: + Version: '2012-10-17' + Statement: + - Effect: Allow + Principal: + Service: sagemaker.amazonaws.com + Action: sts:AssumeRole + Policies: + - PolicyName: MlflowArtifactStoreAccess + PolicyDocument: + Version: '2012-10-17' + Statement: + - Sid: ArtifactStoreReadWrite + Effect: Allow + Action: + - s3:GetObject + - s3:PutObject + - s3:DeleteObject + - s3:ListBucket + Resource: + - !GetAtt PromptsBucket.Arn + - !Sub "${PromptsBucket.Arn}/mlflow/*" + + # Managed MLflow tracking server. Pass its ARN to + # MultiTurnRLTrainer(mlflow_app_arn=...) to log reward + token curves. + MlflowApp: + Type: AWS::SageMaker::MlflowApp + Condition: CreateMlflowApp + Properties: + Name: !Sub "sm-mtrl-${AWS::AccountId}" + ArtifactStoreUri: !Sub "s3://${PromptsBucket}/mlflow" + RoleArn: !GetAtt MlflowServiceRole.Arn + ModelRegistrationMode: AutoModelRegistrationDisabled + + # Attach these two managed policies (plus the AWS managed AmazonSageMakerFullAccess and + # BedrockAgentCoreFullAccess) to the IAM identity that runs the notebooks, e.g. your + # SageMaker Studio execution role. They cover: launching MTRL jobs, deploying the AgentCore + # agent (CodeBuild + ECR), the prompts bucket, MLflow, and passing the roles above. + MtrlNotebookPolicy: + Type: AWS::IAM::ManagedPolicy + # checkov:skip=CKV_AWS_111: The Resource:"*" write statements are account/global-level actions with no resource-level support -- ecr:GetAuthorizationToken, the CloudWatch Logs delivery APIs (AgentCoreObservabilityDelivery), cloudformation:ListStacks/DescribeStacks (notebook stack-output lookup), and codebuild:ListProjects. IAM permission-management actions are scoped to *BedrockAgentCore* verbatim from the AgentCore starter-toolkit docs. Verified narrowing breaks `agentcore deploy` and the training notebook. + Properties: + ManagedPolicyName: !Sub "sm-mtrl-notebook-${AWS::Region}" + Description: Runtime permissions for the SageMaker AI MTRL SOP-Bench example notebooks. + PolicyDocument: + Version: '2012-10-17' + Statement: + - Sid: PromptsBucket + Effect: Allow + Action: + - s3:GetObject + - s3:PutObject + - s3:ListBucket + - s3:GetBucketLocation + # Cleanup (notebook 3): empty the versioned bucket before deleting the stack. + - s3:ListBucketVersions + - s3:DeleteObject + - s3:DeleteObjectVersion + Resource: + - !GetAtt PromptsBucket.Arn + - !Sub "${PromptsBucket.Arn}/*" + # Cleanup (notebook 3): `agentcore destroy --delete-ecr-repo` deletes the agent's ECR + # repository. + - Sid: AgentCoreEcrCleanup + Effect: Allow + Action: ecr:DeleteRepository + Resource: !Sub "arn:aws:ecr:${AWS::Region}:${AWS::AccountId}:repository/bedrock-agentcore-*" + # MTRL SDK (MultiTurnRLTrainer.__init__) reads the gated model's + # fine-tuning-options JSON from the public JumpStart metadata bucket + # at trainer-construction time. AmazonSageMakerFullAccess only covers + # *sagemaker* buckets, so JumpStart (no "sagemaker" substring) needs an + # explicit GetObject grant. Gated model *artifacts* live in the private + # JumpStart bucket and are pulled by the JOB role at run time via the + # model subscription -- not IAM-grantable here. + - Sid: JumpStartModelOptionsRead + Effect: Allow + Action: s3:GetObject + Resource: !Sub "arn:aws:s3:::jumpstart-cache-prod-${AWS::Region}/*" + # bedrock-agentcore:* is granted via the BedrockAgentCoreFullAccess + # managed policy above (covers runtime, runtime-endpoint, workload-identity, + # memory, etc.) -- do not re-enumerate control-plane actions here. + - Sid: EcrPush + Effect: Allow + Action: + - ecr:CreateRepository + - ecr:DescribeRepositories + - ecr:BatchCheckLayerAvailability + - ecr:GetDownloadUrlForLayer + - ecr:BatchGetImage + - ecr:InitiateLayerUpload + - ecr:UploadLayerPart + - ecr:CompleteLayerUpload + - ecr:PutImage + Resource: !Sub "arn:aws:ecr:${AWS::Region}:${AWS::AccountId}:repository/*" + - Sid: EcrAuthToken + Effect: Allow + Action: ecr:GetAuthorizationToken + Resource: "*" + - Sid: CodeBuildForArm64AgentBuild + Effect: Allow + Action: + - codebuild:CreateProject + - codebuild:UpdateProject + - codebuild:StartBuild + - codebuild:BatchGetBuilds + - codebuild:BatchGetProjects + Resource: !Sub "arn:aws:codebuild:${AWS::Region}:${AWS::AccountId}:project/bedrock-agentcore-*" + - Sid: MlflowTracking + Effect: Allow + Action: + - sagemaker-mlflow:* + - sagemaker:CreatePresignedMlflowTrackingServerUrl + - sagemaker:DescribeMlflowTrackingServer + Resource: !Sub "arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:mlflow-app/*" + - Sid: AgentCoreObservabilityDelivery + Effect: Allow + Action: + - logs:PutDeliverySource + - logs:GetDeliverySource + - logs:DeleteDeliverySource + - logs:PutDeliveryDestination + - logs:GetDeliveryDestination + - logs:DeleteDeliveryDestination + - logs:PutDeliveryDestinationPolicy + - logs:CreateDelivery + - logs:GetDelivery + - logs:DeleteDelivery + - logs:DescribeDeliveries + - logs:DescribeDeliverySources + - logs:DescribeDeliveryDestinations + - logs:CreateLogGroup + - logs:PutRetentionPolicy + Resource: "*" + # The notebooks auto-resolve stack outputs (bucket/role/MlflowAppArn). ListStacks/DescribeStacks do not support + # resource-level scoping (must be "*"); read-only. + - Sid: CloudFormationOutputScan + Effect: Allow + Action: + - cloudformation:ListStacks + - cloudformation:DescribeStacks + Resource: "*" + # iam:PassRole is scoped to the three provisioned role ARNs (same-account). + # We intentionally do NOT add an iam:PassedToService condition: the SageMaker + # AgentRFT CreateJob API does not populate that context key (verified via IAM + # policy simulation -- the call implicitDenies when the condition is present), + # so it would block the legitimate trainer.train() call. Confused-deputy + # protection remains enforced by each target role's TRUST POLICY, which only + # allows its intended service principal (sagemaker / bedrock-agentcore) to + # assume it -- passing these roles to any other service is inert. + - Sid: PassProvisionedServiceRoles + Effect: Allow + Action: iam:PassRole + Resource: + - !GetAtt SageMakerFineTuningJobRole.Arn + - !GetAtt AgentCoreExecRole.Arn + - !If [CreateMlflowApp, !GetAtt MlflowServiceRole.Arn, !Ref AWS::NoValue] + # The MTRL SDK's IamRoleResolver pre-flight-verifies the job role's + # permissions via iam:SimulatePrincipalPolicy before submitting. Without + # this the notebook logs a scary (non-fatal) warning and skips the check. + # Read-only action, scoped to the three provisioned roles it may simulate. + - Sid: SimulateProvisionedRoles + Effect: Allow + Action: iam:SimulatePrincipalPolicy + Resource: + - !GetAtt SageMakerFineTuningJobRole.Arn + - !GetAtt AgentCoreExecRole.Arn + - !If [CreateMlflowApp, !GetAtt MlflowServiceRole.Arn, !Ref AWS::NoValue] + + # Verbatim from the AgentCore starter-toolkit docs (runtime-permissions + # #runtime-permissions-starter-toolkit). Required because `agentcore deploy` + # CREATES its own CodeBuild + runtime IAM roles, ECR repo, and source bucket. + AgentCoreStarterToolkitPolicy: + Type: AWS::IAM::ManagedPolicy + Properties: + ManagedPolicyName: !Sub "sm-mtrl-agentcore-toolkit-${AWS::Region}" + Description: Permissions for the AgentCore starter toolkit CLI (agentcore configure/deploy). + PolicyDocument: + Version: '2012-10-17' + Statement: + - Sid: IAMRoleManagement + Effect: Allow + Action: + - iam:CreateRole + - iam:DeleteRole + - iam:GetRole + - iam:PutRolePolicy + - iam:DeleteRolePolicy + - iam:AttachRolePolicy + - iam:DetachRolePolicy + - iam:TagRole + - iam:ListRolePolicies + - iam:ListAttachedRolePolicies + Resource: + - arn:aws:iam::*:role/*BedrockAgentCore* + - Sid: CodeBuildProjectAccess + Effect: Allow + Action: + - codebuild:StartBuild + - codebuild:BatchGetBuilds + - codebuild:ListBuildsForProject + - codebuild:CreateProject + - codebuild:UpdateProject + - codebuild:BatchGetProjects + Resource: + - arn:aws:codebuild:*:*:project/bedrock-agentcore-* + - arn:aws:codebuild:*:*:build/bedrock-agentcore-* + - Sid: CodeBuildListAccess + Effect: Allow + Action: codebuild:ListProjects + Resource: "*" + - Sid: IAMPassRoleAccess + Effect: Allow + Action: iam:PassRole + Resource: + - arn:aws:iam::*:role/AmazonBedrockAgentCore* + - arn:aws:iam::*:role/service-role/AmazonBedrockAgentCore* + - Sid: CloudWatchLogsAccess + Effect: Allow + Action: + - logs:GetLogEvents + - logs:DescribeLogGroups + - logs:DescribeLogStreams + Resource: + - arn:aws:logs:*:*:log-group:/aws/bedrock-agentcore/* + - arn:aws:logs:*:*:log-group:/aws/codebuild/* + - Sid: S3Access + Effect: Allow + Action: + - s3:GetObject + - s3:PutObject + - s3:ListBucket + - s3:CreateBucket + - s3:PutLifecycleConfiguration + Resource: + - arn:aws:s3:::bedrock-agentcore-* + - arn:aws:s3:::bedrock-agentcore-*/* + - Sid: ECRRepositoryAccess + Effect: Allow + Action: + - ecr:CreateRepository + - ecr:DescribeRepositories + - ecr:GetRepositoryPolicy + - ecr:InitiateLayerUpload + - ecr:CompleteLayerUpload + - ecr:PutImage + - ecr:UploadLayerPart + - ecr:BatchCheckLayerAvailability + - ecr:GetDownloadUrlForLayer + - ecr:BatchGetImage + - ecr:ListImages + - ecr:TagResource + Resource: arn:aws:ecr:*:*:repository/bedrock-agentcore-* + - Sid: ECRAuthorizationAccess + Effect: Allow + Action: ecr:GetAuthorizationToken + Resource: "*" + +Outputs: + PromptsBucketName: + Description: S3 bucket for training prompts and job output. + Value: !Ref PromptsBucket + JobRoleArn: + Description: Pass this as the SageMaker AgentRFT job role. + Value: !GetAtt SageMakerFineTuningJobRole.Arn + AgentCoreExecRoleArn: + Description: Use with `agentcore configure --execution-role`. + Value: !GetAtt AgentCoreExecRole.Arn + MlflowAppArn: + Description: Pass this as MultiTurnRLTrainer(mlflow_app_arn=...). + Value: !If [CreateMlflowApp, !GetAtt MlflowApp.Arn, !Ref ExistingMlflowAppArn] + NotebookPolicyArn: + Description: Attach to the identity that runs the notebooks. + Value: !Ref MtrlNotebookPolicy + AgentCoreToolkitPolicyArn: + Description: Attach to the identity that runs `agentcore configure/deploy`. + Value: !Ref AgentCoreStarterToolkitPolicy diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/.dockerignore b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/.dockerignore new file mode 100644 index 0000000000..01251e0b59 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/.dockerignore @@ -0,0 +1,8 @@ +# Only src/, pyproject.toml and requirements.txt are needed to build the agent image. +.venv/ +**/__pycache__/ +*.py[cod] +.bedrock_agentcore/ +.bedrock_agentcore.yaml +data/ +scripts/ diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/.gitignore b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/.gitignore new file mode 100644 index 0000000000..4cf843b59d --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/.gitignore @@ -0,0 +1,7 @@ +# Generated by `agentcore configure` (absolute, machine-specific paths) +.bedrock_agentcore.yaml +.bedrock_agentcore/ +# Prompt parquets written by scripts/preprocess.py +data/ +*.egg-info/ +__pycache__/ diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/Dockerfile b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/Dockerfile new file mode 100644 index 0000000000..11c330bc8f --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/Dockerfile @@ -0,0 +1,66 @@ +# Multi-stage build to fit AgentCore Runtime's hard 2048 MB image cap. +# Stage 1 installs deps (needs git for the amazon-sop-bench git pin); stage 2 is a slim runtime +# that copies only the installed packages + app code, dropping git, apt lists, pip, and caches. +# All deps resolve from public PyPI + the public SOP-Bench git repo. + +# ---- Stage 1: builder ---- +# No hardcoded --platform: the build host sets the arch (arm64 natively in cloud CodeBuild for the +# AgentCore Runtime; amd64 natively for `agentcore deploy --local` on an amd64 dev box). Hardcoding +# --platform=linux/arm64 breaks local builds with "exec format error" on amd64 hosts. +FROM public.ecr.aws/docker/library/python:3.11-slim AS builder + +RUN apt-get update \ + && apt-get install -y --no-install-recommends git \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /app + +# Install runtime deps and the package into an isolated prefix we can copy wholesale. +COPY requirements.txt ./ +RUN pip install --no-cache-dir --prefix=/install -r requirements.txt +COPY pyproject.toml ./ +COPY src/ ./src/ +# --no-deps: the dependencies were installed by the previous layer. +RUN pip install --no-cache-dir --no-deps --prefix=/install . + +# Prune to shrink the image: bytecode caches, test suites, and packaging metadata docs. +# (Kept: .dist-info top-level so imports relying on metadata still work; only bulky bits removed.) +RUN set -eux; \ + find /install -type d -name '__pycache__' -prune -exec rm -rf {} + ; \ + find /install -type d -name 'tests' -prune -exec rm -rf {} + ; \ + find /install -type d -name 'test' -prune -exec rm -rf {} + ; \ + find /install -name '*.pyc' -delete ; \ + find /install -name '*.pyo' -delete ; \ + find /install -name '*.c' -delete ; \ + find /install -name '*.h' -delete ; \ + find /install -name '*.pyi' -delete ; \ + python -m compileall -q --invalidation-mode unchecked-hash /install || true +# Bytecode is precompiled above so cold starts do not recompile every import. + +# ---- Stage 2: slim runtime ---- +FROM public.ecr.aws/docker/library/python:3.11-slim + +WORKDIR /app + +# Copy the installed packages + console scripts from the builder. +COPY --from=builder /install /usr/local +# App code (the AgentCore entrypoint references src/sopbench_agentcore/main.py). +COPY src/ ./src/ + +# Defaults; overridden via `agentcore deploy --env` or .bedrock_agentcore.yaml. +ENV BENCHMARK_NAME=aircraft_inspection \ + MAX_TURNS=10 \ + PYTHONPATH=/app/src \ + PYTHONDONTWRITEBYTECODE=1 + +# Run as an unprivileged user (port 8080 needs no root). +RUN useradd --create-home --uid 1000 agent +USER agent + +EXPOSE 8080 + +# BedrockAgentCoreApp serves GET /ping on 8080. +HEALTHCHECK --interval=30s --timeout=5s --start-period=30s --retries=3 \ + CMD ["python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8080/ping', timeout=4)"] + +CMD ["python", "-m", "sopbench_agentcore.main"] diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/pyproject.toml b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/pyproject.toml new file mode 100644 index 0000000000..748c48e895 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/pyproject.toml @@ -0,0 +1,37 @@ +[build-system] +requires = ["setuptools>=68.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "sopbench-agentcore" +version = "0.1.0" +description = "SOP-Bench rollout agent for multi-turn RL on Amazon SageMaker AI, hosted on Amazon Bedrock AgentCore" +requires-python = ">=3.10" +dependencies = [ + "bedrock-agentcore", + # Modular SageMaker: train (RFT) + core (token_generator) only, not the full sagemaker + # meta-package (keeps the AgentCore image under the 2048 MB cap). + "sagemaker-train", + "sagemaker-core", + "openai", + # public OSS; bundles benchmark data in-package. Pinned to the tested commit (not on PyPI yet). + "amazon-sop-bench @ git+https://github.com/amazon-science/SOP-Bench.git@f6d43858b94cf439854038101bce0eecd3e16eae", +] + +# Dev/launch/notebook tooling -- NOT in the container image (Dockerfile runs `pip install .`, which +# ignores extras). Install with: pip install -e .[dev] (or: uv pip install -r requirements-dev.txt) +[project.optional-dependencies] +dev = [ + # deploy CLI (`agentcore`); distinct from the bedrock-agentcore runtime SDK + "bedrock-agentcore-starter-toolkit", + # launcher CreateJob floor (deployed agent never calls CreateJob, so this is dev-only) + "boto3>=1.43.37", + "botocore>=1.43.37", + # preprocess.py (prompts -> parquet) + "pandas", + "numpy", + "pyarrow", +] + +[tool.setuptools.packages.find] +where = ["src"] diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/requirements-dev.txt b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/requirements-dev.txt new file mode 100644 index 0000000000..e9a550200c --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/requirements-dev.txt @@ -0,0 +1,25 @@ +# Local development, deployment, and notebook tooling -- NOT baked into the runtime Docker image. +# Installs the runtime deps too, so one command sets up a full dev env: +# uv pip install -r requirements-dev.txt +-r requirements.txt + +# ---- Deploy ---- +# AgentCore Starter Toolkit: the `agentcore` CLI (configure/deploy/status/invoke/destroy). Distinct +# from the `bedrock-agentcore` runtime SDK in requirements.txt (which the agent imports). +bedrock-agentcore-starter-toolkit + +# ---- Launch a training job (scripts/run_rft_training.py, 02_train.ipynb) ---- +# The launcher calls SageMaker CreateJob via MultiTurnRLTrainer.train(). Those AgentRFT ops +# (CreateJob/DescribeJob/StopJob/ListJobs) ship in PUBLIC botocore only as of 1.43.37 (verified; +# 1.43.13 lacked them). Flooring here -- NOT in the container -- because the deployed agent never +# calls CreateJob. This is why AWS_DATA_PATH (bundled service model) is no longer needed. +boto3>=1.43.37 +botocore>=1.43.37 + +# ---- Preprocess prompts (scripts/preprocess.py) ---- +pandas +numpy +pyarrow # parquet engine for pandas.to_parquet + +# Notebook kernel (skip if your Jupyter already provides one). +ipykernel diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/requirements.txt b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/requirements.txt new file mode 100644 index 0000000000..498c0b5197 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/requirements.txt @@ -0,0 +1,25 @@ +# Runtime (container) dependencies -- installed into the AgentCore Docker image (see Dockerfile). +# These are EXACTLY what the deployed agent imports at rollout time. Dev/launch/notebook tooling +# lives in requirements-dev.txt (which installs this file too via `-r requirements.txt`). +# +# Note there is deliberately NO explicit boto3/botocore floor here: the deployed agent only does +# inference + bearer-token reward report-back, never SageMaker CreateJob, so any transitive botocore +# is fine. The CreateJob >=1.43.37 floor is a launcher concern and lives in requirements-dev.txt. + +# amazon-sop-bench is public OSS (github.com/amazon-science/SOP-Bench, CC-BY-NC-4.0). It bundles its +# benchmark data in-package (benchmarks/data/**), so a plain pip install includes the data -- +# you do NOT need to clone the repo. Pinned to the commit this example was tested +# against, for reproducible builds (not on PyPI yet). +amazon-sop-bench @ git+https://github.com/amazon-science/SOP-Bench.git@f6d43858b94cf439854038101bce0eecd3e16eae + +# AgentCore runtime SDK: BedrockAgentCoreApp + @sagemaker_rft_handler (the deployed entrypoint). +bedrock-agentcore + +# Modular SageMaker SDK: only the train (RFT adapters/decorators) + core (token_generator) modules, +# NOT the full `sagemaker` meta-package (which drags in serve/mlops/jumpstart and balloons the image +# past AgentCore's 2048 MB cap). Imports: sagemaker.train.rft, sagemaker.core.token_generator. +sagemaker-train +sagemaker-core + +# OpenAI-compatible client for the agent loop. +openai diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/deploy.sh b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/deploy.sh new file mode 100755 index 0000000000..a2fba71959 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/deploy.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# Deploy the SOP-Bench rollout agent to Amazon Bedrock AgentCore. +# +# Usage (from sopbench_agent/, after `agentcore configure`): +# ./scripts/deploy.sh +# DEBUG=1 ./scripts/deploy.sh # verbose per-turn logs (debugging only: large log volume) +# AGENT_NAME=my_agent ./scripts/deploy.sh +# +# `agentcore deploy` builds the arm64 image from the Dockerfile on AWS CodeBuild (no local Docker +# needed), pushes it to ECR, and creates or updates the runtime. It uses the execution role you +# passed to `agentcore configure --execution-role`. +set -euo pipefail + +AGENT_NAME="${AGENT_NAME:-sopbench_agentcore}" # must match `agentcore configure --name` +DEBUG="${DEBUG:-0}" + +cd "$(dirname "$0")/.." + +echo "Deploying agent '$AGENT_NAME' (DEBUG=$DEBUG) ..." +# --auto-update-on-conflict updates an existing runtime of the same name instead of failing. +AGENTCORE_SUPPRESS_RECOMMENDATION=1 agentcore deploy \ + -a "$AGENT_NAME" \ + --auto-update-on-conflict \ + --env "DEBUG=$DEBUG" + +echo +echo "Next: 'agentcore status' to confirm the runtime is READY." +echo "Tear down: 'agentcore destroy -a $AGENT_NAME --force'." diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/inspect_reward.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/inspect_reward.py new file mode 100644 index 0000000000..caf7730976 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/inspect_reward.py @@ -0,0 +1,49 @@ +"""Inspect the reward on canned answers - no model endpoint needed. + +Scores a few hand-written answers with the same reward.compute_reward used in training, against +one real SOP-Bench instance, so you can see how dense, TSR, and the training reward respond to +the output contract. The live, model-in-the-loop evaluation is 03_evaluate_and_deploy.ipynb. + +Usage: + PYTHONPATH=src python3 scripts/inspect_reward.py +""" + +import json +import sys + + +def tagged(answer: dict) -> str: + return "" + json.dumps(answer) + "" + + +def main() -> int: + from sopbench_agentcore.reward import compute_reward + from sopbench_agentcore.samples import load_instance + + gt = load_instance()["ground_truth"] # aircraft_inspection_110 + not_applicable = [k for k, v in gt.items() if isinstance(v, float) and v != v] + correct = {k: ("NA" if k in not_applicable else v) for k, v in gt.items()} + applicable = {k: v for k, v in correct.items() if k not in not_applicable} + first = next(k for k in applicable if isinstance(gt[k], str)) + + cases = [ + ("1. every field, NA for not-applicable", tagged(correct)), + ("2. not-applicable fields omitted", tagged(applicable)), + ("3. hedged value in one field", tagged(dict(correct, **{first: "success or failed"}))), + ("4. correct fields, no answer tags", json.dumps(correct)), + ] + print(f"Ground truth: {len(gt)} fields; not applicable: {not_applicable}\n") + print(f"{'case':<42} {'dense':>6} {'TSR':>4} {'reward':>7}") + for label, answer in cases: + r = compute_reward([{"role": "assistant", "content": answer}], gt) + print(f"{label:<42} {r['dense']:>6.2f} {r['binary']:>4.0f} {r['reward']:>7.2f}") + print( + '\n* Case 2: every field must be emitted ("NA" when not applicable) or it scores a miss.\n' + "* Case 3: a value that matches several of the field's allowed values never counts.\n" + "* Case 4: without the answer tags the rollout has not answered, so it scores -1." + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/preprocess.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/preprocess.py new file mode 100644 index 0000000000..b9dc468524 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/preprocess.py @@ -0,0 +1,186 @@ +"""Preprocess SOP-Bench tasks into RFT prompt parquets for an MTRL training job. + +Emits one parquet row per task with a single `prompt` column. The MTRL RFT runtime forwards +that `prompt` string verbatim to the deployed AgentCore agent, which deserializes it in +main._parse_instance (json.loads), reading benchmark / instance_id / task_input / ground_truth. +So this schema is the contract between the training data and the rollout agent: do NOT switch +to the verl chat/reward_spec schema. + +Strict-JSON guarantees (must hold or the rollout breaks): + - Upstream tasks carry float NaN for inapplicable ground-truth fields. NaN/Inf are converted + to JSON null and json.dumps uses allow_nan=False, so a bare `NaN` literal can never reach + the parquet (it would parse only because every Python hop is lenient, then fail elsewhere). + The reward comparator treats null and NaN as equivalent, so scoring is unchanged. + - `cross_check_response` is excluded: the SOP output template does not list it, so the model + cannot know to emit it. Keeping it would pin TSR at 0 for those tasks. + +Deterministic 80/20 train/val split (seed 42). + +Usage: + # Local only (writes ./data/aircraft/aircraft_inspection_{train,val}.parquet) + PYTHONPATH=src python3 scripts/preprocess.py + + # Also upload to S3 in the layout the launch notebook expects + # /train/aircraft_inspection_train.parquet + # /validation/aircraft_inspection_val.parquet + PYTHONPATH=src python3 scripts/preprocess.py \ + --s3-prefix s3:///sopbench-mtrl/aircraft + +Requires amazon_sop_bench (in requirements.txt). S3 upload uses boto3 (pulled in transitively). +""" + +import argparse +import json +import math +import os +from urllib.parse import urlparse + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +from amazon_sop_bench.benchmarks import load_benchmark + +from sopbench_agentcore.config import SUPPORTED_BENCHMARKS +from sopbench_agentcore.samples import EXCLUDED_FIELDS + + +def _clean_nan(obj): + """Recursively replace float NaN/Inf with None so the result serializes to valid JSON.""" + if isinstance(obj, dict): + return {k: _clean_nan(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_clean_nan(v) for v in obj] + if isinstance(obj, float): + return None if (math.isnan(obj) or math.isinf(obj)) else obj + if isinstance(obj, np.floating): + f = float(obj) + return None if (math.isnan(f) or math.isinf(f)) else f + if isinstance(obj, np.integer): + return int(obj) + return obj + + +def _dumps_strict(value): + """json.dumps with NaN cleaned first and allow_nan=False as a hard guard (fails loudly).""" + return json.dumps(_clean_nan(value), allow_nan=False, default=str) + + +def split_tasks(tasks, seed=42, train_ratio=0.8): + """Deterministic train/val split.""" + rng = np.random.default_rng(seed) + indices = rng.permutation(len(tasks)) + n_train = int(len(tasks) * train_ratio) + return indices[:n_train], indices[n_train:] + + +def serialize_task(benchmark_name, task_id, task): + """Serialize one task to the {"prompt": } row format. None if invalid.""" + expected = task.expected_output + if not isinstance(expected, dict): + return None # this example supports multi-field benchmarks only + ground_truth = {k: v for k, v in expected.items() if k not in EXCLUDED_FIELDS} + if not ground_truth: + return None + + instance = { + "data_source": "sopbench", + "instance_id": f"{benchmark_name}_{task_id}", + "benchmark": benchmark_name, + "task_input": _dumps_strict(task.inputs), + "ground_truth": _clean_nan(ground_truth), + } + return {"prompt": _dumps_strict(instance)} + + +def write_prompts_parquet(rows, out_path): + """Write the `prompt` column as Arrow `string`. + + The MTRL service rejects `large_string`, which pandas >= 3.0 writes by default for str columns, + so the type is pinned explicitly instead of going through DataFrame.to_parquet. + """ + table = pa.table({"prompt": pa.array([r["prompt"] for r in rows], type=pa.string())}) + pq.write_table(table, out_path) + assert pq.read_schema(out_path).field("prompt").type == pa.string(), out_path + + +def _upload_to_s3(local_path, s3_uri): + """Upload a single file to an s3://bucket/key URI using boto3.""" + import boto3 + + parsed = urlparse(s3_uri) + bucket, key = parsed.netloc, parsed.path.lstrip("/") + boto3.client("s3").upload_file(local_path, bucket, key) + print(f" uploaded {local_path} -> {s3_uri}") + + +def main(): + parser = argparse.ArgumentParser( + description="Preprocess SOP-Bench to RFT prompt parquet (strict JSON, NaN->null)" + ) + parser.add_argument("--output-dir", default="./data/aircraft", help="Local output directory") + parser.add_argument( + "--benchmarks", + nargs="*", + default=list(SUPPORTED_BENCHMARKS), + choices=SUPPORTED_BENCHMARKS, + help=f"Benchmarks to process (supported: {', '.join(SUPPORTED_BENCHMARKS)})", + ) + parser.add_argument("--seed", type=int, default=42, help="Random seed (default: 42)") + parser.add_argument( + "--train-ratio", type=float, default=0.8, help="Train fraction (default: 0.8)" + ) + parser.add_argument( + "--s3-prefix", + default=None, + help="Optional s3://bucket/prefix; uploads /train/_train.parquet and " + "/validation/_val.parquet (the layout the launch notebook expects)", + ) + args = parser.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + + for name in sorted(args.benchmarks): + benchmark = load_benchmark(name) + tasks = list(benchmark.tasks) + train_idx, val_idx = split_tasks(tasks, seed=args.seed, train_ratio=args.train_ratio) + + written = {} + for split_name, indices, suffix in [ + ("train", train_idx, "train"), + ("val", val_idx, "val"), + ]: + rows = [] + for i in indices: + row = serialize_task(name, tasks[i].task_id, tasks[i]) + if row is not None: + rows.append(row) + if not rows: + continue + out_path = os.path.join(args.output_dir, f"{name}_{suffix}.parquet") + write_prompts_parquet(rows, out_path) + written[split_name] = (out_path, len(rows)) + + train_n = written.get("train", ("", 0))[1] + val_n = written.get("val", ("", 0))[1] + filtered = len(tasks) - train_n - val_n + print( + f" {name}: {len(tasks)} tasks -> train={train_n} val={val_n}" + + (f" (filtered {filtered} invalid)" if filtered else "") + ) + + if args.s3_prefix: + prefix = args.s3_prefix.rstrip("/") + # train -> train/ ; val -> validation/ (matches the 02_train.ipynb dataset layout). + s3_subdir = {"train": "train", "val": "validation"} + for split_name, (local_path, _) in written.items(): + fname = os.path.basename(local_path) + _upload_to_s3(local_path, f"{prefix}/{s3_subdir[split_name]}/{fname}") + + print( + f"\nDone. Local parquets in {args.output_dir}/" + + (f"; uploaded under {args.s3_prefix.rstrip('/')}/" if args.s3_prefix else "") + ) + + +if __name__ == "__main__": + main() diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/run_manual_eval.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/run_manual_eval.py new file mode 100644 index 0000000000..772bac548f --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/run_manual_eval.py @@ -0,0 +1,174 @@ +"""Manual evaluation runner: score SOPBench rollouts against YOUR vLLM endpoint. + +Runs the real openai_agent loop against an OpenAI-compatible vLLM server (no AgentCore, no RFT +token path) and scores each instance with the same reward.compute_reward used in training. Use +this to sanity-check the agent + reward + prompt before deploying to AgentCore. + +USAGE +----- + PYTHONPATH=src python3 scripts/run_manual_eval.py \ + --base-url http://:8080/v1 \ + --model \ + --max-turns 10 + + # Multi-instance: point --data at a parquet or jsonl of instances + # (same shape as sopbench_agentcore.samples.load_instance(): instance_id, benchmark, task_input, + # ground_truth). + PYTHONPATH=src python3 scripts/run_manual_eval.py \ + --base-url http://:8080/v1 --model \ + --data data/aircraft/aircraft_inspection_val.parquet \ + --num-samples 23 + +NOTES +----- +- api_key is sent as "EMPTY" (vLLM ignores it). Nothing is logged that contains credentials. +- Model name: pass the bare vLLM name (no "openai/" prefix). +- Scores reported: dense (per-field partial credit), binary/TSR (all fields correct), and the + training reward (dense * 5, or -1 when the rollout never answers). +""" + +import argparse +import json +import sys + + +def load_instances(path: str, limit: int): + """Load instances from .json (single or list), .jsonl, or .parquet. Returns list of dicts with + keys: instance_id, benchmark, task_input (str or dict), ground_truth (dict).""" + instances = [] + if path.endswith(".parquet"): + import pandas as pd + + df = pd.read_parquet(path) + rows = df.to_dict(orient="records") + for row in rows: + # A row may itself be the instance, or carry a JSON 'prompt' column. + if "prompt" in row and isinstance(row["prompt"], str): + try: + inst = json.loads(row["prompt"]) + except json.JSONDecodeError: + inst = row + else: + inst = row + instances.append(inst) + elif path.endswith(".jsonl"): + with open(path) as f: + for line in f: + line = line.strip() + if line: + instances.append(json.loads(line)) + else: # .json: single object or a list + with open(path) as f: + data = json.load(f) + instances = data if isinstance(data, list) else [data] + if limit and limit > 0: + instances = instances[:limit] + return instances + + +def main(): + ap = argparse.ArgumentParser(description="Manual SOPBench eval against a vLLM endpoint") + ap.add_argument("--base-url", required=True, help="vLLM OpenAI-compatible base URL (.../v1)") + ap.add_argument( + "--model", required=True, help="Model name as served by vLLM (no openai/ prefix)" + ) + ap.add_argument( + "--data", + default=None, + help="Instances: .json / .jsonl / .parquet (default: one built-in SOP-Bench instance)", + ) + ap.add_argument("--num-samples", type=int, default=0, help="Limit instances (0 = all)") + ap.add_argument("--max-turns", type=int, default=10) + ap.add_argument("--temperature", type=float, default=1.0) + ap.add_argument("--max-tokens", type=int, default=4096) + ap.add_argument("--top-p", type=float, default=1.0) + ap.add_argument( + "--debug", + action="store_true", + help="Print raw per-turn response (content, reasoning_content, tool_calls)", + ) + args = ap.parse_args() + + model = args.model + + from openai import OpenAI + from amazon_sop_bench.benchmarks import load_benchmark + from sopbench_agentcore.openai_agent import run_openai_agent + from sopbench_agentcore.tools import create_openai_tools, make_tool_executor + from sopbench_agentcore.reward import compute_reward + + client = OpenAI(base_url=args.base_url, api_key="EMPTY", max_retries=2, timeout=180.0) + + if args.data: + instances = load_instances(args.data, args.num_samples) + else: + from sopbench_agentcore.samples import load_instance + + instances = [load_instance()] + print(f"Loaded {len(instances)} instance(s) from {args.data}") + print(f"Endpoint: {args.base_url} Model: {model}\n") + + # Cache loaded benchmarks (load_benchmark can be slow). + bench_cache = {} + + dense_scores, binary_scores = [], [] + for i, inst in enumerate(instances, 1): + bname = inst.get("benchmark", "aircraft_inspection") + task_input_raw = inst.get("task_input", "{}") + task_input = ( + json.loads(task_input_raw) if isinstance(task_input_raw, str) else task_input_raw + ) + ground_truth = inst.get("ground_truth", {}) or {} + instance_id = inst.get("instance_id", f"inst_{i}") + + if bname not in bench_cache: + bench_cache[bname] = load_benchmark(bname) + benchmark = bench_cache[bname] + + tools = create_openai_tools(benchmark) + tool_executor = make_tool_executor(benchmark) + try: + final_output, messages = run_openai_agent( + task_input, + benchmark.sop_text, + tools, + tool_executor, + client, + model, + max_turns=args.max_turns, + temperature=args.temperature, + max_tokens=args.max_tokens, + top_p=args.top_p, + debug=args.debug, + ) + scored = compute_reward(messages, ground_truth) + if args.debug: + print(f" [final_output] {final_output!r}") + last_asst = next( + (m for m in reversed(messages) if m.get("role") == "assistant"), None + ) + print(f" [last assistant msg] {(last_asst or {}).get('content', '')!r}") + print(f" [ground_truth] {ground_truth}") + dense_scores.append(scored["dense"]) + binary_scores.append(scored["binary"]) + n_turns = sum(1 for m in messages if m.get("role") == "assistant") + print( + f"[{i:>3}/{len(instances)}] {instance_id:<28} " + f"dense={scored['dense']:.3f} TSR={scored['binary']:.0f} " + f"turns={n_turns} reward={scored['reward']}" + ) + except Exception as e: + print(f"[{i:>3}/{len(instances)}] {instance_id:<28} ERROR: {e}") + dense_scores.append(0.0) + binary_scores.append(0.0) + + n = len(dense_scores) or 1 + print("\n=== Aggregate ===") + print(f"instances: {len(dense_scores)}") + print(f"mean dense: {sum(dense_scores) / n:.4f}") + print(f"mean TSR (bin): {sum(binary_scores) / n:.4f}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/run_rft_training.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/run_rft_training.py new file mode 100644 index 0000000000..cfbac5ffac --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/run_rft_training.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +"""Launch a full SOP-Bench aircraft_inspection GPT-OSS run against your AgentCore-hosted agent. + +Non-interactive equivalent of ../02_train.ipynb with a longer recipe. Account-specific values are +resolved at runtime: + + * prompts bucket, job role, and MLflow app ARN from the `mtrl-prerequisites` stack outputs + (override with --stack-name or pass the values explicitly); + * the AgentCore runtime ARN from the READY runtime named AGENT_NAME (default + `sopbench_agentcore`), or --agent-runtime-arn. + +Exits non-zero if the job does not complete successfully. + +Run from sopbench_agent/ with PYTHONPATH=src. Prereqs: the prerequisites stack, the deployed +agent (./scripts/deploy.sh), the uploaded prompts +(scripts/preprocess.py --s3-prefix ...), and AWS credentials for the region (AWS_REGION). +""" + +import argparse +import os +import sys + +import boto3 +from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer + +from sopbench_agentcore.reward import FINAL_REWARD_SCALE + +DEFAULT_MODEL = "openai-reasoning-gpt-oss-20b" +DEFAULT_STACK_NAME = "mtrl-prerequisites" +AGENT_NAME = os.environ.get("AGENT_NAME", "sopbench_agentcore") # `agentcore configure --name` +PREFIX = "sopbench-mtrl/aircraft" # must match scripts/preprocess.py --s3-prefix + +HYPERPARAMETERS = { + "global_batch_size": 64, + "group_size": 4, + "advantage_method": "group_based", # GRPO + "loss_fn": "ppo", + "clip_low_threshold": 0.8, + "clip_high_threshold": 1.2, + "temperature": 1.0, + "sampling_top_p": 1.0, + "sampling_max_tokens": 8192, + "max_epochs": 15, + "max_steps": 30, + "val_every": 10, + # pass@k counts a rollout as successful when reward >= success_threshold. With rewards in + # [-1, 5], only a fully correct answer (5.0) is a task success. + "success_threshold": FINAL_REWARD_SCALE, + "lora_rank": 32, + "lora_alpha": 64, + "learning_rate": 4e-5, + # AgentCore has no server-side concurrency cap, so this is the only throttle on concurrent + # inference connections. Lower it toward 32 if you see httpx ConnectTimeout errors. + "rollout_max_concurrency": 64, + "rollout_timeout": 600, + "rollout_max_retries": 3, +} + + +def _stack_outputs(session, stack_name): + try: + stack = session.client("cloudformation").describe_stacks(StackName=stack_name)["Stacks"][0] + except Exception as e: # noqa: BLE001 + print(f"(stack '{stack_name}' not found: {e})") + return {} + return {o["OutputKey"]: o["OutputValue"] for o in stack.get("Outputs", [])} + + +def _find_agent_runtime_arn(session): + """ARN of the READY runtime named AGENT_NAME (exact match; the SDK call paginates).""" + for r in MultiTurnRLTrainer.list_bedrock_agentcore_runtimes(session=session): + if r["name"] == AGENT_NAME and r["status"] == "READY": + return r["arn"] + return None + + +def main(): + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--model", default=DEFAULT_MODEL) + ap.add_argument("--stack-name", default=DEFAULT_STACK_NAME) + ap.add_argument("--agent-runtime-arn", default=None, help="default: auto-discover") + ap.add_argument("--bucket", default=None, help="default: stack output PromptsBucketName") + ap.add_argument("--role-arn", default=None, help="default: stack output JobRoleArn") + ap.add_argument("--mlflow-app-arn", default=None, help="default: stack output MlflowAppArn") + ap.add_argument("--attach", metavar="JOB_NAME", default=None, help="monitor an existing job") + ap.add_argument("--timeout", type=int, default=6 * 3600, help="wait timeout in seconds") + args = ap.parse_args() + + session = boto3.session.Session() + if not session.region_name: + raise SystemExit("Set AWS_REGION (or a default region) to the region of your agent.") + + if args.attach: + job = MultiTurnRLTrainer.attach(job_name=args.attach) + else: + if args.model not in MultiTurnRLTrainer.list_supported_models(): + raise SystemExit(f"{args.model} is not a supported MTRL model") + outs = _stack_outputs(session, args.stack_name) + bucket = args.bucket or outs.get("PromptsBucketName") + cfg = { + "agent_env": args.agent_runtime_arn or _find_agent_runtime_arn(session), + "role": args.role_arn or outs.get("JobRoleArn"), + "mlflow_app_arn": args.mlflow_app_arn or outs.get("MlflowAppArn"), + "bucket": bucket, + } + missing = [k for k, v in cfg.items() if not v] + if missing: + raise SystemExit(f"Could not resolve {missing}; pass them explicitly (see --help).") + trainer = MultiTurnRLTrainer( + model=args.model, + agent_env=cfg["agent_env"], + training_dataset=f"s3://{bucket}/{PREFIX}/train/aircraft_inspection_train.parquet", + validation_dataset=f"s3://{bucket}/{PREFIX}/validation/aircraft_inspection_val.parquet", + mlflow_app_arn=cfg["mlflow_app_arn"], + s3_output_path=f"s3://{bucket}/output/", + role=cfg["role"], + accept_eula=True, + ) + hp = trainer.hyperparameters + for key, value in HYPERPARAMETERS.items(): + setattr(hp, key, value) + job = trainer.train(wait=False) + print(f"Submitted {job.job_name}") + + job.wait(poll=60, timeout=args.timeout) + job.refresh() + print(f"Final status: {job.job_status}") + print(f"Output model package: {job.output_model_package_arn}") + return 0 if job.job_status in ("Completed", "Succeeded") else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/test_invoke_local.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/test_invoke_local.py new file mode 100644 index 0000000000..9a1bf68a63 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/test_invoke_local.py @@ -0,0 +1,60 @@ +"""Run the rollout entrypoint (main.run_rollout) end to end against your own model server. + +Exercises everything the deployed agent does except the container and the SageMaker AI runtime: +payload parsing, benchmark loading, the agent loop (through the LOCAL_TEST client), and scoring. +The payload has the same shape the runtime sends: the prompt is a JSON string, encoded exactly +as scripts/preprocess.py writes it (not-applicable ground-truth values as JSON null). + +Run: + LOCAL_TEST=1 LOCAL_ENDPOINT=http://:8080/v1 MODEL_NAME= \ + PYTHONPATH=src python3 scripts/test_invoke_local.py [--debug] +""" + +import json +import math +import os +import sys + + +def _nan_to_null(value): + if isinstance(value, dict): + return {k: _nan_to_null(v) for k, v in value.items()} + if isinstance(value, float) and math.isnan(value): + return None + return value + + +def main(): + if "--debug" in sys.argv: + os.environ["DEBUG"] = "1" # before importing main, so AgentConfig picks it up + if not all(os.environ.get(k) for k in ("LOCAL_TEST", "LOCAL_ENDPOINT", "MODEL_NAME")): + print("Set LOCAL_TEST=1, LOCAL_ENDPOINT, and MODEL_NAME to run against your server.") + return 1 + + from sopbench_agentcore.main import run_rollout + from sopbench_agentcore.samples import load_instance + + instance = _nan_to_null(load_instance()) + payload = { + "prompt": json.dumps(instance, allow_nan=False), + "inferenceParams": { + "model": os.environ.get("MODEL_NAME", ""), + "temperature": float(os.environ.get("TEMPERATURE", "1.0")), + "topP": float(os.environ.get("TOP_P", "1.0")), + "maxTokens": int(os.environ.get("MAX_TOKENS", "4096")), + }, + "metadata": {}, + } + print(f"endpoint={os.environ['LOCAL_ENDPOINT']} model={payload['inferenceParams']['model']}") + result = run_rollout(payload) + print(json.dumps(result, indent=2, default=str)) + + if result.get("status") != "success": + print(f"FAIL rollout failed: {result.get('error')}") + return 3 + print(f"PASS reward={result['reward']} dense={result['dense']} TSR={result['binary']}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/test_local.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/test_local.py new file mode 100644 index 0000000000..64dfbb5c4e --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/scripts/test_local.py @@ -0,0 +1,415 @@ +"""Offline tests for the rollout agent: no deployment, endpoint, or AWS calls. + +Run: PYTHONPATH=src python3 scripts/test_local.py + +1. Agent loop (no extra dependencies): a scripted fake client drives run_openai_agent. +2. Reward and rollout entrypoint: need amazon_sop_bench (and bedrock_agentcore / sagemaker-train + for the entrypoint); each group reports SKIP if its dependencies are not installed. +Exits non-zero if any test fails. +""" + +import contextlib +import json +import os +import sys + +FAILED = [] + + +# -------------------------------------------------------------------------------------------- +# Fakes +# -------------------------------------------------------------------------------------------- +class _Fn: + def __init__(self, name, arguments): + self.name, self.arguments = name, arguments + + +class _ToolCall: + def __init__(self, id, name, arguments): + self.id, self.type, self.function = id, "function", _Fn(name, arguments) + + +class _Message: + def __init__(self, content="", tool_calls=None, reasoning=None): + self.content, self.tool_calls, self.reasoning_content = content, tool_calls, reasoning + + +class _Choice: + def __init__(self, message): + self.message = message + self.finish_reason = "tool_calls" if message.tool_calls else "stop" + + +class _Response: + def __init__(self, message): + self.choices, self.id, self.usage = [_Choice(message)], "fake", None + + +class _APIError(Exception): + def __init__(self, status_code): + super().__init__(f"HTTP {status_code}") + self.status_code = status_code + + +class ScriptedClient: + """Returns the scripted messages in order; an int in the script raises that HTTP status.""" + + def __init__(self, script): + self.script, self.calls = list(script), 0 + self.chat = self.completions = self + + def create(self, **kwargs): + self.calls += 1 + item = self.script.pop(0) + if isinstance(item, int): + raise _APIError(item) + return _Response(item) + + +TOOLS = [ + { + "type": "function", + "function": {"name": "lookup", "description": "", "parameters": {"type": "object"}}, + } +] +ANSWER = '{"status": "pass"}' + + +def _run_agent(script, executor=None, max_turns=5): + from sopbench_agentcore.openai_agent import run_openai_agent + + calls = [] + + def default_executor(name, args): + calls.append((name, args)) + return '{"record": "found"}' + + client = ScriptedClient(script) + final, messages = run_openai_agent( + {"aircraft": "N1"}, + "Inspect the aircraft.", + TOOLS, + executor or default_executor, + client, + "test-model", + max_turns=max_turns, + ) + return final, messages, calls, client + + +@contextlib.contextmanager +def _env(**values): + """Set env vars for the block, then restore the previous values.""" + old = {k: os.environ.get(k) for k in values} + os.environ.update({k: v for k, v in values.items() if v is not None}) + try: + yield + finally: + for k, v in old.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + + +def check(name, fn): + try: + fn() + print(f"PASS {name}") + except Exception as e: # noqa: BLE001 - report and keep going + FAILED.append(name) + print(f"FAIL {name}: {type(e).__name__}: {e}") + + +# -------------------------------------------------------------------------------------------- +# 1. Agent loop +# -------------------------------------------------------------------------------------------- +def agent_tests(): + tool_turn = _Message("checking", [_ToolCall("c1", "lookup", '{"id": "N1"}')]) + + def loop(): + final, messages, calls, _ = _run_agent([tool_turn, _Message(ANSWER)]) + assert calls == [("lookup", {"id": "N1"})], calls + assert final == ANSWER, final + assert [m["role"] for m in messages] == ["system", "user", "assistant", "tool", "assistant"] + + def reasoning_carried_but_not_scored_on_tool_turns(): + thinking = "thinking: the example shows " + ANSWER + leak = _Message("", [_ToolCall("c1", "lookup", "{}")], reasoning=thinking) + final, messages, _, _ = _run_agent([leak, _Message("", None, reasoning=ANSWER)]) + assert messages[2]["content"] == thinking, "the model must see its own reasoning" + assert final == ANSWER, "tool-free turn reasoning is the answer" + + def bad_tool_json_fed_back(): + bad = _Message("", [_ToolCall("c1", "lookup", '{"id": "N1"')]) + _, messages, calls, _ = _run_agent([bad, _Message(ANSWER)]) + assert calls == [], "tool must not run with invalid arguments" + assert "not a valid JSON object" in messages[3]["content"], messages[3] + + def api_error_fails_rollout(): + tool_turn = _Message("checking", [_ToolCall("c1", "lookup", "{}")]) + for status in (400, 500): + try: + _run_agent([tool_turn, status]) + except _APIError: + continue + raise AssertionError(f"a {status} must fail the rollout (so it is retried)") + + def empty_responses_resampled_then_fail(): + from sopbench_agentcore.openai_agent import ( + EMPTY_RESPONSE_MAX_RETRIES, + EmptyInferenceResponse, + ) + + final, _, _, client = _run_agent([_Message(""), _Message(ANSWER)]) + assert final == ANSWER and client.calls == 2 + try: + _run_agent([_Message("")] * (EMPTY_RESPONSE_MAX_RETRIES + 1)) + except EmptyInferenceResponse: + return + raise AssertionError("persistent empty responses must raise EmptyInferenceResponse") + + def harmony_split_keeps_leading_answer(): + from sopbench_agentcore.openai_agent import _split_harmony_channels + + content, _ = _split_harmony_channels(ANSWER + "<|channel|>commentary") + assert ANSWER in content, content + content, _ = _split_harmony_channels("assistant<|channel|>final<|message|>" + ANSWER) + assert content == ANSWER, content # a bare role prefix is header noise + content, _ = _split_harmony_channels("assistant<|channel|>commentary to=functions.x") + assert content == "", content # header-only: empty, so it is resampled + + for name, fn in [ + ("agent loop: tool call then answer", loop), + ( + "agent loop: reasoning is carried in the transcript, scored only on tool-free turns", + reasoning_carried_but_not_scored_on_tool_turns, + ), + ("agent loop: invalid tool JSON is fed back, tool not run", bad_tool_json_fed_back), + ("agent loop: endpoint errors (400/500) fail the rollout", api_error_fails_rollout), + ("agent loop: empty responses resampled, then fail", empty_responses_resampled_then_fail), + ( + "agent loop: Harmony split keeps content, drops header-only noise", + harmony_split_keeps_leading_answer, + ), + ]: + check(name, fn) + + +# -------------------------------------------------------------------------------------------- +# 2. Reward (needs amazon_sop_bench) +# -------------------------------------------------------------------------------------------- +def reward_tests(): + try: + from sopbench_agentcore.reward import FINAL_REWARD_SCALE, NO_ANSWER_REWARD, compute_reward + from sopbench_agentcore.samples import load_instance + except ImportError as e: + print(f"SKIP reward tests (amazon_sop_bench not importable: {e})") + return + + gt = load_instance()["ground_truth"] # aircraft_inspection_110: 6 fields, 2 not applicable + nan = [k for k, v in gt.items() if isinstance(v, float) and v != v] + full = {k: ("NA" if k in nan else v) for k, v in gt.items()} + + def tag(d): + return "" + json.dumps(d) + "" + + def score(answer, messages=None): + msgs = messages if messages is not None else [{"role": "assistant", "content": answer}] + return compute_reward(msgs, gt) + + def contract_correct(): + r = score(tag(full)) + assert (r["reward"], r["dense"], r["binary"]) == (FINAL_REWARD_SCALE, 1.0, 1.0), r + + def omitted_na_fields(): + r = score(tag({k: v for k, v in full.items() if k not in nan})) + assert r["binary"] == 0.0 and abs(r["dense"] - 4 / 6) < 1e-9, r + + def null_counts_as_na(): + r = score(tag({k: (None if k in nan else v) for k, v in full.items()})) + assert r["binary"] == 1.0, r + + def no_wrapper_is_no_answer(): + r = score(json.dumps(full)) + assert r["reward"] == NO_ANSWER_REWARD and r["answered"] is False, r + + def repeated_answer_credited_once(): + msgs = [{"role": "assistant", "content": tag(full)}] * 2 + assert score(tag(full), msgs)["reward"] == FINAL_REWARD_SCALE + + def last_answer_wins(): + wrong = {k: ("NA" if k in nan else "wrong") for k in gt} + msgs = [ + {"role": "assistant", "content": tag(full)}, + {"role": "assistant", "content": tag(wrong)}, + ] + assert score(tag(wrong), msgs)["binary"] == 0.0 + + def tool_call_turn_not_scored(): + call = [{"id": "c1", "type": "function", "function": {"name": "x", "arguments": "{}"}}] + msgs = [{"role": "assistant", "content": tag(full), "tool_calls": call}] + assert score(None, msgs)["reward"] == NO_ANSWER_REWARD + + def hedged_value_not_credited(): + field = next(k for k in gt if k not in nan and isinstance(gt[k], str)) + hedged = dict(full, **{field: "success or failed or retest"}) + r = score(tag(hedged)) + assert r["binary"] == 0.0 and abs(r["dense"] - 5 / 6) < 1e-9, r + + def partial_single_field_credited(): + field = next(k for k in gt if k not in nan) + r = score(tag({field: full[field]})) + assert abs(r["dense"] - 1 / 6) < 1e-9, r + + def universal_hedge_does_not_pay(): + for value, ready in [ + ("success fail retest failed", "true false"), + ("success-fail-retest", "e"), + ("success (or fail)", "true/false"), + ("successfailretest", "truefalse"), + ]: + hack = { + k: ("NA" if k in nan else ready if k == "aircraft_ready" else value) for k in gt + } + r = score(tag(hack)) + # At most the not-applicable fields plus a coincidental single-value match: far below + # the 5.0 a correct answer earns, and never a task success. + assert r["binary"] == 0.0 and r["reward"] <= FINAL_REWARD_SCALE / 2, (value, r) + + def verbose_single_value_credited(): + verbose = { + k: (v + " (all checks passed)" if isinstance(v, str) and k not in nan else v) + for k, v in full.items() + } + assert score(tag(verbose))["binary"] == 1.0 + + def sop_final_response_format(): + body = ", ".join( + f"{k!r}: {None if k in nan else (str(v).upper() if isinstance(v, bool) else v)!r}" + for k, v in gt.items() + ) + r = score("{" + body + "}") + assert (r["reward"], r["binary"]) == (FINAL_REWARD_SCALE, 1.0), r + + def tags_case_insensitive(): + body = json.dumps(full) + r = score("" + body + "") + assert r["binary"] == 1.0, r + + for name, fn in [ + ("reward: contract-correct answer scores 5 / dense 1 / TSR 1", contract_correct), + ("reward: omitting not-applicable fields loses credit and TSR", omitted_na_fields), + ("reward: JSON null counts as not applicable", null_counts_as_na), + ("reward: an untagged answer is no answer (-1)", no_wrapper_is_no_answer), + ("reward: a repeated answer is credited once", repeated_answer_credited_once), + ("reward: the last answer is the one scored", last_answer_wins), + ("reward: an answer on a tool-call turn is not scored", tool_call_turn_not_scored), + ("reward: a hedged multi-value answer gets no credit", hedged_value_not_credited), + ("reward: a one-field answer keeps partial credit", partial_single_field_credited), + ("reward: listing every option is not a universal hack", universal_hedge_does_not_pay), + ("reward: a single value with an explanation keeps credit", verbose_single_value_credited), + ("reward: the SOP's own format is scored", sop_final_response_format), + ("reward: answer tags are matched case-insensitively", tags_case_insensitive), + ]: + check(name, fn) + + +# -------------------------------------------------------------------------------------------- +# 3. Rollout entrypoint (needs amazon_sop_bench, bedrock_agentcore, sagemaker-train) +# -------------------------------------------------------------------------------------------- +def entrypoint_tests(): + try: + from sopbench_agentcore import main + from sopbench_agentcore.config import AgentConfig + from sopbench_agentcore.openai_agent import EmptyInferenceResponse + from sopbench_agentcore.samples import load_instance + except ImportError as e: + print(f"SKIP entrypoint tests (dependency not importable: {e})") + return + + payload = {"prompt": json.dumps(load_instance())} + + def with_agent(fn): + def run(): + with _env(LOCAL_TEST="1"): + original = main._run_openai + main._run_openai = fn + try: + return main.run_rollout(payload) + finally: + main._run_openai = original + + return run + + def failures_are_errors(): + def empty(*a, **k): + raise EmptyInferenceResponse("empty") + + def boom(*a, **k): + raise ValueError("boom") + + for agent in (empty, boom): + out = with_agent(agent)() + assert out["status"] == "error" and "reward" not in out, out + + def instance_payload_shape(): + content = "{}" + agent = lambda *a, **k: (content, [{"role": "assistant", "content": content}]) # noqa: E731 + with _env(LOCAL_TEST="1"): + original, main._run_openai = main._run_openai, agent + try: + out = main.run_rollout({"instance": {"prompt": payload["prompt"]}}) + finally: + main._run_openai = original + assert out["status"] == "success", out + + def malformed_payloads_are_errors(): + for bad in ({}, {"prompt": "not json"}, {"prompt": json.dumps({"task_input": {"a": 1}})}): + with _env(LOCAL_TEST="1"): + out = main.run_rollout(bad) + assert out["status"] == "error", (bad, out) + + def scored_rollout(): + answer = load_instance()["ground_truth"] + answer = {k: ("NA" if isinstance(v, float) else v) for k, v in answer.items()} + content = "" + json.dumps(answer) + "" + out = with_agent(lambda *a, **k: (content, [{"role": "assistant", "content": content}]))() + assert out["status"] == "success" and out["reward"] == 5.0, out + + def snake_case_params(): + cfg = AgentConfig() + p = main._inference_params({"inference_params": {"max_tokens": 77, "top_p": 0.5}}, cfg) + assert (p["max_tokens"], p["top_p"]) == (77, 0.5), p + p = main._inference_params({"inferenceParams": {"maxTokens": 88, "topP": 0.4}}, cfg) + assert (p["max_tokens"], p["top_p"]) == (88, 0.4), p + + def config_validation(): + for bad in ({"BENCHMARK_NAME": "content_flagging"}, {"MAX_TURNS": "0"}): + with _env(**bad): + try: + AgentConfig.from_env() + except ValueError: + continue + raise AssertionError(f"{bad} must be rejected") + + for name, fn in [ + ( + "entrypoint: agent failures are reported as failed rollouts, not scored", + failures_are_errors, + ), + ("entrypoint: malformed payloads are reported as errors", malformed_payloads_are_errors), + ("entrypoint: the SDK RolloutRequest instance shape is accepted", instance_payload_shape), + ("entrypoint: a correct answer is scored 5", scored_rollout), + ("entrypoint: sampling params accepted in camelCase and snake_case", snake_case_params), + ("entrypoint: unsupported config is rejected", config_validation), + ]: + check(name, fn) + + +if __name__ == "__main__": + agent_tests() + reward_tests() + entrypoint_tests() + print(f"=== {'FAILED: ' + ', '.join(FAILED) if FAILED else 'all passed'} ===") + sys.exit(1 if FAILED else 0) diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/__init__.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/config.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/config.py new file mode 100644 index 0000000000..036beaa554 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/config.py @@ -0,0 +1,69 @@ +"""Configuration for the SOP-Bench AgentCore rollout agent. + +All fields are environment-driven so the deployed image can be reconfigured at deploy time +(`agentcore deploy --env KEY=VALUE`) without code changes. The SageMaker AI Job Runtime endpoint +is resolved as: RFT_RUNTIME_ENDPOINT -> RFT_RUNTIME_URL -> built from the region. +""" + +import os +from dataclasses import dataclass + +# This example is built and tested for one multi-field SOP-Bench benchmark. +SUPPORTED_BENCHMARKS = ("aircraft_inspection",) + + +def _flag(name: str) -> bool: + return os.environ.get(name, "").lower() in ("1", "true", "yes") + + +@dataclass +class AgentConfig: + model_name: str = "" # model id sent to the endpoint (the training job usually fixes it) + max_turns: int = 10 + benchmark_name: str = "aircraft_inspection" + # Sampling fallbacks, used only when the runtime sends no inference params (it normally does). + # They match the trainer's default hyperparameters: temperature 1.0, sampling_top_p 1.0, + # sampling_max_tokens 4096. + temperature: float = 1.0 + top_p: float = 1.0 + max_tokens: int = 4096 + region: str = "" # AWS_REGION; AgentCore sets it in the runtime container + local_test: bool = False # skip the token path and call LOCAL_ENDPOINT (offline testing) + debug: bool = False # verbose per-turn logging + stream: bool = False # stream chat/completions (STREAM_MODE) + verbose_trace: bool = False # log every request, streaming delta, and response (VERBOSE_TRACE) + + def __post_init__(self): + if self.benchmark_name not in SUPPORTED_BENCHMARKS: + raise ValueError( + f"BENCHMARK_NAME={self.benchmark_name!r} is not supported; " + f"expected one of {SUPPORTED_BENCHMARKS}" + ) + if self.max_turns < 1: + raise ValueError(f"MAX_TURNS must be >= 1, got {self.max_turns}") + + @classmethod + def from_env(cls) -> "AgentConfig": + return cls( + model_name=os.environ.get("MODEL_NAME", ""), + max_turns=int(os.environ.get("MAX_TURNS", "10")), + benchmark_name=os.environ.get("BENCHMARK_NAME", "aircraft_inspection"), + temperature=float(os.environ.get("TEMPERATURE", "1.0")), + top_p=float(os.environ.get("TOP_P", "1.0")), + max_tokens=int(os.environ.get("MAX_TOKENS", "4096")), + region=os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION", ""), + local_test=_flag("LOCAL_TEST"), + debug=_flag("DEBUG"), + stream=_flag("STREAM_MODE"), + verbose_trace=_flag("VERBOSE_TRACE"), + ) + + @property + def rft_runtime_endpoint(self) -> str: + """Resolve the SageMaker AI Job Runtime endpoint.""" + env_url = os.environ.get("RFT_RUNTIME_ENDPOINT") or os.environ.get("RFT_RUNTIME_URL") + if env_url: + return env_url.rstrip("/") + if not self.region: + raise ValueError("AWS_REGION is not set; cannot build the Job Runtime endpoint") + return f"https://job-runtime.sagemaker.{self.region}.api.aws" diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/main.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/main.py new file mode 100644 index 0000000000..c51a7077b7 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/main.py @@ -0,0 +1,213 @@ +"""SOP-Bench rollout agent on Amazon Bedrock AgentCore, wired to SageMaker AI MTRL. + +`invoke` is the AgentCore entrypoint. `@sagemaker_rft_handler` reports the result to the +SageMaker AI Job Runtime: + * {"reward": ...} -> CompleteRollout + UpdateReward + * {"status": "error", ...} or an exception + -> CompleteRollout(status="failed") + UpdateReward(0.0). + The trainer retries failed rollouts up to the `rollout_max_retries` hyperparameter. +""" + +import functools +import json +import logging +import os +import re +import time + +from bedrock_agentcore import BedrockAgentCoreApp +from sagemaker.train.rft import sagemaker_rft_handler + +from .config import AgentConfig +from .reward import compute_reward, field_vocabulary +from .tools import create_openai_tools, make_tool_executor + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +app = BedrockAgentCoreApp() + + +@functools.lru_cache(maxsize=None) +def _load_benchmark(name: str): + """Load a benchmark once per process. SOP-Bench's tools are stateless lookups.""" + from amazon_sop_bench.benchmarks import load_benchmark + + return load_benchmark(name) + + +def _parse_instance(payload: dict) -> dict: + """Return the task instance {benchmark, task_input, ground_truth}; raise if malformed. + + The training prompt is a JSON string (see scripts/preprocess.py). The instance may also be + nested under "instance", either inside the prompt or at the top level of the payload. + """ + raw = payload.get("prompt") + instance = payload.get("instance") + if raw is None and isinstance(instance, dict): + # SDK RolloutRequest shape: the dataset row, e.g. {"prompt": ""}. + raw = instance.get("prompt", instance) + if raw is None: + raise ValueError("payload has no 'prompt'") + prompt = json.loads(raw) if isinstance(raw, str) else raw + if isinstance(prompt, dict) and "task_input" not in prompt and "instance" in prompt: + prompt = prompt["instance"] + if not isinstance(prompt, dict): + raise ValueError(f"prompt must be a JSON object, got {type(prompt).__name__}") + + task_input = prompt.get("task_input") + if isinstance(task_input, str): + task_input = json.loads(task_input) + ground_truth = prompt.get("ground_truth") + if isinstance(ground_truth, str): + ground_truth = json.loads(ground_truth) + if not isinstance(task_input, dict) or not task_input: + raise ValueError("prompt has no 'task_input' object") + if not isinstance(ground_truth, dict) or not ground_truth: + raise ValueError("prompt has no 'ground_truth' object") + return { + "benchmark": prompt.get("benchmark"), + "task_input": task_input, + "ground_truth": ground_truth, + } + + +def _inference_params(payload: dict, config: AgentConfig) -> dict: + """Per-rollout sampling params from the runtime (camelCase or snake_case), else config.""" + ip = payload.get("inferenceParams") or payload.get("inference_params") or {} + + def pick(*keys, default): + for k in keys: + if ip.get(k) not in (None, ""): + return ip[k] + return default + + return { + "model": pick("model", default=config.model_name or "default"), + "temperature": pick("temperature", default=config.temperature), + "top_p": pick("topP", "top_p", default=config.top_p), + "max_tokens": pick("maxTokens", "max_tokens", default=config.max_tokens), + } + + +def _resolve_inference_endpoint(payload: dict, config: AgentConfig) -> str: + """The runtime injects a per-rollout endpoint in metadata.endpoint; prefer it.""" + meta = payload.get("metadata") or {} + return (meta.get("endpoint") or config.rft_runtime_endpoint).rstrip("/") + + +def run_rollout(payload: dict) -> dict: + """Run one rollout and score it. Returns {"reward", ...} or {"status": "error", ...}. + + Decorator-free so it can be tested locally (LOCAL_TEST) without AgentCore or the runtime. + """ + # [rollout-trace] stage markers: the platform fails a trajectory if no sampling call arrives + # within a time window, and these show the last stage reached before a stall. + t0 = time.monotonic() + meta = (payload.get("metadata") if isinstance(payload, dict) else None) or {} + # Same key lookup as the SDK's RolloutFeedbackClient; the runtime sends rolloutId. + trajectory_id = meta.get("trajectory_id") or meta.get("trajectoryId") or meta.get("rolloutId") + logger.info("[rollout-trace] RECEIVED trajectory_id=%s", trajectory_id) + try: + config = AgentConfig.from_env() + instance = _parse_instance(payload) + benchmark_name = instance["benchmark"] or config.benchmark_name + if benchmark_name != config.benchmark_name: + raise ValueError( + f"prompt benchmark {benchmark_name!r} != BENCHMARK_NAME {config.benchmark_name!r}" + ) + params = _inference_params(payload, config) + # LOCAL_TEST calls LOCAL_ENDPOINT instead (see _runtime_client). + endpoint = None if config.local_test else _resolve_inference_endpoint(payload, config) + benchmark = _load_benchmark(benchmark_name) + logger.info( + "[rollout-trace] stage=ready benchmark=%s endpoint=%s t=%.2fs", + benchmark_name, + endpoint, + time.monotonic() - t0, + ) + final_output, messages = _run_openai( + benchmark, instance["task_input"], params, config, endpoint, trajectory_id + ) + result = compute_reward( + messages, instance["ground_truth"], field_vocabulary(benchmark_name) + ) + result["status"] = "success" + if config.debug: + logger.info( + "rollout trajectory_id=%s reward=%s dense=%s binary=%s assistant_turns=%s", + trajectory_id, + result["reward"], + result["dense"], + result["binary"], + sum(1 for m in messages if m.get("role") == "assistant"), + ) + return result + except Exception as e: # noqa: BLE001 - every failure is reported as a failed rollout + logger.exception("rollout failed trajectory_id=%s", trajectory_id) + return {"status": "error", "error": f"{type(e).__name__}: {e}"} + + +def _runtime_client(config: AgentConfig, endpoint: str): + """OpenAI-compatible client for the SageMaker AI Job Runtime (or LOCAL_ENDPOINT).""" + from openai import OpenAI + + if config.local_test: + base_url = os.environ.get("LOCAL_ENDPOINT", "http://localhost:8080/v1") + # Placeholder key: local OpenAI-compatible servers ignore it. + return OpenAI(base_url=base_url, api_key="local") # pragma: allowlist secret + from sagemaker.core.token_generator import generate_token + from sagemaker.train.rft import get_inference_headers + + # Token region: AWS_REGION, else the region in the runtime endpoint host + # (job-runtime.sagemaker..api.aws). + match = re.search(r"\.sagemaker\.([a-z0-9-]+)\.", endpoint) + region = config.region or (match.group(1) if match else "") + if not region: + raise ValueError("cannot determine the AWS region for the Job Runtime token") + t0 = time.monotonic() + logger.info("[rollout-trace] stage=generating_token region=%s", region) + api_key = generate_token(region=region) + logger.info("[rollout-trace] stage=token_generated t=%.2fs", time.monotonic() - t0) + return OpenAI( + base_url=endpoint + "/v1", + api_key=api_key, + default_headers=get_inference_headers(), + max_retries=5, # ride out HTTP 429 throttling under high rollout concurrency + timeout=300.0, + ) + + +def _run_openai(benchmark, task_input, params, config, endpoint, trajectory_id): + from .openai_agent import run_openai_agent + + client = _runtime_client(config, endpoint) + tools = create_openai_tools(benchmark) + logger.info("[rollout-trace] stage=agent_start n_tools=%s", len(tools)) + return run_openai_agent( + task_input, + benchmark.sop_text, + tools, + make_tool_executor(benchmark), + client, + params["model"], + max_turns=config.max_turns, + temperature=params["temperature"], + max_tokens=params["max_tokens"], + top_p=params["top_p"], + debug=config.debug, + trajectory_id=trajectory_id, + stream=config.stream, + verbose_trace=config.verbose_trace, + ) + + +@app.entrypoint +@sagemaker_rft_handler +def invoke(payload: dict) -> dict: + """AgentCore entrypoint; @sagemaker_rft_handler reports the result to the trainer.""" + return run_rollout(payload) + + +if __name__ == "__main__": + app.run() diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/openai_agent.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/openai_agent.py new file mode 100644 index 0000000000..337b4428a2 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/openai_agent.py @@ -0,0 +1,684 @@ +"""Multi-turn tool-use agent over an OpenAI-compatible client (the SageMaker AI Job Runtime). + +The loop: send system (SOP) + user (task), let the model call tools, feed the results back, and +repeat until the model answers without calling a tool or max_turns is reached. + +Returns (final_output, messages). `messages` is the OpenAI-format transcript; reward.py scores +the last assistant turn that contains a tagged answer. +""" + +import json +import logging +from typing import Callable + +from .prompts import build_system_prompt + +logger = logging.getLogger(__name__) + + +class EmptyInferenceResponse(RuntimeError): + """The endpoint kept returning an empty HTTP 200 (no content, reasoning, or tool calls) after + EMPTY_RESPONSE_MAX_RETRIES in-loop resamples. run_rollout reports the rollout as failed, and + the trainer retries it up to `rollout_max_retries`.""" + + +# --- Uniform response wrappers so the agent loop is identical for streaming vs non-streaming --- +class _ToolCallFn: + def __init__(self, name, arguments): + self.name = name + self.arguments = arguments + + +class _ToolCall: + def __init__(self, id, name, arguments): + self.id = id + self.type = "function" + self.function = _ToolCallFn(name, arguments) + + +class _Msg: + def __init__(self, content, reasoning_content, tool_calls): + self.content = content + self.reasoning_content = reasoning_content + self.tool_calls = tool_calls or None + + +class _Usage: + def __init__(self, prompt_tokens, completion_tokens): + self.prompt_tokens = prompt_tokens + self.completion_tokens = completion_tokens + + +class _Resp: + """Duck-typed stand-in for a non-streaming ChatCompletion: exposes .choices[0].message, + .choices[0].finish_reason, .id, .usage — the exact surface the agent loop reads.""" + + class _Choice: + def __init__(self, message, finish_reason): + self.message = message + self.finish_reason = finish_reason + + def __init__(self, message, finish_reason, resp_id, usage): + self.choices = [self._Choice(message, finish_reason)] + self.id = resp_id + self.usage = usage + + +# Harmony control markers (gpt-oss). Format: <|start|>role<|channel|>chan<|message|>body<|end|>. +# In streaming these can arrive raw in deltas; they must never reach message role/content. +_HARMONY_MARKERS = ( + "<|start|>", + "<|end|>", + "<|message|>", + "<|channel|>", + "<|return|>", + "<|call|>", + "<|constrain|>", + "<|endoftext|>", + "<|im_start|>", + "<|im_end|>", +) + + +# A bare role name before the first channel header is Harmony header noise, not content. +_HARMONY_ROLES = {"assistant", "user", "system", "developer", "tool"} + + +def _strip_harmony(text): + """Truncate text at the first Harmony control marker and drop any stray markers, so no + `<|...|>` token or its trailing channel/role annotation survives into role/content. + `assistant<|channel|>final_output` -> `assistant`. Used as the fallback when content has no + parseable <|message|> body (the channel-aware split handles the structured case).""" + if not text: + return text or "" + import re + + # Cut at the first control marker: everything after it is Harmony header noise, not body. + cut = re.split(r"<\|[a-z_]+\|>", text, maxsplit=1)[0] + return cut.strip() + + +def _clean_tool_name(name): + """Strip Harmony channel annotations that can bleed into tool NAMES. + + Two forms observed in gpt-oss streaming responses: + - "VerifyAircraftClearance<|channel|>commentary" (raw Harmony control marker) + - "VerifyElectricalSystems%commentary" (marker decoded to '%'; seen in eng logs) + + Tool names in this codebase are ASCII identifiers ([A-Za-z0-9_]) so truncate at the first + non-identifier character. Handles both known forms without needing a marker-form allowlist. + Returns "" for a name that starts with a non-identifier char (caller falls back to error path). + """ + if not name: + return name or "" + import re + + m = re.match(r"[A-Za-z0-9_]+", name) + return m.group(0) if m else "" + + +def _resolve_tool_name(raw, valid_names): + """Resolve a model-emitted tool name to a real tool, recovering from two corruption modes: + + 1. Channel bleed handled by _clean_tool_name (`Verify...<|channel|>commentary` / + `...%commentary`). + 2. Hallucinated SUFFIX on an otherwise-valid identifier, e.g. `ReportComponentMismatchjson` + (the model concatenated `json`) or `VerifyElectricalSystemstool`. _clean_tool_name keeps + these intact because the suffix is valid identifier text, so the executor 404s and the + trajectory can fail. Recover by matching against the KNOWN tool list. + + Resolution order: + - cleaned == a valid name -> use it (fast path, no change) + - cleaned starts with exactly one valid name -> that name (strip the hallucinated suffix) + - a valid name starts with cleaned (truncation) -> that name + - otherwise -> cleaned (caller surfaces "not found") + + Returns (resolved_name, changed_bool). `changed` is True only when we altered the cleaned name, + so the caller can log the recovery without spamming on the common exact-match path. + """ + cleaned = _clean_tool_name(raw) + valid = set(valid_names or ()) + if not cleaned or cleaned in valid: + return cleaned, False + # Longest valid name that `cleaned` starts with (handles a trailing hallucinated suffix). + prefixes = sorted((v for v in valid if cleaned.startswith(v)), key=len, reverse=True) + if prefixes: + return prefixes[0], True + # Valid name that starts with `cleaned` (a truncated name); require a unique hit. + completions = [v for v in valid if v.startswith(cleaned)] + if len(completions) == 1: + return completions[0], True + return cleaned, False + + +def _split_harmony_channels(text): + """If `text` has raw Harmony channel segments, split into (final_content, analysis_reasoning). + Returns (None, None) when there are no channel markers (normal already-parsed content). + + Harmony body looks like: <|channel|>analysis<|message|>...<|channel|>final<|message|>... + We route the `analysis` channel to reasoning and everything else (final/commentary) to content. + """ + if not text or "<|channel|>" not in text: + return None, None + import re + + final_parts, analysis_parts = [], [] + segments = text.split("<|channel|>") + # Text before the first channel header is ordinary content (e.g. an answer followed by a + # stray trailing header); keep it, minus any other control markers. + leading = re.split(r"<\|[a-z_]+\|>", segments[0])[0] + if leading.strip() and leading.strip().lower() not in _HARMONY_ROLES: + final_parts.append(leading) + # Each later segment looks like "<|message|>body". + for seg in segments[1:]: + if "<|message|>" in seg: + chan, body = seg.split("<|message|>", 1) + chan = chan.strip().lower() + # body runs until the next control marker + body = re.split(r"<\|[a-z_]+\|>", body)[0] + (analysis_parts if "analysis" in chan else final_parts).append(body) + # a segment with no <|message|> is header noise -> ignore + return ( + "".join(final_parts) if final_parts else "", + "".join(analysis_parts) if analysis_parts else "", + ) + + +# Delta fields the streaming accumulator KNOWS how to read. Any delta attribute outside this set +# (that carries a value) is an UNSEEN field — a prime suspect when tokens are generated but neither +# content nor reasoning is captured (the server may route the answer to a field we don't drain). +_KNOWN_DELTA_FIELDS = { + "content", + "reasoning_content", + "reasoning", + "tool_calls", + "role", + "function_call", + "refusal", +} + + +def _delta_fields(delta): + """Introspect a streaming delta into (non_empty, empty) field dicts. + + Distinguishing "field present but empty" from "field absent" is essential: a case where tokens + are generated but content/reasoning stay empty could be either the server sending empty-string + deltas (field present, empty) or routing text to a field we don't read (field present, + non-empty, UNSEEN) — the two demand different fixes, so we must tell them apart. + + Also drains pydantic's `model_extra` (fields the SDK model doesn't declare — where a server's + non-standard field would land) so a routed-elsewhere payload cannot hide. + + Returns (non_empty, empty): + non_empty: {field: repr-preview} for attributes carrying a value + empty: {field: } for attributes present but empty ("" / [] / {} / None) + """ + data = None + for getter in ("model_dump", "dict"): + fn = getattr(delta, getter, None) + if callable(fn): + try: + data = fn() + break + except Exception: + data = None + if data is None: + data = { + k: getattr(delta, k, None) + for k in dir(delta) + if not k.startswith("_") and not callable(getattr(delta, k, None)) + } + data = dict(data or {}) + # pydantic v2 stashes undeclared server fields here; surface them explicitly. + extra = getattr(delta, "model_extra", None) + if isinstance(extra, dict): + for k, v in extra.items(): + data.setdefault(k, v) + non_empty, empty = {}, {} + for k, v in data.items(): + if v in (None, "", [], {}): + empty[k] = type(v).__name__ + continue + s = repr(v) + non_empty[k] = s if len(s) <= 200 else s[:200] + f"...(+{len(s)-200})" + return non_empty, empty + + +def _create_completion(client, stream, verbose=False, tid=None, turn=None, **kwargs): + """Call chat.completions.create in either mode and return a uniform _Resp-like object. + + Non-streaming: return the SDK object as-is (already has .choices[0].message etc.). + Streaming: accumulate content / reasoning_content / tool-call deltas across chunks and + reconstruct a _Resp so the rest of the agent loop is byte-for-byte identical. Streaming is + gated by the STREAM_MODE env var (see run_openai_agent) so we can toggle without a code change. + + verbose=True (VERBOSE_TRACE env) logs the full request and, for streaming, a per-call summary of + every raw delta the server sent (field coverage + any UNSEEN field + raw pre-strip previews), so + a "tokens generated but content/reasoning empty" case is fully diagnosable from CloudWatch + alone. + """ + if verbose: + msgs = kwargs.get("messages") or [] + last = msgs[-1] if msgs else {} + logger.info( + "TRACE_REQ tid=%s turn=%s stream=%s model=%s n_messages=%d n_tools=%d " + "last_role=%s last_content_preview=%r", + tid, + turn, + stream, + kwargs.get("model"), + len(msgs), + len(kwargs.get("tools") or []), + last.get("role"), + str(last.get("content"))[:200], + ) + + if not stream: + resp = client.chat.completions.create(**kwargs) + if verbose: + try: + m = resp.choices[0].message + u = getattr(resp, "usage", None) + logger.info( + "TRACE_RESP_NONSTREAM tid=%s turn=%s id=%s finish=%s content_len=%d " + "reasoning_len=%d n_tool_calls=%d prompt_tokens=%s completion_tokens=%s " + "content_preview=%r reasoning_preview=%r", + tid, + turn, + getattr(resp, "id", None), + resp.choices[0].finish_reason, + len(m.content or ""), + len( + ( + getattr(m, "reasoning_content", None) + or getattr(m, "reasoning", None) + or "" + ) + ), + len(m.tool_calls or []), + getattr(u, "prompt_tokens", None), + getattr(u, "completion_tokens", None), + (m.content or "")[:400], + str( + getattr(m, "reasoning_content", None) or getattr(m, "reasoning", None) or "" + )[:400], + ) + except Exception: + logger.exception("TRACE_RESP_NONSTREAM introspection failed") + return resp + + content_parts, reasoning_parts = [], [] + finish_reason, resp_id = None, None + prompt_tokens, completion_tokens = None, None + # tool calls arrive as indexed fragments: {index: {id, name, arguments_parts[]}} + tc_acc = {} + # --- verbose accounting: per-field chunk coverage + any field we don't read --- + n_chunks = 0 + n_empty_delta_chunks = 0 # chunks whose delta had NO non-empty field (only empty/absent) + field_hits = {} # delta field name -> count of chunks carrying a NON-EMPTY value + empty_field_hits = {} # delta field name -> count of chunks where it was present but EMPTY + unseen_fields = {} # UNSEEN non-empty field name -> a sample repr (token-loss suspects) + first_delta_fields = None + finish_delta_fields = None + raw_chunk_samples = [] # first few (non_empty, empty) delta shapes, verbatim + stream_kwargs = dict(kwargs, stream=True) + if verbose: + # Token counts on the final chunk, for the diagnostics below (verbose mode only). + stream_kwargs["stream_options"] = {"include_usage": True} + for chunk in client.chat.completions.create(**stream_kwargs): + n_chunks += 1 + if getattr(chunk, "id", None): + resp_id = chunk.id + u = getattr(chunk, "usage", None) + if u is not None: + prompt_tokens = getattr(u, "prompt_tokens", prompt_tokens) + completion_tokens = getattr(u, "completion_tokens", completion_tokens) + if not chunk.choices: + continue + choice = chunk.choices[0] + if getattr(choice, "finish_reason", None): + finish_reason = choice.finish_reason + delta = getattr(choice, "delta", None) + if delta is None: + continue + if verbose: + non_empty, empty = _delta_fields(delta) + if first_delta_fields is None: + first_delta_fields = non_empty + if getattr(choice, "finish_reason", None): + finish_delta_fields = non_empty + if not non_empty: + n_empty_delta_chunks += 1 + if len(raw_chunk_samples) < 5: + raw_chunk_samples.append({"non_empty": non_empty, "empty": empty}) + for k, v in non_empty.items(): + field_hits[k] = field_hits.get(k, 0) + 1 + if k not in _KNOWN_DELTA_FIELDS and k not in unseen_fields: + unseen_fields[k] = v # capture first sample of an unread field + for k in empty: + empty_field_hits[k] = empty_field_hits.get(k, 0) + 1 + if getattr(delta, "content", None): + content_parts.append(delta.content) + # GPT-OSS reasoning arrives on the same delta under reasoning_content/reasoning + r = getattr(delta, "reasoning_content", None) or getattr(delta, "reasoning", None) + if r: + reasoning_parts.append(r) + for tcd in getattr(delta, "tool_calls", None) or []: + idx = getattr(tcd, "index", 0) + slot = tc_acc.setdefault(idx, {"id": None, "name": None, "args": []}) + if getattr(tcd, "id", None): + slot["id"] = tcd.id + fn = getattr(tcd, "function", None) + if fn is not None: + if getattr(fn, "name", None): + slot["name"] = fn.name + if getattr(fn, "arguments", None): + slot["args"].append(fn.arguments) + tool_calls = [ + _ToolCall(v["id"], _clean_tool_name(v["name"]), "".join(v["args"])) + for _, v in sorted(tc_acc.items()) + if v["name"] + ] + # Harmony-aware reconstruction. Streamed gpt-oss deltas can carry raw Harmony control markers + # (<|channel|>, <|message|>, <|start|>, <|end|>, <|return|>, <|call|>). If left in, they leak + # into role/content on the NEXT request and the engine rejects it with + # "Unknown role: assistant<|channel|>final_output" (openai.APIError), failing the trajectory. + # The non-streaming SDK path parses these out; the streaming path must do it explicitly. + raw_content = "".join(content_parts) + raw_reasoning = "".join(reasoning_parts) + # If the raw content contains channel headers, route analysis->reasoning, final/other->content. + chan_content, chan_reasoning = _split_harmony_channels(raw_content) + content = _strip_harmony(chan_content if chan_content is not None else raw_content) + reasoning = _strip_harmony((raw_reasoning + (chan_reasoning or "")).strip()) + msg = _Msg(content, reasoning, tool_calls) + usage = _Usage(prompt_tokens, completion_tokens) + + if verbose: + # Per-call streaming summary: field coverage across chunks, any UNSEEN field (tokens routed + # to something we don't read), and raw vs post-strip lengths. Comparing completion_tokens + # with the raw lengths diagnoses "tokens generated but everything empty". + logger.info( + "TRACE_STREAM tid=%s turn=%s id=%s finish=%s n_chunks=%d n_empty_delta_chunks=%d " + "completion_tokens=%s field_hits=%s empty_field_hits=%s unseen_fields=%s " + "raw_content_len=%d raw_reasoning_len=%d post_content_len=%d post_reasoning_len=%d " + "n_tool_calls=%d first_delta=%s finish_delta=%s", + tid, + turn, + resp_id, + finish_reason, + n_chunks, + n_empty_delta_chunks, + completion_tokens, + field_hits, + empty_field_hits, + unseen_fields, + len(raw_content), + len(raw_reasoning), + len(content), + len(reasoning or ""), + len(tool_calls), + first_delta_fields, + finish_delta_fields, + ) + # If tokens were generated but we captured nothing, dump the raw text AND the first few raw + # delta shapes so the loss is fully diagnosable: empty_field_hits shows fields present-but- + # empty (server sent ""), unseen_fields shows fields carrying data we don't read, and + # raw_chunk_samples shows the verbatim first deltas. This disambiguates "server sent empty" + # from "we dropped it" from "routed to an unread field". + if ( + (completion_tokens or 0) > 0 + and not content.strip() + and not (reasoning or "").strip() + and not tool_calls + ): + logger.warning( + "TRACE_STREAM_EMPTY tid=%s turn=%s id=%s completion_tokens=%s n_chunks=%d " + "n_empty_delta_chunks=%d empty_field_hits=%s unseen_fields=%s -- tokens generated " + "but nothing captured. raw_content=%r raw_reasoning=%r raw_chunk_samples=%s", + tid, + turn, + resp_id, + completion_tokens, + n_chunks, + n_empty_delta_chunks, + empty_field_hits, + unseen_fields, + raw_content[:1000], + raw_reasoning[:1000], + raw_chunk_samples, + ) + + return _Resp(msg, finish_reason, resp_id, usage) + + +# How many times to re-issue the SAME sampling call when the engine returns a degenerate empty +# HTTP 200 (no content, no reasoning, no tool calls). An empty 200 is a success to the OpenAI +# client, so client.max_retries never covers it. Resampling in-loop keeps the trajectory alive; +# giving up fails the whole rollout (see EmptyInferenceResponse). gpt-oss emits these on ~1% of +# calls (more on the answer turn), and a few resamples almost always recover. +EMPTY_RESPONSE_MAX_RETRIES = 3 + + +def _is_empty_response(msg, content, reasoning): + """A degenerate 'empty 200': no content, no reasoning, and no tool calls.""" + return not (content or "").strip() and not msg.tool_calls and not (reasoning or "").strip() + + +def run_openai_agent( + task_input, + sop_text: str, + tools: list, + tool_executor: Callable[[str, dict], str], + client, + model: str, + max_turns: int = 10, + temperature: float = 1.0, + max_tokens: int = 4096, + top_p: float = 1.0, + debug: bool = False, + trajectory_id=None, + stream: bool = False, + verbose_trace: bool = False, +) -> tuple[str, list]: + user_content = task_input if isinstance(task_input, str) else json.dumps(task_input) + messages = [ + {"role": "system", "content": build_system_prompt(sop_text)}, + {"role": "user", "content": user_content}, + ] + # Known tool names, for _resolve_tool_name suffix-recovery (e.g. "ReportComponentMismatchjson"). + valid_tool_names = { + (t.get("function") or {}).get("name") + for t in (tools or []) + if (t.get("function") or {}).get("name") + } + + final_output = "" + if debug: + print( + f" [request params] model={model!r} temperature={temperature} " + f"max_tokens={max_tokens} top_p={top_p} n_tools={len(tools or [])} " + f"max_turns={max_turns} stream={stream} verbose_trace={verbose_trace}" + ) + for _turn in range(max_turns): + # Resample an empty 200 up to EMPTY_RESPONSE_MAX_RETRIES times before giving up. + response = msg = content = reasoning = finish_reason = None + for _attempt in range(EMPTY_RESPONSE_MAX_RETRIES + 1): + try: + # Gated streaming: when stream=True, _create_completion accumulates chunk deltas and + # returns a uniform object with the same .choices[0].message / .finish_reason / + # .usage surface, so everything below is identical to the non-streaming path. + response = _create_completion( + client, + stream, + verbose=verbose_trace, + tid=trajectory_id, + turn=_turn, + model=model, + messages=messages, + tools=tools or None, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + ) + except Exception as e: # noqa: BLE001 - log, then fail the rollout + # Endpoint errors (400 context overflow, 500, timeouts, throttling) fail the + # rollout; the trainer retries it up to `rollout_max_retries`. + logger.warning( + "turn=%s trajectory_id=%s sampling call failed: %s", _turn, trajectory_id, e + ) + raise + msg = response.choices[0].message + content = msg.content or "" + finish_reason = response.choices[0].finish_reason + reasoning = ( + getattr(msg, "reasoning_content", None) or getattr(msg, "reasoning", None) or "" + ) + + # Always-on per-turn log, greppable by trajectory_id under concurrent rollouts. + usage = getattr(response, "usage", None) + logger.info( + "turn=%s trajectory_id=%s attempt=%s finish_reason=%s content_len=%d " + "reasoning_len=%d " + "n_tool_calls=%d response_id=%s prompt_tokens=%s completion_tokens=%s", + _turn, + trajectory_id, + _attempt, + finish_reason, + len(content), + len(reasoning or ""), + len(msg.tool_calls or []), + getattr(response, "id", None), + getattr(usage, "prompt_tokens", None), + getattr(usage, "completion_tokens", None), + ) + + if not _is_empty_response(msg, content, reasoning): + break # got a usable response + + # Empty 200: resample (same messages) unless we've exhausted retries. + if _attempt < EMPTY_RESPONSE_MAX_RETRIES: + logger.warning( + "turn=%s trajectory_id=%s EMPTY_RESPONSE attempt=%s finish_reason=%s " + "response_id=%s " + "-- resampling in-loop (before any /complete-rollout)", + _turn, + trajectory_id, + _attempt, + finish_reason, + getattr(response, "id", None), + ) + + # Still empty after all resamples: fail the rollout (run_rollout reports it as failed) + # rather than treating the empty turn as a final answer. + if _is_empty_response(msg, content, reasoning): + logger.warning( + "turn=%s trajectory_id=%s EMPTY_RESPONSE_EXHAUSTED after %s resamples " + "finish_reason=%s " + "response_id=%s -- raising", + _turn, + trajectory_id, + EMPTY_RESPONSE_MAX_RETRIES, + finish_reason, + getattr(response, "id", None), + ) + raise EmptyInferenceResponse( + f"empty inference response (trajectory={trajectory_id}, turn={_turn}, " + f"finish_reason={finish_reason}, after {EMPTY_RESPONSE_MAX_RETRIES} resamples)" + ) + + # gpt-oss often leaves `content` empty and puts its text in the reasoning (analysis) + # channel. Carry the reasoning as the turn's content so the model keeps its chain of thought + # across the tool loop (without it, gpt-oss tends to return empty responses mid-task) and + # so a final answer written only in reasoning is not lost. The reward only reads answers + # from turns WITHOUT tool calls, so reasoning on a tool-call turn is never scored. + if not content.strip() and reasoning: + content = reasoning + + if debug: + print(f" [turn {_turn}] finish_reason={finish_reason}") + print(f" [turn {_turn}] content(len={len(content)}) tail={content[-600:]!r}") + if reasoning: + rc = str(reasoning) + print(f" [turn {_turn}] reasoning_content(len={len(rc)}) tail={rc[-600:]!r}") + print( + f" [turn {_turn}] tool_calls=" + f"{[(tc.function.name, tc.function.arguments) for tc in (msg.tool_calls or [])]}" + ) + + # Resolve each model-emitted tool name once: strips Harmony channel bleed + # ("...<|channel|>commentary") and recovers a hallucinated suffix + # ("ReportComponentMismatchjson" -> "ReportComponentMismatch") against the known tools. + resolved = [ + _resolve_tool_name(tc.function.name, valid_tool_names) for tc in (msg.tool_calls or []) + ] + assistant_msg = {"role": "assistant", "content": content} + if msg.tool_calls: + assistant_msg["tool_calls"] = [ + { + "id": tc.id, + "type": "function", + "function": {"name": name, "arguments": tc.function.arguments or "{}"}, + } + for tc, (name, _) in zip(msg.tool_calls, resolved) + ] + messages.append(assistant_msg) + if verbose_trace: + logger.info( + "TRACE_ASSISTANT tid=%s turn=%s content=%r tool_calls=%s", + trajectory_id, + _turn, + content[:1000], + [ + (tc["function"]["name"], tc["function"]["arguments"][:300]) + for tc in assistant_msg.get("tool_calls", []) + ], + ) + + if not msg.tool_calls: + final_output = content + break + + for tc, (tool_name, _name_recovered) in zip(msg.tool_calls, resolved): + if _name_recovered: + logger.warning( + "turn=%s trajectory_id=%s TOOL_NAME_RECOVERED raw=%r -> %r", + _turn, + trajectory_id, + tc.function.name, + tool_name, + ) + # A bad tool call (invalid JSON arguments, unknown name, tool error) is recoverable: + # feed the error back as the tool result so the model can correct itself next turn. + try: + args = json.loads(tc.function.arguments) if tc.function.arguments else {} + if not isinstance(args, dict): + raise json.JSONDecodeError("arguments must be a JSON object", "", 0) + except json.JSONDecodeError as e: + args, result = None, f"Tool call arguments are not a valid JSON object: {e}" + if args is not None: + try: + result = tool_executor(tool_name, args) + except Exception as e: # noqa: BLE001 - surface any tool error back to the model + result = f"Tool execution failed. Error: {e}" + if debug and result.startswith(("Tool execution failed", "Tool call arguments")): + print(f" [turn {_turn}] tool error fed back: {result[:300]!r}") + if verbose_trace: + logger.info( + "TRACE_TOOL tid=%s turn=%s name=%r args=%r result=%r", + trajectory_id, + _turn, + tool_name, + (tc.function.arguments or "")[:300], + str(result)[:600], + ) + messages.append({"role": "tool", "tool_call_id": tc.id, "content": result}) + + if not final_output: + # Ended without a tool-free final turn (max_turns, or a 400/500 cut the loop short): use + # the last assistant content. + for m in reversed(messages): + if m.get("role") == "assistant" and m.get("content"): + final_output = m["content"] + break + + return final_output, messages diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/prompts.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/prompts.py new file mode 100644 index 0000000000..d6c3d14110 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/prompts.py @@ -0,0 +1,28 @@ +"""System prompt and output contract for the rollout agent. + +The contract matters for scoring: SOP-Bench's task success rate (TSR) requires every expected +field to be present in the final answer, and a not-applicable field only scores when it is +emitted with a null-like value. So the prompt asks the model to emit ALL fields and to use the +string "NA" for the ones that do not apply. +""" + + +def build_system_prompt(sop_text: str) -> str: + # The answer tags deliberately avoid the word "final". gpt-oss confuses a / + # tag with its Harmony `final` channel and writes the answer into the + # reasoning channel, so the endpoint returns HTTP 200 with empty content. Renaming the tags to + # / removed the collision (A/B: empty responses 8.7% -> 0.2%). + # reward.py maps them back to SOP-Bench's tag names before scoring. + return ( + "You are an AI assistant that processes requests according to a Standard Operating " + "Procedure (SOP). Follow each SOP step in order, using the available tools to gather " + "information and make decisions.\n\n" + f"\n{sop_text}\n\n\n" + "When ALL SOP steps are complete, stop calling tools and provide your final answer in " + " tags, containing the JSON dictionary as specified in the " + "SOP. The JSON MUST include EVERY decision field the SOP asks for. For any " + 'field that does not apply to this case, set its value to the string "NA" -- do not omit ' + "the field and do not use null. Example:\n" + '{"mechanical_inspection_result": "success", ' + '"component_incident_response": "NA"}' + ) diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/reward.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/reward.py new file mode 100644 index 0000000000..bcea5cda43 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/reward.py @@ -0,0 +1,159 @@ +"""SOP-Bench reward scoring. + +The rollout's answer is the LAST assistant turn that contains a tagged answer: the prompt's +`{...}`, or SOP-Bench's own +`{...}` / `...` formats. It is +parsed with SOP-Bench's OutputParser and scored field by field against the ground truth: + +* dense - fraction of ground-truth fields answered correctly (0.0-1.0) +* binary - 1.0 only if every field is correct (SOP-Bench's task success rate, TSR) +* reward - dense * FINAL_REWARD_SCALE, or NO_ANSWER_REWARD if the rollout never answered + +Fields are compared with SOP-Bench's own value comparison so scores line up with the published +benchmark: case-insensitive, null-like values ("NA", null, NaN) are equivalent, and a value that +contains the expected one (or is contained in it) also matches - so, as in SOP-Bench, +"failed (low battery)" counts as "failed". One guard is added: a value that matches MORE THAN ONE +of the field's allowed values (for example "success or fail", "true false", or "e", which is in +both "true" and "false") never matches. Without it, one answer listing every option would score +full credit on every task. +""" + +import ast +import functools +import json +import math +import re + +from amazon_sop_bench.evaluation.parser import OutputParser + +FINAL_REWARD_SCALE = 5.0 +NO_ANSWER_REWARD = -1.0 + +_parser = OutputParser() + +# The prompt asks for / tags, renamed from SOP-Bench's +# / because gpt-oss confuses a tag containing "final" with its +# Harmony `final` channel (see prompts.py). The SOP-Bench parser only knows the original names, +# so tags are rewritten back at scoring time only. +_ANSWER_TAG_ALIASES = ( + (re.compile(r"<(/?)answer_json>", re.I), r"<\1final_output>"), + (re.compile(r"<(/?)task_result>", re.I), r"<\1final_response>"), +) +_ANSWER_MARKER = re.compile(r"<(final_output|final_response)>", re.I) +_TAGGED_DICT = re.compile(r"<(final_output|final_response)>\s*(\{.*?\})\s*", re.DOTALL | re.I) +# Null-like values, as in SOP-Bench's comparison. +_NULL_LIKE = {"none", "null", "na", "n/a", "nan", "unknown", ""} + + +def normalize_answer_tags(text: str) -> str: + """Rewrite the prompt's answer tags to the SOP-Bench parser's tag names.""" + text = text or "" + for pattern, repl in _ANSWER_TAG_ALIASES: + text = pattern.sub(repl, text) + return text + + +@functools.lru_cache(maxsize=None) +def field_vocabulary(benchmark_name: str = "aircraft_inspection") -> dict: + """{field: frozenset of allowed values (lowercase)} across all of the benchmark's tasks.""" + from amazon_sop_bench.benchmarks import load_benchmark + + vocab = {} + for task in load_benchmark(benchmark_name).tasks: + for field, value in (task.expected_output or {}).items(): + if isinstance(value, float) and math.isnan(value): + continue + if str(value).strip().lower() not in _NULL_LIKE: + vocab.setdefault(field.lower(), set()).add(str(value).strip().lower()) + return {field: frozenset(values) for field, values in vocab.items()} + + +def _clean(value) -> str: + return str(value).strip().lower().replace("_", " ").replace("-", " ") + + +def _matches_several(predicted, allowed) -> bool: + """True if `predicted` matches more than one allowed value under containment matching.""" + pred = _clean(predicted) + hits = {v for v in map(_clean, allowed) if v and (v in pred or pred in v)} + # "fail" inside "failed" is one answer, not two: keep only the longest overlapping values. + hits = {v for v in hits if not any(v != w and v in w for w in hits)} + return len(hits) > 1 + + +def _tagged_dict(answer: str) -> dict: + """The last tagged {...} dict in the answer (JSON or a Python literal), or {}.""" + for _, body in reversed(_TAGGED_DICT.findall(answer)): + for loads in (json.loads, ast.literal_eval): + try: + obj = loads(body) + except (ValueError, SyntaxError): + continue + if isinstance(obj, dict): + return obj + return {} + + +def _decisions(answer: str, ground_truth: dict) -> dict: + """{field (lowercase): predicted value} for the ground-truth fields present in the answer.""" + cols = list(ground_truth) + parsed = _parser.parse_decision(answer, cols) + decisions = parsed.get("decisions") + if not isinstance(decisions, dict): + # The parser collapses an answer with only ONE expected field to a bare "decision" + # without saying which field; read the dict directly so partial answers keep credit. + decisions = _tagged_dict(answer) + return {str(k).lower(): v for k, v in decisions.items()} + + +def _field_matches(field: str, predicted, expected, vocabulary: dict) -> bool: + pred = str(predicted).strip().lower() + if pred not in _NULL_LIKE and _matches_several(pred, vocabulary.get(field.lower(), ())): + return False # hedged: matches several allowed values + return _parser._compare_single_value(predicted, expected) + + +def score_answer(answer: str, ground_truth: dict, vocabulary: dict = None) -> tuple: + """(dense, binary) for one tagged answer.""" + if not ground_truth: + return 0.0, 0.0 + vocabulary = field_vocabulary() if vocabulary is None else vocabulary + decisions = _decisions(answer, ground_truth) + hits = sum( + 1 + for field, expected in ground_truth.items() + if str(field).lower() in decisions + and _field_matches(field, decisions[str(field).lower()], expected, vocabulary) + ) + return hits / len(ground_truth), 1.0 if hits == len(ground_truth) else 0.0 + + +def last_answer(messages: list): + """Normalized content of the last assistant turn that contains a tagged answer, or None. + + Turns with tool calls are skipped: the model is still working, and their content may be + reasoning that quotes the prompt's example answer. + """ + for msg in reversed(messages or []): + if isinstance(msg, dict) and msg.get("role") == "assistant" and not msg.get("tool_calls"): + content = normalize_answer_tags(msg.get("content") or "") + if _ANSWER_MARKER.search(content): + return content + return None + + +def compute_reward(messages: list, ground_truth: dict, vocabulary: dict = None) -> dict: + """Score a rollout transcript. The top-level `reward` key is what @sagemaker_rft_handler + reports. `vocabulary` defaults to the aircraft_inspection field vocabulary.""" + if not ground_truth: + raise ValueError("ground_truth is required to score a rollout") + answer = last_answer(messages) + if answer is None: + return {"reward": NO_ANSWER_REWARD, "dense": 0.0, "binary": 0.0, "answered": False} + dense, binary = score_answer(answer, ground_truth, vocabulary) + return { + "reward": dense * FINAL_REWARD_SCALE, + "dense": dense, + "binary": binary, + "answered": True, + } diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/samples.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/samples.py new file mode 100644 index 0000000000..b14e620d72 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/samples.py @@ -0,0 +1,30 @@ +"""Load a single SOP-Bench task instance for offline tests and reward inspection. + +Instances come from the installed `amazon-sop-bench` package (its benchmark data ships in-package), +so no benchmark data is vendored into this example. The shape matches one training prompt +(see scripts/preprocess.py), except that ground-truth NaN values are kept as float NaN so tests +can exercise the "not applicable" path. +""" + +import json + +# Not in the SOP output template, so the model cannot know to emit it (see scripts/preprocess.py). +EXCLUDED_FIELDS = {"cross_check_response"} + + +def load_instance(benchmark_name: str = "aircraft_inspection", task_id: int = 110) -> dict: + """Return {data_source, instance_id, benchmark, task_input (JSON str), ground_truth (dict)}.""" + from amazon_sop_bench.benchmarks import load_benchmark + + benchmark = load_benchmark(benchmark_name) + task = next((t for t in benchmark.tasks if str(t.task_id) == str(task_id)), None) + if task is None: + raise ValueError(f"task_id {task_id!r} not found in {benchmark_name}") + ground_truth = {k: v for k, v in task.expected_output.items() if k not in EXCLUDED_FIELDS} + return { + "data_source": "sopbench", + "instance_id": f"{benchmark_name}_{task.task_id}", + "benchmark": benchmark_name, + "task_input": json.dumps(task.inputs, default=str), + "ground_truth": ground_truth, + } diff --git a/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/tools.py b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/tools.py new file mode 100644 index 0000000000..95a3253c09 --- /dev/null +++ b/v3-examples/model-customization-examples/serverless/mtrl/sopbench_agent/src/sopbench_agentcore/tools.py @@ -0,0 +1,46 @@ +"""SOP-Bench tool adapters: OpenAI function-calling schemas and an executor. + +Both helpers take a loaded Benchmark (amazon_sop_bench.benchmarks.load_benchmark). Its `tools` +is a ToolManager exposing get_tool_specs() (Bedrock format) and +execute_tool(name, args) -> ToolCall(.result, .success, .error). +""" + +import json +from typing import Any, Callable, Dict, List + + +def create_openai_tools(benchmark) -> List[Dict[str, Any]]: + """Convert SOP-Bench's Bedrock-format tool specs to the OpenAI function-calling schema.""" + tools = [] + for spec in benchmark.tools.get_tool_specs(): + ts = spec.get("toolSpec", spec) + schema = ts.get("inputSchema", {}) + params = schema.get("json", schema) if isinstance(schema, dict) else {} + if not isinstance(params, dict) or "type" not in params: + params = {"type": "object", "properties": {}} + tools.append( + { + "type": "function", + "function": { + "name": ts["name"], + "description": ts.get("description", ""), + "parameters": params, + }, + } + ) + return tools + + +def make_tool_executor(benchmark) -> Callable[[str, Dict[str, Any]], str]: + """Return a closure(tool_name, args) -> result string over the benchmark's ToolManager.""" + tm = benchmark.tools + + def _execute(tool_name: str, arguments: Dict[str, Any]) -> str: + result = tm.execute_tool(tool_name, arguments) + if result.success: + if isinstance(result.result, str): + return result.result + return json.dumps(result.result, default=str) + return f"Tool execution failed. Error: {result.error}" + + return _execute