diff --git a/src/agentcore_rl_toolkit/backends/experimental/sagemaker/SETUP.md b/src/agentcore_rl_toolkit/backends/experimental/sagemaker/SETUP.md
index a5d44728..4cdd31d8 100644
--- a/src/agentcore_rl_toolkit/backends/experimental/sagemaker/SETUP.md
+++ b/src/agentcore_rl_toolkit/backends/experimental/sagemaker/SETUP.md
@@ -48,3 +48,15 @@ cp config.yaml.example config.yaml
 cd src/agentcore_rl_toolkit/backends/experimental/sagemaker
 python train_grpo.py --config config.yaml
 ```
+
+To train with DAPO instead, run `train_dapo.py` with the same config:
+
+```bash
+python train_dapo.py --config config.yaml
+```
+
+`train_dapo.py` is the GRPO loop plus dynamic sampling (each step rolls out batches of
+prompts and keeps only groups whose rewards differ, until `batch_size` groups are kept) and
+Clip-Higher (`clip_low_threshold` / `clip_high_threshold`). Its extra settings are in the
+DAPO section of `config.yaml.example`. Each step logs a `sampling:` line with the number of
+generation batches and the filtered fraction.
diff --git a/src/agentcore_rl_toolkit/backends/experimental/sagemaker/config.py b/src/agentcore_rl_toolkit/backends/experimental/sagemaker/config.py
index 6cc183bb..0c7937eb 100644
--- a/src/agentcore_rl_toolkit/backends/experimental/sagemaker/config.py
+++ b/src/agentcore_rl_toolkit/backends/experimental/sagemaker/config.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Literal
+from typing import Literal, TypeVar
 
 import yaml
 
@@ -58,9 +58,18 @@ class Config:
     rollout_log_path: str = ""  # JSONL file for per-rollout diagnostics; auto-named if empty
 
 
-def load_config(path: str) -> Config:
+ConfigT = TypeVar("ConfigT", bound=Config)
+
+
+def load_config(path: str, config_cls: type[ConfigT] = Config) -> ConfigT:
+    """Load a YAML file into ``config_cls``: ``Config``, or a recipe's subclass that adds fields.
+
+    Keys that ``config_cls`` does not define are ignored, so one YAML file can serve several
+    recipes (e.g. ``train_grpo`` skips the fields only ``train_dapo`` reads). Returns an
+    instance of ``config_cls``.
+    """
     with open(path) as f:
         raw: dict = yaml.safe_load(f) or {}
 
-    fields = Config.__dataclass_fields__
-    return Config(**{k: v for k, v in raw.items() if k in fields})
+    fields = config_cls.__dataclass_fields__
+    return config_cls(**{k: v for k, v in raw.items() if k in fields})
diff --git a/src/agentcore_rl_toolkit/backends/experimental/sagemaker/config.yaml.example b/src/agentcore_rl_toolkit/backends/experimental/sagemaker/config.yaml.example
index d1e30dbf..a58b3265 100644
--- a/src/agentcore_rl_toolkit/backends/experimental/sagemaker/config.yaml.example
+++ b/src/agentcore_rl_toolkit/backends/experimental/sagemaker/config.yaml.example
@@ -1,4 +1,5 @@
-# GRPO training — rollout gateway + SageMaker Training Sessions
+# GRPO / DAPO training — rollout gateway + SageMaker Training Sessions
+# Both train_grpo.py and train_dapo.py read this file; each ignores keys it does not use.
 
 # --- AWS resources (required) ---
 role_arn: "arn:aws:iam::<account-id>:role/<your-sagemaker-role>"
@@ -51,3 +52,8 @@ acr_tps_limit: 5
 # --- Logging ---
 wandb_project: ""
 wandb_run_name: ""  # defaults to the SageMaker training session name
+
+# --- DAPO (train_dapo.py only; train_grpo.py ignores these) ---
+clip_low_threshold: 0.8    # PPO ratio lower bound, 1 - eps_low
+clip_high_threshold: 1.28  # PPO ratio upper bound, 1 + eps_high (Clip-Higher)
+max_num_gen_batches: 10    # dynamic sampling: max generation batches per step (0 = no cap)
diff --git a/src/agentcore_rl_toolkit/backends/experimental/sagemaker/train_dapo.py b/src/agentcore_rl_toolkit/backends/experimental/sagemaker/train_dapo.py
new file mode 100644
index 00000000..64b06845
--- /dev/null
+++ b/src/agentcore_rl_toolkit/backends/experimental/sagemaker/train_dapo.py
@@ -0,0 +1,402 @@
+"""DAPO training with the rollout gateway + SageMaker Training Sessions.
+
+Same loop as ``train_grpo.py`` with two DAPO changes:
+
+- Dynamic sampling: each step rolls out generation batches of ``batch_size`` prompts and keeps
+  only groups whose rewards differ (an all-equal group has zero advantage and no gradient),
+  until ``batch_size`` groups are kept.
+- Clip-Higher: the PPO ratio bounds ``clip_low_threshold`` / ``clip_high_threshold`` are sent
+  to the service as ``loss_fn_config``.
+
+Reads the same config.yaml as ``train_grpo.py``; the DAPO-only fields in ``DapoConfig`` have
+defaults, so they are optional in the YAML.
+
+Usage:
+    python -m agentcore_rl_toolkit.backends.experimental.sagemaker.train_dapo --config path/to/config.yaml
+"""
+
+import argparse
+import asyncio
+import logging
+import signal
+import time
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any
+
+import numpy as np
+from sagemaker.train.training_session import AdamParams, ServiceClient
+
+from agentcore_rl_toolkit.backends.experimental.sagemaker.config import Config, load_config
+from agentcore_rl_toolkit.backends.experimental.sagemaker.datum import trace_record_to_datum
+from agentcore_rl_toolkit.backends.experimental.sagemaker.rollout import load_dataset, local_ip, run_one_rollout
+from agentcore_rl_toolkit.backends.experimental.sagemaker.train_grpo import _init_wandb, compute_grpo_advantages
+from agentcore_rl_toolkit.concurrency.rate_limiter import LocalRateLimiter
+from agentcore_rl_toolkit.rollout_gateway import HfTemplateRenderer, TraceRecord
+from agentcore_rl_toolkit.rollout_gateway.gateway import RolloutGateway
+from agentcore_rl_toolkit.rollout_gateway.sampling_backends.sagemaker_sdk import SageMakerSdkBackend
+from agentcore_rl_toolkit.rollout_gateway.server import ThreadedGatewayServer
+from agentcore_rl_toolkit.runtime import AgentCoreHttpClient
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s: %(message)s")
+
+# One prompt's K rollouts: [(records, reward), ...]
+Group = list[tuple[list[TraceRecord], float]]
+
+
+@dataclass
+class DapoConfig(Config):
+    # Clip-Higher: PPO ratio bounds sent to the service, i.e. 1 - eps_low and 1 + eps_high.
+    clip_low_threshold: float = 0.8
+    clip_high_threshold: float = 1.28
+    # Dynamic sampling: max generation batches per step to fill batch_size informative groups (0 = no cap).
+    max_num_gen_batches: int = 10
+
+
+class InsufficientGroupsError(RuntimeError):
+    """Dynamic sampling hit max_num_gen_batches before keeping batch_size informative groups."""
+
+
+def is_informative(rewards: list[float]) -> bool:
+    """Whether a group can train: its rewards must differ, or every advantage is zero."""
+    return len(rewards) >= 2 and max(rewards) > min(rewards)
+
+
+async def sample_informative_groups(
+    example_queue: list[dict],
+    batch_size: int,
+    max_num_gen_batches: int,
+    collect_group: Callable[[dict], Awaitable[Group]],
+) -> tuple[list[Group] | None, list[float], dict[str, int]]:
+    """Dynamic sampling for one step.
+
+    Pops ``batch_size`` prompts at a time from ``example_queue`` (in place), rolls them out with
+    ``collect_group`` and keeps the informative groups, until ``batch_size`` groups are kept.
+    Only one generation batch is in flight at a time, so at most ``batch_size * K`` rollouts run
+    concurrently, as in ``train_grpo.py``.
+
+    Returns ``(groups, rewards, stats)``:
+        groups: the first ``batch_size`` kept groups (extras are dropped), or None if the epoch
+            ran out of prompts first.
+        rewards: every rollout reward of this step, before filtering.
+        stats: num_gen_batches, num_groups_seen, num_groups_filtered, num_kept, num_empty_rollouts.
+
+    Raises:
+        InsufficientGroupsError: ``max_num_gen_batches`` generation batches ran (0 = no cap)
+            without keeping ``batch_size`` groups.
+    """
+    kept: list[Group] = []
+    rewards: list[float] = []
+    stats = {
+        "num_gen_batches": 0,
+        "num_groups_seen": 0,
+        "num_groups_filtered": 0,
+        "num_kept": 0,
+        "num_empty_rollouts": 0,
+    }
+    while len(kept) < batch_size:
+        # Check the cap before the queue so hitting it raises instead of quietly ending the epoch.
+        if max_num_gen_batches > 0 and stats["num_gen_batches"] >= max_num_gen_batches:
+            raise InsufficientGroupsError(
+                f"kept {len(kept)} of {batch_size} informative groups "
+                f"after {stats['num_gen_batches']} generation batches"
+            )
+        if len(example_queue) < batch_size:
+            return None, rewards, stats
+        prompts = [example_queue.pop(0) for _ in range(batch_size)]
+        groups = await asyncio.gather(*[collect_group(p) for p in prompts])
+        stats["num_gen_batches"] += 1
+        for group in groups:
+            group_rewards = [r for _, r in group]
+            rewards.extend(group_rewards)
+            stats["num_groups_seen"] += 1
+            stats["num_empty_rollouts"] += sum(1 for records, _ in group if not records)
+            if is_informative(group_rewards):
+                kept.append(group)
+            else:
+                stats["num_groups_filtered"] += 1
+        stats["num_kept"] = len(kept)
+    return kept[:batch_size], rewards, stats
+
+
+async def train(config: DapoConfig) -> None:
+    if not config.dataset_path:
+        raise ValueError("dataset_path is required")
+
+    # Load data before creating the session: bad data should not start billing.
+    dataset = load_dataset(config.dataset_path)
+    logger.info("Loaded %d examples from %s", len(dataset), config.dataset_path)
+    if config.max_prompts > 0:
+        dataset = dataset[: config.max_prompts]
+    if len(dataset) < config.batch_size:
+        raise ValueError(
+            f"dataset has {len(dataset)} prompts, fewer than batch_size={config.batch_size}; no step would run"
+        )
+
+    eval_dataset: list[dict] = []
+    if config.eval_dataset_path:
+        eval_dataset = load_dataset(config.eval_dataset_path)
+        if config.eval_max_prompts > 0:
+            eval_dataset = eval_dataset[: config.eval_max_prompts]
+        logger.info("Loaded %d eval examples", len(eval_dataset))
+
+    # Clip-Higher. The SDK takes ratio bounds, not eps.
+    loss_fn_config = {
+        "clip_low_threshold": config.clip_low_threshold,
+        "clip_high_threshold": config.clip_high_threshold,
+    }
+    logger.info("loss_fn=%s loss_fn_config=%s", config.loss, loss_fn_config)
+
+    # Create SageMaker training session client
+    svc = ServiceClient(
+        region=config.region,
+        role_arn=config.role_arn,
+        model_package_group_arn=config.model_package_group_arn or None,
+        s3_output_path=config.s3_output_path or None,
+    )
+
+    if config.resume_from_model_package:
+        logger.info("Resuming from checkpoint: %s", config.resume_from_model_package)
+        tc = svc.create_training_client_from_state_with_optimizer(
+            model_package_arn=config.resume_from_model_package,
+            base_model=config.base_model_arn,
+            rank=config.lora_rank,
+            alpha=config.lora_alpha,
+        )
+    else:
+        logger.info("Creating LoRA training session (base_model=%s)", config.base_model_arn)
+        tc = svc.create_lora_training_client(
+            base_model=config.base_model_arn,
+            rank=config.lora_rank,
+            alpha=config.lora_alpha,
+        )
+    logger.info("Training session ready: %s", tc.training_session_name)
+
+    # The session bills until stopped, so every exit path must reach the finally below.
+    # SIGTERM would otherwise kill the interpreter before it runs; cancellation unwinds
+    # through it instead. (SIGINT already does, via KeyboardInterrupt.)
+    task = asyncio.current_task()
+    assert task is not None  # always set inside a running coroutine
+    asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, task.cancel)
+
+    gateway: ThreadedGatewayServer | None = None
+    rollout_client: AgentCoreHttpClient | None = None
+    wandb_run: Any = None
+    try:
+        wandb_run = _init_wandb(config, tc.training_session_name)
+
+        tokenizer = tc.get_tokenizer()
+        backend = SageMakerSdkBackend(tc.create_sampling_client())
+        _gw = RolloutGateway(
+            backend=backend,
+            renderer=HfTemplateRenderer(tokenizer, chat_template_kwargs=config.chat_template_kwargs),
+            tokenizer=tokenizer,
+            adapters=["openai", "anthropic"],
+            history_mode=config.history_mode,
+        )
+        gateway = ThreadedGatewayServer(_gw, host=local_ip(), port=config.gateway_port)
+        gateway.start()
+        logger.info("Gateway base_url: %s", gateway.base_url)
+
+        rollout_client = AgentCoreHttpClient(
+            config.agent_runtime_arn,
+            max_pool_connections=config.batch_size * config.responses_per_prompt + 10,
+            request_rate_limiter=LocalRateLimiter(config.acr_tps_limit),
+        )
+
+        sampling_defaults = {"max_new_tokens": config.max_new_tokens, "temperature": config.temperature}
+        max_context_tokens = config.max_context_tokens
+
+        async def collect_group(payload: dict) -> Group:
+            results = await asyncio.gather(
+                *[
+                    run_one_rollout(
+                        client=rollout_client,
+                        gateway=_gw,
+                        payload=payload,
+                        base_url=gateway.base_url,
+                        model_id=config.base_model_arn,
+                        max_rollout_time=config.max_rollout_time,
+                        sampling_defaults=sampling_defaults,
+                        max_context_tokens=max_context_tokens,
+                    )
+                    for _ in range(config.responses_per_prompt)
+                ]
+            )
+            return [(records, reward) for records, reward, _ in results]
+
+        async def run_eval(step: int) -> None:
+            if not eval_dataset:
+                return
+            eval_sd = {**sampling_defaults, "temperature": config.eval_temperature}
+            logger.info("  Evaluating %d held-out examples…", len(eval_dataset))
+            results = await asyncio.gather(
+                *[
+                    run_one_rollout(
+                        client=rollout_client,
+                        gateway=_gw,
+                        payload=p,
+                        base_url=gateway.base_url,
+                        model_id=config.base_model_arn,
+                        max_rollout_time=config.max_rollout_time,
+                        sampling_defaults=eval_sd,
+                        max_context_tokens=max_context_tokens,
+                    )
+                    for p in eval_dataset
+                ]
+            )
+            scores = [reward for _, reward, _ in results]
+            acc = float(np.mean(scores))
+            logger.info("  eval/accuracy=%.4f (n=%d, step=%d)", acc, len(scores), step)
+            if wandb_run:
+                wandb_run.log({"eval/accuracy": acc, "eval/n": len(scores)}, step=step)
+
+        global_step = 0
+        last_eval_step: int | None = None  # step of the latest eval; decides whether a final eval is needed
+
+        for epoch in range(config.epochs):
+            example_queue = list(dataset)
+            logger.info(
+                "Epoch %d: %d examples, batch=%d, K=%d",
+                epoch,
+                len(example_queue),
+                config.batch_size,
+                config.responses_per_prompt,
+            )
+
+            while len(example_queue) >= config.batch_size:
+                logger.info("=== Epoch %d | Step %d [%s] ===", epoch, global_step, datetime.now().strftime("%H:%M:%S"))
+                metrics: dict[str, Any] = {"step": global_step, "epoch": epoch}
+
+                if config.save_every > 0 and global_step > 0 and global_step % config.save_every == 0:
+                    ckpt = await asyncio.to_thread(lambda: tc.save_state().result(timeout=config.train_op_timeout))
+                    logger.info("Checkpoint saved: %s", ckpt.model_package_arn)
+
+                st = time.monotonic()
+                try:
+                    # Sample prompts, roll out, filter; repeat until batch_size informative groups are kept
+                    all_groups, all_rewards, ds_stats = await sample_informative_groups(
+                        example_queue, config.batch_size, config.max_num_gen_batches, collect_group
+                    )
+                except InsufficientGroupsError:
+                    # The session is healthy here, so keep the progress before stopping.
+                    logger.exception("Dynamic sampling exhausted at step %d; saving before abort", global_step)
+                    ckpt = await asyncio.to_thread(lambda: tc.save_state().result(timeout=config.train_op_timeout))
+                    logger.info("Abort checkpoint: %s", ckpt.model_package_arn)
+                    raise
+                metrics["time/rollout"] = time.monotonic() - st
+                metrics["sampler/mean_reward"] = float(np.mean(all_rewards)) if all_rewards else 0.0
+                metrics.update({f"dynamic_sampling/{k}": v for k, v in ds_stats.items()})
+                seen = ds_stats["num_groups_seen"]
+                filtered_frac = ds_stats["num_groups_filtered"] / seen if seen else 0.0
+                metrics["dynamic_sampling/filtered_frac"] = filtered_frac
+                logger.info(
+                    "  sampling: gen_batches=%d seen=%d filtered=%d (%.0f%%) kept=%d empty_rollouts=%d "
+                    "reward=%.4f rollout_time=%.0fs",
+                    ds_stats["num_gen_batches"],
+                    seen,
+                    ds_stats["num_groups_filtered"],
+                    filtered_frac * 100,
+                    ds_stats["num_kept"],
+                    ds_stats["num_empty_rollouts"],
+                    metrics["sampler/mean_reward"],
+                    metrics["time/rollout"],
+                )
+
+                if all_groups is None:
+                    logger.info("  Epoch %d ran out of prompts during dynamic sampling.", epoch)
+                    if wandb_run:
+                        wandb_run.log(
+                            {k: v for k, v in metrics.items() if k not in ("step", "epoch")}, step=global_step
+                        )
+                    break
+
+                # --- advantage + datums ---
+                training_datums: list[dict] = []
+                for group in all_groups:
+                    # GRPO-normalized advantage within each group of K rollouts
+                    advantages = compute_grpo_advantages([r for _, r in group])
+                    for (records, _), adv in zip(group, advantages, strict=True):
+                        for record in records:
+                            datum = trace_record_to_datum(record, adv)
+                            if datum is not None:
+                                training_datums.append(datum)
+
+                metrics["sampler/num_datums"] = len(training_datums)
+
+                if not training_datums:
+                    logger.warning("Step %d: no valid datums, skipping update.", global_step)
+                    global_step += 1
+                    continue
+
+                logger.info("  Training on %d datums…", len(training_datums))
+                st = time.monotonic()
+                # Server side: forward_backward + optim_step
+                fwd_bwd_op = tc.forward_backward(training_datums, loss_fn=config.loss, loss_fn_config=loss_fn_config)
+                optim_op = tc.optim_step(AdamParams(learning_rate=config.learning_rate))
+                fwd_bwd_result = await asyncio.to_thread(fwd_bwd_op.result, timeout=config.train_op_timeout)
+                await asyncio.to_thread(optim_op.result, timeout=config.train_op_timeout)
+
+                metrics["time/train"] = time.monotonic() - st
+                metrics["timestamp"] = datetime.now().isoformat()
+                metrics.update({f"train/{k}": v for k, v in fwd_bwd_result.metrics.items()})
+                logger.info("  step=%d  reward=%.4f", global_step, metrics["sampler/mean_reward"])
+                if wandb_run:
+                    wandb_run.log(
+                        {k: v for k, v in metrics.items() if k not in ("step", "epoch", "timestamp")}, step=global_step
+                    )
+
+                # rebind to updated weights before next rollout
+                new_sc = await asyncio.to_thread(tc.save_weights_and_get_sampling_client)
+                backend.set_sampling_client(new_sc)
+
+                global_step += 1
+                if config.eval_every > 0 and global_step % config.eval_every == 0:
+                    # Eval every eval_every steps
+                    await run_eval(global_step)
+                    last_eval_step = global_step
+
+            logger.info("Epoch %d complete (%d steps)", epoch, global_step)
+
+        # Compare with the last eval rather than global_step % eval_every: that check skips the
+        # final eval after 0 steps, or when the last step was skipped for having no datums.
+        if eval_dataset and last_eval_step != global_step:
+            await run_eval(global_step)
+
+        logger.info("Training complete. Saving final checkpoint…")
+        final_ckpt = await asyncio.to_thread(lambda: tc.save_state().result(timeout=config.train_op_timeout))
+        logger.info("Final checkpoint: %s", final_ckpt.model_package_arn)
+
+        if config.save_weights_at_end:
+            w = await asyncio.to_thread(
+                lambda: tc.save_weights_for_sampler(description="final weights").result(timeout=config.train_op_timeout)
+            )
+            logger.info("Inference model package: %s", w.model_package_arn)
+    finally:
+        tc.stop()  # first: nothing may preempt stopping the billing session
+        if gateway is not None:
+            gateway.shutdown()
+        if rollout_client is not None:
+            await rollout_client.aclose()
+        if wandb_run:
+            wandb_run.finish()
+
+    logger.info("Done.")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="DAPO training with rollout gateway + SageMaker Training Sessions")
+    parser.add_argument("--config", default="config.yaml", help="Path to config.yaml")
+    args = parser.parse_args()
+    config = load_config(args.config, DapoConfig)
+    try:
+        asyncio.run(train(config))
+    except asyncio.CancelledError:
+        # Raised by the SIGTERM/SIGINT handler after teardown has already run.
+        logger.warning("Training cancelled.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/backends/experimental/sagemaker/test_config.py b/tests/backends/experimental/sagemaker/test_config.py
new file mode 100644
index 00000000..0d9fd7c2
--- /dev/null
+++ b/tests/backends/experimental/sagemaker/test_config.py
@@ -0,0 +1,53 @@
+"""``load_config`` unit tests: unknown keys are ignored, and ``config_cls`` lets a recipe load
+the same YAML into a ``Config`` subclass that adds its own fields."""
+
+from dataclasses import dataclass
+
+from agentcore_rl_toolkit.backends.experimental.sagemaker.config import Config, load_config
+
+
+@dataclass
+class _RecipeConfig(Config):
+    extra_threshold: float = 0.5
+    extra_count: int = 3
+
+
+def _write(tmp_path, text: str) -> str:
+    path = tmp_path / "config.yaml"
+    path.write_text(text)
+    return str(path)
+
+
+def test_default_class_is_config(tmp_path):
+    config = load_config(_write(tmp_path, "batch_size: 8\n"))
+
+    assert type(config) is Config
+    assert config.batch_size == 8
+
+
+def test_empty_file_gives_defaults(tmp_path):
+    assert load_config(_write(tmp_path, "")) == Config()
+
+
+def test_unknown_keys_are_ignored(tmp_path):
+    config = load_config(_write(tmp_path, "batch_size: 8\nnot_a_field: 1\n"))
+
+    assert config.batch_size == 8
+    assert not hasattr(config, "not_a_field")
+
+
+def test_config_cls_reads_subclass_fields(tmp_path):
+    config = load_config(_write(tmp_path, "batch_size: 8\nextra_threshold: 0.9\n"), _RecipeConfig)
+
+    assert type(config) is _RecipeConfig
+    assert config.batch_size == 8  # base field
+    assert config.extra_threshold == 0.9  # subclass field from YAML
+    assert config.extra_count == 3  # subclass field not in YAML keeps its default
+
+
+def test_base_config_ignores_subclass_fields(tmp_path):
+    # One YAML for two recipes: the base loader skips fields only the subclass defines.
+    path = _write(tmp_path, "batch_size: 8\nextra_threshold: 0.9\nextra_count: 7\n")
+
+    assert load_config(path) == Config(batch_size=8)
+    assert load_config(path, _RecipeConfig) == _RecipeConfig(batch_size=8, extra_threshold=0.9, extra_count=7)
diff --git a/tests/backends/experimental/sagemaker/test_train_dapo.py b/tests/backends/experimental/sagemaker/test_train_dapo.py
new file mode 100644
index 00000000..bb0f929d
--- /dev/null
+++ b/tests/backends/experimental/sagemaker/test_train_dapo.py
@@ -0,0 +1,433 @@
+"""``train_dapo`` tests: dynamic sampling on its own, and the full ``train()`` loop against fakes.
+
+No AWS calls: the SageMaker service client, the gateway, the ACR client and the rollouts are
+replaced with fakes. Dynamic sampling, advantages, datum building, the update order, evals and
+teardown are the real code.
+"""
+
+import asyncio
+import logging
+import re
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("sagemaker.train.training_session", reason="needs sagemaker-train>=1.27 (SageMaker extras)")
+
+import agentcore_rl_toolkit.backends.experimental.sagemaker.train_dapo as td  # noqa: E402
+from agentcore_rl_toolkit.rollout_gateway import TraceRecord  # noqa: E402
+
+MIXED = [1.0, 0.0]  # informative: rewards differ
+SAME = [1.0, 1.0]  # filtered: all rewards equal
+
+
+def record(loss_mask: tuple[int, ...] = (1, 1)) -> TraceRecord:
+    n = len(loss_mask)
+    return TraceRecord(token_ids=[1, 2, *range(10, 10 + n)], loss_mask=list(loss_mask), logprobs=[-0.5] * n)
+
+
+def prompts(*rewards: list[float], untrainable: bool = False) -> list[dict]:
+    return [{"id": i, "rewards": r, "untrainable": untrainable} for i, r in enumerate(rewards)]
+
+
+# ---------------------------------------------------------------------------
+# is_informative / sample_informative_groups
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize(
+    ("rewards", "expected"),
+    [
+        ([1.0, 0.0], True),
+        ([0.0, 0.0], False),
+        ([1.0, 1.0], False),
+        ([0.5, 0.5, 0.6], True),
+        ([1.0], False),
+        ([], False),
+    ],
+)
+def test_is_informative(rewards, expected):
+    assert td.is_informative(rewards) is expected
+
+
+class FakeGroups:
+    """Stands in for collect_group: each prompt says its K rewards; one trainable record per rollout."""
+
+    def __init__(self):
+        self.seen: list[int] = []
+        self.in_flight = 0
+        self.max_in_flight = 0
+
+    async def __call__(self, payload: dict) -> td.Group:
+        self.seen.append(payload["id"])
+        self.in_flight += 1
+        self.max_in_flight = max(self.max_in_flight, self.in_flight)
+        await asyncio.sleep(0)
+        self.in_flight -= 1
+        return [([record()], r) for r in payload["rewards"]]
+
+
+def sample(queue: list[dict], batch_size: int = 2, max_num_gen_batches: int = 10, fake: FakeGroups | None = None):
+    fake = fake or FakeGroups()
+    result = asyncio.run(td.sample_informative_groups(queue, batch_size, max_num_gen_batches, fake))
+    return result, fake
+
+
+def test_sample_one_generation_enough():
+    queue = prompts(MIXED, MIXED, MIXED)
+    (groups, rewards, stats), fake = sample(queue)
+    assert groups is not None and len(groups) == 2
+    assert fake.seen == [0, 1]
+    assert len(queue) == 1  # popped in place
+    assert rewards == [1.0, 0.0, 1.0, 0.0]
+    assert stats == {
+        "num_gen_batches": 1,
+        "num_groups_seen": 2,
+        "num_groups_filtered": 0,
+        "num_kept": 2,
+        "num_empty_rollouts": 0,
+    }
+
+
+def test_sample_resamples_until_full():
+    queue = prompts(SAME, MIXED, SAME, SAME, MIXED, SAME)
+    (groups, rewards, stats), fake = sample(queue)
+    assert groups is not None and len(groups) == 2
+    assert fake.seen == [0, 1, 2, 3, 4, 5]
+    assert len(rewards) == 12  # every rollout of every generation, before filtering
+    assert (stats["num_gen_batches"], stats["num_groups_filtered"], stats["num_kept"]) == (3, 4, 2)
+
+
+def test_sample_truncates_extra_groups():
+    (groups, _, stats), _ = sample(prompts(SAME, MIXED, MIXED, MIXED))
+    assert groups is not None and len(groups) == 2
+    assert stats["num_kept"] == 3  # kept before truncation is reported
+
+
+def test_sample_returns_none_when_epoch_runs_out():
+    queue = prompts(SAME, SAME, MIXED)
+    (groups, rewards, stats), fake = sample(queue)
+    assert groups is None
+    assert rewards == [1.0, 1.0, 1.0, 1.0]  # still reported for logging
+    assert len(queue) == 1  # the leftover prompt is not rolled out
+    assert stats["num_gen_batches"] == 1
+
+
+def test_sample_raises_at_cap():
+    fake = FakeGroups()
+    with pytest.raises(td.InsufficientGroupsError, match="after 2 generation batches"):
+        sample(prompts(*[SAME] * 10), max_num_gen_batches=2, fake=fake)
+    assert len(fake.seen) == 4  # exactly 2 generations, no more
+
+
+def test_sample_cap_checked_before_queue():
+    # the cap is reached exactly when the queue runs out: raise, do not quietly end the epoch
+    with pytest.raises(td.InsufficientGroupsError):
+        sample(prompts(SAME, SAME, SAME, SAME), max_num_gen_batches=2)
+
+
+def test_sample_cap_zero_is_unlimited():
+    (groups, _, stats), _ = sample(prompts(*[SAME] * 8, MIXED, MIXED), max_num_gen_batches=0)
+    assert groups is not None
+    assert stats["num_gen_batches"] == 5
+
+
+def test_sample_counts_empty_rollouts():
+    async def one_failed(payload):
+        return [([record()], 1.0), ([], 0.0)]  # rollout 1 failed: no records, reward 0
+
+    queue = prompts(MIXED, MIXED)
+    groups, _, stats = asyncio.run(td.sample_informative_groups(queue, 2, 10, one_failed))
+    assert groups is not None
+    assert stats["num_empty_rollouts"] == 2
+
+
+def test_sample_one_generation_in_flight_at_a_time():
+    fake = FakeGroups()
+    sample(prompts(SAME, SAME, MIXED, MIXED), fake=fake)
+    assert fake.max_in_flight <= 2  # never more than batch_size groups at once
+
+
+# ---------------------------------------------------------------------------
+# train(): fakes for everything that would talk to AWS
+# ---------------------------------------------------------------------------
+class Done:
+    def __init__(self, value=None):
+        self.value = value
+
+    def result(self, timeout=None):
+        return self.value
+
+
+class FakeTrainingClient:
+    training_session_name = "fake-session"
+
+    def __init__(self):
+        self.calls: list[str] = []
+        self.fwd_bwd: list[dict] = []
+        self.datums: list[list[dict]] = []
+
+    def get_tokenizer(self):
+        return object()
+
+    def create_sampling_client(self):
+        return object()
+
+    def forward_backward(self, datums, loss_fn, loss_fn_config=None):
+        self.calls.append("forward_backward")
+        self.fwd_bwd.append({"n": len(datums), "loss_fn": loss_fn, "loss_fn_config": loss_fn_config})
+        self.datums.append(datums)
+        return Done(SimpleNamespace(metrics={"loss": 0.1}))
+
+    def optim_step(self, params):
+        self.calls.append("optim_step")
+        return Done()
+
+    def save_state(self):
+        self.calls.append("save_state")
+        return Done(SimpleNamespace(model_package_arn="ckpt"))
+
+    def save_weights_and_get_sampling_client(self):
+        self.calls.append("rebind")
+        return object()
+
+    def save_weights_for_sampler(self, description):
+        self.calls.append("save_weights_for_sampler")
+        return Done(SimpleNamespace(model_package_arn="weights"))
+
+    def stop(self):
+        self.calls.append("stop")
+
+
+class FakeServer:
+    def __init__(self, gw, host, port):
+        self.base_url = "http://fake"
+        self.closed = False
+
+    def start(self):
+        pass
+
+    def shutdown(self):
+        self.closed = True
+
+
+class FakeHttpClient:
+    def __init__(self, *args, **kwargs):
+        self.closed = False
+
+    async def aclose(self):
+        self.closed = True
+
+
+@pytest.fixture
+def harness(monkeypatch):
+    tc = FakeTrainingClient()
+    state = SimpleNamespace(tc=tc, train=[], eval=[], evals=[], created=0, server=None, client=None, fail_on_call=None)
+    calls = {"n": 0}
+
+    def create_lora_training_client(**kwargs):
+        state.created += 1
+        return tc
+
+    def server(*args, **kwargs):
+        state.server = FakeServer(*args, **kwargs)
+        return state.server
+
+    def http_client(*args, **kwargs):
+        state.client = FakeHttpClient()
+        return state.client
+
+    async def run_one_rollout(*, payload, sampling_defaults, **kwargs):
+        if "eval" in payload:
+            state.evals.append(sampling_defaults["temperature"])
+            return [], payload["eval"], {}
+        if state.fail_on_call is not None and calls["n"] == state.fail_on_call:
+            raise state.fail_on_call_error
+        k = payload.setdefault("_k", 0)
+        payload["_k"] = (k + 1) % len(payload["rewards"])
+        calls["n"] += 1
+        mask = (0, 0) if payload["untrainable"] else (1, 1)
+        return [record(mask)], payload["rewards"][k], {}
+
+    def load_dataset(path):
+        return [dict(p) for p in (state.eval if path == "eval.jsonl" else state.train)]
+
+    monkeypatch.setattr(
+        td, "ServiceClient", lambda **kw: SimpleNamespace(create_lora_training_client=create_lora_training_client)
+    )
+    monkeypatch.setattr(td, "SageMakerSdkBackend", lambda sc: SimpleNamespace(set_sampling_client=lambda new: None))
+    monkeypatch.setattr(td, "RolloutGateway", lambda **kw: object())
+    monkeypatch.setattr(td, "HfTemplateRenderer", lambda *a, **kw: object())
+    monkeypatch.setattr(td, "ThreadedGatewayServer", server)
+    monkeypatch.setattr(td, "AgentCoreHttpClient", http_client)
+    monkeypatch.setattr(td, "LocalRateLimiter", lambda *a: object())
+    monkeypatch.setattr(td, "local_ip", lambda: "127.0.0.1")
+    monkeypatch.setattr(td, "run_one_rollout", run_one_rollout)
+    monkeypatch.setattr(td, "load_dataset", load_dataset)
+    return state
+
+
+def make_config(**overrides) -> td.DapoConfig:
+    base = {"dataset_path": "train.jsonl", "batch_size": 2, "responses_per_prompt": 2}
+    return td.DapoConfig(**{**base, **overrides})
+
+
+def run(config: td.DapoConfig) -> None:
+    asyncio.run(td.train(config))
+
+
+def test_dapo_config_defaults_and_inherits_config():
+    config = td.DapoConfig()
+    assert isinstance(config, td.Config)
+    assert (config.clip_low_threshold, config.clip_high_threshold, config.max_num_gen_batches) == (0.8, 1.28, 10)
+
+
+def test_train_happy_path(harness):
+    # step 0: [MIXED, SAME] keeps 1, resample [MIXED, MIXED] -> full; step 1: [MIXED, SAME] runs out
+    harness.train = prompts(MIXED, SAME, MIXED, MIXED, MIXED, SAME)
+    run(make_config())
+    tc = harness.tc
+    assert tc.calls.count("forward_backward") == 1
+    assert tc.fwd_bwd[0] == {
+        "n": 4,  # 2 groups x K=2, one record each
+        "loss_fn": "ppo",
+        "loss_fn_config": {"clip_low_threshold": 0.8, "clip_high_threshold": 1.28},
+    }
+    assert tc.calls[-3:] == ["rebind", "save_state", "stop"]
+    assert harness.server.closed and harness.client.closed
+
+
+def test_train_passes_configured_clip_bounds(harness):
+    harness.train = prompts(MIXED, MIXED)
+    run(make_config(clip_low_threshold=0.7, clip_high_threshold=1.4, loss="cispo"))
+    assert harness.tc.fwd_bwd[0]["loss_fn"] == "cispo"
+    assert harness.tc.fwd_bwd[0]["loss_fn_config"] == {"clip_low_threshold": 0.7, "clip_high_threshold": 1.4}
+
+
+def test_train_update_order_per_step(harness):
+    harness.train = prompts(MIXED, MIXED, MIXED, MIXED)
+    run(make_config())
+    assert harness.tc.calls == ["forward_backward", "optim_step", "rebind"] * 2 + ["save_state", "stop"]
+
+
+def test_train_advantages_follow_rewards(harness):
+    harness.train = prompts([1.0, 0.0], [0.0, 1.0])
+    run(make_config())
+    per_datum = [
+        {
+            round(a, 6)
+            for a, w in zip(d["lossFnInputs"]["advantages"]["data"], d["lossFnInputs"]["weights"]["data"], strict=True)
+            if w
+        }
+        for d in harness.tc.datums[0]
+    ]
+    # GRPO-normalized within each group: the 1.0 rollout gets +1, the 0.0 rollout gets -1
+    assert per_datum == [{1.0}, {-1.0}, {-1.0}, {1.0}]
+
+
+def test_train_dataset_smaller_than_batch_never_creates_session(harness):
+    harness.train = prompts(MIXED)
+    with pytest.raises(ValueError, match="fewer than batch_size"):
+        run(make_config())
+    assert harness.created == 0  # nothing billed
+
+
+def test_train_insufficient_groups_saves_then_stops(harness):
+    harness.train = prompts(*[SAME] * 20)
+    with pytest.raises(td.InsufficientGroupsError):
+        run(make_config(max_num_gen_batches=2))
+    assert harness.tc.calls == ["save_state", "stop"]  # abort checkpoint, then stop billing
+    assert harness.server.closed and harness.client.closed
+
+
+def test_train_other_errors_do_not_checkpoint(harness):
+    harness.train = prompts(MIXED, MIXED, MIXED, MIXED)
+    harness.fail_on_call = 4  # first rollout of step 1
+    harness.fail_on_call_error = RuntimeError("rollout infrastructure broke")
+    with pytest.raises(RuntimeError, match="infrastructure"):
+        run(make_config())
+    assert "save_state" not in harness.tc.calls
+    assert harness.tc.calls[-1] == "stop"
+
+
+def test_train_cancel_still_stops_session(harness):
+    harness.train = prompts(MIXED, MIXED)
+    harness.fail_on_call = 0
+    harness.fail_on_call_error = asyncio.CancelledError()  # what the SIGTERM handler triggers
+    with pytest.raises(asyncio.CancelledError):
+        run(make_config())
+    assert harness.tc.calls == ["stop"]
+
+
+def test_train_skips_update_when_no_datums(harness):
+    harness.train = prompts(MIXED, MIXED, untrainable=True)
+    run(make_config())
+    assert harness.tc.calls == ["save_state", "stop"]
+
+
+# ---------------------------------------------------------------------------
+# eval schedule
+# ---------------------------------------------------------------------------
+def eval_steps(caplog) -> list[int]:
+    return [int(m) for m in re.findall(r"step=(\d+)\)", caplog.text)]
+
+
+def test_eval_schedule_and_final_eval(harness, caplog):
+    caplog.set_level(logging.INFO, logger=td.logger.name)
+    harness.train = prompts(*[MIXED] * 6)
+    harness.eval = [{"eval": 1.0}]
+    run(make_config(eval_every=2, eval_dataset_path="eval.jsonl"))
+    assert eval_steps(caplog) == [2, 3]  # every 2 steps, then the final step
+
+
+def test_eval_uses_eval_temperature(harness):
+    harness.train = prompts(MIXED, MIXED)
+    harness.eval = [{"eval": 1.0}]
+    run(make_config(eval_dataset_path="eval.jsonl", temperature=1.0, eval_temperature=0.0))
+    assert harness.evals == [0.0]
+
+
+def test_final_eval_runs_when_no_step_ran(harness, caplog):
+    caplog.set_level(logging.INFO, logger=td.logger.name)
+    harness.train = prompts(SAME, SAME, MIXED)  # step 0 filters everything, then runs out
+    harness.eval = [{"eval": 1.0}]
+    run(make_config(eval_every=2, eval_dataset_path="eval.jsonl"))
+    assert "forward_backward" not in harness.tc.calls
+    assert eval_steps(caplog) == [0]  # 0 % eval_every == 0 must not skip it
+
+
+def test_final_eval_runs_when_last_step_was_skipped(harness, caplog):
+    caplog.set_level(logging.INFO, logger=td.logger.name)
+    harness.train = prompts(MIXED, MIXED) + prompts(MIXED, MIXED, untrainable=True)
+    run_config = make_config(eval_every=2, eval_dataset_path="eval.jsonl")
+    harness.eval = [{"eval": 1.0}]
+    run(run_config)
+    assert harness.tc.calls.count("forward_backward") == 1
+    assert eval_steps(caplog) == [2]  # step 1 skipped (no datums) -> global_step 2, no in-loop eval
+
+
+def test_final_eval_not_repeated_after_in_loop_eval(harness, caplog):
+    caplog.set_level(logging.INFO, logger=td.logger.name)
+    harness.train = prompts(*[MIXED] * 4)
+    harness.eval = [{"eval": 1.0}]
+    run(make_config(eval_every=2, eval_dataset_path="eval.jsonl"))
+    assert eval_steps(caplog) == [2]
+
+
+# ---------------------------------------------------------------------------
+# sampling log line
+# ---------------------------------------------------------------------------
+def sampling_lines(caplog) -> list[str]:
+    return [r.getMessage() for r in caplog.records if r.name == td.logger.name and "sampling:" in r.getMessage()]
+
+
+def test_sampling_line_logged_every_step(harness, caplog):
+    caplog.set_level(logging.INFO, logger=td.logger.name)
+    harness.train = prompts(MIXED, SAME, MIXED, MIXED, MIXED, SAME)
+    run(make_config())
+    lines = sampling_lines(caplog)
+    assert len(lines) == 2  # the step that ran out of prompts is logged too
+    assert re.search(
+        r"gen_batches=2 seen=4 filtered=1 \(25%\) kept=3 empty_rollouts=0 reward=0\.6250 rollout_time=\d+s$", lines[0]
+    )
+    assert re.search(
+        r"gen_batches=1 seen=2 filtered=1 \(50%\) kept=1 empty_rollouts=0 reward=0\.7500 rollout_time=\d+s$", lines[1]
+    )