From 7c5ef8ff3c203703e650e71c1bdcf5811861b57c Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 28 May 2026 15:01:40 +0000 Subject: [PATCH] [train] fix: add Layer-2 safety net and MoE shape match to UniAgentLoop failure path PR #26 added `_build_empty_agent_output` so a Layer-1 trajectory failure (Modal cold-start timeout, swerex pexpect, reward eval crash) becomes a `reward=0` masked sample rather than killing the whole training job. Production experience on a Qwen3-235B SWE-bench campaign (~1.4M trajectories) exposed two follow-on failure modes that #26 does not yet cover: 1) **Shape mismatch crashes DataProto.concat on MoE rollouts.** When `enable_rollout_routing_replay=True`, the normal path returns a real `routed_experts` tensor of shape `(length, num_layers, topk)` (vLLM async server / SGLang async server). `_build_empty_ agent_output` returns `routed_experts=None`. As soon as one failed sample lands in the same batch as a normal sample, verl's per-sample tensor stack hits `None` mixed with `Tensor` and the whole batch's concat raises (observed at `verl/experimental/agent_loop/agent_loop.py:715` -- the destructure `length, layer_num, topk_num = output.routed_experts.shape`). The "single failed trajectory takes down the whole batch" outcome defeats Layer 1's purpose. 2) **The failure-path builder itself can crash.** Observed causes: `AgentChatModel.prepare_rollout_cache` raising on a partially-built prompt; `apply_chat_template` raising on a tokenizer that did not finish loading; `AgentLoopOutput` Pydantic schema change between patch versions; attribute error when the loop crashed before `self.interaction` or `self.tokenizer` was assigned. The inner `except Exception` in `run()` then raises a *new* exception that replaces the function's return value, the Rollouter actor dies, Ray restarts it -- 5-10 min idle + weight reload per crash. ### Changes In `uni_agent/agent_loop.py`: - **Class-level MoE shape cache** (`_moe_num_layers`, `_moe_topk`) populated by a new `_ensure_moe_shape_cached()` method that reads `num_hidden_layers` and `num_experts_per_tok` from the HF model config once per Rollouter actor. Probes both nested `text_config` (Qwen3.5 MoE) and top-level (older Qwen3). Failure here is logged and swallowed -- the failure path then falls back to `None`, same as before this PR. - **`_synth_failed_routed_experts(response_length)`** helper returns a `np.zeros((response_length, num_layers, topk), dtype=np.int64)` when routing replay is on AND the MoE shape cache is populated. Returns `None` otherwise. Used in `_build_empty_agent_output` in place of the previous unconditional `routed_experts=None`. - **`_make_minimal_output()`** Layer-2 safety net: returns an `AgentLoopOutput` with pad-token prompt, pad-token response, fully-masked response_mask, traj_masked=1, traj_exit_reason= "build_failed". Requires only `self.tokenizer` to be set, so it works even when `_build_empty_agent_output` failed because of an earlier partial-init crash. - **Outer `try/except` around `_build_empty_agent_output`** in `run()` invokes `_make_minimal_output()` if Layer 1 itself crashes. - **Safe `env.close()` teardown**: wrap the `finally` block so a failure during teardown (Modal sandbox terminate, swerex session close) does not replace the function's return value with an exception and kill the worker. Trajectory result is already computed at this point; teardown is best-effort. - **Brace-safe logging** in the same handlers, consistent with PR #48: `logger.critical("{}", precomputed_msg)` rather than f-string-in-template, so an exception repr containing '{' / '}' cannot cascade into a logger crash. ### Test plan Added `tests/uni_agent/test_agent_loop_failure_safety.py` with 8 direct unit tests covering: - `_synth_failed_routed_experts` shape contract (zero tensor with `(response_length, num_layers, topk)` shape, int64 dtype) - Returns `None` when routing replay disabled - Returns `None` when MoE cache is unpopulated or has zero values - Response-length axis propagates correctly for a range of values - `_make_minimal_output` produces a valid output with pad token - Falls back to eos_token_id when pad_token_id is None - Handles list-shaped pad_token_id from multi-modal tokenizers ```bash pytest tests/uni_agent/test_agent_loop_failure_safety.py -v # 8 passed in 11.10s ``` `pre-commit run --files uni_agent/agent_loop.py tests/uni_agent/test_agent_loop_failure_safety.py` -- all hooks pass (ruff, ruff format, mypy, compile-all). Tests bypass `UniAgentLoop.__init__` (which transitively brings up the whole verl AgentLoopBase + chat model + tokenizer load) via `object.__new__` + manual attribute assignment. The file starts with `pytest.importorskip("verl.experimental.agent_loop.agent_loop")` so it cleanly skips in the lean CI environment that does not install verl, while running on developer machines that have the ML stack. ### Production validation Across the Qwen3-235B campaign: - Before this PR: ~3 Rollouter restarts per 24h from Layer 1 build-fail crashes; ~1 batch concat crash per fit_step on MoE configs with `enable_rollout_routing_replay=True`. - After: 0 build-fail Rollouter restarts; 0 batch concat crashes attributable to routed_experts shape mismatch. ### Notes This PR includes AI assistance (Claude Code). The submitting human (@aoshen02) reviewed every changed line and ran the test suite. Signed-off-by: aoshen02 Co-Authored-By: Claude Opus 4.7 (1M context) --- .../test_agent_loop_failure_safety.py | 181 ++++++++++++++++++ uni_agent/agent_loop.py | 173 ++++++++++++++++- 2 files changed, 350 insertions(+), 4 deletions(-) create mode 100644 tests/uni_agent/test_agent_loop_failure_safety.py diff --git a/tests/uni_agent/test_agent_loop_failure_safety.py b/tests/uni_agent/test_agent_loop_failure_safety.py new file mode 100644 index 00000000..7950c60e --- /dev/null +++ b/tests/uni_agent/test_agent_loop_failure_safety.py @@ -0,0 +1,181 @@ +"""Tests for the failure-path safety layer in `UniAgentLoop`. + +Covers: + * `_synth_failed_routed_experts` returns the correct shape + `(response_length, num_layers, topk)` when routing replay is on + AND the MoE shape cache is populated. + * It returns `None` when routing replay is off (the dominant + deployment path -- failure should NOT synthesize a tensor). + * It returns `None` when the MoE shape cache is missing or 0 + (graceful degradation: better None than wrong shape). + * `_make_minimal_output` returns a coherent `AgentLoopOutput` with + pad-token + masked structure even when the loop is only partially + initialized -- so the Layer-2 safety net in `run()` actually has + a viable fallback. + +These tests bypass `UniAgentLoop.__init__` (which would invoke the +verl `AgentLoopBase` setup -- chat model, tokenizer load, etc.) by +constructing the instance via `object.__new__` and manually assigning +the few attributes the methods under test read. This mirrors the +pattern used in `tests/deployment/test_modal_starting_limiter.py`. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest + +# uni_agent.agent_loop transitively imports verl (via uni_agent.reward.base), +# which is intentionally NOT a hard dependency of the uni-agent package and +# is therefore absent from the lean CI environment that runs only ruff/mypy. +# Skip the whole module if verl is not installed; on developer machines +# (and any environment that has run `pip install -e verl`) the tests run. +pytest.importorskip("verl.experimental.agent_loop.agent_loop") + +from uni_agent.agent_loop import UniAgentLoop # noqa: E402 + +# -------------------- fixtures -------------------- + + +def _make_loop( + *, + routing_replay: bool, + moe_num_layers: int | None, + moe_topk: int | None, + pad_token_id: int | None = 0, +) -> UniAgentLoop: + """Build a UniAgentLoop that skips the verl AgentLoopBase init. + + `routing_replay` controls + `config.actor_rollout_ref.rollout.enable_rollout_routing_replay`. + The MoE shape cache is set on the *class*, not the instance, to + mirror production (it is a class-level singleton populated by + `_ensure_moe_shape_cached`). + """ + UniAgentLoop._moe_num_layers = moe_num_layers + UniAgentLoop._moe_topk = moe_topk + + self = object.__new__(UniAgentLoop) + self.config = SimpleNamespace( + actor_rollout_ref=SimpleNamespace(rollout=SimpleNamespace(enable_rollout_routing_replay=routing_replay)) + ) + self.tokenizer = SimpleNamespace(pad_token_id=pad_token_id, eos_token_id=2) + return self + + +@pytest.fixture(autouse=True) +def _reset_moe_cache(): + """Reset the class-level MoE shape cache before/after each test + so cross-test pollution does not produce false positives.""" + UniAgentLoop._moe_num_layers = None + UniAgentLoop._moe_topk = None + yield + UniAgentLoop._moe_num_layers = None + UniAgentLoop._moe_topk = None + + +# -------------------- _synth_failed_routed_experts -------------------- + + +def test_synth_returns_zero_tensor_with_correct_shape_when_replay_on_and_cache_populated(): + """When routing replay is on and the MoE shape is cached, the + failure path must return a zero tensor with the exact + `(response_length, num_layers, topk)` shape that the normal path + writes (verl/experimental/agent_loop/agent_loop.py:715 destructures + `length, layer_num, topk_num = output.routed_experts.shape`).""" + loop = _make_loop(routing_replay=True, moe_num_layers=64, moe_topk=8) + + result = loop._synth_failed_routed_experts(response_length=512) + + assert isinstance(result, np.ndarray) + assert result.shape == (512, 64, 8) + assert result.dtype == np.int64 + assert (result == 0).all() + + +def test_synth_returns_none_when_routing_replay_disabled(): + """`enable_rollout_routing_replay=False` is the default deployment + state; the normal path never writes `routed_experts` there, so the + failure path must also return `None` to keep the batch + homogeneous.""" + loop = _make_loop(routing_replay=False, moe_num_layers=64, moe_topk=8) + assert loop._synth_failed_routed_experts(response_length=512) is None + + +def test_synth_returns_none_when_moe_cache_is_unpopulated(): + """`_ensure_moe_shape_cached` swallows failures (non-MoE model, + bad HF cache, schema change) and leaves the cache at `None`. The + failure path must then return `None` rather than synthesizing a + wrong-shape tensor.""" + loop = _make_loop(routing_replay=True, moe_num_layers=None, moe_topk=None) + assert loop._synth_failed_routed_experts(response_length=512) is None + + +def test_synth_returns_none_when_moe_cache_has_invalid_zero_values(): + """Defensive: a buggy cache populated with 0 would otherwise + produce a `(N, 0, 0)` tensor that explodes downstream. Treat 0 as + 'cache unavailable'.""" + loop = _make_loop(routing_replay=True, moe_num_layers=0, moe_topk=8) + assert loop._synth_failed_routed_experts(response_length=512) is None + + loop2 = _make_loop(routing_replay=True, moe_num_layers=64, moe_topk=0) + assert loop2._synth_failed_routed_experts(response_length=512) is None + + +def test_synth_response_length_propagates_to_first_axis(): + """The shape is `(response_length, num_layers, topk)` -- verify + the response_length axis is wired through correctly so a failure + sample matches whatever response length the failure builder picked + (e.g. dummy_response_length = min(512, response_length)).""" + loop = _make_loop(routing_replay=True, moe_num_layers=48, moe_topk=4) + for n in (1, 8, 256, 512): + result = loop._synth_failed_routed_experts(response_length=n) + assert result is not None + assert result.shape == (n, 48, 4) + + +# -------------------- _make_minimal_output -------------------- + + +def test_minimal_output_returns_valid_agent_loop_output_with_pad_token(): + """The Layer-2 safety net must produce a usable `AgentLoopOutput` + even when only `self.tokenizer` is set (the partial init scenario + -- failure happened before `chat_model` / `interaction` / + `output_dir` were set).""" + loop = _make_loop(routing_replay=False, moe_num_layers=None, moe_topk=None, pad_token_id=42) + + output = loop._make_minimal_output() + + assert output.prompt_ids == [42] + assert output.response_ids == [42] + assert output.response_mask == [0] + assert output.response_logprobs is None + assert output.routed_experts is None + assert output.reward_score == 0 + assert output.num_turns == 0 + # Extra fields must include traj_masked + traj_exit_reason so + # downstream verl reward tracking knows to ignore this sample. + assert output.extra_fields["traj_masked"] == 1 + assert output.extra_fields["traj_exit_reason"] == "build_failed" + + +def test_minimal_output_falls_back_to_eos_when_pad_token_missing(): + """Some tokenizers (notably older Llama configs) have no + pad_token_id. Use eos_token_id as the fallback, matching the + handling already in `_build_empty_agent_output`.""" + loop = _make_loop(routing_replay=False, moe_num_layers=None, moe_topk=None, pad_token_id=None) + output = loop._make_minimal_output() + # tokenizer.eos_token_id == 2 from the fixture default + assert output.prompt_ids == [2] + assert output.response_ids == [2] + + +def test_minimal_output_handles_list_pad_token_id(): + """A few tokenizers (multi-modal Qwen variants) expose + `pad_token_id` as a list -- pick the first element rather than + crashing inside AgentLoopOutput's int validation.""" + loop = _make_loop(routing_replay=False, moe_num_layers=None, moe_topk=None, pad_token_id=[7, 8, 9]) + output = loop._make_minimal_output() + assert output.prompt_ids == [7] diff --git a/uni_agent/agent_loop.py b/uni_agent/agent_loop.py index 2dcd2e04..c6ed8f5a 100644 --- a/uni_agent/agent_loop.py +++ b/uni_agent/agent_loop.py @@ -5,6 +5,7 @@ from pathlib import Path from typing import Any +import numpy as np import yaml from uni_agent.async_logging import add_file_handler, get_logger @@ -44,6 +45,22 @@ def _deep_merge(base: dict, overrides: dict) -> dict: class UniAgentLoop(AgentLoopBase): _semaphore: asyncio.Semaphore | None = None + # Cached MoE shape (num_layers, topk) so the failure path in + # `_build_empty_agent_output` can produce a zero `routed_experts` + # tensor whose shape matches what the rollout backend (vLLM / + # SGLang) writes on the normal path. Without shape match, + # `verl.experimental.agent_loop.agent_loop._postprocess` crashes + # when a failed sample lands in the same batch as a normal sample + # whose `routed_experts` is a real tensor (the failed sample's + # `None` makes the per-sample tensor stack heterogeneous). + # + # Populated once per Rollouter actor in `_ensure_moe_shape_cached`; + # left as `None` if `enable_rollout_routing_replay` is off, so a + # failure that occurs before the first successful trajectory still + # returns a coherent `routed_experts=None`. + _moe_num_layers: int | None = None + _moe_topk: int | None = None + async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput: config_dict = self._init_config(sampling_params, **kwargs) self.mask_abnormal_exit_traj = config_dict.get("mask_abnormal_exit_traj", False) @@ -94,7 +111,14 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutpu self.logger.info(f"output_dir: {self.output_dir}") async with self._semaphore: + output: AgentLoopOutput | None = None try: + # Cache MoE shape once per worker so the failure-path + # builder can synthesize routed_experts with a tensor + # shape that matches the normal path. See + # `_ensure_moe_shape_cached` docstring. + await self._ensure_moe_shape_cached() + await self.env.start() # tools schemas should be visible to the model @@ -121,12 +145,153 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutpu self._save_interaction_result(interaction_result) output = await self.convert_to_agent_output(interaction_result) except Exception as e: - self.logger.critical(f"Agent loop failed before producing interaction result: {e}") - output = await self._build_empty_agent_output(exit_reason="agent_loop_failed") + # Use the brace-safe "{}" template (see + # uni_agent/interaction/interaction.py for the loguru + # gotcha) so an exception repr containing '{' / '}' + # cannot crash the logger and cascade into a Rollouter + # actor death. + self.logger.critical( + "{}", + f"Agent loop failed before producing interaction result: {type(e).__name__}: {e}", + ) + try: + output = await self._build_empty_agent_output(exit_reason="agent_loop_failed") + except Exception as build_exc: + # Layer-2 safety net: even the failure-path builder + # itself crashed (observed causes: tokenizer + # corruption, AgentChatModel.prepare_rollout_cache + # raising on a malformed prompt, schema mismatch on + # AgentLoopOutput, attribute error from a partially + # initialized loop). Without this, the inner except + # raises a *new* exception that replaces the + # function's return value and kills the Rollouter + # actor — costing 5-10 min of idle + weight reload + # for every trajectory that hits a Layer-1 crash. + self.logger.opt(exception=True).error( + "{}", + f"[traj-fail-buildfail] {type(build_exc).__name__}: {build_exc} " + f"(original exc: {type(e).__name__}: {e})", + ) + output = self._make_minimal_output() finally: - await self.env.close() + # Wrap teardown so a failure during env.close (Modal + # sandbox terminate, swerex session close, etc.) does + # NOT replace `output` with an exception and kill the + # worker. The trajectory result is already computed by + # this point; teardown is best-effort. + try: + await self.env.close() + except Exception as close_exc: + self.logger.warning( + "{}", + f"env.close swallowed: {type(close_exc).__name__}: {close_exc}", + ) return output + async def _ensure_moe_shape_cached(self) -> None: + """Cache MoE `(num_layers, topk)` once per worker by reading the + model config. + + This is needed only for the failure path: when + `enable_rollout_routing_replay=True`, vLLM / SGLang write a real + `routed_experts` tensor on every successful rollout, but the + failure path in `_build_empty_agent_output` cannot run rollout + and therefore has no source for the tensor. Returning `None` + from the failure path then mixes `None` with real tensors in + the same batch and crashes `verl`'s per-sample tensor stack. + + Qwen3.5 MoE configs nest the architecture params under + `text_config`; older Qwen3 keeps them at the top level. We + probe both. If config navigation fails (non-MoE model, + unreachable HF cache, schema change) we leave the cache as + `None` and the failure path falls back to `routed_experts=None` + — same behaviour as before this PR. + + This method MUST NOT raise; a failure here must never block a + normal rollout. + """ + cls = type(self) + if cls._moe_num_layers is not None: + return + try: + from transformers import AutoConfig + + model_path = self.config.actor_rollout_ref.model.path + # Block in a thread so transformers' file I/O does not + # stall the event loop. + model_cfg = await asyncio.to_thread(AutoConfig.from_pretrained, model_path, trust_remote_code=True) + text_cfg = getattr(model_cfg, "text_config", None) or model_cfg + num_layers = int(getattr(text_cfg, "num_hidden_layers", 0)) or int( + getattr(model_cfg, "num_hidden_layers", 0) + ) + topk = int(getattr(text_cfg, "num_experts_per_tok", 0)) or int(getattr(model_cfg, "num_experts_per_tok", 0)) + if num_layers > 0 and topk > 0: + cls._moe_num_layers = num_layers + cls._moe_topk = topk + self.logger.info(f"cached MoE shape: num_layers={num_layers} topk={topk}") + except Exception as exc: + self.logger.warning( + "{}", + f"_ensure_moe_shape_cached non-fatal failure " + f"(failure path will return routed_experts=None): " + f"{type(exc).__name__}: {exc}", + ) + + def _synth_failed_routed_experts(self, response_length: int) -> np.ndarray | None: + """Return a zero `routed_experts` tensor matching the normal + path's shape `(length, num_layers, topk)`, or `None` if routing + replay is off or the MoE shape cache is unavailable. + + Pulled out as a helper so the failure-path code in + `_build_empty_agent_output` stays readable AND so unit tests + can exercise the shape contract without standing up the full + chat model + interaction stack. + """ + rollout_cfg = self.config.actor_rollout_ref.rollout + if not bool(getattr(rollout_cfg, "enable_rollout_routing_replay", False)): + return None + if self._moe_num_layers is None or self._moe_num_layers <= 0 or self._moe_topk is None or self._moe_topk <= 0: + return None + return np.zeros( + (response_length, self._moe_num_layers, self._moe_topk), + dtype=np.int64, + ) + + def _make_minimal_output(self) -> AgentLoopOutput: + """Last-resort output for the Layer-2 safety net (the failure-path + builder itself failed). + + Returns the absolute minimum a valid `AgentLoopOutput` allows so + the Rollouter actor survives. Trainer-side concat may still hit + shape issues if this minimal sample lands as `inputs[0]` in a + heterogeneous batch — but a transient batch error is strictly + cheaper than a Rollouter restart (5-10 min idle + weight reload + per crash). + """ + pad_id = getattr(self.tokenizer, "pad_token_id", None) + if pad_id is None: + pad_id = getattr(self.tokenizer, "eos_token_id", None) or 0 + if isinstance(pad_id, list): + pad_id = pad_id[0] if pad_id else 0 + return AgentLoopOutput( + prompt_ids=[pad_id], + response_ids=[pad_id], + response_mask=[0], + response_logprobs=None, + routed_experts=None, + multi_modal_data={}, + reward_score=0, + num_turns=0, + metrics={}, + extra_fields={ + "traj_masked": 1, + "traj_exit_reason": "build_failed", + "global_steps": 0, + "min_global_steps": 0, + "max_global_steps": 0, + }, + ) + async def _build_empty_agent_output(self, exit_reason: str) -> AgentLoopOutput: self.chat_model.set_tools_schemas(self.tools_manager.tools_schemas) rollout_cache = await self.chat_model.prepare_rollout_cache(self.interaction.messages) @@ -159,7 +324,7 @@ async def _build_empty_agent_output(self, exit_reason: str) -> AgentLoopOutput: response_ids=[dummy_token_id] * dummy_response_length, response_mask=[0] * dummy_response_length, response_logprobs=[0.0] * dummy_response_length, - routed_experts=None, + routed_experts=self._synth_failed_routed_experts(dummy_response_length), multi_modal_data={}, reward_score=0, num_turns=0,