Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -201,5 +201,7 @@ api_key.json
.vscode/
wandb/
outputs/
rl_logs/
.arena-rollout.env
sympy/
!/docs/figures/*
3 changes: 2 additions & 1 deletion areal/api/cli_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
120 changes: 101 additions & 19 deletions areal/experimental/openai/proxy/proxy_rollout_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Comment on lines +631 to +640

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

If messages is a one-shot iterator or generator, iterating over it directly will consume its elements. If an unsupported message type is encountered and the function returns messages early, the caller will receive a partially consumed iterator, leading to silent data loss. Converting messages to a list first avoids this issue.

Suggested change
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)
messages_list = list(messages)
normalized: list[dict] = []
for message in messages_list:
if isinstance(message, BaseModel):
message = message.model_dump()
elif isinstance(message, Mapping):
message = dict(message)
else:
return messages_list
normalized.append(_materialize_iterables(message))
return _preprocess_messages(normalized)



async def _call_client_create(
create_fn,
request: dict[str, Any] | BaseModel,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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}")


Expand Down Expand Up @@ -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(
Expand All @@ -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


Expand Down
7 changes: 7 additions & 0 deletions areal/infra/controller/rollout_controller.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
37 changes: 37 additions & 0 deletions areal/utils/config_utils.py
Original file line number Diff line number Diff line change
@@ -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 = "<redacted>"


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
1 change: 1 addition & 0 deletions areal/utils/logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
14 changes: 12 additions & 2 deletions areal/utils/stats_logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
39 changes: 39 additions & 0 deletions examples/swe/README.md
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Loading
Loading