[train] fix: add Layer-2 safety net and MoE shape match to UniAgentLoop failure path

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]

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,