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,