diff --git a/.gitignore b/.gitignore index dc8f74212f..691e005454 100644 --- a/.gitignore +++ b/.gitignore @@ -201,5 +201,7 @@ api_key.json .vscode/ wandb/ outputs/ +rl_logs/ +.arena-rollout.env sympy/ !/docs/figures/* diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 50d05aa781..f079297b97 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -23,6 +23,7 @@ is_valid_attn_impl, ) from areal.utils import logging, name_resolve, pkg_version +from areal.utils.config_utils import redact_sensitive_config from areal.utils.constants import ( PROX_LOGP_METHOD_RECOMPUTE, PROX_LOGP_METHODS_ALL, @@ -3281,7 +3282,7 @@ def save_config(cfg, log_dir): os.makedirs(log_dir, exist_ok=True) config_save_path = os.path.join(log_dir, "config.yaml") with open(config_save_path, "w") as f: - config_dict: dict = asdict(cfg) + config_dict: dict = redact_sensitive_config(asdict(cfg)) yaml.dump( config_dict, f, diff --git a/areal/experimental/openai/proxy/proxy_rollout_server.py b/areal/experimental/openai/proxy/proxy_rollout_server.py index 5d5453d92d..6dc1a699c0 100644 --- a/areal/experimental/openai/proxy/proxy_rollout_server.py +++ b/areal/experimental/openai/proxy/proxy_rollout_server.py @@ -6,11 +6,12 @@ import asyncio import hmac import inspect +import json import os import secrets import threading import time -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Iterable, Mapping from typing import TYPE_CHECKING, Any import uvicorn @@ -554,6 +555,91 @@ def set_reward( # ============================================================================= +def _contains_image_content(value: Any) -> bool: + """Return whether a nested content value contains an image block.""" + if isinstance(value, (list, tuple)): + return any(_contains_image_content(item) for item in value) + if not isinstance(value, dict): + return False + if value.get("type") in {"image", "image_url", "input_image"}: + return True + return any(_contains_image_content(item) for item in value.values()) + + +def _content_block_text(value: Any) -> str: + """Extract text from nested Claude/OpenAI content blocks.""" + if isinstance(value, str): + return value + if isinstance(value, (list, tuple)): + return "\n".join(filter(None, (_content_block_text(item) for item in value))) + if not isinstance(value, dict): + return str(value) + + text = value.get("text") + if isinstance(text, str): + return text + content = value.get("content") + if isinstance(content, (str, list, tuple, dict)): + return _content_block_text(content) + return json.dumps(value, ensure_ascii=False, sort_keys=True) + + +def _flatten_text_content_lists(messages: list[dict]) -> None: + """Flatten text-only content block lists to strings in-place. + + LiteLLM can forward Claude requests through the OpenAI chat-completions + endpoint while retaining Anthropic-style text block lists. Preserve + multimodal or otherwise structured content for the AReaL client. + """ + for msg in messages: + content = msg.get("content") + if not isinstance(content, list): + continue + + if msg.get("role") == "system" or not _contains_image_content(content): + msg["content"] = _content_block_text(content) + + +def _preprocess_messages(messages: list[dict]) -> list[dict]: + """Normalize messages shared by OpenAI and Anthropic proxy endpoints.""" + _flatten_text_content_lists(messages) + for preprocessor in _message_preprocessors: + messages = preprocessor(messages) + return messages + + +def _materialize_iterables(value: Any) -> Any: + """Recursively convert Pydantic validator iterators to plain containers.""" + if isinstance(value, BaseModel): + return _materialize_iterables(value.model_dump()) + if isinstance(value, Mapping): + return {key: _materialize_iterables(item) for key, item in value.items()} + if isinstance(value, (str, bytes)): + return value + if isinstance(value, Iterable): + return [_materialize_iterables(item) for item in value] + return value + + +def _prepare_request_messages(messages: Any) -> Any: + """Convert request message iterables to mutable dicts and preprocess them.""" + if isinstance(messages, (str, bytes, Mapping)) or not isinstance( + messages, Iterable + ): + return messages + + normalized: list[dict] = [] + for message in messages: + if isinstance(message, BaseModel): + message = message.model_dump() + elif isinstance(message, Mapping): + message = dict(message) + else: + return messages + normalized.append(_materialize_iterables(message)) + return _preprocess_messages(normalized) + + async def _call_client_create( create_fn, request: dict[str, Any] | BaseModel, @@ -587,6 +673,19 @@ async def _call_client_create( ) kwargs = request.model_dump() if isinstance(request, BaseModel) else dict(request) + messages = kwargs.get("messages") + if messages is not None: + prepared_messages = _prepare_request_messages(messages) + kwargs["messages"] = prepared_messages + if isinstance(prepared_messages, list): + logger.debug( + "Preprocessed request messages: container=%s, content_types=%s", + type(messages).__name__, + [ + type(message.get("content")).__name__ + for message in prepared_messages + ], + ) dropped_args = [] for k, v in kwargs.items(): if k not in areal_client_allowed_args: @@ -634,6 +733,7 @@ def _is_default_value(k: str, v: Any) -> bool: except ValueError as e: raise HTTPException(status_code=500, detail=str(e)) except Exception as e: + logger.exception("AReaL client request failed") raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}") @@ -722,19 +822,6 @@ async def responses( ) -def _flatten_content_lists(messages: list[dict]) -> None: - """Flatten Anthropic content block lists to strings in-place.""" - for msg in messages: - if isinstance(msg.get("content"), list): - text_parts = [] - for block in msg["content"]: - if isinstance(block, dict) and block.get("type") == "text": - text_parts.append(block.get("text", "")) - elif isinstance(block, str): - text_parts.append(block) - msg["content"] = "\n".join(text_parts) - - def _translate_anthropic_to_openai_request(anthropic_request: dict[str, Any]) -> dict: """Translate an Anthropic Messages API request to OpenAI format.""" openai_request = _adapter.translate_completion_input_params( @@ -744,11 +831,6 @@ def _translate_anthropic_to_openai_request(anthropic_request: dict[str, Any]) -> raise ValueError("Failed to translate request") openai_request = dict(openai_request) - if "messages" in openai_request: - _flatten_content_lists(openai_request["messages"]) - for preprocessor in _message_preprocessors: - openai_request["messages"] = preprocessor(openai_request["messages"]) - return openai_request diff --git a/areal/infra/controller/rollout_controller.py b/areal/infra/controller/rollout_controller.py index b43f2a7609..63d5745f64 100644 --- a/areal/infra/controller/rollout_controller.py +++ b/areal/infra/controller/rollout_controller.py @@ -917,6 +917,13 @@ def wait( return [r.trajectory if r is not None else None for r in results] + def wait_for_task( + self, task_id: int, timeout: float | None = None, raise_timeout: bool = True + ) -> dict[str, Any] | None: + """Wait for one submitted rollout while preserving its task identity.""" + result = self.dispatcher.wait_for_task(task_id, timeout, raise_timeout) + return result.trajectory if result is not None else None + @trace_perf("rollout_controller.rollout_batch", category="scheduler") def rollout_batch( self, diff --git a/areal/utils/config_utils.py b/areal/utils/config_utils.py new file mode 100644 index 0000000000..abdc3f87b0 --- /dev/null +++ b/areal/utils/config_utils.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Utilities for safely serializing application configuration.""" + +from __future__ import annotations + +from typing import Any + +REDACTED_VALUE = "" + + +def _is_sensitive_key(key: object) -> bool: + if not isinstance(key, str): + return False + normalized = key.lower() + return ( + normalized in {"authorization", "password", "secret", "token"} + or "api_key" in normalized + or normalized.endswith(("_password", "_secret", "_token", "_credential")) + or "private_key" in normalized + ) + + +def redact_sensitive_config(value: Any) -> Any: + """Return a copy with credentials redacted while preserving ordinary fields.""" + if isinstance(value, dict): + return { + key: REDACTED_VALUE + if _is_sensitive_key(key) and item not in (None, "") + else redact_sensitive_config(item) + for key, item in value.items() + } + if isinstance(value, list): + return [redact_sensitive_config(item) for item in value] + if isinstance(value, tuple): + return tuple(redact_sensitive_config(item) for item in value) + return value diff --git a/areal/utils/logging.py b/areal/utils/logging.py index a46e114375..c34e177fc2 100644 --- a/areal/utils/logging.py +++ b/areal/utils/logging.py @@ -112,6 +112,7 @@ "TokenLogpReward": "light_purple", "ProxyUtils": "light_purple", "AReaL-SWEAgent": "light_purple", + "ArenaStreamAgent": "light_purple", "SWETrain": "light_green", # Agent Service - purple "AgentGateway": "light_purple", diff --git a/areal/utils/stats_logger.py b/areal/utils/stats_logger.py index 220e6d3920..6cf5a5f2de 100644 --- a/areal/utils/stats_logger.py +++ b/areal/utils/stats_logger.py @@ -7,18 +7,23 @@ import swanlab import torch.distributed as dist -import trackio import wandb from tensorboardX import SummaryWriter from areal.api import FinetuneSpec from areal.api.cli_args import BaseExperimentConfig, StatsLoggerConfig from areal.utils import logging +from areal.utils.config_utils import redact_sensitive_config from areal.utils.printing import tabulate_stats from areal.version import version_info logger = logging.getLogger("StatsLogger", "system") +# Trackio is optional and older runtime images may not include it. Keep this +# module-level name so tests and callers can patch it, but import it only when +# the backend is enabled. +trackio = None + class StatsLogger: def __init__(self, config: BaseExperimentConfig, ft_spec: FinetuneSpec): @@ -52,7 +57,7 @@ def init(self): if suffix == "timestamp": suffix = time.strftime("%Y_%m_%d_%H_%M_%S") - exp_config_dict = asdict(self.exp_config) + exp_config_dict = redact_sensitive_config(asdict(self.exp_config)) exp_config_dict["version_info"] = { "commit_id": version_info.commit, "branch": version_info.branch, @@ -98,6 +103,11 @@ def init(self): self._trackio_enabled = False trackio_config = self.config.trackio if trackio_config.mode != "disabled": + global trackio + if trackio is None: + import trackio as trackio_module + + trackio = trackio_module trackio.init( project=trackio_config.project or self.config.experiment_name, name=trackio_config.name or self.config.trial_name, diff --git a/examples/swe/README.md b/examples/swe/README.md index ca3985ec40..7ee5a54771 100644 --- a/examples/swe/README.md +++ b/examples/swe/README.md @@ -1,5 +1,44 @@ # SWE-bench RL training with AReaL-SWEAgent +## Arena Stream mode + +The example config can also load data ids from the Arena online Stream OpenAPI and +delegate each rollout to `launch_one_task`, without installing AReaL-SWEAgent. Set: + +```bash +export ARENA_OPENAPI_BASE=https://your-arena-service +export ARENA_OPENAPI_TOKEN=your-token +export ARENA_LLM_API_KEY=your-llm-gateway-key +export SWE_RL_ADMIN_API_KEY=your-rollout-admin-key +``` + +```yaml +econfig: + dataset_source: arena + stream_id: "" # empty selects the first active Stream + arena_base_url: ${oc.env:ARENA_OPENAPI_BASE} + arena_task_envs: + CLAUDE_CODE_DISABLE_TERMINAL_TITLE: "1" +``` + +The trainer first requests one dataset row to discover `total`, then requests +`limit=total` and constructs the training dataset from the returned `data_ids`. The +initial implementation intentionally rejects Streams with more than the API's +single-request limit of 1000 rows; it does not paginate or shard the dataset. + +For every rollout, `ArenaStreamAgentWorkflow` registers the current AReaL proxy through +`/openapi/v1/llm/models` as an OpenAI Chat Completions upstream, then posts the returned +model alias and the row's `data_id` to `launch_one_task`. Arena converts the Harness's +native protocol to Chat before forwarding requests to AReaL. The task environment +receives the managed `MODEL_NAME`, `BASE_URL`, and `API_KEY` variables plus values +configured in `econfig.arena_task_envs`; the workflow polls the returned task id until +it is terminal and returns its numeric score as the reward. Model registrations are +deleted in a `finally` block. Credentials are read only from environment variables and +are not stored in this repository or experiment configs. + +The sections below document the original external AReaL-SWEAgent mode, selected with +`econfig.dataset_source=jsonl`. + This example runs SWE-bench coding-agent RL (GRPO) in AReaL. The actual agent loop, sandboxing and reward computation live in a **separate repository**, [AReaL-SWEAgent](https://github.com/areal-project/AReaL-SWEAgent): for each SWE-bench diff --git a/examples/swe/arena_agent.py b/examples/swe/arena_agent.py new file mode 100644 index 0000000000..0226233642 --- /dev/null +++ b/examples/swe/arena_agent.py @@ -0,0 +1,125 @@ +"""Arena Stream agent workflow using AReaL's OpenAI-compatible proxy.""" + +from __future__ import annotations + +import asyncio +import uuid +from typing import Any + +import httpx + +from examples.swe.arena_client import ArenaOpenAPIClient, LLMProtocol + +from areal.utils import logging + +logger = logging.getLogger("ArenaStreamAgent") + + +class ArenaStreamAgentWorkflow: + """Launch an Arena online task and use its returned reward for RL.""" + + def __init__( + self, + econfig: dict[str, Any] | None = None, + gen_args: dict[str, Any] | None = None, + timeout: float = 3600.0, + ) -> None: + self.econfig = econfig or {} + self.gen_args = gen_args or {} + self.timeout = float(self.econfig.get("timeout", timeout)) + self.registration_timeout = float( + self.econfig.get("arena_registration_timeout", 180.0) + ) + self.task_envs = dict(self.econfig.get("arena_task_envs") or {}) + self.client = ArenaOpenAPIClient( + base_url=str(self.econfig.get("arena_base_url", "")), + timeout=self.timeout, + poll_interval=float(self.econfig.get("arena_poll_interval", 5.0)), + request_retries=int(self.econfig.get("arena_request_retries", 3)), + ) + + async def run( + self, + data: dict[str, Any], + **extra_kwargs: Any, + ) -> float: + """Launch the row's task with the current rollout proxy session.""" + stream_id = str(data.get("stream_id") or self.econfig.get("stream_id") or "") + data_id = str(data.get("data_id") or "") + llm_protocol: LLMProtocol = data.get("llm_protocol", "chat_completions") + proxy_base_url = extra_kwargs.get("base_url") + proxy_api_key = extra_kwargs.get("api_key") + arena_http_client: httpx.AsyncClient | None = extra_kwargs.get( + "arena_http_client" + ) or extra_kwargs.get("http_client") + + if not stream_id: + raise ValueError("stream_id is required for ArenaStreamAgentWorkflow") + if not data_id: + raise ValueError("data_id is required for ArenaStreamAgentWorkflow") + if not proxy_base_url: + raise ValueError("base_url is required for ArenaStreamAgentWorkflow") + if not proxy_api_key: + raise ValueError("api_key is required for ArenaStreamAgentWorkflow") + + suffix = uuid.uuid4().hex[:12] + model_name = f"stream-areal-{suffix}" + deployment_id = str(uuid.uuid4()) + registered_model_id = deployment_id + owns_client = arena_http_client is None + client = arena_http_client or httpx.AsyncClient(timeout=self.timeout) + try: + ( + registered_url, + registered_model_id, + ) = await self.client.register_llm_proxy_async( + model_name=model_name, + upstream_base_url=str(proxy_base_url), + upstream_api_key=str(proxy_api_key), + deployment_id=deployment_id, + protocol=llm_protocol, + client=client, + timeout=self.registration_timeout, + ) + logger.info( + "Registered Arena LLM proxy: model_name=%s, model_id=%s, " + "registered_url=%s, protocol=%s", + model_name, + registered_model_id, + registered_url, + llm_protocol, + ) + logger.info( + f"Launching Arena task: stream_id={stream_id}, data_id={data_id}" + ) + reward = await asyncio.wait_for( + self.client.launch_one_task( + stream_id=stream_id, + data_id=data_id, + model_name=registered_model_id, + proxy_base_url=registered_url, + proxy_api_key=self.client.llm_api_key, + task_envs=self.task_envs, + client=client, + ), + timeout=self.timeout, + ) + finally: + try: + await self.client.delete_llm_proxy_async( + registered_model_id, + client=client, + timeout=self.registration_timeout, + ) + logger.info( + "Deleted Arena LLM proxy registration: model_id=%s", + registered_model_id, + ) + finally: + if owns_client: + await client.aclose() + logger.info( + f"Finished Arena task: stream_id={stream_id}, data_id={data_id}, " + f"reward={reward}" + ) + return reward diff --git a/examples/swe/arena_client.py b/examples/swe/arena_client.py new file mode 100644 index 0000000000..a367c8270f --- /dev/null +++ b/examples/swe/arena_client.py @@ -0,0 +1,568 @@ +"""Client helpers for the Arena online Stream OpenAPI.""" + +from __future__ import annotations + +import asyncio +import os +import time +import uuid +from collections.abc import Mapping +from numbers import Real +from typing import Any, Literal +from urllib.parse import quote + +import httpx + +LLMProtocol = Literal["anthropic", "responses", "chat_completions"] + + +class ArenaAPIError(RuntimeError): + """Raised when the Arena OpenAPI returns an invalid or failed response.""" + + +def _is_retryable_status(status_code: int) -> bool: + return status_code == 429 or status_code >= 500 + + +def resolve_arena_credentials( + base_url: str = "", + api_token: str = "", +) -> tuple[str, str]: + """Resolve Arena connection settings without embedding credentials in configs.""" + resolved_base_url = base_url or os.getenv("ARENA_OPENAPI_BASE", "") + resolved_api_token = api_token or os.getenv("ARENA_OPENAPI_TOKEN", "") + if not resolved_base_url: + raise ValueError( + "Arena OpenAPI base URL is required; set econfig.arena_base_url or " + "ARENA_OPENAPI_BASE" + ) + if not resolved_api_token: + raise ValueError("ARENA_OPENAPI_TOKEN is required") + return resolved_base_url.rstrip("/"), resolved_api_token + + +def _response_json(response: httpx.Response) -> Any: + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + body = response.text[:500] + raise ArenaAPIError( + f"Arena OpenAPI returned HTTP {response.status_code}: {body}" + ) from exc + try: + return response.json() + except ValueError as exc: + raise ArenaAPIError("Arena OpenAPI returned a non-JSON response") from exc + + +def _extract_reward(value: Any) -> float | None: + """Extract a numeric reward from common response envelope shapes.""" + if isinstance(value, Real) and not isinstance(value, bool): + return float(value) + if not isinstance(value, Mapping): + return None + + for key in ("reward", "score"): + reward = value.get(key) + if isinstance(reward, Real) and not isinstance(reward, bool): + return float(reward) + + for key in ("result", "output", "task", "data"): + nested_reward = _extract_reward(value.get(key)) + if nested_reward is not None: + return nested_reward + return None + + +def infer_llm_protocol(stream: Mapping[str, Any]) -> LLMProtocol: + """Select the client protocol used by a Stream's default Harness.""" + harness_ref = stream.get("default_harness_ref") + harness_key = "" + if isinstance(harness_ref, Mapping): + key = harness_ref.get("key") + if isinstance(key, str): + harness_key = key.lower() + + if "claude" in harness_key: + return "anthropic" + if "codex" in harness_key: + return "responses" + return "chat_completions" + + +def _llm_registration_payload( + model_name: str, + upstream_base_url: str, + upstream_api_key: str, + deployment_id: str, + protocol: LLMProtocol, +) -> dict[str, Any]: + if protocol not in ("anthropic", "responses", "chat_completions"): + raise ValueError(f"Unsupported Arena LLM protocol: {protocol!r}") + + # ``protocol`` describes the Arena Harness client. The AReaL rollout + # proxy itself exposes an OpenAI-compatible Chat Completions endpoint, so + # advertise only that native upstream capability. Arena converts Messages + # or Responses requests to Chat before forwarding them to this endpoint. + return { + "model_name": model_name, + "endpoints": [ + { + "endpoint_id": deployment_id, + "upstream_model": model_name, + "base_url": upstream_base_url.rstrip("/"), + "api_key": upstream_api_key, + "inbound_protos": ["chat"], + "enabled": True, + } + ], + "enabled": True, + "metadata": {"deployment_id": deployment_id}, + } + + +class ArenaOpenAPIClient: + """Small sync/async client for Stream discovery, datasets, and task launch.""" + + MAX_DATASET_LIMIT = 1000 + FAILED_TASK_STATUSES = { + "CANCELLED", + "COLLECT_FAILED", + "EVAL_FAILED", + "FAILED", + "HARNESS_FAILED", + "NO_OUTPUT", + "SETUP_FAILED", + "TIMEOUT", + } + + def __init__( + self, + base_url: str = "", + api_token: str = "", + llm_api_key: str = "", + timeout: float = 60.0, + poll_interval: float = 5.0, + request_retries: int = 3, + ) -> None: + self.base_url, self.api_token = resolve_arena_credentials( + base_url=base_url, + api_token=api_token, + ) + self.llm_api_key = llm_api_key or os.getenv("ARENA_LLM_API_KEY", "") + self.timeout = timeout + self.poll_interval = poll_interval + self.request_retries = request_retries + + @property + def _headers(self) -> dict[str, str]: + return {"Authorization": f"Bearer {self.api_token}"} + + @property + def _llm_headers(self) -> dict[str, str]: + if not self.llm_api_key: + raise ValueError("ARENA_LLM_API_KEY is required for LLM gateway traffic") + return {"Authorization": f"Bearer {self.llm_api_key}"} + + def list_streams( + self, + status: str | None = "ACTIVE", + *, + client: httpx.Client | None = None, + ) -> list[dict[str, Any]]: + """Return online Streams, optionally filtered by status.""" + params = {"status": status} if status else None + response = self._sync_request( + "GET", + f"{self.base_url}/openapi/v1/streams", + client=client, + params=params, + headers=self._headers, + ) + payload = _response_json(response) + items = payload.get("items") if isinstance(payload, Mapping) else None + if not isinstance(items, list): + raise ArenaAPIError("Stream list response is missing an 'items' array") + return [dict(item) for item in items if isinstance(item, Mapping)] + + def resolve_stream( + self, + stream_id: str = "", + *, + client: httpx.Client | None = None, + ) -> dict[str, Any]: + """Resolve a Stream and retain metadata needed by the rollout Harness.""" + streams = self.list_streams(client=client) + if not streams: + raise ArenaAPIError("Arena OpenAPI returned no active Streams") + + if stream_id: + for stream in streams: + if stream.get("stream_id") == stream_id: + return stream + raise ArenaAPIError(f"Active Stream {stream_id!r} was not found") + return streams[0] + + def resolve_stream_id( + self, + stream_id: str = "", + *, + client: httpx.Client | None = None, + ) -> str: + """Use an explicit Stream id or fall back to the first active Stream.""" + stream = self.resolve_stream(stream_id, client=client) + selected = stream.get("stream_id") + if not isinstance(selected, str) or not selected: + raise ArenaAPIError("The selected Stream is missing 'stream_id'") + return selected + + def _select_dataset_page( + self, + stream_id: str, + limit: int, + offset: int = 0, + *, + client: httpx.Client | None = None, + ) -> dict[str, Any]: + encoded_stream_id = quote(stream_id, safe="") + url = f"{self.base_url}/openapi/v1/streams/{encoded_stream_id}/dataset" + kwargs = { + "params": {"limit": limit, "offset": offset}, + "headers": self._headers, + } + response = self._sync_request("POST", url, client=client, **kwargs) + payload = _response_json(response) + if not isinstance(payload, Mapping): + raise ArenaAPIError("Dataset response must be a JSON object") + return dict(payload) + + def get_all_dataset_rows( + self, + stream_id: str, + llm_protocol: LLMProtocol = "chat_completions", + *, + client: httpx.Client | None = None, + ) -> list[dict[str, str]]: + """Load one Stream's complete dataset in one page after a size probe.""" + first_page = self._select_dataset_page( + stream_id, + limit=1, + client=client, + ) + total = first_page.get("total") + if not isinstance(total, int) or isinstance(total, bool) or total < 0: + raise ArenaAPIError("Dataset response has an invalid 'total'") + if total == 0: + raise ArenaAPIError(f"Stream {stream_id!r} contains no dataset rows") + if total > self.MAX_DATASET_LIMIT: + raise ArenaAPIError( + f"Stream {stream_id!r} contains {total} rows, exceeding the " + f"single-request limit {self.MAX_DATASET_LIMIT}; pagination is not " + "implemented yet" + ) + + page = ( + first_page + if total == 1 + else self._select_dataset_page( + stream_id, + limit=total, + client=client, + ) + ) + data_ids = page.get("data_ids") + if not isinstance(data_ids, list) or not all( + isinstance(data_id, str) and data_id for data_id in data_ids + ): + raise ArenaAPIError("Dataset response has an invalid 'data_ids' array") + if len(data_ids) != total: + raise ArenaAPIError( + f"Dataset response returned {len(data_ids)} rows, expected {total}" + ) + return [ + { + "data_id": data_id, + "stream_id": stream_id, + "llm_protocol": llm_protocol, + } + for data_id in data_ids + ] + + def register_llm_proxy( + self, + model_name: str, + upstream_base_url: str, + upstream_api_key: str, + *, + deployment_id: str | None = None, + protocol: LLMProtocol = "chat_completions", + client: httpx.Client | None = None, + ) -> tuple[str, str]: + """Register an AReaL OpenAI proxy and return its URL and model alias.""" + if not model_name.startswith("stream-areal-"): + raise ValueError("Arena model_name must start with 'stream-areal-'") + if not upstream_base_url: + raise ValueError("upstream_base_url is required") + if not upstream_api_key: + raise ValueError("upstream_api_key is required") + + resolved_deployment_id = deployment_id or str(uuid.uuid4()) + payload = _llm_registration_payload( + model_name=model_name, + upstream_base_url=upstream_base_url, + upstream_api_key=upstream_api_key, + deployment_id=resolved_deployment_id, + protocol=protocol, + ) + response = self._sync_request( + "POST", + f"{self.base_url}/openapi/v1/llm/models", + client=client, + headers=self._headers, + json=payload, + ) + registration = _response_json(response) + registered_url, registered_model_id = self._registered_llm_target( + registration, + model_name, + ) + return registered_url, registered_model_id + + def _registered_llm_target( + self, + registration: Any, + model_name: str, + ) -> tuple[str, str]: + """Resolve the LLM gateway API base and model alias from registration.""" + if not isinstance(registration, Mapping): + raise ArenaAPIError("LLM registration response must be a JSON object") + registered_model_id = registration.get("model_name") + if not isinstance(registered_model_id, str) or not registered_model_id: + raise ArenaAPIError("LLM registration response is missing model_name") + if registered_model_id != model_name: + raise ArenaAPIError( + "LLM registration response returned an unexpected model_name: " + f"{registered_model_id!r}" + ) + return f"{self.base_url}/api", registered_model_id + + async def register_llm_proxy_async( + self, + model_name: str, + upstream_base_url: str, + upstream_api_key: str, + deployment_id: str, + *, + protocol: LLMProtocol = "chat_completions", + client: httpx.AsyncClient, + timeout: float = 180.0, + ) -> tuple[str, str]: + """Asynchronously register one rollout proxy session.""" + if not model_name.startswith("stream-areal-"): + raise ValueError("Arena model_name must start with 'stream-areal-'") + if not upstream_base_url: + raise ValueError("upstream_base_url is required") + if not upstream_api_key: + raise ValueError("upstream_api_key is required") + payload = _llm_registration_payload( + model_name=model_name, + upstream_base_url=upstream_base_url, + upstream_api_key=upstream_api_key, + deployment_id=deployment_id, + protocol=protocol, + ) + response = await self._async_request( + client, + "POST", + f"{self.base_url}/openapi/v1/llm/models", + headers=self._headers, + json=payload, + timeout=timeout, + ) + registration = _response_json(response) + return self._registered_llm_target(registration, model_name) + + def delete_llm_proxy( + self, + model_name: str, + *, + client: httpx.Client | None = None, + ) -> None: + """Delete one registered LLM model; missing models are clean.""" + response = self._sync_request( + "DELETE", + f"{self.base_url}/openapi/v1/llm/models/{quote(model_name, safe='')}", + client=client, + headers=self._headers, + ) + if response.status_code == 404: + return + if 200 <= response.status_code < 300: + return + _response_json(response) + + async def delete_llm_proxy_async( + self, + model_name: str, + *, + client: httpx.AsyncClient, + timeout: float = 180.0, + ) -> None: + """Asynchronously delete one registered LLM model.""" + response = await self._async_request( + client, + "DELETE", + f"{self.base_url}/openapi/v1/llm/models/{quote(model_name, safe='')}", + headers=self._headers, + timeout=timeout, + ) + if response.status_code == 404: + return + if 200 <= response.status_code < 300: + return + _response_json(response) + + def _sync_request( + self, + method: str, + url: str, + *, + client: httpx.Client | None = None, + **kwargs: Any, + ) -> httpx.Response: + for attempt in range(self.request_retries + 1): + try: + if client is not None: + response = client.request(method, url, **kwargs) + else: + with httpx.Client(timeout=self.timeout) as owned_client: + response = owned_client.request(method, url, **kwargs) + except httpx.RequestError as exc: + if attempt == self.request_retries: + raise ArenaAPIError( + f"Arena OpenAPI request failed after {attempt + 1} attempts: " + f"{type(exc).__name__}" + ) from exc + time.sleep(min(2**attempt, 10)) + continue + if not _is_retryable_status(response.status_code): + return response + if attempt == self.request_retries: + return response + response.close() + time.sleep(min(2**attempt, 10)) + raise AssertionError("unreachable") + + async def _async_request( + self, + client: httpx.AsyncClient, + method: str, + url: str, + **kwargs: Any, + ) -> httpx.Response: + for attempt in range(self.request_retries + 1): + try: + response = await client.request(method, url, **kwargs) + except httpx.RequestError as exc: + if attempt == self.request_retries: + raise ArenaAPIError( + f"Arena OpenAPI request failed after {attempt + 1} attempts: " + f"{type(exc).__name__}" + ) from exc + await asyncio.sleep(min(2**attempt, 10)) + continue + if not _is_retryable_status(response.status_code): + return response + if attempt == self.request_retries: + return response + await response.aclose() + await asyncio.sleep(min(2**attempt, 10)) + raise AssertionError("unreachable") + + async def launch_one_task( + self, + stream_id: str, + data_id: str, + model_name: str, + proxy_base_url: str, + proxy_api_key: str, + *, + task_envs: Mapping[str, str] | None = None, + client: httpx.AsyncClient | None = None, + ) -> float: + """Launch one Arena task through the rollout proxy and return its reward.""" + encoded_stream_id = quote(stream_id, safe="") + url = f"{self.base_url}/openapi/v1/streams/{encoded_stream_id}/launch_one_task" + envs = dict(task_envs or {}) + if not all( + isinstance(key, str) and key and isinstance(value, str) + for key, value in envs.items() + ): + raise ValueError("Arena task environment variables must be strings") + envs.update( + { + "MODEL_NAME": model_name, + "BASE_URL": proxy_base_url, + "API_KEY": proxy_api_key, + } + ) + request = { + "data_id": data_id, + "model_name": model_name, + "base_url": proxy_base_url, + "api_key": proxy_api_key, + "envs": envs, + } + if client is not None: + return await self._launch_and_wait(client, url, request) + async with httpx.AsyncClient(timeout=self.timeout) as owned_client: + return await self._launch_and_wait(owned_client, url, request) + + async def _launch_and_wait( + self, + client: httpx.AsyncClient, + url: str, + request: dict[str, Any], + ) -> float: + response = await self._async_request( + client, + "POST", + url, + json=request, + headers=self._headers, + ) + payload = _response_json(response) + immediate_reward = _extract_reward(payload) + if immediate_reward is not None: + return immediate_reward + task_id = payload.get("task_id") if isinstance(payload, Mapping) else None + if not isinstance(task_id, str) or not task_id: + raise ArenaAPIError( + "launch_one_task response contains neither a reward nor task_id" + ) + return await self._poll_task_result(client, task_id) + + async def _poll_task_result( + self, + client: httpx.AsyncClient, + task_id: str, + ) -> float: + encoded_task_id = quote(task_id, safe="") + url = f"{self.base_url}/openapi/v1/tasks/{encoded_task_id}/result" + while True: + response = await self._async_request( + client, + "GET", + url, + headers=self._headers, + ) + payload = _response_json(response) + if not isinstance(payload, Mapping): + raise ArenaAPIError("Task result response must be a JSON object") + status = str(payload.get("status") or "").upper() + if status in self.FAILED_TASK_STATUSES: + raise ArenaAPIError(f"Arena task {task_id!r} failed with {status}") + reward = _extract_reward(payload) + if reward is not None and status in {"DONE", "OK"}: + return reward + await asyncio.sleep(self.poll_interval) diff --git a/examples/swe/arena_rollout_only.py b/examples/swe/arena_rollout_only.py new file mode 100644 index 0000000000..021da5f472 --- /dev/null +++ b/examples/swe/arena_rollout_only.py @@ -0,0 +1,374 @@ +"""Run Arena Stream rollouts without constructing a training engine.""" + +from __future__ import annotations + +import argparse +import getpass +import os +import random +import sys +import time +import uuid +from dataclasses import asdict +from pathlib import Path +from typing import Any + +import httpx +import torch + +from examples.swe.arena_client import ArenaOpenAPIClient, infer_llm_protocol +from examples.swe.utils import SWEPPOConfig + +from areal.api.alloc_mode import ModelAllocation +from areal.api.cli_args import SGLangConfig, load_expr_config +from areal.engine import RemoteSGLangEngine +from areal.infra import LocalScheduler +from areal.infra.rpc.rtensor import RTensor +from areal.utils import logging +from areal.utils.config_utils import redact_sensitive_config + +logger = logging.getLogger("ArenaRolloutOnly") + + +def _trajectory_reward(trajectory: dict[str, Any]) -> float: + """Extract the episode reward from a localized proxy trajectory.""" + localized = RTensor.localize(trajectory) + rewards = localized.get("rewards") + if torch.is_tensor(rewards): + return float(rewards.sum().item()) + + interactions = localized.get("interactions") + if isinstance(interactions, list): + return float( + sum( + interaction.get("reward", 0.0) + for interaction in interactions + if isinstance(interaction, dict) + ) + ) + raise ValueError("Rollout trajectory contains neither rewards nor interactions") + + +def _run_rollout_tasks( + controller: Any, + rows: list[dict[str, str]], + workflow_kwargs: dict[str, Any], +) -> tuple[list[tuple[str, float]], list[str]]: + """Run queued rollouts while retaining row identity for rejected tasks.""" + submitted: list[tuple[dict[str, str], int]] = [] + failed_data_ids: list[str] = [] + for row in rows: + data_id = row["data_id"] + try: + task_id = controller.submit( + data=row, + workflow="examples.swe.arena_agent.ArenaStreamAgentWorkflow", + workflow_kwargs=workflow_kwargs, + group_size=1, + ) + except Exception: + logger.exception("Failed to submit Arena rollout: data_id=%s", data_id) + failed_data_ids.append(data_id) + continue + submitted.append((row, task_id)) + + completed: list[tuple[str, float]] = [] + for row, task_id in submitted: + data_id = row["data_id"] + try: + trajectory = controller.wait_for_task(task_id) + if trajectory is None: + logger.warning("Arena rollout rejected: data_id=%s", data_id) + failed_data_ids.append(data_id) + continue + reward = _trajectory_reward(trajectory) + except Exception: + logger.exception("Failed to collect Arena rollout: data_id=%s", data_id) + failed_data_ids.append(data_id) + continue + completed.append((data_id, reward)) + logger.info("Collected Arena rollout: data_id=%s, reward=%s", data_id, reward) + + return completed, failed_data_ids + + +def _init_wandb(config: SWEPPOConfig): + """Initialize the configured W&B run without starting trainer infrastructure.""" + import wandb + + wandb_config = config.stats_logger.wandb + log_dir = ( + Path(config.cluster.fileroot) + / "logs" + / getpass.getuser() + / config.experiment_name + / config.trial_name + ) + log_dir.mkdir(parents=True, exist_ok=True) + return wandb.init( + mode=wandb_config.mode, + entity=wandb_config.entity, + project=wandb_config.project or config.experiment_name, + name=wandb_config.name or config.trial_name, + job_type="rollout-only", + group=wandb_config.group or f"{config.experiment_name}_{config.trial_name}", + notes=wandb_config.notes, + tags=wandb_config.tags, + config=redact_sensitive_config(asdict(config)), + dir=str(log_dir), + id=f"{config.experiment_name}_{config.trial_name}_rollout", + resume="allow", + ) + + +def _run_registry_smoke( + arena_client: ArenaOpenAPIClient, + proxy_base_url: str, + proxy_admin_api_key: str, + trial_name: str, +) -> tuple[str, str]: + """Create a proxy session and exercise Arena LLM registration lifecycle.""" + suffix = uuid.uuid4().hex[:10] + model_prefix = "".join( + character if character.isalnum() else "-" for character in trial_name.lower() + ).strip("-")[-40:] + model_name = f"stream-areal-{model_prefix}-{suffix}" + deployment_id = str(uuid.uuid4()) + registered_model_id = deployment_id + admin_headers = {"Authorization": f"Bearer {proxy_admin_api_key}"} + + with httpx.Client(timeout=30.0) as proxy_client: + grant_response = proxy_client.post( + f"{proxy_base_url.rstrip('/')}/grant_capacity", + headers=admin_headers, + ) + grant_response.raise_for_status() + start_response = proxy_client.post( + f"{proxy_base_url.rstrip('/')}/rl/start_session", + headers=admin_headers, + json={"task_id": f"registry-smoke-{suffix}"}, + ) + start_response.raise_for_status() + session_api_key = start_response.json()["api_key"] + + try: + registered_url, registered_model_id = arena_client.register_llm_proxy( + model_name=model_name, + upstream_base_url=proxy_base_url, + upstream_api_key=session_api_key, + deployment_id=deployment_id, + ) + logger.info( + "Arena LLM registration succeeded: model_name=%s, " + "model_id=%s, registered_url=%s", + model_name, + registered_model_id, + registered_url, + ) + return registered_url, registered_model_id + finally: + active_exception = sys.exc_info()[0] is not None + delete_error: Exception | None = None + try: + arena_client.delete_llm_proxy(registered_model_id) + logger.info( + "Arena LLM registration deleted: model_id=%s", + registered_model_id, + ) + except Exception as exc: + delete_error = exc + logger.error( + "Failed to delete Arena LLM registration %s: %s", + deployment_id, + exc, + ) + finally: + end_response = proxy_client.post( + f"{proxy_base_url.rstrip('/')}/rl/end_session", + headers={"Authorization": f"Bearer {session_api_key}"}, + ) + end_response.raise_for_status() + if delete_error is not None and not active_exception: + raise delete_error + + +def main(args: list[str]) -> None: + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--num-rollouts", type=int, default=1) + parser.add_argument("--registry-smoke", action="store_true") + parser.add_argument("--serve-after-smoke-seconds", type=int, default=0) + rollout_args, remaining = parser.parse_known_args(args) + if rollout_args.num_rollouts < 1: + raise ValueError("--num-rollouts must be positive") + + config, _ = load_expr_config(remaining, SWEPPOConfig) + if config.scheduler.type != "local": + raise ValueError("arena_rollout_only.py requires scheduler.type=local") + + log_path = ( + Path(config.cluster.fileroot) + / "logs" + / getpass.getuser() + / config.experiment_name + / config.trial_name + / "rollout_only.log" + ) + log_path.parent.mkdir(parents=True, exist_ok=True) + logging.setup_file_logging(str(log_path)) + + econfig = config.econfig + arena_client = ArenaOpenAPIClient( + base_url=econfig.arena_base_url, + timeout=econfig.arena_request_timeout, + poll_interval=econfig.arena_poll_interval, + request_retries=econfig.arena_request_retries, + ) + stream_id = "" + selected_rows: list[dict[str, str]] = [] + if not rollout_args.registry_smoke: + smoke_data_id = os.getenv("ARENA_ROLLOUT_DATA_ID", "") + if smoke_data_id: + stream = arena_client.resolve_stream(econfig.stream_id) + stream_id = str(stream["stream_id"]) + llm_protocol = infer_llm_protocol(stream) + selected_rows = [ + { + "data_id": smoke_data_id, + "stream_id": stream_id, + "llm_protocol": llm_protocol, + } + ] + logger.info( + "Using preselected Arena smoke row: stream_id=%s, data_id=%s, " + "protocol=%s", + stream_id, + smoke_data_id, + llm_protocol, + ) + else: + logger.info("Loading Arena Stream dataset") + stream = arena_client.resolve_stream(econfig.stream_id) + stream_id = str(stream["stream_id"]) + llm_protocol = infer_llm_protocol(stream) + rows = arena_client.get_all_dataset_rows(stream_id, llm_protocol) + selection_rng = random.Random(config.seed) + selected_rows = selection_rng.sample( + rows, k=min(rollout_args.num_rollouts, len(rows)) + ) + if len(selected_rows) < rollout_args.num_rollouts: + raise ValueError( + f"Stream {stream_id!r} has only {len(rows)} rows, but " + f"{rollout_args.num_rollouts} were requested" + ) + logger.info( + "Loaded %d rows from Arena Stream %s using %s; running %d rollout(s)", + len(rows), + stream_id, + llm_protocol, + len(selected_rows), + ) + + rollout_alloc = ModelAllocation.from_str(config.rollout.backend, name="rollout") + if rollout_alloc.backend != "sglang": + raise ValueError( + f"arena_rollout_only.py requires an SGLang backend, got " + f"{rollout_alloc.backend!r}" + ) + server_args = SGLangConfig.build_args( + sglang_config=config.sglang, + tp_size=rollout_alloc.parallel.tp_size, + pp_size=rollout_alloc.parallel.pp_size, + base_gpu_id=0, + ) + config.rollout.max_head_offpolicyness = int(1e12) + + # These values are shell-expansion placeholders for Slurm workers. Local + # workers already inherit the real container environment; keeping the + # placeholders would overwrite credentials with literal "$..." strings. + for scheduling_spec in config.rollout.scheduling_spec: + scheduling_spec.env_vars.pop("ARENA_OPENAPI_BASE", None) + scheduling_spec.env_vars.pop("ARENA_OPENAPI_TOKEN", None) + scheduling_spec.env_vars.pop("ARENA_LLM_API_KEY", None) + + scheduler = LocalScheduler( + gpu_devices=list(range(config.cluster.n_gpus_per_node)), + exp_config=config, + ) + controller = RemoteSGLangEngine.as_controller(config.rollout, scheduler) + wandb_run = _init_wandb(config) + try: + controller.initialize(role="arena-rollout-smoke", server_args=server_args) + controller.start_proxy() + if rollout_args.registry_smoke: + agent_config = config.rollout.agent + if agent_config is None: + raise ValueError("rollout.agent is required for registry smoke") + registered_url, deployment_id = _run_registry_smoke( + arena_client=arena_client, + proxy_base_url=controller.get_proxy_addr(0), + proxy_admin_api_key=agent_config.admin_api_key, + trial_name=config.trial_name, + ) + wandb_run.log({"registry_smoke/succeeded": 1}) + logger.info( + "Registry smoke complete: deployment_id=%s, registered_url=%s", + deployment_id, + registered_url, + ) + if rollout_args.serve_after_smoke_seconds > 0: + logger.info( + "Keeping rollout service alive for %d seconds", + rollout_args.serve_after_smoke_seconds, + ) + time.sleep(rollout_args.serve_after_smoke_seconds) + return + rollout_results, failed_data_ids = _run_rollout_tasks( + controller=controller, + rows=selected_rows, + workflow_kwargs={ + "econfig": asdict(econfig), + "gen_args": { + "temperature": config.gconfig.temperature, + "max_completion_tokens": config.gconfig.max_new_tokens, + }, + "timeout": econfig.timeout, + }, + ) + rewards = [reward for _, reward in rollout_results] + completed = len(rewards) + failed = len(failed_data_ids) + mean_reward = sum(rewards) / completed if completed else 0.0 + max_reward = max(rewards, default=0.0) + reward_one_count = sum(reward >= 1.0 for reward in rewards) + logger.info( + "Arena rollout-only completed: stream_id=%s, completed=%d, failed=%d, " + "mean_reward=%.4f, max_reward=%.4f, reward_one_count=%d, results=%s, " + "failed_data_ids=%s", + stream_id, + completed, + failed, + mean_reward, + max_reward, + reward_one_count, + rollout_results, + failed_data_ids, + ) + wandb_run.log( + { + "rollout/mean_reward": mean_reward, + "rollout/max_reward": max_reward, + "rollout/reward_one_count": reward_one_count, + "rollout/completed": completed, + "rollout/failed": failed, + } + ) + finally: + try: + controller.destroy() + finally: + scheduler.delete_workers(None) + wandb_run.finish() + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/examples/swe/qwen3_30b_a3b_grpo.yaml b/examples/swe/qwen3_30b_a3b_grpo.yaml index a9fc7c3c34..92d62081b7 100644 --- a/examples/swe/qwen3_30b_a3b_grpo.yaml +++ b/examples/swe/qwen3_30b_a3b_grpo.yaml @@ -1,23 +1,20 @@ -experiment_name: qwen3_30b_a3b_grpo -trial_name: trial0 +experiment_name: ${oc.env:EXPERIMENT_NAME} +trial_name: ${oc.env:TRIAL_NAME} seed: 42 -enable_offload: true +enable_offload: false total_train_epochs: 10 total_train_steps: 100 tokenizer_path: ${actor.path} -# Clean up external SWE sandbox instances after the run (LOG_DIR is injected). -post_exit_hook: 'AWE_ROOT=\${AWEAGENT_ROOT:-\${SWE_AGENT_ROOT:-/path/to/AReaL-SWEAgent}}; cd "$AWE_ROOT" && PYTHONPATH=$PWD:$PYTHONPATH pip install -q -r requirements.txt && AENV_SYSTEM_URL=\${AENV_SYSTEM_URL} python -m aweagent.maintenance.clean_instances' - cluster: - # 4 nodes x 8 GPUs: 2 for rollout (SGLang), 2 for actor (Megatron). - n_nodes: 4 + # 2 nodes x 8 GPUs: 1 for rollout (SGLang), 1 for actor (Megatron). + n_nodes: 2 n_gpus_per_node: 8 - fileroot: /tmp/areal/experiments + fileroot: ${oc.env:AREAL_FILEROOT} name_resolve: type: nfs - nfs_record_root: /tmp/areal/name_resolve + nfs_record_root: ${cluster.fileroot}/name_resolve/${experiment_name} scheduler: type: slurm @@ -35,7 +32,7 @@ gconfig: rollout: experiment_name: ${experiment_name} trial_name: ${trial_name} - backend: "sglang:d4t8p1" + backend: "sglang:d2t4p1" max_concurrent_rollouts: 32 queue_size: null consumer_batch_size: ${train_dataset.batch_size} @@ -70,9 +67,9 @@ rollout: actor: experiment_name: ${experiment_name} trial_name: ${trial_name} - # attn DP4 PP1 TP4 CP2; ffn DP4 PP1 EP8. CP=2 shards long sequences. - backend: "megatron:(attn:d4p1t4c2|ffn:d4p1e8)" - path: /path/to/Qwen3-Coder-30B-A3B-Instruct + # attn DP1 PP1 TP2 CP4; ffn DP1 PP1 EP8. CP=4 shards long sequences. + backend: "megatron:(attn:d1p1t2c4|ffn:d1p1e8)" + path: ${oc.env:MODEL_PATH} init_from_scratch: false disable_dropout: true gradient_checkpointing: true @@ -142,32 +139,24 @@ actor: gpu: 1 cpu: 4 mem: 32 - exclusive: true - reservation: null nodelist: null exclude: null - image: /path/to/areal.sif + image: ${oc.env:AREAL_IMAGE} cmd: python3 -m areal.infra.rpc.rpc_server env_vars: NCCL_TIMEOUT: 600 PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" # Point HOME and caches at shared storage to avoid filling the # container's small default HOME during JIT compilation. - HOME: "/path/to/cache/home" - HF_HOME: "/path/to/cache/hf" - XDG_CACHE_HOME: "/path/to/cache/xdg" - # Workers must import both AReaL and the external AReaL-SWEAgent. - PYTHONPATH: "/path/to/AReaL-SWEAgent:/path/to/AReaL" - AENV_SYSTEM_URL: "http://your-aenv-service:8080" - SWE_AGENT_ROOT: "/path/to/AReaL-SWEAgent" - AWEAGENT_ROOT: "/path/to/AReaL-SWEAgent" + HOME: "${oc.env:AREAL_CACHE_ROOT}/home" + HF_HOME: "${oc.env:AREAL_CACHE_ROOT}/hf" + XDG_CACHE_HOME: "${oc.env:AREAL_CACHE_ROOT}/xdg" + PYTHONPATH: ${oc.env:AREAL_DIR} + # Expand on the worker at launch time; do not store credentials in config. + ARENA_OPENAPI_BASE: "$ARENA_OPENAPI_BASE" + ARENA_OPENAPI_TOKEN: "$ARENA_OPENAPI_TOKEN" + ARENA_LLM_API_KEY: "$ARENA_LLM_API_KEY" OPENAI_MODEL: "Qwen3-Coder-30B-A3B-Instruct" - additional_bash_cmds: - # Install AReaL-SWEAgent's aenv/fastmcp into the venv the workers use. - - "/opt/.venv/bin/python -m ensurepip --upgrade >/dev/null 2>&1 || true" - - "/opt/.venv/bin/python -m pip install --no-deps aenvironment==0.1.8rc2 || echo 'WARN: aenvironment install failed'" - - "/opt/.venv/bin/python -m pip install 'fastmcp<3' || echo 'WARN: fastmcp install failed'" - - "/opt/.venv/bin/python -c 'import aenv, aweagent' || echo 'WARN: aenv/aweagent import failed'" ref: experiment_name: ${experiment_name} @@ -188,20 +177,20 @@ ref: target: actor scheduling_spec: ${actor.scheduling_spec} -# SWE-bench agent environment. Set agent_type=cc to train Claude Code through -# the same proxy path. agent_root points at the external AReaL-SWEAgent checkout. +# Arena Stream dataset and task environment. Leave stream_id empty to use the +# first active Stream returned by GET /openapi/v1/streams. econfig: - dataset_path: "" - agent_type: swe - agent_config: "" - swe_agent_config: 1_0_0/min-swe-agent-train-top1 - cc_agent_config: train_cc_time3600 - agent_root: /path/to/AReaL-SWEAgent - swe_agent_root: /path/to/AReaL-SWEAgent - llm_model: "" - opencode_provider: "" - codex_provider: "" - max_completion_tokens: 131071 + dataset_source: arena + stream_id: swe-bench-verified + arena_base_url: ${oc.env:ARENA_OPENAPI_BASE} + arena_request_timeout: 180.0 + arena_request_retries: 3 + arena_poll_interval: 5.0 + arena_registration_timeout: 180.0 + # Extra environment variables injected into the Arena Harness/Agent sandbox. + arena_task_envs: + # Prevent Claude Code from changing the terminal title during rollouts. + CLAUDE_CODE_DISABLE_TERMINAL_TITLE: "1" timeout: 3600.0 sglang: @@ -226,14 +215,14 @@ sglang: disable_custom_all_reduce: true disable_cuda_graph: false -# Datasets: SWE-bench style jsonl with instance_id / problem_statement / -# eval_script per line. +# Dataset paths are unused when econfig.dataset_source=arena; the data ids are +# loaded from the selected Stream before PPOTrainer starts. train_dataset: batch_size: 4 shuffle: true pin_memory: true num_workers: 4 - path: /path/to/swe_bench_rl.jsonl + path: "" type: rl max_length: 32768 drop_last: true @@ -243,7 +232,7 @@ valid_dataset: shuffle: true pin_memory: true num_workers: 4 - path: /path/to/swe_bench_rl.jsonl + path: "" type: rl # Utilities @@ -280,7 +269,9 @@ stats_logger: trial_name: ${trial_name} fileroot: ${cluster.fileroot} wandb: - mode: disabled + mode: online + project: ${experiment_name} + name: ${trial_name} perf_tracer: enabled: false diff --git a/examples/swe/run_arena_rollout_only.sh b/examples/swe/run_arena_rollout_only.sh new file mode 100755 index 0000000000..89b912ce2e --- /dev/null +++ b/examples/swe/run_arena_rollout_only.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +AREAL_DIR="${AREAL_DIR:-$(cd -- "${SCRIPT_DIR}/../.." && pwd)}" +ARENA_RUN_ENV_FILE="${ARENA_RUN_ENV_FILE:-${AREAL_DIR}/.arena-rollout.env}" + +if [[ ! -f "${ARENA_RUN_ENV_FILE}" ]]; then + echo "Missing ${ARENA_RUN_ENV_FILE}. Create it or set ARENA_RUN_ENV_FILE." >&2 + exit 1 +fi + +set -a +# shellcheck disable=SC1090 +source "${ARENA_RUN_ENV_FILE}" +set +a + +export AREAL_DIR +export EXPERIMENT_NAME="${EXPERIMENT_NAME:-swe-arena-rollout-only}" +export TRIAL_NAME="${TRIAL_NAME:-qwen3-coder-30b-a3b-swebench-128k-batch16-concurrency4-$(date +%Y%m%d-%H%M%S)}" +export ARENA_NUM_ROLLOUTS="${ARENA_NUM_ROLLOUTS:-16}" +export ARENA_MAX_CONCURRENT_ROLLOUTS="${ARENA_MAX_CONCURRENT_ROLLOUTS:-4}" + +required_variables=( + AREAL_IMAGE + AREAL_FILEROOT + AREAL_CACHE_ROOT + MODEL_PATH + ARENA_OPENAPI_BASE + ARENA_OPENAPI_TOKEN + ARENA_LLM_API_KEY + SWE_RL_ADMIN_API_KEY + WANDB_API_KEY + WANDB_BASE_URL +) +for variable_name in "${required_variables[@]}"; do + if [[ -z "${!variable_name:-}" ]]; then + echo "${variable_name} must be set in ${ARENA_RUN_ENV_FILE}" >&2 + exit 1 + fi +done + +if [[ "${1:-}" == "--check" ]]; then + echo "Arena rollout environment is ready." + echo "trial_name=${TRIAL_NAME}" + exit 0 +fi +if [[ $# -ne 0 ]]; then + echo "Usage: $0 [--check]" >&2 + exit 2 +fi + +unset ARENA_ROLLOUT_DATA_ID +mkdir -p "${AREAL_DIR}/rl_logs" + +job_id=$(sbatch --parsable --export=ALL "${SCRIPT_DIR}/sbatch_arena_rollout_only.sh") +echo "Submitted Arena rollout-only job ${job_id}" +echo "trial_name=${TRIAL_NAME}" +echo "log=${AREAL_DIR}/rl_logs/swe-arena-rollout-only-${job_id}.out" diff --git a/examples/swe/sbatch_arena_rollout_only.sh b/examples/swe/sbatch_arena_rollout_only.sh new file mode 100755 index 0000000000..cf610ec7f6 --- /dev/null +++ b/examples/swe/sbatch_arena_rollout_only.sh @@ -0,0 +1,92 @@ +#!/usr/bin/env bash +#SBATCH -J swe-arena-rollout-only +#SBATCH -N 1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=32 +#SBATCH --mem=256G +#SBATCH -t 12:00:00 +#SBATCH -o rl_logs/swe-arena-rollout-only-%j.out +#SBATCH -e rl_logs/swe-arena-rollout-only-%j.err + +set -euo pipefail + +: "${AREAL_DIR:?AREAL_DIR must point to this AReaL checkout}" +: "${AREAL_IMAGE:?AREAL_IMAGE must point to the training image}" +: "${AREAL_FILEROOT:?AREAL_FILEROOT must be shared storage}" +: "${AREAL_CACHE_ROOT:?AREAL_CACHE_ROOT must be shared storage}" +: "${MODEL_PATH:?MODEL_PATH must point to the Qwen3-Coder checkpoint}" +: "${EXPERIMENT_NAME:?EXPERIMENT_NAME is required}" +: "${TRIAL_NAME:?TRIAL_NAME is required}" +: "${ARENA_OPENAPI_BASE:?ARENA_OPENAPI_BASE is required}" +: "${ARENA_OPENAPI_TOKEN:?ARENA_OPENAPI_TOKEN is required}" +: "${ARENA_LLM_API_KEY:?ARENA_LLM_API_KEY is required}" +: "${SWE_RL_ADMIN_API_KEY:?SWE_RL_ADMIN_API_KEY is required}" +: "${WANDB_API_KEY:?WANDB_API_KEY is required}" +: "${WANDB_BASE_URL:?WANDB_BASE_URL is required}" + +AREAL_PYTHON="${AREAL_PYTHON:-/opt/.venv/bin/python3}" +ARENA_NUM_ROLLOUTS="${ARENA_NUM_ROLLOUTS:-16}" +ARENA_MAX_CONCURRENT_ROLLOUTS="${ARENA_MAX_CONCURRENT_ROLLOUTS:-4}" +ROLLOUT_ONLY_MODE_ARG="" +if [[ "${ARENA_REGISTRY_SMOKE:-0}" == "1" ]]; then + ROLLOUT_ONLY_MODE_ARG="--registry-smoke" +fi + +mkdir -p \ + "${AREAL_DIR}/rl_logs" \ + "${AREAL_FILEROOT}/name_resolve/${EXPERIMENT_NAME}" \ + "${AREAL_CACHE_ROOT}/home" \ + "${AREAL_CACHE_ROOT}/hf" \ + "${AREAL_CACHE_ROOT}/xdg" + +export APPTAINERENV_AREAL_DIR="${AREAL_DIR}" +export APPTAINERENV_AREAL_FILEROOT="${AREAL_FILEROOT}" +export APPTAINERENV_AREAL_CACHE_ROOT="${AREAL_CACHE_ROOT}" +export APPTAINERENV_MODEL_PATH="${MODEL_PATH}" +export APPTAINERENV_EXPERIMENT_NAME="${EXPERIMENT_NAME}" +export APPTAINERENV_TRIAL_NAME="${TRIAL_NAME}" +export APPTAINERENV_ARENA_OPENAPI_BASE="${ARENA_OPENAPI_BASE}" +export APPTAINERENV_ARENA_OPENAPI_TOKEN="${ARENA_OPENAPI_TOKEN}" +export APPTAINERENV_ARENA_LLM_API_KEY="${ARENA_LLM_API_KEY:-}" +export APPTAINERENV_ARENA_ROLLOUT_DATA_ID="${ARENA_ROLLOUT_DATA_ID:-}" +export APPTAINERENV_SWE_RL_ADMIN_API_KEY="${SWE_RL_ADMIN_API_KEY}" +export APPTAINERENV_WANDB_API_KEY="${WANDB_API_KEY}" +export APPTAINERENV_WANDB_BASE_URL="${WANDB_BASE_URL}" +export APPTAINERENV_AREAL_PYTHON="${AREAL_PYTHON}" +export APPTAINERENV_HOME="${AREAL_CACHE_ROOT}/home" +export APPTAINERENV_HF_HOME="${AREAL_CACHE_ROOT}/hf" +export APPTAINERENV_XDG_CACHE_HOME="${AREAL_CACHE_ROOT}/xdg" +export APPTAINERENV_PYTHONPATH="${AREAL_DIR}" + +echo "Starting single-node Arena rollout-only job ${SLURM_JOB_ID}" + +srun --mpi=pmi2 --ntasks=1 --cpus-per-task=32 --mem=256G \ + singularity exec --nv --pid --writable-tmpfs \ + --bind /storage:/storage \ + "${AREAL_IMAGE}" \ + bash -lc " + export PATH='${AREAL_PYTHON%/*}':\${PATH} + cd '${AREAL_DIR}' + '${AREAL_PYTHON}' -m examples.swe.arena_rollout_only \ + ${ROLLOUT_ONLY_MODE_ARG} \ + --num-rollouts "${ARENA_NUM_ROLLOUTS}" \ + --config examples/swe/qwen3_30b_a3b_grpo.yaml \ + scheduler.type=local \ + cluster.n_nodes=1 \ + cluster.n_gpus_per_node=8 \ + rollout.backend=sglang:d1t8p1 \ + gconfig.n_samples=1 \ + rollout.consumer_batch_size="${ARENA_NUM_ROLLOUTS}" \ + rollout.max_concurrent_rollouts="${ARENA_MAX_CONCURRENT_ROLLOUTS}" \ + rollout.setup_timeout=3600.0 \ + sglang.context_length="${SGLANG_CONTEXT_LENGTH:-133120}" \ + sglang.max_prefill_tokens="${SGLANG_MAX_PREFILL_TOKENS:-133119}" \ + sglang.max_running_requests="${ARENA_MAX_CONCURRENT_ROLLOUTS}" \ + sglang.cuda_graph_max_bs="${ARENA_MAX_CONCURRENT_ROLLOUTS}" \ + econfig.arena_request_timeout=30.0 \ + train_dataset.batch_size="${ARENA_NUM_ROLLOUTS}" \ + train_dataset.num_workers=0 \ + valid_dataset.batch_size="${ARENA_NUM_ROLLOUTS}" \ + valid_dataset.num_workers=0 + " diff --git a/examples/swe/sbatch_arena_stream.sh b/examples/swe/sbatch_arena_stream.sh new file mode 100755 index 0000000000..22415dcf15 --- /dev/null +++ b/examples/swe/sbatch_arena_stream.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +#SBATCH -J swe-arena-stream +#SBATCH -N 1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=10G +#SBATCH -t 7-00:00:00 +#SBATCH -o rl_logs/swe-arena-stream-%j.out +#SBATCH -e rl_logs/swe-arena-stream-%j.err + +set -euo pipefail + +: "${AREAL_DIR:?AREAL_DIR must point to this AReaL checkout}" +: "${AREAL_IMAGE:?AREAL_IMAGE must point to the training image}" +: "${AREAL_FILEROOT:?AREAL_FILEROOT must be shared storage}" +: "${AREAL_CACHE_ROOT:?AREAL_CACHE_ROOT must be shared storage}" +: "${MODEL_PATH:?MODEL_PATH must point to the Qwen3-Coder checkpoint}" +: "${EXPERIMENT_NAME:?EXPERIMENT_NAME is required}" +: "${TRIAL_NAME:?TRIAL_NAME is required}" +: "${ARENA_OPENAPI_BASE:?ARENA_OPENAPI_BASE is required}" +: "${ARENA_OPENAPI_TOKEN:?ARENA_OPENAPI_TOKEN is required}" +: "${ARENA_LLM_API_KEY:?ARENA_LLM_API_KEY is required}" +: "${SWE_RL_ADMIN_API_KEY:?SWE_RL_ADMIN_API_KEY is required}" +: "${WANDB_API_KEY:?WANDB_API_KEY is required}" +: "${WANDB_BASE_URL:?WANDB_BASE_URL is required}" + +AREAL_PYTHON="${AREAL_PYTHON:-/opt/.venv/bin/python3}" + +mkdir -p \ + "${AREAL_DIR}/rl_logs" \ + "${AREAL_FILEROOT}/name_resolve/${EXPERIMENT_NAME}" \ + "${AREAL_CACHE_ROOT}/home" \ + "${AREAL_CACHE_ROOT}/hf" \ + "${AREAL_CACHE_ROOT}/xdg" + +export APPTAINERENV_AREAL_DIR="${AREAL_DIR}" +export APPTAINERENV_AREAL_IMAGE="${AREAL_IMAGE}" +export APPTAINERENV_AREAL_FILEROOT="${AREAL_FILEROOT}" +export APPTAINERENV_AREAL_CACHE_ROOT="${AREAL_CACHE_ROOT}" +export APPTAINERENV_MODEL_PATH="${MODEL_PATH}" +export APPTAINERENV_EXPERIMENT_NAME="${EXPERIMENT_NAME}" +export APPTAINERENV_TRIAL_NAME="${TRIAL_NAME}" +export APPTAINERENV_ARENA_OPENAPI_BASE="${ARENA_OPENAPI_BASE}" +export APPTAINERENV_ARENA_OPENAPI_TOKEN="${ARENA_OPENAPI_TOKEN}" +export APPTAINERENV_ARENA_LLM_API_KEY="${ARENA_LLM_API_KEY}" +export APPTAINERENV_SWE_RL_ADMIN_API_KEY="${SWE_RL_ADMIN_API_KEY}" +export APPTAINERENV_WANDB_API_KEY="${WANDB_API_KEY}" +export APPTAINERENV_WANDB_BASE_URL="${WANDB_BASE_URL}" +export APPTAINERENV_AREAL_PYTHON="${AREAL_PYTHON}" +export APPTAINERENV_HOME="${AREAL_CACHE_ROOT}/home" +export APPTAINERENV_HF_HOME="${AREAL_CACHE_ROOT}/hf" +export APPTAINERENV_XDG_CACHE_HOME="${AREAL_CACHE_ROOT}/xdg" +export APPTAINERENV_PYTHONPATH="${AREAL_DIR}" + +echo "Starting Arena Stream controller job ${SLURM_JOB_ID}" + +srun --mpi=pmi2 --ntasks=1 --cpus-per-task=4 --mem=10G \ + singularity exec --pid --writable-tmpfs \ + --bind /storage:/storage \ + --bind /etc/slurm/:/etc/slurm/ \ + --bind /etc/passwd:/etc/passwd:ro \ + --bind /etc/group:/etc/group:ro \ + --bind /etc/munge:/etc/munge:ro \ + --bind /var/run/munge:/var/run/munge \ + --bind /usr/bin/sbatch:/usr/bin/sbatch \ + --bind /usr/bin/srun:/usr/bin/srun \ + --bind /usr/bin/squeue:/usr/bin/squeue \ + --bind /usr/bin/scancel:/usr/bin/scancel \ + --bind /usr/bin/scontrol:/usr/bin/scontrol \ + --bind /usr/lib64/slurm:/usr/lib64/slurm \ + "${AREAL_IMAGE}" \ + bash -lc " + (/usr/sbin/munged 2>/dev/null || true) + cd '${AREAL_DIR}' + '${AREAL_PYTHON}' -m examples.swe.train_swe_rl \ + --config examples/swe/qwen3_30b_a3b_grpo.yaml + " diff --git a/examples/swe/train_swe_rl.py b/examples/swe/train_swe_rl.py index 57db957ef7..317879ee3d 100644 --- a/examples/swe/train_swe_rl.py +++ b/examples/swe/train_swe_rl.py @@ -8,6 +8,7 @@ from datasets import Dataset +from examples.swe.arena_client import ArenaOpenAPIClient, infer_llm_protocol from examples.swe.utils import SWEPPOConfig from areal import PPOTrainer @@ -81,6 +82,25 @@ def group_filter(x: dict[str, Any]): return x["rewards"].mean() <= 0.95 +def get_arena_dataset(econfig) -> tuple[Dataset, str]: + """Load all data ids from one Arena Stream without dataset sharding.""" + client = ArenaOpenAPIClient( + base_url=econfig.arena_base_url, + timeout=econfig.arena_request_timeout, + request_retries=econfig.arena_request_retries, + ) + stream = client.resolve_stream(econfig.stream_id) + stream_id = str(stream["stream_id"]) + llm_protocol = infer_llm_protocol(stream) + rows = client.get_all_dataset_rows(stream_id, llm_protocol) + dataset = Dataset.from_list(rows) + logger.info( + f"Created Arena dataset with {len(dataset)} items from stream {stream_id} " + f"using {llm_protocol}" + ) + return dataset, stream_id + + def _install_aweagent_deps_on_ray_nodes(aweagent_root: str): """Install AReaL-SWEAgent dependencies on all Ray GPU nodes. @@ -151,34 +171,45 @@ def main(args): config, _ = load_expr_config(args, SWEPPOConfig) econfig = config.econfig - # When using Ray scheduler, ensure SWEAgent deps are on all nodes - if config.scheduler.type == "ray": + # When using Ray with the external agent, ensure its deps are on all nodes. + if config.scheduler.type == "ray" and econfig.dataset_source == "jsonl": import ray ray.init(address="auto", ignore_reinit_error=True) _install_aweagent_deps_on_ray_nodes(_resolve_aweagent_root(econfig)) - # Resolve dataset paths from config - train_path = config.train_dataset.path - valid_path = config.valid_dataset.path - - def resolve_path(p: str) -> str: - if Path(p).is_absolute() or Path(p).exists(): + if econfig.dataset_source == "arena": + train_dataset, resolved_stream_id = get_arena_dataset(econfig) + valid_dataset = train_dataset + econfig.stream_id = resolved_stream_id + workflow = "examples.swe.arena_agent.ArenaStreamAgentWorkflow" + elif econfig.dataset_source == "jsonl": + # Resolve dataset paths from config + train_path = config.train_dataset.path + valid_path = config.valid_dataset.path + + def resolve_path(p: str) -> str: + if Path(p).is_absolute() or Path(p).exists(): + return p + if econfig.dataset_path: + candidate = Path(econfig.dataset_path) / p + if candidate.exists(): + return str(candidate) return p - if econfig.dataset_path: - candidate = Path(econfig.dataset_path) / p - if candidate.exists(): - return str(candidate) - return p - - train_dataset = get_swe_dataset( - dataset_path=resolve_path(train_path), - split="train", - ) - valid_dataset = get_swe_dataset( - dataset_path=resolve_path(valid_path), - split="test", - ) + + train_dataset = get_swe_dataset( + dataset_path=resolve_path(train_path), + split="train", + ) + valid_dataset = get_swe_dataset( + dataset_path=resolve_path(valid_path), + split="test", + ) + workflow = "examples.swe.agent.SWEAgentWorkflow" + else: + raise ValueError( + f"Unsupported econfig.dataset_source: {econfig.dataset_source!r}" + ) # Build workflow kwargs from dataclasses import asdict @@ -206,7 +237,7 @@ def resolve_path(p: str) -> str: valid_dataset=valid_dataset, ) as trainer: trainer.train( - workflow="examples.swe.agent.SWEAgentWorkflow", + workflow=workflow, workflow_kwargs=workflow_kwargs, eval_workflow=None, eval_workflow_kwargs=eval_workflow_kwargs, diff --git a/examples/swe/utils.py b/examples/swe/utils.py index c53881cbff..9dc978ae20 100644 --- a/examples/swe/utils.py +++ b/examples/swe/utils.py @@ -10,7 +10,15 @@ class SWEEnvConfig: """Environment configuration for AReaL-SWEAgent-backed SWE-bench training. Attributes: + dataset_source: Dataset and agent backend, either ``jsonl`` or ``arena``. dataset_path: Path to the SWE-bench JSONL dataset file. + stream_id: Target Arena Stream. The first active Stream is used when empty. + arena_base_url: Arena OpenAPI base URL. Defaults to ARENA_OPENAPI_BASE. + arena_request_timeout: Timeout for Stream and dataset discovery requests. + arena_request_retries: Retries for transient Arena HTTP failures. + arena_poll_interval: Delay between asynchronous task-result requests. + arena_registration_timeout: Timeout for LLM registration and deletion. + arena_task_envs: Extra environment variables injected into Arena tasks. agent_type: AReaL-SWEAgent agent type to train, e.g. ``swe`` or ``cc``. agent_config: Generic AReaL-SWEAgent config name. When set, this overrides the compatibility fields below. @@ -26,10 +34,60 @@ class SWEEnvConfig: timeout: Maximum time allowed for a single episode in seconds. """ + dataset_source: str = field( + default="jsonl", + metadata={ + "help": "Dataset and agent backend: 'jsonl' or 'arena'.", + "choices": ["jsonl", "arena"], + }, + ) dataset_path: str = field( default="", metadata={"help": "Path to the SWE-bench JSONL dataset file."}, ) + stream_id: str = field( + default="", + metadata={ + "help": ( + "Arena Stream id. When empty, the first active Stream returned by " + "the Arena OpenAPI is used." + ) + }, + ) + arena_base_url: str = field( + default="", + metadata={ + "help": ( + "Arena OpenAPI base URL. Defaults to the ARENA_OPENAPI_BASE " + "environment variable." + ) + }, + ) + arena_request_timeout: float = field( + default=60.0, + metadata={"help": "Arena Stream and dataset request timeout in seconds."}, + ) + arena_request_retries: int = field( + default=3, + metadata={"help": "Retries for transient Arena HTTP request failures."}, + ) + arena_poll_interval: float = field( + default=5.0, + metadata={"help": "Arena task-result polling interval in seconds."}, + ) + arena_registration_timeout: float = field( + default=180.0, + metadata={"help": "Arena LLM registration request timeout in seconds."}, + ) + arena_task_envs: dict[str, str] = field( + default_factory=dict, + metadata={ + "help": ( + "Additional environment variables passed to launch_one_task. " + "MODEL_NAME, BASE_URL, and API_KEY are managed by AReaL." + ) + }, + ) agent_type: str = field( default="swe", metadata={ diff --git a/tests/experimental/openai/test_proxy_rollout_server.py b/tests/experimental/openai/test_proxy_rollout_server.py index 45e3660598..4c62cdc8e3 100644 --- a/tests/experimental/openai/test_proxy_rollout_server.py +++ b/tests/experimental/openai/test_proxy_rollout_server.py @@ -41,6 +41,67 @@ def _admin_headers(): return {"Authorization": f"Bearer {_ADMIN_KEY}"} +# --------------------------------------------------------------------------- +# Tests: message preprocessing +# --------------------------------------------------------------------------- + + +def test_preprocess_messages_flattens_text_blocks_and_preserves_images(monkeypatch): + """OpenAI-routed Claude text blocks should be flattened before inference.""" + + class RemoveReminder: + def __call__(self, messages): + for message in messages: + if isinstance(message.get("content"), str): + message["content"] = message["content"].replace( + "reminder", "processed" + ) + return messages + + monkeypatch.setattr(srv, "_message_preprocessors", [RemoveReminder()]) + image_content = [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": "https://example/image.png"}}, + ] + messages = [ + {"role": "system", "content": image_content}, + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + { + "type": "tool_result", + "content": [{"type": "text", "text": "reminder"}], + }, + ], + }, + {"role": "user", "content": image_content}, + ] + + result = srv._preprocess_messages(messages) + + assert isinstance(result[0]["content"], str) + assert result[1]["content"] == "hello\nprocessed" + assert result[2]["content"] == image_content + + +def test_prepare_request_messages_normalizes_tuple_and_system_content(monkeypatch): + """Tuple-backed request messages should not bypass system normalization.""" + monkeypatch.setattr(srv, "_message_preprocessors", []) + messages = ( + { + "role": "system", + "content": ({"type": "text", "text": "system prompt"},), + }, + {"role": "user", "content": "hello"}, + ) + + result = srv._prepare_request_messages(messages) + + assert isinstance(result, list) + assert result[0]["content"] == "system prompt" + + # --------------------------------------------------------------------------- # Tests: start_session with provided api_key # --------------------------------------------------------------------------- diff --git a/tests/test_config_redaction.py b/tests/test_config_redaction.py new file mode 100644 index 0000000000..fc7c58b097 --- /dev/null +++ b/tests/test_config_redaction.py @@ -0,0 +1,30 @@ +"""Tests for safe experiment configuration serialization.""" + +from areal.utils.config_utils import REDACTED_VALUE, redact_sensitive_config + + +def test_redact_sensitive_config_redacts_nested_credentials_only(): + """Credentials should be removed without hiding token-count configuration.""" + config = { + "admin_api_key": "admin-secret", + "arena": {"access_token": "arena-secret"}, + "workers": [{"password": "worker-secret", "max_tokens": 131072}], + "tokenizer_path": "/models/tokenizer", + } + + redacted = redact_sensitive_config(config) + + assert redacted == { + "admin_api_key": REDACTED_VALUE, + "arena": {"access_token": REDACTED_VALUE}, + "workers": [{"password": REDACTED_VALUE, "max_tokens": 131072}], + "tokenizer_path": "/models/tokenizer", + } + assert config["admin_api_key"] == "admin-secret" + + +def test_redact_sensitive_config_preserves_empty_optional_credentials(): + """Empty optional credential fields should remain empty for readability.""" + config = {"wandb_api_key": "", "refresh_token": None} + + assert redact_sensitive_config(config) == config diff --git a/tests/test_rollout_controller.py b/tests/test_rollout_controller.py index 749a4d193d..1d59eaf307 100644 --- a/tests/test_rollout_controller.py +++ b/tests/test_rollout_controller.py @@ -87,6 +87,7 @@ async def async_call_engine(self, worker_id, method, *args, **kwargs): self._task_counter += 1 # Simulate a successful rollout result result = { + "source_id": kwargs["data"].get("id"), "input_ids": torch.randint(0, 100, (1, 10)), "attention_mask": torch.ones(1, 10, dtype=torch.bool), "loss_mask": torch.tensor( @@ -425,6 +426,34 @@ def test_wait_returns_distributed_batch(self): controller.destroy() + def test_wait_for_task_returns_matching_trajectory(self): + config = create_test_config(consumer_batch_size=16, max_concurrent_rollouts=50) + scheduler = MockScheduler() + controller = RolloutController( + inf_engine=MockInferenceEngine, + config=config, + scheduler=scheduler, + ) + controller.initialize(role="rollout", server_args={}) + + first_task_id = controller.submit( + {"id": 1}, + workflow="tests.utils.TestWorkflow", + workflow_kwargs={}, + ) + second_task_id = controller.submit( + {"id": 2}, + workflow="tests.utils.TestWorkflow", + workflow_kwargs={}, + ) + + second_result = controller.wait_for_task(second_task_id, timeout=5.0) + first_result = controller.wait_for_task(first_task_id, timeout=5.0) + + assert second_result is not None and second_result["source_id"] == 2 + assert first_result is not None and first_result["source_id"] == 1 + controller.destroy() + def test_wait_timeout_when_insufficient_results(self): config = create_test_config(consumer_batch_size=16, max_concurrent_rollouts=10) scheduler = MockScheduler() diff --git a/tests/test_swe_arena.py b/tests/test_swe_arena.py new file mode 100644 index 0000000000..c475747975 --- /dev/null +++ b/tests/test_swe_arena.py @@ -0,0 +1,451 @@ +"""Tests for the Arena Stream dataset and proxy agent integration.""" + +import asyncio +import json + +import httpx +import pytest +import torch + +from examples.swe.arena_agent import ArenaStreamAgentWorkflow +from examples.swe.arena_client import ( + ArenaAPIError, + ArenaOpenAPIClient, + infer_llm_protocol, +) +from examples.swe.arena_rollout_only import _run_rollout_tasks, _trajectory_reward + + +def test_resolve_stream_id_when_unspecified_returns_first_active(monkeypatch): + """The first active Stream should be selected when no id is configured.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.params["status"] == "ACTIVE" + assert request.headers["Authorization"] == "Bearer test-token" + return httpx.Response( + 200, + json={ + "items": [ + {"stream_id": "stream-first", "status": "ACTIVE"}, + {"stream_id": "stream-second", "status": "ACTIVE"}, + ] + }, + ) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + stream_id = client.resolve_stream_id(client=http_client) + + assert stream_id == "stream-first" + + +def test_resolve_stream_when_id_is_explicit_returns_matching_metadata(monkeypatch): + """An explicit Stream should still be discovered so its Harness is available.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + + def handler(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "items": [ + {"stream_id": "stream-first"}, + { + "stream_id": "stream-selected", + "default_harness_ref": {"key": "claude-code"}, + }, + ] + }, + ) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + stream = client.resolve_stream("stream-selected", client=http_client) + + assert stream["default_harness_ref"] == {"key": "claude-code"} + + +@pytest.mark.parametrize( + ("harness_key", "expected_protocol"), + [ + ("Claude-Code-With-Skills", "anthropic"), + ("openai-codex", "responses"), + ("swe-agent", "chat_completions"), + (None, "chat_completions"), + ], +) +def test_infer_llm_protocol_from_harness_key(harness_key, expected_protocol): + """Harness names should select their native protocol case-insensitively.""" + stream = ( + {"default_harness_ref": {"key": harness_key}} if harness_key is not None else {} + ) + + assert infer_llm_protocol(stream) == expected_protocol + + +def test_list_streams_transient_timeout_retries(monkeypatch): + """Transient read timeouts should be retried before failing discovery.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + attempts = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise httpx.ReadTimeout("transient", request=request) + return httpx.Response(200, json={"items": [{"stream_id": "stream-1"}]}) + + client = ArenaOpenAPIClient( + base_url="https://arena.example", + request_retries=1, + ) + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + streams = client.list_streams(client=http_client) + + assert streams == [{"stream_id": "stream-1"}] + assert attempts == 2 + + +def test_list_streams_transient_gateway_error_retries(monkeypatch): + """Transient gateway errors should be retried before parsing the response.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + attempts = 0 + + def handler(_: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts == 1: + return httpx.Response(504, text="gateway timeout") + return httpx.Response(200, json={"items": [{"stream_id": "stream-1"}]}) + + client = ArenaOpenAPIClient( + base_url="https://arena.example", + request_retries=1, + ) + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + streams = client.list_streams(client=http_client) + + assert streams == [{"stream_id": "stream-1"}] + assert attempts == 2 + + +def test_get_all_dataset_rows_uses_total_as_limit(monkeypatch): + """Dataset loading should probe total then request all rows in one page.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + requested_limits: list[int] = [] + + def handler(request: httpx.Request) -> httpx.Response: + assert request.method == "POST" + assert request.url.path.endswith("/streams/stream-1/dataset") + limit = int(request.url.params["limit"]) + requested_limits.append(limit) + if limit == 1: + return httpx.Response( + 200, + json={ + "data_ids": ["data-1"], + "count": 1, + "total": 3, + "offset": 0, + "limit": 1, + }, + ) + return httpx.Response( + 200, + json={ + "data_ids": ["data-1", "data-2", "data-3"], + "count": 3, + "total": 3, + "offset": 0, + "limit": 3, + }, + ) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + rows = client.get_all_dataset_rows("stream-1", client=http_client) + + assert requested_limits == [1, 3] + assert rows == [ + { + "data_id": "data-1", + "stream_id": "stream-1", + "llm_protocol": "chat_completions", + }, + { + "data_id": "data-2", + "stream_id": "stream-1", + "llm_protocol": "chat_completions", + }, + { + "data_id": "data-3", + "stream_id": "stream-1", + "llm_protocol": "chat_completions", + }, + ] + + +def test_get_all_dataset_rows_over_api_limit_raises(monkeypatch): + """The initial implementation should reject Streams that require pagination.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + + def handler(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "data_ids": ["data-1"], + "count": 1, + "total": 1001, + "offset": 0, + "limit": 1, + }, + ) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + with pytest.raises(ArenaAPIError, match="pagination is not implemented"): + client.get_all_dataset_rows("stream-1", client=http_client) + + +def test_register_and_delete_llm_proxy(monkeypatch): + """Registry calls should forward one deployment id and the proxy session.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + deployment_id = "deployment-1" + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + assert request.headers["Authorization"] == "Bearer arena-token" + if request.method == "POST": + assert request.url.path.endswith("/openapi/v1/llm/models") + payload = json.loads(request.content) + assert payload == { + "model_name": "stream-areal-test-1", + "endpoints": [ + { + "endpoint_id": deployment_id, + "upstream_model": "stream-areal-test-1", + "base_url": "http://rollout-proxy", + "api_key": "session-key", + "inbound_protos": ["chat"], + "enabled": True, + } + ], + "enabled": True, + "metadata": {"deployment_id": deployment_id}, + } + return httpx.Response( + 200, + json={"model_name": "stream-areal-test-1"}, + ) + assert request.method == "DELETE" + assert request.url.path.endswith("/openapi/v1/llm/models/stream-areal-test-1") + return httpx.Response(204) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + registered_url, returned_id = client.register_llm_proxy( + model_name="stream-areal-test-1", + upstream_base_url="http://rollout-proxy", + upstream_api_key="session-key", + deployment_id=deployment_id, + client=http_client, + ) + client.delete_llm_proxy(returned_id, client=http_client) + + assert registered_url == "https://arena.example/api" + assert returned_id == "stream-areal-test-1" + assert len(requests) == 2 + + +@pytest.mark.parametrize( + "protocol", + ["anthropic", "responses", "chat_completions"], +) +def test_register_llm_proxy_always_advertises_openai_chat(monkeypatch, protocol): + """All Harness protocols should use the AReaL proxy's native Chat API.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + + def handler(request: httpx.Request) -> httpx.Response: + payload = json.loads(request.content) + assert payload["endpoints"][0]["inbound_protos"] == ["chat"] + return httpx.Response(200, json={"model_name": payload["model_name"]}) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + client.register_llm_proxy( + model_name="stream-areal-test-1", + upstream_base_url="http://rollout-proxy", + upstream_api_key="session-key", + deployment_id="deployment-1", + protocol=protocol, + client=http_client, + ) + + +def test_agent_launches_task_through_proxy_and_returns_reward(monkeypatch): + """The Arena agent should forward proxy credentials and return task reward.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + + result_polls = 0 + deployment_id = "" + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal deployment_id, result_polls + if request.method == "POST" and request.url.path.endswith( + "/openapi/v1/llm/models" + ): + assert request.headers["Authorization"] == "Bearer arena-token" + payload = json.loads(request.content) + assert payload["model_name"].startswith("stream-areal-") + assert payload["endpoints"][0]["upstream_model"] == payload["model_name"] + assert payload["endpoints"][0]["base_url"] == "http://rollout-proxy" + assert payload["endpoints"][0]["api_key"] == "session-key" + assert payload["endpoints"][0]["inbound_protos"] == ["chat"] + deployment_id = payload["model_name"] + return httpx.Response( + 200, + json={"model_name": payload["model_name"]}, + ) + if request.url.path.endswith("/streams/stream-1/launch_one_task"): + assert request.headers["Authorization"] == "Bearer arena-token" + assert json.loads(request.content) == { + "data_id": "data-1", + "model_name": deployment_id, + "base_url": "https://arena.example/api", + "api_key": "test-llm-key", + "envs": { + "MODEL_NAME": deployment_id, + "BASE_URL": "https://arena.example/api", + "API_KEY": "test-llm-key", + "CLAUDE_CODE_DISABLE_TERMINAL_TITLE": "1", + }, + } + return httpx.Response( + 202, + json={"accepted": True, "task_id": "task-1", "status": "PENDING"}, + ) + if request.method == "DELETE" and request.url.path.endswith( + f"/openapi/v1/llm/models/{deployment_id}" + ): + assert request.headers["Authorization"] == "Bearer arena-token" + return httpx.Response(204) + assert request.headers["Authorization"] == "Bearer arena-token" + assert request.url.path.endswith("/tasks/task-1/result") + result_polls += 1 + if result_polls == 1: + return httpx.Response(200, json={"status": "AGENT_RUNNING"}) + return httpx.Response(200, json={"status": "OK", "score": 0.75}) + + async def run_agent() -> float: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_poll_interval": 0.0, + "arena_task_envs": {"CLAUDE_CODE_DISABLE_TERMINAL_TITLE": "1"}, + "timeout": 10.0, + } + ) + return await workflow.run( + { + "stream_id": "stream-1", + "data_id": "data-1", + "llm_protocol": "anthropic", + }, + base_url="http://rollout-proxy", + api_key="session-key", + arena_http_client=http_client, + ) + + reward = asyncio.run(run_agent()) + assert reward == 0.75 + assert result_polls == 2 + + +def test_launch_one_task_failed_result_raises(monkeypatch): + """Terminal infrastructure failures must not silently become zero reward.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return httpx.Response( + 202, + json={"task_id": "task-1", "status": "PENDING"}, + ) + return httpx.Response(200, json={"status": "HARNESS_FAILED"}) + + async def launch_task() -> None: + client = ArenaOpenAPIClient( + base_url="https://arena.example", + poll_interval=0.0, + ) + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler) + ) as http_client: + await client.launch_one_task( + stream_id="stream-1", + data_id="data-1", + model_name="deployment-1", + proxy_base_url="http://rollout-proxy/v1", + proxy_api_key="session-key", + client=http_client, + ) + + with pytest.raises(ArenaAPIError, match="HARNESS_FAILED"): + asyncio.run(launch_task()) + + +def test_trajectory_reward_sums_tensor_rewards(): + """Rollout-only logging should recover the episode reward tensor.""" + reward = _trajectory_reward({"rewards": torch.tensor([0.0, 0.75])}) + + assert reward == 0.75 + + +def test_trajectory_reward_sums_string_interactions(): + """External string trajectories should also produce one episode reward.""" + reward = _trajectory_reward( + { + "interactions": [ + {"reward": 0.0}, + {"reward": 1.0}, + ] + } + ) + + assert reward == 1.0 + + +def test_run_rollout_tasks_keeps_successes_when_one_task_is_rejected(): + """A failed harness task should not discard successful batch results.""" + + class FakeController: + def __init__(self) -> None: + self.rows: dict[int, dict[str, str]] = {} + + def submit(self, data, **_kwargs): + task_id = len(self.rows) + self.rows[task_id] = data + return task_id + + def wait_for_task(self, task_id): + if self.rows[task_id]["data_id"] == "data-failed": + return None + return {"rewards": torch.tensor([1.0])} + + completed, failed = _run_rollout_tasks( + controller=FakeController(), + rows=[ + {"stream_id": "stream-1", "data_id": "data-success"}, + {"stream_id": "stream-1", "data_id": "data-failed"}, + ], + workflow_kwargs={}, + ) + + assert completed == [("data-success", 1.0)] + assert failed == ["data-failed"]