diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 761496b949..185017282e 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -744,6 +744,18 @@ jobs: { "num_gpus": 0, "test_file": "test_agent/test_agent_rollout_cpu.py" + }, + { + "num_gpus": 0, + "test_file": "test_rollout_buffer/test_ags_prompt_source.py" + }, + { + "num_gpus": 0, + "test_file": "test_rollout_buffer/test_ags_empty_patch_guard.py" + }, + { + "num_gpus": 0, + "test_file": "test_rollout_buffer/test_ags_generator.py" } ] defaults: @@ -774,6 +786,8 @@ jobs: pip install torch --index-url https://download.pytorch.org/whl/cpu pip install pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors + pip install openai transformers wandb fastapi uvicorn + - name: Install shell: bash diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 415586cb84..d1578f90bb 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -97,11 +97,15 @@ 'label': 'run-ci-agent', 'always': True, 'cpu': True, + 'extra_pip_deps': 'openai transformers wandb fastapi uvicorn', 'tests': [ {'test_file': 'test_agent/test_trajectory_manager_branching.py', 'num_gpus': 0}, {'test_file': 'test_agent/test_adapters.py', 'num_gpus': 0}, {'test_file': 'test_agent/test_harness.py', 'num_gpus': 0}, {'test_file': 'test_agent/test_agent_rollout_cpu.py', 'num_gpus': 0}, + {'test_file': 'test_rollout_buffer/test_ags_prompt_source.py', 'num_gpus': 0}, + {'test_file': 'test_rollout_buffer/test_ags_empty_patch_guard.py', 'num_gpus': 0}, + {'test_file': 'test_rollout_buffer/test_ags_generator.py', 'num_gpus': 0}, ], }, diff --git a/examples/coding_agent_rl_ags/eval_cbc_qwen35_35b_a3b_swe_2nodes.sh b/examples/coding_agent_rl_ags/eval_cbc_qwen35_35b_a3b_swe_2nodes.sh new file mode 100644 index 0000000000..27d16eaa3c --- /dev/null +++ b/examples/coding_agent_rl_ags/eval_cbc_qwen35_35b_a3b_swe_2nodes.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash +# Score one checkpoint on SWE-bench Verified with the CodeBuddy Code (cbc) +# harness, 2 nodes, no training. +# +# Thin wrapper: the launcher logic lives in eval_cc_qwen35_35b_a3b_swe_2nodes.sh +# and is harness-agnostic -- SWE_AGENT selects the harness, and that script sets +# both harnesses' EXTRA_ARGS/EXTRA_ENVS knobs -- so there is one copy to keep correct. +# +# Every variable the CC eval script reads works here too: +# LOAD_DIR= CKPT_STEP=99 EVAL_DATA= \ +# bash examples/coding_agent_rl_ags/eval_cbc_qwen35_35b_a3b_swe_2nodes.sh +# +# Omit LOAD_DIR/CKPT_STEP to score the untrained HF model as the baseline. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" + +export SWE_AGENT=codebuddy_code +export EXP_TAG="${EXP_TAG:-coding_agent_rl_ags_eval_cbc_qwen35_35b_a3b${CKPT_STEP:+_step${CKPT_STEP}}}" + +exec bash "${SCRIPT_DIR}/eval_cc_qwen35_35b_a3b_swe_2nodes.sh" "$@" diff --git a/examples/coding_agent_rl_ags/eval_cc_qwen35_35b_a3b_swe_2nodes.sh b/examples/coding_agent_rl_ags/eval_cc_qwen35_35b_a3b_swe_2nodes.sh new file mode 100755 index 0000000000..7042123a1d --- /dev/null +++ b/examples/coding_agent_rl_ags/eval_cc_qwen35_35b_a3b_swe_2nodes.sh @@ -0,0 +1,598 @@ +#!/usr/bin/env bash +# Eval-only: score ONE checkpoint on the full SWE-bench Verified set with Claude +# Code on AGS. No training, no rollout buffer, no optimizer step. +# +# Derived from run_qwen35_35b_a3b_swe_2nodes.sh. The differences that matter: +# +# --num-rollout 0 train.py's eval-only branch runs eval once at rollout_id=0 +# and then exits, skipping the whole train loop. +# no buffer the eval path (ags_generator.generate -> AGSRolloutRunner) +# never talks to the rollout buffer, so buffer.py is not +# started and --rollout-function-path / --rollout-buffer-url +# are omitted. The AGS adapter is started by the eval path +# itself (see ADAPTER_PORT below). +# scheduler args train_iters = num_rollout * ... = 0 makes Megatron's +# OptimizerParamScheduler assert; see LR_SCHED_ARGS. +# +# Run from a long-lived shell / tmux session on the Ray head node. Budget ~95min +# for 500 samples at SWE_ROLLOUT_CONCURRENCY=32. +# +# Required: +# EXP=/data_train//experiments/ +# LOAD_DIR= +# E2B_API_KEY= +# Strongly recommended: +# CKPT_STEP= pin the exact step; without it Megatron reads +# latest_checkpointed_iteration.txt and you score +# whatever happens to be newest. +# TRAIN_NUM_ROLLOUT= the --num-rollout of the run that WROTE the +# checkpoint (see LR_SCHED_ARGS). +# +# Example: +# EXP=/data_train/ericxjzheng/experiments/eval_step79 \ +# LOAD_DIR=/data_train/ericxjzheng/experiments//checkpoints \ +# CKPT_STEP=79 TRAIN_NUM_ROLLOUT=100 \ +# E2B_API_KEY=... WANDB_API_KEY=... WANDB_ENTITY=... \ +# bash examples/coding_agent_rl_ags/eval_cc_qwen35_35b_a3b_swe_2nodes.sh + +# Best-effort cleanup so a rerun does not collide with stale workers/services. +pkill -9 sglang || true +pkill -f "slime_plugins.rollout_buffer.buffer" || true +pkill -f "slime_plugins/rollout_buffer/buffer.py" || true +sleep 3 +ray stop --force || true +pkill -9 ray || true +sleep 3 +pkill -9 ray || true + +set -ex + +export PYTHONUNBUFFERED=1 + +EXP="${EXP:?set EXP to an experiment directory, e.g. /data_train/ericxjzheng/experiments/}" +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +SLIME_DIR="${SLIME_DIR:-/data_train/ericxjzheng/workspace/slime}" + +# ============ cluster size ============ +ACTOR_NUM_NODES="${ACTOR_NUM_NODES:-${MLP_WORKER_NUM:-2}}" +ACTOR_NUM_GPUS_PER_NODE="${ACTOR_NUM_GPUS_PER_NODE:-8}" +TOTAL_NUM_GPUS=$((ACTOR_NUM_NODES * ACTOR_NUM_GPUS_PER_NODE)) + +# ============ model parallelism ============ +export TP_SIZE="${TP_SIZE:-2}" +export PP_SIZE="${PP_SIZE:-1}" +export CP_SIZE="${CP_SIZE:-8}" +export EP_SIZE="${EP_SIZE:-8}" +export ETP_SIZE="${ETP_SIZE:-1}" + +# ============ rollout engine ============ +ROLLOUT_NUM_GPUS="${ROLLOUT_NUM_GPUS:-${TOTAL_NUM_GPUS}}" +ROLLOUT_TP_SIZE="${ROLLOUT_TP_SIZE:-8}" +ROLLOUT_DP_SIZE="${ROLLOUT_DP_SIZE:-2}" +ROLLOUT_EP_SIZE="${ROLLOUT_EP_SIZE:-8}" +ROLLOUT_MEM_UTILIZATION="${ROLLOUT_MEM_UTILIZATION:-0.75}" + +# num_rollout 0 selects train.py's eval-only branch. The rollout_batch_size / +# n_samples_per_prompt below are unused by eval (it iterates the eval dataset) +# but must stay positive: model.py divides by global_batch_size. +NUM_ROLLOUT=0 +ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-8}" +N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-8}" +GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-$((ROLLOUT_BATCH_SIZE * N_SAMPLES_PER_PROMPT))}" +MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}" + +# ============ context length ============ +MAX_CONTEXT_LEN="${MAX_CONTEXT_LEN:-96000}" +MAX_GEN_LEN="${MAX_GEN_LEN:-32768}" +ROLLOUT_MAX_PROMPT_LEN="${ROLLOUT_MAX_PROMPT_LEN:-${MAX_CONTEXT_LEN}}" + +# ============ eval ============ +# eval_interval must be non-None for the eval-only branch to fire; with +# num_rollout=0 the value itself is never used as a cadence. +EVAL_INTERVAL="${EVAL_INTERVAL:-1}" +EVAL_DATA="${EVAL_DATA:-/data_train/ericxjzheng/data/SWE-bench_Verified_slime_rl_format_from_harbor/swebench_verified_slime_instruction_tcr.jsonl}" +# Names the dataset in slime's metrics, so the score lands on eval/. Only +# worth overriding when EVAL_DATA is not SWE-bench Verified -- a difficulty scan +# over the *training* pool, say, which would otherwise log a "swebench_verified" +# number that is nothing of the sort. +EVAL_DATASET_NAME="${EVAL_DATASET_NAME:-swebench_verified}" +# 1 attempt x 500 rows = the full set, scored pass@1. +N_SAMPLES_PER_EVAL_PROMPT="${N_SAMPLES_PER_EVAL_PROMPT:-1}" + +# Benchmark sampling, matching the periodic eval in the training launcher so the +# two are comparable. Overridable because measuring how *trainable* a prompt is +# needs the training sampling instead (temperature 1.0, top_p 1.0, top_k -1 -- +# slime's --rollout-top-p/--rollout-top-k defaults, which the training launcher +# leaves untouched): a prompt's solve rate is a property of the sampling +# distribution, so a scan run at 0.7/0.8/20 would not describe the rollouts GRPO +# actually sees. +EVAL_TEMPERATURE="${EVAL_TEMPERATURE:-0.7}" +EVAL_TOP_P="${EVAL_TOP_P:-0.8}" +EVAL_TOP_K="${EVAL_TOP_K:-20}" + +# ============ paths — override before launching ============ +HF_CHECKPOINT="${HF_CHECKPOINT:-/data_train/ericxjzheng/models/Qwen3.5-35B-A3B}" +REF_MODEL_PATH="${REF_MODEL_PATH:-/data_train/ericxjzheng/models/Qwen3.5-35B-A3B_torch_dist}" + +# The checkpoint under test. Without LOAD_DIR this scores the base model, which +# is a valid baseline but almost certainly not what you meant. +LOAD_DIR="${LOAD_DIR:-}" +CKPT_STEP="${CKPT_STEP:-}" + +EXP_TAG="${EXP_TAG:-coding_agent_rl_ags_eval_cc_qwen35_35b_a3b${CKPT_STEP:+_step${CKPT_STEP}}}" +STAMP="$(date +%Y%m%d_%H%M%S)" +RUN_ROOT="${RUN_ROOT:-${EXP}/runs/${EXP_TAG}_${STAMP}}" + +# ============ logging/artifacts ============ +LOG_DIR="${RUN_ROOT}" +mkdir -p "${LOG_DIR}/ags_artifacts" +LOG_FILE="${LOG_DIR}/run.log" +export TRAJECTORY_DUMP_DIR="${TRAJECTORY_DUMP_DIR:-${LOG_DIR}/ags_artifacts}" +export EXPERIMENT_NAME="${EXPERIMENT_NAME:-${EXP_TAG}}" +echo "======================================================================" +echo "Eval log: ${LOG_FILE}" +echo "RUN_ROOT= ${RUN_ROOT}" +echo "Checkpoint: ${LOAD_DIR:-}" +echo "Step: ${CKPT_STEP:-}" +echo "Eval data: ${EVAL_DATA}" +echo "Dataset name: ${EVAL_DATASET_NAME}" +echo "Attempts/row: ${N_SAMPLES_PER_EVAL_PROMPT}" +echo "Sampling: temperature=${EVAL_TEMPERATURE} top_p=${EVAL_TOP_P} top_k=${EVAL_TOP_K}" +echo "Prompt style: ${SWE_EVAL_PROMPT_STYLE:-dataset}" +echo "======================================================================" + +if [[ -n "${LOAD_DIR}" && ! -f "${LOAD_DIR}/latest_checkpointed_iteration.txt" ]]; then + echo "ERROR: LOAD_DIR=${LOAD_DIR} has no latest_checkpointed_iteration.txt." + echo " slime treats such a path as 'no Megatron checkpoint' and silently" + echo " falls back to --ref-load, so you would score the base model." + exit 1 +fi +if [[ -n "${LOAD_DIR}" && -z "${CKPT_STEP}" ]]; then + echo "WARNING: CKPT_STEP unset; Megatron will load whichever step" + echo " latest_checkpointed_iteration.txt names:" + cat "${LOAD_DIR}/latest_checkpointed_iteration.txt" || true +fi + +# ============ ray cluster network ============ +# Set MASTER_ADDR before AGS/SWE blocks: ADAPTER_PUBLIC_HOST below falls back to it. +export MASTER_ADDR="${MASTER_ADDR:-${MLP_WORKER_0_HOST:-$(hostname -I | awk '{print $1}')}}" +export MASTER_PORT="${MASTER_PORT:-${MLP_WORKER_0_PORT:-6379}}" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" + +# ============ SWE / Claude Code / AGS rollout knobs ============ +export SWE_AGENT="${SWE_AGENT:-claude_code}" + +# AGS uses the E2B-compatible SDK surface. Export E2B_API_KEY in the launch +# environment (for Tencent AGS this is typically the AGS gateway key). +export E2B_DOMAIN="${E2B_DOMAIN:-ap-shanghai.tencentags.com}" +export AGS_BASE_TOOL="${AGS_BASE_TOOL:-sdt-3fzh6mv6}" +export AGS_IMAGE_REGISTRY_TYPE="${AGS_IMAGE_REGISTRY_TYPE:-enterprise}" +export AGS_SANDBOX_RESOURCES_JSON=${AGS_SANDBOX_RESOURCES_JSON:-'{"cpu":"4","memory":"16Gi"}'} + +# No training rollout runs here, so nothing binds ADAPTER_PORT ahead of eval. +# get_adapter_service health-probes this port and, finding nothing, starts a +# local eval adapter on an ephemeral port inside the RolloutManager actor -- +# advertising that actor's own node IP, since Ray does not pin it to the head. +# ADAPTER_PUBLIC_HOST is therefore only a probe target here, not the address +# handed to sandboxes. +export ADAPTER_PUBLIC_HOST="${ADAPTER_PUBLIC_HOST:-${MASTER_ADDR:-${MLP_WORKER_0_HOST:-127.0.0.1}}}" +export ADAPTER_BIND_HOST="${ADAPTER_BIND_HOST:-0.0.0.0}" +export ADAPTER_PORT="${ADAPTER_PORT:-18001}" + +export SWE_AGENT_TIME_BUDGET_SEC="${SWE_AGENT_TIME_BUDGET_SEC:-1800}" +export SWE_EVAL_TIMEOUT_SEC="${SWE_EVAL_TIMEOUT_SEC:-600}" +# false: grade in the agent sandbox; true: boot a second clean sandbox for grading. +export SWE_EVAL_ISOLATED_SANDBOX="${SWE_EVAL_ISOLATED_SANDBOX:-false}" +export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-32}" +export SWE_BOOT_RETRIES="${SWE_BOOT_RETRIES:-10}" +export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-32}" + +# The default EVAL_DATA is the instruction.md-sourced set, whose `prompt` already +# carries the task text, so eval uses the dataset style: hand that prompt straight +# to the agent and write no PROBLEM_STATEMENT.md. Override to "instruction" when +# pointing EVAL_DATA at the training pool -- SWE-rebench rows carry a raw issue +# body rather than a rendered instruction, and the training rollouts they must be +# compared against run under SWE_PROMPT_STYLE=instruction. +# SWE_PROMPT_STYLE is set too because AGSGeneratorConfig.from_env validates both, +# though only the eval one is read here. +export SWE_PROMPT_STYLE="${SWE_PROMPT_STYLE:-instruction}" +export SWE_EVAL_PROMPT_STYLE="${SWE_EVAL_PROMPT_STYLE:-dataset}" + +# ---- auto-compaction ------------------------------------------------------- +# Compact before a segment crosses the training-side context cap, otherwise the +# adapter returns finish_reason="length" with zero output tokens once the prompt +# reaches rollout_max_context_len (slime/agent/adapters/common.py) and the turn is +# wasted. +# +# Expressed as a PERCENTAGE of the model's context window rather than an absolute +# token count: an absolute auto-compact window is clamped to [100k, 1M] by both +# CLIs, which is above our 96k cap, so it could never fire in time. Both CLIs are +# also told the real window, so the percentage is of MAX_CONTEXT_LEN on each side. +# +# CodeBuddy CODEBUDDY_AUTOCOMPACT_PCT_OVERRIDE (cbc 1.9.33) is a percentage of +# the model's maxInputTokens, which the harness writes into +# models.json from SLIME_AGENT_MAX_INPUT_TOKENS. Without that key +# resolveCompactTriggerAt() falls back to the clamped absolute +# window. Source: agent-cli src/node/context/context-protocol.ts. +# Claude Code CLAUDE_AUTOCOMPACT_PCT_OVERRIDE, as a percentage of the context +# window, which CLAUDE_CODE_MAX_CONTEXT_TOKENS sets. That variable +# applies DIRECTLY for model names claude does not recognise as a +# Claude model -- ours is "slime-actor", so it does (the binary +# gates on `!normalize(model).startsWith("claude-")`; for a real +# claude-* name it would need DISABLE_COMPACT too, which would +# disable the compaction we want). Without it the CLI assumes its +# 200000 default and would compact at pct% of 200k, i.e. never +# before our 96k cap. +# 60%, not the CLI's 70% default, because the per-turn output reservation eats +# into what is reachable: claude assumes MAX_OUTPUT_TOKENS (32000 by default for +# model ids it does not recognise, which includes ours) is available on top of the +# prompt, so the prompt cannot grow past MAX_CONTEXT_LEN - MAX_GEN_LEN ~= 63k +# before the adapter's hard stop. A 70% trigger (67200) sits ABOVE that and would +# never be reached; 60% (57600) fires with room to spare. Raising +# AGENT_AUTOCOMPACT_PCT re-opens that gap -- the check below says so out loud. +AGENT_AUTOCOMPACT_PCT="${AGENT_AUTOCOMPACT_PCT:-60}" +AGENT_MAX_CONTEXT_TOKENS="${AGENT_MAX_CONTEXT_TOKENS:-${MAX_CONTEXT_LEN}}" + +# cbc: models.json maxInputTokens (written by CodeBuddyCodeHarness). +export SLIME_AGENT_MAX_INPUT_TOKENS="${SLIME_AGENT_MAX_INPUT_TOKENS:-${AGENT_MAX_CONTEXT_TOKENS}}" +# Built in a separate variable, not inline in ${VAR:-...}: a JSON default inside +# that expansion is mis-parsed -- the value's own "}" closes the expansion early +# and the trailing brace leaks into the result ("{...}}"). +CBC_AUTOCOMPACT_ENVS="{\"CODEBUDDY_AUTOCOMPACT_PCT_OVERRIDE\":\"${AGENT_AUTOCOMPACT_PCT}\"}" +export SLIME_AGENT_CBC_EXTRA_ENVS="${SLIME_AGENT_CBC_EXTRA_ENVS:-${CBC_AUTOCOMPACT_ENVS}}" + +# claude: declare the window and the per-turn output budget, then set the trigger +# percentage. MAX_OUTPUT_TOKENS is pinned to MAX_GEN_LEN rather than left at the +# CLI's 32000 default-for-unknown-model-ids so the reservation matches what the +# adapter will actually serve (--rollout-max-response-len). +CC_AUTOCOMPACT_ENVS="{\"CLAUDE_CODE_MAX_CONTEXT_TOKENS\":\"${AGENT_MAX_CONTEXT_TOKENS}\",\"CLAUDE_CODE_MAX_OUTPUT_TOKENS\":\"${MAX_GEN_LEN}\",\"CLAUDE_AUTOCOMPACT_PCT_OVERRIDE\":\"${AGENT_AUTOCOMPACT_PCT}\"}" +export SLIME_AGENT_CC_EXTRA_ENVS="${SLIME_AGENT_CC_EXTRA_ENVS:-${CC_AUTOCOMPACT_ENVS}}" + +AGENT_AUTOCOMPACT_AT=$((AGENT_AUTOCOMPACT_PCT * AGENT_MAX_CONTEXT_TOKENS / 100)) +# The largest prompt that can still be served: the adapter caps prompt+output at +# rollout_max_context_len, so a turn needs MAX_GEN_LEN of headroom. +AGENT_PROMPT_CEILING=$((MAX_CONTEXT_LEN - MAX_GEN_LEN)) +echo "Auto-compact: ${AGENT_AUTOCOMPACT_PCT}% of ${AGENT_MAX_CONTEXT_TOKENS} = ${AGENT_AUTOCOMPACT_AT} tokens" \ + "(prompt ceiling ${AGENT_PROMPT_CEILING} = ${MAX_CONTEXT_LEN} - ${MAX_GEN_LEN})" +if (( AGENT_AUTOCOMPACT_AT >= AGENT_PROMPT_CEILING )); then + echo "WARNING: the compaction trigger (${AGENT_AUTOCOMPACT_AT}) is at or above the prompt ceiling" \ + "(${AGENT_PROMPT_CEILING}); turns will hit finish_reason=length before compaction fires." \ + "Lower AGENT_AUTOCOMPACT_PCT to <= $((AGENT_PROMPT_CEILING * 100 / AGENT_MAX_CONTEXT_TOKENS))." +fi +if (( AGENT_MAX_CONTEXT_TOKENS > MAX_CONTEXT_LEN )); then + echo "WARNING: declared context ${AGENT_MAX_CONTEXT_TOKENS} > MAX_CONTEXT_LEN ${MAX_CONTEXT_LEN};" \ + "the adapter hard-stops at the latter." +fi + +# The only two harness knobs per agent: extra CLI flags, and extra env vars as +# JSON. Everything else (denied tools, the launch flags) is a class attribute in +# slime_plugins/.../harnesses.py, because it is a property of the harness rather +# than of a run. Both are applied LAST -- EXTRA_ARGS after the harness's own +# flags (claude takes the last occurrence of a repeated flag, verified against +# the real CLI) and EXTRA_ENVS after static_env -- so either can override a +# harness default. The *_EXTRA_ENVS pair is set in the auto-compaction block +# above; these two are the remaining passthroughs, declared so SWE_AGENT can be +# either harness from this script (run_cbc_*.sh just sets it and re-execs). +export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:-}" +export SLIME_AGENT_CBC_EXTRA_ARGS="${SLIME_AGENT_CBC_EXTRA_ARGS:-}" + +# ============ proxy bypass for in-cluster/AGS traffic ============ +export no_proxy="127.0.0.1,${MASTER_ADDR},${ADAPTER_PUBLIC_HOST},${E2B_DOMAIN},.tencentags.com" +export NO_PROXY="${no_proxy}" + +cd "${SLIME_DIR}" +source "${SLIME_DIR}/scripts/models/qwen3.5-35B-A3B.sh" + +CKPT_ARGS=( + --hf-checkpoint "${HF_CHECKPOINT}" + --ref-load "${REF_MODEL_PATH}" +) +if [[ -n "${LOAD_DIR}" ]]; then + CKPT_ARGS+=(--load "${LOAD_DIR}") +fi +if [[ -n "${CKPT_STEP}" ]]; then + # Megatron's get_load_checkpoint_path_by_args honours ckpt_step over the + # tracker file, which is how one specific step gets scored. + CKPT_ARGS+=(--ckpt-step "${CKPT_STEP}") +fi +# No --save/--save-interval: nothing is trained, so nothing should be written. + +ROLLOUT_ARGS=( + # Eval reaches AGS through slime's standard sglang eval loop plus this hook. + # --rollout-function-path and --rollout-buffer-url are deliberately absent: + # the eval path never writes to the rollout buffer, so buffer.py is not run. + --custom-generate-function-path slime_plugins.rollout_buffer.generator.ags_generator.generate + --custom-eval-rollout-log-function-path slime_plugins.rollout_buffer.generator.ags_generator.wandb_metrics.log_eval_rollout_data + --rollout-task-type ags + # NOTE: no --apply-chat-template, unlike the training script. Under + # SWE_EVAL_PROMPT_STYLE=dataset the agent is handed sample.prompt verbatim + # (swe_task._coerce_prompt), and the flag makes slime's Dataset run + # tokenizer.apply_chat_template over the row first -- which would hand Claude + # Code a prompt literally containing "<|im_start|>user ... <|im_end|>". + # The eval path renders its own chat template per turn inside the adapter. + --input-key prompt + --label-key label + --metadata-key metadata + --num-rollout "${NUM_ROLLOUT}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" + --rollout-max-context-len "${MAX_CONTEXT_LEN}" + --rollout-max-response-len "${MAX_GEN_LEN}" + --rollout-stop-token-ids 248046 248044 + --global-batch-size "${GLOBAL_BATCH_SIZE}" + --micro-batch-size "${MICRO_BATCH_SIZE}" + --loss-mask-type qwen3_5 +) + +EVAL_ARGS=( + --eval-function-path slime.rollout.sglang_rollout.generate_rollout + --eval-interval "${EVAL_INTERVAL}" + --eval-prompt-data "${EVAL_DATASET_NAME}" "${EVAL_DATA}" + --n-samples-per-eval-prompt "${N_SAMPLES_PER_EVAL_PROMPT}" + --eval-max-prompt-len "${ROLLOUT_MAX_PROMPT_LEN}" + --eval-max-response-len "${MAX_GEN_LEN}" + --eval-temperature "${EVAL_TEMPERATURE}" + --eval-top-p "${EVAL_TOP_P}" + --eval-top-k "${EVAL_TOP_K}" +) + +# train_iters = num_rollout * rollout_batch_size * n_samples_per_prompt // +# global_batch_size = 0 here, and Megatron's OptimizerParamScheduler asserts on +# a zero-length schedule. Three asserts fire in sequence otherwise: +# 1. lr_decay_steps > 0 -> --lr-decay-iters +# 2. total-iterations mismatch vs the ckpt -> the value must equal the +# *original* run's num_rollout, not any positive number +# 3. weight-decay-iterations mismatch -> no override arg exists, so take +# the whole scheduler state from the checkpoint +# TRAIN_NUM_ROLLOUT must therefore be the --num-rollout of the run that wrote +# this checkpoint. Only needed when actually loading a Megatron checkpoint. +# --lr-decay-iters is needed even with no --load: train_iters is 0 either way, so +# `assert self.lr_decay_steps > 0` fires while building the scheduler, before any +# checkpoint is touched. For the base model the value is arbitrary (nothing is +# trained and no scheduler state is compared), so default it to 1. +LR_SCHED_ARGS=(--lr-decay-iters "${LR_DECAY_ITERS_BASE:-1}") +if [[ -n "${LOAD_DIR}" ]]; then + TRAIN_NUM_ROLLOUT="${TRAIN_NUM_ROLLOUT:?set TRAIN_NUM_ROLLOUT to the --num-rollout of the run that wrote this checkpoint (e.g. 100)}" + LR_SCHED_ARGS=( + --lr-decay-iters $((TRAIN_NUM_ROLLOUT * ROLLOUT_BATCH_SIZE * N_SAMPLES_PER_PROMPT / GLOBAL_BATCH_SIZE)) + # Only the weights are needed: nothing is trained here, so the optimizer and + # RNG state in the checkpoint are dead weight -- and loading them actually + # crashes. Megatron's generate_state_dict -> DistributedOptimizer + # .sharded_state_dict -> load_state_dict -> dummy_step() reaches + # hybrid_optimizer._set_sub_optimizer_grads, whose torch.empty for the CPU + # offload copy map dies with "CUDA error: invalid argument" (this + # checkpoint was written with --optimizer-cpu-offload). slime itself takes + # this same route whenever it wants weights only (see arguments.py where it + # sets no_load_optim/no_load_rng/finetune together). + # + # --finetune also stops Megatron restoring the iteration counter, which is + # what --use-checkpoint-opt-param-scheduler was working around; keep the + # explicit --lr-decay-iters above so the scheduler still constructs. + --no-load-optim + --no-load-rng + --finetune + ) +fi + +PERF_ARGS=( + --tensor-model-parallel-size "${TP_SIZE}" + --sequence-parallel + --pipeline-model-parallel-size "${PP_SIZE}" + --context-parallel-size "${CP_SIZE}" + --expert-model-parallel-size "${EP_SIZE}" + --expert-tensor-parallel-size "${ETP_SIZE}" + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + --max-tokens-per-gpu $((MAX_CONTEXT_LEN / CP_SIZE)) + --log-probs-chunk-size 1024 + --use-dynamic-batch-size +) + +# Kept because slime constructs the optimizer/scheduler even for eval-only; none +# of these values affect the score. +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +ALGO_ARGS=( + --advantage-estimator grpo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +SGLANG_ARGS=( + --rollout-num-gpus "${ROLLOUT_NUM_GPUS}" + --rollout-num-gpus-per-engine "${ROLLOUT_TP_SIZE}" + --sglang-mem-fraction-static "${ROLLOUT_MEM_UTILIZATION}" + --sglang-enable-dp-attention + --sglang-dp-size "${ROLLOUT_DP_SIZE}" + --sglang-ep-size "${ROLLOUT_EP_SIZE}" + --sglang-enable-dp-lm-head + --sglang-moe-dense-tp-size 1 + --sglang-tool-call-parser qwen3_coder + --sglang-reasoning-parser qwen3 +) + +if [[ -n "${WANDB_API_KEY:-}" ]]; then + WANDB_ARGS=( + --use-wandb + --wandb-team "${WANDB_ENTITY:?WANDB_ENTITY is required when WandB is enabled}" + --wandb-project "${WANDB_PROJECT:-slime-claude-code-ags}" + --wandb-group "${WANDB_GROUP:-${EXP_TAG}}" + --wandb-key "${WANDB_API_KEY}" + --wandb-dir "${LOG_DIR}" + --disable-wandb-random-suffix + ) +else + WANDB_ARGS=() +fi + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --moe-token-dispatcher-type flex + --moe-enable-deepep + --colocate + --log-passrate +) + +# ============ bring up ray cluster ============ +# No rollout buffer here: the eval path talks to AGS directly. +HOSTFILE="${HOSTFILE:-/root/mpi_rack_hostfile}" + +# Every Ray port is overridable because these H20 nodes routinely host other +# users' Ray clusters. The dashboard AGENT port matters most: it defaults to +# 52365, and when that is taken the agent dies with "address already in use" +# while `ray status` still reports a healthy cluster -- `ray job submit` then +# fails with "No available agent to submit job", which looks nothing like a port +# conflict. Ports below match the training launcher's defaults. +# +# --port must be MASTER_PORT: the worker loop below dials +# ${MASTER_ADDR}:${MASTER_PORT}, but without --port the head's GCS always binds +# 6379, so any non-default MASTER_PORT leaves workers dialing a dead port and +# timing out on "Failed to connect to GCS". +ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" \ + --port "${MASTER_PORT}" \ + --disable-usage-stats --dashboard-host=0.0.0.0 \ + --dashboard-port "${RAY_DASHBOARD_PORT:-8265}" \ + --dashboard-agent-listen-port "${RAY_DASHBOARD_AGENT_LISTEN_PORT:-28065}" \ + --dashboard-agent-grpc-port "${RAY_DASHBOARD_AGENT_GRPC_PORT:-28066}" \ + --runtime-env-agent-port "${RAY_RUNTIME_ENV_AGENT_PORT:-28067}" \ + --metrics-export-port "${RAY_METRICS_EXPORT_PORT:-28068}" + +if [[ -f "${HOSTFILE}" ]]; then + WORKER_LIMIT=$((ACTOR_NUM_NODES - 1)) + STARTED_WORKERS=0 + for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do + [[ -z "${WORKER_IP}" ]] && continue + [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]] && continue + if (( STARTED_WORKERS >= WORKER_LIMIT )); then + break + fi + echo "Starting Ray worker on ${WORKER_IP}" + ssh -o StrictHostKeyChecking=no "root@${WORKER_IP}" \ + "pkill -9 sglang ; ray stop --force ; pkill -9 python ; \ + ray start --address=${MASTER_ADDR}:${MASTER_PORT} --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} \ + --node-ip-address ${WORKER_IP} --disable-usage-stats" & + STARTED_WORKERS=$((STARTED_WORKERS + 1)) + done + for pid in $(jobs -pr); do + wait "${pid}" + done + if (( STARTED_WORKERS < WORKER_LIMIT )); then + echo "WARNING: requested ${ACTOR_NUM_NODES} nodes but only started $((STARTED_WORKERS + 1)) including head." + fi +else + echo "WARNING: HOSTFILE=${HOSTFILE} not found; only the head node was started." +fi + +# Wait for every GPU to actually register instead of sleeping a fixed 30s. Slime +# asks for a placement group of ACTOR_NUM_NODES x ACTOR_NUM_GPUS_PER_NODE GPUs and +# hangs on "1+ pending placement groups" if a worker is late or never joined -- +# which is easy to miss, because `ray status` reports a perfectly healthy cluster +# with only the head's GPUs. This also covers externally started workers (the +# container has no root ssh key, so workers may be joined from outside). +EXPECTED_GPUS=$((ACTOR_NUM_NODES * ACTOR_NUM_GPUS_PER_NODE)) +RAY_WAIT_SEC="${RAY_WAIT_SEC:-300}" +echo "Waiting for ${EXPECTED_GPUS} GPUs to register with Ray (timeout ${RAY_WAIT_SEC}s)..." +for ((waited = 0; waited < RAY_WAIT_SEC; waited += 10)); do + TOTAL_GPUS=$(ray status 2>/dev/null | grep -oE '[0-9.]+/[0-9]+\.[0-9]+ GPU' | head -1 | sed -E 's|.*/([0-9]+)\.[0-9]+ GPU|\1|') + if [[ "${TOTAL_GPUS:-0}" -ge "${EXPECTED_GPUS}" ]]; then + echo "Ray has ${TOTAL_GPUS} GPUs after ${waited}s." + break + fi + echo " ${waited}s: ${TOTAL_GPUS:-0}/${EXPECTED_GPUS} GPUs registered" + sleep 10 +done +if [[ "${TOTAL_GPUS:-0}" -lt "${EXPECTED_GPUS}" ]]; then + echo "ERROR: only ${TOTAL_GPUS:-0}/${EXPECTED_GPUS} GPUs registered after ${RAY_WAIT_SEC}s." + echo " Join the missing worker(s), or the job will hang on a pending placement group." + ray status || true + exit 1 +fi +ray status + +# ============ runtime env propagated to ray workers ============ +export SLIME_DIR +RUNTIME_ENV_JSON=$(python3 - <&1 | tee "${LOG_FILE}" + +echo "======================================================================" +echo "RUN_ROOT=${RUN_ROOT}" +echo "Score: grep -E 'eval/${EVAL_DATASET_NAME}' ${LOG_FILE}" +echo "Trajectories/patches: ${TRAJECTORY_DUMP_DIR}" +echo "======================================================================" diff --git a/examples/coding_agent_rl_ags/run_cbc_qwen35_35b_a3b_swe_2nodes.sh b/examples/coding_agent_rl_ags/run_cbc_qwen35_35b_a3b_swe_2nodes.sh new file mode 100644 index 0000000000..b3ef8cac67 --- /dev/null +++ b/examples/coding_agent_rl_ags/run_cbc_qwen35_35b_a3b_swe_2nodes.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash +# GRPO training on SWE-rebench with the CodeBuddy Code (cbc) harness, 2 nodes. +# +# Thin wrapper: the launcher logic lives in run_cc_qwen35_35b_a3b_swe_2nodes.sh +# and is harness-agnostic -- SWE_AGENT selects the harness, and that script sets +# both harnesses' EXTRA_ARGS/EXTRA_ENVS knobs. Duplicating ~370 lines per harness +# would guarantee the two copies drift, which is the failure this rename cleans up. +# +# Everything the CC script accepts works here too, e.g. +# PROMPT_DATA=... NUM_ROLLOUT=50 bash run_cbc_qwen35_35b_a3b_swe_2nodes.sh +# +# Usage (from a long-lived shell / tmux on the head node, inside the container): +# bash examples/coding_agent_rl_ags/run_cbc_qwen35_35b_a3b_swe_2nodes.sh +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" + +export SWE_AGENT=codebuddy_code +export EXP_TAG="${EXP_TAG:-coding_agent_rl_ags_cbc_qwen35_35b_a3b_2nodes}" + +exec bash "${SCRIPT_DIR}/run_cc_qwen35_35b_a3b_swe_2nodes.sh" "$@" diff --git a/examples/coding_agent_rl_ags/run_cbc_qwen35_35b_a3b_swe_4nodes.sh b/examples/coding_agent_rl_ags/run_cbc_qwen35_35b_a3b_swe_4nodes.sh new file mode 100644 index 0000000000..cc4774b4da --- /dev/null +++ b/examples/coding_agent_rl_ags/run_cbc_qwen35_35b_a3b_swe_4nodes.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash +# GRPO training on SWE-rebench with the CodeBuddy Code (cbc) harness, 4 nodes. +# +# Thin wrapper around the harness-agnostic 4-node launcher; see +# run_cbc_qwen35_35b_a3b_swe_2nodes.sh for why these are wrappers rather than +# copies. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" + +export SWE_AGENT=codebuddy_code +export EXP_TAG="${EXP_TAG:-coding_agent_rl_ags_cbc_qwen35_35b_a3b_4nodes}" + +exec bash "${SCRIPT_DIR}/run_cc_qwen35_35b_a3b_swe_4nodes.sh" "$@" diff --git a/examples/coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_2nodes.sh b/examples/coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_2nodes.sh new file mode 100644 index 0000000000..3a1d23aaa5 --- /dev/null +++ b/examples/coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_2nodes.sh @@ -0,0 +1,444 @@ +#!/usr/bin/env bash +# End-to-end SWE coding-agent RL with Claude Code on AGS, using rollout_buffer +# on a 2-node Ray cluster. Run from a long-lived shell / tmux session on the +# Ray head node. + +# Best-effort cleanup so a rerun does not collide with stale workers/services. +pkill -9 sglang || true +pkill -f "slime_plugins.rollout_buffer.buffer" || true +pkill -f "slime_plugins/rollout_buffer/buffer.py" || true +sleep 3 +ray stop --force || true +pkill -9 ray || true +sleep 3 +pkill -9 ray || true + +set -ex + +export PYTHONUNBUFFERED=1 + +EXP="${EXP:?set EXP to an experiment directory, e.g. /data_train/ericxjzheng/experiments/}" +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +SLIME_DIR="${SLIME_DIR:-/data_train/ericxjzheng/workspace/slime}" + +# ============ cluster size ============ +ACTOR_NUM_NODES="${ACTOR_NUM_NODES:-${MLP_WORKER_NUM:-2}}" +ACTOR_NUM_GPUS_PER_NODE="${ACTOR_NUM_GPUS_PER_NODE:-8}" +TOTAL_NUM_GPUS=$((ACTOR_NUM_NODES * ACTOR_NUM_GPUS_PER_NODE)) + +# ============ model parallelism ============ +export TP_SIZE="${TP_SIZE:-2}" +export PP_SIZE="${PP_SIZE:-1}" +export CP_SIZE="${CP_SIZE:-8}" +export EP_SIZE="${EP_SIZE:-8}" +export ETP_SIZE="${ETP_SIZE:-1}" + +# ============ rollout engine ============ +ROLLOUT_NUM_GPUS="${ROLLOUT_NUM_GPUS:-${TOTAL_NUM_GPUS}}" +ROLLOUT_TP_SIZE="${ROLLOUT_TP_SIZE:-8}" +ROLLOUT_DP_SIZE="${ROLLOUT_DP_SIZE:-2}" +ROLLOUT_EP_SIZE="${ROLLOUT_EP_SIZE:-8}" +ROLLOUT_MEM_UTILIZATION="${ROLLOUT_MEM_UTILIZATION:-0.75}" +NUM_ROLLOUT="${NUM_ROLLOUT:-100}" +ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-8}" +N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-8}" +NUM_STEPS_PER_ROLLOUT="${NUM_STEPS_PER_ROLLOUT:-1}" +GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-$((ROLLOUT_BATCH_SIZE * N_SAMPLES_PER_PROMPT))}" +MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}" + +# ============ context length ============ +MAX_CONTEXT_LEN="${MAX_CONTEXT_LEN:-96000}" +MAX_GEN_LEN="${MAX_GEN_LEN:-32768}" +ROLLOUT_MAX_PROMPT_LEN="${ROLLOUT_MAX_PROMPT_LEN:-${MAX_CONTEXT_LEN}}" + +# ============ eval ============ +EVAL_INTERVAL="${EVAL_INTERVAL:-20}" +EVAL_DATA="${EVAL_DATA:-/data_train/ericxjzheng/data/SWE-bench_Verified_slime_rl_format_from_harbor/swebench_verified_slime_tcr.jsonl}" +SKIP_EVAL_BEFORE_TRAIN="${SKIP_EVAL_BEFORE_TRAIN:-1}" +N_SAMPLES_PER_EVAL_PROMPT="${N_SAMPLES_PER_EVAL_PROMPT:-1}" + +# ============ paths — override before launching ============ +HF_CHECKPOINT="${HF_CHECKPOINT:-/data_train/ericxjzheng/models/Qwen3.5-35B-A3B}" +REF_MODEL_PATH="${REF_MODEL_PATH:-/data_train/ericxjzheng/models/Qwen3.5-35B-A3B_torch_dist}" +PROMPT_DATA="${PROMPT_DATA:-/data_train/ericxjzheng/data/SWE-rebench-filtered/filtered.jsonl}" + +EXP_TAG="${EXP_TAG:-coding_agent_rl_ags_cc_qwen35_35b_a3b_2nodes}" +STAMP="$(date +%Y%m%d_%H%M%S)" +RUN_ROOT="${RUN_ROOT:-${EXP}/runs/${EXP_TAG}_${STAMP}}" + +# ============ logging/artifacts ============ +LOG_DIR="${RUN_ROOT}" +mkdir -p "${LOG_DIR}/rollout_dumps" "${LOG_DIR}/ags_artifacts" +LOG_FILE="${LOG_DIR}/run.log" +BUFFER_LOG_FILE="${LOG_DIR}/rollout_buffer.log" +export TRAJECTORY_DUMP_DIR="${TRAJECTORY_DUMP_DIR:-${LOG_DIR}/ags_artifacts}" +export EXPERIMENT_NAME="${EXPERIMENT_NAME:-${EXP_TAG}}" +echo "======================================================================" +echo "Training log: ${LOG_FILE}" +echo "Rollout buffer log: ${BUFFER_LOG_FILE}" +echo "RUN_ROOT=${RUN_ROOT}" +echo "======================================================================" + +# ============ ray cluster network ============ +# Set MASTER_ADDR before AGS/SWE blocks: ADAPTER_PUBLIC_HOST below falls back to it. +export MASTER_ADDR="${MASTER_ADDR:-${MLP_WORKER_0_HOST:-$(hostname -I | awk '{print $1}')}}" +export MASTER_PORT="${MASTER_PORT:-${MLP_WORKER_0_PORT:-6379}}" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" + +# ============ SWE / Claude Code / AGS rollout knobs ============ +export SWE_AGENT="${SWE_AGENT:-claude_code}" + +# AGS uses the E2B-compatible SDK surface. Export E2B_API_KEY in the launch +# environment (for Tencent AGS this is typically the AGS gateway key). +export E2B_DOMAIN="${E2B_DOMAIN:-ap-shanghai.tencentags.com}" +export AGS_BASE_TOOL="${AGS_BASE_TOOL:-sdt-3fzh6mv6}" +export AGS_IMAGE_REGISTRY_TYPE="${AGS_IMAGE_REGISTRY_TYPE:-enterprise}" +export AGS_SANDBOX_RESOURCES_JSON=${AGS_SANDBOX_RESOURCES_JSON:-'{"cpu":"4","memory":"16Gi"}'} + +# ADAPTER_PUBLIC_HOST must be routable from inside the AGS sandbox (not 127.0.0.1). +export ADAPTER_PUBLIC_HOST="${ADAPTER_PUBLIC_HOST:-${MASTER_ADDR:-${MLP_WORKER_0_HOST:-127.0.0.1}}}" +export ADAPTER_BIND_HOST="${ADAPTER_BIND_HOST:-0.0.0.0}" +export ADAPTER_PORT="${ADAPTER_PORT:-18001}" + +export SWE_AGENT_TIME_BUDGET_SEC="${SWE_AGENT_TIME_BUDGET_SEC:-1800}" +export SWE_EVAL_TIMEOUT_SEC="${SWE_EVAL_TIMEOUT_SEC:-600}" +# false: grade in the agent sandbox; true: boot a second clean sandbox for grading. +export SWE_EVAL_ISOLATED_SANDBOX="${SWE_EVAL_ISOLATED_SANDBOX:-false}" +export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-32}" +export SWE_BOOT_RETRIES="${SWE_BOOT_RETRIES:-10}" +export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-32}" + +# How the agent receives the task, set separately for training and periodic eval. +# instruction: send SWE_CC_PROMPT and write PROBLEM_STATEMENT.md; the agent +# spends its first turns finding and reading that file. +# dataset: hand the row's prompt over directly and write no file. +# Eval uses dataset so the score reflects the model working the task itself, +# the way Harbor measures it, rather than its file-discovery turns. For converted +# Harbor prompt data both paths carry the same instruction.md text, so this +# changes when the agent sees the task, not what it reads. +export SWE_PROMPT_STYLE="${SWE_PROMPT_STYLE:-instruction}" +export SWE_EVAL_PROMPT_STYLE="${SWE_EVAL_PROMPT_STYLE:-dataset}" + +# ---- auto-compaction ------------------------------------------------------- +# Compact before a segment crosses the training-side context cap, otherwise the +# adapter returns finish_reason="length" with zero output tokens once the prompt +# reaches rollout_max_context_len (slime/agent/adapters/common.py) and the turn is +# wasted. +# +# Expressed as a PERCENTAGE of the model's context window rather than an absolute +# token count: an absolute auto-compact window is clamped to [100k, 1M] by both +# CLIs, which is above our 96k cap, so it could never fire in time. Both CLIs are +# also told the real window, so the percentage is of MAX_CONTEXT_LEN on each side. +# +# CodeBuddy CODEBUDDY_AUTOCOMPACT_PCT_OVERRIDE (cbc 1.9.33) is a percentage of +# the model's maxInputTokens, which the harness writes into +# models.json from SLIME_AGENT_MAX_INPUT_TOKENS. Without that key +# resolveCompactTriggerAt() falls back to the clamped absolute +# window. Source: agent-cli src/node/context/context-protocol.ts. +# Claude Code CLAUDE_AUTOCOMPACT_PCT_OVERRIDE, as a percentage of the context +# window, which CLAUDE_CODE_MAX_CONTEXT_TOKENS sets. That variable +# applies DIRECTLY for model names claude does not recognise as a +# Claude model -- ours is "slime-actor", so it does (the binary +# gates on `!normalize(model).startsWith("claude-")`; for a real +# claude-* name it would need DISABLE_COMPACT too, which would +# disable the compaction we want). Without it the CLI assumes its +# 200000 default and would compact at pct% of 200k, i.e. never +# before our 96k cap. +# 60%, not the CLI's 70% default, because the per-turn output reservation eats +# into what is reachable: claude assumes MAX_OUTPUT_TOKENS (32000 by default for +# model ids it does not recognise, which includes ours) is available on top of the +# prompt, so the prompt cannot grow past MAX_CONTEXT_LEN - MAX_GEN_LEN ~= 63k +# before the adapter's hard stop. A 70% trigger (67200) sits ABOVE that and would +# never be reached; 60% (57600) fires with room to spare. Raising +# AGENT_AUTOCOMPACT_PCT re-opens that gap -- the check below says so out loud. +AGENT_AUTOCOMPACT_PCT="${AGENT_AUTOCOMPACT_PCT:-60}" +AGENT_MAX_CONTEXT_TOKENS="${AGENT_MAX_CONTEXT_TOKENS:-${MAX_CONTEXT_LEN}}" + +# cbc: models.json maxInputTokens (written by CodeBuddyCodeHarness). +export SLIME_AGENT_MAX_INPUT_TOKENS="${SLIME_AGENT_MAX_INPUT_TOKENS:-${AGENT_MAX_CONTEXT_TOKENS}}" +# Built in a separate variable, not inline in ${VAR:-...}: a JSON default inside +# that expansion is mis-parsed -- the value's own "}" closes the expansion early +# and the trailing brace leaks into the result ("{...}}"). +CBC_AUTOCOMPACT_ENVS="{\"CODEBUDDY_AUTOCOMPACT_PCT_OVERRIDE\":\"${AGENT_AUTOCOMPACT_PCT}\"}" +export SLIME_AGENT_CBC_EXTRA_ENVS="${SLIME_AGENT_CBC_EXTRA_ENVS:-${CBC_AUTOCOMPACT_ENVS}}" + +# claude: declare the window and the per-turn output budget, then set the trigger +# percentage. MAX_OUTPUT_TOKENS is pinned to MAX_GEN_LEN rather than left at the +# CLI's 32000 default-for-unknown-model-ids so the reservation matches what the +# adapter will actually serve (--rollout-max-response-len). +CC_AUTOCOMPACT_ENVS="{\"CLAUDE_CODE_MAX_CONTEXT_TOKENS\":\"${AGENT_MAX_CONTEXT_TOKENS}\",\"CLAUDE_CODE_MAX_OUTPUT_TOKENS\":\"${MAX_GEN_LEN}\",\"CLAUDE_AUTOCOMPACT_PCT_OVERRIDE\":\"${AGENT_AUTOCOMPACT_PCT}\"}" +export SLIME_AGENT_CC_EXTRA_ENVS="${SLIME_AGENT_CC_EXTRA_ENVS:-${CC_AUTOCOMPACT_ENVS}}" + +AGENT_AUTOCOMPACT_AT=$((AGENT_AUTOCOMPACT_PCT * AGENT_MAX_CONTEXT_TOKENS / 100)) +# The largest prompt that can still be served: the adapter caps prompt+output at +# rollout_max_context_len, so a turn needs MAX_GEN_LEN of headroom. +AGENT_PROMPT_CEILING=$((MAX_CONTEXT_LEN - MAX_GEN_LEN)) +echo "Auto-compact: ${AGENT_AUTOCOMPACT_PCT}% of ${AGENT_MAX_CONTEXT_TOKENS} = ${AGENT_AUTOCOMPACT_AT} tokens" \ + "(prompt ceiling ${AGENT_PROMPT_CEILING} = ${MAX_CONTEXT_LEN} - ${MAX_GEN_LEN})" +if (( AGENT_AUTOCOMPACT_AT >= AGENT_PROMPT_CEILING )); then + echo "WARNING: the compaction trigger (${AGENT_AUTOCOMPACT_AT}) is at or above the prompt ceiling" \ + "(${AGENT_PROMPT_CEILING}); turns will hit finish_reason=length before compaction fires." \ + "Lower AGENT_AUTOCOMPACT_PCT to <= $((AGENT_PROMPT_CEILING * 100 / AGENT_MAX_CONTEXT_TOKENS))." +fi +if (( AGENT_MAX_CONTEXT_TOKENS > MAX_CONTEXT_LEN )); then + echo "WARNING: declared context ${AGENT_MAX_CONTEXT_TOKENS} > MAX_CONTEXT_LEN ${MAX_CONTEXT_LEN};" \ + "the adapter hard-stops at the latter." +fi + +# # `investigator` is a read-only sub-agent, dispatched via the Agent tool. +# AGENTS_JSON='{"investigator":{"description":"Searches the repo for relevant files before any edit","prompt":"You are an investigator sub-agent. Use Grep/Read/Glob to find every file relevant to the user task, then return a short bulleted summary. Do NOT edit anything.","tools":["Grep","Read","Glob"]}}' +# export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS} --disable-slash-commands --agents '${AGENTS_JSON}'" +# (WebFetch/WebSearch are already denied by the harness default; EXTRA_ARGS is +# appended last, so anything set here overrides a repeated default flag.) +# The only two harness knobs per agent: extra CLI flags, and extra env vars as +# JSON. Everything else (denied tools, the launch flags) is a class attribute in +# slime_plugins/.../harnesses.py, because it is a property of the harness rather +# than of a run. Both are applied LAST -- EXTRA_ARGS after the harness's own +# flags (claude takes the last occurrence of a repeated flag, verified against +# the real CLI) and EXTRA_ENVS after static_env -- so either can override a +# harness default. The *_EXTRA_ENVS pair is set in the auto-compaction block +# above; these two are the remaining passthroughs, declared so SWE_AGENT can be +# either harness from this script (run_cbc_*.sh just sets it and re-execs). +export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:-}" +export SLIME_AGENT_CBC_EXTRA_ARGS="${SLIME_AGENT_CBC_EXTRA_ARGS:-}" + +# Optional: require dispatching the investigator before any edit, to maximize sub-agent fan-out. +# export SWE_CC_PROMPT="Read PROBLEM_STATEMENT.md. BEFORE editing any file, dispatch the 'investigator' sub-agent (via the Agent tool with subagent_type=investigator) to locate every file relevant to the issue. Then fix the issue and run the tests." + +# ============ proxy bypass for in-cluster/AGS traffic ============ +export no_proxy="127.0.0.1,${MASTER_ADDR},${ADAPTER_PUBLIC_HOST},${E2B_DOMAIN},.tencentags.com" +export NO_PROXY="${no_proxy}" + +cd "${SLIME_DIR}" +source "${SLIME_DIR}/scripts/models/qwen3.5-35B-A3B.sh" + +SAVE_DIR="${SAVE_DIR:-${EXP}/checkpoints}" +SAVE_INTERVAL="${SAVE_INTERVAL:-5}" +mkdir -p "${SAVE_DIR}" + +CKPT_ARGS=( + --hf-checkpoint "${HF_CHECKPOINT}" + --ref-load "${REF_MODEL_PATH}" + --load "${LOAD_DIR:-${SAVE_DIR}}" + --save "${SAVE_DIR}" + --save-interval "${SAVE_INTERVAL}" +) + +ROLLOUT_ARGS=( + --rollout-function-path slime_plugins.rollout_buffer.rollout_buffer_example.generate_rollout + # Used by periodic eval, which runs AGS through slime's standard sglang eval loop. + --custom-generate-function-path slime_plugins.rollout_buffer.generator.ags_generator.generate + --custom-rollout-log-function-path slime_plugins.rollout_buffer.generator.ags_generator.wandb_metrics.log_rollout_data + --custom-eval-rollout-log-function-path slime_plugins.rollout_buffer.generator.ags_generator.wandb_metrics.log_eval_rollout_data + --rollout-task-type ags + --rollout-buffer-url "http://${MASTER_ADDR}:8889" + --prompt-data "${PROMPT_DATA}" + --input-key prompt + --label-key label + --metadata-key metadata + --apply-chat-template + --rollout-shuffle + --num-rollout "${NUM_ROLLOUT}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" + --rollout-buffer-num-epoch 1 + --rollout-buffer-stop-timeout-sec "${ROLLOUT_BUFFER_STOP_TIMEOUT_SEC:-120}" + --rollout-max-context-len "${MAX_CONTEXT_LEN}" + --rollout-max-response-len "${MAX_GEN_LEN}" + --rollout-temperature 1.0 + --rollout-stop-token-ids 248046 248044 + --num-steps-per-rollout "${NUM_STEPS_PER_ROLLOUT}" + --global-batch-size "${GLOBAL_BATCH_SIZE}" + --micro-batch-size "${MICRO_BATCH_SIZE}" + --loss-mask-type qwen3_5 + --save-debug-rollout-data "${RUN_ROOT}/rollout_dumps/rollout_{rollout_id}.pt" +) + +EVAL_ARGS=( + --eval-function-path slime.rollout.sglang_rollout.generate_rollout + --eval-interval "${EVAL_INTERVAL}" + --eval-prompt-data swebench_verified "${EVAL_DATA}" + --n-samples-per-eval-prompt "${N_SAMPLES_PER_EVAL_PROMPT}" + --eval-max-prompt-len "${ROLLOUT_MAX_PROMPT_LEN}" + --eval-max-response-len "${MAX_GEN_LEN}" + --eval-temperature 0.7 + --eval-top-p 0.8 + --eval-top-k 20 +) + +if [[ "${SKIP_EVAL_BEFORE_TRAIN}" == "1" || "${SKIP_EVAL_BEFORE_TRAIN}" == "true" ]]; then + EVAL_ARGS+=(--skip-eval-before-train) +fi + +PERF_ARGS=( + --tensor-model-parallel-size "${TP_SIZE}" + --sequence-parallel + --pipeline-model-parallel-size "${PP_SIZE}" + --context-parallel-size "${CP_SIZE}" + --expert-model-parallel-size "${EP_SIZE}" + --expert-tensor-parallel-size "${ETP_SIZE}" + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + # max-tokens-per-gpu is one CP rank's slice of MAX_CONTEXT_LEN; log-probs are + # chunked along T to avoid OOM on long single trajectories. + --max-tokens-per-gpu $((MAX_CONTEXT_LEN / CP_SIZE)) + --log-probs-chunk-size 1024 + --use-dynamic-batch-size +) + +ALGO_ARGS=( + --advantage-estimator grpo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +SGLANG_ARGS=( + --rollout-num-gpus "${ROLLOUT_NUM_GPUS}" + --rollout-num-gpus-per-engine "${ROLLOUT_TP_SIZE}" + --sglang-mem-fraction-static "${ROLLOUT_MEM_UTILIZATION}" + --sglang-enable-dp-attention + --sglang-dp-size "${ROLLOUT_DP_SIZE}" + --sglang-ep-size "${ROLLOUT_EP_SIZE}" + --sglang-enable-dp-lm-head + --sglang-moe-dense-tp-size 1 + --sglang-tool-call-parser qwen3_coder + --sglang-reasoning-parser qwen3 +) + +if [[ -n "${WANDB_API_KEY:-}" ]]; then + WANDB_ARGS=( + --use-wandb + --wandb-team "${WANDB_ENTITY:?WANDB_ENTITY is required when WandB is enabled}" + --wandb-project "${WANDB_PROJECT:-slime-claude-code-ags}" + --wandb-group "${WANDB_GROUP:-${EXP_TAG}}" + --wandb-key "${WANDB_API_KEY}" + --wandb-dir "${LOG_DIR}" + --disable-wandb-random-suffix + ) +else + WANDB_ARGS=() +fi + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --moe-token-dispatcher-type flex + --moe-enable-deepep + --colocate + --log-passrate +) + +# ============ bring up rollout buffer ============ +python3 -u -m slime_plugins.rollout_buffer.buffer >"${BUFFER_LOG_FILE}" 2>&1 & +BUFFER_PID=$! +trap 'kill ${BUFFER_PID} 2>/dev/null || true' EXIT +sleep 5 + +# ============ bring up ray cluster ============ +HOSTFILE="${HOSTFILE:-/root/mpi_rack_hostfile}" + +ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" \ + --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +if [[ -f "${HOSTFILE}" ]]; then + WORKER_LIMIT=$((ACTOR_NUM_NODES - 1)) + STARTED_WORKERS=0 + for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do + [[ -z "${WORKER_IP}" ]] && continue + [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]] && continue + if (( STARTED_WORKERS >= WORKER_LIMIT )); then + break + fi + echo "Starting Ray worker on ${WORKER_IP}" + ssh -o StrictHostKeyChecking=no "root@${WORKER_IP}" \ + "pkill -9 sglang ; ray stop --force ; pkill -9 python ; \ + ray start --address=${MASTER_ADDR}:${MASTER_PORT} --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} \ + --node-ip-address ${WORKER_IP} --disable-usage-stats" & + STARTED_WORKERS=$((STARTED_WORKERS + 1)) + done + for pid in $(jobs -pr); do + [[ "${pid}" == "${BUFFER_PID}" ]] && continue + wait "${pid}" + done + if (( STARTED_WORKERS < WORKER_LIMIT )); then + echo "WARNING: requested ${ACTOR_NUM_NODES} nodes but only started $((STARTED_WORKERS + 1)) including head." + fi +else + echo "WARNING: HOSTFILE=${HOSTFILE} not found; only the head node was started." +fi + +echo "Waiting for Ray cluster to stabilize..." +sleep 30 +ray status + +# ============ runtime env propagated to ray workers ============ +export SLIME_DIR +RUNTIME_ENV_JSON=$(python3 - <&1 | tee "${LOG_FILE}" + +echo "RUN_ROOT=${RUN_ROOT}" diff --git a/examples/coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_4nodes.sh b/examples/coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_4nodes.sh new file mode 100644 index 0000000000..f6b29b4d7e --- /dev/null +++ b/examples/coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_4nodes.sh @@ -0,0 +1,444 @@ +#!/usr/bin/env bash +# End-to-end SWE coding-agent RL with Claude Code on AGS, using rollout_buffer +# on a 4-node Ray cluster. Run from a long-lived shell / tmux session on the +# Ray head node. + +# Best-effort cleanup so a rerun does not collide with stale workers/services. +pkill -9 sglang || true +pkill -f "slime_plugins.rollout_buffer.buffer" || true +pkill -f "slime_plugins/rollout_buffer/buffer.py" || true +sleep 3 +ray stop --force || true +pkill -9 ray || true +sleep 3 +pkill -9 ray || true + +set -ex + +export PYTHONUNBUFFERED=1 + +EXP="${EXP:?set EXP to an experiment directory, e.g. /data_train/ericxjzheng/experiments/}" +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +SLIME_DIR="${SLIME_DIR:-/data_train/ericxjzheng/workspace/slime}" + +# ============ cluster size ============ +ACTOR_NUM_NODES="${ACTOR_NUM_NODES:-${MLP_WORKER_NUM:-4}}" +ACTOR_NUM_GPUS_PER_NODE="${ACTOR_NUM_GPUS_PER_NODE:-8}" +TOTAL_NUM_GPUS=$((ACTOR_NUM_NODES * ACTOR_NUM_GPUS_PER_NODE)) + +# ============ model parallelism ============ +export TP_SIZE="${TP_SIZE:-2}" +export PP_SIZE="${PP_SIZE:-1}" +export CP_SIZE="${CP_SIZE:-8}" +export EP_SIZE="${EP_SIZE:-8}" +export ETP_SIZE="${ETP_SIZE:-1}" + +# ============ rollout engine ============ +ROLLOUT_NUM_GPUS="${ROLLOUT_NUM_GPUS:-${TOTAL_NUM_GPUS}}" +ROLLOUT_TP_SIZE="${ROLLOUT_TP_SIZE:-8}" +ROLLOUT_DP_SIZE="${ROLLOUT_DP_SIZE:-4}" +ROLLOUT_EP_SIZE="${ROLLOUT_EP_SIZE:-8}" +ROLLOUT_MEM_UTILIZATION="${ROLLOUT_MEM_UTILIZATION:-0.75}" +NUM_ROLLOUT="${NUM_ROLLOUT:-100}" +ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-8}" +N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-8}" +NUM_STEPS_PER_ROLLOUT="${NUM_STEPS_PER_ROLLOUT:-1}" +GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-$((ROLLOUT_BATCH_SIZE * N_SAMPLES_PER_PROMPT))}" +MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}" + +# ============ context length ============ +MAX_CONTEXT_LEN="${MAX_CONTEXT_LEN:-96000}" +MAX_GEN_LEN="${MAX_GEN_LEN:-32768}" +ROLLOUT_MAX_PROMPT_LEN="${ROLLOUT_MAX_PROMPT_LEN:-${MAX_CONTEXT_LEN}}" + +# ============ eval ============ +EVAL_INTERVAL="${EVAL_INTERVAL:-20}" +EVAL_DATA="${EVAL_DATA:-/data_train/ericxjzheng/data/SWE-bench_Verified_slime_rl_format_from_harbor/swebench_verified_slime_tcr.jsonl}" +SKIP_EVAL_BEFORE_TRAIN="${SKIP_EVAL_BEFORE_TRAIN:-1}" +N_SAMPLES_PER_EVAL_PROMPT="${N_SAMPLES_PER_EVAL_PROMPT:-1}" + +# ============ paths — override before launching ============ +HF_CHECKPOINT="${HF_CHECKPOINT:-/data_train/ericxjzheng/models/Qwen3.5-35B-A3B}" +REF_MODEL_PATH="${REF_MODEL_PATH:-/data_train/ericxjzheng/models/Qwen3.5-35B-A3B_torch_dist}" +PROMPT_DATA="${PROMPT_DATA:-/data_train/ericxjzheng/data/SWE-rebench-filtered/filtered.jsonl}" + +EXP_TAG="${EXP_TAG:-coding_agent_rl_ags_cc_qwen35_35b_a3b_4nodes}" +STAMP="$(date +%Y%m%d_%H%M%S)" +RUN_ROOT="${RUN_ROOT:-${EXP}/runs/${EXP_TAG}_${STAMP}}" + +# ============ logging/artifacts ============ +LOG_DIR="${RUN_ROOT}" +mkdir -p "${LOG_DIR}/rollout_dumps" "${LOG_DIR}/ags_artifacts" +LOG_FILE="${LOG_DIR}/run.log" +BUFFER_LOG_FILE="${LOG_DIR}/rollout_buffer.log" +export TRAJECTORY_DUMP_DIR="${TRAJECTORY_DUMP_DIR:-${LOG_DIR}/ags_artifacts}" +export EXPERIMENT_NAME="${EXPERIMENT_NAME:-${EXP_TAG}}" +echo "======================================================================" +echo "Training log: ${LOG_FILE}" +echo "Rollout buffer log: ${BUFFER_LOG_FILE}" +echo "RUN_ROOT=${RUN_ROOT}" +echo "======================================================================" + +# ============ ray cluster network ============ +# Set MASTER_ADDR before AGS/SWE blocks: ADAPTER_PUBLIC_HOST below falls back to it. +export MASTER_ADDR="${MASTER_ADDR:-${MLP_WORKER_0_HOST:-$(hostname -I | awk '{print $1}')}}" +export MASTER_PORT="${MASTER_PORT:-${MLP_WORKER_0_PORT:-6379}}" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" + +# ============ SWE / Claude Code / AGS rollout knobs ============ +export SWE_AGENT="${SWE_AGENT:-claude_code}" + +# AGS uses the E2B-compatible SDK surface. Export E2B_API_KEY in the launch +# environment (for Tencent AGS this is typically the AGS gateway key). +export E2B_DOMAIN="${E2B_DOMAIN:-ap-shanghai.tencentags.com}" +export AGS_BASE_TOOL="${AGS_BASE_TOOL:-sdt-3fzh6mv6}" +export AGS_IMAGE_REGISTRY_TYPE="${AGS_IMAGE_REGISTRY_TYPE:-enterprise}" +export AGS_SANDBOX_RESOURCES_JSON=${AGS_SANDBOX_RESOURCES_JSON:-'{"cpu":"4","memory":"16Gi"}'} + +# ADAPTER_PUBLIC_HOST must be routable from inside the AGS sandbox (not 127.0.0.1). +export ADAPTER_PUBLIC_HOST="${ADAPTER_PUBLIC_HOST:-${MASTER_ADDR:-${MLP_WORKER_0_HOST:-127.0.0.1}}}" +export ADAPTER_BIND_HOST="${ADAPTER_BIND_HOST:-0.0.0.0}" +export ADAPTER_PORT="${ADAPTER_PORT:-18001}" + +export SWE_AGENT_TIME_BUDGET_SEC="${SWE_AGENT_TIME_BUDGET_SEC:-1800}" +export SWE_EVAL_TIMEOUT_SEC="${SWE_EVAL_TIMEOUT_SEC:-600}" +# false: grade in the agent sandbox; true: boot a second clean sandbox for grading. +export SWE_EVAL_ISOLATED_SANDBOX="${SWE_EVAL_ISOLATED_SANDBOX:-false}" +export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-32}" +export SWE_BOOT_RETRIES="${SWE_BOOT_RETRIES:-10}" +export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-32}" + +# How the agent receives the task, set separately for training and periodic eval. +# instruction: send SWE_CC_PROMPT and write PROBLEM_STATEMENT.md; the agent +# spends its first turns finding and reading that file. +# dataset: hand the row's prompt over directly and write no file. +# Eval uses dataset so the score reflects the model working the task itself, +# the way Harbor measures it, rather than its file-discovery turns. For converted +# Harbor prompt data both paths carry the same instruction.md text, so this +# changes when the agent sees the task, not what it reads. +export SWE_PROMPT_STYLE="${SWE_PROMPT_STYLE:-instruction}" +export SWE_EVAL_PROMPT_STYLE="${SWE_EVAL_PROMPT_STYLE:-dataset}" + +# ---- auto-compaction ------------------------------------------------------- +# Compact before a segment crosses the training-side context cap, otherwise the +# adapter returns finish_reason="length" with zero output tokens once the prompt +# reaches rollout_max_context_len (slime/agent/adapters/common.py) and the turn is +# wasted. +# +# Expressed as a PERCENTAGE of the model's context window rather than an absolute +# token count: an absolute auto-compact window is clamped to [100k, 1M] by both +# CLIs, which is above our 96k cap, so it could never fire in time. Both CLIs are +# also told the real window, so the percentage is of MAX_CONTEXT_LEN on each side. +# +# CodeBuddy CODEBUDDY_AUTOCOMPACT_PCT_OVERRIDE (cbc 1.9.33) is a percentage of +# the model's maxInputTokens, which the harness writes into +# models.json from SLIME_AGENT_MAX_INPUT_TOKENS. Without that key +# resolveCompactTriggerAt() falls back to the clamped absolute +# window. Source: agent-cli src/node/context/context-protocol.ts. +# Claude Code CLAUDE_AUTOCOMPACT_PCT_OVERRIDE, as a percentage of the context +# window, which CLAUDE_CODE_MAX_CONTEXT_TOKENS sets. That variable +# applies DIRECTLY for model names claude does not recognise as a +# Claude model -- ours is "slime-actor", so it does (the binary +# gates on `!normalize(model).startsWith("claude-")`; for a real +# claude-* name it would need DISABLE_COMPACT too, which would +# disable the compaction we want). Without it the CLI assumes its +# 200000 default and would compact at pct% of 200k, i.e. never +# before our 96k cap. +# 60%, not the CLI's 70% default, because the per-turn output reservation eats +# into what is reachable: claude assumes MAX_OUTPUT_TOKENS (32000 by default for +# model ids it does not recognise, which includes ours) is available on top of the +# prompt, so the prompt cannot grow past MAX_CONTEXT_LEN - MAX_GEN_LEN ~= 63k +# before the adapter's hard stop. A 70% trigger (67200) sits ABOVE that and would +# never be reached; 60% (57600) fires with room to spare. Raising +# AGENT_AUTOCOMPACT_PCT re-opens that gap -- the check below says so out loud. +AGENT_AUTOCOMPACT_PCT="${AGENT_AUTOCOMPACT_PCT:-60}" +AGENT_MAX_CONTEXT_TOKENS="${AGENT_MAX_CONTEXT_TOKENS:-${MAX_CONTEXT_LEN}}" + +# cbc: models.json maxInputTokens (written by CodeBuddyCodeHarness). +export SLIME_AGENT_MAX_INPUT_TOKENS="${SLIME_AGENT_MAX_INPUT_TOKENS:-${AGENT_MAX_CONTEXT_TOKENS}}" +# Built in a separate variable, not inline in ${VAR:-...}: a JSON default inside +# that expansion is mis-parsed -- the value's own "}" closes the expansion early +# and the trailing brace leaks into the result ("{...}}"). +CBC_AUTOCOMPACT_ENVS="{\"CODEBUDDY_AUTOCOMPACT_PCT_OVERRIDE\":\"${AGENT_AUTOCOMPACT_PCT}\"}" +export SLIME_AGENT_CBC_EXTRA_ENVS="${SLIME_AGENT_CBC_EXTRA_ENVS:-${CBC_AUTOCOMPACT_ENVS}}" + +# claude: declare the window and the per-turn output budget, then set the trigger +# percentage. MAX_OUTPUT_TOKENS is pinned to MAX_GEN_LEN rather than left at the +# CLI's 32000 default-for-unknown-model-ids so the reservation matches what the +# adapter will actually serve (--rollout-max-response-len). +CC_AUTOCOMPACT_ENVS="{\"CLAUDE_CODE_MAX_CONTEXT_TOKENS\":\"${AGENT_MAX_CONTEXT_TOKENS}\",\"CLAUDE_CODE_MAX_OUTPUT_TOKENS\":\"${MAX_GEN_LEN}\",\"CLAUDE_AUTOCOMPACT_PCT_OVERRIDE\":\"${AGENT_AUTOCOMPACT_PCT}\"}" +export SLIME_AGENT_CC_EXTRA_ENVS="${SLIME_AGENT_CC_EXTRA_ENVS:-${CC_AUTOCOMPACT_ENVS}}" + +AGENT_AUTOCOMPACT_AT=$((AGENT_AUTOCOMPACT_PCT * AGENT_MAX_CONTEXT_TOKENS / 100)) +# The largest prompt that can still be served: the adapter caps prompt+output at +# rollout_max_context_len, so a turn needs MAX_GEN_LEN of headroom. +AGENT_PROMPT_CEILING=$((MAX_CONTEXT_LEN - MAX_GEN_LEN)) +echo "Auto-compact: ${AGENT_AUTOCOMPACT_PCT}% of ${AGENT_MAX_CONTEXT_TOKENS} = ${AGENT_AUTOCOMPACT_AT} tokens" \ + "(prompt ceiling ${AGENT_PROMPT_CEILING} = ${MAX_CONTEXT_LEN} - ${MAX_GEN_LEN})" +if (( AGENT_AUTOCOMPACT_AT >= AGENT_PROMPT_CEILING )); then + echo "WARNING: the compaction trigger (${AGENT_AUTOCOMPACT_AT}) is at or above the prompt ceiling" \ + "(${AGENT_PROMPT_CEILING}); turns will hit finish_reason=length before compaction fires." \ + "Lower AGENT_AUTOCOMPACT_PCT to <= $((AGENT_PROMPT_CEILING * 100 / AGENT_MAX_CONTEXT_TOKENS))." +fi +if (( AGENT_MAX_CONTEXT_TOKENS > MAX_CONTEXT_LEN )); then + echo "WARNING: declared context ${AGENT_MAX_CONTEXT_TOKENS} > MAX_CONTEXT_LEN ${MAX_CONTEXT_LEN};" \ + "the adapter hard-stops at the latter." +fi + +# # `investigator` is a read-only sub-agent, dispatched via the Agent tool. +# AGENTS_JSON='{"investigator":{"description":"Searches the repo for relevant files before any edit","prompt":"You are an investigator sub-agent. Use Grep/Read/Glob to find every file relevant to the user task, then return a short bulleted summary. Do NOT edit anything.","tools":["Grep","Read","Glob"]}}' +# export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS} --disable-slash-commands --agents '${AGENTS_JSON}'" +# (WebFetch/WebSearch are already denied by the harness default; EXTRA_ARGS is +# appended last, so anything set here overrides a repeated default flag.) +# The only two harness knobs per agent: extra CLI flags, and extra env vars as +# JSON. Everything else (denied tools, the launch flags) is a class attribute in +# slime_plugins/.../harnesses.py, because it is a property of the harness rather +# than of a run. Both are applied LAST -- EXTRA_ARGS after the harness's own +# flags (claude takes the last occurrence of a repeated flag, verified against +# the real CLI) and EXTRA_ENVS after static_env -- so either can override a +# harness default. The *_EXTRA_ENVS pair is set in the auto-compaction block +# above; these two are the remaining passthroughs, declared so SWE_AGENT can be +# either harness from this script (run_cbc_*.sh just sets it and re-execs). +export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:-}" +export SLIME_AGENT_CBC_EXTRA_ARGS="${SLIME_AGENT_CBC_EXTRA_ARGS:-}" + +# Optional: require dispatching the investigator before any edit, to maximize sub-agent fan-out. +# export SWE_CC_PROMPT="Read PROBLEM_STATEMENT.md. BEFORE editing any file, dispatch the 'investigator' sub-agent (via the Agent tool with subagent_type=investigator) to locate every file relevant to the issue. Then fix the issue and run the tests." + +# ============ proxy bypass for in-cluster/AGS traffic ============ +export no_proxy="127.0.0.1,${MASTER_ADDR},${ADAPTER_PUBLIC_HOST},${E2B_DOMAIN},.tencentags.com" +export NO_PROXY="${no_proxy}" + +cd "${SLIME_DIR}" +source "${SLIME_DIR}/scripts/models/qwen3.5-35B-A3B.sh" + +SAVE_DIR="${SAVE_DIR:-${EXP}/checkpoints}" +SAVE_INTERVAL="${SAVE_INTERVAL:-5}" +mkdir -p "${SAVE_DIR}" + +CKPT_ARGS=( + --hf-checkpoint "${HF_CHECKPOINT}" + --ref-load "${REF_MODEL_PATH}" + --load "${LOAD_DIR:-${SAVE_DIR}}" + --save "${SAVE_DIR}" + --save-interval "${SAVE_INTERVAL}" +) + +ROLLOUT_ARGS=( + --rollout-function-path slime_plugins.rollout_buffer.rollout_buffer_example.generate_rollout + # Used by periodic eval, which runs AGS through slime's standard sglang eval loop. + --custom-generate-function-path slime_plugins.rollout_buffer.generator.ags_generator.generate + --custom-rollout-log-function-path slime_plugins.rollout_buffer.generator.ags_generator.wandb_metrics.log_rollout_data + --custom-eval-rollout-log-function-path slime_plugins.rollout_buffer.generator.ags_generator.wandb_metrics.log_eval_rollout_data + --rollout-task-type ags + --rollout-buffer-url "http://${MASTER_ADDR}:8889" + --prompt-data "${PROMPT_DATA}" + --input-key prompt + --label-key label + --metadata-key metadata + --apply-chat-template + --rollout-shuffle + --num-rollout "${NUM_ROLLOUT}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" + --rollout-buffer-num-epoch 1 + --rollout-buffer-stop-timeout-sec "${ROLLOUT_BUFFER_STOP_TIMEOUT_SEC:-120}" + --rollout-max-context-len "${MAX_CONTEXT_LEN}" + --rollout-max-response-len "${MAX_GEN_LEN}" + --rollout-temperature 1.0 + --rollout-stop-token-ids 248046 248044 + --num-steps-per-rollout "${NUM_STEPS_PER_ROLLOUT}" + --global-batch-size "${GLOBAL_BATCH_SIZE}" + --micro-batch-size "${MICRO_BATCH_SIZE}" + --loss-mask-type qwen3_5 + --save-debug-rollout-data "${RUN_ROOT}/rollout_dumps/rollout_{rollout_id}.pt" +) + +EVAL_ARGS=( + --eval-function-path slime.rollout.sglang_rollout.generate_rollout + --eval-interval "${EVAL_INTERVAL}" + --eval-prompt-data swebench_verified "${EVAL_DATA}" + --n-samples-per-eval-prompt "${N_SAMPLES_PER_EVAL_PROMPT}" + --eval-max-prompt-len "${ROLLOUT_MAX_PROMPT_LEN}" + --eval-max-response-len "${MAX_GEN_LEN}" + --eval-temperature 0.7 + --eval-top-p 0.8 + --eval-top-k 20 +) + +if [[ "${SKIP_EVAL_BEFORE_TRAIN}" == "1" || "${SKIP_EVAL_BEFORE_TRAIN}" == "true" ]]; then + EVAL_ARGS+=(--skip-eval-before-train) +fi + +PERF_ARGS=( + --tensor-model-parallel-size "${TP_SIZE}" + --sequence-parallel + --pipeline-model-parallel-size "${PP_SIZE}" + --context-parallel-size "${CP_SIZE}" + --expert-model-parallel-size "${EP_SIZE}" + --expert-tensor-parallel-size "${ETP_SIZE}" + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + # max-tokens-per-gpu is one CP rank's slice of MAX_CONTEXT_LEN; log-probs are + # chunked along T to avoid OOM on long single trajectories. + --max-tokens-per-gpu $((MAX_CONTEXT_LEN / CP_SIZE)) + --log-probs-chunk-size 1024 + --use-dynamic-batch-size +) + +ALGO_ARGS=( + --advantage-estimator grpo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +SGLANG_ARGS=( + --rollout-num-gpus "${ROLLOUT_NUM_GPUS}" + --rollout-num-gpus-per-engine "${ROLLOUT_TP_SIZE}" + --sglang-mem-fraction-static "${ROLLOUT_MEM_UTILIZATION}" + --sglang-enable-dp-attention + --sglang-dp-size "${ROLLOUT_DP_SIZE}" + --sglang-ep-size "${ROLLOUT_EP_SIZE}" + --sglang-enable-dp-lm-head + --sglang-moe-dense-tp-size 1 + --sglang-tool-call-parser qwen3_coder + --sglang-reasoning-parser qwen3 +) + +if [[ -n "${WANDB_API_KEY:-}" ]]; then + WANDB_ARGS=( + --use-wandb + --wandb-team "${WANDB_ENTITY:?WANDB_ENTITY is required when WandB is enabled}" + --wandb-project "${WANDB_PROJECT:-slime-claude-code-ags}" + --wandb-group "${WANDB_GROUP:-${EXP_TAG}}" + --wandb-key "${WANDB_API_KEY}" + --wandb-dir "${LOG_DIR}" + --disable-wandb-random-suffix + ) +else + WANDB_ARGS=() +fi + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --moe-token-dispatcher-type flex + --moe-enable-deepep + --colocate + --log-passrate +) + +# ============ bring up rollout buffer ============ +python3 -u -m slime_plugins.rollout_buffer.buffer >"${BUFFER_LOG_FILE}" 2>&1 & +BUFFER_PID=$! +trap 'kill ${BUFFER_PID} 2>/dev/null || true' EXIT +sleep 5 + +# ============ bring up ray cluster ============ +HOSTFILE="${HOSTFILE:-/root/mpi_rack_hostfile}" + +ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" \ + --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +if [[ -f "${HOSTFILE}" ]]; then + WORKER_LIMIT=$((ACTOR_NUM_NODES - 1)) + STARTED_WORKERS=0 + for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do + [[ -z "${WORKER_IP}" ]] && continue + [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]] && continue + if (( STARTED_WORKERS >= WORKER_LIMIT )); then + break + fi + echo "Starting Ray worker on ${WORKER_IP}" + ssh -o StrictHostKeyChecking=no "root@${WORKER_IP}" \ + "pkill -9 sglang ; ray stop --force ; pkill -9 python ; \ + ray start --address=${MASTER_ADDR}:${MASTER_PORT} --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} \ + --node-ip-address ${WORKER_IP} --disable-usage-stats" & + STARTED_WORKERS=$((STARTED_WORKERS + 1)) + done + for pid in $(jobs -pr); do + [[ "${pid}" == "${BUFFER_PID}" ]] && continue + wait "${pid}" + done + if (( STARTED_WORKERS < WORKER_LIMIT )); then + echo "WARNING: requested ${ACTOR_NUM_NODES} nodes but only started $((STARTED_WORKERS + 1)) including head." + fi +else + echo "WARNING: HOSTFILE=${HOSTFILE} not found; only the head node was started." +fi + +echo "Waiting for Ray cluster to stabilize..." +sleep 30 +ray status + +# ============ runtime env propagated to ray workers ============ +export SLIME_DIR +RUNTIME_ENV_JSON=$(python3 - <&1 | tee "${LOG_FILE}" + +echo "RUN_ROOT=${RUN_ROOT}" diff --git a/requirements.txt b/requirements.txt index 95d0a842d5..6558e8abc5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -22,5 +22,6 @@ sglang-router>=0.2.3 tensorboard transformers wandb +weave xxhash # disk delta weight sync (checksum + codec) zstandard diff --git a/slime/agent/adapters/common.py b/slime/agent/adapters/common.py index 6f01d4ed9d..1964118983 100644 --- a/slime/agent/adapters/common.py +++ b/slime/agent/adapters/common.py @@ -172,6 +172,9 @@ def __init__( self.app.router.add_get("/healthz", _health) self.app.router.add_get("/v1/models", _health) + self.app.router.add_post("/_slime/open_session", self._control_open_session) + self.app.router.add_post("/_slime/finish_session", self._control_finish_session) + self.app.router.add_post("/_slime/drop_session", self._control_drop_session) self._register_routes(self.app) # -- wire hooks (subclass overrides) ------------------------------------- @@ -222,6 +225,34 @@ def open_session( max_context_tokens=int(max_context_tokens or 0), ) + async def _control_open_session(self, request: web.Request) -> web.Response: + body = await request.json() + sid = body["sid"] + self.open_session( + sid, + sampling_defaults=body.get("sampling_defaults") or {}, + max_context_tokens=int(body.get("max_context_tokens") or 0), + ) + return web.json_response({"ok": True}) + + async def _control_finish_session(self, request: web.Request) -> web.Response: + from slime.utils.types import Sample + + body = await request.json() + samples = await self.finish_session( + body["sid"], + base_sample=Sample.from_dict(body["base_sample"]), + reward=float(body.get("reward", 0.0)), + extra_metadata=body.get("extra_metadata") or {}, + wait_timeout=float(body.get("wait_timeout", 5.0)), + ) + return web.json_response({"ok": True, "samples": [sample.to_dict() for sample in samples]}) + + async def _control_drop_session(self, request: web.Request) -> web.Response: + body = await request.json() + await self.drop_session(body["sid"], wait_timeout=float(body.get("wait_timeout", 5.0))) + return web.json_response({"ok": True}) + async def shutdown_session(self, sid: str, *, wait_timeout: float = 5.0) -> None: """Mark a sid closed and drain its in-flight turn tasks.""" self.closed.add(sid) @@ -427,9 +458,17 @@ def _sampling_params(session: Any, body: dict, *, max_token_keys: tuple[str, ... sp["max_new_tokens"] = min(int(sp.get("max_new_tokens", body[key])), int(body[key])) break - for src_k, dst_k in (("temperature", "temperature"), ("top_p", "top_p"), ("top_k", "top_k")): - if src_k in body: - sp[dst_k] = body[src_k] + # The CLI's own sampling knobs only apply where the caller left that key + # unset. open_session's sampling_defaults carry the trainer's temperature / + # top_p / top_k (eval uses different values than training), and a harness + # that hardcodes its own would otherwise silently override them: codebuddy + # sends temperature=1 on every /v1/chat/completions request, so an + # --eval-temperature would never reach sglang. Claude Code sends none of + # these three, which is why this only ever bit the codebuddy path. + defaults = session.sampling_defaults or {} + for key in ("temperature", "top_p", "top_k"): + if key in body and key not in defaults: + sp[key] = body[key] for key in stop_keys: if body.get(key): diff --git a/slime/agent/adapters/openai.py b/slime/agent/adapters/openai.py index ad3d2e4d87..90440257a9 100644 --- a/slime/agent/adapters/openai.py +++ b/slime/agent/adapters/openai.py @@ -99,12 +99,28 @@ def _arguments_as_dict(arguments: Any) -> dict[str, Any]: return {"_raw_arguments": str(arguments)} +# Keys an OpenAI-compatible client may echo reasoning back under. The spec never +# standardised one, so CodeBuddy Code replays ours verbatim as "reasoning" while +# we send "reasoning_content"; reading only the canonical key would make every +# reasoning turn's echo compare unequal to our leaf. +_REASONING_KEYS = ("reasoning_content", "reasoning") + + +def _reasoning_text(msg: dict) -> str: + """Reasoning text from an echoed assistant message, under any known key.""" + for key in _REASONING_KEYS: + value = msg.get(key) + if isinstance(value, str) and value: + return value + return "" + + def _translate_messages(messages: list[dict]) -> list[dict]: """OpenAI chat messages -> tokenizer chat-template messages. Mirrors anthropic._translate_messages so a replayed assistant turn compares equal (dict equality) to the leaf the manager appended on the previous - request. Two invariants must hold: + request. Three invariants must hold: * tool_calls[i].function.arguments is a dict (not a JSON string): the chat template needs a mapping, and the manager matches history by dict @@ -112,6 +128,8 @@ def _translate_messages(messages: list[dict]) -> list[dict]: * Wire-only correlation ids are dropped (tool_call_id on tool messages, tool_calls[i].id on echoed assistant messages). Fresh ids are minted on each response, so keeping the wire ids would diverge the replay match. + * Reasoning is normalised onto reasoning_content whichever key it arrived + under (see _REASONING_KEYS). """ translated: list[dict] = [] for msg in messages: @@ -132,7 +150,7 @@ def _translate_messages(messages: list[dict]) -> list[dict]: "role": "assistant", "content": flatten_content(content), } - reasoning = msg.get("reasoning_content") + reasoning = _reasoning_text(msg) if reasoning: assistant["reasoning_content"] = reasoning tool_calls = msg.get("tool_calls") or [] @@ -249,24 +267,29 @@ def _build_reply_parts(parsed: ParsedModelOutput, finish: str) -> tuple[dict[str wire_message: dict[str, Any] = { "role": "assistant", - # send content=null when there are tool_calls: some OpenAI clients split - # a mixed text+tool_calls turn into two echoed messages otherwise, which - # diverges the history match against our leaf - "content": None if wire_tool_calls else (parsed.text or None), + "content": parsed.text or None, } - # manager_message must match what the client echoes on the next request, or - # the manager's history match (dict equality) diverges and every turn forks. - # Differences from wire_message, each needed to match the echo: - # * no reasoning_content -- some clients strip it on echo (the reasoning - # token ids are still kept in the trained tokens, only the text drops) + # manager_message is what re-renders as history on the next turn, so it must + # carry everything the model actually generated -- text and reasoning + # included, even alongside tool_calls. Dropping a field here does not just + # lose it from the history: the re-render then no longer reproduces the ids we + # sampled, so _SampleBuilder sees token drift and either REALIGNs (rewriting + # a real trained response as loss_mask=0) or forks. Measured on a 500-instance + # CodeBuddy run, dropping reasoning_content alone split each rollout into ~25 + # samples where Claude Code produced 1. + # + # It must also equal what the client echoes back, or the manager's history + # match (dict equality) diverges instead. The two remaining differences from + # wire_message are both required for that: + # * "" rather than None for empty content -- flatten_content's shape # * only the first tool_call -- some clients drop extra parallel tool_calls - # * empty content when tool_calls are present -- mirrors content=null above manager_message: dict[str, Any] = { "role": "assistant", - "content": "" if wire_tool_calls else (parsed.text or ""), + "content": parsed.text or "", } if parsed.reasoning: wire_message["reasoning_content"] = parsed.reasoning + manager_message["reasoning_content"] = parsed.reasoning if wire_tool_calls: wire_message["tool_calls"] = wire_tool_calls[:1] manager_message["tool_calls"] = manager_tool_calls[:1] diff --git a/slime/agent/sandbox.py b/slime/agent/sandbox.py index b81fff6e0c..ffb0ef2719 100644 --- a/slime/agent/sandbox.py +++ b/slime/agent/sandbox.py @@ -13,6 +13,7 @@ import logging import os import random +import shlex import time from pathlib import Path from typing import Protocol, runtime_checkable @@ -60,6 +61,53 @@ async def read_file(self, sandbox_path: str, *, user: str = "root") -> str: ... EXIT_TIME_BUDGET_EXCEEDED = -1 +async def terminate_process_group( + sb: Sandbox, + *, + pid_file: str, + user: str, + grace_sec: int = 5, +) -> None: + """Terminate a detached ``setsid`` process group and wait until it exits.""" + + pid_path = shlex.quote(pid_file) + missing_pid_message = shlex.quote(f"missing process-group pid file: {pid_file}") + invalid_pid_prefix = shlex.quote(f"invalid process-group id in {pid_file}:") + command = f""" +set -u +if [ ! -s {pid_path} ]; then + echo {missing_pid_message} >&2 + exit 1 +fi +pgid=$(cat {pid_path}) +case "$pgid" in + ""|*[!0-9]*) + echo {invalid_pid_prefix} "$pgid" >&2 + exit 1 + ;; +esac +kill -TERM -- "-$pgid" 2>/dev/null || true +remaining={grace_sec} +while [ "$remaining" -gt 0 ]; do + if ! kill -0 -- "-$pgid" 2>/dev/null; then + exit 0 + fi + sleep 1 + remaining=$((remaining - 1)) +done +kill -KILL -- "-$pgid" 2>/dev/null || true +for _ in 1 2 3 4 5; do + if ! kill -0 -- "-$pgid" 2>/dev/null; then + exit 0 + fi + sleep 0.2 +done +echo "process group $pgid is still alive after SIGKILL" >&2 +exit 1 +""".strip() + await sb.exec(command, user=user, timeout=grace_sec + 10, check=True, idempotent=False) + + async def _await_done_marker(sb: Sandbox, done_file: str, *, user: str, time_budget_sec: int) -> int: """Poll a detached command's exit-code marker until it appears, returning the exit code (or ``EXIT_TIME_BUDGET_EXCEEDED`` if the budget runs out first). @@ -106,6 +154,7 @@ async def exec_and_wait( done_file = f"/tmp/.{tag}.done" launcher = f"/tmp/.{tag}.sh" lock_dir = f"/tmp/.{tag}.spawned" + pid_file = f"/tmp/.{tag}.pid" prefix = f"cd {workdir}\nexport HOME=/home/{user}\n" if workdir else "" launcher_body = f"#!/bin/bash\n{prefix}{cmd}\necho $? > {done_file}\n" await sb.write_file(launcher, launcher_body, user=user) @@ -113,8 +162,9 @@ async def exec_and_wait( await sb.exec( f"chmod +x {launcher}; " f"mkdir {lock_dir} 2>/dev/null || exit 0; " - f"rm -f {out_file} {done_file}; " - f"setsid bash {launcher} < /dev/null > {out_file} 2>&1 &", + f"rm -f {out_file} {done_file} {pid_file}; " + f"setsid bash {launcher} < /dev/null > {out_file} 2>&1 & " + f"echo $! > {pid_file}", user=user, env=env, timeout=30, @@ -122,6 +172,9 @@ async def exec_and_wait( idempotent=True, ) exit_code = await _await_done_marker(sb, done_file, user=user, time_budget_sec=time_budget_sec) + if exit_code == EXIT_TIME_BUDGET_EXCEEDED: + logger.warning("Detached command %s exceeded %ss; terminating process group", tag, time_budget_sec) + await terminate_process_group(sb, pid_file=pid_file, user=user) if exit_code == 0 and not want_output: return exit_code, "" if want_output: diff --git a/slime/backends/megatron_utils/data.py b/slime/backends/megatron_utils/data.py index 51b008d111..859076b05e 100644 --- a/slime/backends/megatron_utils/data.py +++ b/slime/backends/megatron_utils/data.py @@ -11,7 +11,7 @@ from slime.utils import train_metric_utils from slime.utils.flops_utils import calculate_fwd_flops -from slime.utils.metric_utils import compute_pass_rate, compute_rollout_step +from slime.utils.metric_utils import compute_grouped_pass_rate, compute_pass_rate, compute_rollout_step from slime.utils.types import RolloutBatch from ...utils import logging_utils @@ -283,6 +283,8 @@ def log_rollout_data( "loss_masks", "sample_indices", "rollout_ids", + "raw_reward_group_indices", + "raw_reward_rollout_ids", "rollout_mask_sums", "rollout_top_p_token_ids", "rollout_top_p_token_offsets", @@ -481,12 +483,20 @@ def log_passrate(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) - """ if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): log_dict = {} - for key, val in rollout_data.items(): - if key != "raw_reward": - continue - + raw_rewards = rollout_data.get("raw_reward") + if raw_rewards is None: + return + + if "raw_reward_group_indices" in rollout_data and "raw_reward_rollout_ids" in rollout_data: + log_dict |= compute_grouped_pass_rate( + flat_rewards=raw_rewards, + group_indices=rollout_data["raw_reward_group_indices"], + rollout_ids=rollout_data["raw_reward_rollout_ids"], + group_size=args.n_samples_per_prompt, + ) + else: log_dict |= compute_pass_rate( - flat_rewards=val, + flat_rewards=raw_rewards, group_size=args.n_samples_per_prompt, num_groups=args.rollout_batch_size, ) diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py index f4636d3b24..3b49ee0f79 100644 --- a/slime/backends/sglang_utils/sglang_engine.py +++ b/slime/backends/sglang_utils/sglang_engine.py @@ -10,6 +10,7 @@ import requests import sglang_router from packaging.version import parse +from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.server_args import ServerArgs from sglang.srt.utils import kill_process_tree from urllib3.exceptions import NewConnectionError @@ -99,6 +100,12 @@ def _wait_server_healthy(base_url, api_key, is_process_alive): time.sleep(2) +def _resume_restores_generation(tags: list[str] | None) -> bool: + if tags is None: + return True + return GPU_MEMORY_TYPE_KV_CACHE in tags or GPU_MEMORY_TYPE_CUDA_GRAPH in tags + + class SGLangEngine(RayActor): def __init__( self, @@ -115,6 +122,7 @@ def __init__( self.base_gpu_id = base_gpu_id self.sglang_overrides = sglang_overrides or {} self.num_gpus_per_engine = num_gpus_per_engine + self._generation_paused_by_offload = False def init( self, @@ -378,6 +386,8 @@ def set_weight_version(self, new_version: str): return self._make_request("update_weight_version", {"new_version": str(new_version)}) def release_memory_occupation(self): + self.pause_generation() + self._generation_paused_by_offload = True self.flush_cache() return self._make_request("release_memory_occupation") @@ -385,10 +395,14 @@ def resume_memory_occupation(self, tags: list[str] = None): """ Available tags for multi-stage resume: weights, kv_cache """ - return self._make_request( + result = self._make_request( "resume_memory_occupation", {"tags": tags}, ) + if self._generation_paused_by_offload and _resume_restores_generation(tags): + self.continue_generation() + self._generation_paused_by_offload = False + return result def check_weights(self, action: str): return self._make_request("weights_checker", {"action": action}) diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index be8f42438c..d1291373fc 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -566,8 +566,14 @@ def load(self, rollout_id=None): def offload(self): self.health_monitoring_pause() + handles = [] for srv in self.servers.values(): - srv.offload() + for group in srv.server_groups: + handles.extend(group.offload()) + if handles: + logger.info("Waiting for %d rollout engine offload requests to finish.", len(handles)) + return ray.get(handles) + return [] def onload(self, tags: list[str] | None = None): for srv in self.servers.values(): @@ -721,6 +727,11 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl "truncated": [1 if sample.status == Sample.Status.TRUNCATED else 0 for sample in samples], "sample_indices": [sample.index for sample in samples], "rollout_ids": rollout_ids, + "raw_reward_group_indices": [ + sample.group_index if sample.group_index is not None else i // self.args.n_samples_per_prompt + for i, sample in enumerate(samples) + ], + "raw_reward_rollout_ids": rollout_ids, } # loss mask @@ -860,7 +871,7 @@ def _split_train_data_by_dp(self, data): continue rollout_data[key] = [data[key][j] for j in partition] # keys that need to be splited at train side - for key in ["raw_reward", "total_lengths"]: + for key in ["raw_reward", "raw_reward_group_indices", "raw_reward_rollout_ids", "total_lengths"]: if key not in data: continue rollout_data[key] = data[key] diff --git a/slime/rollout/data_source.py b/slime/rollout/data_source.py index ca7171ecae..536d054afc 100644 --- a/slime/rollout/data_source.py +++ b/slime/rollout/data_source.py @@ -204,10 +204,12 @@ def add_samples(self, samples: list[list[Sample]]): assert isinstance(samples, list), f"samples must be a list, got {type(samples)}" assert isinstance(samples[0], list), f"the elements of samples must be list, got {type(samples[0])}" for i in range(0, len(samples)): - assert ( - len(samples[i]) == self.args.n_samples_per_prompt - ), f"the length of the elements of samples must be equal to n_samples_per_prompt, got {len(samples[i])} != {self.args.n_samples_per_prompt}" + assert len(samples[i]) >= self.args.n_samples_per_prompt, ( + "the elements of samples must include at least n_samples_per_prompt entries, " + f"got {len(samples[i])} < {self.args.n_samples_per_prompt}" + ) group = samples[i] # type: ignore + assert all(isinstance(sample, Sample) for sample in group), "sample groups must contain Sample instances" self.buffer.append(group) # TODO remove diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index ccb121f35d..6bf5c0ff25 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -1397,6 +1397,25 @@ def add_rollout_buffer_arguments(parser): default=-1, help="Number of times to retry fetching trajectory, -1 means unlimited retry", ) + parser.add_argument( + "--rollout-buffer-num-epoch", + type=int, + default=1, + help=( + "Generator-side epochs to run for each rollout-buffer start request. " + "Keep this at 1 for colocated/offloaded rollout so background generation " + "does not continue after one training rollout has returned." + ), + ) + parser.add_argument( + "--rollout-buffer-stop-timeout-sec", + type=float, + default=120.0, + help=( + "Seconds to wait for the rollout-buffer background job to stop after " + "enough data has been collected for the current training rollout." + ), + ) parser.add_argument( "--min-batch-collection-ratio", type=float, @@ -1408,6 +1427,12 @@ def add_rollout_buffer_arguments(parser): type=str, default="math", ) + parser.add_argument( + "--enable-token2text", + action="store_true", + default=False, + help="Decode token IDs to readable text in rollout traces. Disabled by default.", + ) parser.add_argument( "--loss-mask-type", type=str, diff --git a/slime/utils/data.py b/slime/utils/data.py index 102b5ef44e..58dfec6459 100644 --- a/slime/utils/data.py +++ b/slime/utils/data.py @@ -244,12 +244,21 @@ def __init__( else: output_prompt = prompt - if processor: + # Vision extraction is gated on multimodal_keys, not on the mere + # existence of a processor. AutoProcessor returns one for any + # VL-capable checkpoint -- Qwen3.5-35B-A3B yields a Qwen3VLProcessor + # -- so keying off `processor` alone forced a conversation-shaped + # prompt on text-only datasets. That made --apply-chat-template + # mandatory, which in turn rewrote the prompt into a templated + # "<|im_start|>user ..." string; agent rollouts that hand + # Sample.prompt straight to a CLI then fed it those literal markers. + # This mirrors the `as_conversation` condition computed above. + if processor and multimodal_keys is not None: from slime.utils.processing_utils import process_vision_info assert isinstance( prompt, list - ), f"prompt must be a list when processor is not None, got {type(prompt)} instead" + ), f"prompt must be a list when multimodal_keys is set, got {type(prompt)} instead" multimodal_inputs = process_vision_info(prompt, processor) else: multimodal_inputs = None diff --git a/slime/utils/dp_schedule.py b/slime/utils/dp_schedule.py index 1735fad3b4..37f8bf3531 100644 --- a/slime/utils/dp_schedule.py +++ b/slime/utils/dp_schedule.py @@ -17,7 +17,12 @@ single first-fit pass (dynamic batch) or fixed-size chunking (static batch). 3. Adjust ``K`` to a multiple of ``dp_size * (mb_group if vpp>1 else 1)`` - by splitting the largest multi-sample bins (dynamic only). + by splitting the largest multi-sample bins (dynamic only). When no bin + can be split -- every bin holds one sample, the usual case once a + single sample fills ``max_per_bin`` -- round ``K`` *down* instead, + repacking so the largest bin stays as small as possible. Rounding down + needs at least ``align_to`` bins to land on; below that the step cannot + be aligned at all and :func:`build_dp_schedule` raises. 4. Distribute the ``K`` mbs across ``dp_size`` ranks, ``K / dp_size`` each, with either a strided round-robin or a Karmarkar-Karp pass on estimated mbs FLOPs. @@ -26,15 +31,21 @@ - every DP rank runs the **same** ``num_microbatches`` per training step (required for PP sync); - every mbs (dynamic path without ``balance_by_flops``) holds - ``<= max_tokens_per_gpu * cp_size`` tokens, with one exception — an - individual sample larger than that cap lands alone in its own mbs (and - that mbs is the only one allowed to exceed the cap); + ``<= max_tokens_per_gpu * cp_size`` tokens, with two exceptions — an + individual sample larger than that cap lands alone in its own mbs, and + a step whose bin count had to be repacked down for alignment (step 3 + above) may hold mbs above the cap. On that path ``max_tokens_per_gpu`` + is a target, not a hard bound, and the repack is logged at WARNING; - the union of per-rank sample indices equals the set of samples kept after trimming trailing rollouts (every kept sample placed exactly once); - flattening ``micro_batch_indices`` for a rank yields ``range(num_samples_rank)`` (each rank's samples are tiled exactly once by its mbs schedule). + +Note that an mbs may hold samples from different rollouts after a repack. +That is harmless: the loss is aggregated by ``rollout_id`` across the whole +step, not per micro-batch, so mbs composition only affects memory. """ from __future__ import annotations @@ -43,7 +54,12 @@ from typing import Any from slime.utils.flops_utils import calculate_fwd_flops -from slime.utils.seqlen_balancing import expand_bins_by_splitting, first_fit_pack, get_seqlen_balanced_partitions +from slime.utils.seqlen_balancing import ( + expand_bins_by_splitting, + first_fit_pack, + get_seqlen_balanced_partitions, + shrink_bins_by_merging, +) logger = logging.getLogger(__name__) @@ -169,11 +185,63 @@ def build_dp_schedule( if target_K != len(step_mbs): if args.use_dynamic_batch_size: expand_bins_by_splitting(step_mbs, target_K, step_lengths) - assert len(step_mbs) == target_K, ( - f"dynamic path: could only produce {len(step_mbs)} mbs after maximal splitting; " - f"need {target_K}. step {step_i} has {len(sample_indices)} samples, below the " - f"alignment threshold ({align_to})." - ) + if len(step_mbs) != target_K: + # Splitting up to ``target_K`` was impossible: every bin holds a + # single sample, so there is nothing left to divide. This is the + # normal state for long-context runs, where one sample alone fills + # ``max_per_bin`` (e.g. 69k-token SWE trajectories against a 96k + # cap) -- first-fit then emits one bin per sample and an odd sample + # count can never reach an even ``target_K``. + # + # Rounding DOWN to the nearest multiple is the fallback, but it only + # exists when there are at least ``align_to`` bins to round down to. + # Below that the floor is 0, which is not a schedule -- the step + # genuinely cannot be aligned and the run must stop. Report that + # directly rather than letting the merge below fail with a message + # that blames merging for a shortage of samples. + if len(step_mbs) < align_to: + raise AssertionError( + f"dynamic path: step {step_i} produced {len(step_mbs)} unsplittable " + f"single-sample mbs, below the alignment threshold ({align_to}); " + f"cannot align by splitting (all bins are singletons) or by merging " + f"down (the floor is 0 mbs). step has {len(sample_indices)} samples, " + f"dp_size={dp_size}, mb_group={mb_group if vpp_size > 1 else 1}. " + f"Raise global_batch_size so each step holds >= {align_to} samples, " + f"or lower dp_size / disable VPP to reduce the alignment requirement." + ) + # Round DOWN to the nearest multiple, repacking the bins so the + # LARGEST one stays as small as possible -- that bin sets peak + # activation memory. Rounding down keeps every sample in the step; + # the alternative, dropping the odd sample, would silently discard + # completed rollout work. + # + # The cost is real: merged bins can exceed ``max_per_bin``, because + # first-fit already produced a maximal packing, so any reduction in + # bin count must combine samples that did not fit together. We only + # shrink by the minimum needed for alignment. Callers must therefore + # treat max_tokens_per_gpu as a target rather than a hard guarantee + # on this path. + aligned_K = (len(step_mbs) // align_to) * align_to + logger.warning( + "step %d: %d samples packed into %d unsplittable single-sample mbs, which is not " + "a multiple of %d; repacking down to %d mbs. Some mbs will exceed " + "max_tokens_per_gpu * cp_size (%s), which lowering max_tokens_per_gpu cannot " + "prevent -- the merged mbs holds whole samples, so its token count is set by the " + "data. If this OOMs, lower rollout_max_context_len (shorter samples) or choose a " + "global_batch_size whose per-step sample count is a multiple of %d.", + step_i, + len(sample_indices), + len(step_mbs), + align_to, + aligned_K, + max_per_bin, + align_to, + ) + shrink_bins_by_merging(step_mbs, aligned_K, step_lengths) + assert len(step_mbs) == aligned_K, ( + f"dynamic path: repacking produced {len(step_mbs)} mbs, expected {aligned_K}. " + f"step {step_i} has {len(sample_indices)} samples." + ) else: raise AssertionError( f"static path: num_mbs ({len(step_mbs)}) is not a multiple of " diff --git a/slime/utils/metric_utils.py b/slime/utils/metric_utils.py index 46e42d73b0..40f894819e 100644 --- a/slime/utils/metric_utils.py +++ b/slime/utils/metric_utils.py @@ -24,7 +24,14 @@ def compute_pass_rate( pass_rate_name_list = [2**i for i in range(int(math.log2(group_size)) + 1)] - assert len(flat_rewards) == num_groups * group_size, f"{len(flat_rewards)=} {num_groups=} {group_size=}" + if len(flat_rewards) != num_groups * group_size: + logger.warning( + "skip fixed-shape passrate: len(flat_rewards)=%d num_groups=%d group_size=%d", + len(flat_rewards), + num_groups, + group_size, + ) + return {} rewards_of_group = np.array(flat_rewards).reshape(num_groups, group_size) log_dict = {} @@ -40,6 +47,68 @@ def compute_pass_rate( return log_dict +def compute_grouped_pass_rate( + flat_rewards: list[float], + group_indices: list[int | str], + rollout_ids: list[int | str], + group_size: int, +): + """Compute pass@k on prompt groups while tolerating fan-out samples. + + Agentic rollouts can emit multiple train samples for one rollout attempt + (for example one sample per root-to-leaf chain in a tool-use tree). Those + fan-out siblings are training segments, not independent pass@k samples. So + this metric first deduplicates by ``(group_index, rollout_id)`` and then + computes pass@k from the attempt-level rewards in each prompt group. + """ + + if group_size == 1: + return {} + + if not (len(flat_rewards) == len(group_indices) == len(rollout_ids)): + logger.warning( + "skip grouped passrate: rewards=%d group_indices=%d rollout_ids=%d", + len(flat_rewards), + len(group_indices), + len(rollout_ids), + ) + return {} + + grouped_attempt_rewards: dict[int | str, dict[int | str, float]] = {} + for position, (reward, group_index, rollout_id) in enumerate( + zip(flat_rewards, group_indices, rollout_ids, strict=True) + ): + # Missing rollout ids are not expected on the normal train path because + # rollout.py fills them in before packaging. Keep this fallback local so + # a malformed custom path does not merge unrelated attempts. + attempt_key = rollout_id if rollout_id is not None else f"position:{position}" + attempts = grouped_attempt_rewards.setdefault(group_index, {}) + reward = float(reward) + attempts[attempt_key] = max(attempts.get(attempt_key, reward), reward) + + if not grouped_attempt_rewards: + return {} + + pass_rate_name_list = [2**i for i in range(int(math.log2(group_size)) + 1)] + log_dict = {} + for k in pass_rate_name_list: + num_samples = [] + num_correct = [] + for attempts in grouped_attempt_rewards.values(): + rewards = list(attempts.values()) + if len(rewards) < k: + continue + num_samples.append(len(rewards)) + num_correct.append(sum(1 for reward in rewards if reward == 1)) + if not num_samples: + continue + + pass_k_estimates = _estimate_pass_at_k(np.array(num_samples), np.array(num_correct), k) + log_dict[f"pass@{k}"] = np.mean(pass_k_estimates).item() + + return log_dict + + def _estimate_pass_at_k(num_samples, num_correct, k): """ Estimates pass@k of each problem and returns them in an array. diff --git a/slime/utils/seqlen_balancing.py b/slime/utils/seqlen_balancing.py index 5736d8850e..dd6ab54ddd 100644 --- a/slime/utils/seqlen_balancing.py +++ b/slime/utils/seqlen_balancing.py @@ -229,6 +229,48 @@ def expand_bins_by_splitting(bins: list[list[int]], target_count: int, lengths) bins.append(right) +def shrink_bins_by_merging(bins: list[list[int]], target_count: int, lengths) -> None: + """Shrink ``bins`` in place down to exactly ``target_count`` bins. + + The counterpart to :func:`expand_bins_by_splitting`, for when a bin packing has + too *many* bins to align and cannot be split any further (every bin is a + singleton). + + Minimises the **largest** resulting bin, because that bin is what sets peak + activation memory. Uses longest-processing-time-first: walk the bins from largest + to smallest and drop each into the currently lightest target slot. + + Repeatedly merging the two smallest bins is the intuitive alternative and is + *worse* whenever more than one merge is needed: it keeps stacking samples onto the + same growing bin. With ``lengths = [49000, 49000, 90000, 90000, 90000]`` merged to + 3 bins, that greedy rule peaks at 180000 tokens where this one peaks at 139000 — + a 41000-token difference in peak memory, i.e. 1.88x vs 1.45x a 96000 cap. The two + rules agree when exactly one merge is needed (``target_count == len(bins) - 1``), + which is the common ``dp_size == 2`` case. + + Unlike splitting, merging **can push a bin past the token cap** — that is + unavoidable, because first-fit already produced a maximal packing, so any + reduction in bin count must combine samples that did not fit together. Callers + that care about the cap must budget for the returned bins exceeding it, and + should only shrink by the minimum needed to satisfy an alignment constraint. + + ``target_count`` must be >= 1 and <= ``len(bins)``; shrinking below 1 has no + meaning and growing is :func:`expand_bins_by_splitting`'s job. + """ + assert target_count >= 1, f"target_count {target_count} must be >= 1" + if len(bins) <= target_count: + return + + slots: list[list[int]] = [[] for _ in range(target_count)] + slot_sums = [0] * target_count + for bin_ in sorted(bins, key=lambda b: -sum(lengths[i] for i in b)): + lightest = min(range(target_count), key=lambda i: slot_sums[i]) + slots[lightest].extend(bin_) + slot_sums[lightest] += sum(lengths[i] for i in bin_) + + bins[:] = slots + + def get_reverse_idx(idx_map): reverse_idx_map = copy.deepcopy(idx_map) diff --git a/slime_plugins/rollout_buffer/buffer.py b/slime_plugins/rollout_buffer/buffer.py index db4e561124..487859852c 100644 --- a/slime_plugins/rollout_buffer/buffer.py +++ b/slime_plugins/rollout_buffer/buffer.py @@ -2,13 +2,16 @@ import glob import importlib.util import json +import os import pathlib import threading import time +import traceback +import uuid from typing import Any import uvicorn -from fastapi import BackgroundTasks, FastAPI, HTTPException, Request +from fastapi import FastAPI, HTTPException, Request from pydantic import BaseModel app = FastAPI(title="Rollout Buffer Server", debug=True) @@ -221,6 +224,7 @@ def __init__( transform_group_func=None, is_valid_group_func=None, get_group_data_meta_info_func=None, + rollout_job_id: str | None = None, ): self.buffer = BufferQueue( group_size=group_size, @@ -234,9 +238,14 @@ def __init__( self.total_written = 0 self.total_read = 0 self.task_type = task_type + self.rollout_job_id = rollout_job_id def write(self, data): with self.lock: + item_job_id = data.get("rollout_job_id") + if self.rollout_job_id and item_job_id is not None and item_job_id != self.rollout_job_id: + print(f"Ignore stale rollout item from job {item_job_id}; " f"current job is {self.rollout_job_id}") + return None self.buffer.append(data) self.total_written += 1 self.not_empty.notify_all() @@ -256,11 +265,87 @@ def read(self): buffer = RolloutBuffer() +class RolloutJob: + def __init__(self, payload: dict[str, Any]): + self.job_id = str(uuid.uuid4()) + self.payload = dict(payload) + self.stop_event = threading.Event() + self.started_at = time.time() + self.finished_at: float | None = None + self.status = "starting" + self.error: str | None = None + self.thread = threading.Thread( + target=self._run, + name=f"rollout-buffer-{self.job_id[:8]}", + daemon=True, + ) + + def start(self) -> None: + self.thread.start() + + def stop(self) -> None: + self.stop_event.set() + + def is_alive(self) -> bool: + return self.thread.is_alive() + + def join(self, timeout: float | None = None) -> bool: + self.thread.join(timeout) + return not self.thread.is_alive() + + def snapshot(self) -> dict[str, Any]: + return { + "job_id": self.job_id, + "status": self.status, + "alive": self.thread.is_alive(), + "stop_requested": self.stop_event.is_set(), + "started_at": self.started_at, + "finished_at": self.finished_at, + "error": self.error, + } + + def _run(self) -> None: + self.status = "running" + try: + run_rollout(self.payload, stop_event=self.stop_event, rollout_job_id=self.job_id) + self.status = "stopped" if self.stop_event.is_set() else "finished" + except Exception as exc: + self.status = "failed" + self.error = f"{type(exc).__name__}: {exc}" + print(f"Rollout job {self.job_id} failed: {self.error}") + traceback.print_exc() + finally: + self.finished_at = time.time() + + +current_rollout_job: RolloutJob | None = None +rollout_job_lock = threading.RLock() + + +def _stop_current_rollout_job(wait: bool, timeout_sec: float | None) -> dict[str, Any]: + with rollout_job_lock: + job = current_rollout_job + if job is None: + return {"had_job": False, "stopped": True, "job": None} + + job.stop() + stopped = True + if wait and job.is_alive(): + stopped = job.join(timeout_sec) + return {"had_job": True, "stopped": stopped, "job": job.snapshot()} + + @app.post("/buffer/write", response_model=BufferResponse) async def write_to_buffer(request: Request): try: data = await request.json() item = buffer.write(data) + if item is None: + return BufferResponse( + success=False, + message="Ignored stale rollout item", + data={"data": [], "meta_info": "stale rollout job"}, + ) return BufferResponse( success=True, message="Data has been successfully written to buffer", @@ -295,7 +380,11 @@ async def get_rollout_data(request: Request): ) -def run_rollout(data: dict): +def run_rollout( + data: dict, + stop_event: threading.Event | None = None, + rollout_job_id: str | None = None, +): global buffer # Auto-discover generators generator_map = discover_generators() @@ -315,25 +404,65 @@ def run_rollout(data: dict): transform_group_func=generator_info.get("transform_group", None), is_valid_group_func=generator_info.get("is_valid_group"), get_group_data_meta_info_func=generator_info.get("get_group_data_meta_info"), + rollout_job_id=rollout_job_id, ) # Call the run_rollout function from the appropriate generator module - generator_info["run_rollout"](data) + generator_payload = dict(data) + if stop_event is not None: + generator_payload["_stop_event"] = stop_event + if rollout_job_id is not None: + generator_payload["_rollout_job_id"] = rollout_job_id + generator_info["run_rollout"](generator_payload) print(f"Rollout completed successfully for task_type: {task_type}") @app.post("/start_rollout") -async def start_rollout(request: Request, background: BackgroundTasks): +async def start_rollout(request: Request): + global current_rollout_job + payload = await request.json() + + stop_timeout_sec = float(payload.get("stop_previous_timeout_sec", 60)) + stopped = _stop_current_rollout_job(wait=True, timeout_sec=stop_timeout_sec) + if stopped["had_job"] and not stopped["stopped"]: + raise HTTPException( + status_code=409, + detail={ + "message": "Previous rollout is still running after stop request; refusing to start a new one.", + "previous": stopped["job"], + }, + ) + + job = RolloutJob(payload) + with rollout_job_lock: + current_rollout_job = job + job.start() + return {"message": "Rollout started", "rollout_job_id": job.job_id, "previous": stopped} + + +@app.post("/stop_rollout") +async def stop_rollout(request: Request): payload = await request.json() - background.add_task(run_rollout, payload) - return {"message": "Rollout started"} + wait = bool(payload.get("wait", True)) + timeout_sec = payload.get("timeout_sec", 60) + timeout_sec = None if timeout_sec is None else float(timeout_sec) + return _stop_current_rollout_job(wait=wait, timeout_sec=timeout_sec) + + +@app.get("/rollout_status") +async def rollout_status(): + with rollout_job_lock: + job = current_rollout_job + if job is None: + return {"job": None} + return {"job": job.snapshot()} if __name__ == "__main__": uvicorn.run( app, host="0.0.0.0", - port=8889, + port=int(os.environ.get("ROLLOUT_BUFFER_PORT", "8889")), limit_concurrency=1000, # Connection concurrency limit # limit_max_requests=1000000, # Maximum request limit timeout_keep_alive=5, # Keep-alive timeout, diff --git a/slime_plugins/rollout_buffer/generator/ags_generator.py b/slime_plugins/rollout_buffer/generator/ags_generator.py new file mode 100644 index 0000000000..d620a4e5a2 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator.py @@ -0,0 +1,10 @@ +"""Compatibility shim so buffer.py can discover the ags package.""" + +from slime_plugins.rollout_buffer.generator.ags_generator.entry import ( # noqa: F401 + TASK_TYPE, + generate, + get_group_data_meta_info, + is_valid_group, + run_rollout, + transform_group, +) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/__init__.py b/slime_plugins/rollout_buffer/generator/ags_generator/__init__.py new file mode 100644 index 0000000000..eacbbd947f --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/__init__.py @@ -0,0 +1,5 @@ +"""AGS coding-agent rollout-buffer generator package.""" + +from .entry import TASK_TYPE, generate, run_rollout + +__all__ = ["TASK_TYPE", "run_rollout", "generate"] diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py b/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py new file mode 100644 index 0000000000..bad050c837 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py @@ -0,0 +1,270 @@ +"""SGLang-to-agent adapter service for AGS generator workers.""" + +from __future__ import annotations + +import asyncio +import logging +import os +from argparse import Namespace + +import requests + +from slime.agent.aiohttp_threaded import FilteredAccessLogger, run_app_in_thread +from slime.utils.http_utils import get_host_info +from slime.utils.misc import SingletonMeta +from slime.utils.processing_utils import load_tokenizer +from slime.utils.types import Sample + +from .config import AGSGeneratorConfig + +logger = logging.getLogger(__name__) + + +class RemoteAdapterProxy: + """Control an already-running adapter service via slime control endpoints.""" + + def __init__(self, control_url: str) -> None: + self.control_url = control_url.rstrip("/") + + def _post(self, path: str, payload: dict) -> dict: + response = requests.post(f"{self.control_url}{path}", json=payload, timeout=30) + response.raise_for_status() + return response.json() + + def open_session( + self, + sid: str, + *, + sampling_defaults: dict | None = None, + max_context_tokens: int = 0, + ) -> None: + self._post( + "/_slime/open_session", + { + "sid": sid, + "sampling_defaults": sampling_defaults or {}, + "max_context_tokens": int(max_context_tokens or 0), + }, + ) + + async def finish_session( + self, + sid: str, + *, + base_sample, + reward: float = 0.0, + extra_metadata: dict | None = None, + wait_timeout: float = 5.0, + ) -> list[Sample]: + result = await asyncio.to_thread( + self._post, + "/_slime/finish_session", + { + "sid": sid, + "base_sample": base_sample.to_dict(), + "reward": float(reward), + "extra_metadata": extra_metadata or {}, + "wait_timeout": float(wait_timeout), + }, + ) + return [Sample.from_dict(item) for item in result.get("samples", [])] + + async def drop_session(self, sid: str, *, wait_timeout: float = 5.0) -> None: + await asyncio.to_thread( + self._post, + "/_slime/drop_session", + {"sid": sid, "wait_timeout": float(wait_timeout)}, + ) + + +def _adapter_build_key( + args: Namespace, + config: AGSGeneratorConfig, + adapter_cls: type, + sglang_url: str, +) -> tuple: + """The construction inputs baked into a live adapter, for staleness checks. + + AdapterService is a singleton, so only the first caller's values take + effect; everything here is ignored on later calls. get_adapter_service + compares this key and warns rather than silently serving an adapter built + from a different tokenizer or context budget. + """ + return ( + getattr(args, "hf_checkpoint", None), + int(getattr(args, "rollout_max_context_len", 0) or 0), + getattr(args, "sglang_tool_call_parser", None) or None, + getattr(args, "sglang_reasoning_parser", None) or None, + sglang_url, + adapter_cls.__name__, + config.fork_merge_threshold, + ) + + +class AdapterService(metaclass=SingletonMeta): + def __init__( + self, + args: Namespace, + config: AGSGeneratorConfig, + adapter_cls: type, + *, + port: int | None = None, + ) -> None: + self.tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + self.max_context_len = int(getattr(args, "rollout_max_context_len", 0) or 0) + self.tool_parser = getattr(args, "sglang_tool_call_parser", None) or None + self.reasoning_parser = getattr(args, "sglang_reasoning_parser", None) or None + sglang_url = ( + os.environ.get("SWE_SGLANG_URL") + or os.environ.get("AGS_GENERATOR_SGLANG_URL") + or f"http://{args.sglang_router_ip}:{args.sglang_router_port}" + ) + public_base_url = (os.environ.get("ADAPTER_PUBLIC_BASE_URL") or "").strip().rstrip("/") + if not public_base_url and not config.adapter_public_host: + raise RuntimeError( + "ADAPTER_PUBLIC_HOST or ADAPTER_PUBLIC_BASE_URL is not set; " + "AGS sandboxes need it to reach the adapter" + ) + + self.adapter = adapter_cls( + tokenizer=self.tokenizer, + sglang_url=sglang_url, + tool_parser=self.tool_parser, + reasoning_parser=self.reasoning_parser, + fork_threshold_tokens=config.fork_merge_threshold, + ) + bind_port = config.adapter_port if port is None else port + self.app_handle = run_app_in_thread( + self.adapter.app, + host=config.adapter_bind_host, + port=bind_port, + thread_name="ags-rollout-adapter", + runner_kwargs={"handler_cancellation": True, "access_log_class": FilteredAccessLogger}, + ) + # An ephemeral bind (port=0) means this is the eval fallback adapter, + # which lives in the RolloutManager actor rather than the rollout-buffer + # process. That actor is not pinned to the head node, so neither + # ADAPTER_PUBLIC_BASE_URL nor ADAPTER_PUBLIC_HOST -- both of which name + # the head and a fixed port -- describe where this adapter is listening. + # Advertise the local node's own routable IP and the port actually bound. + public_host = config.adapter_public_host + if bind_port == 0: + local_ip = get_host_info()[1] + if public_base_url: + logger.warning( + "[ags_generator] ignoring ADAPTER_PUBLIC_BASE_URL=%s for the ephemeral eval adapter; " + "advertising %s instead", + public_base_url, + local_ip, + ) + public_base_url = "" + public_host = local_ip + self.adapter_url = public_base_url or f"http://{public_host}:{self.app_handle.port}" + self.build_key = _adapter_build_key(args, config, adapter_cls, sglang_url) + logger.info( + "[ags_generator] tokenizer=%s adapter=%s sglang_url=%s max_context_len=%s tool_parser=%s reasoning_parser=%s", + args.hf_checkpoint, + self.adapter_url, + sglang_url, + self.max_context_len, + self.tool_parser, + self.reasoning_parser, + ) + + +class RemoteAdapterService(metaclass=SingletonMeta): + def __init__(self, args: Namespace, config: AGSGeneratorConfig) -> None: + self.tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + self.max_context_len = int(getattr(args, "rollout_max_context_len", 0) or 0) + public_base_url = (os.environ.get("ADAPTER_PUBLIC_BASE_URL") or "").strip().rstrip("/") + if not public_base_url and not config.adapter_public_host: + raise RuntimeError( + "ADAPTER_PUBLIC_HOST or ADAPTER_PUBLIC_BASE_URL is not set; " + "AGS sandboxes need it to reach the adapter" + ) + control_url = ( + os.environ.get("AGS_EVAL_ADAPTER_CONTROL_URL") + or os.environ.get("ADAPTER_CONTROL_BASE_URL") + or f"http://{config.adapter_public_host}:{config.adapter_port}" + ) + self.adapter = RemoteAdapterProxy(control_url) + self.adapter_url = public_base_url or f"http://{config.adapter_public_host}:{config.adapter_port}" + self.control_url = control_url + logger.info( + "[ags_generator] using remote adapter control=%s public=%s max_context_len=%s", + control_url, + self.adapter_url, + self.max_context_len, + ) + + +def _remote_adapter_alive(control_url: str, timeout: float = 5.0) -> bool: + """True when something is already serving the adapter at control_url.""" + try: + response = requests.get(f"{control_url.rstrip('/')}/healthz", timeout=timeout) + return response.ok + except Exception as exc: + logger.info("[ags_generator] no adapter reachable at %s (%s)", control_url, type(exc).__name__) + return False + + +def get_adapter_service( + args: Namespace, + config: AGSGeneratorConfig, + adapter_cls: type, + *, + evaluation: bool = False, +): + """Return the adapter service backing one AGS rollout or eval pass. + + Training always owns a local adapter. Eval prefers to reuse it -- the + trajectory trees live in that process, and a second adapter on the same port + would fail to bind -- but the training adapter only exists once + entry.run_rollout has built it inside the rollout-buffer process. Under + --num-rollout 0, or --skip-eval-before-train=0 on rollout 0, eval runs + first and there is nothing to reuse, so fall back to a local adapter on an + ephemeral port instead of failing every prompt with a connection error. + """ + if not evaluation: + return _local_adapter_service(args, config, adapter_cls) + + control_url = ( + os.environ.get("AGS_EVAL_ADAPTER_CONTROL_URL") + or os.environ.get("ADAPTER_CONTROL_BASE_URL") + or f"http://{config.adapter_public_host}:{config.adapter_port}" + ) + if _remote_adapter_alive(control_url): + return RemoteAdapterService(args, config) + + # Port 0: the training adapter may still claim config.adapter_port later in + # this run, and two binds on one port would collide. + logger.info( + "[ags_generator] no training adapter at %s; starting a local eval adapter on an ephemeral port", + control_url, + ) + return _local_adapter_service(args, config, adapter_cls, port=0) + + +def _local_adapter_service( + args: Namespace, + config: AGSGeneratorConfig, + adapter_cls: type, + *, + port: int | None = None, +): + """AdapterService singleton, warning when a later caller's args are ignored.""" + service = AdapterService(args, config, adapter_cls, port=port) + sglang_url = ( + os.environ.get("SWE_SGLANG_URL") + or os.environ.get("AGS_GENERATOR_SGLANG_URL") + or f"http://{args.sglang_router_ip}:{args.sglang_router_port}" + ) + wanted = _adapter_build_key(args, config, adapter_cls, sglang_url) + if wanted != service.build_key: + logger.warning( + "[ags_generator] reusing the live adapter built with %s; this call asked for %s. " + "AdapterService is a singleton, so the requested values are ignored.", + service.build_key, + wanted, + ) + return service diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/ags_sandbox.py b/slime_plugins/rollout_buffer/generator/ags_generator/ags_sandbox.py new file mode 100644 index 0000000000..2f49cbffb8 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/ags_sandbox.py @@ -0,0 +1,267 @@ +"""Tencent AGS sandbox backend for the slime coding-agent RL experiment. + +This module intentionally lives in the experiment directory so the slime source +checkout stays unchanged. It implements the same async Sandbox contract as +``slime.agent.sandbox.E2BSandbox`` but creates sandboxes through Tencent AGS's +E2B-compatible gateway and sidecar mount. +""" + +from __future__ import annotations + +import asyncio +import io +import json +import logging +import os +from pathlib import Path +from typing import Any + +from slime.agent.sandbox import ExecResult, FileContent + +logger = logging.getLogger(__name__) + + +def _env(name: str, default: str = "") -> str: + val = os.environ.get(name) + return default if val is None or val == "" else val + + +def _append_no_proxy(host: str) -> None: + existing = os.environ.get("NO_PROXY") or os.environ.get("no_proxy") or "" + parts = [p.strip() for p in existing.split(",") if p.strip()] + for item in (host, ".tencentags.com"): + if item and item not in parts: + parts.append(item) + value = ",".join(parts) + os.environ["NO_PROXY"] = value + os.environ["no_proxy"] = value + + +def _disable_proxy_for_ags() -> None: + # The AGS E2B-compatible endpoint is on Tencent internal networking. In this + # environment routing it through the generic HTTP(S)_PROXY intermittently + # returns STGW 502 during sandbox create/readiness polling. Disable proxy + # for SDK-side AGS RPCs in this Ray worker process. + for key in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"): + os.environ.pop(key, None) + _append_no_proxy(_env("E2B_DOMAIN", "ap-shanghai.tencentags.com")) + + +ENVD_CMD = r""" +set -e +ln -sfn /proc/self/fd /dev/fd + +for p in \ + runtimes/python runtimes/node \ + agents/craft agents/mini-swe agents/swe agents/claude agents/cbc + do + mkdir -p "/opt/${p%/*}" + ln -sfn "/envd-mount/opt/$p" "/opt/$p" + done + +for t in uv uvx; do + ln -sfn "/envd-mount/usr/local/bin/$t" "/usr/local/bin/$t" +done + +for t in node npm npx; do + ln -sfn "/opt/runtimes/node/bin/$t" "/usr/local/bin/$t" +done + +ln -sfn /opt/agents/claude/bin/claude /usr/local/bin/claude + +mkdir -p /etc/pip /root/.config/uv /etc/xdg/uv +printf '[global]\nindex-url = https://mirrors.cloud.tencent.com/pypi/simple\ntrusted-host = mirrors.cloud.tencent.com\n' > /etc/pip.conf +printf '[[index]]\nurl = "https://mirrors.cloud.tencent.com/pypi/simple"\ndefault = true\n' \ + | tee /root/.config/uv/uv.toml /etc/xdg/uv/uv.toml >/dev/null + +/envd-mount/usr/bin/envd & exec sleep infinity +""".strip() + + +class AGSSandbox: + """Async AGS sandbox wrapper compatible with ``slime.agent.sandbox.Sandbox``.""" + + default_lifetime_sec = 3600 + default_rpc_retries = 3 + rpc_backoff_base_sec = 1.0 + + def __init__(self, image: str, *, timeout: int | None = None, rpc_retries: int | None = None) -> None: + self.image = image + self.timeout = timeout or int(_env("SLIME_AGENT_SANDBOX_LIFETIME_SEC", str(self.default_lifetime_sec))) + self.rpc_retries = rpc_retries or int(_env("SLIME_AGENT_SANDBOX_RPC_RETRIES", str(self.default_rpc_retries))) + self._sb = None + self.sandbox_id = "" + + @staticmethod + def _resources() -> dict[str, str]: + raw = _env("AGS_SANDBOX_RESOURCES_JSON", '{"cpu":"2","memory":"4Gi"}') + try: + parsed = json.loads(raw) + return parsed if isinstance(parsed, dict) else {"cpu": "2", "memory": "4Gi"} + except Exception: + return {"cpu": "2", "memory": "4Gi"} + + def _custom_config(self) -> dict[str, Any]: + return { + "image": self.image, + "imageRegistryType": _env("AGS_IMAGE_REGISTRY_TYPE", "enterprise"), + "command": ["/bin/sh", "-c"], + "args": [ENVD_CMD], + "ports": [{"name": "envd", "port": 49983, "protocol": "TCP"}], + "probe": { + "httpGet": {"path": "/health", "port": 49983, "scheme": "HTTP"}, + "readyTimeoutMs": 30000, + "probeTimeoutMs": 1000, + "probePeriodMs": 2000, + "successThreshold": 1, + "failureThreshold": 15, + }, + "resources": self._resources(), + } + + @staticmethod + def _is_transient_rpc_error(e: BaseException) -> bool: + name = type(e).__name__ + if name in { + "ProtocolError", + "LocalProtocolError", + "WriteError", + "ReadError", + "ConnectError", + "ConnectTimeout", + "ReadTimeout", + "WriteTimeout", + "PoolTimeout", + "RemoteProtocolError", + "SSLError", + }: + return True + msg = str(e) + if name == "SandboxException": + return not ("does not exist" in msg or "STOPPED state" in msg) + return False + + async def _rpc_retry(self, op_name: str, coro_factory, *, idempotent: bool = True): + last_err = None + for attempt in range(self.rpc_retries): + try: + return await coro_factory() + except Exception as e: + if not self._is_transient_rpc_error(e): + raise + if not idempotent: + raise + last_err = e + if attempt + 1 < self.rpc_retries: + backoff = self.rpc_backoff_base_sec * (2**attempt) + logger.debug( + "[ags_sandbox] %s transient %s retry %d/%d in %.1fs: %s", + op_name, + type(e).__name__, + attempt + 1, + self.rpc_retries, + backoff, + str(e)[:200], + ) + await asyncio.sleep(backoff) + assert last_err is not None + raise last_err + + async def __aenter__(self) -> AGSSandbox: + os.environ.setdefault("E2B_DOMAIN", _env("E2B_DOMAIN", "ap-shanghai.tencentags.com")) + _disable_proxy_for_ags() + + # The upstream E2B SDK now validates API keys locally and only accepts + # the public e2b_... format. Tencent AGS intentionally uses ark_... + # gateway keys while keeping the E2B-compatible HTTP surface, so bypass + # only this client-side format check and still send the configured key + # as X-API-KEY to AGS. Keep this experiment-local; do not patch slime. + import e2b.api as _e2b_api # type: ignore + + def _allow_ags_api_key(_api_key: str) -> None: + return None + + _e2b_api.validate_api_key = _allow_ags_api_key + from e2b import AsyncSandbox # type: ignore + + template = _env("AGS_BASE_TOOL", "sdt-3fzh6mv6") + md = { + "x-custom-config": json.dumps(self._custom_config(), ensure_ascii=False), + "environment_name": _env("EXPERIMENT_NAME", "slime-coding-agent-rl"), + "session_id": f"slime-coding-agent-rl-{os.getpid()}-{id(self)}", + } + envs = { + "IS_SANDBOX": "1", + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + } + self._sb = await AsyncSandbox.create(template=template, timeout=self.timeout, metadata=md, envs=envs) + self.sandbox_id = getattr(self._sb, "sandbox_id", getattr(self._sb, "id", "")) + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + try: + if self._sb is not None: + await self._sb.kill() + except Exception as e: + logger.warning("[ags_sandbox] kill %s failed: %s", self.sandbox_id[:8], e) + + async def exec( + self, + cmd: str, + *, + user: str = "root", + env: dict[str, str] | None = None, + timeout: int = 120, + check: bool = False, + idempotent: bool = True, + ) -> ExecResult: + from e2b.sandbox.commands.command_handle import CommandExitException + + # AGS sidecar commands are root-friendly. Keep the contract's user arg, + # but fall back to root if a base image lacks the requested user. + async def _run(): + return await self._sb.commands.run( + cmd, + user=user, + envs=env, + timeout=timeout, + on_stdout=lambda _s: None, + on_stderr=lambda _s: None, + ) + + try: + res = await self._rpc_retry(f"exec({cmd[:60]!r})", _run, idempotent=idempotent) + return res.exit_code, res.stdout or "", res.stderr or "" + except CommandExitException as e: + if check: + raise RuntimeError( + f"ags exec failed (exit={e.exit_code}): {cmd[:160]}\n{(e.stderr or '')[:800]}" + ) from None + return e.exit_code, e.stdout or "", e.stderr or "" + + async def write_file(self, sandbox_path: str, content: FileContent, *, user: str = "root") -> None: + if isinstance(content, Path): + with open(content, "rb") as fp: + data = fp.read() + await self.write_file(sandbox_path, data, user=user) + return + if isinstance(content, bytes): + await self._rpc_retry( + f"write_file({sandbox_path}, bytes={len(content)})", + lambda: self._sb.files.write( + sandbox_path, io.BytesIO(content), user=user, gzip=False, use_octet_stream=True + ), + ) + return + await self._rpc_retry( + f"write_file({sandbox_path})", + lambda: self._sb.files.write(sandbox_path, content, user=user), + ) + + async def read_file(self, sandbox_path: str, *, user: str = "root") -> str: + try: + return await self._rpc_retry( + f"read_file({sandbox_path})", lambda: self._sb.files.read(sandbox_path, user=user) + ) + except Exception: + return "" diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/artifacts.py b/slime_plugins/rollout_buffer/generator/ags_generator/artifacts.py new file mode 100644 index 0000000000..572883041c --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/artifacts.py @@ -0,0 +1,70 @@ +"""Local artifact helpers for AGS rollout-buffer generation.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +from slime.agent.sandbox import Sandbox +from slime.utils.types import Sample + + +def safe_artifact_name(value: str) -> str: + return "".join(c if c.isalnum() or c in "._-" else "_" for c in value)[:200] or "unknown" + + +def sample_artifact_id(instance_id: str, sample: Sample) -> str: + parts = [instance_id] + if sample.group_index is not None: + parts.append(f"g{sample.group_index}") + if sample.index is not None: + parts.append(f"i{sample.index}") + if sample.session_id: + parts.append(sample.session_id[-8:]) + return "__".join(parts) + + +class ArtifactWriter: + def __init__(self, root: str | os.PathLike[str] | None = None) -> None: + self.root = Path(root) if root else None + if self.root is not None: + self.root.mkdir(parents=True, exist_ok=True) + + @classmethod + def from_env(cls) -> ArtifactWriter: + return cls(os.environ.get("TRAJECTORY_DUMP_DIR", "").strip() or None) + + def enabled(self) -> bool: + return self.root is not None + + def _path(self, artifact_id: str, suffix: str) -> Path | None: + if self.root is None: + return None + return self.root / f"{safe_artifact_name(artifact_id)}{suffix}" + + async def dump_trajectory(self, sb: Sandbox, workdir: str, artifact_id: str) -> str | None: + path = self._path(artifact_id, ".trajectory.jsonl") + if path is None: + return None + try: + content = await sb.read_file(f"{workdir}/.harness/trajectory.jsonl", user="root") + except Exception: + content = "" + path.write_text(content or "", encoding="utf-8") + return str(path) + + def dump_patch(self, diff_text: str, artifact_id: str) -> str | None: + path = self._path(artifact_id, ".patch") + if path is None: + return None + path.write_text(diff_text or "", encoding="utf-8") + return str(path) + + def dump_rollout(self, payload: dict[str, Any], artifact_id: str) -> str | None: + path = self._path(artifact_id, ".rollout.json") + if path is None: + return None + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2, default=str), encoding="utf-8") + return str(path) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/config.py b/slime_plugins/rollout_buffer/generator/ags_generator/config.py new file mode 100644 index 0000000000..c6d7ad2dbc --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/config.py @@ -0,0 +1,124 @@ +"""Configuration for the AGS coding-agent rollout-buffer generator.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + +# What to do when the agent exits 0 but leaves no diff. +# "off" -- skip the check entirely +# "metrics" -- label the samples and count it; no behaviour change (default) +# "abort" -- additionally drop the rollout via Sample.Status.ABORTED +# +# "abort" is deliberately not the default. An empty-patch trajectory still holds +# real on-policy tokens whose reward of 0 is correct, i.e. the negative half of +# a GRPO group; at the observed ~33% empty-patch rate, masking those out biases +# the group baseline toward successes. It also only requeues under +# slime.rollout.fully_async_rollout -- on the rollout_buffer path the AGS +# scripts actually use, an aborted sample is shipped with its loss masked and +# never retried. +EMPTY_PATCH_GUARD_POLICIES = frozenset({"off", "metrics", "abort"}) + +# How the agent learns what to do. +# "instruction" -- send SWE_CC_PROMPT, which points at PROBLEM_STATEMENT.md +# "dataset" -- send the row's own prompt field verbatim +# +# Set independently for training (SWE_PROMPT_STYLE) and periodic eval +# (SWE_EVAL_PROMPT_STYLE), because the two want different things: eval should +# measure the model the way a benchmark would, handing over the task text +# directly the way Harbor does, while training may prefer the agent to work for +# it. Converted Harbor rows carry the same instruction.md text in both the prompt +# field and metadata.problem_statement, so the styles differ in *when* the agent +# sees the task, not in what it reads. +# +# Only "instruction" writes PROBLEM_STATEMENT.md into the workspace (see +# swe_task.prepare_workspace) -- under "dataset" the prompt already carries the +# task text, so the file would just be a stray artifact in the repo. +PROMPT_STYLES = frozenset({"dataset", "instruction"}) + + +@dataclass(frozen=True) +class AGSGeneratorConfig: + agent_name: str + adapter_public_host: str | None + adapter_bind_host: str + adapter_port: int + fork_merge_threshold: int | None + agent_time_budget_sec: int + eval_timeout_sec: int + eval_bootstrap_cmd: str | None + eval_isolated_sandbox: bool + rollout_guard_sec: int + boot_concurrency: int + rollout_concurrency: int + boot_retries: int + artifact_dir: str | None + enable_token2text: bool + prompt: str + empty_patch_guard: str + prompt_style: str + eval_prompt_style: str + + @classmethod + def from_env(cls, *, enable_token2text: bool = False) -> AGSGeneratorConfig: + agent_time_budget = int(os.environ.get("SWE_AGENT_TIME_BUDGET_SEC", "1800")) + eval_timeout = int(os.environ.get("SWE_EVAL_TIMEOUT_SEC", "600")) + guard = int(os.environ.get("SWE_ROLLOUT_GUARD_SEC", "0") or 0) or (agent_time_budget + eval_timeout + 180) + fork = int(v) if (v := os.environ.get("SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS")) else None + rollout_concurrency = int(os.environ.get("SWE_ROLLOUT_CONCURRENCY", "1")) + return cls( + agent_name=os.environ.get("SWE_AGENT", "claude_code"), + adapter_public_host=os.environ.get("ADAPTER_PUBLIC_HOST"), + adapter_bind_host=os.environ.get("ADAPTER_BIND_HOST", "0.0.0.0"), + adapter_port=int(os.environ.get("ADAPTER_PORT", "18001")), + fork_merge_threshold=fork, + agent_time_budget_sec=agent_time_budget, + eval_timeout_sec=eval_timeout, + eval_bootstrap_cmd=os.environ.get("SWE_EVAL_BOOTSTRAP_CMD") or None, + eval_isolated_sandbox=_env_flag("SWE_EVAL_ISOLATED_SANDBOX", default=False), + rollout_guard_sec=guard, + boot_concurrency=int(os.environ.get("SWE_BOOT_CONCURRENCY", "16")), + rollout_concurrency=max(1, rollout_concurrency), + boot_retries=int(os.environ.get("SWE_BOOT_RETRIES", "2")), + artifact_dir=os.environ.get("TRAJECTORY_DUMP_DIR", "").strip() or None, + enable_token2text=enable_token2text, + prompt=os.environ.get( + "SWE_CC_PROMPT", + "Read PROBLEM_STATEMENT.md in the current directory and resolve the issue. Edit source files only (do NOT touch tests). After editing, run the relevant tests to verify your fix passes. Do NOT modify PROBLEM_STATEMENT.md and do NOT commit. When finished, print a one-line summary and exit.", + ), + empty_patch_guard=_empty_patch_guard_policy(os.environ.get("SWE_EMPTY_PATCH_GUARD")), + prompt_style=_prompt_style("SWE_PROMPT_STYLE", default="instruction"), + eval_prompt_style=_prompt_style("SWE_EVAL_PROMPT_STYLE", default="dataset"), + ) + + def prompt_style_for(self, *, evaluation: bool) -> str: + """Prompt style for this rollout: eval and training are set separately.""" + return self.eval_prompt_style if evaluation else self.prompt_style + + +def _empty_patch_guard_policy(raw: str | None) -> str: + """Validate SWE_EMPTY_PATCH_GUARD, defaulting to "metrics". + + Raising here fails the run at construction time; silently falling back to a + default would disable the guard on a typo without anyone noticing. + """ + policy = (raw or "metrics").strip().lower() + if policy not in EMPTY_PATCH_GUARD_POLICIES: + raise ValueError(f"SWE_EMPTY_PATCH_GUARD={raw!r} is not one of {sorted(EMPTY_PATCH_GUARD_POLICIES)}") + return policy + + +def _prompt_style(env_name: str, *, default: str) -> str: + """Validate a prompt-style env var, naming it in the error.""" + raw = os.environ.get(env_name) + style = (raw or default).strip().lower() + if style not in PROMPT_STYLES: + raise ValueError(f"{env_name}={raw!r} is not one of {sorted(PROMPT_STYLES)}") + return style + + +def _env_flag(name: str, *, default: bool) -> bool: + raw = os.environ.get(name) + if raw is None or raw == "": + return default + return raw.lower() in {"1", "true", "yes", "on"} diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/empty_patch_guard.py b/slime_plugins/rollout_buffer/generator/ags_generator/empty_patch_guard.py new file mode 100644 index 0000000000..0497c1972d --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/empty_patch_guard.py @@ -0,0 +1,111 @@ +"""Detect coding-agent rollouts that exit successfully without editing anything. + +Some agent CLIs treat a text-only final turn as a finished answer: the model +stops mid-task ("Let me look at query.py:"), no tool call is parsed, the CLI +exits 0, and the rollout yields an empty diff. Measured on SWE-bench Verified +with CodeBuddy Code, that accounts for roughly a third of all rollouts, so it is +worth counting and attributing rather than leaving it indistinguishable from a +genuine "no change needed" outcome. + +This module only classifies. The policy decision (see +``AGSGeneratorConfig.empty_patch_guard``) lives in the caller, and the default +does not change rollout behaviour -- an empty-patch trajectory is still real +on-policy data whose zero reward is correct. +""" + +from __future__ import annotations + +import dataclasses +import logging +import re + +from .weave_trace import iter_trajectory_events + +logger = logging.getLogger(__name__) + +# Reasons are attribution labels, not trigger conditions. The guard fires on +# "exit 0 + empty diff" alone; on SWE-bench an empty patch is a failure no +# matter how the agent phrased its last message, and gating the counter on +# string heuristics would make it brittle. +REASON_MIDTASK_NARRATION = "empty_patch_midtask_narration" +REASON_CLAIMED_COMPLETE = "empty_patch_claimed_complete" +REASON_NO_FINAL_TEXT = "empty_patch_no_final_text" +REASON_UNCLASSIFIED = "empty_patch_unclassified" + +# Checked against the tail of the final message, where the intent actually sits. +_CONTINUATION_RE = re.compile( + r"\b(?:let me|let's|now let me|next,?|i'll|i will|i'm going to|i need to|i should)\b", + re.IGNORECASE, +) +_COMPLETION_RE = re.compile( + r"\b(?:fix is complete|all tests? pass|issue is resolved|changes? (?:are|is) complete|" + r"successfully (?:fixed|resolved)|the fix works)\b", + re.IGNORECASE, +) +_FINAL_TEXT_TAIL_CHARS = 400 +_FINAL_TEXT_EXCERPT_CHARS = 400 + + +@dataclasses.dataclass(frozen=True) +class GuardVerdict: + """Outcome of the empty-patch check for one rollout.""" + + triggered: bool + reason: str = "" + final_text: str = "" + + +def final_agent_text(trajectory_path: str | None, *, agent: str) -> str | None: + """Return the agent's last user-visible text, or None if unavailable. + + None means "could not tell" rather than "empty": the trajectory dump is + disabled (``TRAJECTORY_DUMP_DIR`` unset, so ``dump_trajectory`` returned + None) or the agent has no trajectory parser (codex yields no events). + """ + if not trajectory_path: + return None + result_text: str | None = None + last_assistant_text: str | None = None + for event in iter_trajectory_events(trajectory_path, agent=agent): + kind = event.get("kind") + if kind == "result": + output = event.get("output") + if isinstance(output, dict) and output.get("result") is not None: + result_text = str(output["result"]) + elif kind == "text": + output = event.get("output") + if isinstance(output, dict) and output.get("text"): + last_assistant_text = str(output["text"]) + if result_text is not None: + return result_text + return last_assistant_text + + +def classify_empty_patch( + *, + agent_exit_code: int, + diff_text: str, + trajectory_path: str | None, + agent: str, +) -> GuardVerdict: + """Flag a rollout that reported success but produced no diff.""" + # git_diff returns "" (not None) when nothing changed. + if agent_exit_code != 0 or (diff_text or "").strip(): + return GuardVerdict(triggered=False) + + text = final_agent_text(trajectory_path, agent=agent) + if text is None: + return GuardVerdict(triggered=True, reason=REASON_UNCLASSIFIED) + if not text.strip(): + return GuardVerdict(triggered=True, reason=REASON_NO_FINAL_TEXT) + + excerpt = text[-_FINAL_TEXT_EXCERPT_CHARS:] + tail = text.replace("<|im_end|>", " ")[-_FINAL_TEXT_TAIL_CHARS:] + # Completion claims win: "all tests pass. Let me summarise." is a claimed + # completion, not mid-task narration. With no diff it is a distinct failure + # mode -- the agent believes it edited something it never edited. + if _COMPLETION_RE.search(tail): + return GuardVerdict(triggered=True, reason=REASON_CLAIMED_COMPLETE, final_text=excerpt) + if _CONTINUATION_RE.search(tail): + return GuardVerdict(triggered=True, reason=REASON_MIDTASK_NARRATION, final_text=excerpt) + return GuardVerdict(triggered=True, reason=REASON_UNCLASSIFIED, final_text=excerpt) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py new file mode 100644 index 0000000000..10f951bd7c --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -0,0 +1,412 @@ +"""rollout_buffer generator entry point for AGS coding-agent tasks.""" + +from __future__ import annotations + +import asyncio +import copy +import logging +from argparse import Namespace +from typing import Any + +import requests + +from slime.utils.misc import SingletonMeta +from slime.utils.types import Sample + +from .config import AGSGeneratorConfig +from .rollout import AGSRolloutRunner +from .serialization import output_item_from_samples, samples_from_payload +from .source import AGSPromptSource + +TASK_TYPE = "ags" + +logger = logging.getLogger(__name__) + + +class _AGSGenerateState(metaclass=SingletonMeta): + """Process-local state for the per-sample AGS generate hook. + + The default slime eval loop calls the custom generate function once per + sample. Keep the runner and concurrency semaphore process-local so periodic + eval reuses the same adapter service instead of rebuilding it for every + prompt. + """ + + def __init__(self, args: Namespace, *, evaluation: bool = False) -> None: + self.config = AGSGeneratorConfig.from_env( + enable_token2text=_as_bool(getattr(args, "enable_token2text", False)) + ) + # This singleton is keyed on the class, not on `evaluation`, so the first + # caller's mode does fix the runner's adapter for the process. That is + # harmless: one adapter serves both modes (sessions are keyed by sid), + # and get_adapter_service falls back to a local adapter when the training + # one is absent. The prompt style, which must differ per call, is passed + # to generate() instead of being read off the runner. + self.runner = AGSRolloutRunner(args, self.config, use_remote_adapter=evaluation) + self.semaphore = asyncio.Semaphore(self.config.rollout_concurrency) + + +async def generate( + args: Namespace, + base_sample: Sample, + sampling_params: dict[str, Any], + evaluation: bool = False, +) -> list[Sample]: + """Generate one AGS rollout for slime's default rollout/eval loop. + + This hook lets periodic eval use the AGS runner via: + + --eval-function-path slime.rollout.sglang_rollout.generate_rollout + --custom-generate-function-path slime_plugins.rollout_buffer.generator.ags_generator.generate + + Training keeps AGSRolloutRunner's trainable segment output. Eval collapses + the possibly multi-segment trajectory into one scored sample so pass-rate + metrics count one eval attempt per prompt. + + ``evaluation`` is forwarded to each generate() call rather than baked into the + runner: _AGSGenerateState is a singleton, so the first caller's mode would + otherwise stick for the life of the process and silently give eval the + training prompt style (or vice versa). + """ + + state = _AGSGenerateState(args, evaluation=evaluation) + async with state.semaphore: + samples = await state.runner.generate(base_sample, sampling_params, evaluation=evaluation) + + if not evaluation: + return samples + return [_collapse_eval_samples(base_sample, samples)] + + +def run_rollout(data: dict[str, Any]) -> str: + """Generate AGS coding-agent trajectories and stream them into buffer.py.""" + + logging.basicConfig(level=getattr(logging, data.get("log_level", "INFO"), logging.INFO)) + args = _build_args(data) + config = AGSGeneratorConfig.from_env(enable_token2text=_as_bool(data.get("enable_token2text", False))) + source = AGSPromptSource(args) + runner = AGSRolloutRunner(args, config) + remote_buffer_url = data["remote_buffer_url"].rstrip("/") + "/buffer/write" + num_epoch = int(data.get("num_epoch", 1)) + groups_per_epoch = int( + data.get("num_groups_per_epoch") or data.get("rollout_batch_size") or args.rollout_batch_size + ) + skip_instance_ids = data.get("skip_instance_ids") or [] + stop_event = data.get("_stop_event") + rollout_job_id = data.get("_rollout_job_id") + + logger.info( + "[ags_generator] start task_type=%s groups_per_epoch=%s repeats=%s epochs=%s concurrency=%s buffer=%s", + TASK_TYPE, + groups_per_epoch, + args.n_samples_per_prompt, + num_epoch, + config.rollout_concurrency, + remote_buffer_url, + ) + + async def _run_sample(epoch: int, sample: Sample) -> None: + outputs = await runner.generate(sample, args.sampling_params) + first = outputs[0] + instance_id = _instance_id(first) + item = output_item_from_samples( + outputs, + instance_id=instance_id, + extra_info={ + "epoch": epoch, + "task_type": TASK_TYPE, + "reward": first.reward, + **(first.metadata or {}), + }, + ) + if rollout_job_id is not None: + item["rollout_job_id"] = rollout_job_id + await asyncio.to_thread(_send_data_to_buffer, remote_buffer_url, item) + + async def _run_epoch(epoch: int, samples: list[Sample]) -> None: + async def _guarded(sample: Sample) -> None: + try: + await _run_sample(epoch, sample) + except asyncio.CancelledError: + raise + except Exception as exc: + instance_id = _instance_id(sample) + logger.exception( + "[ags_generator] %s: sample task failed; writing aborted rollout: %s", + instance_id, + exc, + ) + outputs = runner._abort_result(sample, f"task_exception:{type(exc).__name__}", instance_id) + first = outputs[0] + item = output_item_from_samples( + outputs, + instance_id=instance_id, + extra_info={ + "epoch": epoch, + "task_type": TASK_TYPE, + "reward": first.reward, + **(first.metadata or {}), + }, + ) + if rollout_job_id is not None: + item["rollout_job_id"] = rollout_job_id + try: + await asyncio.to_thread(_send_data_to_buffer, remote_buffer_url, item) + except Exception as send_exc: + logger.exception( + "[ags_generator] %s: failed to write aborted rollout after task failure: %s", + instance_id, + send_exc, + ) + + pending: set[asyncio.Task[None]] = set() + sample_iter = iter(samples) + exhausted = False + + while pending or not exhausted: + while not exhausted and len(pending) < config.rollout_concurrency and not _stop_requested(stop_event): + try: + sample = next(sample_iter) + except StopIteration: + exhausted = True + break + pending.add(asyncio.create_task(_guarded(sample))) + + if not pending: + break + + done, pending = await asyncio.wait(pending, timeout=1.0, return_when=asyncio.FIRST_COMPLETED) + for task in done: + task.result() + + if _stop_requested(stop_event): + logger.info( + "[ags_generator] stop requested; cancelling %d in-flight tasks for epoch %d", + len(pending), + epoch, + ) + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + pending.clear() + break + + for epoch in range(num_epoch): + if _stop_requested(stop_event): + logger.info("[ags_generator] stop requested before epoch %d; exiting", epoch) + break + samples = source.get_repeated_samples(groups_per_epoch, skip_instance_ids=skip_instance_ids) + skip_instance_ids = [] + asyncio.run(_run_epoch(epoch, samples)) + if _stop_requested(stop_event): + logger.info("[ags_generator] stop requested after epoch %d; exiting", epoch) + break + return "finished" + + +def _collapse_eval_samples(base_sample: Sample, samples: list[Sample]) -> Sample: + """Return a single eval Sample from AGSRolloutRunner's segment output.""" + + sample = copy.deepcopy(base_sample) + first = samples[0] if samples else None + metadata = dict(getattr(sample, "metadata", None) or {}) + if first is not None: + metadata.update(first.metadata or {}) + metadata["eval_collapsed_segments"] = len(samples) + + sample.tokens = [0, 0] + sample.response = "" + sample.response_length = 1 + sample.loss_mask = [0] + sample.rollout_log_probs = [0.0] + sample.reward = 0.0 if first is None or first.reward is None else float(first.reward) + sample.remove_sample = True + sample.status = Sample.Status.ABORTED if first is None else first.status + sample.metadata = metadata + return sample + + +def transform_group(group, task_type: str = TASK_TYPE): + return group + + +def is_valid_group(group, min_valid_group_size: int, task_type: str = TASK_TYPE) -> bool: + _instance_id, items = group + valid = 0 + for item in items: + try: + samples = samples_from_payload(item) + except Exception: + continue + if samples and all(sample.response_length > 0 and sample.tokens for sample in samples): + valid += 1 + return len(items) >= min_valid_group_size and valid >= min_valid_group_size + + +def get_group_data_meta_info(temp_data: dict[str, list[dict[str, Any]]]) -> dict[str, Any]: + rewards = [] + status_counts: dict[str, int] = {} + artifact_counts = {"trajectory": 0, "patch": 0, "rollout_dump": 0, "complete": 0} + elapsed_secs = [] + agent_exit_nonzero = 0 + applied_cleanly = 0 + rollout_concurrency = 0 + total_samples = 0 + total_rollouts = sum(len(items) for items in temp_data.values()) + + for items in temp_data.values(): + for item in items: + samples = samples_from_payload(item) + for sample in samples: + total_samples += 1 + if sample.reward is not None: + rewards.append(float(sample.reward)) + status_counts[sample.status.value] = status_counts.get(sample.status.value, 0) + 1 + metadata = sample.metadata or {} + has_trajectory = bool(metadata.get("trajectory_path")) + has_patch = bool(metadata.get("patch_path")) + has_rollout_dump = bool(metadata.get("rollout_dump_path")) + artifact_counts["trajectory"] += int(has_trajectory) + artifact_counts["patch"] += int(has_patch) + artifact_counts["rollout_dump"] += int(has_rollout_dump) + artifact_counts["complete"] += int(has_trajectory and has_patch and has_rollout_dump) + elapsed_sec = _float_or_none(metadata.get("ags_elapsed_sec")) + if elapsed_sec is not None: + elapsed_secs.append(elapsed_sec) + agent_exit_nonzero += int((metadata.get("agent_exit_code") or 0) != 0) + applied_cleanly += int(bool(metadata.get("applied_cleanly"))) + rollout_concurrency = max(rollout_concurrency, int(metadata.get("ags_rollout_concurrency") or 0)) + + completed = status_counts.get(Sample.Status.COMPLETED.value, 0) + aborted = status_counts.get(Sample.Status.ABORTED.value, 0) + solved = sum(1 for reward in rewards if reward == 1.0) + nonzero = sum(1 for reward in rewards if reward != 0) + return { + "total_samples": total_samples, + "total_rollouts": total_rollouts, + "num_groups": len(temp_data), + "avg_group_size": total_rollouts / len(temp_data) if temp_data else 0, + "avg_samples_per_group": total_samples / len(temp_data) if temp_data else 0, + "avg_reward": sum(rewards) / len(rewards) if rewards else 0, + "nonzero_reward_samples": nonzero, + "solved_samples": solved, + "solve_rate": solved / len(rewards) if rewards else 0, + "nonzero_reward_rate": nonzero / len(rewards) if rewards else 0, + "completed_rate": completed / total_samples if total_samples else 0, + "abort_rate": aborted / total_samples if total_samples else 0, + "artifact_complete_rate": artifact_counts["complete"] / total_samples if total_samples else 0, + "status_counts": status_counts, + "artifact_counts": artifact_counts, + "performance": { + "rollout_concurrency": rollout_concurrency, + "avg_elapsed_sec": sum(elapsed_secs) / len(elapsed_secs) if elapsed_secs else 0, + "p50_elapsed_sec": _percentile(elapsed_secs, 0.50), + "p95_elapsed_sec": _percentile(elapsed_secs, 0.95), + "max_elapsed_sec": max(elapsed_secs) if elapsed_secs else 0, + "elapsed_sec_values": elapsed_secs, + "agent_exit_nonzero_count": agent_exit_nonzero, + "applied_cleanly_count": applied_cleanly, + }, + } + + +def _float_or_none(value: Any) -> float | None: + if value is None: + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def _percentile(values: list[float], quantile: float) -> float: + if not values: + return 0 + ordered = sorted(values) + index = min(len(ordered) - 1, max(0, int(round((len(ordered) - 1) * quantile)))) + return ordered[index] + + +def _build_args(data: dict[str, Any]) -> Namespace: + sampling_params = dict(data.get("sampling_params") or {}) + if "max_tokens" not in sampling_params and "max_tokens" in data: + sampling_params["max_tokens"] = int(data["max_tokens"]) + max_tokens = int(sampling_params.get("max_tokens") or data.get("max_tokens") or 4096) + prompt_data = data["input_file"] + return Namespace( + hf_checkpoint=data["tokenizer_path"], + prompt_data=prompt_data, + input_key=data.get("input_key", "prompt"), + label_key=data.get("label_key", "label"), + metadata_key=data.get("metadata_key", "metadata"), + tool_key=data.get("tool_key"), + multimodal_keys=data.get("multimodal_keys"), + apply_chat_template=_as_bool(data.get("apply_chat_template", False)), + apply_chat_template_kwargs=data.get("apply_chat_template_kwargs") or {}, + rollout_global_dataset=True, + rollout_shuffle=_as_bool(data.get("rollout_shuffle", False)), + rollout_seed=int(data.get("rollout_seed", 42)), + rollout_max_prompt_len=data.get("rollout_max_prompt_len"), + dump_details=None, + rollout_max_context_len=int(data.get("rollout_max_context_len", 0) or 0), + rollout_batch_size=int(data.get("rollout_batch_size", 1)), + rollout_start_group=int(data.get("rollout_start_group", 0) or 0), + n_samples_per_prompt=int(data["num_repeat_per_sample"]), + sglang_router_ip=_router_ip(data["remote_engine_url"]), + sglang_router_port=_router_port(data["remote_engine_url"]), + sglang_tool_call_parser=data.get("sglang_tool_call_parser"), + sglang_reasoning_parser=data.get("sglang_reasoning_parser"), + use_wandb=_as_bool(data.get("use_wandb", False)), + wandb_mode=data.get("wandb_mode"), + wandb_project=data.get("wandb_project"), + wandb_team=data.get("wandb_team"), + wandb_run_id=data.get("wandb_run_id"), + wandb_group=data.get("wandb_group"), + sampling_params=sampling_params | {"max_tokens": max_tokens}, + ) + + +def _router_ip(url: str) -> str: + from urllib.parse import urlparse + + parsed = urlparse(url if "://" in url else f"http://{url}") + return parsed.hostname or "127.0.0.1" + + +def _router_port(url: str) -> int: + from urllib.parse import urlparse + + parsed = urlparse(url if "://" in url else f"http://{url}") + return int(parsed.port or 80) + + +def _instance_id(sample: Sample) -> str: + metadata = sample.metadata or {} + remote = metadata.get("remote_env_info") or {} + label = sample.label if isinstance(sample.label, str) and len(sample.label) < 256 else None + return str(metadata.get("instance_id") or remote.get("instance_id") or label or sample.index or "unknown") + + +def _send_data_to_buffer(remote_buffer_url: str, data: dict[str, Any]) -> None: + last_err = None + for _ in range(3): + try: + response = requests.post(remote_buffer_url, json=data, timeout=30) + if response.status_code == 200: + return + last_err = RuntimeError(f"status={response.status_code} body={response.text[:200]}") + except Exception as exc: + last_err = exc + raise RuntimeError(f"send data to buffer failed: {last_err}") + + +def _as_bool(value: Any) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "y", "on"} + return bool(value) + + +def _stop_requested(stop_event: Any) -> bool: + return bool(stop_event is not None and stop_event.is_set()) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py new file mode 100644 index 0000000000..fc61c26867 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py @@ -0,0 +1,303 @@ +"""Harness registry for AGS-backed coding agents.""" + +from __future__ import annotations + +import base64 +import json +import os +import shlex + +from slime.agent.adapters import AnthropicAdapter, OpenAIAdapter +from slime.agent.harness import CodexHarness +from slime.agent.harness.common import BaseHarness, HarnessContext +from slime.agent.sandbox import Sandbox + +from .runner import run_root_command + + +class AGSSidecarClaudeCodeHarness(BaseHarness): + """Claude Code harness using the AGS sidecar binary instead of npm install.""" + + name = "claude_code" + extra_args_env = "SLIME_AGENT_CC_EXTRA_ARGS" + extra_envs_env = "SLIME_AGENT_CC_EXTRA_ENVS" + launch_flags = "--dangerously-skip-permissions --verbose --output-format stream-json --include-partial-messages --include-hook-events" + + # Baseline flags, emitted BEFORE extra_args so a caller can override any of + # them (claude, like cbc, takes the last occurrence of a repeated flag). + default_flags = ("--disallowedTools WebSearch WebFetch",) + + static_env = { + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", + "CLAUDE_CODE_ATTRIBUTION_HEADER": "0", + } + + async def install_cli(self, sb: Sandbox) -> None: + await sb.exec( + "command -v node && node --version && command -v claude && claude --version", + user="root", + check=True, + timeout=120, + ) + + async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: + settings = json.dumps({"hasCompletedOnboarding": True, "bypassPermissionsModeAccepted": True}) + await sb.exec( + f"mkdir -p /root/.claude /home/agent/.claude && echo {shlex.quote(settings)} | tee /root/.claude.json /root/.claude/settings.json /home/agent/.claude.json /home/agent/.claude/settings.json > /dev/null && chown -R agent:agent /home/agent/.claude /home/agent/.claude.json", + user="root", + check=True, + timeout=60, + ) + + async def launch_and_wait(self, sb: Sandbox, ctx: HarnessContext, prompt: str, time_budget_sec: int) -> int: + cmd = f"/usr/local/bin/claude -p {shlex.quote(prompt)} {self.launch_flags} {' '.join(self.default_flags)}" + extra = os.environ.get(self.extra_args_env, "").strip() + if extra: + # Last, so it overrides default_flags. + cmd = f"{cmd} {extra}" + + env = { + "ANTHROPIC_BASE_URL": ctx.adapter_url, + "ANTHROPIC_AUTH_TOKEN": ctx.session_id, + "ANTHROPIC_API_KEY": ctx.session_id, + "ANTHROPIC_MODEL": ctx.model_label, + "ANTHROPIC_DEFAULT_SONNET_MODEL": ctx.model_label, + "ANTHROPIC_DEFAULT_OPUS_MODEL": ctx.model_label, + "ANTHROPIC_DEFAULT_HAIKU_MODEL": ctx.model_label, + "CLAUDE_CODE_SUBAGENT_MODEL": ctx.model_label, + **self.static_env, + "IS_SANDBOX": "1", + } + extra_envs = os.environ.get(self.extra_envs_env, "").strip() + if extra_envs: + # Applied last so it can override static_env. + env.update(json.loads(extra_envs)) + + return await run_root_command( + sb, + workdir=ctx.workdir, + start_cmd=cmd, + env=env, + time_budget_sec=time_budget_sec, + ) + + async def run( + self, + sb: Sandbox, + *, + workdir: str, + session_id: str, + adapter_url: str, + time_budget_sec: int, + prompt: str, + ) -> int: + from slime.agent import sandbox as agent_sandbox + + await agent_sandbox.ensure_agent_user(sb, workdir) + ctx = HarnessContext(workdir=workdir, session_id=session_id, adapter_url=adapter_url) + await self.write_config(sb, ctx) + return await self.launch_and_wait(sb, ctx, prompt, time_budget_sec) + + +class CodeBuddyCodeHarness(BaseHarness): + """CodeBuddy Code (cbc) harness using the AGS sidecar binary.""" + + name = "codebuddy_code" + extra_args_env = "SLIME_AGENT_CBC_EXTRA_ARGS" + extra_envs_env = "SLIME_AGENT_CBC_EXTRA_ENVS" + # The serving-side context cap, mirrored into models.json as maxInputTokens so + # auto-compaction can be expressed as a percentage of it. Read from the env + # rather than passed in because HarnessContext carries no run config. + max_input_tokens_env = "SLIME_AGENT_MAX_INPUT_TOKENS" + + # Baseline flags, emitted BEFORE extra_args so a caller can override any of + # them: commander's last-flag-wins was verified against cbc 2.125.0 + # (`--max-turns 99 --max-turns 1` stops after 1). + # + # --disallowedTools rather than --tools: `--tools` does NOT restrict the + # surface, whereas --disallowedTools is enforced. + # Denying web access keeps rollouts reproducible and stops the agent looking + # up the fix being graded. The flag is variadic, so it has to stay ahead of + # the non-variadic tail or it swallows --max-turns' value and the prompt. + default_flags = ("--disallowedTools WebSearch WebFetch",) + + async def install_cli(self, sb: Sandbox) -> None: + await sb.exec( + "set -e\n" + "if [ -x /opt/runtimes/node/bin/node ]; then\n" + " ln -sf /opt/runtimes/node/bin/node /usr/local/bin/node\n" + " ln -sf /opt/runtimes/node/bin/npm /usr/local/bin/npm\n" + " ln -sf /opt/runtimes/node/bin/npx /usr/local/bin/npx\n" + "fi\n" + "python_bin=$(ls /opt/runtimes/python/cpython-*/bin/python3 " + "/envd-mount/opt/runtimes/python/cpython-*/bin/python3 2>/dev/null | head -1 || true)\n" + 'if [ -n "$python_bin" ]; then\n' + ' ln -sf "$python_bin" /usr/local/bin/python3\n' + ' pip_bin="${python_bin%/python3}/pip3"\n' + ' [ -x "$pip_bin" ] && ln -sf "$pip_bin" /usr/local/bin/pip3\n' + "fi\n" + "test -x /opt/agents/cbc/bin/cbc\n" + "ln -sf /opt/agents/cbc/bin/cbc /usr/local/bin/cbc\n" + "command -v node && node --version\n" + "command -v cbc && cbc --version", + user="root", + check=True, + timeout=120, + ) + + async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: + model_entry = { + "id": ctx.model_label, + "name": ctx.model_label, + "vendor": "OpenAI", + "apiKey": ctx.session_id, + "url": self._chat_completions_url(ctx.adapter_url), + "supportsToolCall": True, + "supportsImages": False, + "supportsReasoning": True, + } + # maxInputTokens is what makes percentage-based auto-compaction usable. + # resolveCompactTriggerAt (agent-cli src/node/context/context-protocol.ts): + # modelMaxInputTokens ? modelMaxInputTokens * percentThreshold + # : getAutoCompactWindow() + # Without it the CLI falls back to CODEBUDDY_AUTO_COMPACT_WINDOW, which is + # clamped to [100k, 1M] -- above our 96k context cap, so compaction could + # never fire in time. With it, CODEBUDDY_AUTOCOMPACT_PCT_OVERRIDE becomes a + # percentage OF THIS VALUE and there is no clamp. The launcher passes the + # run's rollout_max_context_len here. + max_input_tokens = _positive_int_env(self.max_input_tokens_env) + if max_input_tokens: + model_entry["maxInputTokens"] = max_input_tokens + + models_json = { + "models": [model_entry], + "availableModels": [ctx.model_label], + } + settings_json = { + "cleanupPeriodDays": 30, + "includeCoAuthoredBy": False, + "autoCompactEnabled": True, + # Reasoning on: the trained policy emits blocks, and the + # adapter records them, so disabling it would train on a different + # distribution than it serves. + "alwaysThinkingEnabled": True, + "showTokensCounter": False, + "enablePasteImageFromClipboard": False, + "enableTerminalProgressBar": False, + "fileCheckpointingEnabled": False, + "promptSuggestionEnabled": False, + "enableAllProjectMcpServers": False, + "memory": {"autoMemoryEnabled": False}, + } + models_b64 = _json_b64(models_json) + settings_b64 = _json_b64(settings_json) + await sb.exec( + "set -e\n" + "mkdir -p /root/.codebuddy/debug /root/.codebuddy/projects /root/.codebuddy/statsig " + "/home/agent/.codebuddy/debug /home/agent/.codebuddy/projects /home/agent/.codebuddy/statsig\n" + f"printf %s {shlex.quote(models_b64)} | base64 -d " + "| tee /root/.codebuddy/models.json /home/agent/.codebuddy/models.json >/dev/null\n" + f"printf %s {shlex.quote(settings_b64)} | base64 -d " + "| tee /root/.codebuddy/settings.json /home/agent/.codebuddy/settings.json >/dev/null\n" + "chown -R agent:agent /home/agent/.codebuddy", + user="root", + check=True, + timeout=60, + ) + + async def launch_and_wait(self, sb: Sandbox, ctx: HarnessContext, prompt: str, time_budget_sec: int) -> int: + parts: list[str] = [ + f"--model {shlex.quote(ctx.model_label)}", + "--verbose", + "--output-format stream-json", + "--include-partial-messages", + *self.default_flags, + ] + extra = os.environ.get(self.extra_args_env, "").strip() + if extra: + # After default_flags so it can override them, before the prompt so the + # prompt stays positional. + parts.append(extra) + parts.append("-y") + + session_log_dir = f"{ctx.workdir}/.harness/codebuddy_sessions" + raw_cmd = f"cbc {' '.join(parts)} {shlex.quote(prompt)}; rc=$?; mkdir -p {shlex.quote(session_log_dir)}/projects; cp -r /root/.codebuddy/projects/. {shlex.quote(session_log_dir)}/projects/ 2>/dev/null || true; exit $rc" + cmd = f"bash -lc {shlex.quote(raw_cmd)}" + + env = { + "OPENAI_API_KEY": ctx.session_id, + "OPENAI_BASE_URL": f"{ctx.adapter_url}/v1", + "CBC_API_KEY": ctx.session_id, + "CBC_BASE_URL": self._chat_completions_url(ctx.adapter_url), + "NO_COLOR": "1", + "CI": "1", + "TERM": "dumb", + "IS_SANDBOX": "1", + } + extra_envs = os.environ.get(self.extra_envs_env, "").strip() + if extra_envs: + env.update(json.loads(extra_envs)) + return await run_root_command( + sb, + workdir=ctx.workdir, + start_cmd=cmd, + env=env, + time_budget_sec=time_budget_sec, + ) + + async def run( + self, + sb: Sandbox, + *, + workdir: str, + session_id: str, + adapter_url: str, + time_budget_sec: int, + prompt: str, + ) -> int: + from slime.agent import sandbox as agent_sandbox + + await agent_sandbox.ensure_agent_user(sb, workdir) + ctx = HarnessContext(workdir=workdir, session_id=session_id, adapter_url=adapter_url) + await self.write_config(sb, ctx) + return await self.launch_and_wait(sb, ctx, prompt, time_budget_sec) + + @staticmethod + def _chat_completions_url(adapter_url: str) -> str: + url = adapter_url.rstrip("/") + if url.endswith("/chat/completions"): + return url + if url.endswith("/v1"): + return f"{url}/chat/completions" + return f"{url}/v1/chat/completions" + + +def _positive_int_env(name: str) -> int | None: + """Read a positive int from the environment, ignoring unset/blank/invalid.""" + raw = (os.environ.get(name) or "").strip() + if not raw: + return None + try: + value = int(raw) + except ValueError: + return None + return value if value > 0 else None + + +def _json_b64(value: dict) -> str: + payload = json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + return base64.b64encode(payload).decode("ascii") + + +HARNESS_REGISTRY: dict[str, tuple[type[BaseHarness], type]] = { + "claude_code": (AGSSidecarClaudeCodeHarness, AnthropicAdapter), + "codex": (CodexHarness, OpenAIAdapter), + "codebuddy_code": (CodeBuddyCodeHarness, OpenAIAdapter), +} + + +def resolve_agent(agent_name: str) -> tuple[type[BaseHarness], type]: + if agent_name not in HARNESS_REGISTRY: + raise ValueError(f"SWE_AGENT={agent_name!r} not in {sorted(HARNESS_REGISTRY)}") + return HARNESS_REGISTRY[agent_name] diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py new file mode 100644 index 0000000000..49913bdbff --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py @@ -0,0 +1,398 @@ +"""Per-sample AGS coding-agent rollout implementation.""" + +from __future__ import annotations + +import asyncio +import copy +import logging +import secrets +import time +import traceback +from argparse import Namespace +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +from slime.utils.types import Sample + +from .adapter_service import get_adapter_service +from .ags_sandbox import AGSSandbox +from .artifacts import ArtifactWriter, sample_artifact_id +from .config import AGSGeneratorConfig +from .empty_patch_guard import GuardVerdict, classify_empty_patch +from .harnesses import resolve_agent +from .sampling import normalize_sampling_params +from .swe_task import evaluate, get_metadata, git_diff, prepare_workspace +from .weave_trace import AGSWeaveTrace + +logger = logging.getLogger(__name__) + + +class AGSRolloutRunner: + def __init__( + self, + args: Namespace, + config: AGSGeneratorConfig | None = None, + *, + use_remote_adapter: bool = False, + ) -> None: + self.args = args + self.config = config or AGSGeneratorConfig.from_env() + self.harness_cls, self.adapter_cls = resolve_agent(self.config.agent_name) + # use_remote_adapter asks to reuse the training adapter rather than bind + # a second one on the same port. get_adapter_service downgrades that to a + # local adapter when no training adapter is actually running, so + # eval-before-train and --num-rollout 0 work instead of failing every + # prompt with a connection error. + self.adapter_service = get_adapter_service( + args, + self.config, + self.adapter_cls, + evaluation=use_remote_adapter, + ) + self.artifacts = ArtifactWriter(self.config.artifact_dir) + self.weave_trace = AGSWeaveTrace( + args, + self.adapter_service.tokenizer, + agent=self.config.agent_name, + enable_token2text=self.config.enable_token2text, + ) + self._boot_sem = asyncio.Semaphore(self.config.boot_concurrency) + + async def generate(self, base_sample: Sample, sampling_params: dict, *, evaluation: bool = False) -> list[Sample]: + md = get_metadata(base_sample) + instance_id = md["instance_id"] + base_sample = copy.deepcopy(base_sample) + session_id = _session_id(base_sample, instance_id) + base_sample.session_id = session_id + artifact_id = sample_artifact_id(instance_id, base_sample) + normalized_sampling = normalize_sampling_params(sampling_params) + trace_call = self.weave_trace.start_rollout( + instance_id=instance_id, + session_id=session_id, + sample=base_sample, + sampling_params=normalized_sampling, + agent=self.config.agent_name, + ) + if not md["image"] or not md["workdir"]: + samples = self._abort_result(base_sample, "missing_image_or_workdir", instance_id) + self.weave_trace.finish_rollout(trace_call, samples=samples) + return samples + t0 = time.time() + session_opened = False + trajectory_path = None + evaluation_args = { + "image": md["image"], + "workdir": md["workdir"], + "swepro": md["swepro"], + "eval_cmd": md["eval_cmd"], + "f2p_script": md["f2p_script"], + "pre_commands": md["pre_commands"], + "eval_bootstrap_cmd": self.config.eval_bootstrap_cmd, + "timeout_sec": self.config.eval_timeout_sec, + } + try: + self.adapter_service.adapter.open_session( + session_id, + sampling_defaults=normalized_sampling, + max_context_tokens=self.adapter_service.max_context_len, + ) + session_opened = True + async with asyncio.timeout(self.config.rollout_guard_sec): + async with self._boot_agent_sandbox(md["image"], instance_id) as sb: + prompt_style = self.config.prompt_style_for(evaluation=evaluation) + await prepare_workspace( + sb, + md["workdir"], + md, + write_problem_statement=prompt_style == "instruction", + ) + agent_exit_code = await self.harness_cls().run( + sb, + workdir=md["workdir"], + session_id=session_id, + adapter_url=self.adapter_service.adapter_url, + time_budget_sec=self.config.agent_time_budget_sec, + prompt=self._agent_prompt(md, prompt_style), + ) + trajectory_path = await self.artifacts.dump_trajectory(sb, md["workdir"], artifact_id) + diff_text = await git_diff(sb, md["workdir"]) + guard_verdict = self._check_empty_patch( + agent_exit_code=agent_exit_code, + diff_text=diff_text, + trajectory_path=trajectory_path, + instance_id=instance_id, + ) + patch_path = self.artifacts.dump_patch(diff_text, artifact_id) + if not self.config.eval_isolated_sandbox: + reward, applied_cleanly = await evaluate( + sandbox=sb, + diff_text=diff_text, + **evaluation_args, + ) + + if self.config.eval_isolated_sandbox: + reward, applied_cleanly = await evaluate( + sandbox=None, + diff_text=diff_text, + **evaluation_args, + ) + # Must come before finish_session: that call drains the session + # and is idempotent, so a later return would yield no samples. + if guard_verdict.triggered and self.config.empty_patch_guard == "abort": + samples = self._abort_result(base_sample, guard_verdict.reason, instance_id) + self.weave_trace.finish_rollout(trace_call, samples=samples, trajectory_path=trajectory_path) + return samples + + samples = await self.adapter_service.adapter.finish_session( + session_id, + base_sample=base_sample, + reward=float(reward), + extra_metadata={ + "grading_solved": float(reward) == 1.0, + "instance_id": instance_id, + }, + ) + if not samples: + samples = self._abort_result(base_sample, "adapter_session_empty", instance_id) + self.weave_trace.finish_rollout(trace_call, samples=samples, trajectory_path=trajectory_path) + return samples + + rollout_path = self.artifacts.dump_rollout( + { + "instance_id": instance_id, + "session_id": session_id, + "agent": self.config.agent_name, + "reward": float(reward), + "applied_cleanly": bool(applied_cleanly), + "eval_isolated_sandbox": self.config.eval_isolated_sandbox, + "agent_exit_code": agent_exit_code, + "elapsed_sec": time.time() - t0, + "num_samples": len(samples), + "patch_path": patch_path, + "trajectory_path": trajectory_path, + "empty_patch_guard_triggered": guard_verdict.triggered, + "empty_patch_guard_reason": guard_verdict.reason, + # Excerpt goes in the dump only; sample metadata stays small. + "empty_patch_final_text": guard_verdict.final_text, + }, + artifact_id, + ) + elapsed_sec = time.time() - t0 + for sample in samples: + sample.metadata = { + **(sample.metadata or {}), + "agent": self.config.agent_name, + "agent_exit_code": agent_exit_code, + "applied_cleanly": bool(applied_cleanly), + "eval_isolated_sandbox": self.config.eval_isolated_sandbox, + "trajectory_path": trajectory_path, + "patch_path": patch_path, + "rollout_dump_path": rollout_path, + "ags_elapsed_sec": elapsed_sec, + "ags_num_samples": len(samples), + "ags_rollout_concurrency": self.config.rollout_concurrency, + "empty_patch_guard_triggered": guard_verdict.triggered, + "empty_patch_guard_reason": guard_verdict.reason, + } + logger.info( + "[ags_generator] %s: reward=%.2f applied=%s eval_isolated=%s exit=%s elapsed=%.1fs segments=%d", + instance_id, + float(reward), + bool(applied_cleanly), + self.config.eval_isolated_sandbox, + agent_exit_code, + elapsed_sec, + len(samples), + ) + self.weave_trace.finish_rollout( + trace_call, + samples=samples, + trajectory_path=trajectory_path, + output={ + "reward": float(reward), + "applied_cleanly": bool(applied_cleanly), + "agent_exit_code": agent_exit_code, + "elapsed_sec": elapsed_sec, + "patch_path": patch_path, + "rollout_dump_path": rollout_path, + }, + ) + return samples + except asyncio.TimeoutError: + _log_timeout_diagnostic(t0, instance_id, self.config.rollout_guard_sec) + samples = self._abort_result(base_sample, "wall_clock_timeout", instance_id) + self.weave_trace.finish_rollout(trace_call, samples=samples, trajectory_path=trajectory_path) + return samples + except Exception as exc: + logger.warning("[ags_generator] %s: rollout failed: %s\n%s", instance_id, exc, traceback.format_exc()) + samples = self._abort_result(base_sample, f"exception:{type(exc).__name__}", instance_id) + self.weave_trace.finish_rollout( + trace_call, + samples=samples, + trajectory_path=trajectory_path, + exception=exc, + ) + return samples + finally: + if session_opened: + try: + await self.adapter_service.adapter.drop_session(session_id) + except Exception: + logger.warning( + "[ags_generator] %s: failed to drop session %s\n%s", + instance_id, + session_id, + traceback.format_exc(), + ) + + def _agent_prompt(self, md: dict, prompt_style: str) -> str: + """Build the prompt handed to the coding agent. + + "dataset" forwards the row's own prompt untouched, which for converted + Harbor data is instruction.md byte for byte -- the same string Harbor hands + its agents. "instruction" instead sends SWE_CC_PROMPT, which tells the + agent to go read PROBLEM_STATEMENT.md and costs it turns before it even + knows the task. The caller picks the style per rollout (training vs eval); + see AGSGeneratorConfig.prompt_style_for. + + An empty prompt field falls back to SWE_CC_PROMPT rather than sending the + agent nothing, but prepare_workspace only writes PROBLEM_STATEMENT.md for + the "instruction" style -- so under this fallback the file the prompt + names is absent. It should not happen (the converter always writes a + prompt), hence the warning. + """ + if prompt_style == "instruction": + return self.config.prompt + + text = (md.get("dataset_prompt") or "").strip() + if not text: + logger.warning( + "[ags_generator] %s: prompt_style=dataset but the row's prompt is empty; " + "falling back to SWE_CC_PROMPT, whose PROBLEM_STATEMENT.md was not written", + md.get("instance_id"), + ) + return self.config.prompt + return text + + def _check_empty_patch( + self, + *, + agent_exit_code: int, + diff_text: str, + trajectory_path: str | None, + instance_id: str, + ) -> GuardVerdict: + """Classify an exit-0-but-no-diff rollout; never raises. + + Classification failure must not cost us an otherwise usable rollout, so + anything unexpected degrades to "not triggered" rather than propagating + into the caller's generic exception handler (which would abort). + """ + if self.config.empty_patch_guard == "off": + return GuardVerdict(triggered=False) + try: + verdict = classify_empty_patch( + agent_exit_code=agent_exit_code, + diff_text=diff_text, + trajectory_path=trajectory_path, + agent=self.config.agent_name, + ) + except Exception: + logger.warning("[ags_generator] %s: empty patch guard failed", instance_id, exc_info=True) + return GuardVerdict(triggered=False) + if verdict.triggered: + logger.warning( + "[ags_generator] %s: empty patch guard: %s (policy=%s)", + instance_id, + verdict.reason, + self.config.empty_patch_guard, + ) + return verdict + + @asynccontextmanager + async def _boot_agent_sandbox(self, image: str, instance_id: str) -> AsyncIterator[AGSSandbox]: + sb = None + last_err: Exception | None = None + for attempt in range(self.config.boot_retries): + cand = AGSSandbox(image) + try: + async with self._boot_sem: + await cand.__aenter__() + try: + await self.harness_cls().install_cli(cand) + except BaseException: + await cand.__aexit__(None, None, None) + raise + sb = cand + break + except Exception as exc: + last_err = exc + logger.warning( + "[ags_generator] %s: AGS provision attempt %d/%d failed: %s: %s", + instance_id, + attempt + 1, + self.config.boot_retries, + type(exc).__name__, + str(exc)[:200], + ) + await asyncio.sleep(1 + attempt) + if sb is None: + assert last_err is not None + raise last_err + try: + yield sb + finally: + await sb.__aexit__(None, None, None) + + def _abort_result(self, sample: Sample, reason: str, instance_id: str) -> list[Sample]: + sample = copy.deepcopy(sample) + sample.tokens = [0, 0] + sample.response = "" + sample.response_length = 1 + sample.loss_mask = [0] + sample.rollout_log_probs = [0.0] + sample.reward = 0.0 + sample.remove_sample = True + sample.status = Sample.Status.ABORTED + sample.metadata = {**(sample.metadata or {}), "abort_reason": reason, "instance_id": instance_id} + logger.warning("[ags_generator] %s aborted: %s", instance_id, reason) + return [sample] + + +def _session_id(sample: Sample, instance_id: str) -> str: + """Return a fresh adapter session id for one AGS attempt. + + RolloutDataSource can hand out deep copies of Samples that already carry a + session_id, and failed/partial reruns can also revisit the same + (instance_id, index, group_index). Adapter sessions are process-global for + one generator run, so every AGS attempt must get a unique id instead of + reusing the sample's existing session_id. + """ + + parts = ["cagent", instance_id] + if sample.index is not None: + parts.append(str(sample.index)) + if sample.group_index is not None: + parts.append(str(sample.group_index)) + parts.append(secrets.token_hex(4)) + return "-".join(parts) + + +def _log_timeout_diagnostic(t0: float, instance_id: str, guard_sec: int) -> None: + try: + elapsed = time.time() - t0 + pending = [task for task in asyncio.all_tasks() if not task.done()] + stuck = [] + for task in pending[:5]: + coro = getattr(task, "_coro", None) + stuck.append(getattr(coro, "__qualname__", repr(coro))) + logger.warning( + "[ags_generator] %s: wall_clock_timeout after %.1fs (guard=%ds); %d tasks pending; sample=%s", + instance_id, + elapsed, + guard_sec, + len(pending), + stuck, + ) + except Exception: + pass diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/runner.py b/slime_plugins/rollout_buffer/generator/ags_generator/runner.py new file mode 100644 index 0000000000..594cc5d26a --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/runner.py @@ -0,0 +1,64 @@ +"""Command-running helpers for AGS sidecar harnesses.""" + +from __future__ import annotations + +import asyncio +import logging +import shlex +import time + +from slime.agent.sandbox import EXIT_TIME_BUDGET_EXCEEDED, Sandbox, terminate_process_group + +logger = logging.getLogger(__name__) + + +async def run_root_command( + sb: Sandbox, + *, + workdir: str, + start_cmd: str, + env: dict[str, str], + time_budget_sec: int, +) -> int: + """Run an AGS sidecar command as root and persist its stream-json output.""" + + meta_dir = f"{workdir}/.harness" + done = f"{meta_dir}/done" + launcher = f"{meta_dir}/run.sh" + traj = f"{meta_dir}/trajectory.jsonl" + pid_file = f"{meta_dir}/pid" + lock_dir = f"{meta_dir}/spawned" + launcher_body = ( + "#!/bin/bash\n" + f"cd {workdir}\n" + "export HOME=/root\n" + f"{start_cmd} 2>&1 | tee {shlex.quote(traj)}\n" + f"echo ${{PIPESTATUS[0]}} > {done}\n" + ) + await sb.exec(f"mkdir -p {meta_dir} && chmod 777 {meta_dir}", user="root", check=True, timeout=30) + await sb.write_file(launcher, launcher_body, user="root") + await sb.exec(f"chmod +x {launcher}", user="root", timeout=30, check=True) + + export_lines = " ".join(f"{k}={shlex.quote(str(v))}" for k, v in env.items()) + await sb.exec( + f"mkdir {lock_dir} 2>/dev/null || exit 0; " + f"rm -f {done} {pid_file}; " + f"env {export_lines} setsid {launcher} < /dev/null > /dev/null 2>&1 & " + f"echo $! > {pid_file}", + user="root", + timeout=30, + check=True, + ) + + deadline = time.time() + time_budget_sec + exit_code = EXIT_TIME_BUDGET_EXCEEDED + while time.time() < deadline: + await asyncio.sleep(5) + ec, out, _ = await sb.exec(f"test -f {done} && cat {done}", user="root", timeout=15, check=False) + if ec == 0 and (out or "").strip(): + exit_code = int((out or "").strip()) + break + if exit_code == EXIT_TIME_BUDGET_EXCEEDED: + logger.warning("AGS agent exceeded %ss; terminating process group", time_budget_sec) + await terminate_process_group(sb, pid_file=pid_file, user="root") + return exit_code diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/sampling.py b/slime_plugins/rollout_buffer/generator/ags_generator/sampling.py new file mode 100644 index 0000000000..a7f40d4218 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/sampling.py @@ -0,0 +1,21 @@ +"""Sampling-parameter normalization for OpenAI-style agent adapters.""" + +from __future__ import annotations + +from typing import Any + + +def normalize_sampling_params(params: dict[str, Any]) -> dict[str, Any]: + out = dict(params or {}) + if "max_tokens" in out and "max_new_tokens" not in out: + out["max_new_tokens"] = out.pop("max_tokens") + else: + out.pop("max_tokens", None) + if "max_response_len" in out and "max_new_tokens" not in out: + out["max_new_tokens"] = out.pop("max_response_len") + else: + out.pop("max_response_len", None) + out.pop("skip_special_tokens", None) + out.pop("no_stop_trim", None) + out.pop("spaces_between_special_tokens", None) + return out diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/serialization.py b/slime_plugins/rollout_buffer/generator/ags_generator/serialization.py new file mode 100644 index 0000000000..d4545d3545 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/serialization.py @@ -0,0 +1,43 @@ +"""JSON-safe conversion between rollout-buffer payloads and slime Samples.""" + +from __future__ import annotations + +from typing import Any + +from slime.utils.types import Sample + + +SAMPLE_MARKER = "sample_dict_v1" + + +def sample_to_payload(sample: Sample) -> dict[str, Any]: + return {"__type__": SAMPLE_MARKER, "sample": sample.to_dict()} + + +def samples_from_payload(payload: dict[str, Any]) -> list[Sample]: + if "samples" in payload and isinstance(payload["samples"], list): + return [Sample.from_dict(item) for item in payload["samples"]] + return [sample_from_payload(payload)] + + +def sample_from_payload(payload: dict[str, Any]) -> Sample: + if payload.get("__type__") == SAMPLE_MARKER and "sample" in payload: + return Sample.from_dict(payload["sample"]) + if "sample" in payload and isinstance(payload["sample"], dict): + return Sample.from_dict(payload["sample"]) + return Sample.from_dict(payload) + + +def output_item_from_samples( + samples: list[Sample], *, instance_id: str, extra_info: dict[str, Any] | None = None +) -> dict[str, Any]: + first = samples[0] + return { + "uid": first.session_id or f"sample-{first.index}", + "instance_id": instance_id, + "messages": [], + "reward": first.reward, + "extra_info": extra_info or {}, + "samples": [sample.to_dict() for sample in samples], + "__type__": SAMPLE_MARKER, + } diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/source.py b/slime_plugins/rollout_buffer/generator/ags_generator/source.py new file mode 100644 index 0000000000..7579fe8615 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/source.py @@ -0,0 +1,68 @@ +"""Prompt-data source used by the standalone rollout-buffer AGS generator.""" + +from __future__ import annotations + +import copy +from argparse import Namespace + +from slime.rollout.data_source import RolloutDataSource +from slime.utils.types import Sample + + +class AGSPromptSource: + """Small wrapper around RolloutDataSource that yields one repeat at a time.""" + + def __init__(self, args: Namespace) -> None: + self.args = args + self.data_source = RolloutDataSource(args) + start_group = int(getattr(args, "rollout_start_group", 0) or 0) + dataset = self.data_source.dataset + if start_group > 0 and dataset is not None and len(dataset) > 0: + # There is no generator-side checkpoint: the trainer identifies the + # position by absolute group index (rollout_id * rollout_batch_size), + # so seek to it. Group index g lives at offset g % len in epoch + # g // len, mirroring how RolloutDataSource.get_samples advances the + # offset by one per group and bumps the epoch on each wraparound. + dataset_len = len(dataset) + self.data_source.sample_offset = start_group % dataset_len + self.data_source.epoch_id = start_group // dataset_len + self.data_source.sample_group_index = start_group + self.data_source.sample_index = start_group * int(args.n_samples_per_prompt) + if args.rollout_shuffle: + # Each epoch has its own seeded permutation; without this the + # samples stay on epoch 0's order after a wraparound and the + # prompt sequence diverges from the uninterrupted run. + dataset.shuffle(self.data_source.epoch_id) + + def get_groups(self, num_groups: int) -> list[list[Sample]]: + groups = self.data_source.get_samples(num_groups) + for group in groups: + for sample in group: + if sample.rollout_id is None: + sample.rollout_id = sample.index + return groups + + def get_repeated_samples(self, num_groups: int, skip_instance_ids: list[str] | None = None) -> list[Sample]: + skip = list(skip_instance_ids or []) + samples: list[Sample] = [] + while len(samples) < num_groups * self.args.n_samples_per_prompt: + groups = self.get_groups(num_groups) + for group in groups: + instance_id = _instance_id(group[0]) + for sample in group: + if instance_id in skip: + skip.remove(instance_id) + continue + samples.append(copy.deepcopy(sample)) + if len(samples) >= num_groups * self.args.n_samples_per_prompt: + break + if len(samples) >= num_groups * self.args.n_samples_per_prompt: + break + return samples + + +def _instance_id(sample: Sample) -> str: + metadata = sample.metadata or {} + remote = metadata.get("remote_env_info") or {} + label = sample.label if isinstance(sample.label, str) and len(sample.label) < 256 else None + return str(metadata.get("instance_id") or remote.get("instance_id") or label or sample.index or "unknown") diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py b/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py new file mode 100644 index 0000000000..b4ff2d87fd --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py @@ -0,0 +1,272 @@ +"""SWE task operations used by the AGS rollout-buffer generator.""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Any + +from slime.agent import sandbox as agent_sandbox +from slime.agent.adapters.common import flatten_content +from slime.agent.sandbox import Sandbox +from slime.utils.types import Sample + +from .ags_sandbox import AGSSandbox + +logger = logging.getLogger(__name__) + +_PATCH = "/workspace/__cagent_patch__.diff" +_PRE = "/workspace/__cagent_pre__.sh" +_F2P = "/workspace/__cagent_f2p__.py" +_SWEPRO_DIR = "/workspace/swepro_eval" + + +def get_metadata(sample: Sample) -> dict[str, Any]: + m = sample.metadata or {} + rem = m.get("remote_env_info") or {} + label = sample.label if (isinstance(sample.label, str) and len(sample.label) < 256) else None + return { + "instance_id": m.get("instance_id") or rem.get("instance_id") or label or "unknown", + "image": m.get("image") or rem.get("image_url"), + "workdir": m.get("workdir") or rem.get("workdir"), + "problem_statement": m.get("problem_statement") or _coerce_prompt(sample.prompt), + # The row's own prompt text. Kept separate from problem_statement so the + # "dataset" prompt style forwards exactly what the data carries, even for + # rows whose metadata.problem_statement was written by some other producer. + "dataset_prompt": _coerce_prompt(sample.prompt), + "swepro": m.get("swepro"), + "eval_cmd": m.get("eval_cmd"), + "f2p_script": rem.get("f2p_script"), + "pre_commands": m.get("pre_commands") or rem.get("pre_commands"), + } + + +def _coerce_prompt(prompt) -> str: + if isinstance(prompt, str): + return prompt + if isinstance(prompt, list): + for message in prompt: + if isinstance(message, dict) and message.get("role") == "user": + return flatten_content(message.get("content")) + return "" + + +async def prepare_workspace( + sb: Sandbox, workdir: str, md: dict[str, Any], *, write_problem_statement: bool = True +) -> None: + """Set up the agent's workspace before the harness runs. + + ``write_problem_statement`` is False under the "dataset" prompt style, whose + prompt already carries the task text: writing the file anyway would leave an + untracked artifact in the repo that only ``git_diff``'s exclude pathspec keeps + out of the patch, and would tempt the agent into spending turns reading a file + that merely restates its prompt. Harbor writes no such file either. + """ + await agent_sandbox.ensure_agent_user(sb, workdir) + swepro = md.get("swepro") + if swepro: + await apply_before_repo_set_cmd(sb, workdir, swepro) + pre_commands = md.get("pre_commands") + if pre_commands: + await apply_pre_commands(sb, workdir, pre_commands) + if write_problem_statement: + await sb.write_file(f"{workdir}/PROBLEM_STATEMENT.md", md.get("problem_statement") or "", user="agent") + + +async def apply_before_repo_set_cmd(sb: Sandbox, workdir: str, swepro: dict[str, Any]) -> None: + before = swepro.get("before_repo_set_cmd") + if not before: + return + payload = f"set -e\ncd {workdir}\n{before}\n" + await sb.exec( + "mkdir -p /workspace/swepro_setup && chown agent:agent /workspace/swepro_setup", user="root", check=True + ) + await sb.write_file("/workspace/swepro_setup/before.sh", payload, user="agent") + await sb.exec("bash /workspace/swepro_setup/before.sh", user="agent", check=False, timeout=600) + + +async def apply_pre_commands(sb: Sandbox, workdir: str, pre: list[str] | str) -> None: + body = pre.replace("\\n", "\n") if isinstance(pre, str) else "\n".join(c for c in (pre or []) if c) + await sb.write_file(_PRE, "set -e\n" + body, user="agent") + await sb.exec(f"chmod 755 {_PRE} && cd {workdir} && bash {_PRE}", user="agent", check=False, timeout=600) + + +async def git_diff(sb: Sandbox, workdir: str) -> str: + cmd = f"cd {workdir} && git add -N . && git diff -- . ':(exclude)PROBLEM_STATEMENT.md' ':(exclude).harness/'" + _, out, _ = await sb.exec(cmd, user="agent", timeout=120) + return out + + +async def evaluate( + *, + sandbox: Sandbox | None = None, + image: str, + workdir: str, + diff_text: str, + swepro: dict[str, Any] | None = None, + eval_cmd: str | None = None, + f2p_script: str | None = None, + pre_commands: list[str] | str | None = None, + eval_bootstrap_cmd: str | None = None, + timeout_sec: int = 600, +) -> tuple[float, bool]: + if not (swepro or eval_cmd or f2p_script): + logger.warning("[ags_generator.evaluate] no swepro/eval_cmd/f2p_script; reward=0") + return 0.0, True + + if sandbox is not None: + return await _evaluate_in_sandbox( + sandbox, + workdir=workdir, + diff_text=diff_text, + swepro=swepro, + eval_cmd=eval_cmd, + f2p_script=f2p_script, + pre_commands=pre_commands, + eval_bootstrap_cmd=eval_bootstrap_cmd, + timeout_sec=timeout_sec, + isolated=False, + ) + + async with AGSSandbox(image) as ev: + return await _evaluate_in_sandbox( + ev, + workdir=workdir, + diff_text=diff_text, + swepro=swepro, + eval_cmd=eval_cmd, + f2p_script=f2p_script, + pre_commands=pre_commands, + eval_bootstrap_cmd=eval_bootstrap_cmd, + timeout_sec=timeout_sec, + isolated=True, + ) + + +async def _evaluate_in_sandbox( + ev: Sandbox, + *, + workdir: str, + diff_text: str, + swepro: dict[str, Any] | None, + eval_cmd: str | None, + f2p_script: str | None, + pre_commands: list[str] | str | None, + eval_bootstrap_cmd: str | None, + timeout_sec: int, + isolated: bool, +) -> tuple[float, bool]: + if isolated: + await agent_sandbox.ensure_agent_user(ev, workdir) + if swepro: + await _setup_swepro_assets(ev, swepro) + if isolated: + await apply_before_repo_set_cmd(ev, workdir, swepro) + if isolated: + if pre_commands: + await apply_pre_commands(ev, workdir, pre_commands) + if eval_bootstrap_cmd: + await _run_eval_bootstrap(ev, workdir, eval_bootstrap_cmd, timeout=min(600, max(120, timeout_sec))) + + if isolated: + applied = await _apply_diff(ev, workdir, diff_text) + if not applied: + return 0.0, False + + if swepro: + reward = await _run_swepro(ev, workdir, swepro, timeout_sec) + elif eval_cmd: + reward = await _run_eval_cmd(ev, workdir, eval_cmd, timeout_sec) + else: + reward = await _run_f2p_script(ev, workdir, f2p_script or "", timeout_sec) + return reward, True + + +async def _setup_swepro_assets(ev: Sandbox, swepro: dict[str, Any]) -> None: + await ev.exec(f"mkdir -p {_SWEPRO_DIR} && chmod 777 {_SWEPRO_DIR}", user="root", check=True) + for key, dst in [("run_script_path", "run_script.sh"), ("parser_script_path", "parser.py")]: + host_path = swepro.get(key) + if host_path: + await ev.write_file(f"{_SWEPRO_DIR}/{dst}", Path(host_path), user="root") + await ev.exec(f"chmod 755 {_SWEPRO_DIR}/* && chown -R agent:agent {_SWEPRO_DIR}", user="root", check=True) + + +async def _apply_diff(ev: Sandbox, workdir: str, diff_text: str) -> bool: + if not diff_text.strip(): + return True + await ev.write_file(_PATCH, diff_text, user="agent") + for cmd in [ + f"cd {workdir} && git apply --3way --whitespace=nowarn {_PATCH}", + f"cd {workdir} && git apply --whitespace=nowarn {_PATCH}", + f"cd {workdir} && patch -p1 --no-backup-if-mismatch < {_PATCH}", + ]: + ec, _, _ = await ev.exec(cmd, user="agent", check=False, timeout=120) + if ec == 0: + return True + return False + + +async def _run_swepro(ev: Sandbox, workdir: str, swepro: dict[str, Any], timeout: int) -> float: + test_arg = ",".join(swepro.get("selected_test_files") or []) + stdout_f = f"{_SWEPRO_DIR}/stdout.log" + stderr_f = f"{_SWEPRO_DIR}/stderr.log" + result_f = f"{_SWEPRO_DIR}/result.json" + await ev.exec( + f"cd {workdir} && bash {_SWEPRO_DIR}/run_script.sh {json.dumps(test_arg)} > {stdout_f} 2> {stderr_f} || true", + user="agent", + check=False, + timeout=timeout, + ) + await ev.exec( + f"python3 {_SWEPRO_DIR}/parser.py {stdout_f} {stderr_f} {result_f}", user="agent", check=False, timeout=120 + ) + raw = await ev.read_file(result_f, user="agent") + parsed = json.loads(raw) if raw else {"tests": []} + passed = {t["name"] for t in parsed.get("tests", []) if t.get("status") == "PASSED"} + required = set(swepro.get("fail_to_pass") or []) | set(swepro.get("pass_to_pass") or []) + solved = bool(required) and required.issubset(passed) + return 1.0 if solved else 0.0 + + +async def _run_eval_cmd(ev: Sandbox, workdir: str, cmd: str, timeout: int) -> float: + # Run eval_cmd as root. Harbor-derived SWE-bench verifier scripts preserve + # Harbor's canonical container layout and intentionally materialize + # /tests/config.json, /tests/test.sh, /logs/verifier, and a parser next to + # /testbed (via `cd ..`). The verified SWE-bench images do not pre-create + # those root-owned paths for the unprivileged agent user, and we want to run + # Harbor's test.sh verbatim rather than patching its paths. Evaluation starts + # only after the agent process exits; depending on SWE_EVAL_ISOLATED_SANDBOX, + # it runs either in that agent sandbox or in a separate clean sandbox. + ec, _, _ = await ev.exec(f"cd {workdir} && {cmd}", user="root", check=False, timeout=timeout) + return 1.0 if ec == 0 else 0.0 + + +async def _run_eval_bootstrap(ev: Sandbox, workdir: str, cmd: str, timeout: int) -> None: + ec, out, err = await ev.exec(f"cd {workdir} && {cmd}", user="agent", check=False, timeout=timeout) + logger.info( + "[ags_generator.evaluate] bootstrap exit=%s stdout_tail=%r stderr_tail=%r", + ec, + (out or "")[-2000:], + (err or "")[-2000:], + ) + + +async def _run_f2p_script(ev: Sandbox, workdir: str, script: str, timeout: int) -> float: + await ev.write_file(_F2P, script, user="agent") + ec, out, err = await ev.exec( + f"cd {workdir} && export PATH=/opt/conda/bin:/usr/local/bin:$PATH; " + f"if [ -x /opt/conda/bin/python ]; then /opt/conda/bin/python {_F2P}; " + f"elif command -v python >/dev/null 2>&1; then python {_F2P}; " + f"else python3 {_F2P}; fi", + user="agent", + check=False, + timeout=timeout, + ) + logger.info( + "[ags_generator.evaluate] f2p exit=%s stdout_tail=%r stderr_tail=%r", + ec, + (out or "")[-4000:], + (err or "")[-4000:], + ) + return 1.0 if ec == 0 else 0.0 diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/wandb_metrics.py b/slime_plugins/rollout_buffer/generator/ags_generator/wandb_metrics.py new file mode 100644 index 0000000000..77e5f1f782 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/wandb_metrics.py @@ -0,0 +1,396 @@ +"""Minimal AGS-specific rollout/eval logging hooks for slime. + +Wire these through:: + + --custom-rollout-log-function-path \ + slime_plugins.rollout_buffer.generator.ags_generator.wandb_metrics.log_rollout_data + --custom-eval-rollout-log-function-path \ + slime_plugins.rollout_buffer.generator.ags_generator.wandb_metrics.log_eval_rollout_data + +The hooks intentionally keep slime's default rollout/eval metrics and add a +small AGS block derived from ``Sample.metadata`` produced by +``ags_generator.rollout.AGSRolloutRunner``. +""" + +from __future__ import annotations + +import logging +import re +from collections import Counter, defaultdict +from collections.abc import Iterable +from typing import Any + +import numpy as np + +from slime.utils import logging_utils +from slime.utils.metric_utils import compute_pass_rate, compute_rollout_step, compute_statistics, dict_add_prefix +from slime.utils.types import Sample + +logger = logging.getLogger(__name__) + +_BUCKET_RE = re.compile(r"[^A-Za-z0-9_.-]+") + + +def log_rollout_data(rollout_id: int, args: Any, samples: list[Sample], rollout_extra_metrics, rollout_time) -> bool: + """Log default rollout metrics plus a compact AGS diagnostics block. + + Return ``True`` after logging so slime's default logger is not called a + second time. If anything goes wrong, return ``False`` and let the default + logger handle the batch. + """ + + try: + if getattr(args, "load_debug_rollout_data", None): + return False + + log_dict = {**(rollout_extra_metrics or {})} + log_dict |= dict_add_prefix(_compute_default_metrics_from_samples(args, samples), "rollout/") + log_dict |= dict_add_prefix(_compute_default_perf_metrics_from_samples(args, samples, rollout_time), "perf/") + log_dict |= dict_add_prefix(_compute_ags_metrics(args, samples), "rollout/ags/") + + step = compute_rollout_step(args, rollout_id) + log_dict["rollout/step"] = step + logger.info("ags rollout %s: %s", rollout_id, log_dict) + logging_utils.log(args, log_dict, step_key="rollout/step") + return True + except Exception: # noqa: BLE001 - logging must never break rollout/training + logger.exception("AGS custom rollout logger failed; falling back to slime default logger") + return False + + +def log_eval_rollout_data(rollout_id: int, args: Any, data: dict[str, dict[str, Any]], extra_metrics=None) -> bool: + """Log default eval metrics plus per-dataset and pooled AGS diagnostics.""" + + try: + log_dict = {**(extra_metrics or {})} + all_rewards: list[float] = [] + all_samples: list[Sample] = [] + + for dataset_name, dataset_data in data.items(): + rewards = [_safe_float(r) for r in dataset_data.get("rewards", [])] + all_rewards.extend(rewards) + if rewards: + log_dict[f"eval/{dataset_name}"] = sum(rewards) / len(rewards) + + samples = dataset_data.get("samples") + if samples is not None: + samples = list(samples) + all_samples.extend(samples) + log_dict |= dict_add_prefix( + _compute_default_metrics_from_samples(args, samples), + f"eval/{dataset_name}/", + ) + log_dict |= dict_add_prefix(_compute_ags_metrics(args, samples), f"eval/{dataset_name}/ags/") + + if "truncated" in dataset_data: + truncated = dataset_data["truncated"] + if truncated: + log_dict[f"eval/{dataset_name}-truncated_ratio"] = sum(bool(x) for x in truncated) / len(truncated) + + if getattr(args, "log_passrate", False) and rewards: + log_dict |= dict_add_prefix( + compute_pass_rate( + flat_rewards=rewards, + group_size=getattr(args, "n_samples_per_eval_prompt", 1), + ), + f"eval/{dataset_name}-", + ) + + if all_rewards: + log_dict["eval/overall"] = sum(all_rewards) / len(all_rewards) + log_dict |= dict_add_prefix(_stats(all_rewards), "eval/overall/reward/") + if all_samples: + log_dict |= dict_add_prefix(_compute_ags_metrics(args, all_samples), "eval/overall/ags/") + + step = compute_rollout_step(args, rollout_id) + log_dict["eval/step"] = step + logger.info("ags eval %s: %s", rollout_id, log_dict) + logging_utils.log(args, log_dict, step_key="eval/step") + return True + except Exception: # noqa: BLE001 - logging must never break eval + logger.exception("AGS custom eval logger failed; falling back to slime default logger") + return False + + +def _compute_default_metrics_from_samples(args: Any, samples: list[Sample]) -> dict[str, Any]: + # Lazy import keeps this hook importable in lightweight environments where + # rollout-only optional dependencies (for example sglang) are not installed. + from slime.ray.rollout import compute_metrics_from_samples + + return compute_metrics_from_samples(args, samples) + + +def _compute_default_perf_metrics_from_samples( + args: Any, + samples: list[Sample], + rollout_time: float, +) -> dict[str, Any]: + # See _compute_default_metrics_from_samples for why this is lazy. + from slime.ray.rollout import compute_perf_metrics_from_samples + + return compute_perf_metrics_from_samples(args, samples, rollout_time) + + +def _compute_ags_metrics(args: Any, samples: Iterable[Sample]) -> dict[str, float | int]: + samples = list(samples) + if not samples: + return {} + + n = len(samples) + rewards = [_reward_value(args, sample) for sample in samples] + statuses = Counter(_status_value(sample) for sample in samples) + metadata = [_metadata(sample) for sample in samples] + + valid_response_count = sum(1 for sample in samples if getattr(sample, "response_length", 0) > 0 and sample.tokens) + remove_sample_count = sum(1 for sample in samples if bool(getattr(sample, "remove_sample", False))) + nonzero_reward_count = sum(1 for reward in rewards if reward != 0.0) + solve_count = sum(1 for sample, reward in zip(samples, rewards, strict=True) if _is_solved(sample, reward)) + + metrics: dict[str, float | int] = { + "samples/count": n, + "reward/nonzero_count": nonzero_reward_count, + "reward/nonzero_rate": _ratio(nonzero_reward_count, n), + "solve/count": solve_count, + "remove_sample/count": remove_sample_count, + "remove_sample/rate": _ratio(remove_sample_count, n), + "response/valid_count": valid_response_count, + "response/valid_rate": _ratio(valid_response_count, n), + } + metrics |= dict_add_prefix(_stats(rewards), "reward/") + metrics["solve/rate"] = _ratio(metrics["solve/count"], n) + + for status in sorted(statuses): + count = statuses[status] + metrics[f"status/{_bucket(status)}/count"] = count + metrics[f"status/{_bucket(status)}/rate"] = _ratio(count, n) + for status in Sample.Status: + metrics.setdefault(f"status/{status.value}/count", 0) + metrics.setdefault(f"status/{status.value}/rate", 0.0) + + metrics |= _artifact_metrics(metadata, n) + metrics |= _runtime_metrics(metadata, n) + metrics |= _abort_reason_metrics(metadata, n) + metrics |= _guardrail_metrics(metadata, n) + metrics |= _rollout_level_metrics(samples, rewards) + metrics |= _agent_metrics(samples, metadata, rewards) + return metrics + + +def _artifact_metrics(metadata: list[dict[str, Any]], n: int) -> dict[str, float | int]: + has_trajectory = [bool(md.get("trajectory_path")) for md in metadata] + has_patch = [bool(md.get("patch_path")) for md in metadata] + has_rollout_dump = [bool(md.get("rollout_dump_path")) for md in metadata] + complete = [t and p and d for t, p, d in zip(has_trajectory, has_patch, has_rollout_dump, strict=True)] + + metrics: dict[str, float | int] = {} + for name, values in { + "trajectory": has_trajectory, + "patch": has_patch, + "rollout_dump": has_rollout_dump, + "complete": complete, + }.items(): + count = sum(values) + metrics[f"artifact/{name}/count"] = count + metrics[f"artifact/{name}/rate"] = _ratio(count, n) + return metrics + + +def _runtime_metrics(metadata: list[dict[str, Any]], n: int) -> dict[str, float | int]: + metrics: dict[str, float | int] = {} + + elapsed = [_safe_float(md.get("ags_elapsed_sec"), default=None) for md in metadata] + elapsed = [value for value in elapsed if value is not None] + if elapsed: + metrics |= dict_add_prefix(_stats(elapsed), "elapsed_sec/") + metrics["elapsed_sec/p50"] = _percentile(elapsed, 50) + metrics["elapsed_sec/p95"] = _percentile(elapsed, 95) + metrics["elapsed_sec/count"] = len(elapsed) + + segments = [_safe_float(md.get("ags_num_samples"), default=None) for md in metadata] + segments = [value for value in segments if value is not None] + if segments: + metrics |= dict_add_prefix(_stats(segments), "segments/") + + concurrencies = [_safe_float(md.get("ags_rollout_concurrency"), default=None) for md in metadata] + concurrencies = [value for value in concurrencies if value is not None] + if concurrencies: + metrics["rollout_concurrency/max"] = max(concurrencies) + + agent_exit_codes = [_safe_float(md.get("agent_exit_code"), default=None) for md in metadata] + agent_exit_codes = [value for value in agent_exit_codes if value is not None] + if agent_exit_codes: + nonzero = sum(1 for code in agent_exit_codes if int(code) != 0) + metrics["agent_exit/nonzero_count"] = nonzero + metrics["agent_exit/nonzero_rate"] = _ratio(nonzero, len(agent_exit_codes)) + + applied = [md.get("applied_cleanly") for md in metadata if "applied_cleanly" in md] + if applied: + count = sum(1 for value in applied if _safe_bool(value)) + metrics["patch/applied_cleanly_count"] = count + metrics["patch/applied_cleanly_rate"] = _ratio(count, len(applied)) + else: + metrics["patch/applied_cleanly_count"] = 0 + metrics["patch/applied_cleanly_rate"] = 0.0 + + # ``n`` is kept as an explicit denominator for dashboards where some fields + # are absent on aborted samples. + metrics["metadata/coverage_rate"] = _ratio(sum(1 for md in metadata if md), n) + return metrics + + +def _abort_reason_metrics(metadata: list[dict[str, Any]], n: int) -> dict[str, float | int]: + reasons = Counter(str(md.get("abort_reason") or "unknown") for md in metadata if md.get("abort_reason")) + metrics: dict[str, float | int] = {} + for reason, count in sorted(reasons.items()): + key = _bucket(reason) + metrics[f"abort_reason/{key}/count"] = count + metrics[f"abort_reason/{key}/rate"] = _ratio(count, n) + metrics["abort_reason/total_count"] = sum(reasons.values()) + metrics["abort_reason/total_rate"] = _ratio(sum(reasons.values()), n) + return metrics + + +def _guardrail_metrics(metadata: list[dict[str, Any]], n: int) -> dict[str, float | int]: + """Counters for silent-failure guardrails. + + ``empty_patch`` counts rollouts that reported success with no diff, split by + the reason label -- the failure mode where a coding-agent CLI accepts a + text-only mid-task turn as a finished answer. ``ill_formed`` is surfaced + alongside it because the trajectory manager has always recorded it and it + never had a metrics outlet. + """ + metrics: dict[str, float | int] = {} + + triggered = [md for md in metadata if _safe_bool(md.get("empty_patch_guard_triggered"))] + metrics["guardrail/empty_patch/count"] = len(triggered) + metrics["guardrail/empty_patch/rate"] = _ratio(len(triggered), n) + for reason, count in sorted( + Counter(str(md.get("empty_patch_guard_reason") or "unknown") for md in triggered).items() + ): + metrics[f"guardrail/empty_patch/{_bucket(reason)}/count"] = count + metrics[f"guardrail/empty_patch/{_bucket(reason)}/rate"] = _ratio(count, n) + + ill_formed = sum(1 for md in metadata if _safe_bool(md.get("ill_formed"))) + metrics["guardrail/ill_formed/count"] = ill_formed + metrics["guardrail/ill_formed/rate"] = _ratio(ill_formed, n) + return metrics + + +def _rollout_level_metrics(samples: list[Sample], rewards: list[float]) -> dict[str, float | int]: + by_rollout: dict[str, list[tuple[Sample, float]]] = defaultdict(list) + for position, (sample, reward) in enumerate(zip(samples, rewards, strict=True)): + by_rollout[_rollout_key(sample, position)].append((sample, reward)) + + rollout_rewards = [max(reward for _, reward in items) for items in by_rollout.values()] + rollout_solved = [any(_is_solved(sample, reward) for sample, reward in items) for items in by_rollout.values()] + metrics: dict[str, float | int] = { + "rollout/count": len(by_rollout), + "rollout/solve_count": sum(rollout_solved), + "rollout/solve_rate": _ratio(sum(rollout_solved), len(by_rollout)), + } + metrics |= dict_add_prefix(_stats(rollout_rewards), "rollout/reward/") + return metrics + + +def _agent_metrics( + samples: list[Sample], + metadata: list[dict[str, Any]], + rewards: list[float], +) -> dict[str, float | int]: + by_agent: dict[str, list[tuple[Sample, float]]] = defaultdict(list) + for sample, md, reward in zip(samples, metadata, rewards, strict=True): + by_agent[str(md.get("agent") or "unknown")].append((sample, reward)) + + metrics: dict[str, float | int] = {"agent/count": len(by_agent)} + for agent, items in sorted(by_agent.items()): + key = _bucket(agent) + count = len(items) + agent_rewards = [reward for _, reward in items] + metrics[f"agent/{key}/count"] = count + metrics[f"agent/{key}/rate"] = _ratio(count, len(samples)) + metrics[f"agent/{key}/solve_rate"] = _ratio( + sum(1 for sample, reward in items if _is_solved(sample, reward)), count + ) + metrics[f"agent/{key}/reward_mean"] = sum(agent_rewards) / count if count else 0.0 + return metrics + + +def _reward_value(args: Any, sample: Sample) -> float: + reward = getattr(sample, "reward", 0.0) + if isinstance(reward, dict): + reward_key = getattr(args, "eval_reward_key", None) or getattr(args, "reward_key", None) + if reward_key and reward_key in reward: + return _safe_float(reward[reward_key]) or 0.0 + for key in ("score", "reward", "acc"): + if key in reward: + return _safe_float(reward[key]) or 0.0 + for value in reward.values(): + coerced = _safe_float(value, default=None) + if coerced is not None: + return coerced + return 0.0 + return _safe_float(reward) or 0.0 + + +def _is_solved(sample: Sample, reward: float) -> bool: + md = _metadata(sample) + if "grading_solved" in md: + return _safe_bool(md["grading_solved"]) + return reward == 1.0 + + +def _metadata(sample: Sample) -> dict[str, Any]: + md = getattr(sample, "metadata", None) + return md if isinstance(md, dict) else {} + + +def _status_value(sample: Sample) -> str: + status = getattr(sample, "status", None) + return status.value if isinstance(status, Sample.Status) else str(status or "unknown") + + +def _rollout_key(sample: Sample, position: int) -> str: + rollout_id = getattr(sample, "rollout_id", None) + if rollout_id is not None: + return f"rollout:{rollout_id}" + session_id = getattr(sample, "session_id", None) + if session_id is not None: + return f"session:{session_id}" + return f"position:{position}" + + +def _stats(values: list[float]) -> dict[str, float]: + values = [float(value) for value in values] + if not values: + return {} + return compute_statistics(values) + + +def _percentile(values: list[float], q: float) -> float: + if not values: + return 0.0 + return float(np.percentile(np.asarray(values, dtype=float), q)) + + +def _safe_float(value: Any, default: float | None = 0.0) -> float | None: + try: + if value is None: + return default + return float(value) + except (TypeError, ValueError): + return default + + +def _safe_bool(value: Any) -> bool: + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "y", "on"} + return bool(value) + + +def _ratio(count: float | int, total: float | int) -> float: + return float(count) / float(total) if total else 0.0 + + +def _bucket(value: str) -> str: + value = value.strip() or "unknown" + return _BUCKET_RE.sub("_", value).strip("_") or "unknown" diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/weave_trace.py b/slime_plugins/rollout_buffer/generator/ags_generator/weave_trace.py new file mode 100644 index 0000000000..d43c27148c --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/weave_trace.py @@ -0,0 +1,288 @@ +"""Best-effort Weave tracing for AGS rollouts.""" + +from __future__ import annotations + +import json +import logging +import os +from argparse import Namespace +from datetime import datetime +from pathlib import Path +from typing import Any + +from slime.utils.types import Sample + +logger = logging.getLogger(__name__) + + +class AGSWeaveTrace: + def __init__( + self, + args: Namespace, + tokenizer, + *, + agent: str = "claude_code", + enable_token2text: bool = False, + ) -> None: + self.project = _wandb_project(args) + self.tokenizer = tokenizer + self.agent = agent + self.enable_token2text = enable_token2text + self.wandb_run_id = getattr(args, "wandb_run_id", None) + self.wandb_group = getattr(args, "wandb_group", None) + self.client = None + if self.project is None: + return + try: + import weave + + self.client = weave.init(self.project) + if self.wandb_run_id: + self.client.set_wandb_run_context(run_id=self.wandb_run_id) + except Exception: + logger.warning("[ags_generator] failed to initialize Weave tracing", exc_info=True) + + def start_rollout( + self, + *, + instance_id: str, + session_id: str, + sample: Sample, + sampling_params: dict[str, Any], + agent: str, + ): + if self.client is None: + return None + inputs = { + "instance_id": instance_id, + "session_id": session_id, + "prompt": sample.prompt, + "sampling_params": sampling_params, + } + attributes = { + "task_type": "ags", + "agent": agent, + "group_index": sample.group_index, + "sample_index": sample.index, + "rollout_id": sample.rollout_id, + "wandb_run_id": self.wandb_run_id, + "wandb_group": self.wandb_group, + } + try: + return self.client.create_call( + "slime.ags.rollout", + inputs, + attributes=attributes, + display_name=instance_id, + use_stack=False, + ) + except Exception: + logger.warning("[ags_generator] %s: failed to start Weave trace", instance_id, exc_info=True) + return None + + def finish_rollout( + self, + call, + *, + samples: list[Sample], + trajectory_path: str | None = None, + output: dict[str, Any] | None = None, + exception: BaseException | None = None, + ) -> None: + if self.client is None or call is None: + return + result = dict(output or {}) + try: + if trajectory_path: + self._log_trajectory(call, trajectory_path) + result["samples"] = [self._sample_payload(sample) for sample in samples] + except Exception: + logger.warning("[ags_generator] failed to build Weave trace output", exc_info=True) + result["trace_output_error"] = True + try: + self.client.finish_call(call, output=result, exception=exception) + except Exception: + logger.warning("[ags_generator] failed to finish Weave root call", exc_info=True) + + def _log_trajectory(self, parent, trajectory_path: str) -> None: + tool_calls = {} + for event in iter_trajectory_events(trajectory_path, agent=self.agent): + tool_use_id = event.get("tool_use_id") + if event["kind"] == "tool_result" and tool_use_id in tool_calls: + self.client.finish_call( + tool_calls.pop(tool_use_id), + output=event.get("output"), + ended_at=event.get("started_at"), + ) + continue + child = self.client.create_call( + f"slime.ags.{event['kind']}", + event["inputs"], + parent=parent, + attributes=event.get("attributes"), + display_name=event["display_name"], + use_stack=False, + started_at=event.get("started_at"), + ) + if event["kind"] == "tool_call" and tool_use_id: + tool_calls[tool_use_id] = child + else: + self.client.finish_call(child, output=event.get("output"), ended_at=event.get("started_at")) + for child in tool_calls.values(): + self.client.finish_call(child, output={"missing_tool_result": True}) + + def _sample_payload(self, sample: Sample) -> dict[str, Any]: + tokens = [int(token) for token in sample.tokens] + response_length = min(max(int(sample.response_length or 0), 0), len(tokens)) + prompt_tokens = tokens[:-response_length] if response_length else tokens + response_tokens = tokens[-response_length:] if response_length else [] + payload = { + "status": sample.status.value, + "reward": sample.reward, + "remove_sample": sample.remove_sample, + "prompt_token_ids": prompt_tokens, + "response_token_ids": response_tokens, + "response_length": response_length, + "metadata": sample.metadata, + } + if self.enable_token2text: + payload["prompt_text"] = self.tokenizer.decode(prompt_tokens, skip_special_tokens=False) + payload["response_text"] = self.tokenizer.decode(response_tokens, skip_special_tokens=False) + return payload + + +def iter_trajectory_events(path: str | Path, *, agent: str = "claude_code"): + parser = _trajectory_parser(agent) + if parser is None: + logger.debug("[ags_generator] no Weave trajectory parser for agent=%s", agent) + return + yield from parser(path) + + +def iter_claude_code_trajectory_events(path: str | Path): + yield from _iter_anthropic_message_events(path, timestamp_key="timestamp") + + +def iter_codebuddy_code_trajectory_events(path: str | Path): + yield from _iter_anthropic_message_events(path, timestamp_key="__timestamp") + + +def iter_codex_trajectory_events(path: str | Path): + return + yield + + +def _trajectory_parser(agent: str): + return { + "claude_code": iter_claude_code_trajectory_events, + "codebuddy_code": iter_codebuddy_code_trajectory_events, + "codex": iter_codex_trajectory_events, + }.get(agent) + + +def _iter_anthropic_message_events(path: str | Path, *, timestamp_key: str): + try: + lines = Path(path).read_text(encoding="utf-8").splitlines() + except OSError: + logger.warning("[ags_generator] failed to read trajectory for Weave trace: %s", path, exc_info=True) + return + + for line_number, line in enumerate(lines, start=1): + try: + record = json.loads(line) + except json.JSONDecodeError: + logger.warning("[ags_generator] skipping invalid trajectory JSON at %s:%d", path, line_number) + continue + record_type = record.get("type") + if record_type == "assistant": + yield from _assistant_events(record, timestamp_key=timestamp_key) + elif record_type == "user": + yield from _tool_result_events(record, timestamp_key=timestamp_key) + elif record_type == "result": + yield { + "kind": "result", + "display_name": "agent result", + "inputs": _common_fields(record, timestamp_key=timestamp_key), + "output": record, + "started_at": _record_timestamp(record, timestamp_key), + } + + +def _assistant_events(record: dict[str, Any], *, timestamp_key: str): + message = record.get("message") or {} + common = _common_fields(record, timestamp_key=timestamp_key) + for block in message.get("content") or []: + block_type = block.get("type") + if block_type == "tool_use": + yield { + "kind": "tool_call", + "display_name": str(block.get("name") or "tool call"), + "tool_use_id": block.get("id"), + "inputs": {**common, "tool_use_id": block.get("id"), "input": block.get("input")}, + "attributes": {"tool_name": block.get("name")}, + "started_at": _record_timestamp(record, timestamp_key), + } + elif block_type in {"text", "thinking"}: + text = block.get(block_type) + if text: + yield { + "kind": block_type, + "display_name": f"assistant {block_type}", + "inputs": common, + "output": {"text": text, "usage": message.get("usage")}, + "started_at": _record_timestamp(record, timestamp_key), + } + + +def _tool_result_events(record: dict[str, Any], *, timestamp_key: str): + message = record.get("message") or {} + common = _common_fields(record, timestamp_key=timestamp_key) + for block in message.get("content") or []: + if block.get("type") == "tool_result": + yield { + "kind": "tool_result", + "display_name": "tool result", + "tool_use_id": block.get("tool_use_id"), + "inputs": {**common, "tool_use_id": block.get("tool_use_id")}, + "output": { + "content": block.get("content"), + "is_error": block.get("is_error", False), + "tool_use_result": record.get("tool_use_result"), + }, + "started_at": _record_timestamp(record, timestamp_key), + } + + +def _common_fields(record: dict[str, Any], *, timestamp_key: str) -> dict[str, Any]: + return { + "session_id": record.get("session_id"), + "event_id": record.get("uuid"), + "parent_tool_use_id": record.get("parent_tool_use_id"), + "timestamp": record.get(timestamp_key), + } + + +def _record_timestamp(record: dict[str, Any], timestamp_key: str) -> datetime | None: + return _parse_timestamp(record.get(timestamp_key)) + + +def _parse_timestamp(value: Any) -> datetime | None: + if not isinstance(value, str) or not value: + return None + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + + +def _wandb_project(args: Namespace) -> str | None: + if not getattr(args, "use_wandb", False): + return None + wandb_mode = getattr(args, "wandb_mode", None) or os.environ.get("WANDB_MODE") + if wandb_mode in {"disabled", "offline"}: + return None + project = getattr(args, "wandb_project", None) + if not project: + return None + team = getattr(args, "wandb_team", None) + return f"{team}/{project}" if team else project diff --git a/slime_plugins/rollout_buffer/rollout_buffer_example.py b/slime_plugins/rollout_buffer/rollout_buffer_example.py index 74b4b4c46c..02a6721b56 100644 --- a/slime_plugins/rollout_buffer/rollout_buffer_example.py +++ b/slime_plugins/rollout_buffer/rollout_buffer_example.py @@ -16,7 +16,6 @@ # Global variables for evaluation TOKENIZER = None -START_ROLLOUT = True def select_rollout_data(args, results, need_length): @@ -45,8 +44,10 @@ def select_rollout_data(args, results, need_length): print(f"📊 Total groups: {len(groups)}, total samples: {len(results)}") - # If we don't have too many samples, return all - assert need_length < len(results), "need_length must be smaller than results length" + # Return grouped records even when there is no over-collection; downstream + # expects prompt-group shape, not a flat record list. + if len(groups) <= need_length: + return list(groups.values()) # Get timestamp for each group (use the latest timestamp in the group) def get_group_timestamp(group_items): @@ -79,7 +80,7 @@ def get_group_timestamp(group_items): newest_ts = selected_groups[0][1] oldest_ts = selected_groups[-1][1] print( - f"📈 Selected {len(selected_groups)} groups with {len(selected_results)*args.n_samples_per_prompt} samples" + f"📈 Selected {len(selected_groups)} groups with {len(selected_results) * args.n_samples_per_prompt} samples" ) print(f"📈 Group timestamp range: {oldest_ts:.2f} to {newest_ts:.2f}") print(f"📈 Time span: {newest_ts - oldest_ts:.2f} seconds") @@ -87,52 +88,194 @@ def get_group_timestamp(group_items): return selected_results -def log_raw_info(args, all_meta_info, rollout_id): - final_meta_info = {} - if all_meta_info: - final_meta_info = { - "total_samples": sum(meta["total_samples"] for meta in all_meta_info if "total_samples" in meta) - } +def _select_complete_rollout_groups(args, results, need_length): + """Keep only prompt groups with every requested rollout attempt present.""" - total_samples = final_meta_info["total_samples"] - if total_samples > 0: - weighted_reward_sum = sum( - meta["avg_reward"] * meta["total_samples"] - for meta in all_meta_info - if "avg_reward" in meta and "total_samples" in meta - ) + grouped = {} + for record in results: + grouped.setdefault(record["instance_id"], []).append(record) - final_meta_info.update( - { - "avg_reward": weighted_reward_sum / total_samples, - } - ) - if args.use_wandb: - log_dict = { - "rollout/no_filter/total_samples": final_meta_info["total_samples"], - "rollout/no_filter/avg_reward": final_meta_info["avg_reward"], - } - try: - step = ( - rollout_id - if not args.wandb_always_use_train_step - else rollout_id * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size - ) - if args.use_wandb: - log_dict["rollout/step"] = step - wandb.log(log_dict) + expected_size = args.n_samples_per_prompt + complete_records = [] + incomplete_groups = {} + for instance_id, group_records in grouped.items(): + if len(group_records) == expected_size: + complete_records.extend(group_records) + else: + incomplete_groups[instance_id] = len(group_records) - if args.use_tensorboard: - from slime.utils.tensorboard_utils import _TensorboardAdapter + return select_rollout_data(args, complete_records, need_length), incomplete_groups - tb = _TensorboardAdapter(args) - tb.log(data=log_dict, step=step) - print(f"no filter rollout log {rollout_id}: {log_dict}") - except Exception as e: - print(f"Failed to log to wandb: {e}") - print(f"no filter rollout log {rollout_id}: {final_meta_info}") - else: - print(f"no filter rollout log {rollout_id}: {final_meta_info}") + +def log_raw_info(args, all_meta_info, rollout_id): + if not all_meta_info: + return + + final_meta_info = _merge_rollout_meta_info(all_meta_info) + if final_meta_info.get("total_samples", 0) <= 0: + print(f"no filter rollout log {rollout_id}: {final_meta_info}") + return + + try: + step = ( + rollout_id + if not args.wandb_always_use_train_step + else rollout_id * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size + ) + log_dict = _flatten_numeric_metrics("rollout/no_filter", final_meta_info) + if getattr(args, "rollout_task_type", None) == "ags": + log_dict.update(_flatten_numeric_metrics("rollout/ags", final_meta_info)) + log_dict["rollout/step"] = step + if args.use_wandb: + wandb.log(log_dict) + + if args.use_tensorboard: + from slime.utils.tensorboard_utils import _TensorboardAdapter + + tb = _TensorboardAdapter(args) + tb.log(data=log_dict, step=step) + print(f"no filter rollout log {rollout_id}: {log_dict}") + except Exception as e: + print(f"Failed to log rollout metrics: {e}") + print(f"no filter rollout log {rollout_id}: {final_meta_info}") + + +def _merge_rollout_meta_info(all_meta_info: list[dict[str, Any]]) -> dict[str, Any]: + total_samples = sum(int(meta.get("total_samples", 0) or 0) for meta in all_meta_info) + merged: dict[str, Any] = {"total_samples": total_samples} + if total_samples <= 0: + return merged + + merged["num_groups"] = sum(int(meta.get("num_groups", 0) or 0) for meta in all_meta_info) + + if any("total_rollouts" in meta for meta in all_meta_info): + merged["total_rollouts"] = sum( + int(meta.get("total_rollouts", meta.get("total_samples", 0)) or 0) for meta in all_meta_info + ) + + for key in ["nonzero_reward_samples", "solved_samples"]: + if any(key in meta for meta in all_meta_info): + merged[key] = sum(int(meta.get(key, 0) or 0) for meta in all_meta_info) + + weighted_avg_keys = [ + "avg_reward", + "solve_rate", + "nonzero_reward_rate", + "completed_rate", + "abort_rate", + "artifact_complete_rate", + ] + for key in weighted_avg_keys: + if not any(key in meta for meta in all_meta_info): + continue + weighted_sum = sum( + float(meta[key]) * int(meta.get("total_samples", 0) or 0) + for meta in all_meta_info + if key in meta and meta.get("total_samples", 0) + ) + merged[key] = weighted_sum / total_samples + + num_groups = int(merged.get("num_groups", 0) or 0) + if "total_rollouts" in merged: + merged["avg_group_size"] = int(merged["total_rollouts"]) / num_groups if num_groups else 0 + merged["avg_samples_per_group"] = total_samples / num_groups if num_groups else 0 + elif any("avg_group_size" in meta for meta in all_meta_info): + weighted_sum = sum( + float(meta["avg_group_size"]) * int(meta.get("total_samples", 0) or 0) + for meta in all_meta_info + if "avg_group_size" in meta and meta.get("total_samples", 0) + ) + merged["avg_group_size"] = weighted_sum / total_samples + + for key in ["status_counts", "artifact_counts"]: + if any(key in meta for meta in all_meta_info): + merged[key] = _sum_nested_counts(meta.get(key, {}) for meta in all_meta_info) + + performance_items = [meta.get("performance", {}) for meta in all_meta_info] + if any(performance_items): + elapsed_sec_values = [ + float(value) + for item in performance_items + if isinstance(item, dict) + for value in item.get("elapsed_sec_values", []) + ] + merged["performance"] = { + "rollout_concurrency": _max_nested(performance_items, "rollout_concurrency"), + "avg_elapsed_sec": ( + sum(elapsed_sec_values) / len(elapsed_sec_values) + if elapsed_sec_values + else _weighted_average_nested(performance_items, all_meta_info, "avg_elapsed_sec") + ), + "p50_elapsed_sec": ( + _percentile(elapsed_sec_values, 0.50) + if elapsed_sec_values + else _max_nested(performance_items, "p50_elapsed_sec") + ), + "p95_elapsed_sec": ( + _percentile(elapsed_sec_values, 0.95) + if elapsed_sec_values + else _max_nested(performance_items, "p95_elapsed_sec") + ), + "max_elapsed_sec": ( + max(elapsed_sec_values) if elapsed_sec_values else _max_nested(performance_items, "max_elapsed_sec") + ), + "agent_exit_nonzero_count": sum( + int(item.get("agent_exit_nonzero_count", 0) or 0) for item in performance_items + ), + "applied_cleanly_count": sum(int(item.get("applied_cleanly_count", 0) or 0) for item in performance_items), + } + return merged + + +def _sum_nested_counts(dicts) -> dict[str, int]: + output: dict[str, int] = {} + for data in dicts: + if not isinstance(data, dict): + continue + for key, value in data.items(): + if isinstance(value, bool): + value = int(value) + if isinstance(value, int | float): + output[str(key)] = output.get(str(key), 0) + int(value) + return output + + +def _weighted_average_nested(items: list[dict[str, Any]], all_meta_info: list[dict[str, Any]], key: str) -> float: + weighted_sum = 0.0 + total_weight = 0 + for item, meta in zip(items, all_meta_info, strict=False): + if key not in item: + continue + weight = int(meta.get("total_samples", 0) or 0) + weighted_sum += float(item[key]) * weight + total_weight += weight + return weighted_sum / total_weight if total_weight else 0 + + +def _max_nested(items: list[dict[str, Any]], key: str) -> float: + values = [float(item[key]) for item in items if isinstance(item, dict) and key in item] + return max(values) if values else 0 + + +def _percentile(values: list[float], quantile: float) -> float: + if not values: + return 0 + ordered = sorted(values) + index = min(len(ordered) - 1, max(0, int(round((len(ordered) - 1) * quantile)))) + return ordered[index] + + +def _flatten_numeric_metrics(prefix: str, data: dict[str, Any]) -> dict[str, int | float]: + metrics: dict[str, int | float] = {} + for key, value in data.items(): + metric_key = f"{prefix}/{key}" + if isinstance(value, bool): + metrics[metric_key] = int(value) + elif isinstance(value, int | float): + metrics[metric_key] = value + elif isinstance(value, dict): + metrics.update(_flatten_numeric_metrics(metric_key, value)) + return metrics async def get_rollout_data(api_base_url: str) -> tuple[list[dict[str, Any]], dict[str, Any]]: @@ -170,13 +313,20 @@ async def get_rollout_data(api_base_url: str) -> tuple[list[dict[str, Any]], dic return data, meta_info -def start_rollout(api_base_url: str, args, metadata): +def start_rollout(api_base_url: str, args, metadata, num_groups_per_epoch: int | None = None, rollout_id: int = 0): url = f"{api_base_url}/start_rollout" print(f"metadata: {metadata}") finished_groups_instance_id_list = [item for sublist in metadata.values() for item in sublist] + rollout_buffer_num_epoch = int(getattr(args, "rollout_buffer_num_epoch", 1) or 1) + groups_per_epoch = int(num_groups_per_epoch or args.rollout_batch_size) payload = { "num_process": str(getattr(args, "rollout_num_process", 100)), - "num_epoch": str(args.num_epoch or 3), + # This is the generator-side epoch count for one rollout-buffer request, + # not slime's trainer --num-epoch. Keeping it at 1 prevents the + # background AGS job from continuing to call SGLang after this training + # rollout has already returned and the colocated engine is offloaded. + "num_epoch": str(rollout_buffer_num_epoch), + "num_groups_per_epoch": str(groups_per_epoch), "remote_engine_url": f"http://{args.sglang_router_ip}:{args.sglang_router_port}", "remote_buffer_url": args.rollout_buffer_url, "task_type": args.rollout_task_type, @@ -187,8 +337,30 @@ def start_rollout(api_base_url: str, args, metadata): "max_tokens": args.rollout_max_response_len, "temperature": args.rollout_temperature, "top_p": args.rollout_top_p, + "top_k": args.rollout_top_k, }, "tokenizer_path": args.hf_checkpoint, + "input_key": getattr(args, "input_key", "prompt"), + "label_key": getattr(args, "label_key", "label"), + "metadata_key": getattr(args, "metadata_key", "metadata"), + "tool_key": getattr(args, "tool_key", None), + "apply_chat_template": getattr(args, "apply_chat_template", False), + "apply_chat_template_kwargs": getattr(args, "apply_chat_template_kwargs", {}) or {}, + "rollout_batch_size": args.rollout_batch_size, + "rollout_id": int(rollout_id), + "rollout_start_group": int(rollout_id) * int(args.rollout_batch_size), + "rollout_max_context_len": getattr(args, "rollout_max_context_len", 0), + "rollout_seed": getattr(args, "rollout_seed", 42), + "rollout_shuffle": getattr(args, "rollout_shuffle", False), + "sglang_tool_call_parser": getattr(args, "sglang_tool_call_parser", None), + "sglang_reasoning_parser": getattr(args, "sglang_reasoning_parser", None), + "use_wandb": getattr(args, "use_wandb", False), + "wandb_mode": getattr(args, "wandb_mode", None), + "wandb_project": getattr(args, "wandb_project", None), + "wandb_team": getattr(args, "wandb_team", None), + "wandb_run_id": getattr(args, "wandb_run_id", None), + "wandb_group": getattr(args, "wandb_group", None), + "enable_token2text": getattr(args, "enable_token2text", False), "skip_instance_ids": finished_groups_instance_id_list, } print("start rollout with payload: ", payload) @@ -202,21 +374,29 @@ def start_rollout(api_base_url: str, args, metadata): return data except Exception as e: print(f"[start_rollout] Failed to send rollout config: {e}") + time.sleep(3) + + +def stop_rollout(api_base_url: str, timeout_sec: float = 120.0) -> dict[str, Any]: + url = f"{api_base_url}/stop_rollout" + resp = requests.post( + url, + json={"wait": True, "timeout_sec": timeout_sec}, + timeout=timeout_sec + 10, + ) + resp.raise_for_status() + data = resp.json() + print(f"[stop_rollout] {data}") + if not data.get("stopped", False): + raise RuntimeError(f"rollout background job did not stop within {timeout_sec}s: {data}") + return data async def generate_rollout_async(args, rollout_id: int, data_buffer, evaluation: bool = False) -> dict[str, Any]: - global START_ROLLOUT if evaluation: raise NotImplementedError("Evaluation rollout is not implemented") - if START_ROLLOUT: - metadata = data_buffer.get_metadata() - start_inform = start_rollout(args.rollout_buffer_url, args, metadata) - print(f"start rollout with payload: {start_inform}") - print(f"start rollout id: {rollout_id}") - START_ROLLOUT = False - data_number_to_fetch = args.rollout_batch_size * args.n_samples_per_prompt - data_buffer.get_buffer_length() if data_number_to_fetch <= 0: print( @@ -226,80 +406,131 @@ async def generate_rollout_async(args, rollout_id: int, data_buffer, evaluation: assert ( data_number_to_fetch % args.n_samples_per_prompt == 0 ), "data_number_to_fetch must be a multiple of n_samples_per_prompt" - print(f"INFO: buffer length: {data_buffer.get_buffer_length()}, data_number_to_fetch: {data_number_to_fetch}") - base_url = args.rollout_buffer_url - tokenizer = AutoTokenizer.from_pretrained(args.hf_checkpoint, trust_remote_code=True) - retry_times = 0 - results = [] - all_meta_info = [] - - if args.fetch_trajectory_retry_times == -1: - print( - "⚠️ [get_rollout_data] Fetch trajectory retry times set to -1, will retry indefinitely until sufficient data is collected" + + groups_to_fetch = data_number_to_fetch // args.n_samples_per_prompt + rollout_started = False + try: + metadata = data_buffer.get_metadata() + start_inform = start_rollout( + args.rollout_buffer_url, args, metadata, num_groups_per_epoch=groups_to_fetch, rollout_id=rollout_id ) - while args.fetch_trajectory_retry_times == -1 or retry_times < args.fetch_trajectory_retry_times: - try: - while len(results) < data_number_to_fetch: - time.sleep(5) - data, meta_info = await get_rollout_data(api_base_url=base_url) - results.extend(data) - if meta_info: - all_meta_info.append(meta_info) - print(f"get rollout data with length: {len(results)}") - break - except Exception as err: - print(f"[get_rollout_data] Failed to get rollout data: {err}, retry times: {retry_times}") - retry_times += 1 - - log_raw_info(args, all_meta_info, rollout_id) - - # Apply group-based data selection if there are too many samples - results = select_rollout_data(args, results, data_number_to_fetch // args.n_samples_per_prompt) - - if len(all_meta_info) > 0 and "finished_groups" in all_meta_info[0]: - finished_groups_instance_id_list = [] - for item in all_meta_info: - finished_groups_instance_id_list.extend(item["finished_groups"]) - - data_buffer.update_metadata({str(rollout_id): finished_groups_instance_id_list}) - - print("finally get rollout data with length: ", len(results)) - sample_results = [] - - for _i, group_record in enumerate(results): - group_results = [] - for record in group_record: - oai_messages = record["messages"] - - mask_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type=args.loss_mask_type) - token_ids, loss_mask = mask_generator.get_loss_mask(oai_messages) - response_length = mask_generator.get_response_lengths([loss_mask])[0] - - loss_mask = loss_mask[-response_length:] - - group_results.append( - Sample( - index=record["instance_id"], - prompt=record["uid"], - tokens=token_ids, - response_length=response_length, - reward=record["reward"], - status=( - Sample.Status.COMPLETED - if "finish_reason" not in record["extra_info"] - or record["extra_info"]["finish_reason"] != "length" - else Sample.Status.TRUNCATED - ), - loss_mask=loss_mask, - metadata={**record["extra_info"]}, - ) + print(f"start rollout with payload: {start_inform}") + print(f"start rollout id: {rollout_id}") + rollout_started = True + + print(f"INFO: buffer length: {data_buffer.get_buffer_length()}, data_number_to_fetch: {data_number_to_fetch}") + base_url = args.rollout_buffer_url + tokenizer = AutoTokenizer.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + retry_times = 0 + results = [] + all_meta_info = [] + selected_results = [] + incomplete_groups = {} + + if args.fetch_trajectory_retry_times == -1: + print( + "⚠️ [get_rollout_data] Fetch trajectory retry times set to -1, will retry indefinitely until sufficient data is collected" + ) + while args.fetch_trajectory_retry_times == -1 or retry_times < args.fetch_trajectory_retry_times: + try: + while len(selected_results) < groups_to_fetch: + time.sleep(5) + data, meta_info = await get_rollout_data(api_base_url=base_url) + results.extend(data) + if meta_info: + all_meta_info.append(meta_info) + selected_results, incomplete_groups = _select_complete_rollout_groups( + args, results, groups_to_fetch + ) + print( + "get rollout data with " + f"{len(results)} records, {len(selected_results)} complete prompt groups" + ) + break + except Exception as err: + print(f"[get_rollout_data] Failed to get rollout data: {err}, retry times: {retry_times}") + retry_times += 1 + + log_raw_info(args, all_meta_info, rollout_id) + + if len(selected_results) < groups_to_fetch: + raise RuntimeError( + "Insufficient complete rollout groups: " + f"got {len(selected_results)}, expected {groups_to_fetch}; " + f"incomplete_groups={incomplete_groups}" + ) + if incomplete_groups: + print( + "Skipping incomplete rollout groups until they receive all " + f"{args.n_samples_per_prompt} attempts: {incomplete_groups}" ) - sample_results.append(group_results) - data_buffer.add_samples(sample_results) - final_return_results = data_buffer.get_samples(args.rollout_batch_size) # type: ignore + results = selected_results + + if len(all_meta_info) > 0 and "finished_groups" in all_meta_info[0]: + finished_groups_instance_id_list = [] + for item in all_meta_info: + finished_groups_instance_id_list.extend(item["finished_groups"]) + + data_buffer.update_metadata({str(rollout_id): finished_groups_instance_id_list}) + + print("finally get rollout data with length: ", len(results)) + sample_results = [] + + for _i, group_record in enumerate(results): + group_results = [] + for record in group_record: + if "samples" in record: + compact_samples = [Sample.from_dict(item) for item in record["samples"]] + group_results.extend(compact_samples) + continue + + oai_messages = record["messages"] + + mask_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type=args.loss_mask_type) + token_ids, loss_mask = mask_generator.get_loss_mask(oai_messages) + response_length = mask_generator.get_response_lengths([loss_mask])[0] + + loss_mask = loss_mask[-response_length:] + + group_results.append( + Sample( + index=record["instance_id"], + prompt=record["uid"], + tokens=token_ids, + response_length=response_length, + reward=record["reward"], + status=( + Sample.Status.COMPLETED + if "finish_reason" not in record["extra_info"] + or record["extra_info"]["finish_reason"] != "length" + else Sample.Status.TRUNCATED + ), + loss_mask=loss_mask, + metadata={**record["extra_info"]}, + ) + ) + if len(group_results) < args.n_samples_per_prompt: + print( + "Skipping prompt group with insufficient trainable samples: " + f"got {len(group_results)}, expected at least {args.n_samples_per_prompt}" + ) + continue + sample_results.append(group_results) - return final_return_results + if len(sample_results) < groups_to_fetch: + raise RuntimeError( + "Complete rollout groups produced too few trainable sample groups: " + f"got {len(sample_results)}, expected {groups_to_fetch}" + ) + data_buffer.add_samples(sample_results) + final_return_results = data_buffer.get_samples(args.rollout_batch_size) # type: ignore + + return final_return_results + finally: + if rollout_started: + timeout_sec = float(getattr(args, "rollout_buffer_stop_timeout_sec", 120.0)) + await asyncio.to_thread(stop_rollout, args.rollout_buffer_url, timeout_sec) def generate_rollout(args, rollout_id, data_buffer, evaluation=False): diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000000..96e85b325d --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,6 @@ +"""Local test helpers package. + +This file keeps imports such as ``tests.test_agent._fakes`` resolved to this +repository's test tree even when the active Python environment contains an +unrelated third-party package named ``tests``. +""" diff --git a/tests/test_agent/_fakes.py b/tests/test_agent/_fakes.py index 6800652de7..31c4999342 100644 --- a/tests/test_agent/_fakes.py +++ b/tests/test_agent/_fakes.py @@ -9,6 +9,10 @@ round-trips (``decode(encode(t)) == t``), so a scripted model reply survives the encode->generate->decode->parse round trip. + * :class:`RoundTrippingTokenizer` -- same, but renders assistant turns back + into the model's own surface form (tool calls + included), so a clean multi-turn chain stays + drift-free and only real drift shows up. * :class:`ScriptedTokenizer` -- pre-baked prompt-id queue + id->text decode, for adapter unit tests that assert exact ids. * :class:`FakeSGLangServer` -- a real aiohttp ``/generate`` upstream returning @@ -112,6 +116,78 @@ def apply_chat_template(self, messages, tools=None, tokenize=True, add_generatio return out +class RoundTrippingTokenizer: + """Tokenizer whose assistant rendering is reproducible from model output. + + :class:`FakeTokenizer` renders a tool call as the opaque marker + ``toolcall:``, which carries no arguments and which no model output can + ever reproduce. That is fine for tests asserting ids, but it makes every + tool-call turn look like token drift -- so it cannot be used to test whether + the manager sees a *clean* multi-turn chain. + + Here an assistant message renders back into the surface form qwen3-coder + actually generates:: + + {reasoning} {content} + V + + so a manager leaf that preserves reasoning + text + tool calls re-renders to + exactly the ids that were sampled, and any field the leaf *drops* shows up as + genuine drift. Words are whitespace-delimited (the real tokenizer is sub-word, + which changes id counts but not whether drift occurs). + """ + + def __init__(self) -> None: + self._vocab: dict[str, int] = {} + self._inv: dict[int, str] = {} + self._next = _WORD_BASE + + def _id(self, word: str) -> int: + if word not in self._vocab: + self._vocab[word] = self._next + self._inv[self._next] = word + self._next += 1 + return self._vocab[word] + + def encode(self, text: str) -> list[int]: + return [self._id(w) for w in text.split()] if text else [] + + def decode(self, ids, skip_special_tokens: bool = False) -> str: + return " ".join(self._inv[i] for i in ids if i in self._inv) + + @staticmethod + def render_assistant(message: dict) -> str: + """Assistant message -> the model's own surface form (see class docstring).""" + parts: list[str] = [] + reasoning = message.get("reasoning_content") + if reasoning: + parts.append(f" {reasoning} ") + content = message.get("content") + if isinstance(content, str) and content.strip(): + parts.append(content) + for call in message.get("tool_calls") or []: + fn = call.get("function") or {} + args = fn.get("arguments") or {} + inner = " ".join(f" {v} " for k, v in args.items()) + parts.append(f" {inner} ") + return " ".join(parts) + + def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): + out: list[int] = [] + for m in messages: + role = m.get("role", "user") + out.append(_ROLE_BEGIN.get(role, _ROLE_BEGIN["user"])) + if role == "assistant": + out.extend(self.encode(self.render_assistant(m))) + else: + content = m.get("content") + out.extend(self.encode(content if isinstance(content, str) else "")) + out.append(_ROLE_END) + if add_generation_prompt: + out.append(_GEN) + return out + + class ScriptedTokenizer: """Pre-baked prompt-id queue + id->text decode, for adapter unit tests that assert exact token sequences. ``apply_chat_template`` ignores the messages diff --git a/tests/test_agent/test_adapters.py b/tests/test_agent/test_adapters.py index 852a9cf973..98590a05a7 100644 --- a/tests/test_agent/test_adapters.py +++ b/tests/test_agent/test_adapters.py @@ -25,10 +25,10 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -from tests.test_agent._fakes import FakeSGLangServer, FakeTokenizer # noqa: E402 +from tests.test_agent._fakes import FakeSGLangServer, FakeTokenizer, RoundTrippingTokenizer # noqa: E402 from slime.agent.adapters import anthropic, openai # noqa: E402 -from slime.agent.parsing import parse_model_output, parse_xml_tool_uses # noqa: E402 +from slime.agent.parsing import ParsedModelOutput, parse_model_output, parse_xml_tool_uses # noqa: E402 from slime.utils.types import Sample # noqa: E402 NUM_GPUS = 0 @@ -131,6 +131,20 @@ def test_anthropic_translation_keeps_tool_results_thinking_and_tools(): ] +def test_openai_translation_accepts_reasoning_under_either_key(): + """CodeBuddy Code echoes reasoning back as ``reasoning``, not + ``reasoning_content``; both must normalise onto the canonical key or the echo + compares unequal to our leaf and every reasoning turn forks.""" + for key in ("reasoning_content", "reasoning"): + translated = openai._translate_messages([{"role": "assistant", "content": "ok", key: "plan"}]) + assert translated == [{"role": "assistant", "content": "ok", "reasoning_content": "plan"}], key + # canonical key wins when a client sends both + both = openai._translate_messages( + [{"role": "assistant", "content": "ok", "reasoning_content": "canonical", "reasoning": "alias"}] + ) + assert both[0]["reasoning_content"] == "canonical" + + def test_openai_translation_developer_to_system_and_tool_calls_to_dict(): translated = openai._translate_messages( [ @@ -375,11 +389,101 @@ async def run_case(): asyncio.run(run_case()) +def test_openai_manager_message_keeps_text_and_reasoning_with_tool_calls(): + """The manager leaf must carry everything the model generated. + + It is what re-renders as history next turn, so a dropped field stops the + re-render from reproducing the sampled ids -- token drift, which then either + rewrites a trained response as loss_mask=0 or forks the trajectory. See + test_openai_multiturn_thinking_tool_calls_do_not_fork for the consequence. + """ + parsed = ParsedModelOutput( + reasoning="plan", text="Let me look.", tool_uses=[{"name": "lookup", "input": {"q": "x"}}] + ) + wire_message, manager_message, wire_finish = openai._build_reply_parts(parsed, "stop") + + assert manager_message["content"] == "Let me look." + assert manager_message["reasoning_content"] == "plan" + assert manager_message["tool_calls"] == [ + {"type": "function", "function": {"name": "lookup", "arguments": {"q": "x"}}} + ] + # the wire keeps the same content so the client's echo still matches the leaf; + # arguments are a JSON string there (spec) and the id is wire-only. + assert wire_message["content"] == "Let me look." + assert wire_message["reasoning_content"] == "plan" + assert wire_finish == "tool_calls" + assert json.loads(wire_message["tool_calls"][0]["function"]["arguments"]) == {"q": "x"} + + # =========================================================================== -# §6 adapter behaviour: turn cap, mid-list system fold +# §6 adapter behaviour: turn cap, mid-list system fold, sampling precedence # =========================================================================== +def test_open_session_sampling_defaults_outrank_body(): + """A caller-set temperature must survive a harness that sends its own. + + codebuddy puts temperature=1 on every /v1/chat/completions request, so + without this precedence an --eval-temperature would never reach sglang. + """ + + async def run_case(): + async with FakeSGLangServer([[(-0.1, 601)]]) as sglang: + tok = FakeTokenizer(outputs={(601,): "ok"}) + adapter = openai.OpenAIAdapter(tokenizer=tok, sglang_url=sglang.url) + adapter.open_session("sid-sp", sampling_defaults={"temperature": 0.7, "top_p": 0.8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer sid-sp"}, + # temperature/top_k as a harness would send them; top_k is + # absent from the defaults so the body value still applies. + json={ + "model": "m", + "temperature": 1, + "top_k": 40, + "messages": [{"role": "user", "content": "hi"}], + }, + ) + finally: + await client.close() + await _drain(adapter, "sid-sp") + + sp = sglang.requests[0]["sampling_params"] + assert sp["temperature"] == 0.7, "body temperature must not override the caller's default" + assert sp["top_p"] == 0.8 + assert sp["top_k"] == 40, "keys absent from sampling_defaults still come from the body" + + asyncio.run(run_case()) + + +def test_body_sampling_params_apply_without_open_session_defaults(): + """With no caller defaults, the harness's own knobs are still honoured.""" + + async def run_case(): + async with FakeSGLangServer([[(-0.1, 602)]]) as sglang: + tok = FakeTokenizer(outputs={(602,): "ok"}) + adapter = openai.OpenAIAdapter(tokenizer=tok, sglang_url=sglang.url) + adapter.open_session("sid-nd") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer sid-nd"}, + json={"model": "m", "temperature": 0.3, "messages": [{"role": "user", "content": "hi"}]}, + ) + finally: + await client.close() + await _drain(adapter, "sid-nd") + + assert sglang.requests[0]["sampling_params"]["temperature"] == 0.3 + + asyncio.run(run_case()) + + def test_max_turns_per_sid_returns_429(): async def run_case(): async with FakeSGLangServer([[(-0.1, 501)], [(-0.1, 502)]]) as sglang: @@ -418,6 +522,82 @@ def test_mid_list_system_folds_into_user(): assert any(b.get("text", "").startswith("") for b in folded) +@pytest.mark.parametrize( + "shape,turn_text", + [ + ("tool_only", " slime "), + ( + "text_and_tool", + "Let me look. slime ", + ), + ( + "think_and_tool", + " plan slime " + " ", + ), + ], +) +def test_openai_multiturn_thinking_tool_calls_do_not_fork(shape, turn_text): + """A clean multi-turn OpenAI chain must yield ONE sample with every turn trained. + + Regression test for the adapter dropping text / reasoning from + manager_message: on a 500-instance CodeBuddy run that split each rollout into + ~25 samples (versus 1 for the Anthropic path) and silently rewrote real + responses as loss_mask=0. Both symptoms are asserted -- sample count AND + trained-token count -- because a leaf that drops a field can also keep the + count at 1 while training almost nothing. + """ + turns = 4 + + async def run_case(): + tok = RoundTrippingTokenizer() + scripted = [[(-0.1, tid) for tid in tok.encode(turn_text)] for _ in range(turns)] + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + }, + } + ] + async with FakeSGLangServer(scripted) as sglang: + # tool_parser=None routes through the XML fallback, so no sglang needed; + # reasoning_parser=None keeps inside .text, which for this test + # is equivalent -- the point is whether the leaf preserves what was + # generated, not which field it lands in. + adapter = openai.OpenAIAdapter(tokenizer=tok, sglang_url=sglang.url) + sid = f"sid-{shape}" + adapter.open_session(sid) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + messages: list[dict] = [{"role": "user", "content": "fix it"}] + try: + for _ in range(turns): + resp = await client.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {sid}"}, + json={"model": "m", "tools": tools, "messages": messages}, + ) + assert resp.status == 200 + assistant = (await resp.json())["choices"][0]["message"] + # echo history back verbatim, as the real CLI was measured to do + messages = [*messages, assistant] + for call in assistant.get("tool_calls") or []: + messages.append({"role": "tool", "tool_call_id": call["id"], "content": "ok"}) + finally: + await client.close() + return await _drain(adapter, sid) + + samples = asyncio.run(run_case()) + assert len(samples) == 1, f"{shape}: chain forked into {len(samples)} samples" + response_tokens = len(tok_ids := samples[0].loss_mask) + trained = sum(tok_ids) + # every turn's response is trained; the only untrained ids are the tool results + # threaded back in as prompt between turns. + assert trained > 0.5 * response_tokens, f"{shape}: only {trained}/{response_tokens} response tokens trained" + + # =========================================================================== # §7 parsing helpers (slime.agent.parsing) # =========================================================================== diff --git a/tests/test_agent/test_harness.py b/tests/test_agent/test_harness.py index 2fb1e38bd3..c1dfb370d9 100644 --- a/tests/test_agent/test_harness.py +++ b/tests/test_agent/test_harness.py @@ -71,6 +71,7 @@ async def fake_agent(env): assert any("run.sh" in p for p in sb.files) assert _find(sb.exec_log, "setsid") assert any("echo $?" in v for v in sb.files.values()) + assert not _find(sb.exec_log, "kill -TERM") asyncio.run(run_case()) @@ -94,6 +95,8 @@ async def run_case(): with patch.object(hc.asyncio, "sleep", new=_fast_sleep): rc = await hc.run_agent(sb, workdir="/w", start_cmd="x", env={}, time_budget_sec=0) assert rc == sandbox_mod.EXIT_TIME_BUDGET_EXCEEDED + assert _find(sb.exec_log, "kill -TERM") + assert _find(sb.exec_log, "kill -KILL") asyncio.run(run_case()) diff --git a/tests/test_dp_schedule.py b/tests/test_dp_schedule.py index 33a8be53b5..a82624b66f 100644 --- a/tests/test_dp_schedule.py +++ b/tests/test_dp_schedule.py @@ -57,12 +57,18 @@ def assert_invariants( expected_global_sample_indices, total_lengths, max_per_bin=None, + allow_merged_over_cap=False, ): """Check the invariants documented at the top of dp_schedule.py. ``expected_global_sample_indices`` is the set of global sample indices that should end up covered (after trim). Trailing rollouts that don't fit are excluded. + + ``allow_merged_over_cap`` relaxes the per-mbs token cap to allow multi-sample + mbs above it. Pass it for steps that had to merge bins down to satisfy the + ``dp_size`` alignment (the path where every bin was an unsplittable singleton); + on that path the cap is a target rather than a bound. """ seen_global: set[int] = set() for r in range(dp_size): @@ -84,12 +90,13 @@ def assert_invariants( if max_per_bin is None: return - # Every mbs <= max_per_bin tokens, EXCEPT a singleton bin holding an oversized sample. + # Every mbs <= max_per_bin tokens, EXCEPT a singleton bin holding an oversized + # sample (or, when allow_merged_over_cap, an mbs produced by alignment merging). for r in range(dp_size): partition = partitions[r] for mbs in micro_batch_indices[r]: bin_total = sum(total_lengths[partition[i]] for i in mbs) - if bin_total > max_per_bin: + if bin_total > max_per_bin and not (allow_merged_over_cap and len(mbs) > 1): assert len(mbs) == 1, f"rank {r}: mbs sum {bin_total} > {max_per_bin} but contains {len(mbs)} samples" @@ -322,5 +329,241 @@ def test_rejects_when_fewer_rollouts_than_gbs(): build_dp_schedule(args, tp, [3] * 6, global_batch_size=4, rollout_indices=[0, 0, 1, 1, 2, 2]) +# --------------------------------------------------------------------------- +# Alignment when no bin can be split: long-context runs where one sample alone +# fills max_per_bin, so first-fit emits one bin per sample. An odd sample count +# can then never be rounded UP to an even target_K, and the schedule must round +# down by merging instead of failing. +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_odd_unsplittable_singletons_merge_instead_of_asserting(): + """Every sample fills a whole bin, and there is an odd number of them. + + Splitting is impossible (all bins are singletons), so the schedule merges the two + smallest bins to reach an even mbs count rather than raising. Every sample must + still be placed. + """ + total_lengths = [10] * 9 + rollout_indices = list(range(9)) + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=10) + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=9, rollout_indices=rollout_indices + ) + + # 9 singleton bins -> merged down to 8 -> 4 mbs per rank. + assert nmb == [4] + assert gbs_per_step == [9] + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(9), + total_lengths=total_lengths, + max_per_bin=10, + allow_merged_over_cap=True, + ) + # Exactly one mbs holds two samples; the rest stay singletons. + sizes = sorted(len(mbs) for r in range(2) for mbs in mbi[r]) + assert sizes == [1, 1, 1, 1, 1, 1, 1, 2], sizes + + +@pytest.mark.unit +@pytest.mark.parametrize("dp_size,num_samples", [(2, 3), (4, 5), (4, 7), (8, 9)]) +def test_unsplittable_alignment_across_dp_sizes(dp_size, num_samples): + """Merging down aligns to dp_size for any misaligned all-singleton bin count.""" + total_lengths = [10] * num_samples + rollout_indices = list(range(num_samples)) + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=10) + tp = make_tp(dp_size=dp_size) + + partitions, mbi, nmb, _ = build_dp_schedule( + args, tp, total_lengths, global_batch_size=num_samples, rollout_indices=rollout_indices + ) + + total_mbs = sum(len(mbi[r]) for r in range(dp_size)) + assert total_mbs % dp_size == 0, f"{total_mbs} mbs not aligned to dp_size {dp_size}" + assert_invariants( + partitions, + mbi, + nmb, + dp_size=dp_size, + expected_global_sample_indices=range(num_samples), + total_lengths=total_lengths, + max_per_bin=10, + allow_merged_over_cap=True, + ) + + +# Real per-sample token lengths and rollout ids captured from the rollout dumps of run +# nspr3_stress_qwen35_35b_a3b_4nodes_20260810 (Qwen3.5-35B-A3B, 4 nodes, TP2/CP8 -> +# dp_size 2, max_tokens_per_gpu 12000 -> max_per_bin 96000). Rollout 1 crashed the +# unfixed scheduler with "could only produce 39 mbs ...; need 40" at step 2; rollout 0 +# happened to survive. 48 rollouts each, but 77 / 74 training samples, because compact +# segmentation lets one rollout emit several samples. +# +# Stored as whitespace-separated strings rather than list literals purely so the +# formatter keeps them on one line instead of exploding to one element per line. +PROD_ROLLOUT0_LENGTHS = """ +41660 48307 44888 37079 44934 46423 49780 55550 37387 46094 38929 34033 34547 39827 46419 50051 48041 51804 +73102 46217 52259 72876 86641 44319 56019 46280 66098 86083 96000 53259 62065 65596 56250 73241 59030 84499 +80569 88055 96000 96000 96000 96000 96000 96000 96000 96000 96000 81223 73704 96000 96000 96000 96000 96000 +95456 96000 96000 96000 96000 96000 96000 90658 96000 96000 96000 44700 48326 59807 62062 44375 49897 61361 +76566 68744 48080 47574 47846 +""" +PROD_ROLLOUT0_ROLLOUT_IDS = """ +1 5 0 7 7 2 6 4 3 3 10 13 8 15 12 11 14 9 44 46 42 47 40 45 45 41 41 43 28 29 27 26 25 25 30 31 24 22 20 20 +20 20 19 19 21 21 21 17 17 16 16 16 16 18 18 18 18 23 23 23 23 23 23 23 23 34 33 36 36 38 37 35 35 39 32 32 +32 +""" +PROD_ROLLOUT1_LENGTHS = """ +41582 76811 51088 70564 35093 67817 51789 46953 60184 62764 70756 40596 40050 40142 41407 40828 41753 42428 +53024 42705 46273 32293 40500 42337 37544 32540 49951 32752 44972 49089 49369 46050 56614 58384 72760 77339 +51852 60275 58984 78804 96000 96000 83523 96000 96000 95609 96000 96000 92476 95364 96000 96000 69612 96000 +96000 96000 49568 57135 82158 59784 72398 90316 96000 96000 46691 59485 83890 80286 96000 96000 96000 96000 +96000 96000 +""" +PROD_ROLLOUT1_ROLLOUT_IDS = """ +74 78 73 72 79 79 75 75 77 77 76 80 81 83 82 84 86 87 85 91 88 89 94 90 95 93 92 63 58 57 61 60 62 56 59 66 +70 68 67 64 69 69 71 65 65 54 54 54 53 53 53 53 55 55 55 55 51 51 51 48 48 48 48 48 52 52 52 49 49 49 49 49 +50 50 +""" + + +def _ints(blob: str) -> list[int]: + return [int(tok) for tok in blob.split()] + + +@pytest.mark.unit +@pytest.mark.parametrize( + "lengths_blob,rollout_ids_blob", + [ + (PROD_ROLLOUT0_LENGTHS, PROD_ROLLOUT0_ROLLOUT_IDS), + (PROD_ROLLOUT1_LENGTHS, PROD_ROLLOUT1_ROLLOUT_IDS), + ], + ids=["rollout0", "rollout1"], +) +def test_production_long_context_rollouts_schedule(lengths_blob, rollout_ids_blob): + """Regression: the real 96k-cap SWE rollouts that crashed the scheduler. + + num_steps_per_rollout=3 over 48 rollouts -> global_batch_size 16. Rollout 1's third + step held 39 samples in 39 unsplittable bins and used to raise. + """ + total_lengths = _ints(lengths_blob) + rollout_indices = _ints(rollout_ids_blob) + assert len(total_lengths) == len(rollout_indices), "fixture lengths/ids out of sync" + assert len(set(rollout_indices)) == 48, "fixture should hold 48 distinct rollouts" + + max_per_bin = 12000 * 8 + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=12000) + tp = make_tp(dp_size=2, cp_size=8) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=16, rollout_indices=rollout_indices + ) + + assert gbs_per_step == [16, 16, 16], "3 steps of 16 rollouts each" + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(len(total_lengths)), + total_lengths=total_lengths, + max_per_bin=max_per_bin, + allow_merged_over_cap=True, + ) + + # Repacking is a last resort, so the overflow must stay marginal. Anything near 2x + # the cap would risk OOM in real training. + worst = max(sum(total_lengths[partitions[r][i]] for i in mbs) for r in range(2) for mbs in mbi[r]) + assert worst <= int(max_per_bin * 1.05), f"worst mbs {worst} exceeds the cap by more than 5%" + + +# --------------------------------------------------------------------------- +# Alignment fallback: failure mode and peak-memory behaviour. +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_rejects_when_fewer_mbs_than_alignment_threshold(): + """K < align_to cannot be aligned in either direction, and must say so. + + Splitting is impossible (all singletons) and rounding down lands on 0 mbs, so the + step genuinely has no valid schedule. The error must name the real cause — too few + samples for the alignment threshold — rather than blaming the repack. + """ + total_lengths = [10] * 3 # each fills max_per_bin -> 3 unsplittable singleton bins + rollout_indices = [0, 1, 2] + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=10) + # vpp_size=2 with mb_group=2 makes align_to = dp_size * mb_group = 4 > 3 bins. + tp = make_tp(dp_size=2, vpp_size=2, microbatch_group_size_per_vp_stage=2) + + with pytest.raises(AssertionError, match="below the alignment threshold"): + build_dp_schedule(args, tp, total_lengths, global_batch_size=3, rollout_indices=rollout_indices) + + +@pytest.mark.unit +def test_repack_minimises_the_largest_mbs(): + """The repack must minimise the PEAK mbs, not just merge the smallest bins. + + Merging the two smallest bins repeatedly is the intuitive rule and is worse as soon + as more than one merge is needed, because it keeps stacking onto the same bin. With + align_to=4 and 7 singleton bins (3 merges) the difference is large enough to decide + whether a real 96k-cap step OOMs. + """ + # Every length is above max_per_bin/2 so no two samples ever share a bin. + total_lengths = [49_000, 49_000, 49_000, 90_000, 90_000, 90_000, 90_000] + rollout_indices = list(range(7)) + max_per_bin = 96_000 + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=max_per_bin) + tp = make_tp(dp_size=4) # align_to = 4, so 7 bins must become 4 + + partitions, mbi, nmb, _ = build_dp_schedule( + args, tp, total_lengths, global_batch_size=7, rollout_indices=rollout_indices + ) + + assert sum(len(mbi[r]) for r in range(4)) == 4, "7 singleton bins should repack to 4" + assert_invariants( + partitions, + mbi, + nmb, + dp_size=4, + expected_global_sample_indices=range(7), + total_lengths=total_lengths, + max_per_bin=max_per_bin, + allow_merged_over_cap=True, + ) + + worst = max(sum(total_lengths[partitions[r][i]] for i in mbs) for r in range(4) for mbs in mbi[r]) + # Optimal here is 139000 (90000+49000). The merge-two-smallest rule would stack all + # three 49000s into one bin and peak at 147000; anything at or above that means the + # peak-minimising property regressed. + assert worst <= 139_000, f"peak mbs {worst} above the achievable minimum 139000" + + +@pytest.mark.unit +def test_single_merge_case_still_takes_the_smallest_pair(): + """With exactly one merge needed (the dp_size=2 production case), the peak-minimising + repack must still combine the two smallest bins — that is optimal there, and it is + what keeps the measured production overflow at ~0.3% of the cap.""" + total_lengths = [60_000, 61_000, 95_000, 96_000, 96_000] + rollout_indices = list(range(5)) + max_per_bin = 96_000 + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=max_per_bin) + tp = make_tp(dp_size=2) # align_to = 2, so 5 bins become 4 + + partitions, mbi, nmb, _ = build_dp_schedule( + args, tp, total_lengths, global_batch_size=5, rollout_indices=rollout_indices + ) + + merged = [sorted(total_lengths[partitions[r][i]] for i in mbs) for r in range(2) for mbs in mbi[r] if len(mbs) > 1] + assert merged == [[60_000, 61_000]], f"expected the two smallest bins merged, got {merged}" + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_harbor_task_to_slime_prompt_data.py b/tests/test_harbor_task_to_slime_prompt_data.py new file mode 100644 index 0000000000..8e487e091a --- /dev/null +++ b/tests/test_harbor_task_to_slime_prompt_data.py @@ -0,0 +1,150 @@ +"""Regression tests for Harbor-to-slime prompt conversion.""" + +from __future__ import annotations + +import ast +import base64 +import gzip +import importlib.util +import json +import sys +from pathlib import Path + +import pytest + + +@pytest.fixture(scope="module") +def converter_module(): + path = Path(__file__).parents[1] / "tools" / "harbor_task_to_slime_prompt_data.py" + spec = importlib.util.spec_from_file_location("harbor_task_to_slime_prompt_data", path) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def harbor_task(tmp_path: Path) -> Path: + task = tmp_path / "sample-task" + (task / "environment").mkdir(parents=True) + (task / "tests").mkdir() + (task / "instruction.md").write_text("Fix the bug.") + (task / "task.toml").write_text("[task]\nname = 'sample-task'\n") + (task / "environment" / "Dockerfile").write_text("FROM registry.example/swe:sample\nWORKDIR /testbed\n") + (task / "tests" / "config.json").write_text( + json.dumps( + { + "instance_id": "sample__1", + "repo": "example/sample", + "base_commit": "deadbeef", + "problem_statement": "Fix the bug.", + "FAIL_TO_PASS": [], + "PASS_TO_PASS": [], + } + ) + ) + (task / "tests" / "test.sh").write_text( + "#!/bin/bash\npython -m pip install -e .[test] --verbose\npytest tests/test_sample.py || true\n" + ) + return task + + +def _embedded_test_script(eval_cmd: str) -> str: + payloads_text = eval_cmd.split("payloads = [", 1)[1].split("]", 1)[0] + payloads = ast.literal_eval(f"[{payloads_text}]") + return gzip.decompress(base64.b64decode(payloads[1])).decode("utf-8") + + +def test_converter_preserves_image_head_by_default(converter_module, harbor_task: Path): + row = converter_module.task_to_row( + harbor_task, + dataset_root=harbor_task.parent, + source="test", + input_key="prompt", + prompt_alias_key="", + label_key="label", + metadata_key="metadata", + image_override=None, + default_workdir="/testbed", + include_eval_cmd=True, + include_inline_files=False, + inline_files=(), + provenance_root=False, + ) + + metadata = row["metadata"] + assert "pre_commands" not in metadata + assert _embedded_test_script(metadata["eval_cmd"]) == (harbor_task / "tests" / "test.sh").read_text() + + +def test_prompt_source_option_is_removed(converter_module, monkeypatch: pytest.MonkeyPatch): + """--prompt-source no longer exists; instruction.md is now unconditional. + + The flag's only other value fed the raw tests/config.json text to the agent, + which keeps upstream CRLF on ~50% of SWE-bench Verified and drops Harbor's + header and provenance block. Accepting the flag silently would let an old + command line look like it still selected that behaviour. + """ + monkeypatch.setattr( + sys, + "argv", + ["x", "--input", ".", "--output", "out.jsonl", "--prompt-source", "problem_statement"], + ) + with pytest.raises(SystemExit): + converter_module.parse_args() + + +def test_prompt_and_problem_statement_are_harbor_instruction(converter_module, harbor_task: Path): + """Prompt and metadata.problem_statement are both instruction.md; the raw + tests/config.json text survives under metadata.harbor.problem_statement.""" + instruction = "# Task\n\nFix the bug.\n\n---\n\n**Repo:** example/sample\n" + (harbor_task / "instruction.md").write_text(instruction) + row = converter_module.task_to_row( + harbor_task, + dataset_root=harbor_task.parent, + source="test", + input_key="prompt", + prompt_alias_key="", + label_key="label", + metadata_key="metadata", + image_override=None, + default_workdir="/testbed", + include_eval_cmd=True, + include_inline_files=False, + inline_files=(), + provenance_root=False, + ) + assert row["prompt"] == instruction + assert row["metadata"]["problem_statement"] == instruction + # the raw upstream field differs from instruction.md and must not be lost + assert row["metadata"]["harbor"]["problem_statement"] == "Fix the bug." + + +def test_reset_to_base_commit_option_is_removed(converter_module, monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + input_path = tmp_path / "input" + output_path = tmp_path / "output.jsonl" + + monkeypatch.setattr( + sys, + "argv", + [ + "harbor_task_to_slime_prompt_data.py", + "--input", + str(input_path), + "--output", + str(output_path), + "--reset-to-base-commit", + ], + ) + + with pytest.raises(SystemExit): + converter_module.parse_args() + + +def test_gzip_payload_has_zero_mtime_and_is_reproducible(converter_module): + first = base64.b64decode(converter_module._gzip_base64("same content")) + second = base64.b64decode(converter_module._gzip_base64("same content")) + + assert first == second + assert first[4:8] == b"\0\0\0\0" + assert gzip.decompress(first) == b"same content" diff --git a/tests/test_passrate_metrics.py b/tests/test_passrate_metrics.py new file mode 100644 index 0000000000..75396b2e58 --- /dev/null +++ b/tests/test_passrate_metrics.py @@ -0,0 +1,97 @@ +import pytest + +from slime.utils.metric_utils import compute_grouped_pass_rate, compute_pass_rate + +NUM_GPUS = 0 + + +@pytest.mark.unit +def test_compute_pass_rate_fixed_shape(): + metrics = compute_pass_rate( + flat_rewards=[ + 1, + 0, + 0, + 0, + 1, + 1, + 0, + 0, + ], + group_size=4, + num_groups=2, + ) + + assert metrics["pass@1"] == pytest.approx(0.375) + assert metrics["pass@2"] == pytest.approx((0.5 + 5 / 6) / 2) + assert metrics["pass@4"] == pytest.approx(1.0) + + +@pytest.mark.unit +def test_compute_pass_rate_skips_mismatched_fixed_shape(): + assert compute_pass_rate(flat_rewards=[0] * 23, group_size=4, num_groups=4) == {} + + +@pytest.mark.unit +def test_compute_grouped_pass_rate_deduplicates_fanout_segments(): + # Two prompt groups, four rollout attempts per prompt. Rollout attempts 0 + # and 6 fan out into multiple train samples; they should still count as one + # independent pass@k sample each. + metrics = compute_grouped_pass_rate( + flat_rewards=[ + 1, + 1, # group 0, rollout 0 fan-out sibling + 0, + 0, + 0, + 0, + 0, + 1, + 1, # group 1, rollout 6 fan-out sibling + 1, + ], + group_indices=[ + 0, + 0, + 0, + 0, + 0, + 1, + 1, + 1, + 1, + 1, + ], + rollout_ids=[ + 0, + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 6, + 7, + ], + group_size=4, + ) + + assert metrics["pass@1"] == pytest.approx(0.375) + assert metrics["pass@2"] == pytest.approx((0.5 + 5 / 6) / 2) + assert metrics["pass@4"] == pytest.approx(1.0) + + +@pytest.mark.unit +def test_compute_grouped_pass_rate_skips_k_when_group_has_too_few_attempts(): + metrics = compute_grouped_pass_rate( + flat_rewards=[1, 0, 1, 0, 0], + group_indices=[0, 0, 1, 1, 1], + rollout_ids=[0, 1, 2, 3, 4], + group_size=4, + ) + + # pass@4 is not estimable because no group has four independent attempts. + assert "pass@4" not in metrics + assert metrics["pass@1"] == pytest.approx((0.5 + 1 / 3) / 2) + assert metrics["pass@2"] == pytest.approx((1.0 + 2 / 3) / 2) diff --git a/tests/test_rollout_buffer/test_ags_empty_patch_guard.py b/tests/test_rollout_buffer/test_ags_empty_patch_guard.py new file mode 100644 index 0000000000..2cbf5f3786 --- /dev/null +++ b/tests/test_rollout_buffer/test_ags_empty_patch_guard.py @@ -0,0 +1,239 @@ +"""Unit tests for the AGS empty-patch guardrail and its config knob. + +The guard flags rollouts where the coding agent exited 0 but left no diff -- +the signature of a CLI that accepted a text-only mid-task turn as a finished +answer. Trajectories here are written as real CodeBuddy Code stream-json JSONL +(``{"type": "assistant", "message": {...}}`` / ``{"type": "result", ...}``) so +the tests exercise the same reader the runner uses. + +Only ``empty_patch_guard`` and ``config`` are loaded, and they are loaded from +their files rather than through the package: ``rollout_buffer.generator.__init__`` +imports ``openai`` and ``ags_generator.__init__`` imports the whole entry module, +neither of which the CPU CI jobs install. The sibling ``test_ags_generator.py`` +does import through the package, which is why it cannot run in those jobs. +""" + +from __future__ import annotations + +import importlib +import json +import sys +import types +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +AGS_DIR = REPO_ROOT / "slime_plugins" / "rollout_buffer" / "generator" / "ags_generator" +_PRIVATE_PACKAGE = "_ags_guard_under_test" + + +def _load_ags_module(name: str): + """Load one ags_generator module through a private synthetic package. + + The modules are mounted under ``_ags_guard_under_test`` -- a package whose + ``__path__`` is the real ags_generator directory -- rather than under their + true dotted path. That keeps relative imports working (``empty_patch_guard`` + does ``from .weave_trace import ...``) while leaving the real + ``slime_plugins...`` entries in ``sys.modules`` untouched: stubbing those + ancestors would linger for the whole pytest session and break sibling + modules that import through the genuine package. + """ + if _PRIVATE_PACKAGE not in sys.modules: + package = types.ModuleType(_PRIVATE_PACKAGE) + package.__path__ = [str(AGS_DIR)] + sys.modules[_PRIVATE_PACKAGE] = package + return importlib.import_module(f".{name}", package=_PRIVATE_PACKAGE) + + +_load_ags_module("weave_trace") # empty_patch_guard imports it relatively +guard = _load_ags_module("empty_patch_guard") +AGSGeneratorConfig = _load_ags_module("config").AGSGeneratorConfig + +NUM_GPUS = 0 + +AGENT = "codebuddy_code" + + +def _write_trajectory(tmp_path: Path, *, result_text: str | None, assistant_text: str | None = None) -> str: + """Write a minimal CBC-shaped stream-json trajectory and return its path.""" + events: list[dict] = [{"type": "system", "subtype": "init", "session_id": "s"}] + if assistant_text is not None: + events.append( + { + "type": "assistant", + "message": {"role": "assistant", "content": [{"type": "text", "text": assistant_text}]}, + "__timestamp": "2026-07-27T00:00:00.000Z", + } + ) + if result_text is not None: + events.append( + { + "type": "result", + "subtype": "success", + "is_error": False, + "result": result_text, + "__timestamp": "2026-07-27T00:00:01.000Z", + } + ) + path = tmp_path / "agent.trajectory.jsonl" + path.write_text("\n".join(json.dumps(e) for e in events) + "\n", encoding="utf-8") + return str(path) + + +def test_guard_flags_midtask_narration(tmp_path): + # verbatim-style mid-task stop from a real CBC run + path = _write_trajectory( + tmp_path, result_text="Let me look at the query.py file to understand values().<|im_end|>" + ) + verdict = guard.classify_empty_patch(agent_exit_code=0, diff_text="", trajectory_path=path, agent=AGENT) + assert verdict.triggered + assert verdict.reason == guard.REASON_MIDTASK_NARRATION + assert "query.py" in verdict.final_text + + +def test_guard_flags_claimed_completion_without_diff(tmp_path): + # Claiming success with nothing to show is a distinct failure mode from + # stopping mid-task, so it must not collapse into the narration bucket. + path = _write_trajectory(tmp_path, result_text="The fix is complete. All tests pass.") + verdict = guard.classify_empty_patch(agent_exit_code=0, diff_text="", trajectory_path=path, agent=AGENT) + assert verdict.triggered + assert verdict.reason == guard.REASON_CLAIMED_COMPLETE + + +def test_guard_prefers_completion_claim_over_trailing_narration(tmp_path): + path = _write_trajectory(tmp_path, result_text="All tests pass. Let me provide a summary.") + verdict = guard.classify_empty_patch(agent_exit_code=0, diff_text="", trajectory_path=path, agent=AGENT) + assert verdict.reason == guard.REASON_CLAIMED_COMPLETE + + +def test_guard_flags_blank_final_text(tmp_path): + path = _write_trajectory(tmp_path, result_text=" ") + verdict = guard.classify_empty_patch(agent_exit_code=0, diff_text="", trajectory_path=path, agent=AGENT) + assert verdict.triggered + assert verdict.reason == guard.REASON_NO_FINAL_TEXT + + +def test_guard_falls_back_to_last_assistant_text(tmp_path): + path = _write_trajectory(tmp_path, result_text=None, assistant_text="Let me check the tests first.") + verdict = guard.classify_empty_patch(agent_exit_code=0, diff_text="", trajectory_path=path, agent=AGENT) + assert verdict.reason == guard.REASON_MIDTASK_NARRATION + + +def test_guard_ignores_nonempty_diff(tmp_path): + path = _write_trajectory(tmp_path, result_text="Let me look at query.py.") + verdict = guard.classify_empty_patch( + agent_exit_code=0, + diff_text="--- a/x.py\n+++ b/x.py\n@@ -1 +1 @@\n-a\n+b\n", + trajectory_path=path, + agent=AGENT, + ) + assert not verdict.triggered + assert verdict.reason == "" + + +def test_guard_ignores_whitespace_only_diff_as_empty(tmp_path): + path = _write_trajectory(tmp_path, result_text="Let me look at query.py.") + verdict = guard.classify_empty_patch(agent_exit_code=0, diff_text=" \n\t\n", trajectory_path=path, agent=AGENT) + assert verdict.triggered + + +def test_guard_ignores_nonzero_exit(tmp_path): + path = _write_trajectory(tmp_path, result_text="Let me look at query.py.") + verdict = guard.classify_empty_patch(agent_exit_code=1, diff_text="", trajectory_path=path, agent=AGENT) + assert not verdict.triggered + + +def test_guard_triggers_without_trajectory_path(): + # TRAJECTORY_DUMP_DIR unset makes dump_trajectory return None; an empty patch + # is still an empty patch, we just cannot attribute it. + verdict = guard.classify_empty_patch(agent_exit_code=0, diff_text="", trajectory_path=None, agent=AGENT) + assert verdict.triggered + assert verdict.reason == guard.REASON_UNCLASSIFIED + + +def test_guard_is_inert_for_agent_without_parser(tmp_path): + path = _write_trajectory(tmp_path, result_text="Let me look at query.py.") + verdict = guard.classify_empty_patch(agent_exit_code=0, diff_text="", trajectory_path=path, agent="codex") + assert verdict.triggered + assert verdict.reason == guard.REASON_UNCLASSIFIED + + +def test_guard_survives_malformed_trajectory(tmp_path): + path = tmp_path / "bad.trajectory.jsonl" + path.write_text('not json\n{"type": "result"\n\n', encoding="utf-8") + verdict = guard.classify_empty_patch(agent_exit_code=0, diff_text="", trajectory_path=str(path), agent=AGENT) + assert verdict.triggered # no exception + + +def test_guard_survives_missing_trajectory_file(tmp_path): + verdict = guard.classify_empty_patch( + agent_exit_code=0, diff_text="", trajectory_path=str(tmp_path / "absent.jsonl"), agent=AGENT + ) + assert verdict.triggered + + +def test_empty_patch_guard_policy_defaults_to_metrics(monkeypatch): + monkeypatch.delenv("SWE_EMPTY_PATCH_GUARD", raising=False) + assert AGSGeneratorConfig.from_env().empty_patch_guard == "metrics" + + +@pytest.mark.parametrize("value,expected", [("off", "off"), ("abort", "abort"), ("METRICS", "metrics")]) +def test_empty_patch_guard_policy_from_env(monkeypatch, value, expected): + monkeypatch.setenv("SWE_EMPTY_PATCH_GUARD", value) + assert AGSGeneratorConfig.from_env().empty_patch_guard == expected + + +def test_empty_patch_guard_policy_rejects_unknown_value(monkeypatch): + monkeypatch.setenv("SWE_EMPTY_PATCH_GUARD", "retry") + with pytest.raises(ValueError, match="SWE_EMPTY_PATCH_GUARD"): + AGSGeneratorConfig.from_env() + + +def test_prompt_styles_default_to_instruction_for_train_and_dataset_for_eval(monkeypatch): + monkeypatch.delenv("SWE_PROMPT_STYLE", raising=False) + monkeypatch.delenv("SWE_EVAL_PROMPT_STYLE", raising=False) + config = AGSGeneratorConfig.from_env() + assert config.prompt_style == "instruction" + assert config.eval_prompt_style == "dataset" + + +@pytest.mark.parametrize( + "env_name,attr", [("SWE_PROMPT_STYLE", "prompt_style"), ("SWE_EVAL_PROMPT_STYLE", "eval_prompt_style")] +) +@pytest.mark.parametrize( + "value,expected", [("dataset", "dataset"), ("DATASET", "dataset"), ("instruction", "instruction")] +) +def test_prompt_style_from_env(monkeypatch, env_name, attr, value, expected): + monkeypatch.setenv(env_name, value) + assert getattr(AGSGeneratorConfig.from_env(), attr) == expected + + +def test_prompt_styles_are_independent(monkeypatch): + """The two knobs must not read each other's env var: sharing one would make + the train/eval split silently collapse to whichever was set.""" + monkeypatch.setenv("SWE_PROMPT_STYLE", "dataset") + monkeypatch.setenv("SWE_EVAL_PROMPT_STYLE", "instruction") + config = AGSGeneratorConfig.from_env() + assert (config.prompt_style, config.eval_prompt_style) == ("dataset", "instruction") + assert config.prompt_style_for(evaluation=False) == "dataset" + assert config.prompt_style_for(evaluation=True) == "instruction" + + +@pytest.mark.parametrize("env_name", ["SWE_PROMPT_STYLE", "SWE_EVAL_PROMPT_STYLE"]) +@pytest.mark.parametrize("value", ["harbor", "inline"]) +def test_prompt_style_rejects_unknown_value(monkeypatch, env_name, value): + # "inline" was removed: it rebuilt the prompt from the raw problem_statement, + # which "dataset" now supersedes. Failing loudly beats silently changing the + # prompt for a run that still asks for it. The error must name the var that + # was actually wrong, or a typo in one sends you looking at the other. + monkeypatch.setenv(env_name, value) + with pytest.raises(ValueError, match=env_name): + AGSGeneratorConfig.from_env() + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/test_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py new file mode 100644 index 0000000000..9135acfc5d --- /dev/null +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -0,0 +1,1170 @@ +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import json +import re +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest +from aiohttp import web + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +# rollout.generate() uses asyncio.timeout(), a 3.11+ API, and swallows the +# resulting AttributeError into an abort -- so on 3.10 the tests below fail on a +# missing capture rather than on the real cause. CI pins 3.10, so shim it onto a +# pass-through: the wall-clock guard never fires here (every case finishes well +# under rollout_guard_sec). Mirrors the shim in test_agent/test_agent_rollout_cpu.py. +if not hasattr(asyncio, "timeout"): + + @contextlib.asynccontextmanager + async def _timeout_shim(_delay): + yield + + asyncio.timeout = _timeout_shim + +from tests.test_agent._fakes import FakeSandbox # noqa: E402 + +from slime.utils.misc import SingletonMeta +from slime.utils.types import Sample +from slime_plugins.rollout_buffer.generator.ags_generator import adapter_service, swe_task +from slime_plugins.rollout_buffer.generator.ags_generator.config import AGSGeneratorConfig +from slime_plugins.rollout_buffer.generator.ags_generator.entry import ( + _collapse_eval_samples, + get_group_data_meta_info, + is_valid_group, + transform_group, +) +from slime_plugins.rollout_buffer.generator.ags_generator.harnesses import ( + AGSSidecarClaudeCodeHarness, + CodeBuddyCodeHarness, + resolve_agent, +) +from slime_plugins.rollout_buffer.generator.ags_generator.rollout import AGSRolloutRunner +from slime_plugins.rollout_buffer.generator.ags_generator.runner import run_root_command +from slime_plugins.rollout_buffer.generator.ags_generator.sampling import normalize_sampling_params +from slime_plugins.rollout_buffer.generator.ags_generator.serialization import ( + output_item_from_samples, + samples_from_payload, +) +from slime_plugins.rollout_buffer.generator.ags_generator.weave_trace import AGSWeaveTrace, iter_trajectory_events +from slime_plugins.rollout_buffer.rollout_buffer_example import start_rollout + + +def _sample(*, reward=1.0, status=Sample.Status.COMPLETED): + return Sample( + index=3, + group_index=1, + rollout_id=3, + prompt="p", + tokens=[1, 2, 3], + response_length=2, + loss_mask=[1, 1], + rollout_log_probs=[0.0, 0.0], + reward=reward, + status=status, + metadata={"trajectory_path": "/tmp/t.jsonl", "patch_path": "/tmp/p.patch", "rollout_dump_path": "/tmp/r.json"}, + ) + + +def test_output_item_round_trips_compact_samples(): + samples = [_sample(), _sample(reward=1.0)] + item = output_item_from_samples(samples, instance_id="inst-1") + + restored = samples_from_payload(item) + + assert len(restored) == 2 + assert restored[0].status == Sample.Status.COMPLETED + assert restored[0].reward == 1.0 + assert restored[0].metadata["patch_path"] == "/tmp/p.patch" + + +def test_group_hooks_accept_complete_sample_payloads(): + item = output_item_from_samples([_sample()], instance_id="inst-1") + group = ("inst-1", [item]) + + assert is_valid_group(group, min_valid_group_size=1) + assert transform_group(group) is group + + meta = get_group_data_meta_info({"inst-1": [item]}) + assert meta["total_samples"] == 1 + assert meta["avg_reward"] == 1.0 + assert meta["nonzero_reward_samples"] == 1 + assert meta["artifact_counts"] == {"trajectory": 1, "patch": 1, "rollout_dump": 1, "complete": 1} + + +def test_eval_collapse_keeps_one_scored_sample_per_attempt(): + base = Sample(index=7, prompt="p", metadata={"dataset": "eval"}) + segments = [ + _sample(reward=1.0), + _sample(reward=1.0), + ] + segments[0].metadata = {**segments[0].metadata, "instance_id": "inst-1"} + + collapsed = _collapse_eval_samples(base, segments) + + assert collapsed.reward == 1.0 + assert collapsed.status == Sample.Status.COMPLETED + assert collapsed.remove_sample is True + assert collapsed.tokens == [0, 0] + assert collapsed.response_length == 1 + assert collapsed.metadata["dataset"] == "eval" + assert collapsed.metadata["instance_id"] == "inst-1" + assert collapsed.metadata["eval_collapsed_segments"] == 2 + + +def test_eval_collapse_marks_empty_output_aborted(): + collapsed = _collapse_eval_samples(Sample(index=7, metadata={"dataset": "eval"}), []) + + assert collapsed.reward == 0.0 + assert collapsed.status == Sample.Status.ABORTED + assert collapsed.metadata["eval_collapsed_segments"] == 0 + + +def test_sampling_params_use_sglang_generate_names(): + assert normalize_sampling_params({"max_tokens": 128, "temperature": 1.0}) == { + "max_new_tokens": 128, + "temperature": 1.0, + } + + +def test_eval_isolated_sandbox_defaults_to_false(monkeypatch): + monkeypatch.delenv("SWE_EVAL_ISOLATED_SANDBOX", raising=False) + + assert AGSGeneratorConfig.from_env().eval_isolated_sandbox is False + + +def test_eval_isolated_sandbox_can_be_enabled(monkeypatch): + monkeypatch.setenv("SWE_EVAL_ISOLATED_SANDBOX", "true") + + assert AGSGeneratorConfig.from_env().eval_isolated_sandbox is True + + +def _adapter_args(**overrides): + args = SimpleNamespace( + hf_checkpoint="/models/fake", + rollout_max_context_len=4096, + sglang_tool_call_parser="qwen3_coder", + sglang_reasoning_parser="qwen3", + sglang_router_ip="127.0.0.1", + sglang_router_port=30000, + ) + for key, value in overrides.items(): + setattr(args, key, value) + return args + + +def test_eval_falls_back_to_a_local_adapter_when_none_is_running(monkeypatch): + """Eval must not require the training rollout to have built the adapter. + + Under --num-rollout 0, or eval-before-train on rollout 0, nothing has bound + ADAPTER_PORT yet, so a RemoteAdapterProxy would fail every prompt with a + connection error. Fall back to a local adapter on an ephemeral port. + """ + monkeypatch.setenv("ADAPTER_PUBLIC_HOST", "10.0.0.1") + monkeypatch.delenv("ADAPTER_PUBLIC_BASE_URL", raising=False) + monkeypatch.delenv("AGS_EVAL_ADAPTER_CONTROL_URL", raising=False) + monkeypatch.delenv("ADAPTER_CONTROL_BASE_URL", raising=False) + monkeypatch.setattr(adapter_service, "_remote_adapter_alive", lambda *a, **k: False) + built: dict = {} + + def _fake_local(args, config, adapter_cls, *, port=None): + built["port"] = port + return SimpleNamespace(kind="local") + + monkeypatch.setattr(adapter_service, "_local_adapter_service", _fake_local) + + service = adapter_service.get_adapter_service( + _adapter_args(), AGSGeneratorConfig.from_env(), object, evaluation=True + ) + + assert service.kind == "local" + # Ephemeral: the training adapter may still claim ADAPTER_PORT later in the run. + assert built["port"] == 0 + + +def test_eval_reuses_the_training_adapter_when_one_is_live(monkeypatch): + """The live training adapter owns the trajectory trees, so prefer it.""" + monkeypatch.setenv("ADAPTER_PUBLIC_HOST", "10.0.0.1") + monkeypatch.delenv("ADAPTER_PUBLIC_BASE_URL", raising=False) + monkeypatch.setattr(adapter_service, "_remote_adapter_alive", lambda *a, **k: True) + monkeypatch.setattr( + adapter_service, + "_local_adapter_service", + lambda *a, **k: pytest.fail("must not start a local adapter when one is already live"), + ) + monkeypatch.setattr(adapter_service, "RemoteAdapterService", lambda args, config: SimpleNamespace(kind="remote")) + + service = adapter_service.get_adapter_service( + _adapter_args(), AGSGeneratorConfig.from_env(), object, evaluation=True + ) + + assert service.kind == "remote" + + +def test_training_never_probes_for_a_remote_adapter(monkeypatch): + """Training owns its adapter; it must not depend on a health probe.""" + monkeypatch.setenv("ADAPTER_PUBLIC_HOST", "10.0.0.1") + monkeypatch.setattr( + adapter_service, + "_remote_adapter_alive", + lambda *a, **k: pytest.fail("training must not probe for a remote adapter"), + ) + monkeypatch.setattr( + adapter_service, "_local_adapter_service", lambda *a, **k: SimpleNamespace(kind="local", port=k.get("port")) + ) + + service = adapter_service.get_adapter_service( + _adapter_args(), AGSGeneratorConfig.from_env(), object, evaluation=False + ) + + assert service.kind == "local" + + +def test_ephemeral_eval_adapter_advertises_its_own_node_and_port(monkeypatch): + """The fallback adapter must not advertise the head node's host:port. + + It runs in the RolloutManager actor, which Ray does not pin to the head, so + ADAPTER_PUBLIC_HOST/ADAPTER_PUBLIC_BASE_URL (head + fixed port) would send + sandboxes to the wrong address. Only the fixed-port path may use them. + """ + from slime.utils.http_utils import get_host_info + + monkeypatch.setenv("ADAPTER_PUBLIC_HOST", "10.255.255.1") + monkeypatch.setenv("ADAPTER_PUBLIC_BASE_URL", "http://10.255.255.1:18001") + monkeypatch.setenv("ADAPTER_PORT", "18903") + monkeypatch.setattr(adapter_service, "load_tokenizer", lambda *a, **k: SimpleNamespace()) + + class _StubAdapter: + def __init__(self, **kwargs): + self.app = web.Application() + + config = AGSGeneratorConfig.from_env() + SingletonMeta.clear_instances(adapter_service.AdapterService) + try: + ephemeral = adapter_service._local_adapter_service(_adapter_args(), config, _StubAdapter, port=0) + assert ephemeral.adapter_url == f"http://{get_host_info()[1]}:{ephemeral.app_handle.port}" + assert ephemeral.app_handle.port not in (0, config.adapter_port) + + SingletonMeta.clear_instances(adapter_service.AdapterService) + fixed = adapter_service._local_adapter_service(_adapter_args(), config, _StubAdapter) + assert fixed.adapter_url == "http://10.255.255.1:18001" + finally: + SingletonMeta.clear_instances(adapter_service.AdapterService) + + +def test_reusing_the_adapter_singleton_warns_when_args_differ(monkeypatch, caplog): + """AdapterService is a singleton: later callers' args are silently ignored. + + A mismatch means the live adapter was built with a different tokenizer or + context budget than this caller asked for, which would otherwise be invisible. + """ + monkeypatch.setenv("ADAPTER_PUBLIC_HOST", "10.0.0.1") + config = AGSGeneratorConfig.from_env() + first = _adapter_args() + live = SimpleNamespace( + build_key=adapter_service._adapter_build_key(first, config, object, "http://127.0.0.1:30000") + ) + monkeypatch.setattr(adapter_service, "AdapterService", lambda *a, **k: live) + + with caplog.at_level("WARNING"): + same = adapter_service._local_adapter_service(first, config, object) + assert same is live + assert not caplog.records, "identical args must not warn" + + with caplog.at_level("WARNING"): + adapter_service._local_adapter_service(_adapter_args(rollout_max_context_len=131072), config, object) + + assert any("singleton" in r.message for r in caplog.records) + + +_MD = {"instance_id": "x__1", "dataset_prompt": "# Task\n\nFix the bug.", "problem_statement": "Fix the bug."} + + +def _async_return(value): + """Async callable ignoring its arguments and returning ``value``.""" + + async def _call(*args, **kwargs): + return value + + return _call + + +@pytest.mark.parametrize( + "style,expected_prompt,expects_file", + [ + ("dataset", "# Task\n\nFix the bug.", False), + ("instruction", "Read PROBLEM_STATEMENT.md and fix it.", True), + ], +) +def test_agent_prompt_and_statement_file_agree_per_style(monkeypatch, style, expected_prompt, expects_file): + """The prompt style must decide the prompt and the statement file together. + + Wiring these two independently is how you get "instruction" without the file + it names, or "dataset" with a stray file in the repo, so both are asserted + from one style. + """ + monkeypatch.setenv("SWE_CC_PROMPT", "Read PROBLEM_STATEMENT.md and fix it.") + runner = AGSRolloutRunner.__new__(AGSRolloutRunner) # no sandbox/adapter needed + runner.config = AGSGeneratorConfig.from_env() + + assert runner._agent_prompt(_MD, style) == expected_prompt + + async def run_case(): + sb = FakeSandbox() + await swe_task.prepare_workspace(sb, "/testbed", _MD, write_problem_statement=style == "instruction") + return sb + + assert ("/testbed/PROBLEM_STATEMENT.md" in asyncio.run(run_case()).files) is expects_file + + +@pytest.mark.parametrize( + "evaluation,expected_prompt", + [(False, "Read PROBLEM_STATEMENT.md and fix it."), (True, "# Task\n\nFix the bug.")], +) +def test_generate_picks_prompt_style_by_rollout_mode(monkeypatch, evaluation, expected_prompt): + """Training and eval must resolve to their own style from one config. + + This is the wiring generate() does, exercised end to end from the env vars: + the default split is instruction for training, dataset for eval. + """ + monkeypatch.delenv("SWE_PROMPT_STYLE", raising=False) + monkeypatch.delenv("SWE_EVAL_PROMPT_STYLE", raising=False) + monkeypatch.setenv("SWE_CC_PROMPT", "Read PROBLEM_STATEMENT.md and fix it.") + runner = AGSRolloutRunner.__new__(AGSRolloutRunner) + runner.config = AGSGeneratorConfig.from_env() + + style = runner.config.prompt_style_for(evaluation=evaluation) + assert runner._agent_prompt(_MD, style) == expected_prompt + + +@pytest.mark.parametrize( + "evaluation,expected_prompt,expects_file", + [(False, "Read PROBLEM_STATEMENT.md and fix it.", True), (True, "# Task\n\nFix the bug.", False)], +) +def test_generate_threads_evaluation_flag_to_harness_and_workspace( + monkeypatch, evaluation, expected_prompt, expects_file +): + """Drive the real generate() and capture what the harness was handed. + + The helper tests above verify the style→prompt mapping; this one verifies the + plumbing, which is the part that silently breaks: generate() must resolve the + style from its own `evaluation` argument and use that same value for both the + harness prompt and the PROBLEM_STATEMENT.md decision. + """ + monkeypatch.delenv("SWE_PROMPT_STYLE", raising=False) + monkeypatch.delenv("SWE_EVAL_PROMPT_STYLE", raising=False) + monkeypatch.setenv("SWE_CC_PROMPT", "Read PROBLEM_STATEMENT.md and fix it.") + + captured: dict = {} + sandbox = FakeSandbox() + + class _Harness: + async def run(self, sb, *, workdir, session_id, adapter_url, time_budget_sec, prompt): + captured["prompt"] = prompt + return 0 + + @contextlib.asynccontextmanager + async def _fake_boot(self, image, instance_id): + yield sandbox + + runner = AGSRolloutRunner.__new__(AGSRolloutRunner) + runner.config = AGSGeneratorConfig.from_env() + runner.harness_cls = _Harness + runner.artifacts = SimpleNamespace( + dump_trajectory=_async_return(None), dump_patch=lambda *a, **k: None, dump_rollout=lambda *a, **k: None + ) + runner.weave_trace = SimpleNamespace(start_rollout=lambda **k: None, finish_rollout=lambda *a, **k: None) + # finish_session returning [] short-circuits into _abort_result, which is fine: + # the prompt and the workspace file are already decided by then. + runner.adapter_service = SimpleNamespace( + adapter=SimpleNamespace( + open_session=lambda *a, **k: None, + finish_session=_async_return([]), + drop_session=_async_return(None), + ), + max_context_len=4096, + adapter_url="http://127.0.0.1:1", + ) + monkeypatch.setattr(AGSRolloutRunner, "_boot_agent_sandbox", _fake_boot) + monkeypatch.setattr("slime_plugins.rollout_buffer.generator.ags_generator.rollout.git_diff", _async_return("")) + monkeypatch.setattr( + "slime_plugins.rollout_buffer.generator.ags_generator.rollout.evaluate", _async_return((0.0, True)) + ) + + sample = Sample( + index=0, + prompt="# Task\n\nFix the bug.", + metadata={"instance_id": "x__1", "image": "img", "workdir": "/testbed"}, + ) + asyncio.run(runner.generate(sample, {}, evaluation=evaluation)) + + assert captured["prompt"] == expected_prompt + assert ("/testbed/PROBLEM_STATEMENT.md" in sandbox.files) is expects_file + + +def test_agent_prompt_falls_back_when_dataset_prompt_is_empty(monkeypatch): + monkeypatch.setenv("SWE_CC_PROMPT", "fallback prompt") + runner = AGSRolloutRunner.__new__(AGSRolloutRunner) + runner.config = AGSGeneratorConfig.from_env() + + assert runner._agent_prompt({"instance_id": "x__1", "dataset_prompt": " "}, "dataset") == "fallback prompt" + + +@pytest.mark.parametrize("write_problem_statement", [True, False]) +def test_prepare_workspace_writes_problem_statement_only_when_asked(write_problem_statement): + """Under the "dataset" prompt style the prompt already carries the task text, + so PROBLEM_STATEMENT.md must not be created: it would be an untracked file in + the repo that only git_diff's exclude pathspec keeps out of the patch.""" + + async def run_case(): + sb = FakeSandbox() + await swe_task.prepare_workspace( + sb, + "/testbed", + {"problem_statement": "Fix the bug."}, + write_problem_statement=write_problem_statement, + ) + return sb + + sb = asyncio.run(run_case()) + written = "/testbed/PROBLEM_STATEMENT.md" in sb.files + assert written is write_problem_statement + if written: + assert sb.files["/testbed/PROBLEM_STATEMENT.md"] == "Fix the bug." + + +def test_prepare_workspace_writes_problem_statement_by_default(): + """Callers that predate the flag keep the old behaviour.""" + + async def run_case(): + sb = FakeSandbox() + await swe_task.prepare_workspace(sb, "/testbed", {"problem_statement": "Fix the bug."}) + return sb + + assert "/testbed/PROBLEM_STATEMENT.md" in asyncio.run(run_case()).files + + +def test_evaluate_can_reuse_agent_sandbox(monkeypatch): + async def run_case(): + monkeypatch.setattr( + swe_task, + "AGSSandbox", + lambda _image: (_ for _ in ()).throw(AssertionError("must not boot an isolated sandbox")), + ) + sb = FakeSandbox() + + reward, applied = await swe_task.evaluate( + sandbox=sb, + image="unused-image", + workdir="/workspace/repo", + diff_text="diff --git a/a.py b/a.py", + eval_cmd="pytest -q", + pre_commands=["echo must-not-rerun"], + eval_bootstrap_cmd="echo bootstrap", + ) + + assert reward == 1.0 + assert applied is True + commands = [cmd for cmd, _user in sb.exec_log] + assert "cd /workspace/repo && pytest -q" in commands + assert "cd /workspace/repo && echo bootstrap" in commands + assert not any("must-not-rerun" in cmd for cmd in commands) + assert not any("git apply" in cmd or "patch -p1" in cmd for cmd in commands) + + asyncio.run(run_case()) + + +def test_evaluate_isolated_sandbox_keeps_clean_apply_flow(monkeypatch): + async def run_case(): + sandboxes = [] + + def sandbox_factory(image): + sb = FakeSandbox(image) + sandboxes.append(sb) + return sb + + monkeypatch.setattr(swe_task, "AGSSandbox", sandbox_factory) + + reward, applied = await swe_task.evaluate( + sandbox=None, + image="eval-image", + workdir="/workspace/repo", + diff_text="diff --git a/a.py b/a.py", + eval_cmd="pytest -q", + ) + + assert reward == 1.0 + assert applied is True + assert len(sandboxes) == 1 + commands = [cmd for cmd, _user in sandboxes[0].exec_log] + assert any("git apply --3way" in cmd for cmd in commands) + assert "cd /workspace/repo && pytest -q" in commands + + asyncio.run(run_case()) + + +class _FakeTokenizer: + def __init__(self): + self.calls = [] + + def decode(self, tokens, skip_special_tokens=False): + self.calls.append((tokens, skip_special_tokens)) + return "|".join(str(token) for token in tokens) + + +class _FakeWeaveClient: + def __init__(self): + self.created = [] + self.finished = [] + self.wandb_contexts = [] + + def create_call(self, op, inputs, **kwargs): + call = SimpleNamespace(op=op, inputs=inputs, kwargs=kwargs) + self.created.append(call) + return call + + def finish_call(self, call, output=None, exception=None, **kwargs): + self.finished.append((call, output, exception, kwargs)) + + def set_wandb_run_context(self, run_id, step=None): + self.wandb_contexts.append((run_id, step)) + + +def _trace(enable_token2text=False): + trace = AGSWeaveTrace(_trace_args(), _FakeTokenizer(), enable_token2text=enable_token2text) + trace.client = _FakeWeaveClient() + return trace + + +def _trace_for_agent(agent): + trace = AGSWeaveTrace(_trace_args(), _FakeTokenizer(), agent=agent) + trace.client = _FakeWeaveClient() + return trace + + +def _trace_args(**overrides): + data = { + "use_wandb": False, + "wandb_mode": None, + "wandb_team": "team", + "wandb_project": "project", + "wandb_run_id": "run-1", + "wandb_group": "group-1", + } + data.update(overrides) + return SimpleNamespace(**data) + + +def test_weave_trace_disabled_without_use_wandb(monkeypatch): + monkeypatch.setitem( + sys.modules, "weave", types.SimpleNamespace(init=lambda project: (_ for _ in ()).throw(AssertionError)) + ) + + trace = AGSWeaveTrace(_trace_args(use_wandb=False, wandb_mode="online"), _FakeTokenizer()) + + assert trace.client is None + + +def test_weave_trace_disabled_for_offline_or_disabled_wandb(monkeypatch): + monkeypatch.setitem( + sys.modules, "weave", types.SimpleNamespace(init=lambda project: (_ for _ in ()).throw(AssertionError)) + ) + + assert AGSWeaveTrace(_trace_args(use_wandb=True, wandb_mode="disabled"), _FakeTokenizer()).client is None + assert AGSWeaveTrace(_trace_args(use_wandb=True, wandb_mode="offline"), _FakeTokenizer()).client is None + + +def test_weave_trace_uses_wandb_project_and_run_context(monkeypatch): + client = _FakeWeaveClient() + seen = {} + + def fake_init(project): + seen["project"] = project + return client + + monkeypatch.setitem(sys.modules, "weave", types.SimpleNamespace(init=fake_init)) + + trace = AGSWeaveTrace( + _trace_args(use_wandb=True, wandb_mode="online", wandb_team="entity", wandb_project="train-proj"), + _FakeTokenizer(), + ) + + assert trace.client is client + assert seen["project"] == "entity/train-proj" + assert client.wandb_contexts == [("run-1", None)] + + +def test_weave_trace_keeps_token_ids_without_decoding_by_default(): + trace = _trace() + + payload = trace._sample_payload(_sample()) + + assert payload["prompt_token_ids"] == [1] + assert payload["response_token_ids"] == [2, 3] + assert "prompt_text" not in payload + assert "response_text" not in payload + assert trace.tokenizer.calls == [] + + +def test_weave_trace_decodes_prompt_and_response_when_enabled(): + trace = _trace(enable_token2text=True) + + payload = trace._sample_payload(_sample()) + + assert payload["prompt_text"] == "1" + assert payload["response_text"] == "2|3" + assert trace.tokenizer.calls == [([1], False), ([2, 3], False)] + + +def test_weave_trace_pairs_tool_call_and_result(tmp_path): + trajectory = tmp_path / "trajectory.jsonl" + trajectory.write_text( + "\n".join( + [ + json.dumps( + { + "type": "assistant", + "timestamp": "2026-07-13T01:02:03Z", + "message": { + "content": [ + {"type": "thinking", "thinking": "inspect"}, + {"type": "tool_use", "id": "tool-1", "name": "Read", "input": {"file": "a.py"}}, + ] + }, + } + ), + json.dumps( + { + "type": "user", + "timestamp": "2026-07-13T01:02:04Z", + "message": { + "content": [{"type": "tool_result", "tool_use_id": "tool-1", "content": "source"}] + }, + } + ), + json.dumps({"type": "stream_event", "event": {"type": "content_block_delta"}}), + ] + ), + encoding="utf-8", + ) + trace = _trace() + parent = SimpleNamespace() + + trace._log_trajectory(parent, str(trajectory)) + + assert [call.op for call in trace.client.created] == ["slime.ags.thinking", "slime.ags.tool_call"] + tool_call = trace.client.created[1] + tool_finish = next(item for item in trace.client.finished if item[0] is tool_call) + assert tool_call.inputs["input"] == {"file": "a.py"} + assert tool_finish[1]["content"] == "source" + assert tool_finish[3]["ended_at"].isoformat() == "2026-07-13T01:02:04+00:00" + assert len(list(iter_trajectory_events(trajectory))) == 3 + + +def test_weave_trace_parses_codebuddy_code_stream_json(tmp_path): + trajectory = tmp_path / "trajectory.jsonl" + trajectory.write_text( + "\n".join( + [ + json.dumps( + { + "type": "stream_event", + "event": {"type": "content_block_delta"}, + "__timestamp": "2026-07-14T03:24:03.333Z", + } + ), + json.dumps( + { + "type": "assistant", + "uuid": "think-1", + "session_id": "sess-cbc", + "message": { + "content": [{"type": "thinking", "thinking": "inspect cbc"}], + "usage": {"input_tokens": 0, "output_tokens": 0}, + }, + "__timestamp": "2026-07-14T03:24:03.343Z", + } + ), + json.dumps( + { + "type": "assistant", + "uuid": "msg-1", + "session_id": "sess-cbc", + "message": { + "content": [ + { + "type": "tool_use", + "id": "call-1", + "name": "Read", + "input": {"file_path": "/testbed/PROBLEM_STATEMENT.md"}, + } + ], + "usage": {"input_tokens": 10, "output_tokens": 2}, + }, + "__timestamp": "2026-07-14T03:24:03.352Z", + } + ), + json.dumps( + { + "type": "user", + "uuid": "result-1", + "session_id": "sess-cbc", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "call-1", + "content": [{"type": "text", "text": "source"}], + "is_error": False, + } + ] + }, + "parent_tool_use_id": "call-1", + "__timestamp": "2026-07-14T03:24:03.379Z", + } + ), + json.dumps( + { + "type": "result", + "subtype": "success", + "is_error": False, + "result": "done", + "session_id": "sess-cbc", + "__timestamp": "2026-07-14T03:24:04.000Z", + } + ), + ] + ), + encoding="utf-8", + ) + trace = _trace_for_agent("codebuddy_code") + parent = SimpleNamespace() + + trace._log_trajectory(parent, str(trajectory)) + + assert [call.op for call in trace.client.created] == [ + "slime.ags.thinking", + "slime.ags.tool_call", + "slime.ags.result", + ] + thinking, tool_call, result_call = trace.client.created + tool_finish = next(item for item in trace.client.finished if item[0] is tool_call) + assert thinking.inputs["timestamp"] == "2026-07-14T03:24:03.343Z" + assert tool_call.inputs["input"] == {"file_path": "/testbed/PROBLEM_STATEMENT.md"} + assert tool_finish[1]["content"] == [{"type": "text", "text": "source"}] + assert tool_finish[3]["ended_at"].isoformat() == "2026-07-14T03:24:03.379000+00:00" + assert next(item for item in trace.client.finished if item[0] is result_call)[1]["result"] == "done" + assert len(list(iter_trajectory_events(trajectory, agent="codebuddy_code"))) == 4 + + +def test_weave_trace_codex_parser_placeholder_noops(tmp_path): + trajectory = tmp_path / "trajectory.jsonl" + trajectory.write_text( + json.dumps( + { + "type": "assistant", + "timestamp": "2026-07-13T01:02:03Z", + "message": {"content": [{"type": "text", "text": "not parsed yet"}]}, + } + ), + encoding="utf-8", + ) + trace = _trace_for_agent("codex") + + trace._log_trajectory(SimpleNamespace(), str(trajectory)) + + assert trace.client.created == [] + assert list(iter_trajectory_events(trajectory, agent="codex")) == [] + + +def test_weave_trace_finishes_root_when_child_logging_fails(monkeypatch): + trace = _trace(enable_token2text=True) + root = SimpleNamespace() + + def fail_child_logging(parent, trajectory_path): + raise RuntimeError("trace backend unavailable") + + monkeypatch.setattr(trace, "_log_trajectory", fail_child_logging) + trace.finish_rollout(root, samples=[_sample()], trajectory_path="trajectory.jsonl") + + assert len(trace.client.finished) == 1 + call, output, exception, _ = trace.client.finished[0] + assert call is root + assert output == {"trace_output_error": True} + assert exception is None + + +def test_start_rollout_forwards_enable_token2text(monkeypatch): + captured = {} + + class _Response: + def raise_for_status(self): + return None + + def json(self): + return {"message": "Rollout started"} + + def fake_post(url, json, timeout): + captured.update(json) + return _Response() + + monkeypatch.setattr("slime_plugins.rollout_buffer.rollout_buffer_example.requests.post", fake_post) + args = SimpleNamespace( + rollout_num_process=1, + num_epoch=1, + sglang_router_ip="127.0.0.1", + sglang_router_port=30000, + rollout_buffer_url="http://127.0.0.1:8889", + rollout_task_type="ags", + prompt_data="smoke.jsonl", + n_samples_per_prompt=1, + rollout_max_response_len=16, + rollout_temperature=1.0, + rollout_top_p=1.0, + rollout_top_k=-1, + hf_checkpoint="model", + rollout_batch_size=1, + enable_token2text=True, + use_wandb=True, + wandb_mode="online", + wandb_project="train-proj", + wandb_team="entity", + wandb_run_id="run-1", + wandb_group="group-1", + ) + + start_rollout(args.rollout_buffer_url, args, {}) + + assert captured["num_epoch"] == "1" + assert captured["num_groups_per_epoch"] == "1" + assert captured["enable_token2text"] is True + assert captured["use_wandb"] is True + assert captured["wandb_mode"] == "online" + assert captured["wandb_project"] == "train-proj" + assert captured["wandb_team"] == "entity" + assert captured["wandb_run_id"] == "run-1" + assert captured["wandb_group"] == "group-1" + + +def test_start_rollout_uses_one_buffer_epoch_even_when_trainer_num_epoch_is_larger(monkeypatch): + captured = {} + + class _Response: + def raise_for_status(self): + return None + + def json(self): + return {"message": "Rollout started"} + + def fake_post(url, json, timeout): + captured.update(json) + return _Response() + + monkeypatch.setattr("slime_plugins.rollout_buffer.rollout_buffer_example.requests.post", fake_post) + args = SimpleNamespace( + rollout_num_process=1, + num_epoch=3, + sglang_router_ip="127.0.0.1", + sglang_router_port=30000, + rollout_buffer_url="http://127.0.0.1:8889", + rollout_task_type="ags", + prompt_data="smoke.jsonl", + n_samples_per_prompt=4, + rollout_max_response_len=16, + rollout_temperature=1.0, + rollout_top_p=1.0, + rollout_top_k=-1, + hf_checkpoint="model", + rollout_batch_size=8, + ) + + start_rollout(args.rollout_buffer_url, args, {}, num_groups_per_epoch=2) + + assert captured["num_epoch"] == "1" + assert captured["num_groups_per_epoch"] == "2" + + +def test_rollout_buffer_ignores_stale_ags_job_items(): + from slime_plugins.rollout_buffer.buffer import RolloutBuffer + + current = output_item_from_samples([_sample()], instance_id="inst-1") + current["rollout_job_id"] = "job-current" + stale = output_item_from_samples([_sample()], instance_id="inst-2") + stale["rollout_job_id"] = "job-stale" + + buffer = RolloutBuffer(group_size=1, rollout_job_id="job-current") + + assert buffer.write(stale) is None + assert buffer.write(current) == current + + data = buffer.read()["data"] + assert len(data) == 1 + assert data[0]["instance_id"] == "inst-1" + + +def _ctx(workdir="/workspace/repo", sid="sess-1", url="http://host:18001"): + from slime.agent.harness.common import HarnessContext + + return HarnessContext(workdir=workdir, session_id=sid, adapter_url=url) + + +def _decode_first_b64(cmd: str, path: str) -> dict: + pattern = rf"printf %s ([^ ]+) \| base64 -d \| tee .*{re.escape(path)}" + m = re.search(pattern, cmd) + assert m, cmd + return json.loads(base64.b64decode(m.group(1)).decode()) + + +def test_codebuddy_code_registry_uses_openai_adapter(): + from slime.agent.adapters import OpenAIAdapter + + harness_cls, adapter_cls = resolve_agent("codebuddy_code") + assert harness_cls is CodeBuddyCodeHarness + assert adapter_cls is OpenAIAdapter + + +def test_codebuddy_code_install_uses_ags_sidecar_binary(): + async def run_case(): + sb = FakeSandbox() + await CodeBuddyCodeHarness().install_cli(sb) + + cmd = "\n".join(c for c, _ in sb.exec_log) + assert "/opt/agents/cbc/bin/cbc" in cmd + assert "ln -sf /opt/agents/cbc/bin/cbc /usr/local/bin/cbc" in cmd + assert "cbc --version" in cmd + + asyncio.run(run_case()) + + +def test_codebuddy_code_write_config_points_to_adapter(): + async def run_case(): + sb = FakeSandbox() + await CodeBuddyCodeHarness().write_config(sb, _ctx(sid="sess-cbc", url="http://host:18001")) + + cmd = next(c for c, _ in sb.exec_log if "/root/.codebuddy/models.json" in c) + models = _decode_first_b64(cmd, "/root/.codebuddy/models.json") + settings = _decode_first_b64(cmd, "/root/.codebuddy/settings.json") + assert models["models"][0]["id"] == "slime-actor" + assert models["models"][0]["apiKey"] == "sess-cbc" + assert models["models"][0]["url"] == "http://host:18001/v1/chat/completions" + assert models["models"][0]["supportsToolCall"] is True + # No maxOutputTokens: the harness sets no per-turn cap, so the CLI applies + # its own default. The adapter still bounds a turn via max_new_tokens. + assert "maxOutputTokens" not in models["models"][0] + assert settings["alwaysThinkingEnabled"] is True + + asyncio.run(run_case()) + + +def test_codebuddy_code_launch_command_and_env(): + async def run_case(): + sb = FakeSandbox() + rc = await CodeBuddyCodeHarness().launch_and_wait( + sb, + _ctx(sid="sess-cbc", url="http://host:18001"), + prompt="solve it", + time_budget_sec=0, + ) + + assert rc != 0 # time_budget=0 avoids waiting; launch still happens. + body = next(v for k, v in sb.files.items() if k.endswith("run.sh")) + assert "cbc --model slime-actor --verbose --output-format stream-json --include-partial-messages" in body + assert "-y" in body and "solve it" in body + # Tool restriction uses --disallowedTools, which the CLI enforces; --tools + # does not restrict the surface, so it is never passed. + assert "--disallowedTools WebSearch WebFetch" in body + assert "--tools" not in body + assert "codebuddy_sessions" in body + # No turn cap from the harness: the run is bounded by + # SWE_AGENT_TIME_BUDGET_SEC, and a caller who wants one passes it in + # SLIME_AGENT_CBC_EXTRA_ARGS. + assert "--max-turns" not in body + # --disallowedTools is variadic, so the non-variadic tail and the prompt + # must follow it, or it would swallow them. + assert body.index("--disallowedTools") < body.index("-y") < body.index("solve it") + + launch_cmd = next(c for c, _ in sb.exec_log if "setsid" in c) + assert "OPENAI_API_KEY=sess-cbc" in launch_cmd + assert "OPENAI_BASE_URL=http://host:18001/v1" in launch_cmd + assert "CBC_API_KEY=sess-cbc" in launch_cmd + assert "CBC_BASE_URL=http://host:18001/v1/chat/completions" in launch_cmd + assert "IS_SANDBOX=1" in launch_cmd + assert any("kill -TERM" in cmd for cmd, _user in sb.exec_log) + assert any("kill -KILL" in cmd for cmd, _user in sb.exec_log) + + asyncio.run(run_case()) + + +def test_extra_args_come_after_defaults_so_they_win(monkeypatch): + """EXTRA_ARGS is the only flag knob, so it must be able to beat a default. + + Both CLIs take the LAST occurrence of a repeated flag (verified against the + real binaries: `--max-turns 99 --max-turns 1` stops after 1 turn on each), so + "wins" here means "appears later in the command". + """ + + async def run_case(): + monkeypatch.setenv("SLIME_AGENT_CBC_EXTRA_ARGS", "--disallowedTools ImageGen --max-turns 7") + sb = FakeSandbox() + await CodeBuddyCodeHarness().launch_and_wait( + sb, _ctx(sid="s", url="http://host:18001"), prompt="go", time_budget_sec=0 + ) + body = next(v for k, v in sb.files.items() if k.endswith("run.sh")) + assert "--max-turns 7" in body + assert ( + body.index("--disallowedTools WebSearch WebFetch") + < body.index("--disallowedTools ImageGen") + < body.index("go") + ) + + monkeypatch.setenv("SLIME_AGENT_CC_EXTRA_ARGS", "--disallowedTools ImageGen --max-turns 7") + sb2 = FakeSandbox() + await AGSSidecarClaudeCodeHarness().launch_and_wait( + sb2, _ctx(sid="s", url="http://host:18001"), prompt="go", time_budget_sec=0 + ) + body2 = next(v for k, v in sb2.files.items() if k.endswith("run.sh")) + assert "--max-turns 7" in body2 + assert body2.index("--disallowedTools WebSearch WebFetch") < body2.index("--disallowedTools ImageGen") + + asyncio.run(run_case()) + + +def test_web_tools_denied_by_default_on_both_harnesses(): + """Web access makes a rollout unreproducible and can leak the graded fix. + + A 2026-07-29 eval matrix had this denied for CodeBuddy but not for Claude + Code, which then used WebSearch/WebFetch on 1-2% of instances -- a difference + of the same order as the training effect being measured. + """ + + async def run_case(): + for harness in (AGSSidecarClaudeCodeHarness(), CodeBuddyCodeHarness()): + sb = FakeSandbox() + await harness.launch_and_wait(sb, _ctx(sid="s", url="http://host:18001"), prompt="go", time_budget_sec=0) + body = next(v for k, v in sb.files.items() if k.endswith("run.sh")) + assert "--disallowedTools WebSearch WebFetch" in body, harness.name + + asyncio.run(run_case()) + + +def test_claude_code_extra_envs_override_static_env(monkeypatch): + """EXTRA_ENVS is merged after static_env, so it can override any of it.""" + + async def run_case(): + sb = FakeSandbox() + await AGSSidecarClaudeCodeHarness().launch_and_wait( + sb, _ctx(sid="s", url="http://host:18001"), prompt="go", time_budget_sec=0 + ) + launch_cmd = next(c for c, _ in sb.exec_log if "setsid" in c) + assert "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1" in launch_cmd + # The harness sets no per-turn output cap of its own. + assert "CLAUDE_CODE_MAX_OUTPUT_TOKENS" not in launch_cmd + + monkeypatch.setenv( + "SLIME_AGENT_CC_EXTRA_ENVS", + '{"CLAUDE_CODE_MAX_OUTPUT_TOKENS":"4096","CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC":"0"}', + ) + sb2 = FakeSandbox() + await AGSSidecarClaudeCodeHarness().launch_and_wait( + sb2, _ctx(sid="s", url="http://host:18001"), prompt="go", time_budget_sec=0 + ) + launch_cmd2 = next(c for c, _ in sb2.exec_log if "setsid" in c) + assert "CLAUDE_CODE_MAX_OUTPUT_TOKENS=4096" in launch_cmd2 + assert "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=0" in launch_cmd2 + + asyncio.run(run_case()) + + +def test_ags_timeout_fails_closed_if_agent_process_group_cannot_be_stopped(): + class StopFailureSandbox(FakeSandbox): + async def exec(self, cmd, **kwargs): + if "kill -TERM" in cmd: + raise RuntimeError("cannot stop agent") + return await super().exec(cmd, **kwargs) + + async def run_case(): + with pytest.raises(RuntimeError, match="cannot stop agent"): + await run_root_command( + StopFailureSandbox(), + workdir="/workspace/repo", + start_cmd="claude -p solve", + env={}, + time_budget_sec=0, + ) + + asyncio.run(run_case()) + + +def test_autocompact_percentage_reaches_each_cli(monkeypatch): + """Auto-compaction is configured as a PERCENTAGE, not an absolute window. + + Both CLIs clamp an absolute auto-compact window to [100k, 1M], which is above + the default 96k context cap -- so an absolute setting could never fire in time. + The percentage path has no clamp. + + For CodeBuddy the percentage is of the model's maxInputTokens, so models.json + MUST carry that key: resolveCompactTriggerAt() in agent-cli + src/node/context/context-protocol.ts reads + `modelMaxInputTokens ? modelMaxInputTokens * pct : getAutoCompactWindow()`, + and the fallback is the clamped absolute window. + """ + + async def run_case(): + monkeypatch.setenv("SLIME_AGENT_MAX_INPUT_TOKENS", "96000") + sb = FakeSandbox() + await CodeBuddyCodeHarness().write_config(sb, _ctx(sid="s", url="http://host:18001")) + cmd = next(c for c, _ in sb.exec_log if "/root/.codebuddy/models.json" in c) + models = _decode_first_b64(cmd, "/root/.codebuddy/models.json") + assert models["models"][0]["maxInputTokens"] == 96000 + # The on/off switch lives in settings.json, the threshold in the env var. + assert _decode_first_b64(cmd, "/root/.codebuddy/settings.json")["autoCompactEnabled"] is True + + # Absent/blank/garbage must leave the key out rather than write a bogus + # value, since a wrong maxInputTokens silently moves the trigger point. + for bad in ("", "0", "-1", "not-a-number"): + monkeypatch.setenv("SLIME_AGENT_MAX_INPUT_TOKENS", bad) + sb_bad = FakeSandbox() + await CodeBuddyCodeHarness().write_config(sb_bad, _ctx(sid="s", url="http://host:18001")) + cmd_bad = next(c for c, _ in sb_bad.exec_log if "/root/.codebuddy/models.json" in c) + entry = _decode_first_b64(cmd_bad, "/root/.codebuddy/models.json")["models"][0] + assert "maxInputTokens" not in entry, bad + + # The thresholds themselves ride the two EXTRA_* knobs. + monkeypatch.setenv("SLIME_AGENT_CBC_EXTRA_ENVS", '{"CODEBUDDY_AUTOCOMPACT_PCT_OVERRIDE":"70"}') + sb2 = FakeSandbox() + await CodeBuddyCodeHarness().launch_and_wait( + sb2, _ctx(sid="s", url="http://host:18001"), prompt="go", time_budget_sec=0 + ) + assert "CODEBUDDY_AUTOCOMPACT_PCT_OVERRIDE=70" in next(c for c, _ in sb2.exec_log if "setsid" in c) + + # Claude Code needs the window declared too, or the percentage is of its + # 200000 default -- which our 96k cap is never reached from. It applies + # directly because "slime-actor" is not a claude-* name. + monkeypatch.setenv( + "SLIME_AGENT_CC_EXTRA_ENVS", + '{"CLAUDE_CODE_MAX_CONTEXT_TOKENS":"96000",' + '"CLAUDE_CODE_MAX_OUTPUT_TOKENS":"32768",' + '"CLAUDE_AUTOCOMPACT_PCT_OVERRIDE":"60"}', + ) + sb3 = FakeSandbox() + await AGSSidecarClaudeCodeHarness().launch_and_wait( + sb3, _ctx(sid="s", url="http://host:18001"), prompt="go", time_budget_sec=0 + ) + launch = next(c for c, _ in sb3.exec_log if "setsid" in c) + assert "CLAUDE_CODE_MAX_CONTEXT_TOKENS=96000" in launch + assert "CLAUDE_CODE_MAX_OUTPUT_TOKENS=32768" in launch + assert "CLAUDE_AUTOCOMPACT_PCT_OVERRIDE=60" in launch + + asyncio.run(run_case()) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/test_rollout_buffer/test_ags_prompt_source.py b/tests/test_rollout_buffer/test_ags_prompt_source.py new file mode 100644 index 0000000000..b9b59e0bb0 --- /dev/null +++ b/tests/test_rollout_buffer/test_ags_prompt_source.py @@ -0,0 +1,156 @@ +"""Unit tests for AGSPromptSource's seek-to-absolute-group-index behaviour. + +The generator has no checkpoint of its own: the trainer tells it where to start +by absolute group index (``rollout_start_group = rollout_id * +rollout_batch_size``), and AGSPromptSource seeks ``RolloutDataSource`` there. +So the property under test is that seeking to group g yields exactly the +prompts an uninterrupted run would have produced at group g -- which is what +makes a resumed run replay the same per-step prompt set. The interesting case +is a seek past the end of the dataset, where the epoch (and with shuffle on, +the permutation) has to advance too. + +``source`` is loaded from its file through a private synthetic package rather +than through ``slime_plugins...``: that package's ``__init__`` imports the whole +entry module (and ``openai``), which the CPU CI jobs do not install. See the +sibling ``test_ags_empty_patch_guard.py`` for the same pattern. +""" + +from __future__ import annotations + +import importlib +import json +import sys +import types +from argparse import Namespace +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +# slime.rollout.data_source -> slime.utils.processing_utils -> transformers, a +# heavy dep absent from the CPU-only CI env. No real tokenizer is ever used +# here (load_tokenizer/load_processor are patched, and Dataset skips tokenizing +# when apply_chat_template is off and max_length is None), so stub it out. +if "transformers" not in sys.modules: + _tf_stub = types.ModuleType("transformers") + for _name in ("AutoProcessor", "AutoTokenizer", "PreTrainedTokenizerBase", "ProcessorMixin"): + setattr(_tf_stub, _name, type(_name, (), {})) + sys.modules["transformers"] = _tf_stub + +AGS_DIR = REPO_ROOT / "slime_plugins" / "rollout_buffer" / "generator" / "ags_generator" +_PRIVATE_PACKAGE = "_ags_source_under_test" + +if _PRIVATE_PACKAGE not in sys.modules: + _package = types.ModuleType(_PRIVATE_PACKAGE) + _package.__path__ = [str(AGS_DIR)] + sys.modules[_PRIVATE_PACKAGE] = _package + +source_mod = importlib.import_module(".source", package=_PRIVATE_PACKAGE) +AGSPromptSource = source_mod.AGSPromptSource + +NUM_GPUS = 0 + +DATASET_LEN = 5 + + +@pytest.fixture(autouse=True) +def _no_tokenizer(monkeypatch): + """RolloutDataSource always loads a tokenizer/processor; neither is used.""" + data_source_mod = sys.modules["slime.rollout.data_source"] + monkeypatch.setattr(data_source_mod, "load_tokenizer", lambda *a, **k: None) + monkeypatch.setattr(data_source_mod, "load_processor", lambda *a, **k: None) + + +@pytest.fixture +def prompt_data(tmp_path) -> str: + path = tmp_path / "prompts.jsonl" + lines = [ + json.dumps({"prompt": f"p{i}", "label": f"l{i}", "metadata": {"instance_id": f"i{i}"}}) + for i in range(DATASET_LEN) + ] + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return str(path) + + +def _args(prompt_data: str, *, start_group: int, shuffle: bool) -> Namespace: + return Namespace( + hf_checkpoint="unused", + prompt_data=prompt_data, + input_key="prompt", + label_key="label", + metadata_key="metadata", + tool_key=None, + multimodal_keys=None, + apply_chat_template=False, + apply_chat_template_kwargs={}, + rollout_global_dataset=True, + rollout_shuffle=shuffle, + rollout_seed=42, + rollout_max_prompt_len=None, + dump_details=None, + rollout_start_group=start_group, + n_samples_per_prompt=2, + ) + + +def _prompts_from_seek(prompt_data: str, *, start_group: int, num_groups: int, shuffle: bool) -> list[str]: + """Prompts a generator started at ``start_group`` produces for its first request.""" + source = AGSPromptSource(_args(prompt_data, start_group=start_group, shuffle=shuffle)) + return [group[0].prompt for group in source.get_groups(num_groups)] + + +def _prompts_from_start(prompt_data: str, *, num_groups: int, shuffle: bool) -> list[str]: + """Prompts an uninterrupted run produces, drawn one group at a time from 0.""" + source = AGSPromptSource(_args(prompt_data, start_group=0, shuffle=shuffle)) + return [source.get_groups(1)[0][0].prompt for _ in range(num_groups)] + + +@pytest.mark.parametrize("shuffle", [False, True]) +def test_seek_within_first_epoch_matches_uninterrupted_run(prompt_data, shuffle): + uninterrupted = _prompts_from_start(prompt_data, num_groups=DATASET_LEN, shuffle=shuffle) + resumed = _prompts_from_seek(prompt_data, start_group=2, num_groups=2, shuffle=shuffle) + assert resumed == uninterrupted[2:4] + + +@pytest.mark.parametrize("shuffle", [False, True]) +def test_seek_past_dataset_end_matches_uninterrupted_run(prompt_data, shuffle): + """The regression: a seek that wraps must land in the right epoch. + + Group DATASET_LEN + 1 is the second group of epoch 1. Restoring only + ``sample_offset`` left the source on epoch 0's permutation, so a resumed run + diverged from the original once training passed one pass over the data. + """ + start_group = DATASET_LEN + 1 + uninterrupted = _prompts_from_start(prompt_data, num_groups=start_group + 2, shuffle=shuffle) + resumed = _prompts_from_seek(prompt_data, start_group=start_group, num_groups=2, shuffle=shuffle) + assert resumed == uninterrupted[start_group : start_group + 2] + + +def test_seek_past_dataset_end_advances_epoch(prompt_data): + source = AGSPromptSource(_args(prompt_data, start_group=2 * DATASET_LEN + 3, shuffle=True)) + assert source.data_source.epoch_id == 2 + assert source.data_source.sample_offset == 3 + assert source.data_source.dataset.epoch_id == 2 + + +def test_shuffled_epochs_differ(prompt_data): + """Guards the test above: epoch 1's order must actually differ from epoch 0's.""" + epoch0 = _prompts_from_seek(prompt_data, start_group=0, num_groups=DATASET_LEN, shuffle=True) + epoch1 = _prompts_from_seek(prompt_data, start_group=DATASET_LEN, num_groups=DATASET_LEN, shuffle=True) + assert sorted(epoch0) == sorted(epoch1) + assert epoch0 != epoch1 + + +def test_seek_is_noop_at_group_zero(prompt_data): + source = AGSPromptSource(_args(prompt_data, start_group=0, shuffle=True)) + assert source.data_source.epoch_id == 0 + assert source.data_source.sample_offset == 0 + assert source.data_source.sample_group_index == 0 + assert source.data_source.sample_index == 0 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/test_rollout_buffer/test_rollout_buffer_example.py b/tests/test_rollout_buffer/test_rollout_buffer_example.py new file mode 100644 index 0000000000..7e860ef3ce --- /dev/null +++ b/tests/test_rollout_buffer/test_rollout_buffer_example.py @@ -0,0 +1,34 @@ +from types import SimpleNamespace + +from slime.rollout.data_source import RolloutDataSourceWithBuffer +from slime.utils.types import Sample +from slime_plugins.rollout_buffer.rollout_buffer_example import _select_complete_rollout_groups + + +def _record(instance_id, timestamp): + return {"instance_id": instance_id, "timestamp": timestamp} + + +def test_select_complete_rollout_groups_skips_partial_prompt_groups(): + args = SimpleNamespace(n_samples_per_prompt=8) + results = [_record("complete", index) for index in range(8)] + results += [_record("partial", index) for index in range(2)] + + selected, incomplete = _select_complete_rollout_groups(args, results, need_length=1) + + assert len(selected) == 1 + assert len(selected[0]) == 8 + assert {record["instance_id"] for record in selected[0]} == {"complete"} + assert incomplete == {"partial": 2} + + +def test_buffer_accepts_flattened_agent_segments_above_requested_attempt_count(): + data_source = RolloutDataSourceWithBuffer.__new__(RolloutDataSourceWithBuffer) + data_source.args = SimpleNamespace(n_samples_per_prompt=8) + data_source.buffer = [] + + segments = [Sample(index=index) for index in range(9)] + + data_source.add_samples([segments]) + + assert data_source.buffer == [segments] diff --git a/tools/harbor_task_to_slime_prompt_data.py b/tools/harbor_task_to_slime_prompt_data.py new file mode 100644 index 0000000000..e792a76b88 --- /dev/null +++ b/tools/harbor_task_to_slime_prompt_data.py @@ -0,0 +1,678 @@ +#!/usr/bin/env python3 +"""Convert Harbor SWE-style tasks to semantic slime/AGS prompt-data JSONL. + +The default output is not a mirror of Harbor task files. It extracts the small +set of fields that slime's AGS rollout-buffer generator already consumes: + + - prompt text + - metadata.instance_id + - metadata.image + - metadata.workdir + - metadata.problem_statement + - metadata.eval_cmd + +Both the prompt and metadata.problem_statement are Harbor's instruction.md, the +exact string Harbor hands its agents. Upstream's raw tests/config.json text is +kept under metadata.harbor.problem_statement for provenance. + +The generated rows are still ordinary slime JSONL prompt data: use --input-key +prompt, --label-key label, and --metadata-key metadata. The eval command is +built from Harbor's tests/test.sh plus tests/config.json so the row can be used +by ags_generator without a harbor_task_path. + +Important: Harbor verifier assets are created inside metadata.eval_cmd, not +metadata.pre_commands, so they are materialized only after the foreground agent +process returns. With ``SWE_EVAL_ISOLATED_SANDBOX=true`` this happens in a +separate clean sandbox. With the default ``false`` it happens in the existing +agent sandbox; that mode avoids a second sandbox boot but is not a security +boundary against a root-capable agent or background processes it leaves behind. +Converted Harbor tasks preserve the prebuilt task image's Git HEAD by default: +that image can contain task-specific environment compatibility commits beyond +``tests/config.json.base_commit``. + +Example: + python tools/harbor_task_to_slime_prompt_data.py \ + --input /path/to/harbor-datasets/datasets/swebench-verified \ + --output ./local/swebench-verified-harbor-dataset/prompt_data.jsonl \ + --task astropy__astropy-12907 \ + --source swebench-verified \ + --pretty-output ./local/swebench-verified-harbor-dataset/example_pretty.json \ + --schema-output ./local/swebench-verified-harbor-dataset/schema.json +""" + +from __future__ import annotations + +import argparse +import base64 +import fnmatch +import gzip +import json +import os +import re +import shlex +from concurrent.futures import ThreadPoolExecutor, as_completed + +try: + import tomllib +except ModuleNotFoundError: # pragma: no cover - Python 3.10 compatibility. + import tomli as tomllib # type: ignore[no-redef] + +from collections.abc import Iterable +from pathlib import Path +from typing import Any, TypeVar + +try: + from tqdm import tqdm +except ModuleNotFoundError: # pragma: no cover - tqdm is optional at runtime. + tqdm = None # type: ignore[assignment] + +INLINE_TASK_FORMAT = "harbor_task_inline_v1" +DEFAULT_INLINE_FILES = ( + "instruction.md", + "task.toml", + "environment/Dockerfile", + "tests/test.sh", + "tests/config.json", + "solution/solve.sh", +) +T = TypeVar("T") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--input", + type=Path, + required=True, + help="A Harbor task directory, or a dataset directory containing Harbor task subdirectories.", + ) + parser.add_argument("--output", type=Path, required=True, help="Output JSONL file path.") + parser.add_argument( + "--task", + action="append", + default=[], + help="Task name or glob to include. Repeatable. If omitted, include all tasks under --input.", + ) + parser.add_argument( + "--exclude-task", + action="append", + default=[], + help="Task name or glob to exclude. Repeatable.", + ) + parser.add_argument("--limit", type=int, default=None, help="Maximum number of tasks to write after filtering.") + parser.add_argument("--offset", type=int, default=0, help="Number of filtered tasks to skip before writing.") + parser.add_argument( + "--source", + default=None, + help="Dataset/source name stored in metadata. Defaults to --input basename for dataset inputs.", + ) + parser.add_argument( + "--input-key", + default="prompt", + help="Primary slime prompt key to write. Defaults to prompt.", + ) + parser.add_argument( + "--prompt-alias-key", + default="", + help="Optional prompt alias key. Use '' to disable. Disabled by default.", + ) + parser.add_argument("--label-key", default="label", help="Label key to write. Use '' to disable.") + parser.add_argument("--metadata-key", default="metadata", help="Metadata key to write.") + parser.add_argument( + "--default-workdir", + default="/testbed", + help="Workdir fallback when environment/Dockerfile has no WORKDIR.", + ) + parser.add_argument( + "--image", + default=None, + help="Override image for all rows. By default it is extracted from the active Dockerfile FROM line.", + ) + parser.add_argument( + "--no-eval-cmd", + action="store_true", + help="Do not derive metadata.eval_cmd from tests/test.sh and tests/config.json.", + ) + parser.add_argument( + "--include-inline-files", + action="store_true", + help="Also include metadata.harbor_task.files for Harbor materialization/debugging. Disabled by default.", + ) + parser.add_argument( + "--inline-file", + action="append", + default=None, + help="Relative Harbor task file to include when --include-inline-files is set. Repeatable.", + ) + parser.add_argument( + "--provenance-root", + action="store_true", + help="Record the source dataset root and task name in metadata.source_provenance.", + ) + parser.add_argument( + "--workers", + type=int, + default=min(32, max(4, (os.cpu_count() or 4) * 4)), + help=( + "Number of worker threads used to read/convert tasks. Harbor conversion is I/O-bound, so the default uses several threads." + ), + ) + parser.add_argument( + "--no-progress", + action="store_true", + help="Disable tqdm progress bars.", + ) + parser.add_argument( + "--pretty-output", type=Path, default=None, help="Optional pretty JSON file for the first row." + ) + parser.add_argument("--schema-output", type=Path, default=None, help="Optional JSON schema output path.") + return parser.parse_args() + + +def find_task_dirs(input_path: Path, include_patterns: list[str], exclude_patterns: list[str]) -> list[Path]: + input_path = input_path.expanduser().resolve() + if _is_harbor_task_dir(input_path): + tasks = [input_path] + else: + tasks = _find_dataset_task_dirs(input_path, include_patterns) + + if include_patterns: + tasks = [path for path in tasks if _matches_any(path.name, include_patterns)] + if exclude_patterns: + tasks = [path for path in tasks if not _matches_any(path.name, exclude_patterns)] + return tasks + + +def _find_dataset_task_dirs(input_path: Path, include_patterns: list[str]) -> list[Path]: + """Find Harbor task directories under a dataset root. + + Exact ``--task`` names are resolved directly to avoid stat-ing every task on + slow shared filesystems. Glob patterns still require scanning the dataset + root. + """ + + exact_patterns = [pattern for pattern in include_patterns if not _has_glob(pattern)] + glob_patterns = [pattern for pattern in include_patterns if _has_glob(pattern)] + + tasks_by_name: dict[str, Path] = {} + for name in exact_patterns: + path = input_path / name + if path.is_dir() and _is_harbor_task_dir(path): + tasks_by_name[path.name] = path + + if not include_patterns or glob_patterns: + for path in sorted(input_path.iterdir()): + if not path.is_dir(): + continue + if glob_patterns and not _matches_any(path.name, glob_patterns): + continue + if _is_harbor_task_dir(path): + tasks_by_name[path.name] = path + + return [tasks_by_name[name] for name in sorted(tasks_by_name)] + + +def task_to_row( + task_dir: Path, + *, + dataset_root: Path, + source: str | None, + input_key: str, + prompt_alias_key: str, + label_key: str, + metadata_key: str, + image_override: str | None, + default_workdir: str, + include_eval_cmd: bool, + include_inline_files: bool, + inline_files: tuple[str, ...], + provenance_root: bool, +) -> dict[str, Any]: + task_dir = task_dir.resolve() + instruction = (task_dir / "instruction.md").read_text() + task_toml = read_task_toml(task_dir) + swe_config = read_swe_config(task_dir) + dockerfile = (task_dir / "environment" / "Dockerfile").read_text() + + instance_id = str(swe_config.get("instance_id") or task_dir.name) + source_name = source or dataset_root.name + # Both the prompt and metadata.problem_statement are Harbor's instruction.md: + # that is the exact string Harbor hands its agents, and having the two agree + # means a consumer reading either one gets the same text. The raw + # tests/config.json field is preserved under metadata.harbor.problem_statement + # -- it is upstream's own text, still carrying CRLF on about half of SWE-bench + # Verified, whereas instruction.md is dedent-normalised and wrapped in a + # "# Task" header plus a Repo/Version/Base commit/Instance ID block. + prompt = instruction + image = image_override or extract_dockerfile_image(dockerfile) + if not image: + raise ValueError(f"Cannot extract Docker image from {task_dir / 'environment' / 'Dockerfile'}") + workdir = extract_dockerfile_workdir(dockerfile) or default_workdir + row: dict[str, Any] = {input_key: prompt} + if prompt_alias_key and prompt_alias_key != input_key: + row[prompt_alias_key] = prompt + if label_key: + row[label_key] = instance_id + + metadata: dict[str, Any] = { + "instance_id": instance_id, + "source": source_name, + "image": image, + "workdir": workdir, + "problem_statement": instruction, + "harbor": harbor_metadata(task_dir, source_name, task_toml, swe_config, image, workdir), + } + if include_eval_cmd: + metadata["eval_cmd"] = build_eval_cmd(task_dir, swe_config) + if include_inline_files: + metadata["harbor_task"] = { + "format": INLINE_TASK_FORMAT, + "name": task_dir.name, + "source": source_name, + "files": read_inline_task_files(task_dir, inline_files), + } + if provenance_root: + metadata["source_provenance"] = { + "dataset_root": str(dataset_root), + "task_name": task_dir.name, + } + row[metadata_key] = metadata + return row + + +def read_task_toml(task_dir: Path) -> dict[str, Any]: + return tomllib.loads((task_dir / "task.toml").read_text()) + + +def read_swe_config(task_dir: Path) -> dict[str, Any]: + path = task_dir / "tests" / "config.json" + if not path.is_file(): + return {} + config = json.loads(path.read_text()) + for key in ("FAIL_TO_PASS", "PASS_TO_PASS"): + if isinstance(config.get(key), str): + try: + config[key] = json.loads(config[key]) + except json.JSONDecodeError: + pass + return config + + +def extract_dockerfile_image(dockerfile: str) -> str | None: + image = None + for raw_line in dockerfile.splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + match = re.match(r"FROM\s+([^\s]+)", line, flags=re.IGNORECASE) + if match: + image = match.group(1) + return image + + +def extract_dockerfile_workdir(dockerfile: str) -> str | None: + workdir = None + for raw_line in dockerfile.splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + match = re.match(r"WORKDIR\s+(.+)", line, flags=re.IGNORECASE) + if match: + workdir = match.group(1).strip().strip("\"'") + return workdir + + +def harbor_metadata( + task_dir: Path, + source_name: str, + task_toml: dict[str, Any], + swe_config: dict[str, Any], + image: str, + workdir: str, +) -> dict[str, Any]: + return _drop_none( + { + "task_name": task_dir.name, + "source": source_name, + # Upstream's raw text, as opposed to the top-level problem_statement, + # which is Harbor's rendered instruction.md (see task_to_row). + "problem_statement": swe_config.get("problem_statement"), + "repo": swe_config.get("repo"), + "version": swe_config.get("version"), + "base_commit": swe_config.get("base_commit"), + "difficulty": swe_config.get("difficulty") or (task_toml.get("metadata") or {}).get("difficulty"), + "docker_image": image, + "docker_workdir": workdir, + "fail_to_pass": swe_config.get("FAIL_TO_PASS"), + "pass_to_pass": swe_config.get("PASS_TO_PASS"), + "test_patch": swe_config.get("test_patch"), + "reference_patch": swe_config.get("patch"), + "verifier_timeout_sec": (task_toml.get("verifier") or {}).get("timeout_sec"), + "agent_timeout_sec": (task_toml.get("agent") or {}).get("timeout_sec"), + } + ) + + +def build_eval_cmd(task_dir: Path, swe_config: dict[str, Any]) -> str: + """Build an eval-only command that materializes Harbor verifier assets. + + AGS does not mount the Harbor task directory, so tests/config.json and + tests/test.sh must be embedded in the prompt-data row. Keep these grading + assets inside eval_cmd rather than pre_commands: pre_commands run before the + agent, while eval_cmd runs only after the foreground agent process returns. + A separate clean grading environment and a strong no-test-cheating boundary + require SWE_EVAL_ISOLATED_SANDBOX=true. With the default false, eval_cmd + materializes the assets later in the same root-capable agent sandbox. + + The embedded files intentionally use Harbor's canonical absolute paths + (/tests/config.json, /tests/test.sh, and /logs/verifier) instead of /tmp + paths because Harbor-generated tests/test.sh and downstream tooling expect + that layout. The payloads are gzip+base64 encoded so large SWE-bench + PASS_TO_PASS lists do not make the shell command exceed AGS/E2B argv limits. + """ + test_script_path = task_dir / "tests" / "test.sh" + if not test_script_path.is_file(): + raise FileNotFoundError(f"Cannot derive eval_cmd without {test_script_path}") + task_slug = _safe_slug(task_dir.name) + config_path = "/tests/config.json" + script_path = "/tests/test.sh" + test_script = test_script_path.read_text() + config_json = json.dumps(swe_config, ensure_ascii=False, separators=(",", ":")) + config_payload = _gzip_base64(config_json) + test_payload = _gzip_base64(test_script) + decoder_delim = f"SLIME_AGS_DECODE_{task_slug}" + decoder = "\n".join( + [ + "import base64, gzip, pathlib", + f"paths = [{config_path!r}, {script_path!r}]", + "payloads = [", + repr(config_payload), + ",", + repr(test_payload), + "]", + "for path, payload in zip(paths, payloads):", + " pathlib.Path(path).write_bytes(gzip.decompress(base64.b64decode(payload)))", + ] + ) + return "\n".join( + [ + "set -euo pipefail", + # /tests and /logs are part of Harbor's verifier contract. Create + # them only when eval_cmd starts; whether that is the existing agent + # sandbox or a clean sandbox is selected at rollout runtime. + "mkdir -p /tests /logs/verifier", + _python_heredoc(decoder, decoder_delim), + f"chmod +x {shlex.quote(script_path)}", + f"bash {shlex.quote(script_path)}", + ] + ) + + +def patch_harbor_test_script_for_ags(test_script: str, config_path: str, verifier_dir: str) -> str: + """Deprecated compatibility helper. + + Older converter output rewrote Harbor absolute paths to /tmp. Current + output preserves /tests/config.json and /logs/verifier and creates them in + eval_cmd, so this helper is intentionally a no-op except for callers that + still import it. + """ + return test_script + + +def read_inline_task_files(task_dir: Path, rel_paths: tuple[str, ...]) -> dict[str, str]: + files: dict[str, str] = {} + for rel_path in sorted(set(rel_paths)): + _validate_relative_file_path(rel_path) + path = task_dir / rel_path + if not path.is_file(): + raise FileNotFoundError(f"Inline Harbor task file not found: {path}") + files[rel_path] = path.read_text() + return files + + +def write_jsonl(rows: list[dict[str, Any]], output: Path) -> None: + output.parent.mkdir(parents=True, exist_ok=True) + with output.open("w") as f: + for row in rows: + f.write(json.dumps(row, ensure_ascii=False) + "\n") + + +def write_pretty_example(rows: list[dict[str, Any]], output: Path) -> None: + if not rows: + return + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(rows[0], ensure_ascii=False, indent=2) + "\n") + + +def write_schema(output: Path, *, input_key: str, prompt_alias_key: str, label_key: str, metadata_key: str) -> None: + row_required = [input_key, metadata_key] + properties: dict[str, Any] = { + input_key: { + "type": "string", + "description": "Primary slime prompt key; Harbor's instruction.md verbatim.", + }, + metadata_key: { + "type": "object", + "required": ["instance_id", "image", "workdir", "problem_statement"], + "properties": { + "instance_id": {"type": "string"}, + "source": {"type": "string"}, + "image": {"type": "string", "description": "Sandbox image consumed by ags_generator."}, + "workdir": {"type": "string", "description": "Repository path inside the sandbox."}, + "problem_statement": { + "type": "string", + "description": ( + "Harbor's instruction.md, same as the prompt. Upstream's raw " + "tests/config.json text is under harbor.problem_statement." + ), + }, + "eval_cmd": { + "type": "string", + "description": ( + "Reward command run after the agent; the rollout's " + "SWE_EVAL_ISOLATED_SANDBOX setting selects the sandbox." + ), + }, + "harbor": {"type": "object", "description": "Extracted Harbor/SWE provenance and grading fields."}, + "harbor_task": { + "type": "object", + "description": "Optional inline files, present only with --include-inline-files.", + }, + }, + }, + } + if prompt_alias_key and prompt_alias_key != input_key: + properties[prompt_alias_key] = {"type": "string", "description": "Prompt alias for --input-key prompt."} + if label_key: + properties[label_key] = {"type": "string", "description": "Optional label, usually the Harbor task name."} + + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text( + json.dumps( + { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "title": "slime AGS prompt-data row converted from Harbor task", + "type": "object", + "required": row_required, + "properties": properties, + }, + ensure_ascii=False, + indent=2, + ) + + "\n" + ) + + +def convert_tasks( + task_dirs: list[Path], + *, + dataset_root: Path, + source: str | None, + input_key: str, + prompt_alias_key: str, + label_key: str, + metadata_key: str, + image_override: str | None, + default_workdir: str, + include_eval_cmd: bool, + include_inline_files: bool, + inline_files: tuple[str, ...], + provenance_root: bool, + workers: int, + show_progress: bool, +) -> list[dict[str, Any]]: + """Convert tasks concurrently while preserving deterministic output order.""" + + if workers <= 1 or len(task_dirs) <= 1: + iterator = _progress(task_dirs, total=len(task_dirs), desc="Converting Harbor tasks", enabled=show_progress) + return [ + task_to_row( + task_dir, + dataset_root=dataset_root, + source=source, + input_key=input_key, + prompt_alias_key=prompt_alias_key, + label_key=label_key, + metadata_key=metadata_key, + image_override=image_override, + default_workdir=default_workdir, + include_eval_cmd=include_eval_cmd, + include_inline_files=include_inline_files, + inline_files=inline_files, + provenance_root=provenance_root, + ) + for task_dir in iterator + ] + + rows: list[dict[str, Any] | None] = [None] * len(task_dirs) + max_workers = min(max(1, workers), len(task_dirs)) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = { + executor.submit( + task_to_row, + task_dir, + dataset_root=dataset_root, + source=source, + input_key=input_key, + prompt_alias_key=prompt_alias_key, + label_key=label_key, + metadata_key=metadata_key, + image_override=image_override, + default_workdir=default_workdir, + include_eval_cmd=include_eval_cmd, + include_inline_files=include_inline_files, + inline_files=inline_files, + provenance_root=provenance_root, + ): index + for index, task_dir in enumerate(task_dirs) + } + for future in _progress( + as_completed(futures), + total=len(futures), + desc=f"Converting Harbor tasks ({max_workers} threads)", + enabled=show_progress, + ): + index = futures[future] + try: + rows[index] = future.result() + except Exception as exc: + raise RuntimeError(f"Failed to convert Harbor task {task_dirs[index]}") from exc + + return [row for row in rows if row is not None] + + +def _progress(iterable: Iterable[T], *, total: int, desc: str, enabled: bool) -> Iterable[T]: + if enabled and tqdm is not None: + return tqdm(iterable, total=total, desc=desc, unit="task") + return iterable + + +def _is_harbor_task_dir(path: Path) -> bool: + return (path / "instruction.md").is_file() and (path / "task.toml").is_file() + + +def _matches_any(name: str, patterns: list[str]) -> bool: + return any(fnmatch.fnmatch(name, pattern) for pattern in patterns) + + +def _has_glob(pattern: str) -> bool: + return any(ch in pattern for ch in "*?[") + + +def _validate_relative_file_path(rel_path: str) -> None: + path = Path(rel_path) + if path.is_absolute() or any(part in {"", ".", ".."} for part in path.parts): + raise ValueError(f"Unsafe task file path: {rel_path!r}") + + +def _safe_slug(value: str) -> str: + return re.sub(r"[^A-Za-z0-9_.-]+", "_", value)[:80] or "task" + + +def _heredoc(path: str, content: str, delimiter: str) -> str: + while delimiter in content: + delimiter += "_END" + return f"cat > {shlex.quote(path)} <<'{delimiter}'\n{content.rstrip()}\n{delimiter}" + + +def _gzip_base64(content: str) -> str: + return base64.b64encode(gzip.compress(content.encode("utf-8"), mtime=0)).decode("ascii") + + +def _python_heredoc(script: str, delimiter: str) -> str: + while delimiter in script: + delimiter += "_END" + return f"python3 - <<'{delimiter}'\n{script.rstrip()}\n{delimiter}" + + +def _drop_none(data: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in data.items() if value is not None} + + +def main() -> None: + args = parse_args() + input_path = args.input.expanduser().resolve() + task_dirs = find_task_dirs(input_path, args.task, args.exclude_task) + if args.offset: + task_dirs = task_dirs[args.offset :] + if args.limit is not None: + task_dirs = task_dirs[: args.limit] + if not task_dirs: + raise ValueError(f"No Harbor tasks found under {input_path}") + + dataset_root = input_path if not _is_harbor_task_dir(input_path) else input_path.parent + inline_files = tuple(args.inline_file) if args.inline_file else DEFAULT_INLINE_FILES + rows = convert_tasks( + task_dirs, + dataset_root=dataset_root, + source=args.source, + input_key=args.input_key, + prompt_alias_key=args.prompt_alias_key, + label_key=args.label_key, + metadata_key=args.metadata_key, + image_override=args.image, + default_workdir=args.default_workdir, + include_eval_cmd=not args.no_eval_cmd, + include_inline_files=args.include_inline_files, + inline_files=inline_files, + provenance_root=args.provenance_root, + workers=args.workers, + show_progress=not args.no_progress, + ) + + write_jsonl(rows, args.output) + if args.pretty_output: + write_pretty_example(rows, args.pretty_output) + if args.schema_output: + write_schema( + args.schema_output, + input_key=args.input_key, + prompt_alias_key=args.prompt_alias_key, + label_key=args.label_key, + metadata_key=args.metadata_key, + ) + print(f"Wrote {len(rows)} rows to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/train.py b/train.py index 9cb2968866..35b4282a0b 100644 --- a/train.py +++ b/train.py @@ -80,14 +80,18 @@ def offload_train(actor_trains_this_step): ray.get(rollout_manager.save.remote(rollout_id)) offload_train(actor_trains) - if args.offload_rollout and not release_train: - ray.get(rollout_manager.onload_weights.remote()) - actor_model.update_weights() - if args.offload_rollout: - ray.get(rollout_manager.onload_kv.remote()) + will_eval = should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch) + needs_rollout_after_step = (rollout_id + 1 < args.num_rollout) or will_eval + if needs_rollout_after_step: + if args.offload_rollout and not release_train: + ray.get(rollout_manager.onload_weights.remote()) + actor_model.update_weights() + + if args.offload_rollout: + ray.get(rollout_manager.onload_kv.remote()) - if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch): + if will_eval: ray.get(rollout_manager.eval.remote(rollout_id)) ray.get(rollout_manager.dispose.remote())