From ca97d216468e93038a45709f9d6d888f40fd0433 Mon Sep 17 00:00:00 2001 From: MaxLEAF3824 Date: Tue, 11 Aug 2026 12:56:22 +0000 Subject: [PATCH 1/3] feat: support RAO metadata and FA4 attention Carry recursive rollout metadata through the OpenAI proxy and training data path so downstream RAO workloads can preserve node identity and depth. Add the FA4 backend required by the team's B200 environment. Key changes: - Preserve session, parent, depth, and interaction metadata - Export recursive rollout structure for downstream training - Register FlashAttention 4 as a supported FSDP backend - Refresh generated CLI references Co-Authored-By: Claude --- areal/engine/fsdp_engine.py | 6 + areal/engine/fsdp_utils/attn_impl.py | 3 + areal/engine/fsdp_utils/fa4_attn.py | 179 ++++++++++++++++++ areal/experimental/openai/cache.py | 16 ++ areal/experimental/openai/client.py | 32 ++++ .../openai/proxy/proxy_rollout_server.py | 16 ++ areal/experimental/openai/proxy/server.py | 14 ++ areal/experimental/openai/types.py | 90 +++++++++ areal/infra/controller/rollout_callback.py | 10 +- areal/infra/workflow_executor.py | 76 +++++++- docs/en/cli_reference.md | 8 +- docs/zh/cli_reference.md | 8 +- 12 files changed, 448 insertions(+), 10 deletions(-) create mode 100644 areal/engine/fsdp_utils/fa4_attn.py diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index 27fc005f86..c71e6af1bf 100644 --- a/areal/engine/fsdp_engine.py +++ b/areal/engine/fsdp_engine.py @@ -1033,6 +1033,12 @@ def _create_device_model(self): self.get_device_stats().log("before model creation/loading") + # FA4 不是 transformers 自带的 backend，用到时先注册进 ALL_ATTENTION_FUNCTIONS。 + # Blackwell(sm100) 上 FA2/FA3 都不可用，只有 FA4 能跑。 + from areal.engine.fsdp_utils.fa4_attn import maybe_register_fa4 + + maybe_register_fa4(self.config.attn_impl) + # Note: VLMs often have vision_tower in fp32 already; loading whole # model in optimizer_dtype (fp32 default) is consistent. if self.is_vision_model: diff --git a/areal/engine/fsdp_utils/attn_impl.py b/areal/engine/fsdp_utils/attn_impl.py index 5bbc92a45f..11ca9dbcd5 100644 --- a/areal/engine/fsdp_utils/attn_impl.py +++ b/areal/engine/fsdp_utils/attn_impl.py @@ -7,6 +7,9 @@ "sdpa", "flash_attention_2", "flash_attention_3", + # FA4（flash_attn.cute）。transformers 没有这个 backend，由 fa4_attn.py 注册； + # 名字不能含 "flash"（否则被 transformers 抢走）。Blackwell(sm100) 上唯一可用的 FA。 + "fa4", "flex_attention", ) diff --git a/areal/engine/fsdp_utils/fa4_attn.py b/areal/engine/fsdp_utils/fa4_attn.py new file mode 100644 index 0000000000..7bbc246221 --- /dev/null +++ b/areal/engine/fsdp_utils/fa4_attn.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +"""把 FlashAttention-4（`flash_attn.cute`）注册成 transformers 的 attention backend。 + +**为什么需要这个文件** + +Blackwell（B200 / sm_100）上： + - FA2 没有 sm100 kernel，编不出来； + - FA3 是 Hopper-only； + - 只有 FA4（`flash-attn-4` 包，import 路径 `flash_attn.cute`）支持 sm100， + 且前向/反向都有。 + +但 transformers（截至 5.3.0）没有 `flash_attention_4` 这个 backend—— +它的 `flash_attention_2` 路径走的是 FA2 的老 API `from flash_attn import flash_attn_func`， +而 `flash-attn-4` 装出来的 `flash_attn` 是个只含 `cute` 子模块的 namespace package， +于是报 `cannot import name 'flash_attn_func' from 'flash_attn' (unknown location)`。 +参见 huggingface/transformers#44559。 + +本模块用 transformers 的 `AttentionInterface.register()` 补上这个 backend， +之后 `attn_implementation="flash_attention_4"` 即可正常使用。 + +AReaL 侧的输入是**打包过的变长序列**（batch=1，附带 `cu_seq_lens_q/k`、 +`max_length_q/k`），正好对应 FA4 的 `flash_attn_varlen_func`。 +""" + +from typing import Any + +import torch + +from areal.utils import logging + +logger = logging.getLogger("FA4") + +ATTN_IMPL_NAME = "fa4" # 名字不能含 "flash"：transformers 的 is_flash_attention_requested 是 `"flash" in name`， +# 一旦命中就会去 lazy-import FA2/FA3/hub-kernels，绕不过我们注册的实现。 + +_registered = False + + +def fa4_available() -> bool: + """FA4 是否可用（装了 flash-attn-4 且能 import 到 cute 后端）。""" + try: + from flash_attn.cute import flash_attn_func # noqa: F401 + from flash_attn.cute import flash_attn_varlen_func # noqa: F401 + + return True + except Exception: + return False + + +def _unwrap(out: Any) -> torch.Tensor: + """FA4 在 return_lse=True 时返回 (out, lse)，统一取出 out。""" + return out[0] if isinstance(out, tuple) else out + + +def fa4_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + dropout: float = 0.0, + scaling: float | None = None, + sliding_window: int | None = None, + softcap: float | None = None, + is_causal: bool | None = None, + **kwargs, +) -> tuple[torch.Tensor, None]: + """transformers attention-interface 约定的签名。 + + 入参 q/k/v 形状是 (batch, num_heads, seq, head_dim)； + FA4 和 FA2 一样吃 (…, seq, num_heads, head_dim)，所以先 transpose。 + """ + # ── 树训练分支 ───────────────────────────────────────────────────── + # enable_tree_training 打开时，AReaL 会把 trie 打包后的注意力掩码作为 + # tree_block_mask / tree_triton_data 透传进来。它官方的做法是 monkey-patch + # transformers.integrations.flash_attention._flash_attention_forward + # （module_fsdp.patch_fsdp_for_tree_training），但那条路径只在 attn_impl 是 + # flash_attention_2 时才被调用 —— 而 B200 上 FA2 根本导入不了 + # （`cannot import name 'flash_attn_func'`），我们走的是自注册的 fa4。 + # 不在这里接管的话，树打包后的序列会被当成普通因果序列算，**分支之间会互相 + # 看见**，是静默的错误结果。 + if kwargs.get("tree_block_mask") is not None or kwargs.get("tree_triton_data") is not None: + from areal.models.tree_attn.module_fsdp import _tree_attn_fwd_func + + out = _tree_attn_fwd_func( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attention_mask, + scaling, + **kwargs, + ) + return out, None + + from flash_attn.cute import flash_attn_func, flash_attn_varlen_func + + if dropout not in (0.0, None): + raise ValueError("FlashAttention-4 不支持 attention dropout（dropout != 0）。") + + # (b, h, s, d) -> (b, s, h, d) + q = query.transpose(1, 2) + k = key.transpose(1, 2) + v = value.transpose(1, 2) + + # 精度：FA4 只接受 bf16/fp16。混精下 query 可能是 fp32，跟随 value 的 dtype。 + if q.dtype not in (torch.bfloat16, torch.float16): + target = v.dtype if v.dtype in (torch.bfloat16, torch.float16) else torch.bfloat16 + q, k, v = q.to(target), k.to(target), v.to(target) + + causal = module.is_causal if is_causal is None else is_causal + window = (sliding_window - 1, 0) if sliding_window else (None, None) + + common = dict( + softmax_scale=scaling, + causal=causal, + window_size=window, + softcap=0.0 if softcap is None else float(softcap), + ) + + cu_q = kwargs.get("cu_seq_lens_q") + cu_k = kwargs.get("cu_seq_lens_k") + + if cu_q is not None and cu_k is not None: + # 打包变长路径：AReaL 的常规路径（batch 恒为 1） + b, s = q.shape[0], q.shape[1] + out = flash_attn_varlen_func( + q.reshape(b * s, *q.shape[2:]), + k.reshape(b * s, *k.shape[2:]), + v.reshape(b * s, *v.shape[2:]), + cu_seqlens_q=cu_q.to(torch.int32), + cu_seqlens_k=cu_k.to(torch.int32), + max_seqlen_q=int(kwargs.get("max_length_q", s)), + max_seqlen_k=int(kwargs.get("max_length_k", s)), + **common, + ) + attn_output = _unwrap(out).reshape(b, s, *q.shape[2:]) + else: + # 稠密路径（没打包时，比如某些 eval / VLM 分支） + if attention_mask is not None: + raise ValueError( + "fa4 目前只支持打包变长输入或无 padding 的稠密输入，" + "收到了非空 attention_mask。请改用 sdpa。" + ) + attn_output = _unwrap(flash_attn_func(q, k, v, **common)) + + return attn_output.to(value.dtype), None + + +def register_fa4() -> bool: + """把 FA4 注册进 transformers 的 ALL_ATTENTION_FUNCTIONS。幂等。 + + Returns + ------- + bool + 注册成功（或此前已注册）返回 True；FA4 不可用返回 False。 + """ + global _registered + if _registered: + return True + if not fa4_available(): + return False + + from transformers import AttentionInterface + + AttentionInterface.register(ATTN_IMPL_NAME, fa4_attention_forward) + _registered = True + logger.info(f"Registered transformers attention backend '{ATTN_IMPL_NAME}'.") + return True + + +def maybe_register_fa4(attn_impl: str) -> None: + """建模型前调用：若用户选了 flash_attention_4，就确保它已注册。""" + if attn_impl != ATTN_IMPL_NAME: + return + if not register_fa4(): + raise RuntimeError( + "attn_impl='fa4' 但 FA4 不可用。" + "请安装 flash-attn-4（`from flash_attn.cute import flash_attn_func` 需能 import）。" + ) diff --git a/areal/experimental/openai/cache.py b/areal/experimental/openai/cache.py index ba6e22bbbb..1bc894a961 100644 --- a/areal/experimental/openai/cache.py +++ b/areal/experimental/openai/cache.py @@ -263,6 +263,15 @@ def __setitem__( value: InteractionWithTokenLogpReward, ) -> None: """Add a new interaction to the cache, automatically building parent-child relationships.""" + # ⭐ RAO：把本次请求原始的 `model`（含 `__dN` 深度后缀）钉到 interaction 上。 + # proxy 在转发前会丢掉 model 字段，这里是唯一还拿得到它的地方。 + if getattr(value, "rao_model", None) is None: + try: + from areal.experimental.openai.types import RAO_REQ_MODEL + + value.rao_model = RAO_REQ_MODEL.get() + except Exception: + pass if value.messages is None: raise ValueError( "Interaction messages must be set to find parent relationship." @@ -468,6 +477,13 @@ def export_interactions( if len(complete_cache) == 0: return {} + # 盖上 session 章：一个 InteractionCache 就是一次 agent 运行（一条 rollout + # session）。下游 dump 靠它还原"哪几轮属于同一条会话"——不能靠 parent 链， + # 因为 harness 丢弃被截断的回复时，parent_data 不再是 child 的前缀， + # 父子关系会断（见本文件 _is_similar_on_last_message 的 Prefix mismatch 告警）。 + for interaction in complete_cache.values(): + interaction.session_id = self._session_id + if style == "concat": for interaction in complete_cache.values(): if interaction.chat_template_type != "concat": diff --git a/areal/experimental/openai/client.py b/areal/experimental/openai/client.py index dc68117fed..c7dd8741d2 100644 --- a/areal/experimental/openai/client.py +++ b/areal/experimental/openai/client.py @@ -357,6 +357,33 @@ def _resolve_max_total_tokens( return min(prompt_len + max_new_tokens, cap) +def _flatten_text_content_parts(messages: list[dict]) -> list[dict]: + """把纯文本的 OpenAI content-parts 列表压回普通字符串。 + + 有些 harness（如 pi）按 OpenAI 多模态格式发消息，``content`` 是 + ``[{"type": "text", "text": "..."}]`` 而不是字符串。而不少 chat template + 对非字符串 content 直接置空 —— Qwen3 的模板就是： + + {%- if message.content is string %}{%- set content = message.content %} + {%- else %}{%- set content = '' %}{%- endif %} + + 结果是**用户消息被静默渲染成空串**，模型收到一个没有内容的请求， + 还会一本正经地回"你还没给我任务"。这里在套模板之前压平，只动纯文本的情况； + 含图片等其它 part 的消息原样保留，交给 _extract_images_from_messages。 + """ + out = [] + for m in messages: + c = m.get("content") + if ( + isinstance(c, list) + and c + and all(isinstance(x, dict) and x.get("type") == "text" for x in c) + ): + m = {**m, "content": "".join(x.get("text") or "" for x in c)} + out.append(m) + return out + + def _parse_tool_call_arguments(messages: list[dict]) -> list[dict]: """Return a new message list with tool_call arguments parsed from JSON strings to dicts. Some chat templates (e.g. GLM-5.1) iterate over arguments with @@ -662,6 +689,7 @@ async def create( has_images = len(image_data) > 0 tokenizer_messages = messages_for_tokenizer if has_images else messages_list + tokenizer_messages = _flatten_text_content_parts(tokenizer_messages) tokenizer_messages = _parse_tool_call_arguments(tokenizer_messages) if self.chat_template_type == "hf": prompt_token_ids = apply_chat_template( @@ -684,6 +712,10 @@ async def create( ) else: concat_tok_messages = concat_messages + # ★ concat 分支用的是 remaining_messages，绕开了上面那次扁平化， + # 这里必须再压一次，否则 pi 这类发 content-parts 的 harness + # 会被模板渲染成空消息。 + concat_tok_messages = _flatten_text_content_parts(concat_tok_messages) prompt_token_ids = concat_prompt_token_ids_with_parent( concat_tok_messages, interaction.parent if interaction is not None else None, diff --git a/areal/experimental/openai/proxy/proxy_rollout_server.py b/areal/experimental/openai/proxy/proxy_rollout_server.py index b42ecb6522..2e3649827b 100644 --- a/areal/experimental/openai/proxy/proxy_rollout_server.py +++ b/areal/experimental/openai/proxy/proxy_rollout_server.py @@ -30,6 +30,7 @@ from areal.api.cli_args import NameResolveConfig from areal.experimental.openai.client import ArealOpenAI +from areal.experimental.openai.types import RAO_REQ_MODEL from areal.infra.rpc.serialization import deserialize_value, serialize_value from areal.infra.utils.http import validate_admin_api_key from areal.utils import name_resolve, names, seeding @@ -578,6 +579,21 @@ async def _call_client_create( session_data.update_last_access() sig = inspect.signature(create_fn) + # ⭐ RAO：`model` 马上就要被丢掉（AReaL 的 client 只服务单一模型、不接受该参数）， + # 但 harness 把**递归深度**编码在 model id 尾部（`__dN`）。丢掉之后 + # 深度信息就永久消失了 —— 实测表现是训练数据里每条序列的 rao_depth 都是 0。 + # 所以在丢弃之前先存进 ContextVar，interaction 落库时抄到 `rao_model` 上。 + # 用 ContextVar 而非实例属性：同一 session 的多个 sub-agent 是并行请求的。 + try: + _req_model = ( + request.get("model") + if isinstance(request, dict) + else getattr(request, "model", None) + ) + if _req_model: + RAO_REQ_MODEL.set(str(_req_model)) + except Exception: + pass # 观测失败绝不能影响请求 areal_client_ignored_args = ["model"] + (extra_ignored_args or []) areal_client_disallowed_args = ["areal_cache"] areal_client_allowed_args = list( diff --git a/areal/experimental/openai/proxy/server.py b/areal/experimental/openai/proxy/server.py index 9c7f6fe4f0..f9e20ff5ac 100644 --- a/areal/experimental/openai/proxy/server.py +++ b/areal/experimental/openai/proxy/server.py @@ -140,11 +140,22 @@ def serialize_interactions( result = {} for key, interaction in interactions.items(): + # 会话身份必须一起过河：subproc/online 模式下 interaction 要跨进程传回， + # 不带上的话下游只能看到一堆扁平的 turn，无法还原"哪几轮属于同一条会话"。 + identity = { + "session_id": interaction.session_id, + "parent_interaction_id": ( + interaction.parent.interaction_id + if interaction.parent is not None + else interaction.parent_interaction_id + ), + } if interaction.has_tensor_data: result[key] = { "tensor_dict": interaction.to_tensor_dict(), "reward": interaction.reward, "interaction_id": interaction.interaction_id, + **identity, } else: result[key] = { @@ -152,6 +163,7 @@ def serialize_interactions( "output_message_list": interaction.output_message_list, "reward": interaction.reward, "interaction_id": interaction.interaction_id, + **identity, } return serialize_value(result) @@ -174,6 +186,8 @@ def deserialize_interactions( interaction.output_message_list = item["output_message_list"] interaction.reward = item["reward"] interaction.interaction_id = item["interaction_id"] + interaction.session_id = item.get("session_id") + interaction.parent_interaction_id = item.get("parent_interaction_id") result[key] = interaction return result diff --git a/areal/experimental/openai/types.py b/areal/experimental/openai/types.py index 2582c37e3b..72e1d5b216 100644 --- a/areal/experimental/openai/types.py +++ b/areal/experimental/openai/types.py @@ -5,6 +5,9 @@ from dataclasses import dataclass, field from enum import Enum +import contextvars +import hashlib +import re import torch from openai.types.chat import ChatCompletion from openai.types.responses.response import Response @@ -32,6 +35,38 @@ class InputName(str, Enum): NONE = "none" +# RAO：请求体里原始的 `model` 字段（含 `__dN` 深度后缀）。 +# proxy 在丢弃 model 之前 set，interaction 落库时读。见下方 `rao_depth` 的说明。 +RAO_REQ_MODEL: contextvars.ContextVar = contextvars.ContextVar( + "rao_req_model", default=None +) + + +def rao_iid_hash(interaction_id: str | None) -> int: + """`interaction_id` → 48 位整数，用于把训练序列**精确**认回它的 rollout 行。 + + rollout dump（文本，带 session/interaction id）和 train_batch dump（token， + 带 loss_mask）之间原本**没有共享 key** —— 训练侧只有 token id，谁生成的、 + 在树的哪个节点上，全丢了。以前只能靠 `seqlen == n_tokens` 反推，实测 + force50 的 5654 条训练序列里有 602 条歧义（多条 rollout 撞同一个长度）。 + + 做法沿用 `rao_depth` 趟通的那条路：per-token 常量张量搭训练数据的顺风车。 + + ⚠️ **48 位不是随手取的**： + - 张量装不下字符串，只能塞整数； + - 用 `dtype=torch.long`（和 `input_ids` 同款）而不是 `rao_depth` 那样的 + float32 —— float32 只有 24 位尾数，48 位的哈希塞进去会被截得面目全非； + - 48 < 53，所以中途任何一次 `float()` 往返（tap 里就有）都还是精确的。 + + ~3 万条序列下碰撞期望约 0.05 条，可以忽略。 + """ + if not interaction_id: + return 0 + return int.from_bytes( + hashlib.blake2b(interaction_id.encode(), digest_size=6).digest(), "big" + ) + + @dataclass class InteractionWithTokenLogpReward: """Internal structure to store completions/responses with their rewards.""" @@ -50,12 +85,48 @@ class InteractionWithTokenLogpReward: # Completion fields (optional for response) completion: ChatCompletion | None = None + # RAO：递归深度。harness 把它编码在 model id 尾部（`__d`）—— + # 因为 sub-agent 的 system prompt 与 root 不同、prefix_matcher 匹配不上、 + # parent 恒为 None，树结构在 AReaL 侧根本推不出来，只能由 harness 显式传。 + # 走 model id 是因为 pi 既没有 --base-url 也不给自定义 header 的口子， + # 而 model 必然出现在请求体里。换 API key 不行 —— 会把每层拆成不同 session。 + # + # ⚠️ **不能只读 `completion.model`** —— proxy 在转发前把 `model` 字段整个丢掉了 + # （`proxy_rollout_server.py`：`areal_client_ignored_args = ["model"]`， + # 因为 AReaL 的 client 只服务单一模型、不接受 model 参数）。 + # 响应里的 model 名是服务端改写过的，`__dN` 早没了。 + # 实测：真实跑里**每条序列的 rao_depth 都是 0，连 sub-agent 也是**。 + # （mock 测试没抓到 —— 它只数派生次数，从没验证过深度解析结果。） + # + # 修法：proxy 在丢掉 model **之前**写进 `RAO_REQ_MODEL` 这个 ContextVar， + # interaction 落库时抄到 `rao_model`。 + # ⚠️ 必须是 ContextVar 不能是实例属性 —— 同一 session 里多个 sub-agent 是 + # **并行**请求的，实例属性会串味；ContextVar 在 asyncio 任务间天然隔离。 + rao_model: str | None = None + + @property + def rao_depth(self) -> int: + model = ( + self.rao_model + or getattr(self.completion, "model", None) + or getattr(self.response, "model", None) + or "" + ) + m = re.search(r"__d(\d+)$", str(model)) + return int(m.group(1)) if m else 0 + # Response fields (optional for completion) response: Response | None = None input_data: str | ResponseInputParam = field(default_factory=lambda: "") # Interaction ID cache (used for deserialization) _interaction_id: str | None = None + # 该 interaction 属于哪一次 agent 运行（= 一条 rollout session）。 + # 由 InteractionCache.export_interactions 盖章，供 dump 落盘还原会话归属。 + session_id: str | None = None + # 父节点的 id。subproc/online 模式下 interaction 要跨进程传回，`parent` + # 这个对象引用没法序列化，只能带 id 过来。 + parent_interaction_id: str | None = None @property def has_tensor_data(self) -> bool: @@ -201,6 +272,25 @@ def to_tensor_dict(self) -> dict[str, torch.Tensor]: attention_mask=torch.ones(len(seq), dtype=torch.bool).unsqueeze(0), # reward rewards=torch.tensor([float(reward)]), + # RAO：递归深度。 + # + # ⚠️ 这里原本是 per-seq 的 `[1]` 形张量（像 `rewards` 一样）。**那样是错的**： + # `split_padded_tensor_dict_into_mb_list` 只切分**带序列维**的张量， + # `[B]` 形的会被**整份复制**给每个微批 —— 于是到 loss 现场 + # `rao_depth.numel()` 是整个 batch 的条数，而 `cu_seqlens` 只描述本微批， + # 两边对不上，`rao_depth_weighting` 的守卫判否后静默跳过， + # train_batch dump 里也就一直看不到 `rao_depth`。 + # （`rewards` 没暴露这个问题，是因为它在 actor.py:343 算完 advantage 就被 pop 掉了。） + # + # 改成 per-token 的 `[1, L]`：跟着 `loss_mask` 一起切分、一起 pack， + # loss 现场取每条序列的首 token（`dep[cu_seqlens[i]]`）就是该条的深度。 + rao_depth=torch.full((1, len(seq)), float(self.rao_depth)), + # RAO：这条序列出自哪个 interaction（见 `rao_iid_hash`）。 + # 和 rollout dump 的 `interaction_id` 一 join 就能把"模型说了什么" + # 和"这段话进没进梯度"接起来。同样是 per-token，理由同上。 + rao_iid=torch.full( + (1, len(seq)), rao_iid_hash(self.interaction_id), dtype=torch.long + ), ) self._cache = result return result diff --git a/areal/infra/controller/rollout_callback.py b/areal/infra/controller/rollout_callback.py index 1f106a5373..8c0b6e7d5c 100644 --- a/areal/infra/controller/rollout_callback.py +++ b/areal/infra/controller/rollout_callback.py @@ -29,7 +29,15 @@ class RolloutCallback: """ controller_addr: str - request_timeout: float = 600.0 + # Raised from 600s (2026-08-02). The LoRA weight-update path fans out + # `/load_lora_adapter` to every inference engine, and that request queues behind + # in-flight `/generate` work on the SGLang scheduler -- unlike the full-model path, + # which passes `abort_all_requests=True`. Under a saturated rollout queue the load + # can starve well past 600s; observed on a 32-concurrency agentic RL run where ten + # updates succeeded and the eleventh timed out at exactly 600s, killing the job. + # A longer ceiling only costs time in the pathological case and never in the happy + # path, so prefer it over losing the run. + request_timeout: float = 1800.0 def _post(self, endpoint: str, payload: dict[str, Any] | None = None) -> dict: """Make synchronous HTTP POST to controller callback endpoint. diff --git a/areal/infra/workflow_executor.py b/areal/infra/workflow_executor.py index 8d00e68877..a6966cd392 100644 --- a/areal/infra/workflow_executor.py +++ b/areal/infra/workflow_executor.py @@ -906,11 +906,76 @@ def _split_trajectory_for_dump( "segments": segments, } + @staticmethod + def _build_interaction_meta(traj: dict[str, Any]) -> list[dict[str, Any]] | None: + """为每个 interaction 算出它在会话树里的身份，供 dump 落盘。 + + `traj` 是 ``dict[completion_id, InteractionWithTokenLogpReward]``， + 每个 value 自带 ``parent`` 指针，因此整棵会话树在这里是完整的。但下游 + ``concat_padded_tensors([v.to_tensor_dict() for v in traj.values()])`` + 只取 values，身份随即丢失 —— 落盘后就再也无法可靠还原"哪几轮属于同一条 + 会话"（实测靠 prompt 前缀反推只有 29% 的任务能还原正确）。 + + 约定： + - ``session_id`` 由 ``InteractionCache.export_interactions`` 盖章， + 一个 cache = 一次 agent 运行 = 一条 rollout session，形如 ``"2-0"`` + - ``session_idx`` 按各 session 首次出现的顺序编号（cache 是 OrderedDict， + 迭代序即真实发起顺序） + - ``turn_idx`` = 该 session 内的第几轮（同 session 内的迭代序） + + 为什么不用 ``parent`` 链：harness 会把超长被截断的回复整个丢弃，此时 + ``parent.messages + parent.output_message_list`` 不再是 child 的前缀， + 父子关系直接断掉（cache 里会打 "Prefix mismatch" 告警）。parent 仍然照常 + 导出，只是当作辅助信息，不作为会话归属的依据。 + """ + items = list(traj.values()) + if not items: + return None + + def safe_id(v: Any, fallback: int) -> str: + try: + return v.interaction_id or f"anon-{fallback}" + except Exception: + return f"anon-{fallback}" + + def safe_created(v: Any) -> float | None: + try: + return v.created_at + except Exception: + return None + + session_order: dict[str, int] = {} + turn_counter: dict[str, int] = {} + meta = [] + for i, v in enumerate(items): + sid = getattr(v, "session_id", None) or "unknown" + if sid not in session_order: + session_order[sid] = len(session_order) + turn = turn_counter.get(sid, 0) + turn_counter[sid] = turn + 1 + parent = getattr(v, "parent", None) + meta.append( + { + "interaction_id": safe_id(v, i), + "parent_id": ( + safe_id(parent, -1) + if parent is not None + else getattr(v, "parent_interaction_id", None) + ), + "session_id": sid, + "session_idx": session_order[sid], + "turn_idx": turn, + "created_at": safe_created(v), + } + ) + return meta + async def _dump_trajectory( self, traj: dict[str, Any] | None, task_id: int, is_eval: bool, + interaction_meta: list[dict[str, Any]] | None = None, ) -> tuple[bool, str]: if traj is None: return False, "trajectory is None" @@ -1005,6 +1070,12 @@ async def _dump_trajectory( if split["segments"] is not None: record["segments"] = split["segments"] + # 会话身份：让下游能精确还原"哪几轮属于同一条 session"。 + # 每个 interaction 的 to_tensor_dict() 只产生一行，因此 + # 批内行号 i 与 traj.values() 的第 i 项严格一一对应。 + if interaction_meta is not None and i < len(interaction_meta): + record.update(interaction_meta[i]) + await f.write(json.dumps(record) + "\n") return True, "" except Exception as e: @@ -1164,9 +1235,12 @@ async def _execute_workflow() -> _RolloutResult | None: # External-API interactions have no tensor data; fall back to # concat_string_interactions which produces a plain dict of # request/response strings instead of padded tensors. + interaction_meta = None if isinstance(traj, dict) and all( isinstance(v, InteractionWithTokenLogpReward) for v in traj.values() ): + # 拍平之前先把会话树的身份抓出来，之后 .values() 就丢了 + interaction_meta = self._build_interaction_meta(traj) if all(v.has_tensor_data for v in traj.values()): traj = concat_padded_tensors( [v.to_tensor_dict() for v in traj.values()] @@ -1191,7 +1265,7 @@ async def _execute_workflow() -> _RolloutResult | None: # Dump trajectory to file if self.config.dump_to_file: dump_success, dump_reason = await self._dump_trajectory( - traj, task_id, pending_task.is_eval + traj, task_id, pending_task.is_eval, interaction_meta ) if not dump_success: self.logger.warning( diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index 883738179c..22440d2907 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -352,7 +352,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `experiment_name` | string | **Required** | - | | `trial_name` | string | **Required** | - | | `path` | string | `""` | Path to HuggingFace checkpoint | -| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `flex_attention` | +| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `fa4`, `flex_attention` | | `use_kernels` | boolean | `False` | Enable Hugging Face kernels model kernelization after model creation. | | `init_from_scratch` | boolean | `False` | Initialize model weights randomly | | `is_critic` | boolean | `False` | Whether to use a critic/reward model | @@ -427,7 +427,7 @@ Configuration for PPO critic model, a subclass of a TrainEngine. | `experiment_name` | string | **Required** | - | | `trial_name` | string | **Required** | - | | `path` | string | `""` | Path to HuggingFace checkpoint | -| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `flex_attention` | +| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `fa4`, `flex_attention` | | `use_kernels` | boolean | `False` | Enable Hugging Face kernels model kernelization after model creation. | | `init_from_scratch` | boolean | `False` | Initialize model weights randomly | | `is_critic` | boolean | `False` | Whether to use a critic/reward model | @@ -475,7 +475,7 @@ Core configuration for model training, including optimization and backend settin | `experiment_name` | string | **Required** | - | | `trial_name` | string | **Required** | - | | `path` | string | `""` | Path to HuggingFace checkpoint | -| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `flex_attention` | +| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `fa4`, `flex_attention` | | `use_kernels` | boolean | `False` | Enable Hugging Face kernels model kernelization after model creation. | | `init_from_scratch` | boolean | `False` | Initialize model weights randomly | | `is_critic` | boolean | `False` | Whether to use a critic/reward model | @@ -989,7 +989,7 @@ fields. | `experiment_name` | string | **Required** | - | | `trial_name` | string | **Required** | - | | `path` | string | `""` | Path to HuggingFace checkpoint | -| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `flex_attention` | +| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `fa4`, `flex_attention` | | `use_kernels` | boolean | `False` | Enable Hugging Face kernels model kernelization after model creation. | | `init_from_scratch` | boolean | `False` | Initialize model weights randomly | | `is_critic` | boolean | `False` | Whether to use a critic/reward model | diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index 09c3167ff2..3d84810116 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -350,7 +350,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `experiment_name` | string | **Required** | - | | `trial_name` | string | **Required** | - | | `path` | string | `""` | Path to HuggingFace checkpoint | -| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `flex_attention` | +| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `fa4`, `flex_attention` | | `use_kernels` | boolean | `False` | Enable Hugging Face kernels model kernelization after model creation. | | `init_from_scratch` | boolean | `False` | Initialize model weights randomly | | `is_critic` | boolean | `False` | Whether to use a critic/reward model | @@ -425,7 +425,7 @@ Configuration for PPO critic model, a subclass of a TrainEngine. | `experiment_name` | string | **Required** | - | | `trial_name` | string | **Required** | - | | `path` | string | `""` | Path to HuggingFace checkpoint | -| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `flex_attention` | +| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `fa4`, `flex_attention` | | `use_kernels` | boolean | `False` | Enable Hugging Face kernels model kernelization after model creation. | | `init_from_scratch` | boolean | `False` | Initialize model weights randomly | | `is_critic` | boolean | `False` | Whether to use a critic/reward model | @@ -473,7 +473,7 @@ Core configuration for model training, including optimization and backend settin | `experiment_name` | string | **Required** | - | | `trial_name` | string | **Required** | - | | `path` | string | `""` | Path to HuggingFace checkpoint | -| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `flex_attention` | +| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `fa4`, `flex_attention` | | `use_kernels` | boolean | `False` | Enable Hugging Face kernels model kernelization after model creation. | | `init_from_scratch` | boolean | `False` | Initialize model weights randomly | | `is_critic` | boolean | `False` | Whether to use a critic/reward model | @@ -987,7 +987,7 @@ fields. | `experiment_name` | string | **Required** | - | | `trial_name` | string | **Required** | - | | `path` | string | `""` | Path to HuggingFace checkpoint | -| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `flex_attention` | +| `attn_impl` | string | `"flash_attention_2"` | Attention implementation for huggingface transformers model. Accepts builtin transformers backends or a Hugging Face kernels repo ID formatted as org/repo\[@revision\]\[:entrypoint\]. **Choices:** `eager`, `sdpa`, `flash_attention_2`, `flash_attention_3`, `fa4`, `flex_attention` | | `use_kernels` | boolean | `False` | Enable Hugging Face kernels model kernelization after model creation. | | `init_from_scratch` | boolean | `False` | Initialize model weights randomly | | `is_critic` | boolean | `False` | Whether to use a critic/reward model | From fa65238b233efee7b83d9a4d2d4e841e6b8c37fc Mon Sep 17 00:00:00 2001 From: MaxLEAF3824 Date: Thu, 13 Aug 2026 09:47:51 +0000 Subject: [PATCH 2/3] feat(workflow): support pre-merge group credit Allow agent workflows to attach node-level rollout metadata and process complete rollout groups before flattening, enabling strict root LOO and zero-variance rejection without changing PPO.\n\nCo-Authored-By: Claude --- areal/experimental/openai/proxy/workflow.py | 46 ++++++++++++++++++++- areal/infra/remote_inf_engine.py | 5 +++ 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/areal/experimental/openai/proxy/workflow.py b/areal/experimental/openai/proxy/workflow.py index 862b43887b..24176c008b 100644 --- a/areal/experimental/openai/proxy/workflow.py +++ b/areal/experimental/openai/proxy/workflow.py @@ -120,6 +120,10 @@ def __init__( self.subproc_max_workers = subproc_max_workers self.drop_retry_orphans = drop_retry_orphans + def process_group_results(self, results): + processor = getattr(self.agent, "process_group_results", None) + return processor(results) if callable(processor) else results + @trace_session("run_agent") async def _run_agent(self, session_api_key: str, data: dict): if self.mode == "inline": @@ -236,8 +240,16 @@ async def arun_episode( ) raise - # Assign rewards back according to user code output - if isinstance(rewards, dict): + # Assign rewards back according to user code output. RAO returns a + # picklable rollout result with one raw reward per node-final interaction. + rollout_result = None + if hasattr(rewards, "interaction_rewards") and hasattr(rewards, "root_reward"): + rollout_result = rewards + if not getattr(rewards, "valid", True): + raise ValueError(f"Invalid RAO rollout: {getattr(rewards, 'error', '')}") + for completion_id, reward in rewards.interaction_rewards.items(): + await proxy_client.set_reward(completion_id, reward) + elif isinstance(rewards, dict): for completion_id, reward in rewards.items(): await proxy_client.set_reward(completion_id, reward) elif isinstance(rewards, float): @@ -251,6 +263,36 @@ async def arun_episode( style=self.export_style, drop_retry_orphans=self.drop_retry_orphans, ) + if rollout_result is not None: + node_by_interaction = rollout_result.node_by_interaction + missing = set(node_by_interaction) - set(interactions) + if missing: + raise ValueError( + f"RAO node-final interactions missing after export: {sorted(missing)[:4]}" + ) + for interaction_id, interaction in interactions.items(): + node = node_by_interaction.get(interaction_id) + if node is None: + raise ValueError(f"Exported interaction {interaction_id} has no RAO node") + interaction.rao_episode_id = rollout_result.episode_id + interaction.rao_node_id = node.node_id + interaction.rao_node_reward = node.judge.score + interaction.rao_root_reward = rollout_result.root_reward + interaction.rao_is_root = node.is_root + interaction.rao_node_depth = node.depth + cache = getattr(interaction, "_cache", None) + if cache is not None and "input_ids" in cache: + import torch + + from areal.experimental.openai.types import rao_iid_hash + + input_ids = cache["input_ids"] + cache["rao_node_id"] = torch.full_like( + input_ids, rao_iid_hash(node.node_id), dtype=torch.long + ) + node_start = torch.zeros_like(input_ids, dtype=torch.float32) + node_start[:, 0] = 1.0 + cache["rao_node_start"] = node_start # Record stats last_id = list(interactions.keys())[-1] if interactions else None diff --git a/areal/infra/remote_inf_engine.py b/areal/infra/remote_inf_engine.py index 6d443a3f45..ac33415838 100644 --- a/areal/infra/remote_inf_engine.py +++ b/areal/infra/remote_inf_engine.py @@ -88,6 +88,11 @@ async def arun_episode( results = await asyncio.gather( *[self.workflow.arun_episode(engine, data) for _ in range(self.group_size)] ) + group_processor = getattr(self.workflow, "process_group_results", None) + if callable(group_processor): + results = group_processor(results) + if results is None: + return None valid_results = [r for r in results if r is not None] From 751748cff7d36d3bf25813a1a1de354cb514fb47 Mon Sep 17 00:00:00 2001 From: MaxLEAF3824 Date: Mon, 17 Aug 2026 06:22:41 +0000 Subject: [PATCH 3/3] fix: harden grouped agent rollout contracts Keep valid siblings when one rollout fails, propagate stable group and rejection metadata, and bind Agent proxy limits to the active train or eval generation context. Key changes: - isolate per-child grouped rollout failures - reject invalid RAO exports without aborting the worker - preserve group metadata through tensor export - apply role-specific Agent token ceilings Co-Authored-By: Claude --- areal/experimental/openai/proxy/workflow.py | 42 +++++++- areal/experimental/openai/types.py | 3 + areal/infra/remote_inf_engine.py | 25 ++++- areal/trainer/rl_trainer.py | 17 ++- tests/test_eval_agent_token_limit.py | 28 +++++ tests/test_grouped_rollout_workflow.py | 112 ++++++++++++++++++++ 6 files changed, 219 insertions(+), 8 deletions(-) create mode 100644 tests/test_eval_agent_token_limit.py create mode 100644 tests/test_grouped_rollout_workflow.py diff --git a/areal/experimental/openai/proxy/workflow.py b/areal/experimental/openai/proxy/workflow.py index 24176c008b..e07324bbc7 100644 --- a/areal/experimental/openai/proxy/workflow.py +++ b/areal/experimental/openai/proxy/workflow.py @@ -243,10 +243,22 @@ async def arun_episode( # Assign rewards back according to user code output. RAO returns a # picklable rollout result with one raw reward per node-final interaction. rollout_result = None - if hasattr(rewards, "interaction_rewards") and hasattr(rewards, "root_reward"): + if hasattr(rewards, "interaction_rewards") and hasattr( + rewards, "root_reward" + ): rollout_result = rewards if not getattr(rewards, "valid", True): - raise ValueError(f"Invalid RAO rollout: {getattr(rewards, 'error', '')}") + rejection_reason = ( + getattr(rewards, "rejection_reason", None) + or getattr(rewards, "error", None) + or "invalid_rao_rollout" + ) + logger.warning( + "Rejecting invalid RAO rollout group_id=%s: %s", + getattr(rewards, "group_id", None), + rejection_reason, + ) + return None for completion_id, reward in rewards.interaction_rewards.items(): await proxy_client.set_reward(completion_id, reward) elif isinstance(rewards, dict): @@ -270,11 +282,25 @@ async def arun_episode( raise ValueError( f"RAO node-final interactions missing after export: {sorted(missing)[:4]}" ) + extras = set(interactions) - set(node_by_interaction) + if extras: + logger.warning( + "Dropping %d exported retry/orphan interactions without RAO nodes: %s", + len(extras), + sorted(extras)[:4], + ) + interactions = { + interaction_id: interaction + for interaction_id, interaction in interactions.items() + if interaction_id in node_by_interaction + } for interaction_id, interaction in interactions.items(): - node = node_by_interaction.get(interaction_id) - if node is None: - raise ValueError(f"Exported interaction {interaction_id} has no RAO node") + node = node_by_interaction[interaction_id] interaction.rao_episode_id = rollout_result.episode_id + interaction.rao_group_id = getattr(rollout_result, "group_id", None) + interaction.rao_rejection_reason = getattr( + rollout_result, "rejection_reason", None + ) interaction.rao_node_id = node.node_id interaction.rao_node_reward = node.judge.score interaction.rao_root_reward = rollout_result.root_reward @@ -290,6 +316,12 @@ async def arun_episode( cache["rao_node_id"] = torch.full_like( input_ids, rao_iid_hash(node.node_id), dtype=torch.long ) + if rollout_result.group_id is not None: + cache["rao_group_id"] = torch.full_like( + input_ids, + rao_iid_hash(rollout_result.group_id), + dtype=torch.long, + ) node_start = torch.zeros_like(input_ids, dtype=torch.float32) node_start[:, 0] = 1.0 cache["rao_node_start"] = node_start diff --git a/areal/experimental/openai/types.py b/areal/experimental/openai/types.py index 72e1d5b216..4d0012ec98 100644 --- a/areal/experimental/openai/types.py +++ b/areal/experimental/openai/types.py @@ -127,6 +127,9 @@ def rao_depth(self) -> int: # 父节点的 id。subproc/online 模式下 interaction 要跨进程传回，`parent` # 这个对象引用没法序列化，只能带 id 过来。 parent_interaction_id: str | None = None + # RAO 组 credit 元数据。即使幸存组整体被拒绝，processor 输入里仍保留原因。 + rao_group_id: str | None = None + rao_rejection_reason: str | None = None @property def has_tensor_data(self) -> bool: diff --git a/areal/infra/remote_inf_engine.py b/areal/infra/remote_inf_engine.py index ac33415838..5c4216640b 100644 --- a/areal/infra/remote_inf_engine.py +++ b/areal/infra/remote_inf_engine.py @@ -85,8 +85,29 @@ async def arun_episode( ) -> dict[str, Any] | None: from areal.experimental.openai import InteractionWithTokenLogpReward + task_id = workflow_context.get().task_id + group_id = str(task_id) if task_id is not None else str(uuid.uuid4()) + + async def run_child(index: int): + child_data = dict(data) + child_data["group_id"] = group_id + try: + return await self.workflow.arun_episode(engine, child_data) + except Exception as exc: + if self.logger is not None: + self.logger.warning( + "GroupedRolloutWorkflow: child %d/%d failed (%s: %s); " + "rejecting only this trajectory", + index + 1, + self.group_size, + type(exc).__name__, + exc, + exc_info=True, + ) + return None + results = await asyncio.gather( - *[self.workflow.arun_episode(engine, data) for _ in range(self.group_size)] + *[run_child(index) for index in range(self.group_size)] ) group_processor = getattr(self.workflow, "process_group_results", None) if callable(group_processor): @@ -101,7 +122,7 @@ async def arun_episode( return None # Some results None -> warn and continue with valid ones - if len(valid_results) < len(results): + if len(valid_results) < len(results) and self.logger is not None: self.logger.warning( f"GroupedRolloutWorkflow: {len(results) - len(valid_results)}/{len(results)} " "trajectories returned None, using remaining results" diff --git a/areal/trainer/rl_trainer.py b/areal/trainer/rl_trainer.py index 2ec47d7d94..7c44afaf5d 100644 --- a/areal/trainer/rl_trainer.py +++ b/areal/trainer/rl_trainer.py @@ -75,6 +75,18 @@ logger = logging.getLogger("RLTrainer") +def _set_default_agent_engine_max_tokens( + config: InferenceEngineConfig, max_tokens: int | None +) -> None: + """Bind an unset Agent proxy ceiling to the role's generation context.""" + if ( + config.agent is not None + and config.agent.engine_max_tokens is None + and max_tokens is not None + ): + config.agent.engine_max_tokens = max_tokens + + class _EmptyDataLoader: """Minimal dataloader for online mode that yields empty dicts. @@ -1026,8 +1038,11 @@ def _init_rollout( "Use `python3 train.py scheduler.type=local` instead of " "`python3 -m areal.infra.launcher.local`." ) - # Create a working copy of config + # Create a working copy of config. Train and eval share inference servers, + # but their proxy workers need independent total-token ceilings. config = deepcopy(rollout_config) + generation_config = self.config.eval_gconfig if is_eval else self.config.gconfig + _set_default_agent_engine_max_tokens(config, generation_config.max_tokens) if is_eval: # NOTE: eval does not have any offpolicyness control config.max_head_offpolicyness = int(1e12) diff --git a/tests/test_eval_agent_token_limit.py b/tests/test_eval_agent_token_limit.py new file mode 100644 index 0000000000..37d4e38af3 --- /dev/null +++ b/tests/test_eval_agent_token_limit.py @@ -0,0 +1,28 @@ +from types import SimpleNamespace + +from areal.trainer.rl_trainer import _set_default_agent_engine_max_tokens + + +def test_role_context_sets_unconfigured_agent_proxy_limit(): + config = SimpleNamespace(agent=SimpleNamespace(engine_max_tokens=None)) + _set_default_agent_engine_max_tokens(config, 40_000) + assert config.agent.engine_max_tokens == 40_000 + + eval_config = SimpleNamespace(agent=SimpleNamespace(engine_max_tokens=None)) + _set_default_agent_engine_max_tokens(eval_config, 262_144) + assert eval_config.agent.engine_max_tokens == 262_144 + + +def test_explicit_agent_proxy_limit_is_preserved(): + config = SimpleNamespace(agent=SimpleNamespace(engine_max_tokens=65_536)) + _set_default_agent_engine_max_tokens(config, 262_144) + assert config.agent.engine_max_tokens == 65_536 + + +def test_missing_agent_or_generation_limit_is_ignored(): + no_agent = SimpleNamespace(agent=None) + _set_default_agent_engine_max_tokens(no_agent, 40_000) + + no_limit = SimpleNamespace(agent=SimpleNamespace(engine_max_tokens=None)) + _set_default_agent_engine_max_tokens(no_limit, None) + assert no_limit.agent.engine_max_tokens is None diff --git a/tests/test_grouped_rollout_workflow.py b/tests/test_grouped_rollout_workflow.py new file mode 100644 index 0000000000..33e440c366 --- /dev/null +++ b/tests/test_grouped_rollout_workflow.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from areal.experimental.openai import InteractionWithTokenLogpReward +from areal.experimental.openai.proxy import workflow as proxy_workflow +from areal.infra import workflow_context +from areal.infra.remote_inf_engine import GroupedRolloutWorkflow + + +class RecordingWorkflow: + def __init__(self, *, failing_indices: set[int] | None = None): + self.calls: list[dict] = [] + self.failing_indices = failing_indices or set() + self.processor_input = None + + async def arun_episode(self, _engine, data): + index = len(self.calls) + self.calls.append(data) + if index in self.failing_indices: + raise RuntimeError(f"child {index} failed") + interaction = InteractionWithTokenLogpReward() + interaction.interaction_id = f"interaction-{index}" + return {interaction.interaction_id: interaction} + + def process_group_results(self, results): + self.processor_input = results + return results + + +@pytest.mark.asyncio +async def test_grouped_workflow_isolates_child_exception_and_keeps_siblings(): + """Test that one child exception becomes None without discarding its siblings.""" + child = RecordingWorkflow(failing_indices={3}) + logger = MagicMock() + workflow = GroupedRolloutWorkflow(child, 8, logger=logger) + + merged = await workflow.arun_episode(None, {"task": "x"}) + + assert len(child.calls) == 8 + assert child.processor_input is not None + assert child.processor_input[3] is None + assert list(merged) == [f"interaction-{index}" for index in range(8) if index != 3] + logger.warning.assert_called() + + +@pytest.mark.asyncio +async def test_grouped_workflow_injects_stable_task_local_group_id(): + """Test that siblings receive one stable group ID without mutating input data.""" + child = RecordingWorkflow() + workflow = GroupedRolloutWorkflow(child, 3, logger=MagicMock()) + data = {"task": "x"} + workflow_context.set(workflow_context.WorkflowContext(task_id=41)) + + await workflow.arun_episode(None, data) + + assert data == {"task": "x"} + assert [call["group_id"] for call in child.calls] == ["41", "41", "41"] + assert all(call is not data for call in child.calls) + assert len({id(call) for call in child.calls}) == 3 + + +class InvalidRAOAgent: + async def run(self, _data): + return None + + +class FakeProxyClient: + def __init__(self, **_kwargs): + self.session_api_key = "session-key" + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + +@pytest.mark.asyncio +async def test_proxy_rejects_invalid_rao_result_without_raising(monkeypatch): + """Test that invalid RAO output returns None and logs its rejection reason.""" + workflow = proxy_workflow.OpenAIProxyWorkflow( + mode="inline", agent=InvalidRAOAgent() + ) + invalid = SimpleNamespace( + interaction_rewards={}, + root_reward=0.0, + valid=False, + group_id="task-41", + rejection_reason="judge_unavailable", + error=None, + ) + monkeypatch.setattr(workflow, "_grant_capacity", AsyncMock()) + monkeypatch.setattr(workflow, "_run_agent", AsyncMock(return_value=invalid)) + monkeypatch.setattr(proxy_workflow, "OpenAIProxyClient", FakeProxyClient) + monkeypatch.setattr( + workflow_context, "get_aiohttp_session", AsyncMock(return_value=MagicMock()) + ) + logger = MagicMock() + monkeypatch.setattr(proxy_workflow, "logger", logger) + + result = await workflow.arun_episode(None, {"task": "x"}) + + assert result is None + logger.warning.assert_called_once_with( + "Rejecting invalid RAO rollout group_id=%s: %s", + "task-41", + "judge_unavailable", + )