feat: add Arena Stream rollout integration

.gitignore

@@ -201,5 +201,6 @@ api_key.json
 .vscode/
 wandb/
 outputs/
+rl_logs/
 sympy/
 !/docs/figures/*

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,

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
 
 

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,

areal/utils/config_utils.py

@@ -0,0 +1,35 @@
+"""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

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",

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,

examples/swe/README.md

@@ -1,5 +1,40 @@
 # 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}
+```
+
+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
+`/llm/model/new`, then posts the returned model id and the row's `data_id` to
+`launch_one_task`. The task environment receives `MODEL_NAME`, `BASE_URL`, and
+`API_KEY`; 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