From ac1a8ba0999c2ead81c72c6e7722cd1d29028105 Mon Sep 17 00:00:00 2001 From: Linbo Liu Date: Fri, 31 Jul 2026 22:16:40 +0000 Subject: [PATCH] feat(slime): Add integration for slime + rollout gateway. --- .gitignore | 6 +- AGENTS.md | 118 +++- examples/strands_appworld_agent/rl_app.py | 7 +- examples/strands_migration_agent/rl_app.py | 5 +- examples/strands_officebench_agent/rl_app.py | 7 +- .../run_local_eval.py | 3 +- .../strands_officebench_agent/test_local.py | 3 +- examples/strands_taubench_agent/rl_app.py | 5 +- .../backends/experimental/slime/SETUP.md | 162 +++++ .../backends/experimental/slime/__init__.py | 0 .../examples/math_agent/.wandb.env.example | 23 + .../examples/math_agent/config.yaml.example | 27 + .../slime/examples/math_agent/train.sh | 171 +++++ .../slime/integration/__init__.py | 0 .../experimental/slime/integration/rewards.py | 102 +++ .../experimental/slime/integration/rollout.py | 639 ++++++++++++++++++ .../slime/integration/sglang_parsing.py | 105 +++ .../slime/scripts/install_slime.sh | 101 +++ 18 files changed, 1455 insertions(+), 29 deletions(-) create mode 100644 src/agentcore_rl_toolkit/backends/experimental/slime/SETUP.md create mode 100644 src/agentcore_rl_toolkit/backends/experimental/slime/__init__.py create mode 100644 src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/.wandb.env.example create mode 100644 src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/config.yaml.example create mode 100755 src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/train.sh create mode 100644 src/agentcore_rl_toolkit/backends/experimental/slime/integration/__init__.py create mode 100644 src/agentcore_rl_toolkit/backends/experimental/slime/integration/rewards.py create mode 100644 src/agentcore_rl_toolkit/backends/experimental/slime/integration/rollout.py create mode 100644 src/agentcore_rl_toolkit/backends/experimental/slime/integration/sglang_parsing.py create mode 100644 src/agentcore_rl_toolkit/backends/experimental/slime/scripts/install_slime.sh diff --git a/.gitignore b/.gitignore index 62a8d0f..6fcb72e 100644 --- a/.gitignore +++ b/.gitignore @@ -18,9 +18,11 @@ build/ # Logs and experiment tracking wandb/ *.log +*.jsonl -# Local slime backend config (contains ARN/bucket; commit config.yaml.example only) -src/agentcore_rl_toolkit/backends/slime/examples/**/config.yaml +# Local backend config (contains ARN/bucket; commit config.yaml.example only). +# Covers backends/slime and backends/experimental/*. +src/agentcore_rl_toolkit/backends/**/examples/**/config.yaml # Training/example run artifacts, wherever generated. Convention-level patterns: # hydra run dirs, verl checkpoint dirs (scripts use CKPTS_DIR=exp_*), wandb run diff --git a/AGENTS.md b/AGENTS.md index 38e0ff6..07b0e21 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -12,6 +12,8 @@ This document provides context, patterns, and guidelines for AI coding assistant - [Background: BedrockAgentCoreApp](#background-bedrockagentcoreapp) - [What agentcore-rl-toolkit Provides](#what-agentcore-rl-toolkit-provides) - [Rollout Gateway](#rollout-gateway) + - [Experimental verl backend](#experimental-verl-backend-backendsexperimentalverl) + - [Experimental slime backend](#experimental-slime-backend-backendsexperimentalslime) - [Sandbox SDK](#sandbox-sdk) - [Migration Guide (basic_app → rl_app)](#migration-guide-basic_app--rl_app) - [Deployment to ACR](#deployment-to-acr) @@ -53,6 +55,7 @@ cd examples/strands_math_agent && uv sync && uv run python rl_app.py | `src/agentcore_rl_toolkit/reward_function.py` | `RewardFunction` base class | | `src/agentcore_rl_toolkit/rollout_gateway/` | In-repo token-level trajectory capture layer: `RolloutGateway`, `Renderer`, `SamplingBackend`, `TraceRecord` (see [Rollout Gateway](#rollout-gateway)) | | `src/agentcore_rl_toolkit/backends/experimental/verl/` | Experimental verl backend: `AgentCoreAgentLoop` plugged into stock verl main_ppo via the rollout gateway (successor to `backends/verl`) | +| `src/agentcore_rl_toolkit/backends/experimental/slime/` | Experimental slime backend: custom rollout function + episode-level reward normalization plugged into stock slime `train.py` via the rollout gateway (successor to `backends/slime`) | | `src/agentcore_rl_toolkit/sandbox/` | Sandbox SDK: `SandboxClient`, `Sandbox`, `ExecResult` — run shell commands in arbitrary images on ACR (see [Sandbox SDK](#sandbox-sdk)) | | `sandboxd/` | Go health shim (`agentcore-sandboxd`) that makes arbitrary Docker images satisfy the ACR container contract | | `examples/strands_math_agent/` | GSM8K math agent example | @@ -83,12 +86,21 @@ agentcore-rl-toolkit/ │ │ ├── server.py # ThreadedGatewayServer — serve the gateway from sync trainers │ │ ├── adapters/ # OpenAI + Anthropic wire protocol adapters │ │ └── sampling_backends/ # SamplingBackend impls (vLLM/SGLang HTTP, Tinker SDK) -│ └── backends/experimental/verl/ # Experimental verl backend on the rollout gateway -│ ├── sampling_backend.py # VerlSamplingBackend over verl's LLMServerClient -│ ├── gateway_host.py # per-worker-process gateway singleton (threaded aiohttp) -│ ├── agent_loop.py # AgentCoreAgentLoop (verl custom agent loop) -│ ├── dataset.py # PayloadDataset (payload-column dataset contract) -│ └── examples/ # GSM8K example: run script + per-run agent-loop YAML +│ └── backends/experimental/ +│ ├── verl/ # Experimental verl backend on the rollout gateway +│ │ ├── sampling_backend.py # VerlSamplingBackend over verl's LLMServerClient +│ │ ├── gateway_host.py # per-worker-process gateway singleton (threaded aiohttp) +│ │ ├── agent_loop.py # AgentCoreAgentLoop (verl custom agent loop) +│ │ ├── dataset.py # PayloadDataset (payload-column dataset contract) +│ │ └── examples/ # GSM8K example: run script + per-run agent-loop YAML +│ └── slime/ # Experimental slime backend on the rollout gateway +│ ├── SETUP.md # env install + train/eval walkthrough, tested versions +│ ├── integration/ +│ │ ├── rollout.py # generate_rollout (slime --rollout-function-path) +│ │ ├── rewards.py # normalize_episode_rewards (episode-level GRPO) +│ │ └── sglang_parsing.py # derender seams built from SGLang's own detectors +│ ├── scripts/install_slime.sh # bare-metal slime install (CUDA 13) +│ └── examples/math_agent/ # GSM8K example: train.sh + config templates ├── examples/ │ ├── strands_math_agent/ # GSM8K example │ │ ├── .bedrock_agentcore/ # Dockerfiles for deployment @@ -318,11 +330,12 @@ drains the tree into `list[TraceRecord]`. - `pip install agentcore-rl-toolkit[gateway]` → `aiohttp` + `transformers`. - Tool/reasoning parsing defaults to a dependency-free regex (the `` XML format) and `` split; the gateway itself never imports an inference engine. - For any other model format (e.g. Qwen3's JSON ``), inject an `output_parser` - callable (`(raw_output, tools_schema) -> ParsedOutput`) into `HfTemplateRenderer` — it - replaces the built-in derender entirely. The slime backend ships such a parser built from - SGLang's own detectors (`backends/slime/integration/sglang_parsing.py`, composing - `FunctionCallParser` + `ReasoningParser`) and wires it automatically from slime's + For any other model format (e.g. Qwen3's JSON ``), inject `tool_parser` / + `reasoning_parser` callables into `HfTemplateRenderer` — each replaces that derender + stage, and the two are independent (leave one unset and it keeps the built-in default). + The experimental slime backend ships such parsers built from SGLang's own detectors + (`backends/experimental/slime/integration/sglang_parsing.py`, wrapping + `FunctionCallParser` + `ReasoningParser`) and wires them automatically from slime's `--sglang-tool-call-parser` / `--sglang-reasoning-parser` args (names must match the served model); sglang is always importable there because the trainer serves SGLang. - For the Tinker backend (`TinkerSdkBackend` + `TinkerRenderer`), install `tinker` and @@ -333,10 +346,13 @@ The core (`TraceRecord`, `TrajectoryManager`, `Renderer` protocol, `SamplingBack protocol) imports torch-free and aiohttp-free; `RolloutGateway` is exposed lazily so importing the package never requires aiohttp. Tests live in `tests/rollout_gateway/`. -**Status.** The capture layer above is implemented and tested. Its first training-backend -consumer is the **experimental verl backend** (`backends/experimental/verl/`, see below); -other backends' dispatch/reward-join glue is not yet on the main branch — a prototype -dispatcher is parked on the `wip/online-rl-dispatch` branch. +**Status.** The capture layer above is implemented and tested. Its training-backend +consumers are the **experimental verl backend** (`backends/experimental/verl/`) and the +**experimental slime backend** (`backends/experimental/slime/`), both described below and +both validated end to end on GSM8K. The cross-agent dispatch/reward-join glue (routing one +batch across multiple agents, stamping a shared `rollout_id` across sub-agent sessions) is +not yet on the main branch — a prototype dispatcher is parked on the +`wip/online-rl-dispatch` branch. ### Experimental verl backend (`backends/experimental/verl/`) @@ -375,13 +391,70 @@ Key pieces (see `backends/experimental/verl/README.md` for the full design): doesn't provide. - Agent-side contract: the app sets `api_key = context.session_id or "EMPTY"` so the gateway can key capture off the Bearer/api-key slot (`"EMPTY"` keeps local runs and - the legacy per-session-URL gateways working). Adopted by `strands_math_agent` - (validated end to end); the other examples still send `"EMPTY"` and migrate as - they're validated against this backend. + the legacy per-session-URL gateways working). `strands_math_agent` does this and is + validated end to end; the other examples read the key from `_rollout["api_key"]` + instead (see the slime backend below) — equivalent for capture, since both backends + use the ACR session id as the gateway session id. - Validated end to end: `examples/math_agent/fsdp_fft_sync_grpo.sh` (GRPO, Qwen3-4B full-FT, TIS + KL trust region) reaches ~0.93 GSM8K val reward in one epoch against a live ACR agent. +### Experimental slime backend (`backends/experimental/slime/`) + +The successor to `backends/slime` (which stays untouched — it still depends on the external +`rllm-model-gateway` — until this graduates). It plugs into **stock** `slime/train.py` +through slime's public extension points only (`--rollout-function-path`, +`--custom-reward-post-process-path`, `--custom-config-path`): no forked trainer, no custom +entrypoint. Setup, config reference, and the validated version pins live in +`backends/experimental/slime/SETUP.md`. + +Key pieces: +- `integration/rollout.py` — `generate_rollout`, slime's rollout-function hook. Serves one + `RolloutGateway` in-process via `ThreadedGatewayServer`, sampling token-in/token-out + through SGLang's native `/generate` (`SglangHttpBackend` against slime's router) and + rendering with the HF chat template of `--hf-checkpoint`. Per episode: create a gateway + session (sid = uuid4), invoke ACR via `RolloutClient.invoke_async`, await the S3 result, + drain the session into `TraceRecord`s, and convert each record to a slime `Sample` + (records arrive already merged + loss-masked, so conversion is direct). Training and eval + share one path — training pulls GRPO-grouped batches from `data_source`, eval reads + `args.eval_datasets` with their own resolved params (e.g. greedy at + `--eval-temperature 0`), cached per dataset to avoid re-tokenizing the JSONL every eval. + Episode failures never abort a batch: they yield a zero-gradient no-op Sample + (`loss_mask=[0]`) tagged `episode_error`, still counted in its GRPO group. +- `integration/rewards.py` — `normalize_episode_rewards`, replacing slime's default + reshape-based normalization (which assumes a fixed turn count per row). Aggregates rows to + one reward per episode by `(group_index, gateway_session_id)`, normalizes across episodes + within a task group, then writes the result back to every row — so a 3-turn success and a + 2-turn failure each count as one data point. Strategies are pluggable via + `reward_postprocessing` (`grpo` default, `identity`); padding rows (`group_index=-1`) skip. +- `integration/sglang_parsing.py` — see the parsing bullet under + [Rollout Gateway](#rollout-gateway) dependencies. +- Config: ACR pointers and tunables come from a `--custom-config-path` YAML that slime merges + into its args namespace (`SlimeArtConfig.from_args`); every field also honors an uppercase + env-var override. `examples/math_agent/config.yaml.example` is the template + (`config.yaml` itself is gitignored — it carries the runtime ARN and bucket). +- **Networking constraint:** the gateway binds to `args.sglang_router_ip`, not loopback, + because the ACR-deployed agent dials back into it on every LLM call. `gateway_port` + (default 9090) must be reachable from the ACR VPC. +- **Concurrency:** a shared semaphore caps in-flight episodes at `max_concurrent`. The + client's ACR TPS limiter only paces session *starts*, so without this cap a large batch + (e.g. a full 1319-prompt eval set) launches every episode at once — saturating the + gateway/router and S3 result polling until episodes miss `acr_timeout`, and + over-pressuring the colocated SGLang KV cache to token-pool exhaustion. +- Agent-side contract: same api-key-slot session identity as the verl backend, but sent + explicitly as `_rollout["api_key"]` (the rollout function passes the sid), so the app reads + `payload["_rollout"].get("api_key", "EMPTY")`. All `examples/*/rl_app.py` read this; + `strands_math_agent` then overrides it with `context.session_id`, which is the same value + because the rollout function uses one uuid as both the ACR session id and the gateway sid. +- Validated end to end: `examples/math_agent/train.sh` (GSM8K GRPO, Qwen3, colocated + train+rollout on 8×B200) against a live ACR deployment of `strands_math_agent`. +- **CUDA 13 only.** `scripts/install_slime.sh` plus the env-pinning preamble in `train.sh` + work around several cu12/cu13 collisions (TE's vendored cuDNN Frontend probing for + `libcudart.so.12`, cuDNN main/sublib version mismatch, slime's hardcoded cu12 + `torch_memory_saver` `.so`, and slime dropping the pinned paths from Megatron actors' + `runtime_env`). Each workaround is commented inline where it lives — read those before + touching the library paths. + **Vendored from slime (upstream baselines).** Several files are adapted from [slime](https://github.com/THUDM/slime) (Apache-2.0; see `NOTICE`). To check what changed upstream before re-syncing, diff the source file against the baseline commit below: @@ -479,7 +552,7 @@ See `examples/strands_math_agent` for a complete example adapting from `basic_ap - Model config (`base_url`, `model_id`) comes from the `_rollout` payload, not environment variables - Optional `sampling_params` (e.g., `max_completion_tokens`, `temperature`) can also be passed via `_rollout` for training-engine-controlled generation settings - Use standard `OpenAIModel` — no custom model wrappers needed. For evaluation, `base_url` can point directly to any OpenAI-compatible endpoint (vLLM, SGLang, LiteLLM, etc.), or you can use `BedrockModel` directly -- `api_key` is set from `context.session_id` (the ACR runtime session id, available when the entrypoint declares a second `context` parameter) — trajectory-capture gateways like the experimental verl backend key token capture off the api-key slot. Fall back to `"EMPTY"` (the standard vLLM convention for unauthenticated servers) for local runs and gateways with per-session URLs, which ignore the api key +- `api_key` carries the trajectory-capture session key: gateways key token capture off the api-key / Bearer slot. Two equivalent sources — `context.session_id` (the ACR runtime session id, available when the entrypoint declares a second `context` parameter), which the experimental verl backend uses since it sets the ACR session id itself, or `payload["_rollout"]["api_key"]`, which the experimental slime backend passes explicitly. Fall back to `"EMPTY"` (the standard vLLM convention for unauthenticated servers) for local runs and gateways with per-session URLs, which ignore the api key - Model and agent are created per-invocation inside the entrypoint - This gives flexibility for the training engine to pass runtime configuration (inference address, sampling parameters, system prompt, etc.) to accommodate different learning scenarios - This is safe because RL rollouts are single-invocation — the agent doesn't need persistent conversation history across requests, so there's no need to keep model/agent as global state @@ -650,6 +723,13 @@ they run in `.github/workflows/experimental-verl-integration.yml`, which syncs t `verl-experimental-ci` extra (same pinned verl, but CPU torch and no vllm/flash-attn — the LLM server client is the faked seam, so no inference engine is needed). +The experimental slime backend has no automated tests yet — its seams (`Sample`, +`RolloutFnTrainOutput`, Megatron shapes) are only meaningful against an installed slime, +which the bare-metal CUDA 13 install makes impractical in CI. Verify changes with an +end-to-end run per `backends/experimental/slime/SETUP.md`. The pure functions +(`_record_to_sample`, `normalize_episode_rewards`) are the natural first unit tests, using +faked slime types the way the verl CI job fakes its seam. + ### Building and Pushing Docker Images ```bash diff --git a/examples/strands_appworld_agent/rl_app.py b/examples/strands_appworld_agent/rl_app.py index 77f8069..db1c6d5 100644 --- a/examples/strands_appworld_agent/rl_app.py +++ b/examples/strands_appworld_agent/rl_app.py @@ -95,8 +95,10 @@ def invoke_agent(payload: dict): rollout_config = payload.get("_rollout", {}) + # During training the rollout gateway keys the session off the api-key slot; + # "EMPTY" (the vLLM convention) is fine for plain evaluation endpoints. model = OpenAIModel( - client_args={"api_key": "EMPTY", "base_url": rollout_config["base_url"]}, + client_args={"api_key": rollout_config.get("api_key", "EMPTY"), "base_url": rollout_config["base_url"]}, model_id=rollout_config["model_id"], params=rollout_config.get("sampling_params", {}), ) @@ -157,7 +159,8 @@ def execute(code: str) -> str: ) response = agent(user_message) - logger.info(f"Agent response: {response.message['content'][0]['text']}") + content = response.message.get("content") or [] + logger.info("Agent response: %s", "".join(b["text"] for b in content if "text" in b)) # Save state and evaluate world.save() diff --git a/examples/strands_migration_agent/rl_app.py b/examples/strands_migration_agent/rl_app.py index 81608a2..32a58cd 100644 --- a/examples/strands_migration_agent/rl_app.py +++ b/examples/strands_migration_agent/rl_app.py @@ -46,6 +46,9 @@ def invoke_agent(payload: dict): base_url = payload["_rollout"]["base_url"] model_id = payload["_rollout"]["model_id"] params = payload["_rollout"].get("sampling_params", {}) + # During training the rollout gateway keys the session off the api-key slot; + # "EMPTY" (the vLLM convention) is fine for plain evaluation endpoints. + api_key = payload["_rollout"].get("api_key", "EMPTY") tools = [shell, editor] request = InvocationRequest(**payload) @@ -76,7 +79,7 @@ def invoke_agent(payload: dict): ) tools.append(search_dependency_version) - model = OpenAIModel(client_args={"api_key": "EMPTY", "base_url": base_url}, model_id=model_id, params=params) + model = OpenAIModel(client_args={"api_key": api_key, "base_url": base_url}, model_id=model_id, params=params) agent = Agent( model=model, diff --git a/examples/strands_officebench_agent/rl_app.py b/examples/strands_officebench_agent/rl_app.py index d2dc127..a6e46b8 100644 --- a/examples/strands_officebench_agent/rl_app.py +++ b/examples/strands_officebench_agent/rl_app.py @@ -41,8 +41,10 @@ def invoke_agent(payload: dict): # Choose model based on config if rollout_config.get("base_url"): + # During training the rollout gateway keys the session off the api-key + # slot; "EMPTY" (the vLLM convention) is fine for plain evaluation endpoints. model = OpenAIModel( - client_args={"api_key": "EMPTY", "base_url": rollout_config["base_url"]}, + client_args={"api_key": rollout_config.get("api_key", "EMPTY"), "base_url": rollout_config["base_url"]}, model_id=rollout_config["model_id"], params=rollout_config.get("sampling_params", {}), ) @@ -92,7 +94,8 @@ def invoke_agent(payload: dict): logger.info(f"Task: {user_input}") response = agent(user_input) - logger.info(f"Agent response: {response.message['content'][0]['text']}") + content = response.message.get("content") or [] + logger.info("Agent response: %s", "".join(b["text"] for b in content if "text" in b)) # Collect full conversation history messages = [{"role": msg.get("role", "unknown"), "content": msg.get("content", [])} for msg in agent.messages] diff --git a/examples/strands_officebench_agent/run_local_eval.py b/examples/strands_officebench_agent/run_local_eval.py index 0376a76..0e73960 100644 --- a/examples/strands_officebench_agent/run_local_eval.py +++ b/examples/strands_officebench_agent/run_local_eval.py @@ -146,7 +146,8 @@ def run_single_task(task_id, subtask_id, model_id): ) response = agent(task_config["task"]) - response_text = response.message["content"][0]["text"] + content = response.message.get("content") or [] + response_text = "".join(b["text"] for b in content if "text" in b) reward = reward_fn(testbed_dir=TESTBED_DIR, evaluation_config=task_config["evaluation"]) diff --git a/examples/strands_officebench_agent/test_local.py b/examples/strands_officebench_agent/test_local.py index 1496cec..4d42e50 100644 --- a/examples/strands_officebench_agent/test_local.py +++ b/examples/strands_officebench_agent/test_local.py @@ -109,7 +109,8 @@ def main(): ) response = agent(task_config["task"]) - response_text = response.message["content"][0]["text"] + content = response.message.get("content") or [] + response_text = "".join(b["text"] for b in content if "text" in b) logger.info(f"Agent response: {response_text}") reward_fn = OfficeBenchReward() diff --git a/examples/strands_taubench_agent/rl_app.py b/examples/strands_taubench_agent/rl_app.py index 7481341..7e52d5d 100644 --- a/examples/strands_taubench_agent/rl_app.py +++ b/examples/strands_taubench_agent/rl_app.py @@ -80,6 +80,9 @@ def _setup_rollout(payload: dict) -> RolloutContext: """ base_url = payload["_rollout"]["base_url"] model_id = payload["_rollout"]["model_id"] + # During training the rollout gateway keys the session off the api-key slot; + # "EMPTY" (the vLLM convention) is fine for plain evaluation endpoints. + api_key = payload["_rollout"].get("api_key", "EMPTY") # Copy so we don't mutate the caller's payload while applying defaults. params = copy.deepcopy(payload["_rollout"].get("sampling_params", {})) for k, v in ASSISTANT_MODEL_DEFAULTS.items(): @@ -95,7 +98,7 @@ def _setup_rollout(payload: dict) -> RolloutContext: # Assistant model (RL-trained, served via vLLM) assistant_model = OpenAIModel( - client_args={"api_key": "EMPTY", "base_url": base_url}, + client_args={"api_key": api_key, "base_url": base_url}, model_id=model_id, params=params, ) diff --git a/src/agentcore_rl_toolkit/backends/experimental/slime/SETUP.md b/src/agentcore_rl_toolkit/backends/experimental/slime/SETUP.md new file mode 100644 index 0000000..ba8c6ea --- /dev/null +++ b/src/agentcore_rl_toolkit/backends/experimental/slime/SETUP.md @@ -0,0 +1,162 @@ +# Experimental Slime Backend Setup Guide + +How to train an ACR-deployed agent with the [slime](https://github.com/THUDM/slime) +training backend, using `agentcore_rl_toolkit.rollout_gateway`. + +The guide contains: + +- **[Part 1 — Slime environment](#part-1--slime-environment)**: get a working + slime runtime plus this toolkit, via the bare-metal install script. +- **[Part 2 — Run training and evaluation](#part-2--run-training-and-evaluation)**: + deploy the agent, prepare data, configure `config.yaml`, run `train.sh`, evaluate. +- **[Tested Versions](#tested-versions)** pins the exact environment this was + validated against — check here if you want to reproduce our results. + +--- + +## Prerequisites + +- Hardware requirements: see + [slime's README](https://github.com/THUDM/slime#installation) for tested GPU + configurations per model size. The defaults in `train.sh` target a single + 8-GPU node. +- A GPU cluster with **CUDA 13** installed (`/usr/local/cuda-13.0` by default; + export `CUDA_HOME` if yours differs). +- Python==3.12 and [`uv`](https://docs.astral.sh/uv/). +- AWS credentials with permission to invoke an ACR runtime and read/write an S3 + bucket (`aws sts get-caller-identity` works). +- An ACR deployment of your agent — `rl_app.py` configured and deployed per + [`examples/strands_math_agent/README.md`](../../../../../examples/strands_math_agent/README.md). +- **Network**: the training node's gateway port (`gateway_port`, default 9090) + and the SGLang router port must be reachable *from the ACR VPC* — the deployed + agent dials back into the gateway on every LLM call. Loopback won't work. + +--- + +## Part 1 — Slime environment - Bare-metal + +There is one supported path: the bare-metal install script. + +```bash +# From a clone of this repo, inside your activated python environment +cd /path/to/agentcore-rl-toolkit + +# Install the toolkit with the rollout-gateway extras (aiohttp + transformers). +# NOTE: this backend does NOT need the [slime] extra — that one pulls +# rllm-model-gateway, which only the legacy backends/slime uses. +uv pip install -e ".[gateway]" + +export CUDA_HOME=/usr/local/cuda-13.0 +bash src/agentcore_rl_toolkit/backends/experimental/slime/scripts/install_slime.sh +``` + +Notes: + +- Expect a long build (the flash-attn / TE / apex source compiles dominate). +- Point `SLIME_DIR` (step 2.4) at the `slime` directory the script cloned. + +--- + +## Part 2 — Run training and evaluation + +### 2.1 Deploy the agent to ACR + +Follow the "Run RL App Hosted on ACR" section in +[`examples/strands_math_agent/README.md`](../../../../../examples/strands_math_agent/README.md) +— it covers the `agentcore configure` / `agentcore deploy` flow plus VPC and IAM +setup. + +Save the resulting **runtime ARN** — they go into `config.yaml` in step 2.3. + +### 2.2 Download model and data + +```bash +python -c " +from huggingface_hub import snapshot_download +snapshot_download('Qwen/Qwen3-0.6B', local_dir='/path/to/Qwen3-0.6B') +" + +python -c " +from datasets import load_dataset +import json +ds = load_dataset('openai/gsm8k', 'main', split='train') +with open('/path/to/gsm8k_train.jsonl', 'w') as f: + for i, row in enumerate(ds): + question = row['question'] + answer = row['answer'].split('####')[-1].strip() + # Top-level 'prompt' is read by slime (tokenization, length filter). + # 'metadata' is the agent payload verbatim — shape it however the agent expects. + f.write(json.dumps({ + 'prompt': question, + 'metadata': {'prompt': question, 'answer': answer}, + }) + '\n') +" +``` + +### 2.3 Configure deployment settings + +ACR deployment pointers and toolkit tunables live in `config.yaml`, passed to +slime via `--custom-config-path`. Slime merges every key into its args namespace, +where the rollout function reads them (`SlimeArtConfig.from_args` in +`integration/rollout.py`; each field also honors an uppercase env-var override). + +```bash +cd src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent + +cp config.yaml.example config.yaml +# Edit config.yaml: +# agent_runtime_arn: "arn:aws:bedrock-agentcore:..." (from step 2.1) +# s3_bucket: "your-bucket-name" + +cp .wandb.env.example .wandb.env # optional; skip to disable wandb +# Edit .wandb.env: +# WANDB_API_KEY="..." +# WANDB_ENTITY="your-org" +``` + +| `config.yaml` key | Env override | Default | Meaning | +|---|---|---|---| +| `agent_runtime_arn` | `ACR_AGENT_RUNTIME_ARN` | *(required)* | ACR runtime to invoke | +| `s3_bucket` | `ACR_S3_BUCKET` | *(required)* | bucket the agent writes results to | +| `exp_id` | `EXP_ID` | `slime-training` | S3 key prefix for this experiment | +| `gateway_port` | `GATEWAY_PORT` | `9090` | in-process rollout gateway port (must be reachable from the ACR VPC) | +| `acr_timeout` | `ACR_TIMEOUT` | `900` | per-session ACR invocation timeout (s) | +| `model_id` | `MODEL_ID` | `default` | OpenAI model id served to the agent | +| `acr_tps_limit` | `ACR_TPS_LIMIT` | `25` | ACR invocation rate limit (paces session *starts*) | +| `max_concurrent` | `MAX_CONCURRENT` | `100` | max concurrent in-flight ACR sessions | +| `max_pool_connections` | `MAX_POOL_CONNECTIONS` | `10` | boto3 conn-pool size — caps *reused* connections, not concurrency. Below `max_concurrent` it only logs "Connection pool is full" warnings, which are not errors. | +| `reward_postprocessing` | `REWARD_POSTPROCESSING` | `grpo` | `grpo` (group-relative) or `identity` | + +### 2.4 Run training + +`train.sh` is the only entry point. First fill in the values in +`src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/train.sh` + +```bash +cd /path/to/agentcore-rl-toolkit/ + +bash src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/train.sh +``` + +--- + +## Tested Versions + +For reproducibility, here's the exact environment this integration was validated +against: + +| Component | Version / SHA | +|---|---| +| Instance type | 8 × NVIDIA B200 180GB | +| CUDA | `13.0` (driver 580.159.04) | +| Python | `3.12` | +| PyTorch | `2.11.0+cu130` | +| slime | commit `fa3c990af6f18efd3fd9922698bf4bf4048d1263` | +| SGLang | `0.5.13` (sglang-kernel `0.4.3`, sgl-deep-gemm `0.1.2`, sglang-router `0.3.2`) | +| Megatron-LM | commit `1dcf0dafa884ad52ffb243625717a3471643e087` + slime's `megatron.patch` | +| Megatron-Bridge | `0.5.0+6fde1c85` | +| TransformerEngine | `2.11.0` (`core-cu13`) | +| flash-attn | `2.8.3` | +| Apex | `10417ace` | +| transformers | `5.8.1` | +| numpy | `1.26.4` (`<2`, required by Megatron) | diff --git a/src/agentcore_rl_toolkit/backends/experimental/slime/__init__.py b/src/agentcore_rl_toolkit/backends/experimental/slime/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/.wandb.env.example b/src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/.wandb.env.example new file mode 100644 index 0000000..c1eb4ef --- /dev/null +++ b/src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/.wandb.env.example @@ -0,0 +1,23 @@ +# Wandb configuration — copy to .wandb.env (gitignored) and fill in values. +# +# train.sh sources .wandb.env if it exists and passes WANDB_API_KEY / +# WANDB_ENTITY into the Ray runtime env. When WANDB_API_KEY is set, the +# training job also enables slime's built-in wandb logger via +# --use-wandb --wandb-project ... --wandb-group gsm8k-grpo. +# +# If you don't want wandb logging, skip this file entirely — train.sh +# will fall back to no wandb. +# +# Usage: +# cp .wandb.env.example .wandb.env +# # edit .wandb.env with your API key / entity + +# Wandb API key. Find at https://wandb.ai/settings under "API keys". +export WANDB_API_KEY="" + +# Wandb entity (username or team) that owns the destination project. +export WANDB_ENTITY="" + +# (optional) Override the default project name. train.sh defaults to +# "slime-art" when unset. +# export WANDB_PROJECT="slime-art" diff --git a/src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/config.yaml.example b/src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/config.yaml.example new file mode 100644 index 0000000..2b97aa2 --- /dev/null +++ b/src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/config.yaml.example @@ -0,0 +1,27 @@ +# Toolkit + slime backend settings for the rollout function. +# Copy to config.yaml (gitignored) and fill in your deployment values. +# +# cp config.yaml.example config.yaml +# # edit agent_runtime_arn + s3_bucket +# +# Passed to slime via --custom-config-path config.yaml. +# All keys become attributes on slime's args namespace. + +# ACR deployment pointers (from `bedrock-agentcore-starter-toolkit launch`) +agent_runtime_arn: "arn:aws:bedrock-agentcore:::runtime/-" +s3_bucket: "your-s3-bucket-name" + +# Toolkit tunables +exp_id: "gsm8k-grpo-train" +gateway_port: 9090 # rollout gateway port +acr_timeout: 900 # per-session ACR invocation timeout (s) +model_id: "default" # OpenAI model id served to the agent +acr_tps_limit: 5 # ACR service TPS quota +max_concurrent: 100 # max concurrent ACR sessions + +# Number of network connections the AWS client keeps open (and reuses) to talk +# to ACR/S3. This is NOT max_concurrent: it caps reused connections, not how many +# sessions run at once. Too small value causes "Connection pool is full" warnings, +# they are NOT errors. Setting a too large value wastes resources (idle open sockets). +max_pool_connections: 10 +reward_postprocessing: "grpo" # "grpo" or "identity" diff --git a/src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/train.sh b/src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/train.sh new file mode 100755 index 0000000..f068e80 --- /dev/null +++ b/src/agentcore_rl_toolkit/backends/experimental/slime/examples/math_agent/train.sh @@ -0,0 +1,171 @@ +#!/bin/bash +# Train the strands math agent (examples/strands_math_agent/rl_app.py, deployed to +# ACR first) with slime GRPO: our rollout function submits tasks to ACR, captures +# per-turn token ids + logprobs via the rollout gateway, and feeds Megatron. +# +# Config: config.yaml (ACR ARN + tunables; cp from config.yaml.example), +# .wandb.env (optional wandb creds), env vars below. +# +# Usage: +# export SLIME_DIR=/root/slime \ +# MODEL_DIR=/path/to/Qwen3-0.6B \ +# TRAIN_DATA_PATH=/path/to/gsm8k_tiny.jsonl \ +# MODEL_TYPE=qwen3-0.6B +# bash train.sh +set -euo pipefail + +# === Paths (set these via env) === +SLIME_DIR="${SLIME_DIR:?Set SLIME_DIR (path to the slime repo, e.g. /root/slime)}" +MODEL_DIR="${MODEL_DIR:?Set MODEL_DIR (path to the HF model checkpoint)}" +TRAIN_DATA_PATH="${TRAIN_DATA_PATH:?Set TRAIN_DATA_PATH (path to the training JSONL)}" +VAL_DATA_PATH="${VAL_DATA_PATH:-${TRAIN_DATA_PATH}}" +MODEL_TYPE="${MODEL_TYPE:?Set MODEL_TYPE (slime model-arch name, e.g. qwen3-0.6B)}" +CONFIG="${CONFIG:-$(dirname "$0")/config.yaml}" + +# Set your cuda path. CUDA 13 only — cu12 is not supported. +CUDA_HOME="${CUDA_HOME:-/usr/local/cuda-13.0}" + +# Checkpoint output dir (cleared at start; comment the rm to resume). +CKPTS_DIR="${CKPTS_DIR:-checkpoints/exp_agentcore_grpo}" +rm -rf "${CKPTS_DIR}" + +# GPUs on this node; with --colocate this is also the train+rollout pool. +NUM_GPUS="${NUM_GPUS:-8}" + +# Optional wandb creds (WANDB_API_KEY / _ENTITY / _PROJECT) — never commit real keys. +[ -f "$(dirname "$0")/.wandb.env" ] && source "$(dirname "$0")/.wandb.env" + +# === Setup === +pkill -9 sglang 2>/dev/null || true +ray stop --force 2>/dev/null || true +sleep 3 + +# The cu12 decoy below only works for processes that inherit LD_LIBRARY_PATH at exec +# time (the loader reads it once at startup), so a leftover Ray cluster would run +# actors that never see it. Fail loudly now, not mid-training on "Multiple libcudart +# libraries found". +if pgrep -f 'raylet|gcs_server' >/dev/null 2>&1; then + echo "ERROR: Ray is still running after 'ray stop --force'." >&2 + echo " Kill it before rerunning: pkill -9 -f 'raylet|gcs_server|ray::'" >&2 + exit 1 +fi + +# slime/ray/actor_group.py hardcodes the cu12-named torch_memory_saver preload .so, +# but a cu13 build ships *_cu13.abi3.so. Bridge the filename (idempotent). +TMS_SP="$(python -c 'import os, torch_memory_saver; print(os.path.dirname(os.path.dirname(torch_memory_saver.__file__)))' 2>/dev/null || true)" +if [ -n "$TMS_SP" ] \ + && [ ! -e "$TMS_SP/torch_memory_saver_hook_mode_preload_cu12.abi3.so" ] \ + && [ -e "$TMS_SP/torch_memory_saver_hook_mode_preload_cu13.abi3.so" ]; then + ln -s "$TMS_SP/torch_memory_saver_hook_mode_preload_cu13.abi3.so" \ + "$TMS_SP/torch_memory_saver_hook_mode_preload_cu12.abi3.so" + echo "[setup] linked tms cu13 .so -> cu12 name for slime compatibility" +fi + +export CUDA_HOME +export PATH="${CUDA_HOME}/bin:${PATH}" +NVIDIA_LIBS=$(python -c "import sysconfig, os, glob; base=os.path.join(sysconfig.get_path('purelib'), 'nvidia'); print(':'.join(sorted(glob.glob(os.path.join(base, '*', 'lib')))))") +LD_LIBRARY_PATH="$(echo "${LD_LIBRARY_PATH:-}" | tr ':' '\n' | grep -vE '^/usr/local/cuda(-[0-9.]+)?(/|$)' | paste -sd ':' -)" +export LD_LIBRARY_PATH="${NVIDIA_LIBS}:${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}" + +# Block cu12's libcudart from loading alongside cu13: TE 2.11.0's vendored cuDNN +# Frontend probes both and raises "Multiple libcudart libraries found". Scrubbing +# LD_LIBRARY_PATH isn't enough — /etc/ld.so.cache still resolves cu12. An empty decoy +# at the FRONT of the path makes dlopen fail on that name ("file too short") without +# falling back to the cache. (Newer frontends just warn and honor +# CUDNN_FRONTEND_CUDART_LIB_NAME; TE's vendored copy does neither.) +# +# Path is fixed, not mktemp'd: Ray workers resolve it lazily (first cu12 pull is at +# fused_attn_fwd), so it must still exist whenever an actor gets there, across reruns. +DECOY_DIR="${DECOY_DIR:-/tmp/cudart-decoy-cu13-${USER}}" +mkdir -p "${DECOY_DIR}" +: > "${DECOY_DIR}/libcudart.so.12" +export LD_LIBRARY_PATH="${DECOY_DIR}:${LD_LIBRARY_PATH}" + +# Honored by the standalone cudnn-frontend (>=1.26), which then skips the probe. +export CUDNN_FRONTEND_CUDART_LIB_NAME=libcudart.so.13 + +# Pin cuDNN to the venv wheel. TE loads the main libcudnn by absolute path, globbing +# ${CUDNN_HOME}|${CUDNN_PATH}|${CUDA_HOME} in that order, while its sublibraries +# (libcudnn_graph/_engines_*) come from LD_LIBRARY_PATH — i.e. the wheel. Without +# these vars, CUDA_HOME's own cuDNN wins and the main/sublib version mismatch kills +# fused attention with CUDNN_STATUS_SUBLIBRARY_LOADING_FAILED. +CUDNN_HOME="$(python -c "import sysconfig, os; print(os.path.join(sysconfig.get_path('purelib'), 'nvidia', 'cudnn'))")" +export CUDNN_HOME +export CUDNN_PATH="${CUDNN_HOME}" + +export PYTHONUNBUFFERED=1 +ray start --head --num-gpus ${NUM_GPUS} --disable-usage-stats + +# Source model architecture args (populates MODEL_ARGS) +source ${SLIME_DIR}/scripts/models/${MODEL_TYPE}.sh + +# === Launch training === +export no_proxy=127.0.0.1 + +# Env forwarded to every Ray worker (ACR ARN + bucket come from config.yaml instead). +# WANDB_API_KEY is appended only when set, to avoid injecting an empty value. +RUNTIME_ENV_JSON="{\"env_vars\": {\"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", \"CUDA_HOME\": \"${CUDA_HOME}\", \"CUDNN_HOME\": \"${CUDNN_HOME}\", \"CUDNN_PATH\": \"${CUDNN_PATH}\", \"CUDNN_FRONTEND_CUDART_LIB_NAME\": \"${CUDNN_FRONTEND_CUDART_LIB_NAME}\", \"LD_LIBRARY_PATH\": \"${LD_LIBRARY_PATH}\"${WANDB_API_KEY:+, \"WANDB_API_KEY\": \"${WANDB_API_KEY}\"}}}" + +# slime gives the Megatron train actors their own runtime_env.env_vars +# (slime/ray/actor_group.py), which drops CUDA_HOME / CUDNN_* / LD_LIBRARY_PATH — those +# workers would fall back to the cluster CUDA and hit the failures above. +# --train-env-vars is merged into that actor env, so re-pin the paths there. +TRAIN_ENV_VARS_JSON="{\"CUDA_HOME\": \"${CUDA_HOME}\", \"CUDNN_HOME\": \"${CUDNN_HOME}\", \"CUDNN_PATH\": \"${CUDNN_PATH}\", \"CUDNN_FRONTEND_CUDART_LIB_NAME\": \"${CUDNN_FRONTEND_CUDART_LIB_NAME}\", \"LD_LIBRARY_PATH\": \"${LD_LIBRARY_PATH}\"}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 ${SLIME_DIR}/train.py \ + "${MODEL_ARGS[@]}" \ + --hf-checkpoint ${MODEL_DIR} \ + --ref-load ${MODEL_DIR} \ + --prompt-data ${TRAIN_DATA_PATH} \ + --eval-prompt-data gsm8k ${VAL_DATA_PATH} \ + --num-rollout 100 \ + --tensor-model-parallel-size 2 \ + --rollout-num-gpus-per-engine 2 \ + --input-key prompt \ + --rollout-batch-size 64 \ + --n-samples-per-prompt 4 \ + --num-steps-per-rollout 1 \ + --sglang-context-length 14336 \ + --max-tokens-per-gpu 14336 \ + --rollout-max-response-len 2048 \ + --eval-max-response-len 2048 \ + --rollout-temperature 1.0 \ + --eval-interval 10 \ + --eval-input-key prompt \ + --n-samples-per-eval-prompt 1 \ + --eval-temperature 0.0 \ + --advantage-estimator grpo \ + --use-kl-loss \ + --kl-loss-type low_var_kl \ + --eps-clip 0.2 \ + --eps-clip-high 0.28 \ + --lr 1e-6 \ + --lr-decay-style constant \ + --optimizer-cpu-offload \ + --overlap-cpu-optimizer-d2h-h2d \ + --use-precision-aware-optimizer \ + --sequence-parallel \ + --sglang-mem-fraction-static 0.6 \ + --sglang-cuda-graph-max-bs 32 \ + --sglang-tool-call-parser qwen \ + --sglang-log-level warning \ + --sglang-log-level-http warning \ + --accumulate-allreduce-grads-in-fp32 \ + --attention-softmax-in-fp32 \ + --attention-backend flash \ + --actor-num-gpus-per-node ${NUM_GPUS} \ + --colocate \ + --train-env-vars "${TRAIN_ENV_VARS_JSON}" \ + --megatron-to-hf-mode bridge \ + --rollout-function-path \ + agentcore_rl_toolkit.backends.experimental.slime.integration.rollout.generate_rollout \ + --custom-reward-post-process-path \ + agentcore_rl_toolkit.backends.experimental.slime.integration.rewards.normalize_episode_rewards \ + --custom-config-path ${CONFIG} \ + --use-dynamic-batch-size \ + --save ${CKPTS_DIR} \ + --save-interval 100 \ + --save-hf ${CKPTS_DIR}/hf/{rollout_id} \ + ${WANDB_API_KEY:+--use-wandb --wandb-project ${WANDB_PROJECT:-slime-art} --wandb-group gsm8k-slime-grpo} diff --git a/src/agentcore_rl_toolkit/backends/experimental/slime/integration/__init__.py b/src/agentcore_rl_toolkit/backends/experimental/slime/integration/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/agentcore_rl_toolkit/backends/experimental/slime/integration/rewards.py b/src/agentcore_rl_toolkit/backends/experimental/slime/integration/rewards.py new file mode 100644 index 0000000..040d840 --- /dev/null +++ b/src/agentcore_rl_toolkit/backends/experimental/slime/integration/rewards.py @@ -0,0 +1,102 @@ +"""Custom reward post-processing for multi-turn ACR episodes. + +Replaces slime's default reshape-based normalization with episode-level +normalization that handles variable turn counts. + +Provides pluggable normalization strategies. Default is GRPO (group-relative). + +Usage: + --custom-reward-post-process-path \ + agentcore_rl_toolkit.backends.experimental.slime.integration.rewards.normalize_episode_rewards +""" + +from collections import defaultdict + + +def _grpo_normalize(episode_rewards: list[float], std_normalize: bool) -> list[float]: + """GRPO: mean-center, optionally std-normalize across episodes in a task group.""" + n = len(episode_rewards) + mean = sum(episode_rewards) / n + centered = [r - mean for r in episode_rewards] + + if std_normalize: + variance = sum(c * c for c in centered) / n + std = variance**0.5 + if std > 1e-6: + centered = [c / (std + 1e-6) for c in centered] + else: + centered = [0.0] * n + + return centered + + +def _identity_normalize(episode_rewards: list[float], std_normalize: bool) -> list[float]: + """No normalization — pass raw rewards through.""" + return list(episode_rewards) + + +NORMALIZATION_STRATEGIES = { + "grpo": _grpo_normalize, + "identity": _identity_normalize, +} + + +def normalize_episode_rewards(args, samples): + """Normalize rewards at the episode level within each task group. + + Two-step process: + 1. Aggregate: group turns by (group_index, session_id) to get one reward + per episode. All turns in an episode share the same broadcast reward, + so we take the first turn's reward as the episode reward. + 2. Normalize: within each task group (group_index), apply the chosen + normalization strategy to the per-episode rewards. + 3. Assign: write the normalized episode reward back to all turns. + + This ensures episodes with different turn counts are weighted equally — + a 3-turn success and a 2-turn failure each count as one data point. + + Normalization strategy is selected via reward_postprocessing + (in config.yaml), default "grpo". Options: "grpo", "identity". + + Samples with group_index=-1 (dummy padding) are skipped. + + Args: + args: Slime argument namespace. + samples: Flat list of slime Samples. + + Returns: + Tuple of (raw_rewards, normalized_rewards) as lists of floats. + """ + strategy_name = getattr(args, "reward_postprocessing", "grpo") + normalize_fn = NORMALIZATION_STRATEGIES.get(strategy_name, _grpo_normalize) + std_normalize = getattr(args, "grpo_std_normalization", False) + + raw_rewards = [s.get_reward_value(args) for s in samples] + + # Step 1: Group samples by (group_index, session_id) to identify episodes + episodes = defaultdict(list) + for i, s in enumerate(samples): + if s.group_index == -1: + continue + session_id = s.metadata.get("gateway_session_id", "") if s.metadata else "" + episodes[(s.group_index, session_id)].append(i) + + # Step 2: Group episodes by task (group_index) + task_groups = defaultdict(list) + for (grp_idx, _session_id), sample_indices in episodes.items(): + episode_reward = raw_rewards[sample_indices[0]] + task_groups[grp_idx].append((sample_indices, episode_reward)) + + # Step 3: Normalize per-episode rewards within each task group + rewards = list(raw_rewards) + + for episode_list in task_groups.values(): + episode_rewards = [r for _, r in episode_list] + normalized = normalize_fn(episode_rewards, std_normalize) + + # Step 4: Assign normalized episode reward back to all turns + for (sample_indices, _), norm_reward in zip(episode_list, normalized, strict=True): + for idx in sample_indices: + rewards[idx] = norm_reward + + return raw_rewards, rewards diff --git a/src/agentcore_rl_toolkit/backends/experimental/slime/integration/rollout.py b/src/agentcore_rl_toolkit/backends/experimental/slime/integration/rollout.py new file mode 100644 index 0000000..af8defb --- /dev/null +++ b/src/agentcore_rl_toolkit/backends/experimental/slime/integration/rollout.py @@ -0,0 +1,639 @@ +"""Custom slime rollout function using ACR agents + the in-repo rollout gateway. + +:class:`agentcore_rl_toolkit.rollout_gateway.RolloutGateway` is served in-process +(an aiohttp app in a dedicated thread) and owns tokenization: it renders each turn +with the HF chat template of ``--hf-checkpoint``, samples token-in/token-out via +SGLang's native ``/generate`` (``SglangHttpBackend``), and linearizes each session's +message tree into loss-masked :class:`TraceRecord` rows — multi-turn prefix merging +happens inside the gateway, so this module only converts records to slime Samples. + +Session identity — api-key / Bearer slot: the gateway keys sessions off the +api-key slot of the agent's LLM client (what OpenAI/Anthropic SDKs — and harnesses +like Claude Code / Codex — forward on every request). Each episode sends its +session id as ``_rollout["api_key"]``; the rl_app plugs it into its model client +(``api_key=payload["_rollout"].get("api_key", "EMPTY")``), so every LLM call +arrives at the fixed gateway ``base_url`` tagged ``Authorization: Bearer ``. +No per-session URLs and no model wrapper are needed. + +Usage: + python -m slime.train \ + --rollout-function-path \ + agentcore_rl_toolkit.backends.experimental.slime.integration.rollout.generate_rollout \ + --custom-reward-post-process-path \ + agentcore_rl_toolkit.backends.experimental.slime.integration.rewards.normalize_episode_rewards \ + --custom-config-path config.yaml \ + --use-dynamic-batch-size \ + --max-tokens-per-gpu 9216 \ + ... + + Slime's own --sglang-tool-call-parser / --sglang-reasoning-parser args are + honored: when the tool-call parser is set, the gateway derenders model output + with SGLang's own detectors (see sglang_parsing.py; the names must match the + served model); when unset, the gateway's dependency-free built-in parsing is + used. + + Configuration via --custom-config-path YAML: + agent_runtime_arn: "arn:aws:bedrock-agentcore:..." + s3_bucket: "my-bucket" + exp_id: "slime-training" + gateway_port: 9090 # in-process rollout gateway port + acr_timeout: 900 # per-session ACR invocation timeout + model_id: "default" # OpenAI model id served to the agent + acr_tps_limit: 25 # ACR service TPS quota + max_concurrent: 100 # max concurrent ACR sessions + max_pool_connections: 100 # boto3 conn-pool size (>= max_concurrent) + reward_postprocessing: "grpo" # "grpo" or "identity" +""" + +import asyncio +import copy +import json +import logging +import os +import uuid +from argparse import Namespace +from dataclasses import dataclass +from pathlib import Path + +from agentcore_rl_toolkit.rollout_gateway import BaseTrace, Status + +logger = logging.getLogger(__name__) + +# File-based trace logging for debugging captured trajectories +_TRACE_LOG_PATH = Path(os.environ.get("TRACE_LOG", "trace_log.jsonl")) + +# Module-level singletons (initialized on first call, reused across rollout steps) +_gateway_server = None +_client = None +_config = None + +# Cache of eval datasets keyed by EvalDatasetConfig.cache_key, so repeated +# evaluations don't re-read + re-tokenize the same JSONL every rollout. +_eval_datasets: dict = {} + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + + +@dataclass +class SlimeArtConfig: + """Configuration for ACR-based rollouts with slime. + + All fields come from slime's args namespace via --custom-config-path YAML + (see module docstring). Env var fallbacks are provided as an override path + for CI/dev convenience. + """ + + agent_runtime_arn: str = "" + s3_bucket: str = "" + exp_id: str = "slime-training" + gateway_port: int = 9090 + acr_timeout: float = 900.0 + model_id: str = "default" + acr_tps_limit: int = 25 + max_concurrent: int = 100 + max_pool_connections: int = 10 + reward_postprocessing: str = "grpo" + sglang_tool_call_parser: str | None = None + sglang_reasoning_parser: str | None = None + + @classmethod + def from_args(cls, args: Namespace) -> "SlimeArtConfig": + """Build config from slime args, falling back to env vars then defaults.""" + + def _get(attr: str, env: str, default): + val = getattr(args, attr, None) + if val is not None and val != "" and val != default: + return val + return os.environ.get(env, default) + + return cls( + agent_runtime_arn=_get("agent_runtime_arn", "ACR_AGENT_RUNTIME_ARN", cls.agent_runtime_arn), + s3_bucket=_get("s3_bucket", "ACR_S3_BUCKET", cls.s3_bucket), + exp_id=_get("exp_id", "EXP_ID", cls.exp_id), + gateway_port=int(_get("gateway_port", "GATEWAY_PORT", cls.gateway_port)), + acr_timeout=float(_get("acr_timeout", "ACR_TIMEOUT", cls.acr_timeout)), + model_id=_get("model_id", "MODEL_ID", cls.model_id), + acr_tps_limit=int(_get("acr_tps_limit", "ACR_TPS_LIMIT", cls.acr_tps_limit)), + max_concurrent=int(_get("max_concurrent", "MAX_CONCURRENT", cls.max_concurrent)), + max_pool_connections=int(_get("max_pool_connections", "MAX_POOL_CONNECTIONS", cls.max_pool_connections)), + reward_postprocessing=_get("reward_postprocessing", "REWARD_POSTPROCESSING", cls.reward_postprocessing), + # Read slime's own SGLang server args (args.sglang_tool_call_parser / + # args.sglang_reasoning_parser) so the gateway parses identically. + sglang_tool_call_parser=_get( + "sglang_tool_call_parser", "SGLANG_TOOL_CALL_PARSER", cls.sglang_tool_call_parser + ), + sglang_reasoning_parser=_get( + "sglang_reasoning_parser", "SGLANG_REASONING_PARSER", cls.sglang_reasoning_parser + ), + ) + + +# --------------------------------------------------------------------------- +# Lazy imports +# --------------------------------------------------------------------------- + + +def _import_slime_types(): + try: + from slime.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput + from slime.utils.types import Sample + + return Sample, RolloutFnTrainOutput, RolloutFnEvalOutput + except ImportError as err: + raise ImportError( + "slime is required for this module. Install with: pip install agentcore-rl-toolkit[slime]" + ) from err + + +# --------------------------------------------------------------------------- +# Gateway assembly (slime-specific: SGLang backend + --hf-checkpoint template) +# --------------------------------------------------------------------------- + + +def _start_gateway_server( + *, + host: str, + port: int, + sglang_url: str, + hf_checkpoint: str, + acr_timeout: float, + tool_call_parser: str | None = None, + reasoning_parser: str | None = None, +): + """Assemble the slime-flavored gateway and serve it on a background thread. + + Slime always drives SGLang, so the gateway samples via ``SglangHttpBackend`` + and renders with the HF chat template of ``--hf-checkpoint``. The serving + mechanics live in the shared :class:`ThreadedGatewayServer`. + + ``tool_call_parser`` / ``reasoning_parser`` (slime's --sglang-tool-call-parser / + --sglang-reasoning-parser) select SGLang's engine-grade parsers for the served + model's output format; the gateway samples via the native /generate endpoint, + so that parsing happens here rather than in the SGLang server. Each is + independent: an unset one leaves that derender stage on the gateway's + dependency-free default. + """ + try: + from transformers import AutoTokenizer + + from agentcore_rl_toolkit.rollout_gateway import HfTemplateRenderer, RolloutGateway, ThreadedGatewayServer + from agentcore_rl_toolkit.rollout_gateway.sampling_backends.sglang_http import SglangHttpBackend + except ImportError as err: + raise ImportError( + "The rollout gateway requires aiohttp + transformers. " + "Install with: pip install agentcore-rl-toolkit[gateway]" + ) from err + + # Each stage falls back to the gateway's dependency-free default when its + # slime arg is unset, so the two can be configured independently. + parser_fns = {} + if tool_call_parser: + from .sglang_parsing import build_tool_parser + + parser_fns["tool_parser"] = build_tool_parser(tool_call_parser) + if reasoning_parser: + from .sglang_parsing import build_reasoning_parser + + parser_fns["reasoning_parser"] = build_reasoning_parser(reasoning_parser) + + tokenizer = AutoTokenizer.from_pretrained(hf_checkpoint, trust_remote_code=True) + gateway = RolloutGateway( + backend=SglangHttpBackend(sglang_url, sock_read_timeout=acr_timeout), + renderer=HfTemplateRenderer(tokenizer, **parser_fns), + tokenizer=tokenizer, + ) + server = ThreadedGatewayServer(gateway, host=host, port=port) + server.start() + logger.info("Rollout gateway serving at %s (sglang worker: %s)", server.base_url, sglang_url) + return server + + +# --------------------------------------------------------------------------- +# Initialization +# --------------------------------------------------------------------------- + + +def _ensure_initialized(args: Namespace): + """Lazily initialize the in-process gateway and ACR client on first rollout call.""" + global _gateway_server, _client, _config + + from agentcore_rl_toolkit import RolloutClient + + if _config is None: + _config = SlimeArtConfig.from_args(args) + + cfg = _config + + if _gateway_server is None: + hf_checkpoint = getattr(args, "hf_checkpoint", None) + if not hf_checkpoint: + raise ValueError( + "--hf-checkpoint is required: the rollout gateway renders prompts " + "with the served checkpoint's HF chat template." + ) + _gateway_server = _start_gateway_server( + host=args.sglang_router_ip, # bind to routable IP so VPC agents can reach it + port=cfg.gateway_port, + sglang_url=f"http://{args.sglang_router_ip}:{args.sglang_router_port}", + hf_checkpoint=hf_checkpoint, + acr_timeout=cfg.acr_timeout, + tool_call_parser=cfg.sglang_tool_call_parser, + reasoning_parser=cfg.sglang_reasoning_parser, + ) + + if _client is None: + _client = RolloutClient( + agent_runtime_arn=cfg.agent_runtime_arn, + s3_bucket=cfg.s3_bucket, + exp_id=cfg.exp_id, + tps_limit=cfg.acr_tps_limit, + max_pool_connections=cfg.max_pool_connections, + ) + + return _gateway_server, _client, cfg + + +# --------------------------------------------------------------------------- +# Payload conversion +# --------------------------------------------------------------------------- + + +def _sample_to_payload(sample) -> dict: + """The agent payload is the JSONL row's ``metadata`` dict, verbatim. + + slime's Dataset reads the JSONL row's ``metadata`` field into + ``Sample.metadata``; we hand that dict to the agent unchanged. The JSONL's + top-level ``prompt`` field is for slime (tokenization, length filtering); + the agent's payload shape is entirely defined by whatever the data author + put in ``metadata``. A shallow copy isolates the agent's view from + downstream mutations to ``Sample.metadata`` (e.g. ``task_metadata`` + injection in ``_process_one_episode``). + """ + metadata = getattr(sample, "metadata", None) + if isinstance(metadata, dict): + return dict(metadata) + return {} + + +def _extract_reward(acr_result: dict) -> float: + """Extract scalar reward from an ACR S3 result dict.""" + rewards = acr_result.get("rewards", 0.0) + if isinstance(rewards, list): + return rewards[-1] if rewards else 0.0 + return float(rewards) + + +# --------------------------------------------------------------------------- +# TraceRecord -> slime Sample conversion +# --------------------------------------------------------------------------- + + +def _make_noop_sample(group_index: int = -1, session_id: str = "", status_name: str = "COMPLETED"): + """Create a minimum-valid Sample that contributes zero gradient. + + Used for DP padding and failed episodes. Has 2 tokens (1 prompt + 1 response), + loss_mask=[0] so Megatron processes it without error but produces no gradient. + + For failed episodes, pass session_id so that normalize_episode_rewards counts + this as a separate episode (reward=0) in the GRPO group. + """ + Sample, _, _ = _import_slime_types() + s = Sample() + s.tokens = [0, 0] + s.response_length = 1 + s.loss_mask = [0] + s.reward = 0.0 + s.rollout_log_probs = [0.0] + s.group_index = group_index + s.status = Sample.Status[status_name] + if session_id: + s.session_id = session_id + s.metadata = {"gateway_session_id": session_id, "task_index": group_index, "record_index": 0} + return s + + +def _record_to_sample( + record, + group_index: int, + sample_index: int, + session_id: str, + record_index: int, +): + """Convert one gateway TraceRecord into a slime Sample. + + A TraceRecord is already a merged, loss-masked training row: ``token_ids`` is + the full sequence, ``loss_mask`` / ``logprobs`` cover the response region only + (bridge tokens between turns carry loss_mask=0), so the conversion is direct. + """ + Sample = _import_slime_types()[0] + + tokens = list(record.token_ids) + loss_mask = list(record.loss_mask) + logprobs = list(record.logprobs) + response_length = int(record.response_length or len(loss_mask)) + + # Megatron requires prompt_length >= 1; a record whose first-turn prompt is + # empty cannot occur with a chat template, but guard anyway. + if len(tokens) - response_length < 1: + tokens = [0] + tokens + + # Defensive alignment: mask/logprobs must both span the response region. + if len(loss_mask) != response_length: + loss_mask = (loss_mask + [0] * response_length)[:response_length] + if len(logprobs) != response_length: + logprobs = (logprobs + [0.0] * response_length)[:response_length] + + sample = Sample() + sample.tokens = tokens + sample.response_length = response_length + sample.loss_mask = loss_mask + sample.rollout_log_probs = logprobs + sample.reward = float(record.reward) + sample.group_index = group_index + sample.index = sample_index + sample.session_id = session_id + sample.metadata = { + **(record.metadata or {}), + "task_index": group_index, + "gateway_session_id": session_id, + "record_index": record_index, + } + + truncated = record.status is Status.TRUNCATED or (record.metadata or {}).get("truncated") + sample.status = Sample.Status.TRUNCATED if truncated else Sample.Status.COMPLETED + return sample + + +def _log_records(session_id: str, records: list, reward: float, task_index: int): + """Append captured TraceRecords to a JSONL file for debugging.""" + try: + with open(_TRACE_LOG_PATH, "a") as f: + for i, r in enumerate(records): + record = { + "session_id": session_id, + "task_index": task_index, + "record": i, + "reward": reward, + "rollout_id": r.rollout_id, + "status": r.status.value, + "total_tokens": len(r.token_ids), + "response_length": r.response_length, + "trained_tokens": sum(r.loss_mask), + "response": r.response, + "metadata": r.metadata, + } + f.write(json.dumps(record) + "\n") + except Exception: + logger.warning("Failed to write trace log", exc_info=True) + + +# --------------------------------------------------------------------------- +# Episode processing +# --------------------------------------------------------------------------- + + +def _session_sampling_defaults(sampling_params: dict) -> dict: + """Per-session sampling defaults for the gateway (canonical keys, None-filtered).""" + return {k: v for k, v in sampling_params.items() if v is not None} + + +async def _process_one_episode( + sample, + server, + client, + cfg, + sampling_params: dict, + task_index: int, + sample_counter, +) -> list: + """Run one agent episode, return its slime Samples. + + All returned Samples share task_index (as group_index) so that + normalize_episode_rewards() can group all rows from all episodes + of the same task together for GRPO normalization. + """ + gateway = server.gateway + session_id = str(uuid.uuid4()) + try: + gateway.create_session(session_id, sampling_defaults=_session_sampling_defaults(sampling_params)) + + payload = _sample_to_payload(sample) + # Translate to OpenAI-compatible params (max_new_tokens→max_tokens, drop top_k) + agent_params = { + k if k != "max_new_tokens" else "max_tokens": v for k, v in sampling_params.items() if k != "top_k" + } + # Session identity in the api-key slot: the agent sets + # api_key=_rollout["api_key"] on its LLM client, so every call to the + # fixed gateway base_url carries "Authorization: Bearer ". + future = await client.invoke_async( + payload=payload, + session_id=session_id, + input_id=session_id, + base_url=server.base_url, + api_key=session_id, + model_id=cfg.model_id, + sampling_params=agent_params, + ) + + result = await future.result_async(timeout=cfg.acr_timeout) + episode_reward = _extract_reward(result) + records = await gateway.finish_session( + session_id, + base_sample=BaseTrace(rollout_id=session_id, group_index=task_index), + reward=episode_reward, + ) + _log_records(session_id, records, episode_reward, task_index) + + if not records: + noop = _make_noop_sample(group_index=task_index, session_id=session_id, status_name="FAILED") + noop.metadata["episode_error"] = "no trace records captured" + logger.info("Episode failed (session=%s): %s", session_id, noop.metadata["episode_error"]) + return [noop] + + samples = [ + _record_to_sample(rec, task_index, next(sample_counter), session_id, i) for i, rec in enumerate(records) + ] + for s in samples: + s.prompt = sample.prompt + s.label = sample.label + if sample.metadata: + s.metadata["task_metadata"] = sample.metadata + return samples + + except Exception as e: + # Record the failure on the sample (the per-batch summary counts these) + # and log it at INFO so a failing episode is visible. + noop = _make_noop_sample(group_index=task_index, session_id=session_id, status_name="FAILED") + noop.metadata["episode_error"] = str(e) or type(e).__name__ + logger.info("Episode failed (session=%s): %s", session_id, noop.metadata["episode_error"]) + return [noop] + finally: + # Idempotent: after a successful finish_session this is a no-op; on any + # failure path it drains stragglers and discards the partial trajectory. + await gateway.drop_session(session_id) + + +# --------------------------------------------------------------------------- +# Public entry point +# --------------------------------------------------------------------------- + + +def generate_rollout( + args: Namespace, + rollout_id: int, + data_source, + evaluation: bool = False, +): + """Custom slime rollout function: ACR agents + in-process rollout gateway. + + Implements slime's --rollout-function-path interface. + """ + _, RolloutFnTrainOutput, RolloutFnEvalOutput = _import_slime_types() + server, client, cfg = _ensure_initialized(args) + + # ---- Step 1: resolve the prompt groups to run, each paired with the + # sampling params it should use ---- + # Training samples one batch of GRPO-grouped prompts from the live + # data_source, all sharing the rollout params. Eval reads held-out datasets + # (args.eval_datasets, built by slime from --eval-prompt-data / + # --eval-config), each carrying its own already-resolved params (e.g. + # greedy via --eval-temperature 0); eval rewards are independent, so each + # (prompt, sample) is its own size-1 group. + if evaluation: + groups = [] # list of (prompt_group, sampling_params) + for dataset_cfg in getattr(args, "eval_datasets", None) or []: + params = { + "temperature": dataset_cfg.temperature, + "top_p": dataset_cfg.top_p, + "top_k": dataset_cfg.top_k, + "max_new_tokens": dataset_cfg.max_response_len, + } + dataset = _get_eval_dataset(args, dataset_cfg) + for prompt in dataset.samples: + for _ in range(dataset_cfg.n_samples_per_eval_prompt or 1): + groups.append(([copy.deepcopy(prompt)], params)) + else: + params = { + "temperature": args.rollout_temperature, + "top_p": args.rollout_top_p, + "top_k": args.rollout_top_k, + "max_new_tokens": args.rollout_max_response_len, + } + groups = [(group, params) for group in data_source.get_samples(args.rollout_batch_size)] + + # ---- Step 2 (shared): run every group as parallel ACR episodes ---- + # Each sample in a group becomes one episode tagged with the group index + # (GRPO grouping in training, a unique id in eval); rows are flattened + # back per group. All groups (and all episodes within them) are scheduled + # concurrently, but a shared semaphore caps the number of episodes that are + # actually in flight at once (cfg.max_concurrent). Without this cap, a large + # batch (e.g. a full 1319-prompt eval set) launches every episode at once — + # the 25-TPS client limiter only paces session *starts*, not the live count — + # which saturates the gateway/router + S3 result polling (episodes then miss + # acr_timeout and fail) and over-pressures the colocated SGLang KV cache + # (token-pool exhaustion crash). asyncio.gather preserves argument order, so + # the returned list stays group-ordered (list[list[Sample]]), keeping the + # GRPO group_index tags and slime's nesting-depth contract intact. (Ordering + # is non-semantic anyway: grouping is by explicit group_index/session_id, not + # list position — see rewards.normalize_episode_rewards.) + sample_counter = iter(range(10**9)) + + async def _run(): + # Bound concurrent in-flight episodes across ALL groups. Created inside + # the running loop (asyncio.Semaphore binds to the active event loop). + sem = asyncio.Semaphore(max(1, cfg.max_concurrent)) + + async def _episode(s, group_index, sampling_params): + async with sem: + return await _process_one_episode(s, server, client, cfg, sampling_params, group_index, sample_counter) + + async def _run_group(group_index, group, sampling_params): + results = await asyncio.gather(*(_episode(s, group_index, sampling_params) for s in group)) + return [s for r in results for s in r] + + return list( + await asyncio.gather( + *( + _run_group(group_index, group, sampling_params) + for group_index, (group, sampling_params) in enumerate(groups) + ) + ) + ) + + try: + loop = asyncio.get_event_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + results = loop.run_until_complete(_run()) + + # ---- Step 3: per-batch summary ---- + num_episodes = sum(len(group) for group, _ in groups) + num_sequences = sum(len(g) for g in results) + failed = sum(1 for g in results for s in g if (s.metadata or {}).get("episode_error")) + succeeded = num_episodes - failed + phase = "Eval" if evaluation else "Rollout" + logger.info( + "%s %d batch: episodes=%d (succeeded=%d failed=%d) sequences=%d", + phase, + rollout_id, + num_episodes, + succeeded, + failed, + num_sequences, + ) + + # ---- Step 4: shape the backend-specific output ---- + if evaluation: + # Episode reward is broadcast to every row-Sample; take the first. + rewards = [float(g[0].reward) if g and isinstance(g[0].reward, (int, float)) else 0.0 for g in results] + n = max(len(rewards), 1) + accuracy = sum(1 for r in rewards if r > 0) / n + avg_reward = sum(rewards) / n + return RolloutFnEvalOutput( + data={"eval": {"rewards": rewards}}, + metrics={ + "eval/accuracy": accuracy, + "eval/avg_reward": avg_reward, + "eval/n_samples": len(rewards), + }, + ) + + # Training: pad to a dp_size multiple so no real samples are trimmed. + dp_size = args.actor_num_nodes * args.actor_num_gpus_per_node // args.tensor_model_parallel_size + remainder = sum(len(g) for g in results) % dp_size + if remainder > 0: + results[-1].extend([_make_noop_sample(group_index=-1) for _ in range(dp_size - remainder)]) + return RolloutFnTrainOutput(samples=results) + + +def _get_eval_dataset(args, dataset_cfg): + """Load + cache a held-out eval dataset described by an EvalDatasetConfig. + + Reads the JSONL itself (independent of the training data_source) using + slime's Dataset so the prompt/metadata parsing matches the training path. + """ + from slime.utils.data import Dataset + from slime.utils.processing_utils import load_processor, load_tokenizer + + key = dataset_cfg.cache_key + (args.hf_checkpoint, args.apply_chat_template) + if key not in _eval_datasets: + tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + processor = load_processor(args.hf_checkpoint, trust_remote_code=True) + _eval_datasets[key] = Dataset( + path=dataset_cfg.path, + tokenizer=tokenizer, + processor=processor, + max_length=args.eval_max_prompt_len, + prompt_key=dataset_cfg.input_key, + label_key=dataset_cfg.label_key, + metadata_key=dataset_cfg.metadata_key, + multimodal_keys=args.multimodal_keys, + tool_key=dataset_cfg.tool_key, + apply_chat_template=args.apply_chat_template, + apply_chat_template_kwargs=args.apply_chat_template_kwargs, + ) + return _eval_datasets[key] diff --git a/src/agentcore_rl_toolkit/backends/experimental/slime/integration/sglang_parsing.py b/src/agentcore_rl_toolkit/backends/experimental/slime/integration/sglang_parsing.py new file mode 100644 index 0000000..7ac721c --- /dev/null +++ b/src/agentcore_rl_toolkit/backends/experimental/slime/integration/sglang_parsing.py @@ -0,0 +1,105 @@ +"""SGLang-backed parsers for the rollout gateway's derender seams. + +The gateway samples via SGLang's native ``/generate`` (token-in/token-out), which +bypasses the server's OpenAI layer — so the reasoning/tool-call parsing that layer +would normally do must happen gateway-side. This module builds the two stage +callables :class:`HfTemplateRenderer` accepts, from SGLang's own detectors, so the +gateway derenders exactly what the served engine would have: + +- :func:`build_tool_parser` wraps ``FunctionCallParser`` (slime's + ``--sglang-tool-call-parser``, e.g. ``qwen`` for Qwen2.5/Qwen3 JSON ````); +- :func:`build_reasoning_parser` wraps ``ReasoningParser`` (slime's + ``--sglang-reasoning-parser``, e.g. ``qwen3``). + +Both are independent: pass only the one whose format the built-in default gets wrong +and the renderer keeps its dependency-free default for the other stage. + +Lives in the slime backend (not ``rollout_gateway``) because slime trainers always +have sglang importable — the gateway package itself stays engine-free. +""" + +import json +import logging + +logger = logging.getLogger(__name__) + + +def _import_sglang(what: str, flag: str): + """Import an sglang parser symbol, or raise with actionable context.""" + try: + if what == "tool": + from sglang.srt.entrypoints.openai.protocol import Tool + from sglang.srt.function_call.function_call_parser import FunctionCallParser + + return FunctionCallParser, Tool + from sglang.srt.parser.reasoning_parser import ReasoningParser + + return ReasoningParser, None + except ImportError as err: + raise ImportError( + f"{flag} needs sglang's parsers, which the slime trainer environment normally " + f"provides; unset {flag} to fall back to the gateway's built-in " + "dependency-free parsing." + ) from err + + +def build_tool_parser(parser_name: str): + """Build a renderer ``tool_parser``: ``(body_text, tools_schema) -> (text, tool_uses, ill_formed)``. + + ``parser_name`` is an SGLang ``--tool-call-parser`` name and must match the served + model. Imported and validated eagerly so a bad name or environment fails at gateway + construction, not on the first rollout turn. + """ + FunctionCallParser, Tool = _import_sglang("tool", "--sglang-tool-call-parser") + if parser_name not in FunctionCallParser.ToolCallParserEnum: + raise ValueError( + f"unknown tool_call_parser {parser_name!r}; choose one of {sorted(FunctionCallParser.ToolCallParserEnum)}" + ) + + def parse_tool_uses(body_text: str, tools_schema: list[dict]) -> tuple[str, list[dict], bool]: + # detectors carry streaming state -> fresh parser per turn + parser = FunctionCallParser([Tool.model_validate(t) for t in tools_schema], parser_name) + normal_text, calls = parser.parse_non_stream(body_text) + + tool_uses: list[dict] = [] + ill_formed = False + for call in calls: + if not call.name: + ill_formed = True + continue + try: + # ToolCallItem.parameters is a JSON string; adapters want a dict + args = json.loads(call.parameters) if call.parameters else {} + except json.JSONDecodeError: + ill_formed = True + continue + tool_uses.append({"name": call.name, "input": args if isinstance(args, dict) else {}}) + # the model opened a tool call but the parser extracted nothing usable + if not tool_uses and parser.has_tool_call(body_text): + ill_formed = True + return normal_text or "", tool_uses, ill_formed + + return parse_tool_uses + + +def build_reasoning_parser(parser_name: str): + """Build a renderer ``reasoning_parser``: ``raw_output -> (reasoning, body_text)``. + + ``parser_name`` is an SGLang ``--reasoning-parser`` name (e.g. ``qwen3``). Use this + for models whose reasoning is not a plain ```` block; the gateway's default + handles that common case with no dependency. + """ + ReasoningParser, _ = _import_sglang("reasoning", "--sglang-reasoning-parser") + if parser_name.lower() not in ReasoningParser.DetectorMap: + raise ValueError( + f"unknown reasoning_parser {parser_name!r}; choose one of {sorted(ReasoningParser.DetectorMap)}" + ) + + def split_reasoning(raw_output: str) -> tuple[str, str]: + reasoning, body = ReasoningParser(model_type=parser_name).parse_non_stream(raw_output) + return reasoning or "", body or "" + + return split_reasoning + + +__all__ = ["build_reasoning_parser", "build_tool_parser"] diff --git a/src/agentcore_rl_toolkit/backends/experimental/slime/scripts/install_slime.sh b/src/agentcore_rl_toolkit/backends/experimental/slime/scripts/install_slime.sh new file mode 100644 index 0000000..fcae44b --- /dev/null +++ b/src/agentcore_rl_toolkit/backends/experimental/slime/scripts/install_slime.sh @@ -0,0 +1,101 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Install the slime training backend (CUDA 13 only — cu12 is not supported). +# +# Usage: +# bash src/agentcore_rl_toolkit/backends/experimental/slime/scripts/install_slime.sh + +TORCH_BACKEND=cu130 # uv --torch-backend for the PyTorch ecosystem +: "${CUDA_HOME:=/usr/local/cuda-13.0}" +export CUDA_HOME + +# The source builds below (flash-attn, transformer-engine, apex) write tens of GB +# of nvcc scratch files to $TMPDIR and can fill up a small root volume +# ("No space left on device" from cc1plus/nvcc). Prefer the instance's large +# ephemeral NVMe volume (present on AWS DLAMI/HyperPod nodes) unless the caller +# already set TMPDIR. +if [ -z "${TMPDIR:-}" ] && [ -d /opt/dlami/nvme ]; then + TMPDIR="/opt/dlami/nvme/${USER}/tmp" + mkdir -p "$TMPDIR" +fi +export TMPDIR="${TMPDIR:-/tmp}" + +echo "=== slime installer (cu13): TORCH_BACKEND=$TORCH_BACKEND CUDA_HOME=$CUDA_HOME TMPDIR=$TMPDIR ===" + +# Assumes your python environment is already activated. + +uv pip install torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 --torch-backend="$TORCH_BACKEND" +uv pip install cmake ninja pybind11 "packaging>=24.2" wheel + +MAX_JOBS=64 uv pip install "flash-attn==2.8.3" \ + --no-binary flash-attn --no-build-isolation --no-cache-dir --torch-backend="$TORCH_BACKEND" + +uv pip install "git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c" --no-deps +uv pip install "flash-linear-attention" --torch-backend="$TORCH_BACKEND" + +uv pip install tilelang + +# Explicitly exclude transformer-engine-cu12 to avoid "Multiple libcudart +# libraries found" errors. +echo "transformer-engine-cu12 ; sys_platform == 'never'" | \ +MAX_JOBS=128 uv pip install --no-cache --no-build-isolation \ + --overrides - \ + "transformer_engine[pytorch,core-cu13]==2.11" + +NVCC_APPEND_FLAGS="--threads 4" \ + APEX_CPP_EXT=1 APEX_CUDA_EXT=1 APEX_PARALLEL_BUILD=8 \ + uv pip install -v --no-build-isolation --no-cache-dir \ + "git+https://github.com/NVIDIA/apex.git@10417aceddd7d5d05d7cbf7b0fc2daad1105f8b4" + +# torch_memory_saver's TMS_CUDA_MAJOR sets the compiled .so suffix (_cu13) and +# must match what its runtime detector reads from torch.version.cuda, so we +# derive it from torch rather than hardcoding. +export TMS_CUDA_MAJOR="$(python -c 'import torch; print(torch.version.cuda.split(".")[0])')" +uv pip install -v "git+https://github.com/fzyzcjy/torch_memory_saver.git@a193d9dd1b877d33c64a41cfb3db9f867df2d926" \ + --no-cache-dir --force-reinstall --no-build-isolation + +uv pip install "git+https://github.com/radixark/Megatron-Bridge.git@6fde1c8538ea4ad966c7fba5f759be54f943b598" --no-deps --no-build-isolation +uv pip install "nvidia-modelopt[torch]>=0.37.0" --no-build-isolation + +# sglang's default kernel + deep-gemm builds already target cu13, so no +# wheel-index reinstall is needed here. +uv pip install --prerelease=allow "sglang==0.5.13" --torch-backend="$TORCH_BACKEND" + +# We have to git clone and install from local because wheel file does not expose megatron.training that is required by slime +git clone https://github.com/NVIDIA/Megatron-LM.git +cd Megatron-LM +git checkout "1dcf0dafa884ad52ffb243625717a3471643e087" +uv pip install -e . --no-build-isolation --config-settings editable_mode=compat +cd .. + +uv pip install --reinstall-package nvidia-cutlass-dsl-libs-base --no-deps \ + "nvidia-cutlass-dsl-libs-base==4.5.2" + +uv pip install --reinstall-package pyjwt PyJWT + +# Install slime +git clone https://github.com/THUDM/slime.git +cd slime +git checkout "fa3c990af6f18efd3fd9922698bf4bf4048d1263" +uv pip install -r "requirements.txt" +uv pip install -e . --no-deps +cd .. + +uv pip install "https://github.com/zhuzilin/sgl-router/releases/download/v0.3.2-1117d05/sglang_router-0.3.2-cp38-abi3-manylinux_2_28_x86_64.whl" --force-reinstall + +# numpy<2 for Megatron; scipy<1.14 because scipy>=1.14 requires numpy>=2. +# --no-config so this isn't silently overridden. +uv pip install --no-config "numpy<2" "scipy<1.14" + +# Apply slime's official patches to megatron + sglang. +SLIME_PATCH_DIR="$(cd slime/docker/patch/latest && pwd)" +SITE_PACKAGES="$(python -c 'import sysconfig; print(sysconfig.get_path("purelib"))')" + +# Megatron patch: apply against the repo without leaving cwd (git -C). +git -C Megatron-LM update-index --refresh >/dev/null 2>&1 || true +git -C Megatron-LM apply --3way "$SLIME_PATCH_DIR/megatron.patch" + +# sglang patches: apply into site-packages without leaving cwd (patch -d). +patch -d "$SITE_PACKAGES" -p2 -F0 -N < "$SLIME_PATCH_DIR/sglang.patch" +patch -d "$SITE_PACKAGES" -p2 -F0 -N < "$SLIME_PATCH_DIR/sglang-top_p.patch"