From d4d095223feb5caeb3d0dab5dd1eadb0dbdfd8b3 Mon Sep 17 00:00:00 2001 From: FunJim Date: Tue, 7 Jul 2026 16:22:31 +0800 Subject: [PATCH 01/43] Add AGS rollout buffer generator --- .../rollout_buffer/generator/ags_generator.py | 9 + .../generator/ags_generator/__init__.py | 5 + .../ags_generator/adapter_service.py | 55 ++++ .../generator/ags_generator/ags_sandbox.py | 267 ++++++++++++++++++ .../generator/ags_generator/artifacts.py | 70 +++++ .../generator/ags_generator/config.py | 51 ++++ .../generator/ags_generator/entry.py | 187 ++++++++++++ .../generator/ags_generator/harnesses.py | 134 +++++++++ .../generator/ags_generator/rollout.py | 210 ++++++++++++++ .../generator/ags_generator/runner.py | 57 ++++ .../generator/ags_generator/sampling.py | 21 ++ .../generator/ags_generator/serialization.py | 43 +++ .../generator/ags_generator/source.py | 50 ++++ .../generator/ags_generator/swe_task.py | 203 +++++++++++++ .../rollout_buffer/rollout_buffer_example.py | 23 +- .../test_rollout_buffer/test_ags_generator.py | 62 ++++ 16 files changed, 1445 insertions(+), 2 deletions(-) create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/__init__.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/ags_sandbox.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/artifacts.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/config.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/entry.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/rollout.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/runner.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/sampling.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/serialization.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/source.py create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py create mode 100644 tests/test_rollout_buffer/test_ags_generator.py 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..2879d328fb --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator.py @@ -0,0 +1,9 @@ +"""Compatibility shim so buffer.py can discover the ags package.""" + +from slime_plugins.rollout_buffer.generator.ags_generator.entry import ( # noqa: F401 + TASK_TYPE, + 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..da4f371e11 --- /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, run_rollout + +__all__ = ["TASK_TYPE", "run_rollout"] 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..d63fa099c5 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py @@ -0,0 +1,55 @@ +"""SGLang-to-agent adapter service for AGS generator workers.""" + +from __future__ import annotations + +import logging +import os +from argparse import Namespace + +from slime.agent.aiohttp_threaded import FilteredAccessLogger, run_app_in_thread +from slime.utils.misc import SingletonMeta +from slime.utils.processing_utils import load_tokenizer + +from .config import AGSGeneratorConfig + +logger = logging.getLogger(__name__) + + +class AdapterService(metaclass=SingletonMeta): + def __init__(self, args: Namespace, config: AGSGeneratorConfig, adapter_cls: type) -> 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}" + ) + if not config.adapter_public_host: + raise RuntimeError("ADAPTER_PUBLIC_HOST 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, + ) + self.app_handle = run_app_in_thread( + self.adapter.app, + host=config.adapter_bind_host, + port=config.adapter_port, + thread_name="ags-rollout-adapter", + runner_kwargs={"handler_cancellation": True, "access_log_class": FilteredAccessLogger}, + ) + self.adapter_url = f"http://{config.adapter_public_host}:{self.app_handle.port}" + 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, + ) 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..40a8cbbf24 --- /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":"4","memory":"16Gi"}') + try: + parsed = json.loads(raw) + return parsed if isinstance(parsed, dict) else {"cpu": "4", "memory": "16Gi"} + except Exception: + return {"cpu": "4", "memory": "16Gi"} + + 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..4b272b2edd --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/config.py @@ -0,0 +1,51 @@ +"""Configuration for the AGS coding-agent rollout-buffer generator.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + + +@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 + rollout_guard_sec: int + boot_concurrency: int + boot_retries: int + artifact_dir: str | None + prompt: str + + @classmethod + def from_env(cls) -> 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 + 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, + rollout_guard_sec=guard, + boot_concurrency=int(os.environ.get("SWE_BOOT_CONCURRENCY", "16")), + boot_retries=int(os.environ.get("SWE_BOOT_RETRIES", "2")), + artifact_dir=os.environ.get("TRAJECTORY_DUMP_DIR", "").strip() or None, + 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.", + ), + ) 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..5dea5dfef0 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -0,0 +1,187 @@ +"""rollout_buffer generator entry point for AGS coding-agent tasks.""" + +from __future__ import annotations + +import asyncio +import logging +from argparse import Namespace +from typing import Any + +import requests + +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__) + + +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() + 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("rollout_batch_size") or data.get("num_groups_per_epoch") or args.rollout_batch_size + ) + skip_instance_ids = data.get("skip_instance_ids") or [] + + logger.info( + "[ags_generator] start task_type=%s groups_per_epoch=%s repeats=%s epochs=%s buffer=%s", + TASK_TYPE, + groups_per_epoch, + args.n_samples_per_prompt, + num_epoch, + remote_buffer_url, + ) + + async def _run_epoch(epoch: int, samples: list[Sample]) -> None: + for sample in samples: + 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 {}), + }, + ) + _send_data_to_buffer(remote_buffer_url, item) + + for epoch in range(num_epoch): + samples = source.get_repeated_samples(groups_per_epoch, skip_instance_ids=skip_instance_ids) + skip_instance_ids = [] + asyncio.run(_run_epoch(epoch, samples)) + return "finished" + + +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} + for items in temp_data.values(): + for item in items: + samples = samples_from_payload(item) + for sample in samples: + 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 {} + artifact_counts["trajectory"] += int(bool(metadata.get("trajectory_path"))) + artifact_counts["patch"] += int(bool(metadata.get("patch_path"))) + artifact_counts["rollout_dump"] += int(bool(metadata.get("rollout_dump_path"))) + total = sum(len(items) for items in temp_data.values()) + return { + "total_samples": total, + "num_groups": len(temp_data), + "avg_group_size": total / len(temp_data) if temp_data else 0, + "avg_reward": sum(rewards) / len(rewards) if rewards else 0, + "nonzero_reward_samples": sum(1 for reward in rewards if reward != 0), + "status_counts": status_counts, + "artifact_counts": artifact_counts, + } + + +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)), + 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"), + 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) 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..013381042e --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py @@ -0,0 +1,134 @@ +"""Harness registry for AGS-backed coding agents.""" + +from __future__ import annotations + +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" + ) + 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( + "mkdir -p /root/.claude /home/agent/.claude && " + f"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}" + extra = os.environ.get(self.extra_args_env, "").strip() + if extra: + 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: + env.update(json.loads(extra_envs)) + if os.environ.get("CLAUDE_CODE_MAX_OUTPUT_TOKENS"): + env["CLAUDE_CODE_MAX_OUTPUT_TOKENS"] = os.environ["CLAUDE_CODE_MAX_OUTPUT_TOKENS"] + + 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): + """Placeholder for AGS CodeBuddy Code sidecar integration. + + The package-level registry can add the harness once its non-interactive CLI + contract is finalized without touching the rollout orchestration code. + """ + + name = "codebuddy_code" + + async def install_cli(self, sb: Sandbox) -> None: + raise NotImplementedError("CodeBuddy Code AGS sidecar CLI contract is not configured yet") + + async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: + raise NotImplementedError("CodeBuddy Code AGS sidecar CLI contract is not configured yet") + + async def launch_and_wait(self, sb: Sandbox, ctx: HarnessContext, prompt: str, time_budget_sec: int) -> int: + raise NotImplementedError("CodeBuddy Code AGS sidecar CLI contract is not configured yet") + + +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..35478d259a --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py @@ -0,0 +1,210 @@ +"""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 AdapterService +from .ags_sandbox import AGSSandbox +from .artifacts import ArtifactWriter, sample_artifact_id +from .config import AGSGeneratorConfig +from .harnesses import resolve_agent +from .sampling import normalize_sampling_params +from .swe_task import evaluate, get_metadata, git_diff, prepare_workspace + +logger = logging.getLogger(__name__) + + +class AGSRolloutRunner: + def __init__(self, args: Namespace, config: AGSGeneratorConfig | None = None) -> None: + self.args = args + self.config = config or AGSGeneratorConfig.from_env() + self.harness_cls, self.adapter_cls = resolve_agent(self.config.agent_name) + self.adapter_service = AdapterService(args, self.config, self.adapter_cls) + self.artifacts = ArtifactWriter(self.config.artifact_dir) + self._boot_sem = asyncio.Semaphore(self.config.boot_concurrency) + + async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sample]: + md = get_metadata(base_sample) + instance_id = md["instance_id"] + if not md["image"] or not md["workdir"]: + return self._abort_result(base_sample, "missing_image_or_workdir", instance_id) + + session_id = base_sample.session_id = _session_id(base_sample, instance_id) + artifact_id = sample_artifact_id(instance_id, base_sample) + normalized_sampling = normalize_sampling_params(sampling_params) + self.adapter_service.adapter.open_session( + session_id, + sampling_defaults=normalized_sampling, + max_context_tokens=self.adapter_service.max_context_len, + ) + t0 = time.time() + try: + async with asyncio.timeout(self.config.rollout_guard_sec): + async with self._boot_agent_sandbox(md["image"], instance_id) as sb: + await prepare_workspace(sb, md["workdir"], md) + 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.config.prompt, + ) + trajectory_path = await self.artifacts.dump_trajectory(sb, md["workdir"], artifact_id) + diff_text = await git_diff(sb, md["workdir"]) + patch_path = self.artifacts.dump_patch(diff_text, artifact_id) + + reward, applied_cleanly = await evaluate( + image=md["image"], + workdir=md["workdir"], + diff_text=diff_text, + 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, + ) + 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: + return self._abort_result(base_sample, "adapter_session_empty", instance_id) + + 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), + "agent_exit_code": agent_exit_code, + "elapsed_sec": time.time() - t0, + "num_samples": len(samples), + "patch_path": patch_path, + "trajectory_path": trajectory_path, + }, + artifact_id, + ) + 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), + "trajectory_path": trajectory_path, + "patch_path": patch_path, + "rollout_dump_path": rollout_path, + } + logger.info( + "[ags_generator] %s: reward=%.2f applied=%s exit=%s elapsed=%.1fs segments=%d", + instance_id, + float(reward), + bool(applied_cleanly), + agent_exit_code, + time.time() - t0, + len(samples), + ) + return samples + except asyncio.TimeoutError: + _log_timeout_diagnostic(t0, instance_id, self.config.rollout_guard_sec) + return self._abort_result(base_sample, "wall_clock_timeout", instance_id) + except Exception as exc: + logger.warning("[ags_generator] %s: rollout failed: %s\n%s", instance_id, exc, traceback.format_exc()) + return self._abort_result(base_sample, f"exception:{type(exc).__name__}", instance_id) + finally: + await self.adapter_service.adapter.drop_session(session_id) + + @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: + if sample.session_id: + return sample.session_id + if sample.index is not None and sample.group_index is not None: + return f"cagent-{instance_id}-{sample.index}-{sample.group_index}" + return f"cagent-{instance_id}-{secrets.token_hex(8)}" + + +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..668742fbb2 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/runner.py @@ -0,0 +1,57 @@ +"""Command-running helpers for AGS sidecar harnesses.""" + +from __future__ import annotations + +import asyncio +import shlex +import time + +try: + from slime.agent.sandbox import EXIT_TIME_BUDGET_EXCEEDED +except ImportError: # older slime checkout keeps this constant in harness.common + from slime.agent.harness.common import EXIT_TIME_BUDGET_EXCEEDED +from slime.agent.sandbox import Sandbox + + +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" + 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"env {export_lines} setsid {launcher} < /dev/null > /dev/null 2>&1 &", + 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 + 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..a6096ff388 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/source.py @@ -0,0 +1,50 @@ +"""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) + + 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..eda9f760b0 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py @@ -0,0 +1,203 @@ +"""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), + "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]) -> None: + 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) + 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( + *, + 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 + + async with AGSSandbox(image) as ev: + await agent_sandbox.ensure_agent_user(ev, workdir) + if swepro: + await _setup_swepro_assets(ev, swepro) + await apply_before_repo_set_cmd(ev, workdir, swepro) + 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))) + + 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: + ec, _, _ = await ev.exec(f"cd {workdir} && {cmd}", user="agent", 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/rollout_buffer_example.py b/slime_plugins/rollout_buffer/rollout_buffer_example.py index 74b4b4c46c..eafcef35a0 100644 --- a/slime_plugins/rollout_buffer/rollout_buffer_example.py +++ b/slime_plugins/rollout_buffer/rollout_buffer_example.py @@ -45,8 +45,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): @@ -189,6 +191,18 @@ def start_rollout(api_base_url: str, args, metadata): "top_p": args.rollout_top_p, }, "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_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), "skip_instance_ids": finished_groups_instance_id_list, } print("start rollout with payload: ", payload) @@ -269,6 +283,11 @@ async def generate_rollout_async(args, rollout_id: int, data_buffer, evaluation: 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.append(compact_samples) + continue + oai_messages = record["messages"] mask_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type=args.loss_mask_type) 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..2f289880a9 --- /dev/null +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from slime.utils.types import Sample +from slime_plugins.rollout_buffer.generator.ags_generator.entry import ( + get_group_data_meta_info, + is_valid_group, + transform_group, +) +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, +) + + +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} + + +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, + } From a4b921df58951c69927d50148a5563e3d54f1a60 Mon Sep 17 00:00:00 2001 From: FunJim Date: Wed, 8 Jul 2026 10:41:44 +0800 Subject: [PATCH 02/43] Add AGS rollout concurrency --- .../generator/ags_generator/config.py | 3 ++ .../generator/ags_generator/entry.py | 41 +++++++++++-------- 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/config.py b/slime_plugins/rollout_buffer/generator/ags_generator/config.py index 4b272b2edd..5beeb65f88 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/config.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/config.py @@ -18,6 +18,7 @@ class AGSGeneratorConfig: eval_bootstrap_cmd: str | None rollout_guard_sec: int boot_concurrency: int + rollout_concurrency: int boot_retries: int artifact_dir: str | None prompt: str @@ -28,6 +29,7 @@ def from_env(cls) -> AGSGeneratorConfig: 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"), @@ -39,6 +41,7 @@ def from_env(cls) -> AGSGeneratorConfig: eval_bootstrap_cmd=os.environ.get("SWE_EVAL_BOOTSTRAP_CMD") or None, 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, prompt=os.environ.get( diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py index 5dea5dfef0..46eb1af1b6 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -37,30 +37,39 @@ def run_rollout(data: dict[str, Any]) -> str: skip_instance_ids = data.get("skip_instance_ids") or [] logger.info( - "[ags_generator] start task_type=%s groups_per_epoch=%s repeats=%s epochs=%s buffer=%s", + "[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 {}), + }, + ) + await asyncio.to_thread(_send_data_to_buffer, remote_buffer_url, item) + async def _run_epoch(epoch: int, samples: list[Sample]) -> None: - for sample in samples: - 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 {}), - }, - ) - _send_data_to_buffer(remote_buffer_url, item) + semaphore = asyncio.Semaphore(config.rollout_concurrency) + + async def _guarded(sample: Sample) -> None: + async with semaphore: + await _run_sample(epoch, sample) + + await asyncio.gather(*(_guarded(sample) for sample in samples)) for epoch in range(num_epoch): samples = source.get_repeated_samples(groups_per_epoch, skip_instance_ids=skip_instance_ids) From 29a5e75901743a12adf6e144c59eb0d65ce224ec Mon Sep 17 00:00:00 2001 From: FunJim Date: Wed, 8 Jul 2026 15:05:36 +0800 Subject: [PATCH 03/43] Log AGS rollout metrics --- .../generator/ags_generator/entry.py | 73 +++++- .../generator/ags_generator/rollout.py | 6 +- .../rollout_buffer/rollout_buffer_example.py | 211 ++++++++++++++---- 3 files changed, 237 insertions(+), 53 deletions(-) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py index 46eb1af1b6..4a0ee818f0 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -98,30 +98,87 @@ def is_valid_group(group, min_valid_group_size: int, task_type: str = TASK_TYPE) 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} + 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 {} - artifact_counts["trajectory"] += int(bool(metadata.get("trajectory_path"))) - artifact_counts["patch"] += int(bool(metadata.get("patch_path"))) - artifact_counts["rollout_dump"] += int(bool(metadata.get("rollout_dump_path"))) - total = sum(len(items) for items in temp_data.values()) + 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, + "total_samples": total_samples, + "total_rollouts": total_rollouts, "num_groups": len(temp_data), - "avg_group_size": total / len(temp_data) if temp_data else 0, + "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": sum(1 for reward in rewards if reward != 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: diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py index 35478d259a..5e61ded430 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py @@ -103,6 +103,7 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam }, artifact_id, ) + elapsed_sec = time.time() - t0 for sample in samples: sample.metadata = { **(sample.metadata or {}), @@ -112,6 +113,9 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam "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, } logger.info( "[ags_generator] %s: reward=%.2f applied=%s exit=%s elapsed=%.1fs segments=%d", @@ -119,7 +123,7 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam float(reward), bool(applied_cleanly), agent_exit_code, - time.time() - t0, + elapsed_sec, len(samples), ) return samples diff --git a/slime_plugins/rollout_buffer/rollout_buffer_example.py b/slime_plugins/rollout_buffer/rollout_buffer_example.py index eafcef35a0..15c8f2222c 100644 --- a/slime_plugins/rollout_buffer/rollout_buffer_example.py +++ b/slime_plugins/rollout_buffer/rollout_buffer_example.py @@ -90,51 +90,174 @@ def get_group_timestamp(group_items): 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) - } - - 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 - ) + 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 + ) - 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) - - 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 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}") + 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]]: From 3a7ebad0e9dbbe239a16ca5ad5ede7c652e96d5f Mon Sep 17 00:00:00 2001 From: FunJim Date: Thu, 9 Jul 2026 15:52:53 +0800 Subject: [PATCH 04/43] Fix AGS rollout session isolation Use fresh adapter session IDs for each AGS attempt and isolate per-sample failures so one bad rollout does not stop the generator. --- .../generator/ags_generator/entry.py | 30 +++++++++++- .../generator/ags_generator/rollout.py | 48 ++++++++++++++----- 2 files changed, 65 insertions(+), 13 deletions(-) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py index 4a0ee818f0..97b9e3a54f 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -67,7 +67,35 @@ async def _run_epoch(epoch: int, samples: list[Sample]) -> None: async def _guarded(sample: Sample) -> None: async with semaphore: - await _run_sample(epoch, sample) + try: + await _run_sample(epoch, sample) + 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 {}), + }, + ) + 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, + ) await asyncio.gather(*(_guarded(sample) for sample in samples)) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py index 5e61ded430..48f48022f6 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py @@ -40,16 +40,20 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam if not md["image"] or not md["workdir"]: return self._abort_result(base_sample, "missing_image_or_workdir", instance_id) - session_id = base_sample.session_id = _session_id(base_sample, 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) - self.adapter_service.adapter.open_session( - session_id, - sampling_defaults=normalized_sampling, - max_context_tokens=self.adapter_service.max_context_len, - ) t0 = time.time() + session_opened = False 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: await prepare_workspace(sb, md["workdir"], md) @@ -134,7 +138,16 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam logger.warning("[ags_generator] %s: rollout failed: %s\n%s", instance_id, exc, traceback.format_exc()) return self._abort_result(base_sample, f"exception:{type(exc).__name__}", instance_id) finally: - await self.adapter_service.adapter.drop_session(session_id) + 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(), + ) @asynccontextmanager async def _boot_agent_sandbox(self, image: str, instance_id: str) -> AsyncIterator[AGSSandbox]: @@ -187,11 +200,22 @@ def _abort_result(self, sample: Sample, reason: str, instance_id: str) -> list[S def _session_id(sample: Sample, instance_id: str) -> str: - if sample.session_id: - return sample.session_id - if sample.index is not None and sample.group_index is not None: - return f"cagent-{instance_id}-{sample.index}-{sample.group_index}" - return f"cagent-{instance_id}-{secrets.token_hex(8)}" + """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: From ddc6e2024424068dd633ab758a2a1192ec707dbf Mon Sep 17 00:00:00 2001 From: FunJim Date: Fri, 10 Jul 2026 17:21:34 +0800 Subject: [PATCH 05/43] Add Harbor prompt data converter Convert Harbor SWE-style tasks into AGS prompt data while preserving the canonical verifier layout, and run eval_cmd as root so Harbor test.sh can execute unmodified in the evaluator sandbox. --- .../generator/ags_generator/swe_task.py | 10 +- tools/harbor_task_to_slime_prompt_data.py | 649 ++++++++++++++++++ 2 files changed, 658 insertions(+), 1 deletion(-) create mode 100644 tools/harbor_task_to_slime_prompt_data.py diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py b/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py index eda9f760b0..482edc7dde 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py @@ -169,7 +169,15 @@ async def _run_swepro(ev: Sandbox, workdir: str, swepro: dict[str, Any], timeout async def _run_eval_cmd(ev: Sandbox, workdir: str, cmd: str, timeout: int) -> float: - ec, _, _ = await ev.exec(f"cd {workdir} && {cmd}", user="agent", check=False, timeout=timeout) + # 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. This happens only + # in the separate evaluator sandbox after the agent patch is collected, so + # hidden grading assets are not exposed to the agent sandbox. + ec, _, _ = await ev.exec(f"cd {workdir} && {cmd}", user="root", check=False, timeout=timeout) return 1.0 if ec == 0 else 0.0 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..135bede2ee --- /dev/null +++ b/tools/harbor_task_to_slime_prompt_data.py @@ -0,0 +1,649 @@ +#!/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.pre_commands + - metadata.eval_cmd + +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. pre_commands run in both the agent sandbox and the eval +sandbox, so putting tests/config.json there would leak hidden grading data to +the agent. eval_cmd runs only in the eval sandbox. To preserve Harbor's +expected verifier layout, eval_cmd materializes /tests/config.json and +/logs/verifier before invoking the patched Harbor tests/test.sh. + +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 fnmatch +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( + "--prompt-source", + choices=("problem_statement", "instruction"), + default="problem_statement", + help="Which extracted text to put in the primary prompt field.", + ) + 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-pre-commands", + action="store_true", + help="Do not write metadata.pre_commands to reset the repo to base_commit.", + ) + 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, + prompt_source: str, + image_override: str | None, + default_workdir: str, + include_pre_commands: bool, + 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 + problem_statement = str(swe_config.get("problem_statement") or instruction) + prompt = problem_statement if prompt_source == "problem_statement" else 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 + base_commit = swe_config.get("base_commit") + + 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": problem_statement, + "harbor": harbor_metadata(task_dir, source_name, task_toml, swe_config, image, workdir), + } + if include_pre_commands and base_commit: + metadata["pre_commands"] = [ + f"git checkout {shlex.quote(str(base_commit))} -f", + "git clean -fd", + ] + 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, + "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 hidden + grading assets inside eval_cmd rather than pre_commands: pre_commands are + executed in the agent sandbox before the agent runs, while eval_cmd is only + executed in the separate evaluator sandbox after the agent patch is + collected. This preserves the SWE-bench setting and avoids exposing + FAIL_TO_PASS/PASS_TO_PASS/test_patch/reference_patch to the agent. + + 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. + """ + 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, indent=2) + return "\n".join( + [ + "set -euo pipefail", + # /tests and /logs are part of Harbor's verifier contract. Create + # them here so they exist only in the eval sandbox, not in the + # agent sandbox. + "mkdir -p /tests /logs/verifier", + _heredoc(config_path, config_json, f"SLIME_AGS_CONFIG_{task_slug}"), + _heredoc(script_path, test_script, f"SLIME_AGS_TEST_{task_slug}"), + 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."}, + 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"}, + "pre_commands": {"type": "array", "items": {"type": "string"}}, + "eval_cmd": {"type": "string", "description": "Reward command; exit 0 means reward 1."}, + "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, + prompt_source: str, + image_override: str | None, + default_workdir: str, + include_pre_commands: bool, + 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, + prompt_source=prompt_source, + image_override=image_override, + default_workdir=default_workdir, + include_pre_commands=include_pre_commands, + 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, + prompt_source=prompt_source, + image_override=image_override, + default_workdir=default_workdir, + include_pre_commands=include_pre_commands, + 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 _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, + prompt_source=args.prompt_source, + image_override=args.image, + default_workdir=args.default_workdir, + include_pre_commands=not args.no_pre_commands, + 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() From 615f25d87ef1a912aff92b167b2ae401b4535939 Mon Sep 17 00:00:00 2001 From: FunJim Date: Fri, 10 Jul 2026 19:05:18 +0800 Subject: [PATCH 06/43] Add CodeBuddy Code AGS harness --- .../generator/ags_generator/harnesses.py | 182 +++++++++++++++++- .../test_rollout_buffer/test_ags_generator.py | 95 ++++++++- 2 files changed, 268 insertions(+), 9 deletions(-) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py index 013381042e..dd060848f5 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py @@ -2,6 +2,7 @@ from __future__ import annotations +import base64 import json import os import shlex @@ -103,22 +104,187 @@ async def run( class CodeBuddyCodeHarness(BaseHarness): - """Placeholder for AGS CodeBuddy Code sidecar integration. - - The package-level registry can add the harness once its non-interactive CLI - contract is finalized without touching the rollout orchestration code. - """ + """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" + max_turns_env = "SLIME_AGENT_CBC_MAX_TURNS" + max_output_tokens_env = "SLIME_AGENT_CBC_MAX_OUTPUT_TOKENS" + thinking_enabled_env = "SLIME_AGENT_CBC_THINKING_ENABLED" + tools_env = "SLIME_AGENT_CBC_TOOLS" + + # Keep the default tool surface close to Claude Code's coding-agent use case + # while disabling internet search by default for reproducible SWE rollouts. + allowed_tools = ( + "Bash", + "Read", + "Write", + "Edit", + "Glob", + "Grep", + "TaskCreate", + "TaskUpdate", + "TaskGet", + "TaskList", + "Agent", + ) async def install_cli(self, sb: Sandbox) -> None: - raise NotImplementedError("CodeBuddy Code AGS sidecar CLI contract is not configured yet") + 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: - raise NotImplementedError("CodeBuddy Code AGS sidecar CLI contract is not configured yet") + models_json = { + "models": [ + { + "id": ctx.model_label, + "name": ctx.model_label, + "vendor": "OpenAI", + "apiKey": ctx.session_id, + "url": self._chat_completions_url(ctx.adapter_url), + "maxOutputTokens": int(os.environ.get(self.max_output_tokens_env, "16384")), + "supportsToolCall": True, + "supportsImages": False, + "supportsReasoning": True, + } + ], + "availableModels": [ctx.model_label], + } + settings_json = { + "cleanupPeriodDays": 30, + "includeCoAuthoredBy": False, + "autoCompactEnabled": True, + "alwaysThinkingEnabled": _env_flag(self.thinking_enabled_env, default=True), + "showTokensCounter": False, + "enablePasteImageFromClipboard": False, + "enableTerminalProgressBar": False, + "fileCheckpointingEnabled": False, + "promptSuggestionEnabled": False, + "enableAllProjectMcpServers": 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: - raise NotImplementedError("CodeBuddy Code AGS sidecar CLI contract is not configured yet") + parts: list[str] = [ + f"--model {shlex.quote(ctx.model_label)}", + "--output-format json", + ] + tools = os.environ.get(self.tools_env, ",".join(self.allowed_tools)).strip() + if tools: + parts.append(f"--tools {shlex.quote(tools)}") + parts.append("--disallowedTools WebSearch") + extra = os.environ.get(self.extra_args_env, "").strip() + if extra: + # Keep caller-provided flags before the non-variadic tail and prompt. + parts.append(extra) + if not _env_flag(self.thinking_enabled_env, default=True): + parts.append("--effort none") + parts.append(f"--max-turns {int(os.environ.get(self.max_turns_env, '100'))}") + parts.append("-y") + + session_log_dir = f"{ctx.workdir}/.harness/codebuddy_sessions" + raw_cmd = ( + f"cbc {' '.join(parts)} {shlex.quote(prompt)}; " + "rc=$?; " + f"mkdir -p {shlex.quote(session_log_dir)}/projects; " + f"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 _json_b64(value: dict) -> str: + payload = json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + return base64.b64encode(payload).decode("ascii") + + +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"} HARNESS_REGISTRY: dict[str, tuple[type[BaseHarness], type]] = { diff --git a/tests/test_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py index 2f289880a9..238598c9c4 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -1,11 +1,19 @@ from __future__ import annotations +import asyncio +import base64 +import json +import re + +from tests.test_agent._fakes import FakeSandbox + from slime.utils.types import Sample from slime_plugins.rollout_buffer.generator.ags_generator.entry import ( get_group_data_meta_info, is_valid_group, transform_group, ) +from slime_plugins.rollout_buffer.generator.ags_generator.harnesses import CodeBuddyCodeHarness, resolve_agent 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, @@ -52,7 +60,7 @@ def test_group_hooks_accept_complete_sample_payloads(): 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} + assert meta["artifact_counts"] == {"trajectory": 1, "patch": 1, "rollout_dump": 1, "complete": 1} def test_sampling_params_use_sglang_generate_names(): @@ -60,3 +68,88 @@ def test_sampling_params_use_sglang_generate_names(): "max_new_tokens": 128, "temperature": 1.0, } + + +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(monkeypatch): + async def run_case(): + monkeypatch.setenv("SLIME_AGENT_CBC_MAX_OUTPUT_TOKENS", "8192") + monkeypatch.setenv("SLIME_AGENT_CBC_THINKING_ENABLED", "false") + 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]["maxOutputTokens"] == 8192 + assert models["models"][0]["supportsToolCall"] is True + assert settings["alwaysThinkingEnabled"] is False + + asyncio.run(run_case()) + + +def test_codebuddy_code_launch_command_and_env(monkeypatch): + async def run_case(): + monkeypatch.setenv("SLIME_AGENT_CBC_MAX_TURNS", "7") + monkeypatch.setenv("SLIME_AGENT_CBC_THINKING_ENABLED", "false") + 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 --output-format json" in body + assert "--max-turns 7" in body + assert "-y" in body and "solve it" in body + assert "--effort none" in body + assert "--tools Bash,Read,Write,Edit,Glob,Grep,TaskCreate,TaskUpdate,TaskGet,TaskList,Agent" in body + assert "codebuddy_sessions" in body + + 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 + + asyncio.run(run_case()) From bc76729f083da94871562d70f3beea2c0849791b Mon Sep 17 00:00:00 2001 From: FunJim Date: Mon, 13 Jul 2026 12:47:02 +0800 Subject: [PATCH 07/43] Add top_k to rollout sampling params --- slime_plugins/rollout_buffer/rollout_buffer_example.py | 1 + 1 file changed, 1 insertion(+) diff --git a/slime_plugins/rollout_buffer/rollout_buffer_example.py b/slime_plugins/rollout_buffer/rollout_buffer_example.py index 15c8f2222c..ea0a56ded3 100644 --- a/slime_plugins/rollout_buffer/rollout_buffer_example.py +++ b/slime_plugins/rollout_buffer/rollout_buffer_example.py @@ -312,6 +312,7 @@ 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"), From 5ef52677aa39bc831699cc4c646af3a99cfeec81 Mon Sep 17 00:00:00 2001 From: FunJim Date: Tue, 14 Jul 2026 11:11:27 +0800 Subject: [PATCH 08/43] Add Weave tracing for AGS rollouts Log AGS rollout calls, trajectory events, and sample payloads to Weave when W&B online logging is enabled, with an opt-in flag to decode token IDs for trace readability. --- requirements.txt | 1 + slime/utils/arguments.py | 6 + .../generator/ags_generator/config.py | 9 +- .../generator/ags_generator/entry.py | 8 +- .../generator/ags_generator/rollout.py | 51 +++- .../generator/ags_generator/weave_trace.py | 247 ++++++++++++++++++ .../rollout_buffer/rollout_buffer_example.py | 9 +- .../test_rollout_buffer/test_ags_generator.py | 224 ++++++++++++++++ 8 files changed, 542 insertions(+), 13 deletions(-) create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/weave_trace.py 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/utils/arguments.py b/slime/utils/arguments.py index ccb121f35d..4c56cffcad 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -1408,6 +1408,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_plugins/rollout_buffer/generator/ags_generator/config.py b/slime_plugins/rollout_buffer/generator/ags_generator/config.py index 5beeb65f88..a6d2b271db 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/config.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/config.py @@ -21,10 +21,11 @@ class AGSGeneratorConfig: rollout_concurrency: int boot_retries: int artifact_dir: str | None + enable_token2text: bool prompt: str @classmethod - def from_env(cls) -> AGSGeneratorConfig: + 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) @@ -44,11 +45,9 @@ def from_env(cls) -> AGSGeneratorConfig: 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.", + "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.", ), ) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py index 97b9e3a54f..39790cc113 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -26,7 +26,7 @@ def run_rollout(data: dict[str, Any]) -> str: logging.basicConfig(level=getattr(logging, data.get("log_level", "INFO"), logging.INFO)) args = _build_args(data) - config = AGSGeneratorConfig.from_env() + 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" @@ -235,6 +235,12 @@ def _build_args(data: dict[str, Any]) -> Namespace: 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}, ) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py index 48f48022f6..3dad3b0aca 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py @@ -21,6 +21,7 @@ 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__) @@ -32,21 +33,35 @@ def __init__(self, args: Namespace, config: AGSGeneratorConfig | None = None) -> self.harness_cls, self.adapter_cls = resolve_agent(self.config.agent_name) self.adapter_service = AdapterService(args, self.config, self.adapter_cls) self.artifacts = ArtifactWriter(self.config.artifact_dir) + self.weave_trace = AGSWeaveTrace( + args, + self.adapter_service.tokenizer, + 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) -> list[Sample]: md = get_metadata(base_sample) instance_id = md["instance_id"] - if not md["image"] or not md["workdir"]: - return self._abort_result(base_sample, "missing_image_or_workdir", 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 try: self.adapter_service.adapter.open_session( session_id, @@ -90,7 +105,9 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam }, ) if not samples: - return self._abort_result(base_sample, "adapter_session_empty", instance_id) + 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( { @@ -130,13 +147,35 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam 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) - return self._abort_result(base_sample, "wall_clock_timeout", instance_id) + 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()) - return self._abort_result(base_sample, f"exception:{type(exc).__name__}", instance_id) + 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: 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..6460f06e50 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/weave_trace.py @@ -0,0 +1,247 @@ +"""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, *, enable_token2text: bool = False) -> None: + self.project = _wandb_project(args) + self.tokenizer = tokenizer + 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): + 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): + 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) + elif record_type == "user": + yield from _tool_result_events(record) + elif record_type == "result": + yield { + "kind": "result", + "display_name": "agent result", + "inputs": _common_fields(record), + "output": record, + "started_at": _parse_timestamp(record.get("timestamp")), + } + + +def _assistant_events(record: dict[str, Any]): + message = record.get("message") or {} + common = _common_fields(record) + 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": _parse_timestamp(record.get("timestamp")), + } + 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": _parse_timestamp(record.get("timestamp")), + } + + +def _tool_result_events(record: dict[str, Any]): + message = record.get("message") or {} + common = _common_fields(record) + 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": _parse_timestamp(record.get("timestamp")), + } + + +def _common_fields(record: dict[str, Any]) -> 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"), + } + + +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 ea0a56ded3..b1a74468c4 100644 --- a/slime_plugins/rollout_buffer/rollout_buffer_example.py +++ b/slime_plugins/rollout_buffer/rollout_buffer_example.py @@ -81,7 +81,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") @@ -327,6 +327,13 @@ def start_rollout(api_base_url: str, args, metadata): "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) diff --git a/tests/test_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py index 238598c9c4..45fc7c587b 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -4,6 +4,9 @@ import base64 import json import re +import sys +import types +from types import SimpleNamespace from tests.test_agent._fakes import FakeSandbox @@ -19,6 +22,8 @@ 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): @@ -70,6 +75,225 @@ def test_sampling_params_use_sglang_generate_names(): } +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_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_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["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 _ctx(workdir="/workspace/repo", sid="sess-1", url="http://host:18001"): from slime.agent.harness.common import HarnessContext From c360ac8675c0c90321b345ad2e66047ab18c91ea Mon Sep 17 00:00:00 2001 From: FunJim Date: Tue, 14 Jul 2026 11:37:24 +0800 Subject: [PATCH 09/43] Use stream JSON for CodeBuddy Code harness --- .../rollout_buffer/generator/ags_generator/harnesses.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py index dd060848f5..333b49128a 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py @@ -202,7 +202,9 @@ async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: 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)}", - "--output-format json", + "--verbose", + "--output-format stream-json", + "--include-partial-messages", ] tools = os.environ.get(self.tools_env, ",".join(self.allowed_tools)).strip() if tools: From e3df50ed86ada98afe2f83332f97436daea1b89a Mon Sep 17 00:00:00 2001 From: FunJim Date: Tue, 14 Jul 2026 12:50:25 +0800 Subject: [PATCH 10/43] Parse AGS Weave traces by agent Route trajectory parsing through the configured AGS agent so CodeBuddy Code stream JSON can be traced without assuming Claude Code timestamps. --- .../generator/ags_generator/rollout.py | 1 + .../generator/ags_generator/weave_trace.py | 73 ++++++++--- .../test_rollout_buffer/test_ags_generator.py | 124 +++++++++++++++++- 3 files changed, 181 insertions(+), 17 deletions(-) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py index 3dad3b0aca..bc095f2e9a 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py @@ -36,6 +36,7 @@ def __init__(self, args: Namespace, config: AGSGeneratorConfig | None = None) -> 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) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/weave_trace.py b/slime_plugins/rollout_buffer/generator/ags_generator/weave_trace.py index 6460f06e50..d43c27148c 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/weave_trace.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/weave_trace.py @@ -16,9 +16,17 @@ class AGSWeaveTrace: - def __init__(self, args: Namespace, tokenizer, *, enable_token2text: bool = False) -> None: + 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) @@ -98,7 +106,7 @@ def finish_rollout( def _log_trajectory(self, parent, trajectory_path: str) -> None: tool_calls = {} - for event in iter_trajectory_events(trajectory_path): + 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( @@ -143,7 +151,36 @@ def _sample_payload(self, sample: Sample) -> dict[str, Any]: return payload -def iter_trajectory_events(path: str | Path): +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: @@ -158,22 +195,22 @@ def iter_trajectory_events(path: str | Path): continue record_type = record.get("type") if record_type == "assistant": - yield from _assistant_events(record) + yield from _assistant_events(record, timestamp_key=timestamp_key) elif record_type == "user": - yield from _tool_result_events(record) + 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), + "inputs": _common_fields(record, timestamp_key=timestamp_key), "output": record, - "started_at": _parse_timestamp(record.get("timestamp")), + "started_at": _record_timestamp(record, timestamp_key), } -def _assistant_events(record: dict[str, Any]): +def _assistant_events(record: dict[str, Any], *, timestamp_key: str): message = record.get("message") or {} - common = _common_fields(record) + 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": @@ -183,7 +220,7 @@ def _assistant_events(record: dict[str, Any]): "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": _parse_timestamp(record.get("timestamp")), + "started_at": _record_timestamp(record, timestamp_key), } elif block_type in {"text", "thinking"}: text = block.get(block_type) @@ -193,13 +230,13 @@ def _assistant_events(record: dict[str, Any]): "display_name": f"assistant {block_type}", "inputs": common, "output": {"text": text, "usage": message.get("usage")}, - "started_at": _parse_timestamp(record.get("timestamp")), + "started_at": _record_timestamp(record, timestamp_key), } -def _tool_result_events(record: dict[str, Any]): +def _tool_result_events(record: dict[str, Any], *, timestamp_key: str): message = record.get("message") or {} - common = _common_fields(record) + common = _common_fields(record, timestamp_key=timestamp_key) for block in message.get("content") or []: if block.get("type") == "tool_result": yield { @@ -212,19 +249,23 @@ def _tool_result_events(record: dict[str, Any]): "is_error": block.get("is_error", False), "tool_use_result": record.get("tool_use_result"), }, - "started_at": _parse_timestamp(record.get("timestamp")), + "started_at": _record_timestamp(record, timestamp_key), } -def _common_fields(record: dict[str, Any]) -> dict[str, Any]: +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"), + "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 diff --git a/tests/test_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py index 45fc7c587b..45bc00ef09 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -108,6 +108,12 @@ def _trace(enable_token2text=False): 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, @@ -227,6 +233,122 @@ def test_weave_trace_pairs_tool_call_and_result(tmp_path): 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() @@ -362,7 +484,7 @@ async def run_case(): 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 --output-format json" in body + assert "cbc --model slime-actor --verbose --output-format stream-json --include-partial-messages" in body assert "--max-turns 7" in body assert "-y" in body and "solve it" in body assert "--effort none" in body From dccbcb213291a684cf488339f6f3c8205bd18e22 Mon Sep 17 00:00:00 2001 From: FunJim Date: Wed, 15 Jul 2026 15:35:52 +0800 Subject: [PATCH 11/43] Encode Harbor eval payloads with gzip+base64 Large SWE-bench PASS_TO_PASS lists made the inline test.sh/config.json heredocs blow past AGS/E2B argv limits. Embed the files as gzip+base64 blobs and decode them with a python3 heredoc inside the sandbox instead. --- tools/harbor_task_to_slime_prompt_data.py | 36 ++++++++++++++++++++--- 1 file changed, 32 insertions(+), 4 deletions(-) diff --git a/tools/harbor_task_to_slime_prompt_data.py b/tools/harbor_task_to_slime_prompt_data.py index 135bede2ee..ed5665cc4f 100644 --- a/tools/harbor_task_to_slime_prompt_data.py +++ b/tools/harbor_task_to_slime_prompt_data.py @@ -37,7 +37,9 @@ from __future__ import annotations import argparse +import base64 import fnmatch +import gzip import json import os import re @@ -370,7 +372,8 @@ def build_eval_cmd(task_dir: Path, swe_config: dict[str, Any]) -> str: 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. + 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(): @@ -379,7 +382,23 @@ def build_eval_cmd(task_dir: Path, swe_config: dict[str, Any]) -> str: 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, indent=2) + 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", @@ -387,8 +406,7 @@ def build_eval_cmd(task_dir: Path, swe_config: dict[str, Any]) -> str: # them here so they exist only in the eval sandbox, not in the # agent sandbox. "mkdir -p /tests /logs/verifier", - _heredoc(config_path, config_json, f"SLIME_AGS_CONFIG_{task_slug}"), - _heredoc(script_path, test_script, f"SLIME_AGS_TEST_{task_slug}"), + _python_heredoc(decoder, decoder_delim), f"chmod +x {shlex.quote(script_path)}", f"bash {shlex.quote(script_path)}", ] @@ -594,6 +612,16 @@ def _heredoc(path: str, content: str, delimiter: str) -> str: 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"))).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} From a1d8b67e08718eaf9d23c5dd4205bbf30741d392 Mon Sep 17 00:00:00 2001 From: FunJim Date: Wed, 15 Jul 2026 19:15:28 +0800 Subject: [PATCH 12/43] Add AGS wandb metrics hooks --- .../generator/ags_generator/wandb_metrics.py | 369 ++++++++++++++++++ 1 file changed, 369 insertions(+) create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/wandb_metrics.py 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..797ea21811 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/ags_generator/wandb_metrics.py @@ -0,0 +1,369 @@ +"""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 |= _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 _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" From 1b8d561d5e373c212d39f5f11935a50b8a6b3c57 Mon Sep 17 00:00:00 2001 From: FunJim Date: Thu, 16 Jul 2026 13:26:36 +0800 Subject: [PATCH 13/43] Handle fan-out passrate metrics Compute passrate from prompt groups and rollout attempts when rollout fan-out produces multiple training samples per attempt, and avoid failing training on mismatched fixed-shape passrate inputs. --- slime/backends/megatron_utils/data.py | 22 ++++-- slime/ray/rollout.py | 7 +- slime/utils/metric_utils.py | 71 +++++++++++++++++++- tests/test_passrate_metrics.py | 97 +++++++++++++++++++++++++++ 4 files changed, 189 insertions(+), 8 deletions(-) create mode 100644 tests/test_passrate_metrics.py 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/ray/rollout.py b/slime/ray/rollout.py index be8f42438c..f74dfe7fdc 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -721,6 +721,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 +865,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/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/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) From 7ef7e1f18563962229a1030eda26f87241f2f466 Mon Sep 17 00:00:00 2001 From: FunJim Date: Thu, 16 Jul 2026 16:37:17 +0800 Subject: [PATCH 14/43] Support AGS periodic eval --- .../rollout_buffer/generator/ags_generator.py | 1 + .../generator/ags_generator/__init__.py | 4 +- .../generator/ags_generator/entry.py | 68 +++++++++++++++++++ tests/__init__.py | 6 ++ .../test_rollout_buffer/test_ags_generator.py | 29 ++++++++ 5 files changed, 106 insertions(+), 2 deletions(-) create mode 100644 tests/__init__.py diff --git a/slime_plugins/rollout_buffer/generator/ags_generator.py b/slime_plugins/rollout_buffer/generator/ags_generator.py index 2879d328fb..d620a4e5a2 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator.py @@ -2,6 +2,7 @@ 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, diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/__init__.py b/slime_plugins/rollout_buffer/generator/ags_generator/__init__.py index da4f371e11..eacbbd947f 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/__init__.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/__init__.py @@ -1,5 +1,5 @@ """AGS coding-agent rollout-buffer generator package.""" -from .entry import TASK_TYPE, run_rollout +from .entry import TASK_TYPE, generate, run_rollout -__all__ = ["TASK_TYPE", "run_rollout"] +__all__ = ["TASK_TYPE", "run_rollout", "generate"] diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py index 39790cc113..50b265e0f9 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -3,12 +3,14 @@ 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 @@ -21,6 +23,50 @@ 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) -> None: + self.config = AGSGeneratorConfig.from_env( + enable_token2text=_as_bool(getattr(args, "enable_token2text", False)) + ) + self.runner = AGSRolloutRunner(args, self.config) + 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. + """ + + state = _AGSGenerateState(args) + async with state.semaphore: + samples = await state.runner.generate(base_sample, sampling_params) + + 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.""" @@ -106,6 +152,28 @@ async def _guarded(sample: Sample) -> None: 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 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_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py index 45bc00ef09..9eb906899f 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -12,6 +12,7 @@ from slime.utils.types import Sample from slime_plugins.rollout_buffer.generator.ags_generator.entry import ( + _collapse_eval_samples, get_group_data_meta_info, is_valid_group, transform_group, @@ -68,6 +69,34 @@ def test_group_hooks_accept_complete_sample_payloads(): 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, From 361e74242050385a60783b5ae15febc705463a6e Mon Sep 17 00:00:00 2001 From: FunJim Date: Fri, 17 Jul 2026 15:50:30 +0800 Subject: [PATCH 15/43] Support proxied AGS adapter smoke runs Allow AGS adapters to publish a full public base URL so proxy paths can be used, and skip unnecessary rollout onload/update work after the final rollout when no eval remains. --- .../generator/ags_generator/adapter_service.py | 10 +++++++--- train.py | 16 ++++++++++------ 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py b/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py index d63fa099c5..3e4e773782 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py @@ -26,8 +26,12 @@ def __init__(self, args: Namespace, config: AGSGeneratorConfig, adapter_cls: typ or os.environ.get("AGS_GENERATOR_SGLANG_URL") or f"http://{args.sglang_router_ip}:{args.sglang_router_port}" ) - if not config.adapter_public_host: - raise RuntimeError("ADAPTER_PUBLIC_HOST is not set; AGS sandboxes need it to reach the adapter") + 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, @@ -43,7 +47,7 @@ def __init__(self, args: Namespace, config: AGSGeneratorConfig, adapter_cls: typ thread_name="ags-rollout-adapter", runner_kwargs={"handler_cancellation": True, "access_log_class": FilteredAccessLogger}, ) - self.adapter_url = f"http://{config.adapter_public_host}:{self.app_handle.port}" + self.adapter_url = public_base_url or f"http://{config.adapter_public_host}:{self.app_handle.port}" logger.info( "[ags_generator] tokenizer=%s adapter=%s sglang_url=%s max_context_len=%s tool_parser=%s reasoning_parser=%s", args.hf_checkpoint, 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()) From 349e5e3bcaf67e1d8d6716a2eae952012e0df7ce Mon Sep 17 00:00:00 2001 From: FunJim Date: Fri, 17 Jul 2026 16:07:11 +0800 Subject: [PATCH 16/43] Drain AGS rollout jobs before offload Stop each rollout-buffer generation job once enough samples have been collected so colocated SGLang engines are not offloaded while AGS sandboxes can still issue generation requests. Also pause generation before releasing SGLang memory and add tests for one-epoch buffer requests and stale job writes. --- slime/backends/sglang_utils/sglang_engine.py | 16 +- slime/ray/rollout.py | 8 +- slime/utils/arguments.py | 19 ++ slime_plugins/rollout_buffer/buffer.py | 140 +++++++++++- .../generator/ags_generator/entry.py | 105 ++++++--- .../rollout_buffer/rollout_buffer_example.py | 202 ++++++++++-------- .../test_rollout_buffer/test_ags_generator.py | 58 +++++ 7 files changed, 423 insertions(+), 125 deletions(-) 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 f74dfe7fdc..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(): diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 4c56cffcad..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, diff --git a/slime_plugins/rollout_buffer/buffer.py b/slime_plugins/rollout_buffer/buffer.py index db4e561124..d052661321 100644 --- a/slime_plugins/rollout_buffer/buffer.py +++ b/slime_plugins/rollout_buffer/buffer.py @@ -5,10 +5,12 @@ 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 +223,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 +237,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 +264,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 +379,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,18 +403,58 @@ 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__": diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py index 50b265e0f9..833e5443eb 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -78,9 +78,11 @@ def run_rollout(data: dict[str, Any]) -> str: remote_buffer_url = data["remote_buffer_url"].rstrip("/") + "/buffer/write" num_epoch = int(data.get("num_epoch", 1)) groups_per_epoch = int( - data.get("rollout_batch_size") or data.get("num_groups_per_epoch") or args.rollout_batch_size + 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", @@ -106,49 +108,88 @@ async def _run_sample(epoch: int, sample: Sample) -> None: **(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: - semaphore = asyncio.Semaphore(config.rollout_concurrency) - async def _guarded(sample: Sample) -> None: - async with semaphore: + 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 _run_sample(epoch, sample) - except Exception as exc: - instance_id = _instance_id(sample) + await asyncio.to_thread(_send_data_to_buffer, remote_buffer_url, item) + except Exception as send_exc: logger.exception( - "[ags_generator] %s: sample task failed; writing aborted rollout: %s", + "[ags_generator] %s: failed to write aborted rollout after task failure: %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 {}), - }, + send_exc, ) - 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, - ) - await asyncio.gather(*(_guarded(sample) for sample in samples)) + 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" @@ -353,3 +394,7 @@ def _as_bool(value: Any) -> bool: 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/rollout_buffer_example.py b/slime_plugins/rollout_buffer/rollout_buffer_example.py index b1a74468c4..6f9abae5f7 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): @@ -295,13 +294,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): 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, @@ -347,21 +353,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( @@ -371,85 +385,99 @@ 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" - ) - 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: - if "samples" in record: - compact_samples = [Sample.from_dict(item) for item in record["samples"]] - group_results.append(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"]}, - ) + + 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) + 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 = [] + + 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" ) - sample_results.append(group_results) + 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: + if "samples" in record: + compact_samples = [Sample.from_dict(item) for item in record["samples"]] + group_results.append(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"]}, + ) + ) + sample_results.append(group_results) - data_buffer.add_samples(sample_results) - final_return_results = data_buffer.get_samples(args.rollout_batch_size) # type: ignore + data_buffer.add_samples(sample_results) + final_return_results = data_buffer.get_samples(args.rollout_batch_size) # type: ignore - return final_return_results + 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/test_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py index 9eb906899f..ff455af751 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -436,6 +436,8 @@ def fake_post(url, json, timeout): 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" @@ -445,6 +447,62 @@ def fake_post(url, json, timeout): 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 From fe33cd9af2bb4dcb53134606a8cf5994b1835b6a Mon Sep 17 00:00:00 2001 From: FunJim Date: Fri, 17 Jul 2026 18:03:28 +0800 Subject: [PATCH 17/43] Add two-node AGS rollout buffer training script --- .../run_qwen35_35b_a3b_swe_2nodes.sh | 336 ++++++++++++++++++ 1 file changed, 336 insertions(+) create mode 100644 examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh new file mode 100644 index 0000000000..4b2193ad0e --- /dev/null +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh @@ -0,0 +1,336 @@ +#!/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}" +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_eval/swebench_verified_from_yulei_filtered_slime.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:-claude_code_ags_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}" +export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-16}" +export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-16}" + +# # autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI compacts before any +# # segment crosses the training-side cap. `investigator` is a read-only sub-agent. +# SETTINGS_JSON='{"permissions":{"defaultMode":"bypassPermissions"},"autoCompactEnabled":true,"autoCompactWindow":80000}' +# 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="--settings '${SETTINGS_JSON}' --disable-slash-commands --agents '${AGENTS_JSON}' --disallowedTools WebFetch WebSearch" +export SLIME_AGENT_CC_MAX_TURNS="${SLIME_AGENT_CC_MAX_TURNS:-100}" +export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:---max-turns ${SLIME_AGENT_CC_MAX_TURNS}}" + +# 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" + +CKPT_ARGS=( + --hf-checkpoint "${HF_CHECKPOINT}" + --ref-load "${REF_MODEL_PATH}" +) + +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 + --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 1 + --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.6 + --eval-top-p 0.95 + --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-project "${WANDB_PROJECT:-slime-claude-code-ags}" + --wandb-group "${WANDB_GROUP:-${EXP_TAG}}" + --wandb-key "${WANDB_API_KEY}" + --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 + wait + 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}" From 699f9ae2322b013747a9489f75b67db31a44077a Mon Sep 17 00:00:00 2001 From: FunJim Date: Mon, 20 Jul 2026 13:49:22 +0800 Subject: [PATCH 18/43] Add four-node AGS rollout buffer training script Keep the two-node launcher aligned with the scalable setup by enabling checkpoint saves, shuffled chat-templated rollouts, configurable rollout steps, and run-local W&B logs. --- .../run_qwen35_35b_a3b_swe_2nodes.sh | 16 +- .../run_qwen35_35b_a3b_swe_4nodes.sh | 346 ++++++++++++++++++ 2 files changed, 359 insertions(+), 3 deletions(-) create mode 100644 examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh index 4b2193ad0e..eb3c602192 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh @@ -42,6 +42,7 @@ 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}" @@ -93,7 +94,7 @@ export SWE_AGENT="${SWE_AGENT:-claude_code}" 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\"}}" +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}}}" @@ -123,9 +124,15 @@ 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:-10}" +mkdir -p "${SAVE_DIR}" + CKPT_ARGS=( --hf-checkpoint "${HF_CHECKPOINT}" --ref-load "${REF_MODEL_PATH}" + --save "${SAVE_DIR}" + --save-interval "${SAVE_INTERVAL}" ) ROLLOUT_ARGS=( @@ -140,6 +147,8 @@ ROLLOUT_ARGS=( --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}" @@ -149,7 +158,7 @@ ROLLOUT_ARGS=( --rollout-max-response-len "${MAX_GEN_LEN}" --rollout-temperature 1.0 --rollout-stop-token-ids 248046 248044 - --num-steps-per-rollout 1 + --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 @@ -227,9 +236,10 @@ SGLANG_ARGS=( if [[ -n "${WANDB_API_KEY:-}" ]]; then WANDB_ARGS=( --use-wandb - --wandb-project "${WANDB_PROJECT:-slime-claude-code-ags}" + --wandb-project "${WANDB_PROJECT:-slime-rl-training}" --wandb-group "${WANDB_GROUP:-${EXP_TAG}}" --wandb-key "${WANDB_API_KEY}" + --wandb-dir "${LOG_DIR}/wandb" --disable-wandb-random-suffix ) else diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh new file mode 100644 index 0000000000..40dd1caecc --- /dev/null +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh @@ -0,0 +1,346 @@ +#!/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_eval/swebench_verified_from_yulei_filtered_slime.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:-claude_code_ags_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}" +export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-16}" +export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-16}" + +# # autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI compacts before any +# # segment crosses the training-side cap. `investigator` is a read-only sub-agent. +# SETTINGS_JSON='{"permissions":{"defaultMode":"bypassPermissions"},"autoCompactEnabled":true,"autoCompactWindow":80000}' +# 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="--settings '${SETTINGS_JSON}' --disable-slash-commands --agents '${AGENTS_JSON}' --disallowedTools WebFetch WebSearch" +export SLIME_AGENT_CC_MAX_TURNS="${SLIME_AGENT_CC_MAX_TURNS:-100}" +export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:---max-turns ${SLIME_AGENT_CC_MAX_TURNS}}" + +# 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:-10}" +mkdir -p "${SAVE_DIR}" + +CKPT_ARGS=( + --hf-checkpoint "${HF_CHECKPOINT}" + --ref-load "${REF_MODEL_PATH}" + --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.6 + --eval-top-p 0.95 + --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-project "${WANDB_PROJECT:-slime-rl-training}" + --wandb-group "${WANDB_GROUP:-${EXP_TAG}}" + --wandb-key "${WANDB_API_KEY}" + --wandb-dir "${LOG_DIR}/wandb" + --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 + wait + 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}" From a3f953ae6c04cf58f5bd4e9345b482937a619d64 Mon Sep 17 00:00:00 2001 From: FunJim Date: Mon, 20 Jul 2026 14:01:50 +0800 Subject: [PATCH 19/43] Reuse AGS adapter for periodic eval Add lightweight adapter control endpoints so eval rollouts can reuse the already-running AGS adapter instead of binding a second service on the same port. --- slime/agent/adapters/common.py | 31 +++++++ .../ags_generator/adapter_service.py | 86 +++++++++++++++++++ .../generator/ags_generator/entry.py | 6 +- .../generator/ags_generator/rollout.py | 15 +++- 4 files changed, 132 insertions(+), 6 deletions(-) diff --git a/slime/agent/adapters/common.py b/slime/agent/adapters/common.py index 6f01d4ed9d..117fc9177b 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) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py b/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py index 3e4e773782..da82a88037 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py @@ -2,19 +2,80 @@ 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.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)}, + ) + + class AdapterService(metaclass=SingletonMeta): def __init__(self, args: Namespace, config: AGSGeneratorConfig, adapter_cls: type) -> None: self.tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) @@ -57,3 +118,28 @@ def __init__(self, args: Namespace, config: AGSGeneratorConfig, adapter_cls: typ 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}" + logger.info( + "[ags_generator] using remote adapter control=%s public=%s max_context_len=%s", + control_url, + self.adapter_url, + self.max_context_len, + ) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py index 833e5443eb..56c8dd8fe3 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -32,11 +32,11 @@ class _AGSGenerateState(metaclass=SingletonMeta): prompt. """ - def __init__(self, args: Namespace) -> None: + def __init__(self, args: Namespace, *, evaluation: bool = False) -> None: self.config = AGSGeneratorConfig.from_env( enable_token2text=_as_bool(getattr(args, "enable_token2text", False)) ) - self.runner = AGSRolloutRunner(args, self.config) + self.runner = AGSRolloutRunner(args, self.config, use_remote_adapter=evaluation) self.semaphore = asyncio.Semaphore(self.config.rollout_concurrency) @@ -58,7 +58,7 @@ async def generate( metrics count one eval attempt per prompt. """ - state = _AGSGenerateState(args) + state = _AGSGenerateState(args, evaluation=evaluation) async with state.semaphore: samples = await state.runner.generate(base_sample, sampling_params) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py index bc095f2e9a..68261b40f7 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py @@ -14,7 +14,7 @@ from slime.utils.types import Sample -from .adapter_service import AdapterService +from .adapter_service import AdapterService, RemoteAdapterService from .ags_sandbox import AGSSandbox from .artifacts import ArtifactWriter, sample_artifact_id from .config import AGSGeneratorConfig @@ -27,11 +27,20 @@ class AGSRolloutRunner: - def __init__(self, args: Namespace, config: AGSGeneratorConfig | None = None) -> None: + 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) - self.adapter_service = AdapterService(args, self.config, self.adapter_cls) + if use_remote_adapter: + self.adapter_service = RemoteAdapterService(args, self.config) + else: + self.adapter_service = AdapterService(args, self.config, self.adapter_cls) self.artifacts = ArtifactWriter(self.config.artifact_dir) self.weave_trace = AGSWeaveTrace( args, From e9796a4a64c9d920740f3d5fcdaa51c5df81a504 Mon Sep 17 00:00:00 2001 From: FunJim Date: Mon, 20 Jul 2026 14:02:04 +0800 Subject: [PATCH 20/43] Advance AGS prompt source by rollout id Pass rollout_id through the rollout-buffer payload and initialize the AGS prompt source offset from it so standalone generator runs do not repeatedly start from the first prompt groups. --- .../rollout_buffer/generator/ags_generator/entry.py | 1 + .../rollout_buffer/generator/ags_generator/source.py | 6 ++++++ slime_plugins/rollout_buffer/rollout_buffer_example.py | 8 ++++++-- 3 files changed, 13 insertions(+), 2 deletions(-) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py index 56c8dd8fe3..1e1008566b 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -339,6 +339,7 @@ def _build_args(data: dict[str, Any]) -> Namespace: 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"]), diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/source.py b/slime_plugins/rollout_buffer/generator/ags_generator/source.py index a6096ff388..3f220c4b81 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/source.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/source.py @@ -15,6 +15,12 @@ class AGSPromptSource: 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) + if start_group > 0 and self.data_source.dataset is not None: + dataset_len = len(self.data_source.dataset) + self.data_source.sample_offset = start_group % dataset_len if dataset_len else 0 + self.data_source.sample_group_index = start_group + self.data_source.sample_index = start_group * int(args.n_samples_per_prompt) def get_groups(self, num_groups: int) -> list[list[Sample]]: groups = self.data_source.get_samples(num_groups) diff --git a/slime_plugins/rollout_buffer/rollout_buffer_example.py b/slime_plugins/rollout_buffer/rollout_buffer_example.py index 6f9abae5f7..fc44a4ac42 100644 --- a/slime_plugins/rollout_buffer/rollout_buffer_example.py +++ b/slime_plugins/rollout_buffer/rollout_buffer_example.py @@ -294,7 +294,7 @@ 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, num_groups_per_epoch: int | None = None): +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] @@ -328,6 +328,8 @@ def start_rollout(api_base_url: str, args, metadata, num_groups_per_epoch: int | "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), @@ -390,7 +392,9 @@ async def generate_rollout_async(args, rollout_id: int, data_buffer, evaluation: 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) + start_inform = start_rollout( + args.rollout_buffer_url, args, metadata, num_groups_per_epoch=groups_to_fetch, rollout_id=rollout_id + ) print(f"start rollout with payload: {start_inform}") print(f"start rollout id: {rollout_id}") rollout_started = True From 3b6dd88fea242cc0ada3dc323fbf7402e3fe5591 Mon Sep 17 00:00:00 2001 From: FunJim Date: Mon, 20 Jul 2026 14:23:43 +0800 Subject: [PATCH 21/43] chore: Update save interval and WandB project name --- examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh | 4 ++-- examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh index eb3c602192..2a619e4fa5 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh @@ -125,7 +125,7 @@ cd "${SLIME_DIR}" source "${SLIME_DIR}/scripts/models/qwen3.5-35B-A3B.sh" SAVE_DIR="${SAVE_DIR:-${EXP}/checkpoints}" -SAVE_INTERVAL="${SAVE_INTERVAL:-10}" +SAVE_INTERVAL="${SAVE_INTERVAL:-5}" mkdir -p "${SAVE_DIR}" CKPT_ARGS=( @@ -236,7 +236,7 @@ SGLANG_ARGS=( if [[ -n "${WANDB_API_KEY:-}" ]]; then WANDB_ARGS=( --use-wandb - --wandb-project "${WANDB_PROJECT:-slime-rl-training}" + --wandb-project "${WANDB_PROJECT:-slime-claude-code-ags}" --wandb-group "${WANDB_GROUP:-${EXP_TAG}}" --wandb-key "${WANDB_API_KEY}" --wandb-dir "${LOG_DIR}/wandb" diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh index 40dd1caecc..9719d48346 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh @@ -125,7 +125,7 @@ cd "${SLIME_DIR}" source "${SLIME_DIR}/scripts/models/qwen3.5-35B-A3B.sh" SAVE_DIR="${SAVE_DIR:-${EXP}/checkpoints}" -SAVE_INTERVAL="${SAVE_INTERVAL:-10}" +SAVE_INTERVAL="${SAVE_INTERVAL:-5}" mkdir -p "${SAVE_DIR}" CKPT_ARGS=( @@ -236,7 +236,7 @@ SGLANG_ARGS=( if [[ -n "${WANDB_API_KEY:-}" ]]; then WANDB_ARGS=( --use-wandb - --wandb-project "${WANDB_PROJECT:-slime-rl-training}" + --wandb-project "${WANDB_PROJECT:-slime-claude-code-ags}" --wandb-group "${WANDB_GROUP:-${EXP_TAG}}" --wandb-key "${WANDB_API_KEY}" --wandb-dir "${LOG_DIR}/wandb" From 554b1e72639ce4b83578b2f72ff1f25eb64bebfe Mon Sep 17 00:00:00 2001 From: FunJim Date: Tue, 21 Jul 2026 21:00:32 +0800 Subject: [PATCH 22/43] chore: Harden AGS rollout buffer launch scripts --- .../claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh | 11 +++++++---- .../claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh | 11 +++++++---- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh index 2a619e4fa5..ef7d90b0c1 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh @@ -104,6 +104,7 @@ 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}" export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-16}" +export SWE_BOOT_RETRIES="${SWE_BOOT_RETRIES:-6}" export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-16}" # # autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI compacts before any @@ -131,6 +132,7 @@ 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}" ) @@ -286,7 +288,10 @@ if [[ -f "${HOSTFILE}" ]]; then --node-ip-address ${WORKER_IP} --disable-usage-stats" & STARTED_WORKERS=$((STARTED_WORKERS + 1)) done - wait + 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 @@ -316,9 +321,7 @@ keys = ( env = {k: os.environ[k] for k in keys if k in os.environ} env["MASTER_ADDR"] = os.environ["MASTER_ADDR"] env["MASTER_PORT"] = os.environ.get("MASTER_PORT", "") -env["GLOO_SOCKET_IFNAME"] = os.environ["GLOO_SOCKET_IFNAME"] -env["TP_SOCKET_IFNAME"] = os.environ["GLOO_SOCKET_IFNAME"] -env["NCCL_SOCKET_IFNAME"] = os.environ["NCCL_SOCKET_IFNAME"] +# Keep per-node socket interface env inherited from each Ray node; do not override workers with the head ifname. env["PYTHONPATH"] = f"/root/Megatron-LM/:{os.environ['SLIME_DIR']}" env["CUDA_DEVICE_MAX_CONNECTIONS"] = "1" env["NCCL_NVLS_ENABLE"] = "0" diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh index 9719d48346..53cc37226e 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh @@ -104,6 +104,7 @@ 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}" export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-16}" +export SWE_BOOT_RETRIES="${SWE_BOOT_RETRIES:-6}" export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-16}" # # autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI compacts before any @@ -131,6 +132,7 @@ 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}" ) @@ -286,7 +288,10 @@ if [[ -f "${HOSTFILE}" ]]; then --node-ip-address ${WORKER_IP} --disable-usage-stats" & STARTED_WORKERS=$((STARTED_WORKERS + 1)) done - wait + 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 @@ -316,9 +321,7 @@ keys = ( env = {k: os.environ[k] for k in keys if k in os.environ} env["MASTER_ADDR"] = os.environ["MASTER_ADDR"] env["MASTER_PORT"] = os.environ.get("MASTER_PORT", "") -env["GLOO_SOCKET_IFNAME"] = os.environ["GLOO_SOCKET_IFNAME"] -env["TP_SOCKET_IFNAME"] = os.environ["GLOO_SOCKET_IFNAME"] -env["NCCL_SOCKET_IFNAME"] = os.environ["NCCL_SOCKET_IFNAME"] +# Keep per-node socket interface env inherited from each Ray node; do not override workers with the head ifname. env["PYTHONPATH"] = f"/root/Megatron-LM/:{os.environ['SLIME_DIR']}" env["CUDA_DEVICE_MAX_CONNECTIONS"] = "1" env["NCCL_NVLS_ENABLE"] = "0" From 6ae28a1ff2aa25a10e7a89cf44dbaecba82be569 Mon Sep 17 00:00:00 2001 From: FunJim Date: Wed, 22 Jul 2026 14:30:41 +0800 Subject: [PATCH 23/43] chore: Tune AGS launch defaults and WandB ownership Raise the two- and four-node AGS rollout defaults to the validated concurrency settings, and pass the configured WandB entity explicitly when tracking is enabled. --- examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh | 7 ++++--- examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh | 7 ++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh index ef7d90b0c1..8a7d5e9ce0 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh @@ -103,9 +103,9 @@ 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}" -export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-16}" -export SWE_BOOT_RETRIES="${SWE_BOOT_RETRIES:-6}" -export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-16}" +export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-32}" +export SWE_BOOT_RETRIES="${SWE_BOOT_RETRIES:-10}" +export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-32}" # # autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI compacts before any # # segment crosses the training-side cap. `investigator` is a read-only sub-agent. @@ -238,6 +238,7 @@ SGLANG_ARGS=( 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}" diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh index 53cc37226e..668ddc9b0a 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh @@ -103,9 +103,9 @@ 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}" -export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-16}" -export SWE_BOOT_RETRIES="${SWE_BOOT_RETRIES:-6}" -export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-16}" +export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-32}" +export SWE_BOOT_RETRIES="${SWE_BOOT_RETRIES:-10}" +export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-32}" # # autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI compacts before any # # segment crosses the training-side cap. `investigator` is a read-only sub-agent. @@ -238,6 +238,7 @@ SGLANG_ARGS=( 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}" From 2d956becd716ef0f6a562ad03f38a8c26b21212c Mon Sep 17 00:00:00 2001 From: FunJim Date: Thu, 23 Jul 2026 13:41:33 +0800 Subject: [PATCH 24/43] fix: Handle partial AGS rollout groups Keep incomplete prompt groups out of the training buffer while preserving flattened multi-segment agent outputs. Include the staged regression coverage for AGS buffering and Harbor prompt conversion. --- slime/rollout/data_source.py | 8 +- .../rollout_buffer/rollout_buffer_example.py | 59 +++++++- .../test_harbor_task_to_slime_prompt_data.py | 136 ++++++++++++++++++ .../test_rollout_buffer_example.py | 34 +++++ 4 files changed, 229 insertions(+), 8 deletions(-) create mode 100644 tests/test_harbor_task_to_slime_prompt_data.py create mode 100644 tests/test_rollout_buffer/test_rollout_buffer_example.py 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_plugins/rollout_buffer/rollout_buffer_example.py b/slime_plugins/rollout_buffer/rollout_buffer_example.py index fc44a4ac42..02a6721b56 100644 --- a/slime_plugins/rollout_buffer/rollout_buffer_example.py +++ b/slime_plugins/rollout_buffer/rollout_buffer_example.py @@ -88,6 +88,25 @@ def get_group_timestamp(group_items): return selected_results +def _select_complete_rollout_groups(args, results, need_length): + """Keep only prompt groups with every requested rollout attempt present.""" + + grouped = {} + for record in results: + grouped.setdefault(record["instance_id"], []).append(record) + + 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) + + return select_rollout_data(args, complete_records, need_length), incomplete_groups + + def log_raw_info(args, all_meta_info, rollout_id): if not all_meta_info: return @@ -405,6 +424,8 @@ async def generate_rollout_async(args, rollout_id: int, data_buffer, evaluation: retry_times = 0 results = [] all_meta_info = [] + selected_results = [] + incomplete_groups = {} if args.fetch_trajectory_retry_times == -1: print( @@ -412,13 +433,19 @@ async def generate_rollout_async(args, rollout_id: int, data_buffer, evaluation: ) while args.fetch_trajectory_retry_times == -1 or retry_times < args.fetch_trajectory_retry_times: try: - while len(results) < data_number_to_fetch: + 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) - print(f"get rollout data with length: {len(results)}") + 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}") @@ -426,8 +453,19 @@ async def generate_rollout_async(args, rollout_id: int, data_buffer, evaluation: 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(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}" + ) + + results = selected_results if len(all_meta_info) > 0 and "finished_groups" in all_meta_info[0]: finished_groups_instance_id_list = [] @@ -444,7 +482,7 @@ async def generate_rollout_async(args, rollout_id: int, data_buffer, evaluation: for record in group_record: if "samples" in record: compact_samples = [Sample.from_dict(item) for item in record["samples"]] - group_results.append(compact_samples) + group_results.extend(compact_samples) continue oai_messages = record["messages"] @@ -472,8 +510,19 @@ async def generate_rollout_async(args, rollout_id: int, data_buffer, evaluation: 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) + 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 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..34a97a2001 --- /dev/null +++ b/tests/test_harbor_task_to_slime_prompt_data.py @@ -0,0 +1,136 @@ +"""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", + prompt_source="problem_statement", + image_override=None, + default_workdir="/testbed", + include_pre_commands=False, + 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_converter_can_explicitly_reset_to_base_commit(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", + prompt_source="problem_statement", + image_override=None, + default_workdir="/testbed", + include_pre_commands=True, + include_eval_cmd=False, + include_inline_files=False, + inline_files=(), + provenance_root=False, + ) + + assert row["metadata"]["pre_commands"] == [ + "git checkout deadbeef -f", + "git clean -fd", + ] + + +def test_reset_to_base_commit_is_opt_in(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), + ], + ) + assert converter_module.parse_args().reset_to_base_commit is False + + monkeypatch.setattr( + sys, + "argv", + [ + "harbor_task_to_slime_prompt_data.py", + "--input", + str(input_path), + "--output", + str(output_path), + "--reset-to-base-commit", + ], + ) + assert converter_module.parse_args().reset_to_base_commit is True 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] From b084ddce5d5311c0abea863de12c4289e6c5bb11 Mon Sep 17 00:00:00 2001 From: FunJim Date: Thu, 23 Jul 2026 13:44:39 +0800 Subject: [PATCH 25/43] fix: Preserve Harbor image workspace by default Avoid resetting converted tasks to base_commit unless explicitly requested, so task-image compatibility patches remain available to AGS rollouts. --- tools/harbor_task_to_slime_prompt_data.py | 31 ++++++++++++++++------- 1 file changed, 22 insertions(+), 9 deletions(-) diff --git a/tools/harbor_task_to_slime_prompt_data.py b/tools/harbor_task_to_slime_prompt_data.py index ed5665cc4f..112a07b450 100644 --- a/tools/harbor_task_to_slime_prompt_data.py +++ b/tools/harbor_task_to_slime_prompt_data.py @@ -9,7 +9,6 @@ - metadata.image - metadata.workdir - metadata.problem_statement - - metadata.pre_commands - metadata.eval_cmd The generated rows are still ordinary slime JSONL prompt data: use --input-key @@ -18,11 +17,12 @@ by ags_generator without a harbor_task_path. Important: Harbor verifier assets are created inside metadata.eval_cmd, not -metadata.pre_commands. pre_commands run in both the agent sandbox and the eval -sandbox, so putting tests/config.json there would leak hidden grading data to -the agent. eval_cmd runs only in the eval sandbox. To preserve Harbor's -expected verifier layout, eval_cmd materializes /tests/config.json and -/logs/verifier before invoking the patched Harbor tests/test.sh. +metadata.pre_commands. eval_cmd runs only in the clean evaluator sandbox, so +it keeps hidden grading data out of the agent sandbox. 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``. Use --reset-to-base-commit only for an +explicit legacy/debug workflow that intentionally discards those commits. Example: python tools/harbor_task_to_slime_prompt_data.py \ @@ -129,9 +129,22 @@ def parse_args() -> argparse.Namespace: help="Override image for all rows. By default it is extracted from the active Dockerfile FROM line.", ) parser.add_argument( - "--no-pre-commands", + "--reset-to-base-commit", action="store_true", - help="Do not write metadata.pre_commands to reset the repo to base_commit.", + help=( + "Write metadata.pre_commands that reset the workspace to " + "tests/config.json.base_commit. Disabled by default because it " + "discards task-image environment compatibility commits." + ), + ) + parser.add_argument( + "--no-pre-commands", + action="store_false", + dest="reset_to_base_commit", + help=( + "Deprecated compatibility alias. Harbor conversion now preserves " + "the task image HEAD by default, so this is normally a no-op." + ), ) parser.add_argument( "--no-eval-cmd", @@ -650,7 +663,7 @@ def main() -> None: prompt_source=args.prompt_source, image_override=args.image, default_workdir=args.default_workdir, - include_pre_commands=not args.no_pre_commands, + include_pre_commands=args.reset_to_base_commit, include_eval_cmd=not args.no_eval_cmd, include_inline_files=args.include_inline_files, inline_files=inline_files, From 9e40455c924d4c945e74ace101761f0d6dbea216 Mon Sep 17 00:00:00 2001 From: FunJim Date: Thu, 23 Jul 2026 16:13:24 +0800 Subject: [PATCH 26/43] fix: Align AGS evaluation with Harbor semantics Support same-sandbox grading by default, stop timed-out agent process groups before evaluation, and make Harbor prompt conversion preserve image state with reproducible payloads. --- .../run_qwen35_35b_a3b_swe_2nodes.sh | 5 +- .../run_qwen35_35b_a3b_swe_4nodes.sh | 5 +- slime/agent/sandbox.py | 57 ++++++++++- .../generator/ags_generator/config.py | 9 ++ .../generator/ags_generator/rollout.py | 38 +++++--- .../generator/ags_generator/runner.py | 19 ++-- .../generator/ags_generator/swe_task.py | 74 +++++++++++--- tests/test_agent/test_harness.py | 3 + .../test_harbor_task_to_slime_prompt_data.py | 54 +++-------- .../test_rollout_buffer/test_ags_generator.py | 97 +++++++++++++++++++ tools/harbor_task_to_slime_prompt_data.py | 71 +++++--------- 11 files changed, 308 insertions(+), 124 deletions(-) diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh index 8a7d5e9ce0..5df9051589 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh @@ -103,6 +103,8 @@ 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}" @@ -314,7 +316,8 @@ keys = ( "AGS_IMAGE_REGISTRY_TYPE", "AGS_SANDBOX_RESOURCES_JSON", "EXPERIMENT_NAME", "TRAJECTORY_DUMP_DIR", "ADAPTER_PUBLIC_HOST", "ADAPTER_BIND_HOST", "ADAPTER_PORT", - "SWE_AGENT_TIME_BUDGET_SEC", "SWE_EVAL_TIMEOUT_SEC", "SWE_BOOT_CONCURRENCY", + "SWE_AGENT_TIME_BUDGET_SEC", "SWE_EVAL_TIMEOUT_SEC", "SWE_EVAL_ISOLATED_SANDBOX", + "SWE_BOOT_CONCURRENCY", "SWE_BOOT_RETRIES", "SWE_ROLLOUT_GUARD_SEC", "SWE_ROLLOUT_CONCURRENCY", "SLIME_AGENT_CC_MAX_TURNS", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", "CLAUDE_CODE_MAX_OUTPUT_TOKENS", "SWE_CC_PROMPT", diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh index 668ddc9b0a..6cc099a472 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh @@ -103,6 +103,8 @@ 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}" @@ -314,7 +316,8 @@ keys = ( "AGS_IMAGE_REGISTRY_TYPE", "AGS_SANDBOX_RESOURCES_JSON", "EXPERIMENT_NAME", "TRAJECTORY_DUMP_DIR", "ADAPTER_PUBLIC_HOST", "ADAPTER_BIND_HOST", "ADAPTER_PORT", - "SWE_AGENT_TIME_BUDGET_SEC", "SWE_EVAL_TIMEOUT_SEC", "SWE_BOOT_CONCURRENCY", + "SWE_AGENT_TIME_BUDGET_SEC", "SWE_EVAL_TIMEOUT_SEC", "SWE_EVAL_ISOLATED_SANDBOX", + "SWE_BOOT_CONCURRENCY", "SWE_BOOT_RETRIES", "SWE_ROLLOUT_GUARD_SEC", "SWE_ROLLOUT_CONCURRENCY", "SLIME_AGENT_CC_MAX_TURNS", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", "CLAUDE_CODE_MAX_OUTPUT_TOKENS", "SWE_CC_PROMPT", 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_plugins/rollout_buffer/generator/ags_generator/config.py b/slime_plugins/rollout_buffer/generator/ags_generator/config.py index a6d2b271db..c952dda8e3 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/config.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/config.py @@ -16,6 +16,7 @@ class AGSGeneratorConfig: 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 @@ -40,6 +41,7 @@ def from_env(cls, *, enable_token2text: bool = False) -> AGSGeneratorConfig: 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), @@ -51,3 +53,10 @@ def from_env(cls, *, enable_token2text: bool = False) -> AGSGeneratorConfig: "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.", ), ) + + +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/rollout.py b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py index 68261b40f7..299503b30f 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py @@ -72,6 +72,16 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam 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, @@ -93,18 +103,19 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam trajectory_path = await self.artifacts.dump_trajectory(sb, md["workdir"], artifact_id) diff_text = await git_diff(sb, md["workdir"]) 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, + ) - reward, applied_cleanly = await evaluate( - image=md["image"], - workdir=md["workdir"], - diff_text=diff_text, - 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, - ) + if self.config.eval_isolated_sandbox: + reward, applied_cleanly = await evaluate( + sandbox=None, + diff_text=diff_text, + **evaluation_args, + ) samples = await self.adapter_service.adapter.finish_session( session_id, base_sample=base_sample, @@ -126,6 +137,7 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam "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), @@ -141,6 +153,7 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam "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, @@ -149,10 +162,11 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam "ags_rollout_concurrency": self.config.rollout_concurrency, } logger.info( - "[ags_generator] %s: reward=%.2f applied=%s exit=%s elapsed=%.1fs segments=%d", + "[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), diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/runner.py b/slime_plugins/rollout_buffer/generator/ags_generator/runner.py index 668742fbb2..594cc5d26a 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/runner.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/runner.py @@ -3,14 +3,13 @@ from __future__ import annotations import asyncio +import logging import shlex import time -try: - from slime.agent.sandbox import EXIT_TIME_BUDGET_EXCEEDED -except ImportError: # older slime checkout keeps this constant in harness.common - from slime.agent.harness.common import EXIT_TIME_BUDGET_EXCEEDED -from slime.agent.sandbox import Sandbox +from slime.agent.sandbox import EXIT_TIME_BUDGET_EXCEEDED, Sandbox, terminate_process_group + +logger = logging.getLogger(__name__) async def run_root_command( @@ -27,6 +26,8 @@ async def run_root_command( 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" @@ -40,7 +41,10 @@ async def run_root_command( export_lines = " ".join(f"{k}={shlex.quote(str(v))}" for k, v in env.items()) await sb.exec( - f"env {export_lines} setsid {launcher} < /dev/null > /dev/null 2>&1 &", + 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, @@ -54,4 +58,7 @@ async def run_root_command( 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/swe_task.py b/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py index 482edc7dde..fd9e0697a5 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py @@ -85,6 +85,7 @@ async def git_diff(sb: Sandbox, workdir: str) -> str: async def evaluate( *, + sandbox: Sandbox | None = None, image: str, workdir: str, diff_text: str, @@ -99,27 +100,72 @@ async def evaluate( 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 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 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 + 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: @@ -174,9 +220,9 @@ async def _run_eval_cmd(ev: Sandbox, workdir: str, cmd: str, timeout: int) -> fl # /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. This happens only - # in the separate evaluator sandbox after the agent patch is collected, so - # hidden grading assets are not exposed to the agent sandbox. + # 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 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_harbor_task_to_slime_prompt_data.py b/tests/test_harbor_task_to_slime_prompt_data.py index 34a97a2001..c89c528f5a 100644 --- a/tests/test_harbor_task_to_slime_prompt_data.py +++ b/tests/test_harbor_task_to_slime_prompt_data.py @@ -67,7 +67,6 @@ def test_converter_preserves_image_head_by_default(converter_module, harbor_task prompt_source="problem_statement", image_override=None, default_workdir="/testbed", - include_pre_commands=False, include_eval_cmd=True, include_inline_files=False, inline_files=(), @@ -79,48 +78,10 @@ def test_converter_preserves_image_head_by_default(converter_module, harbor_task assert _embedded_test_script(metadata["eval_cmd"]) == (harbor_task / "tests" / "test.sh").read_text() -def test_converter_can_explicitly_reset_to_base_commit(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", - prompt_source="problem_statement", - image_override=None, - default_workdir="/testbed", - include_pre_commands=True, - include_eval_cmd=False, - include_inline_files=False, - inline_files=(), - provenance_root=False, - ) - - assert row["metadata"]["pre_commands"] == [ - "git checkout deadbeef -f", - "git clean -fd", - ] - - -def test_reset_to_base_commit_is_opt_in(converter_module, monkeypatch: pytest.MonkeyPatch, tmp_path: Path): +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), - ], - ) - assert converter_module.parse_args().reset_to_base_commit is False - monkeypatch.setattr( sys, "argv", @@ -133,4 +94,15 @@ def test_reset_to_base_commit_is_opt_in(converter_module, monkeypatch: pytest.Mo "--reset-to-base-commit", ], ) - assert converter_module.parse_args().reset_to_base_commit is True + + 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_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py index ff455af751..f3caca284d 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -8,9 +8,12 @@ import types from types import SimpleNamespace +import pytest from tests.test_agent._fakes import FakeSandbox from slime.utils.types import Sample +from slime_plugins.rollout_buffer.generator.ags_generator import 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, @@ -18,6 +21,7 @@ transform_group, ) from slime_plugins.rollout_buffer.generator.ags_generator.harnesses import CodeBuddyCodeHarness, resolve_agent +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, @@ -104,6 +108,77 @@ def test_sampling_params_use_sglang_generate_names(): } +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 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 = [] @@ -584,5 +659,27 @@ async def run_case(): 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_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()) diff --git a/tools/harbor_task_to_slime_prompt_data.py b/tools/harbor_task_to_slime_prompt_data.py index 112a07b450..42e1cdc475 100644 --- a/tools/harbor_task_to_slime_prompt_data.py +++ b/tools/harbor_task_to_slime_prompt_data.py @@ -17,12 +17,14 @@ by ags_generator without a harbor_task_path. Important: Harbor verifier assets are created inside metadata.eval_cmd, not -metadata.pre_commands. eval_cmd runs only in the clean evaluator sandbox, so -it keeps hidden grading data out of the agent sandbox. 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``. Use --reset-to-base-commit only for an -explicit legacy/debug workflow that intentionally discards those commits. +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 \ @@ -128,24 +130,6 @@ def parse_args() -> argparse.Namespace: default=None, help="Override image for all rows. By default it is extracted from the active Dockerfile FROM line.", ) - parser.add_argument( - "--reset-to-base-commit", - action="store_true", - help=( - "Write metadata.pre_commands that reset the workspace to " - "tests/config.json.base_commit. Disabled by default because it " - "discards task-image environment compatibility commits." - ), - ) - parser.add_argument( - "--no-pre-commands", - action="store_false", - dest="reset_to_base_commit", - help=( - "Deprecated compatibility alias. Harbor conversion now preserves " - "the task image HEAD by default, so this is normally a no-op." - ), - ) parser.add_argument( "--no-eval-cmd", action="store_true", @@ -242,7 +226,6 @@ def task_to_row( prompt_source: str, image_override: str | None, default_workdir: str, - include_pre_commands: bool, include_eval_cmd: bool, include_inline_files: bool, inline_files: tuple[str, ...], @@ -262,8 +245,6 @@ def task_to_row( if not image: raise ValueError(f"Cannot extract Docker image from {task_dir / 'environment' / 'Dockerfile'}") workdir = extract_dockerfile_workdir(dockerfile) or default_workdir - base_commit = swe_config.get("base_commit") - row: dict[str, Any] = {input_key: prompt} if prompt_alias_key and prompt_alias_key != input_key: row[prompt_alias_key] = prompt @@ -278,11 +259,6 @@ def task_to_row( "problem_statement": problem_statement, "harbor": harbor_metadata(task_dir, source_name, task_toml, swe_config, image, workdir), } - if include_pre_commands and base_commit: - metadata["pre_commands"] = [ - f"git checkout {shlex.quote(str(base_commit))} -f", - "git clean -fd", - ] if include_eval_cmd: metadata["eval_cmd"] = build_eval_cmd(task_dir, swe_config) if include_inline_files: @@ -375,12 +351,12 @@ 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 hidden - grading assets inside eval_cmd rather than pre_commands: pre_commands are - executed in the agent sandbox before the agent runs, while eval_cmd is only - executed in the separate evaluator sandbox after the agent patch is - collected. This preserves the SWE-bench setting and avoids exposing - FAIL_TO_PASS/PASS_TO_PASS/test_patch/reference_patch to the agent. + 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 @@ -416,8 +392,8 @@ def build_eval_cmd(task_dir: Path, swe_config: dict[str, Any]) -> str: [ "set -euo pipefail", # /tests and /logs are part of Harbor's verifier contract. Create - # them here so they exist only in the eval sandbox, not in the - # agent sandbox. + # 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)}", @@ -475,8 +451,13 @@ def write_schema(output: Path, *, input_key: str, prompt_alias_key: str, label_k "image": {"type": "string", "description": "Sandbox image consumed by ags_generator."}, "workdir": {"type": "string", "description": "Repository path inside the sandbox."}, "problem_statement": {"type": "string"}, - "pre_commands": {"type": "array", "items": {"type": "string"}}, - "eval_cmd": {"type": "string", "description": "Reward command; exit 0 means reward 1."}, + "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", @@ -519,7 +500,6 @@ def convert_tasks( prompt_source: str, image_override: str | None, default_workdir: str, - include_pre_commands: bool, include_eval_cmd: bool, include_inline_files: bool, inline_files: tuple[str, ...], @@ -543,7 +523,6 @@ def convert_tasks( prompt_source=prompt_source, image_override=image_override, default_workdir=default_workdir, - include_pre_commands=include_pre_commands, include_eval_cmd=include_eval_cmd, include_inline_files=include_inline_files, inline_files=inline_files, @@ -568,7 +547,6 @@ def convert_tasks( prompt_source=prompt_source, image_override=image_override, default_workdir=default_workdir, - include_pre_commands=include_pre_commands, include_eval_cmd=include_eval_cmd, include_inline_files=include_inline_files, inline_files=inline_files, @@ -626,7 +604,7 @@ def _heredoc(path: str, content: str, delimiter: str) -> str: def _gzip_base64(content: str) -> str: - return base64.b64encode(gzip.compress(content.encode("utf-8"))).decode("ascii") + return base64.b64encode(gzip.compress(content.encode("utf-8"), mtime=0)).decode("ascii") def _python_heredoc(script: str, delimiter: str) -> str: @@ -663,7 +641,6 @@ def main() -> None: prompt_source=args.prompt_source, image_override=args.image, default_workdir=args.default_workdir, - include_pre_commands=args.reset_to_base_commit, include_eval_cmd=not args.no_eval_cmd, include_inline_files=args.include_inline_files, inline_files=inline_files, From 67012728022c55d74d7883796b23f4298e186263 Mon Sep 17 00:00:00 2001 From: FunJim Date: Thu, 23 Jul 2026 16:33:30 +0800 Subject: [PATCH 27/43] chore: Tune AGS evaluation and sandbox defaults Use the Harbor TCR eval set, align evaluation sampling, and reduce default sandbox resources for both launchers. --- examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh | 8 ++++---- examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh | 8 ++++---- .../rollout_buffer/generator/ags_generator/ags_sandbox.py | 6 +++--- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh index 5df9051589..4ce8832140 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh @@ -53,7 +53,7 @@ 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_eval/swebench_verified_from_yulei_filtered_slime.jsonl}" +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}" @@ -94,7 +94,7 @@ export SWE_AGENT="${SWE_AGENT:-claude_code}" 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"}'} +export AGS_SANDBOX_RESOURCES_JSON=${AGS_SANDBOX_RESOURCES_JSON:-'{"cpu":"2","memory":"4Gi"}'} # 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}}}" @@ -176,8 +176,8 @@ EVAL_ARGS=( --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.6 - --eval-top-p 0.95 + --eval-temperature 0.7 + --eval-top-p 0.8 --eval-top-k 20 ) diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh index 6cc099a472..47ead7707b 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh @@ -53,7 +53,7 @@ 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_eval/swebench_verified_from_yulei_filtered_slime.jsonl}" +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}" @@ -94,7 +94,7 @@ export SWE_AGENT="${SWE_AGENT:-claude_code}" 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"}'} +export AGS_SANDBOX_RESOURCES_JSON=${AGS_SANDBOX_RESOURCES_JSON:-'{"cpu":"2","memory":"4Gi"}'} # 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}}}" @@ -176,8 +176,8 @@ EVAL_ARGS=( --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.6 - --eval-top-p 0.95 + --eval-temperature 0.7 + --eval-top-p 0.8 --eval-top-k 20 ) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/ags_sandbox.py b/slime_plugins/rollout_buffer/generator/ags_generator/ags_sandbox.py index 40a8cbbf24..2f49cbffb8 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/ags_sandbox.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/ags_sandbox.py @@ -94,12 +94,12 @@ def __init__(self, image: str, *, timeout: int | None = None, rpc_retries: int | @staticmethod def _resources() -> dict[str, str]: - raw = _env("AGS_SANDBOX_RESOURCES_JSON", '{"cpu":"4","memory":"16Gi"}') + raw = _env("AGS_SANDBOX_RESOURCES_JSON", '{"cpu":"2","memory":"4Gi"}') try: parsed = json.loads(raw) - return parsed if isinstance(parsed, dict) else {"cpu": "4", "memory": "16Gi"} + return parsed if isinstance(parsed, dict) else {"cpu": "2", "memory": "4Gi"} except Exception: - return {"cpu": "4", "memory": "16Gi"} + return {"cpu": "2", "memory": "4Gi"} def _custom_config(self) -> dict[str, Any]: return { From 38200c328271abd910ee463af516d3701d880048 Mon Sep 17 00:00:00 2001 From: FunJim Date: Thu, 23 Jul 2026 21:50:50 +0800 Subject: [PATCH 28/43] chore: Increase sandbox CPU and memory defaults --- examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh | 2 +- examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh index 4ce8832140..088bcbbb25 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh @@ -94,7 +94,7 @@ export SWE_AGENT="${SWE_AGENT:-claude_code}" 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":"2","memory":"4Gi"}'} +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}}}" diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh index 47ead7707b..3e1fc31049 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh @@ -94,7 +94,7 @@ export SWE_AGENT="${SWE_AGENT:-claude_code}" 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":"2","memory":"4Gi"}'} +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}}}" From 8586593cf6fe0de1a518fd96a0b4ea234a533c11 Mon Sep 17 00:00:00 2001 From: FunJim Date: Mon, 27 Jul 2026 15:26:47 +0800 Subject: [PATCH 29/43] chore: Make uvicorn port configurable via env var --- slime_plugins/rollout_buffer/buffer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/slime_plugins/rollout_buffer/buffer.py b/slime_plugins/rollout_buffer/buffer.py index d052661321..487859852c 100644 --- a/slime_plugins/rollout_buffer/buffer.py +++ b/slime_plugins/rollout_buffer/buffer.py @@ -2,6 +2,7 @@ import glob import importlib.util import json +import os import pathlib import threading import time @@ -461,7 +462,7 @@ async def rollout_status(): 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, From 367291a5a54dd783b6d5164854362b72bf6615b3 Mon Sep 17 00:00:00 2001 From: FunJim Date: Tue, 28 Jul 2026 18:40:04 +0800 Subject: [PATCH 30/43] fix: Restore epoch when AGS prompt source seeks AGSPromptSource seeks RolloutDataSource to the absolute group index the trainer sends (rollout_id * rollout_batch_size), but only restored sample_offset. Once that index passed one pass over the dataset, the source stayed on epoch 0's permutation, so a resumed run drew a different prompt sequence than the uninterrupted one. Derive epoch_id from the group index and reshuffle to it when --rollout-shuffle is set. --- .github/workflows/pr-test.yml | 4 + .github/workflows/pr-test.yml.j2 | 1 + .../generator/ags_generator/source.py | 18 ++- .../test_ags_prompt_source.py | 152 ++++++++++++++++++ 4 files changed, 172 insertions(+), 3 deletions(-) create mode 100644 tests/test_rollout_buffer/test_ags_prompt_source.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 761496b949..bd98028295 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -744,6 +744,10 @@ 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" } ] defaults: diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 415586cb84..c19b884ed9 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -102,6 +102,7 @@ {'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}, ], }, diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/source.py b/slime_plugins/rollout_buffer/generator/ags_generator/source.py index 3f220c4b81..7579fe8615 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/source.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/source.py @@ -16,11 +16,23 @@ 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) - if start_group > 0 and self.data_source.dataset is not None: - dataset_len = len(self.data_source.dataset) - self.data_source.sample_offset = start_group % dataset_len if dataset_len else 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) 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..73bbb0431c --- /dev/null +++ b/tests/test_rollout_buffer/test_ags_prompt_source.py @@ -0,0 +1,152 @@ +"""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 From c8a731057e9d51d44d28e43613c29c93e813ec7e Mon Sep 17 00:00:00 2001 From: FunJim Date: Tue, 28 Jul 2026 18:56:17 +0800 Subject: [PATCH 31/43] Disable CodeBuddy auto memory in AGS rollouts CodeBuddy's auto-memory feature injects a `` block into every prompt. In SWE rollouts that block is pure overhead -- there is no cross-session continuity to build up -- and it teaches the model to emit `` narration of its own, which no tool-call parser can read. Verified on a 100-instance SWE-bench Verified subset: 0/100 trajectories carry memory markers afterwards, against 48/50 in the previous run. Note this does not by itself change the empty-patch rate (35.0% -> 33.0%, paired McNemar p=1.0); it is a prompt-hygiene fix, not a scoring fix. Also collapses three multi-line f-string concatenations onto single lines, as the repo formatter produces. --- .../generator/ags_generator/harnesses.py | 21 ++++--------------- 1 file changed, 4 insertions(+), 17 deletions(-) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py index 333b49128a..f7f0052123 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py @@ -21,11 +21,7 @@ class AGSSidecarClaudeCodeHarness(BaseHarness): 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" - ) + launch_flags = "--dangerously-skip-permissions --verbose --output-format stream-json --include-partial-messages --include-hook-events" static_env = { "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", @@ -43,11 +39,7 @@ async def install_cli(self, sb: Sandbox) -> None: async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: settings = json.dumps({"hasCompletedOnboarding": True, "bypassPermissionsModeAccepted": True}) await sb.exec( - "mkdir -p /root/.claude /home/agent/.claude && " - f"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", + 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, @@ -182,6 +174,7 @@ async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: "fileCheckpointingEnabled": False, "promptSuggestionEnabled": False, "enableAllProjectMcpServers": False, + "memory": {"autoMemoryEnabled": False}, } models_b64 = _json_b64(models_json) settings_b64 = _json_b64(settings_json) @@ -220,13 +213,7 @@ async def launch_and_wait(self, sb: Sandbox, ctx: HarnessContext, prompt: str, t parts.append("-y") session_log_dir = f"{ctx.workdir}/.harness/codebuddy_sessions" - raw_cmd = ( - f"cbc {' '.join(parts)} {shlex.quote(prompt)}; " - "rc=$?; " - f"mkdir -p {shlex.quote(session_log_dir)}/projects; " - f"cp -r /root/.codebuddy/projects/. {shlex.quote(session_log_dir)}/projects/ 2>/dev/null || true; " - "exit $rc" - ) + 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 = { From 6df7c14347682c380ed3825976e2becee2d397f5 Mon Sep 17 00:00:00 2001 From: FunJim Date: Tue, 28 Jul 2026 18:58:46 +0800 Subject: [PATCH 32/43] Add empty-patch guardrail to the AGS rollout Some coding-agent CLIs treat a text-only final turn as a finished answer: the model stops mid-task, no tool call is parsed, the CLI exits 0, and the rollout yields an empty diff. Measured with CodeBuddy Code on SWE-bench Verified that accounts for roughly a third of all rollouts, and it is indistinguishable in the artifacts from a genuine "no change needed" outcome. empty_patch_guard.py classifies those rollouts and splits them by why they stopped (mid-task narration / claimed completion / no final text / unclassified). It reuses the existing iter_trajectory_events reader rather than adding a third trajectory parser, and it never raises: a malformed trajectory degrades to "not triggered" so classification can never cost an otherwise usable rollout. SWE_EMPTY_PATCH_GUARD selects the policy and defaults to "metrics", which only labels and counts. "abort" is deliberately not the default: an empty-patch trajectory still holds real on-policy tokens whose zero reward is correct, i.e. the negative half of a GRPO group, so masking those out at a ~33% rate would bias the group baseline toward successes. It also only requeues under fully_async_rollout -- on the rollout_buffer path the AGS scripts actually use, an aborted sample ships with its loss masked and is never retried. Verified across five runs: the empty-patch rate reproduces at 32-36%, and on a full 500-instance Claude Code run the guard fired on exactly the 3 of 42 zero-byte patches that had exit code 0, with no false positives or negatives (a nonzero exit is an honestly reported failure and needs no extra label). Also surfaces the pre-existing ill_formed flag, which had no metrics outlet. --- .github/workflows/pr-test.yml | 4 + .github/workflows/pr-test.yml.j2 | 1 + .../run_qwen35_35b_a3b_swe_2nodes.sh | 1 + .../run_qwen35_35b_a3b_swe_4nodes.sh | 1 + .../generator/ags_generator/config.py | 28 +++ .../ags_generator/empty_patch_guard.py | 111 ++++++++++ .../generator/ags_generator/rollout.py | 55 +++++ .../generator/ags_generator/wandb_metrics.py | 27 +++ .../test_ags_empty_patch_guard.py | 197 ++++++++++++++++++ 9 files changed, 425 insertions(+) create mode 100644 slime_plugins/rollout_buffer/generator/ags_generator/empty_patch_guard.py create mode 100644 tests/test_rollout_buffer/test_ags_empty_patch_guard.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index bd98028295..0d662fb3f5 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -748,6 +748,10 @@ jobs: { "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" } ] defaults: diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index c19b884ed9..e4f2064e48 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -103,6 +103,7 @@ {'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}, ], }, diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh index 088bcbbb25..7a4366bda6 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh @@ -319,6 +319,7 @@ keys = ( "SWE_AGENT_TIME_BUDGET_SEC", "SWE_EVAL_TIMEOUT_SEC", "SWE_EVAL_ISOLATED_SANDBOX", "SWE_BOOT_CONCURRENCY", "SWE_BOOT_RETRIES", "SWE_ROLLOUT_GUARD_SEC", "SWE_ROLLOUT_CONCURRENCY", + "SWE_EMPTY_PATCH_GUARD", "SLIME_AGENT_CC_MAX_TURNS", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", "CLAUDE_CODE_MAX_OUTPUT_TOKENS", "SWE_CC_PROMPT", ) diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh index 3e1fc31049..6d672cc963 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh @@ -319,6 +319,7 @@ keys = ( "SWE_AGENT_TIME_BUDGET_SEC", "SWE_EVAL_TIMEOUT_SEC", "SWE_EVAL_ISOLATED_SANDBOX", "SWE_BOOT_CONCURRENCY", "SWE_BOOT_RETRIES", "SWE_ROLLOUT_GUARD_SEC", "SWE_ROLLOUT_CONCURRENCY", + "SWE_EMPTY_PATCH_GUARD", "SLIME_AGENT_CC_MAX_TURNS", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", "CLAUDE_CODE_MAX_OUTPUT_TOKENS", "SWE_CC_PROMPT", ) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/config.py b/slime_plugins/rollout_buffer/generator/ags_generator/config.py index c952dda8e3..f632fa4cf0 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/config.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/config.py @@ -5,6 +5,20 @@ 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"}) + @dataclass(frozen=True) class AGSGeneratorConfig: @@ -24,6 +38,7 @@ class AGSGeneratorConfig: artifact_dir: str | None enable_token2text: bool prompt: str + empty_patch_guard: str @classmethod def from_env(cls, *, enable_token2text: bool = False) -> AGSGeneratorConfig: @@ -52,9 +67,22 @@ def from_env(cls, *, enable_token2text: bool = False) -> AGSGeneratorConfig: "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")), ) +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 _env_flag(name: str, *, default: bool) -> bool: raw = os.environ.get(name) if raw is None or raw == "": 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/rollout.py b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py index 299503b30f..2fa71f65e0 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py @@ -18,6 +18,7 @@ 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 @@ -102,6 +103,12 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam ) 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( @@ -116,6 +123,13 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam 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, @@ -143,6 +157,10 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam "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, ) @@ -160,6 +178,8 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam "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", @@ -212,6 +232,41 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam traceback.format_exc(), ) + 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 diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/wandb_metrics.py b/slime_plugins/rollout_buffer/generator/ags_generator/wandb_metrics.py index 797ea21811..77e5f1f782 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/wandb_metrics.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/wandb_metrics.py @@ -170,6 +170,7 @@ def _compute_ags_metrics(args: Any, samples: Iterable[Sample]) -> dict[str, floa 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 @@ -249,6 +250,32 @@ def _abort_reason_metrics(metadata: list[dict[str, Any]], n: int) -> dict[str, f 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)): 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..74b282b13a --- /dev/null +++ b/tests/test_rollout_buffer/test_ags_empty_patch_guard.py @@ -0,0 +1,197 @@ +"""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() + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) From ddf52192a28b340ba433182a7b59aafdf8c3917b Mon Sep 17 00:00:00 2001 From: FunJim Date: Wed, 29 Jul 2026 12:26:40 +0800 Subject: [PATCH 33/43] Keep text and reasoning on the OpenAI adapter's manager message manager_message is what re-renders as chat history on the next turn, so any field it drops stops the re-render from reproducing the ids we sampled. _SampleBuilder then sees token drift and either REALIGNs -- rewriting a real trained response as loss_mask=0 -- or forks the chain. _build_reply_parts dropped two fields: reasoning_content was never set on the leaf at all, and content was blanked whenever tool_calls were present. Measured over 150 rollouts of a 500-instance SWE-bench Verified run, that split each CodeBuddy rollout into a median of 25 samples where Claude Code produced 1, with num_samples predicted exactly by (wire responses - reasoning-only responses) in 150/150 cases. The comments justifying both omissions assumed clients rewrite history on echo. Probing the real CodeBuddy CLI over 40 turns shows it echoes text alongside tool_calls verbatim, so neither omission was buying anything. It does replay reasoning under "reasoning" rather than "reasoning_content", so _translate_messages now accepts either: the field survived the round trip all along and we were failing to read it (0/15 echoes matched the leaf before, 15/15 after). Fixing only the leaf would trade token drift for a message mismatch and lose more signal, so text and reasoning go on the wire too. This is a training-signal defect only. Scores from --debug-rollout-only runs never consume loss_mask and are unaffected; it would have corrupted GRPO advantage estimation as soon as CodeBuddy drove real training. Tests drive both real adapters over a real /generate upstream and assert one sample per clean chain plus the trained-token count, since a leaf that drops a field can keep the count at 1 while training almost nothing. They need a tokenizer whose assistant rendering a model could actually emit: FakeTokenizer renders tool calls as the opaque marker toolcall:, which carries no arguments and so fabricates drift on every tool turn. --- slime/agent/adapters/openai.py | 49 ++++++++---- tests/test_agent/_fakes.py | 76 +++++++++++++++++++ tests/test_agent/test_adapters.py | 120 +++++++++++++++++++++++++++++- 3 files changed, 230 insertions(+), 15 deletions(-) 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/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..20da702567 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,6 +389,32 @@ 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 # =========================================================================== @@ -418,6 +458,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) # =========================================================================== From 4c9cf9dd2ed697af33faf77e0a96fe1470dfb947 Mon Sep 17 00:00:00 2001 From: FunJim Date: Wed, 29 Jul 2026 16:41:28 +0800 Subject: [PATCH 34/43] Set the AGS prompt style separately for training and eval The prompt style decides whether the agent is handed the task text directly or told to go read PROBLEM_STATEMENT.md, and training and eval want different answers. Eval should measure the model solving the task, the way Harbor scores it, rather than its file-discovery turns; training may prefer it to work for the context. SWE_PROMPT_STYLE now covers training (instruction) and SWE_EVAL_PROMPT_STYLE covers periodic eval (dataset), both spelled out in the 2- and 4-node example scripts. Only "instruction" writes PROBLEM_STATEMENT.md now. Under "dataset" the prompt already carries the task text, so the file was a stray untracked artifact that only git_diff's exclude pathspec kept out of the patch, and an invitation to spend turns reading a restatement of the prompt. The style is passed to each generate() call rather than stored on the runner, because _AGSGenerateState is a singleton: the first caller's mode would otherwise stick for the life of the process and hand eval the training style, or the reverse. Drop the "inline" style, which rebuilt the prompt from the raw problem_statement and is superseded by "dataset". An explicit SWE_PROMPT_STYLE=inline now fails instead of silently changing the prompt. The converter loses --prompt-source and always uses Harbor's instruction.md, for both the prompt and metadata.problem_statement, so the two styles differ in when the agent sees the task and not in what it reads. The raw tests/config.json text moves to metadata.harbor.problem_statement: it keeps upstream CRLF on 252/500 of SWE-bench Verified, which tokenises differently, so it is provenance rather than something to feed an agent. --- .../run_qwen35_35b_a3b_swe_2nodes.sh | 13 +- .../run_qwen35_35b_a3b_swe_4nodes.sh | 13 +- .../generator/ags_generator/config.py | 34 ++++ .../generator/ags_generator/entry.py | 7 +- .../generator/ags_generator/rollout.py | 41 ++++- .../generator/ags_generator/swe_task.py | 19 +- .../test_harbor_task_to_slime_prompt_data.py | 44 ++++- .../test_ags_empty_patch_guard.py | 42 +++++ .../test_rollout_buffer/test_ags_generator.py | 169 ++++++++++++++++++ tools/harbor_task_to_slime_prompt_data.py | 43 +++-- 10 files changed, 400 insertions(+), 25 deletions(-) diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh index 7a4366bda6..3fcc1e064c 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh @@ -109,6 +109,17 @@ 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}" + # # autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI compacts before any # # segment crosses the training-side cap. `investigator` is a read-only sub-agent. # SETTINGS_JSON='{"permissions":{"defaultMode":"bypassPermissions"},"autoCompactEnabled":true,"autoCompactWindow":80000}' @@ -319,7 +330,7 @@ keys = ( "SWE_AGENT_TIME_BUDGET_SEC", "SWE_EVAL_TIMEOUT_SEC", "SWE_EVAL_ISOLATED_SANDBOX", "SWE_BOOT_CONCURRENCY", "SWE_BOOT_RETRIES", "SWE_ROLLOUT_GUARD_SEC", "SWE_ROLLOUT_CONCURRENCY", - "SWE_EMPTY_PATCH_GUARD", + "SWE_EMPTY_PATCH_GUARD", "SWE_PROMPT_STYLE", "SWE_EVAL_PROMPT_STYLE", "SLIME_AGENT_CC_MAX_TURNS", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", "CLAUDE_CODE_MAX_OUTPUT_TOKENS", "SWE_CC_PROMPT", ) diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh index 6d672cc963..a8795fadb1 100644 --- a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh +++ b/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh @@ -109,6 +109,17 @@ 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}" + # # autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI compacts before any # # segment crosses the training-side cap. `investigator` is a read-only sub-agent. # SETTINGS_JSON='{"permissions":{"defaultMode":"bypassPermissions"},"autoCompactEnabled":true,"autoCompactWindow":80000}' @@ -319,7 +330,7 @@ keys = ( "SWE_AGENT_TIME_BUDGET_SEC", "SWE_EVAL_TIMEOUT_SEC", "SWE_EVAL_ISOLATED_SANDBOX", "SWE_BOOT_CONCURRENCY", "SWE_BOOT_RETRIES", "SWE_ROLLOUT_GUARD_SEC", "SWE_ROLLOUT_CONCURRENCY", - "SWE_EMPTY_PATCH_GUARD", + "SWE_EMPTY_PATCH_GUARD", "SWE_PROMPT_STYLE", "SWE_EVAL_PROMPT_STYLE", "SLIME_AGENT_CC_MAX_TURNS", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", "CLAUDE_CODE_MAX_OUTPUT_TOKENS", "SWE_CC_PROMPT", ) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/config.py b/slime_plugins/rollout_buffer/generator/ags_generator/config.py index f632fa4cf0..c6d7ad2dbc 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/config.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/config.py @@ -19,6 +19,23 @@ # 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: @@ -39,6 +56,8 @@ class AGSGeneratorConfig: 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: @@ -68,8 +87,14 @@ def from_env(cls, *, enable_token2text: bool = False) -> AGSGeneratorConfig: "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". @@ -83,6 +108,15 @@ def _empty_patch_guard_policy(raw: str | None) -> str: 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 == "": diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py index 1e1008566b..5a3175afbf 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -56,11 +56,16 @@ async def 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) + samples = await state.runner.generate(base_sample, sampling_params, evaluation=evaluation) if not evaluation: return samples diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py index 2fa71f65e0..556ff45e71 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py @@ -51,7 +51,7 @@ def __init__( ) self._boot_sem = asyncio.Semaphore(self.config.boot_concurrency) - async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sample]: + 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) @@ -92,14 +92,20 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam session_opened = True async with asyncio.timeout(self.config.rollout_guard_sec): async with self._boot_agent_sandbox(md["image"], instance_id) as sb: - await prepare_workspace(sb, md["workdir"], md) + 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.config.prompt, + 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"]) @@ -232,6 +238,35 @@ async def generate(self, base_sample: Sample, sampling_params: dict) -> list[Sam 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, *, diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py b/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py index fd9e0697a5..b4ff2d87fd 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/swe_task.py @@ -31,6 +31,10 @@ def get_metadata(sample: Sample) -> dict[str, Any]: "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"), @@ -48,7 +52,17 @@ def _coerce_prompt(prompt) -> str: return "" -async def prepare_workspace(sb: Sandbox, workdir: str, md: dict[str, Any]) -> None: +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: @@ -56,7 +70,8 @@ async def prepare_workspace(sb: Sandbox, workdir: str, md: dict[str, Any]) -> No pre_commands = md.get("pre_commands") if pre_commands: await apply_pre_commands(sb, workdir, pre_commands) - await sb.write_file(f"{workdir}/PROBLEM_STATEMENT.md", md.get("problem_statement") or "", user="agent") + 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: diff --git a/tests/test_harbor_task_to_slime_prompt_data.py b/tests/test_harbor_task_to_slime_prompt_data.py index c89c528f5a..8e487e091a 100644 --- a/tests/test_harbor_task_to_slime_prompt_data.py +++ b/tests/test_harbor_task_to_slime_prompt_data.py @@ -64,7 +64,6 @@ def test_converter_preserves_image_head_by_default(converter_module, harbor_task prompt_alias_key="", label_key="label", metadata_key="metadata", - prompt_source="problem_statement", image_override=None, default_workdir="/testbed", include_eval_cmd=True, @@ -78,6 +77,49 @@ def test_converter_preserves_image_head_by_default(converter_module, harbor_task 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" diff --git a/tests/test_rollout_buffer/test_ags_empty_patch_guard.py b/tests/test_rollout_buffer/test_ags_empty_patch_guard.py index 74b282b13a..2cbf5f3786 100644 --- a/tests/test_rollout_buffer/test_ags_empty_patch_guard.py +++ b/tests/test_rollout_buffer/test_ags_empty_patch_guard.py @@ -193,5 +193,47 @@ def test_empty_patch_guard_policy_rejects_unknown_value(monkeypatch): 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 index f3caca284d..ad831cbfb4 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -2,6 +2,7 @@ import asyncio import base64 +import contextlib import json import re import sys @@ -21,6 +22,7 @@ transform_group, ) from slime_plugins.rollout_buffer.generator.ags_generator.harnesses import 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 ( @@ -120,6 +122,173 @@ def test_eval_isolated_sandbox_can_be_enabled(monkeypatch): assert AGSGeneratorConfig.from_env().eval_isolated_sandbox is True +_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( diff --git a/tools/harbor_task_to_slime_prompt_data.py b/tools/harbor_task_to_slime_prompt_data.py index 42e1cdc475..e792a76b88 100644 --- a/tools/harbor_task_to_slime_prompt_data.py +++ b/tools/harbor_task_to_slime_prompt_data.py @@ -11,6 +11,10 @@ - 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 @@ -114,12 +118,6 @@ def parse_args() -> argparse.Namespace: ) 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( - "--prompt-source", - choices=("problem_statement", "instruction"), - default="problem_statement", - help="Which extracted text to put in the primary prompt field.", - ) parser.add_argument( "--default-workdir", default="/testbed", @@ -223,7 +221,6 @@ def task_to_row( prompt_alias_key: str, label_key: str, metadata_key: str, - prompt_source: str, image_override: str | None, default_workdir: str, include_eval_cmd: bool, @@ -239,8 +236,14 @@ def task_to_row( instance_id = str(swe_config.get("instance_id") or task_dir.name) source_name = source or dataset_root.name - problem_statement = str(swe_config.get("problem_statement") or instruction) - prompt = problem_statement if prompt_source == "problem_statement" else instruction + # 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'}") @@ -256,7 +259,7 @@ def task_to_row( "source": source_name, "image": image, "workdir": workdir, - "problem_statement": problem_statement, + "problem_statement": instruction, "harbor": harbor_metadata(task_dir, source_name, task_toml, swe_config, image, workdir), } if include_eval_cmd: @@ -331,6 +334,9 @@ def harbor_metadata( { "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"), @@ -441,7 +447,10 @@ def write_pretty_example(rows: list[dict[str, Any]], output: Path) -> None: 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."}, + 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"], @@ -450,7 +459,13 @@ def write_schema(output: Path, *, input_key: str, prompt_alias_key: str, label_k "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"}, + "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": ( @@ -497,7 +512,6 @@ def convert_tasks( prompt_alias_key: str, label_key: str, metadata_key: str, - prompt_source: str, image_override: str | None, default_workdir: str, include_eval_cmd: bool, @@ -520,7 +534,6 @@ def convert_tasks( prompt_alias_key=prompt_alias_key, label_key=label_key, metadata_key=metadata_key, - prompt_source=prompt_source, image_override=image_override, default_workdir=default_workdir, include_eval_cmd=include_eval_cmd, @@ -544,7 +557,6 @@ def convert_tasks( prompt_alias_key=prompt_alias_key, label_key=label_key, metadata_key=metadata_key, - prompt_source=prompt_source, image_override=image_override, default_workdir=default_workdir, include_eval_cmd=include_eval_cmd, @@ -638,7 +650,6 @@ def main() -> None: prompt_alias_key=args.prompt_alias_key, label_key=args.label_key, metadata_key=args.metadata_key, - prompt_source=args.prompt_source, image_override=args.image, default_workdir=args.default_workdir, include_eval_cmd=not args.no_eval_cmd, From 85b86b1181f586790b4172a92b3d790345190a8d Mon Sep 17 00:00:00 2001 From: FunJim Date: Wed, 29 Jul 2026 21:53:37 +0800 Subject: [PATCH 35/43] Let AGS eval stand up its own adapter, and stop the CLI overriding its sampling Eval reached AGS through RemoteAdapterService, which only proxies an adapter that someone else already bound. The training rollout is what binds it, inside the rollout-buffer process, so eval-before-train on rollout 0 and --num-rollout 0 had nothing to talk to and every prompt died on connection refused. get_adapter_service now health-probes the control URL and reuses the training adapter when it answers -- that process owns the trajectory trees, and a second bind on the same port would fail -- otherwise it starts a local adapter on an ephemeral port. Ephemeral because the training adapter may still claim ADAPTER_PORT later in the same run. That fallback lives in the RolloutManager actor, which Ray does not pin to the head node, so it advertises the local node's own IP and the port it actually bound rather than ADAPTER_PUBLIC_HOST/ADAPTER_PUBLIC_BASE_URL -- both of which name the head and a fixed port, and would send sandboxes to the wrong address. Sampling defaults from open_session now outrank the request body. Probing the real CLIs against a recording stub: codebuddy sends temperature=1 on every /v1/chat/completions call, so --eval-temperature never reached sglang under SWE_AGENT=codebuddy_code; Claude Code sends none of temperature/top_p/top_k, which is why only the codebuddy path was affected. Body values still apply for keys the caller left unset. AdapterService stays a singleton, so a later caller's args are still ignored -- tearing down a live adapter would lose in-flight trajectories. It now records its construction inputs and warns when they differ, instead of silently serving an adapter built with another tokenizer or context budget. --- slime/agent/adapters/common.py | 14 +- .../ags_generator/adapter_service.py | 131 +++++++++++++++- .../generator/ags_generator/entry.py | 6 + .../generator/ags_generator/rollout.py | 17 ++- tests/test_agent/test_adapters.py | 66 +++++++- .../test_rollout_buffer/test_ags_generator.py | 142 +++++++++++++++++- 6 files changed, 363 insertions(+), 13 deletions(-) diff --git a/slime/agent/adapters/common.py b/slime/agent/adapters/common.py index 117fc9177b..1964118983 100644 --- a/slime/agent/adapters/common.py +++ b/slime/agent/adapters/common.py @@ -458,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_plugins/rollout_buffer/generator/ags_generator/adapter_service.py b/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py index da82a88037..bad050c837 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/adapter_service.py @@ -10,6 +10,7 @@ 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 @@ -76,8 +77,39 @@ async def drop_session(self, sid: str, *, wait_timeout: float = 5.0) -> None: ) +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) -> None: + 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 @@ -101,14 +133,34 @@ def __init__(self, args: Namespace, config: AGSGeneratorConfig, adapter_cls: typ 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=config.adapter_port, + port=bind_port, thread_name="ags-rollout-adapter", runner_kwargs={"handler_cancellation": True, "access_log_class": FilteredAccessLogger}, ) - self.adapter_url = public_base_url or f"http://{config.adapter_public_host}:{self.app_handle.port}" + # 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, @@ -137,9 +189,82 @@ def __init__(self, args: Namespace, config: AGSGeneratorConfig) -> None: ) 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/entry.py b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py index 5a3175afbf..10f951bd7c 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/entry.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/entry.py @@ -36,6 +36,12 @@ 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) diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py index 556ff45e71..49913bdbff 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/rollout.py @@ -14,7 +14,7 @@ from slime.utils.types import Sample -from .adapter_service import AdapterService, RemoteAdapterService +from .adapter_service import get_adapter_service from .ags_sandbox import AGSSandbox from .artifacts import ArtifactWriter, sample_artifact_id from .config import AGSGeneratorConfig @@ -38,10 +38,17 @@ def __init__( self.args = args self.config = config or AGSGeneratorConfig.from_env() self.harness_cls, self.adapter_cls = resolve_agent(self.config.agent_name) - if use_remote_adapter: - self.adapter_service = RemoteAdapterService(args, self.config) - else: - self.adapter_service = AdapterService(args, self.config, self.adapter_cls) + # 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, diff --git a/tests/test_agent/test_adapters.py b/tests/test_agent/test_adapters.py index 20da702567..98590a05a7 100644 --- a/tests/test_agent/test_adapters.py +++ b/tests/test_agent/test_adapters.py @@ -416,10 +416,74 @@ def test_openai_manager_message_keeps_text_and_reasoning_with_tool_calls(): # =========================================================================== -# §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: diff --git a/tests/test_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py index ad831cbfb4..bc59fb5646 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -10,10 +10,12 @@ from types import SimpleNamespace import pytest +from aiohttp import web from tests.test_agent._fakes import FakeSandbox +from slime.utils.misc import SingletonMeta from slime.utils.types import Sample -from slime_plugins.rollout_buffer.generator.ags_generator import swe_task +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, @@ -122,6 +124,144 @@ def test_eval_isolated_sandbox_can_be_enabled(monkeypatch): 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."} From 0185e6078665411def8bda33d3101162a17ba3ad Mon Sep 17 00:00:00 2001 From: FunJim Date: Tue, 4 Aug 2026 14:46:41 +0800 Subject: [PATCH 36/43] Add an eval-only AGS launcher for scoring a single checkpoint Scores one Megatron checkpoint on the full SWE-bench Verified set with no training: --num-rollout 0 selects train.py's eval-only branch, so no rollout buffer is started and the eval path stands up its own AGS adapter. The non-obvious parts are commented in place -- the LR scheduler asserts that train_iters=0 triggers, why --no-load-optim is required for checkpoints written with --optimizer-cpu-offload, and the Ray port overrides these H20 nodes need. Also gate vision extraction in Dataset on multimodal_keys rather than on the mere presence of a processor. AutoProcessor returns one for any VL-capable checkpoint, so Qwen3.5-35B-A3B forced a conversation-shaped prompt on text-only data, which made --apply-chat-template mandatory and rewrote Sample.prompt into a templated "<|im_start|>user ..." string. Agent rollouts hand that prompt straight to a CLI, so they were fed the literal markers. --- .../eval_qwen35_35b_a3b_swe_2nodes.sh | 519 ++++++++++++++++++ slime/utils/data.py | 13 +- 2 files changed, 530 insertions(+), 2 deletions(-) create mode 100755 examples/claude_code_ags/eval_qwen35_35b_a3b_swe_2nodes.sh diff --git a/examples/claude_code_ags/eval_qwen35_35b_a3b_swe_2nodes.sh b/examples/claude_code_ags/eval_qwen35_35b_a3b_swe_2nodes.sh new file mode 100755 index 0000000000..e024a4885a --- /dev/null +++ b/examples/claude_code_ags/eval_qwen35_35b_a3b_swe_2nodes.sh @@ -0,0 +1,519 @@ +#!/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/claude_code_ags/eval_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:-claude_code_ags_eval_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}" + +export SLIME_AGENT_CC_MAX_TURNS="${SLIME_AGENT_CC_MAX_TURNS:-100}" +export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:---max-turns ${SLIME_AGENT_CC_MAX_TURNS}}" + +# ============ 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}/wandb" + --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/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 From 3d63df9b6c320ed3c09efbc5820b238932f5cd55 Mon Sep 17 00:00:00 2001 From: FunJim Date: Wed, 5 Aug 2026 10:57:19 +0800 Subject: [PATCH 37/43] Rename the Claude Code AGS examples for multi-harness use The examples directory now holds launchers for more than one coding agent, so move them out of claude_code_ags/ into coding_agent_rl_ags/ and prefix each script with the harness it drives (cc_). --- .../eval_cc_qwen35_35b_a3b_swe_2nodes.sh} | 0 .../run_cc_qwen35_35b_a3b_swe_2nodes.sh} | 0 .../run_cc_qwen35_35b_a3b_swe_4nodes.sh} | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename examples/{claude_code_ags/eval_qwen35_35b_a3b_swe_2nodes.sh => coding_agent_rl_ags/eval_cc_qwen35_35b_a3b_swe_2nodes.sh} (100%) rename examples/{claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh => coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_2nodes.sh} (100%) rename examples/{claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh => coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_4nodes.sh} (100%) diff --git a/examples/claude_code_ags/eval_qwen35_35b_a3b_swe_2nodes.sh b/examples/coding_agent_rl_ags/eval_cc_qwen35_35b_a3b_swe_2nodes.sh similarity index 100% rename from examples/claude_code_ags/eval_qwen35_35b_a3b_swe_2nodes.sh rename to examples/coding_agent_rl_ags/eval_cc_qwen35_35b_a3b_swe_2nodes.sh diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh b/examples/coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_2nodes.sh similarity index 100% rename from examples/claude_code_ags/run_qwen35_35b_a3b_swe_2nodes.sh rename to examples/coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_2nodes.sh diff --git a/examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh b/examples/coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_4nodes.sh similarity index 100% rename from examples/claude_code_ags/run_qwen35_35b_a3b_swe_4nodes.sh rename to examples/coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_4nodes.sh From be939b764300507807163c92852200b720e5ed3d Mon Sep 17 00:00:00 2001 From: FunJim Date: Wed, 5 Aug 2026 11:12:39 +0800 Subject: [PATCH 38/43] Add CodeBuddy launchers and cut the harness knobs to two per agent The AGS examples only had Claude Code entry points, so scoring CodeBuddy meant hand-editing SWE_AGENT. Add run_cbc_* / eval_cbc_* wrappers that set SWE_AGENT and re-exec the harness-agnostic launcher, rather than copying ~370 lines per harness for the two copies to drift apart. Reduce each harness's env surface to EXTRA_ARGS + EXTRA_ENVS. The removed knobs (max turns, per-turn output cap, thinking, tool lists) were either properties of the harness rather than of a run, or unreachable: none of the SLIME_AGENT_CBC_* names were in the launchers' Ray env allowlist, so setting them was silently ignored. Both knobs are now applied last -- EXTRA_ARGS after the harness's own flags, EXTRA_ENVS after static_env -- so either can override a default, which relies on both CLIs taking the last occurrence of a repeated flag. Deny WebSearch/WebFetch on both harnesses. Web access makes a rollout unreproducible and lets the agent look up the fix being graded; it was denied for CodeBuddy but not for Claude Code, which then used those tools on 1-2% of instances. Do it with --disallowedTools, which the CLI enforces, not --tools, which does not restrict the surface. Cover the new contract in tests: that EXTRA_ARGS and EXTRA_ENVS are applied late enough to win, that the web tools are denied on both harnesses, and that no turn or output cap is imposed. Verified by mutation -- clearing default_flags or moving EXTRA_ARGS ahead of it turns these red. --- .../eval_cbc_qwen35_35b_a3b_swe_2nodes.sh | 21 ++++ .../eval_cc_qwen35_35b_a3b_swe_2nodes.sh | 25 +++- .../run_cbc_qwen35_35b_a3b_swe_2nodes.sh | 21 ++++ .../run_cbc_qwen35_35b_a3b_swe_4nodes.sh | 14 +++ .../run_cc_qwen35_35b_a3b_swe_2nodes.sh | 27 ++++- .../run_cc_qwen35_35b_a3b_swe_4nodes.sh | 27 ++++- .../generator/ags_generator/harnesses.py | 66 ++++------ .../test_rollout_buffer/test_ags_generator.py | 114 ++++++++++++++++-- 8 files changed, 245 insertions(+), 70 deletions(-) create mode 100644 examples/coding_agent_rl_ags/eval_cbc_qwen35_35b_a3b_swe_2nodes.sh create mode 100644 examples/coding_agent_rl_ags/run_cbc_qwen35_35b_a3b_swe_2nodes.sh create mode 100644 examples/coding_agent_rl_ags/run_cbc_qwen35_35b_a3b_swe_4nodes.sh 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 index e024a4885a..5753fe6570 100755 --- 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 @@ -33,7 +33,7 @@ # LOAD_DIR=/data_train/ericxjzheng/experiments//checkpoints \ # CKPT_STEP=79 TRAIN_NUM_ROLLOUT=100 \ # E2B_API_KEY=... WANDB_API_KEY=... WANDB_ENTITY=... \ -# bash examples/claude_code_ags/eval_qwen35_35b_a3b_swe_2nodes.sh +# 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 @@ -119,7 +119,7 @@ REF_MODEL_PATH="${REF_MODEL_PATH:-/data_train/ericxjzheng/models/Qwen3.5-35B-A3B LOAD_DIR="${LOAD_DIR:-}" CKPT_STEP="${CKPT_STEP:-}" -EXP_TAG="${EXP_TAG:-claude_code_ags_eval_qwen35_35b_a3b${CKPT_STEP:+_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}}" @@ -199,8 +199,20 @@ export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-32}" export SWE_PROMPT_STYLE="${SWE_PROMPT_STYLE:-instruction}" export SWE_EVAL_PROMPT_STYLE="${SWE_EVAL_PROMPT_STYLE:-dataset}" -export SLIME_AGENT_CC_MAX_TURNS="${SLIME_AGENT_CC_MAX_TURNS:-100}" -export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:---max-turns ${SLIME_AGENT_CC_MAX_TURNS}}" +# The only two harness knobs: 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. +export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:-}" +export SLIME_AGENT_CC_EXTRA_ENVS="${SLIME_AGENT_CC_EXTRA_ENVS:-}" + +# The CodeBuddy equivalents, so SWE_AGENT=codebuddy_code works from this script +# too (run_cbc_*.sh just sets SWE_AGENT and re-execs this one). +export SLIME_AGENT_CBC_EXTRA_ARGS="${SLIME_AGENT_CBC_EXTRA_ARGS:-}" +export SLIME_AGENT_CBC_EXTRA_ENVS="${SLIME_AGENT_CBC_EXTRA_ENVS:-}" # ============ proxy bypass for in-cluster/AGS traffic ============ export no_proxy="127.0.0.1,${MASTER_ADDR},${ADAPTER_PUBLIC_HOST},${E2B_DOMAIN},.tencentags.com" @@ -468,8 +480,9 @@ keys = ( "SWE_BOOT_CONCURRENCY", "SWE_BOOT_RETRIES", "SWE_ROLLOUT_GUARD_SEC", "SWE_ROLLOUT_CONCURRENCY", "SWE_EMPTY_PATCH_GUARD", "SWE_PROMPT_STYLE", "SWE_EVAL_PROMPT_STYLE", - "SLIME_AGENT_CC_MAX_TURNS", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", - "CLAUDE_CODE_MAX_OUTPUT_TOKENS", "SWE_CC_PROMPT", + "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", + "SLIME_AGENT_CBC_EXTRA_ARGS", "SLIME_AGENT_CBC_EXTRA_ENVS", + "SWE_CC_PROMPT", ) env = {k: os.environ[k] for k in keys if k in os.environ} env["MASTER_ADDR"] = os.environ["MASTER_ADDR"] 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 index 3fcc1e064c..3e2ea0eb0a 100644 --- 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 @@ -62,7 +62,7 @@ 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:-claude_code_ags_qwen35_35b_a3b_2nodes}" +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}}" @@ -124,9 +124,23 @@ export SWE_EVAL_PROMPT_STYLE="${SWE_EVAL_PROMPT_STYLE:-dataset}" # # segment crosses the training-side cap. `investigator` is a read-only sub-agent. # SETTINGS_JSON='{"permissions":{"defaultMode":"bypassPermissions"},"autoCompactEnabled":true,"autoCompactWindow":80000}' # 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="--settings '${SETTINGS_JSON}' --disable-slash-commands --agents '${AGENTS_JSON}' --disallowedTools WebFetch WebSearch" -export SLIME_AGENT_CC_MAX_TURNS="${SLIME_AGENT_CC_MAX_TURNS:-100}" -export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:---max-turns ${SLIME_AGENT_CC_MAX_TURNS}}" +# export SLIME_AGENT_CC_EXTRA_ARGS="--settings '${SETTINGS_JSON}' --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: 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. +export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:-}" +export SLIME_AGENT_CC_EXTRA_ENVS="${SLIME_AGENT_CC_EXTRA_ENVS:-}" + +# The CodeBuddy equivalents, so SWE_AGENT=codebuddy_code works from this script +# too (run_cbc_*.sh just sets SWE_AGENT and re-execs this one). +export SLIME_AGENT_CBC_EXTRA_ARGS="${SLIME_AGENT_CBC_EXTRA_ARGS:-}" +export SLIME_AGENT_CBC_EXTRA_ENVS="${SLIME_AGENT_CBC_EXTRA_ENVS:-}" # 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." @@ -331,8 +345,9 @@ keys = ( "SWE_BOOT_CONCURRENCY", "SWE_BOOT_RETRIES", "SWE_ROLLOUT_GUARD_SEC", "SWE_ROLLOUT_CONCURRENCY", "SWE_EMPTY_PATCH_GUARD", "SWE_PROMPT_STYLE", "SWE_EVAL_PROMPT_STYLE", - "SLIME_AGENT_CC_MAX_TURNS", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", - "CLAUDE_CODE_MAX_OUTPUT_TOKENS", "SWE_CC_PROMPT", + "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", + "SLIME_AGENT_CBC_EXTRA_ARGS", "SLIME_AGENT_CBC_EXTRA_ENVS", + "SWE_CC_PROMPT", ) env = {k: os.environ[k] for k in keys if k in os.environ} env["MASTER_ADDR"] = os.environ["MASTER_ADDR"] 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 index a8795fadb1..784e336380 100644 --- 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 @@ -62,7 +62,7 @@ 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:-claude_code_ags_qwen35_35b_a3b_4nodes}" +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}}" @@ -124,9 +124,23 @@ export SWE_EVAL_PROMPT_STYLE="${SWE_EVAL_PROMPT_STYLE:-dataset}" # # segment crosses the training-side cap. `investigator` is a read-only sub-agent. # SETTINGS_JSON='{"permissions":{"defaultMode":"bypassPermissions"},"autoCompactEnabled":true,"autoCompactWindow":80000}' # 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="--settings '${SETTINGS_JSON}' --disable-slash-commands --agents '${AGENTS_JSON}' --disallowedTools WebFetch WebSearch" -export SLIME_AGENT_CC_MAX_TURNS="${SLIME_AGENT_CC_MAX_TURNS:-100}" -export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:---max-turns ${SLIME_AGENT_CC_MAX_TURNS}}" +# export SLIME_AGENT_CC_EXTRA_ARGS="--settings '${SETTINGS_JSON}' --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: 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. +export SLIME_AGENT_CC_EXTRA_ARGS="${SLIME_AGENT_CC_EXTRA_ARGS:-}" +export SLIME_AGENT_CC_EXTRA_ENVS="${SLIME_AGENT_CC_EXTRA_ENVS:-}" + +# The CodeBuddy equivalents, so SWE_AGENT=codebuddy_code works from this script +# too (run_cbc_*.sh just sets SWE_AGENT and re-execs this one). +export SLIME_AGENT_CBC_EXTRA_ARGS="${SLIME_AGENT_CBC_EXTRA_ARGS:-}" +export SLIME_AGENT_CBC_EXTRA_ENVS="${SLIME_AGENT_CBC_EXTRA_ENVS:-}" # 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." @@ -331,8 +345,9 @@ keys = ( "SWE_BOOT_CONCURRENCY", "SWE_BOOT_RETRIES", "SWE_ROLLOUT_GUARD_SEC", "SWE_ROLLOUT_CONCURRENCY", "SWE_EMPTY_PATCH_GUARD", "SWE_PROMPT_STYLE", "SWE_EVAL_PROMPT_STYLE", - "SLIME_AGENT_CC_MAX_TURNS", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", - "CLAUDE_CODE_MAX_OUTPUT_TOKENS", "SWE_CC_PROMPT", + "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", + "SLIME_AGENT_CBC_EXTRA_ARGS", "SLIME_AGENT_CBC_EXTRA_ENVS", + "SWE_CC_PROMPT", ) env = {k: os.environ[k] for k in keys if k in os.environ} env["MASTER_ADDR"] = os.environ["MASTER_ADDR"] diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py index f7f0052123..1f0b11eb4a 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py @@ -22,6 +22,11 @@ class AGSSidecarClaudeCodeHarness(BaseHarness): 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", @@ -46,9 +51,10 @@ async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: ) 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}" + 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 = { @@ -65,9 +71,8 @@ async def launch_and_wait(self, sb: Sandbox, ctx: HarnessContext, prompt: str, t } 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)) - if os.environ.get("CLAUDE_CODE_MAX_OUTPUT_TOKENS"): - env["CLAUDE_CODE_MAX_OUTPUT_TOKENS"] = os.environ["CLAUDE_CODE_MAX_OUTPUT_TOKENS"] return await run_root_command( sb, @@ -101,26 +106,17 @@ class CodeBuddyCodeHarness(BaseHarness): name = "codebuddy_code" extra_args_env = "SLIME_AGENT_CBC_EXTRA_ARGS" extra_envs_env = "SLIME_AGENT_CBC_EXTRA_ENVS" - max_turns_env = "SLIME_AGENT_CBC_MAX_TURNS" - max_output_tokens_env = "SLIME_AGENT_CBC_MAX_OUTPUT_TOKENS" - thinking_enabled_env = "SLIME_AGENT_CBC_THINKING_ENABLED" - tools_env = "SLIME_AGENT_CBC_TOOLS" - - # Keep the default tool surface close to Claude Code's coding-agent use case - # while disabling internet search by default for reproducible SWE rollouts. - allowed_tools = ( - "Bash", - "Read", - "Write", - "Edit", - "Glob", - "Grep", - "TaskCreate", - "TaskUpdate", - "TaskGet", - "TaskList", - "Agent", - ) + + # 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( @@ -155,7 +151,6 @@ async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: "vendor": "OpenAI", "apiKey": ctx.session_id, "url": self._chat_completions_url(ctx.adapter_url), - "maxOutputTokens": int(os.environ.get(self.max_output_tokens_env, "16384")), "supportsToolCall": True, "supportsImages": False, "supportsReasoning": True, @@ -167,7 +162,10 @@ async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: "cleanupPeriodDays": 30, "includeCoAuthoredBy": False, "autoCompactEnabled": True, - "alwaysThinkingEnabled": _env_flag(self.thinking_enabled_env, default=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, @@ -198,18 +196,13 @@ async def launch_and_wait(self, sb: Sandbox, ctx: HarnessContext, prompt: str, t "--verbose", "--output-format stream-json", "--include-partial-messages", + *self.default_flags, ] - tools = os.environ.get(self.tools_env, ",".join(self.allowed_tools)).strip() - if tools: - parts.append(f"--tools {shlex.quote(tools)}") - parts.append("--disallowedTools WebSearch") extra = os.environ.get(self.extra_args_env, "").strip() if extra: - # Keep caller-provided flags before the non-variadic tail and prompt. + # After default_flags so it can override them, before the prompt so the + # prompt stays positional. parts.append(extra) - if not _env_flag(self.thinking_enabled_env, default=True): - parts.append("--effort none") - parts.append(f"--max-turns {int(os.environ.get(self.max_turns_env, '100'))}") parts.append("-y") session_log_dir = f"{ctx.workdir}/.harness/codebuddy_sessions" @@ -269,13 +262,6 @@ def _json_b64(value: dict) -> str: return base64.b64encode(payload).decode("ascii") -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"} - - HARNESS_REGISTRY: dict[str, tuple[type[BaseHarness], type]] = { "claude_code": (AGSSidecarClaudeCodeHarness, AnthropicAdapter), "codex": (CodexHarness, OpenAIAdapter), diff --git a/tests/test_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py index bc59fb5646..7dfa472930 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -23,7 +23,11 @@ is_valid_group, transform_group, ) -from slime_plugins.rollout_buffer.generator.ags_generator.harnesses import CodeBuddyCodeHarness, resolve_agent +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 @@ -921,10 +925,8 @@ async def run_case(): asyncio.run(run_case()) -def test_codebuddy_code_write_config_points_to_adapter(monkeypatch): +def test_codebuddy_code_write_config_points_to_adapter(): async def run_case(): - monkeypatch.setenv("SLIME_AGENT_CBC_MAX_OUTPUT_TOKENS", "8192") - monkeypatch.setenv("SLIME_AGENT_CBC_THINKING_ENABLED", "false") sb = FakeSandbox() await CodeBuddyCodeHarness().write_config(sb, _ctx(sid="sess-cbc", url="http://host:18001")) @@ -934,17 +936,17 @@ async def run_case(): 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]["maxOutputTokens"] == 8192 assert models["models"][0]["supportsToolCall"] is True - assert settings["alwaysThinkingEnabled"] is False + # 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(monkeypatch): +def test_codebuddy_code_launch_command_and_env(): async def run_case(): - monkeypatch.setenv("SLIME_AGENT_CBC_MAX_TURNS", "7") - monkeypatch.setenv("SLIME_AGENT_CBC_THINKING_ENABLED", "false") sb = FakeSandbox() rc = await CodeBuddyCodeHarness().launch_and_wait( sb, @@ -956,11 +958,19 @@ async def run_case(): 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 "--max-turns 7" in body assert "-y" in body and "solve it" in body - assert "--effort none" in body - assert "--tools Bash,Read,Write,Edit,Glob,Grep,TaskCreate,TaskUpdate,TaskGet,TaskList,Agent" 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 @@ -974,6 +984,86 @@ async def run_case(): 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): From 6b24af20356d68bcebca7580b8333a9ccce35587 Mon Sep 17 00:00:00 2001 From: FunJim Date: Wed, 5 Aug 2026 16:53:56 +0800 Subject: [PATCH 39/43] Trigger agent auto-compaction by percentage of the real context window Both CLIs clamp an absolute auto-compact window to [100k, 1M], which sits above our 96k rollout_max_context_len -- so an absolute setting could never fire before the adapter hard-stops a turn with finish_reason="length". Switch to the percentage overrides instead, and tell each CLI the real window so the percentage is of MAX_CONTEXT_LEN on both sides: - CodeBuddy: CODEBUDDY_AUTOCOMPACT_PCT_OVERRIDE is a percentage of the model's maxInputTokens, so CodeBuddyCodeHarness now writes that key into models.json from SLIME_AGENT_MAX_INPUT_TOKENS. Without it resolveCompactTriggerAt() falls back to the clamped absolute window. - Claude Code: CLAUDE_CODE_MAX_CONTEXT_TOKENS applies directly for model names it does not recognise as Claude models, which covers our "slime-actor" label, so no DISABLE_COMPACT is needed. MAX_OUTPUT_TOKENS is pinned to MAX_GEN_LEN rather than left at the 32000 default-for-unknown-ids. Default is 60%, not the CLIs' 70%: the per-turn output reservation means a prompt cannot grow past MAX_CONTEXT_LEN - MAX_GEN_LEN (~63k), so a 70% trigger (67200) would never be reached. The launchers now print the resolved trigger point and warn when it lands above that ceiling. Also fixes a shell bug in the same block: a JSON literal used as the default in ${VAR:-{...}} is mis-parsed, because the value's own closing brace terminates the expansion and leaks a stray "}" whenever the variable is already set. The JSON is built in a named variable first. --- .../eval_cc_qwen35_35b_a3b_swe_2nodes.sh | 82 +++++++++++++++-- .../run_cc_qwen35_35b_a3b_swe_2nodes.sh | 88 ++++++++++++++++--- .../run_cc_qwen35_35b_a3b_swe_4nodes.sh | 88 ++++++++++++++++--- .../generator/ags_generator/harnesses.py | 52 ++++++++--- .../test_rollout_buffer/test_ags_generator.py | 63 +++++++++++++ 5 files changed, 329 insertions(+), 44 deletions(-) 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 index 5753fe6570..fbd72da99a 100755 --- 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 @@ -199,20 +199,83 @@ export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-32}" export SWE_PROMPT_STYLE="${SWE_PROMPT_STYLE:-instruction}" export SWE_EVAL_PROMPT_STYLE="${SWE_EVAL_PROMPT_STYLE:-dataset}" -# The only two harness knobs: extra CLI flags, and extra env vars as JSON. -# Everything else (denied tools, the launch flags) is a class attribute in +# ---- 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. +# 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_CC_EXTRA_ENVS="${SLIME_AGENT_CC_EXTRA_ENVS:-}" - -# The CodeBuddy equivalents, so SWE_AGENT=codebuddy_code works from this script -# too (run_cbc_*.sh just sets SWE_AGENT and re-execs this one). export SLIME_AGENT_CBC_EXTRA_ARGS="${SLIME_AGENT_CBC_EXTRA_ARGS:-}" -export SLIME_AGENT_CBC_EXTRA_ENVS="${SLIME_AGENT_CBC_EXTRA_ENVS:-}" # ============ proxy bypass for in-cluster/AGS traffic ============ export no_proxy="127.0.0.1,${MASTER_ADDR},${ADAPTER_PUBLIC_HOST},${E2B_DOMAIN},.tencentags.com" @@ -482,6 +545,9 @@ keys = ( "SWE_EMPTY_PATCH_GUARD", "SWE_PROMPT_STYLE", "SWE_EVAL_PROMPT_STYLE", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", "SLIME_AGENT_CBC_EXTRA_ARGS", "SLIME_AGENT_CBC_EXTRA_ENVS", + # Read by CodeBuddyCodeHarness.write_config to set models.json maxInputTokens, + # which is what the auto-compact percentage is a percentage OF. + "SLIME_AGENT_MAX_INPUT_TOKENS", "SWE_CC_PROMPT", ) env = {k: os.environ[k] for k in keys if k in os.environ} 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 index 3e2ea0eb0a..ea588ed337 100644 --- 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 @@ -120,27 +120,88 @@ export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-32}" export SWE_PROMPT_STYLE="${SWE_PROMPT_STYLE:-instruction}" export SWE_EVAL_PROMPT_STYLE="${SWE_EVAL_PROMPT_STYLE:-dataset}" -# # autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI compacts before any -# # segment crosses the training-side cap. `investigator` is a read-only sub-agent. -# SETTINGS_JSON='{"permissions":{"defaultMode":"bypassPermissions"},"autoCompactEnabled":true,"autoCompactWindow":80000}' +# ---- 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="--settings '${SETTINGS_JSON}' --disable-slash-commands --agents '${AGENTS_JSON}'" +# 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: extra CLI flags, and extra env vars as JSON. -# Everything else (denied tools, the launch flags) is a class attribute in +# 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. +# 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_CC_EXTRA_ENVS="${SLIME_AGENT_CC_EXTRA_ENVS:-}" - -# The CodeBuddy equivalents, so SWE_AGENT=codebuddy_code works from this script -# too (run_cbc_*.sh just sets SWE_AGENT and re-execs this one). export SLIME_AGENT_CBC_EXTRA_ARGS="${SLIME_AGENT_CBC_EXTRA_ARGS:-}" -export SLIME_AGENT_CBC_EXTRA_ENVS="${SLIME_AGENT_CBC_EXTRA_ENVS:-}" # 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." @@ -347,6 +408,9 @@ keys = ( "SWE_EMPTY_PATCH_GUARD", "SWE_PROMPT_STYLE", "SWE_EVAL_PROMPT_STYLE", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", "SLIME_AGENT_CBC_EXTRA_ARGS", "SLIME_AGENT_CBC_EXTRA_ENVS", + # Read by CodeBuddyCodeHarness.write_config to set models.json maxInputTokens, + # which is what the auto-compact percentage is a percentage OF. + "SLIME_AGENT_MAX_INPUT_TOKENS", "SWE_CC_PROMPT", ) env = {k: os.environ[k] for k in keys if k in os.environ} 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 index 784e336380..66260e9686 100644 --- 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 @@ -120,27 +120,88 @@ export SWE_ROLLOUT_CONCURRENCY="${SWE_ROLLOUT_CONCURRENCY:-32}" export SWE_PROMPT_STYLE="${SWE_PROMPT_STYLE:-instruction}" export SWE_EVAL_PROMPT_STYLE="${SWE_EVAL_PROMPT_STYLE:-dataset}" -# # autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI compacts before any -# # segment crosses the training-side cap. `investigator` is a read-only sub-agent. -# SETTINGS_JSON='{"permissions":{"defaultMode":"bypassPermissions"},"autoCompactEnabled":true,"autoCompactWindow":80000}' +# ---- 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="--settings '${SETTINGS_JSON}' --disable-slash-commands --agents '${AGENTS_JSON}'" +# 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: extra CLI flags, and extra env vars as JSON. -# Everything else (denied tools, the launch flags) is a class attribute in +# 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. +# 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_CC_EXTRA_ENVS="${SLIME_AGENT_CC_EXTRA_ENVS:-}" - -# The CodeBuddy equivalents, so SWE_AGENT=codebuddy_code works from this script -# too (run_cbc_*.sh just sets SWE_AGENT and re-execs this one). export SLIME_AGENT_CBC_EXTRA_ARGS="${SLIME_AGENT_CBC_EXTRA_ARGS:-}" -export SLIME_AGENT_CBC_EXTRA_ENVS="${SLIME_AGENT_CBC_EXTRA_ENVS:-}" # 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." @@ -347,6 +408,9 @@ keys = ( "SWE_EMPTY_PATCH_GUARD", "SWE_PROMPT_STYLE", "SWE_EVAL_PROMPT_STYLE", "SLIME_AGENT_CC_EXTRA_ARGS", "SLIME_AGENT_CC_EXTRA_ENVS", "SLIME_AGENT_CBC_EXTRA_ARGS", "SLIME_AGENT_CBC_EXTRA_ENVS", + # Read by CodeBuddyCodeHarness.write_config to set models.json maxInputTokens, + # which is what the auto-compact percentage is a percentage OF. + "SLIME_AGENT_MAX_INPUT_TOKENS", "SWE_CC_PROMPT", ) env = {k: os.environ[k] for k in keys if k in os.environ} diff --git a/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py index 1f0b11eb4a..fc61c26867 100644 --- a/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py +++ b/slime_plugins/rollout_buffer/generator/ags_generator/harnesses.py @@ -106,6 +106,10 @@ class CodeBuddyCodeHarness(BaseHarness): 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 @@ -143,19 +147,31 @@ async def install_cli(self, sb: Sandbox) -> None: ) 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": [ - { - "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, - } - ], + "models": [model_entry], "availableModels": [ctx.model_label], } settings_json = { @@ -257,6 +273,18 @@ def _chat_completions_url(adapter_url: str) -> str: 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") diff --git a/tests/test_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py index 7dfa472930..9b5b83ea1c 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -1082,3 +1082,66 @@ async def run_case(): ) 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()) From 3d22eee17c9ce7669e2e90fee4574afcb4d865f1 Mon Sep 17 00:00:00 2001 From: FunJim Date: Thu, 6 Aug 2026 14:31:42 +0800 Subject: [PATCH 40/43] fix(scripts): update wandb-dir path to LOG_DIR directly Remove the extra `/wandb` subdirectory from `--wandb-dir` in training and evaluation scripts. --- .../coding_agent_rl_ags/eval_cc_qwen35_35b_a3b_swe_2nodes.sh | 2 +- .../coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_2nodes.sh | 2 +- .../coding_agent_rl_ags/run_cc_qwen35_35b_a3b_swe_4nodes.sh | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) 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 index fbd72da99a..7042123a1d 100755 --- 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 @@ -434,7 +434,7 @@ if [[ -n "${WANDB_API_KEY:-}" ]]; then --wandb-project "${WANDB_PROJECT:-slime-claude-code-ags}" --wandb-group "${WANDB_GROUP:-${EXP_TAG}}" --wandb-key "${WANDB_API_KEY}" - --wandb-dir "${LOG_DIR}/wandb" + --wandb-dir "${LOG_DIR}" --disable-wandb-random-suffix ) else 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 index ea588ed337..3a1d23aaa5 100644 --- 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 @@ -330,7 +330,7 @@ if [[ -n "${WANDB_API_KEY:-}" ]]; then --wandb-project "${WANDB_PROJECT:-slime-claude-code-ags}" --wandb-group "${WANDB_GROUP:-${EXP_TAG}}" --wandb-key "${WANDB_API_KEY}" - --wandb-dir "${LOG_DIR}/wandb" + --wandb-dir "${LOG_DIR}" --disable-wandb-random-suffix ) else 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 index 66260e9686..f6b29b4d7e 100644 --- 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 @@ -330,7 +330,7 @@ if [[ -n "${WANDB_API_KEY:-}" ]]; then --wandb-project "${WANDB_PROJECT:-slime-claude-code-ags}" --wandb-group "${WANDB_GROUP:-${EXP_TAG}}" --wandb-key "${WANDB_API_KEY}" - --wandb-dir "${LOG_DIR}/wandb" + --wandb-dir "${LOG_DIR}" --disable-wandb-random-suffix ) else From fc9348fc6583f2f3a12af2074e40b6e8e19174e0 Mon Sep 17 00:00:00 2001 From: FunJim Date: Thu, 6 Aug 2026 14:50:23 +0800 Subject: [PATCH 41/43] Run the AGS generator tests in CI test_ags_generator.py was never registered, so its 45 cases only ran locally. Registering it needed three fixes, because the CI runner invokes `python ` rather than pytest: * Add the `__main__` guard the runner relies on. Without it the file is merely imported and exits 0 -- a green check that runs nothing. test_ags_prompt_source.py had the same gap while already being registered, so it has been silently passing; fixed too. * Insert REPO_ROOT into sys.path before `from tests.test_agent._fakes`, matching what every other agent test already does. Bare `python ` does not put the repo root on the path. * Install openai / transformers / wandb / fastapi / uvicorn for the agent-test job. The test imports through the real slime_plugins package, whose __init__ chain needs them; all five are in requirements.txt but none are in the job's hardcoded install line. Uses the template's per-job extra_pip_deps hook so the other CPU job's install is untouched. Verified in a clean venv built from exactly the generated dep list: all seven tests in agent-test pass, 45 of them from this file. --- .github/workflows/pr-test.yml | 6 ++++++ .github/workflows/pr-test.yml.j2 | 2 ++ tests/test_rollout_buffer/test_ags_generator.py | 12 +++++++++++- tests/test_rollout_buffer/test_ags_prompt_source.py | 4 ++++ 4 files changed, 23 insertions(+), 1 deletion(-) diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 0d662fb3f5..185017282e 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -752,6 +752,10 @@ jobs: { "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: @@ -782,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 e4f2064e48..d1578f90bb 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -97,6 +97,7 @@ '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}, @@ -104,6 +105,7 @@ {'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/tests/test_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py index 9b5b83ea1c..882a1282bf 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -7,11 +7,17 @@ import re import sys import types +from pathlib import Path from types import SimpleNamespace import pytest from aiohttp import web -from tests.test_agent._fakes import FakeSandbox + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from tests.test_agent._fakes import FakeSandbox # noqa: E402 from slime.utils.misc import SingletonMeta from slime.utils.types import Sample @@ -1145,3 +1151,7 @@ async def run_case(): 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 index 73bbb0431c..b9b59e0bb0 100644 --- a/tests/test_rollout_buffer/test_ags_prompt_source.py +++ b/tests/test_rollout_buffer/test_ags_prompt_source.py @@ -150,3 +150,7 @@ def test_seek_is_noop_at_group_zero(prompt_data): 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"])) From 76fdf32b4912a66d4864af66a695428faa5a93ee Mon Sep 17 00:00:00 2001 From: FunJim Date: Thu, 6 Aug 2026 15:01:02 +0800 Subject: [PATCH 42/43] Shim asyncio.timeout in the AGS generator test for Python 3.10 CI pins 3.10, where asyncio.timeout() does not exist -- it is 3.11+. The two generate() plumbing cases hit rollout.py's `async with asyncio.timeout(...)`, whose AttributeError is then swallowed by the rollout's broad except into an abort, so the harness never runs and the tests fail on a missing capture rather than on the real cause. Shim it onto a pass-through context manager, mirroring test_agent/test_agent_rollout_cpu.py, which already does this for the same reason. The wall-clock guard never fires in these tests. Verified on a 3.10.20 venv with exactly the generated CI dep list: all seven tests in agent-test pass. The earlier local verification used 3.14, which is why this only surfaced on CI. --- tests/test_rollout_buffer/test_ags_generator.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_rollout_buffer/test_ags_generator.py b/tests/test_rollout_buffer/test_ags_generator.py index 882a1282bf..9135acfc5d 100644 --- a/tests/test_rollout_buffer/test_ags_generator.py +++ b/tests/test_rollout_buffer/test_ags_generator.py @@ -17,6 +17,19 @@ 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 19079d6d5f96bef67cb9c1751d7a41c8af8ac270 Mon Sep 17 00:00:00 2001 From: FunJim Date: Tue, 11 Aug 2026 17:19:52 +0800 Subject: [PATCH 43/43] Align the DP schedule by repacking when no bin can be split A training step whose micro-batch count is not a multiple of dp_size was aligned only by splitting multi-sample bins. That is impossible in long-context runs: once a single sample fills max_tokens_per_gpu * cp_size, first-fit emits one bin per sample and there is nothing left to split, so an odd sample count raised and killed the run. This happened after the rollout phase had fully succeeded -- 48 trajectories graded, patches applied cleanly, abort rate 0 -- so a rollout's worth of sandbox compute was discarded, and no AGS metric warned of it. The sample count is unpredictable because auto-compaction lets one rollout emit several training samples (48 rollouts produced 62 to 77 samples), which makes per-step parity a coin flip. Round the count down to the nearest multiple instead, keeping every sample. Rounding down is only possible with at least align_to bins; below that the floor is 0 micro-batches, so raise there with an error naming the real cause rather than blaming the repack. Repack with longest-processing-time-first so the largest resulting bin stays as small as possible, since that bin sets peak activation memory. Merging the two smallest bins repeatedly is the intuitive rule but degrades once more than one merge is needed: [49000, 49000, 90000, 90000, 90000] merged to 3 bins peaks at 180000 tokens versus 139000, i.e. 1.88x against 1.45x a 96000 cap. The two agree at exactly one merge, which is the dp_size=2 case. Merged bins can exceed the token cap, unavoidably, because first-fit is already a maximal packing; max_tokens_per_gpu becomes a target on this path and the repack is logged. Lowering it cannot relieve the overflow -- a merged micro-batch holds whole samples, so its token count is set by the data -- so the warning points at rollout_max_context_len and the batch shape instead. Tests use the token lengths and rollout ids recovered from the run that crashed, and cover the peak-memory property and the below-threshold error. --- slime/utils/dp_schedule.py | 88 ++++++++++-- slime/utils/seqlen_balancing.py | 42 ++++++ tests/test_dp_schedule.py | 247 +++++++++++++++++++++++++++++++- 3 files changed, 365 insertions(+), 12 deletions(-) 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/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/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__]))