[train] feat: Multiple Linear Chains for Gateway Sessions

docs/source/start/agent_train.md

@@ -152,6 +152,8 @@ Start with the script defaults, then tune these first:
 - `max_prompt_length`, `max_response_length`: context budget for the agent trajectory.
 - `staleness_threshold`, `trigger_parameter_sync_step`, `require_batches`, `partial_rollout`: fully async scheduling and weight synchronization behavior.
 
+Gateway sessions keep multiple active linear chains by default for subagent, compaction, and retry-style flows. Backend generation calls within one session can overlap by default, while prepare, decode, and commit remain serialized under the session `request_lock`. Successful outcomes are accepted in completion order, so the final accepted trajectory used as the scoring target can be timing-dependent. Its score is broadcast to all trajectories from that session.
+
 For MoE or large models, also check tensor, pipeline, context, and expert parallelism settings such as `GEN_TP`, `TP`, `PP`, `CP`, and `EP` in `train_qwen3p5_moe.sh`.
 
 ---

tests/uni_agent/framework/test_generate_sequences_on_cpu.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+from types import SimpleNamespace
+
 import pytest
 
 from uni_agent.framework.framework import OpenAICompatibleAgentFramework
-from uni_agent.gateway.session import SessionHandle, Trajectory
+from uni_agent.gateway.session import GatewaySession, MessageCodec, SessionHandle, Trajectory
 from verl.utils import tensordict_utils as tu
 
 _RUNNER_CALLS = []
@@ -83,7 +85,16 @@ async def _build_framework_with_agent_runners(
             "actor_rollout_ref": {
                 "rollout": {
                     "n": n,
-                    "val_kwargs": {"n": val_n},
+                    "temperature": 0.7,
+                    "top_p": 0.8,
+                    "top_k": 20,
+                    "calculate_log_probs": True,
+                    "val_kwargs": {
+                        "n": val_n,
+                        "temperature": 0,
+                        "top_p": 0.95,
+                        "top_k": -1,
+                    },
                     "custom": {"agent_framework": agent_framework_cfg},
                 }
             }
@@ -96,6 +107,65 @@ async def _build_framework_with_agent_runners(
     )
 
 
+@pytest.mark.parametrize(
+    ("data_config", "expected_chat_template_kwargs"),
+    [
+        ({}, {}),
+        ({"apply_chat_template_kwargs": {"thinking": True}}, {"thinking": True}),
+    ],
+)
+def test_build_gateway_manager_wires_gateway_config_defaults(
+    monkeypatch,
+    data_config,
+    expected_chat_template_kwargs,
+):
+    from omegaconf import OmegaConf
+
+    from uni_agent.framework import entry as entry_module
+
+    class _ModelConfig:
+        tokenizer = object()
+        processor = None
+
+    captured = {}
+
+    class _FakeGatewayManager:
+        def __init__(self, *, llm_client, gateway_count, gateway_actor_config):
+            captured["llm_client"] = llm_client
+            captured["gateway_count"] = gateway_count
+            captured["gateway_actor_config"] = gateway_actor_config
+
+    monkeypatch.setattr(entry_module, "omega_conf_to_dataclass", lambda _config: _ModelConfig())
+    monkeypatch.setattr(entry_module, "GatewayManager", _FakeGatewayManager)
+
+    llm_client = object()
+    config = OmegaConf.create(
+        {
+            "data": data_config,
+            "actor_rollout_ref": {
+                "model": {},
+                "rollout": {
+                    "prompt_length": 128,
+                    "response_length": 64,
+                    "multi_turn": {"format": "hermes"},
+                    "custom": {"agent_framework": {"gateway_count": 2}},
+                },
+            },
+        }
+    )
+
+    manager = entry_module.build_gateway_manager(config=config, llm_client=llm_client)
+
+    assert isinstance(manager, _FakeGatewayManager)
+    assert captured["llm_client"] is llm_client
+    assert captured["gateway_count"] == 2
+    assert captured["gateway_actor_config"].prompt_length == 128
+    assert captured["gateway_actor_config"].response_length == 64
+    assert captured["gateway_actor_config"].tool_parser_name == "hermes"
+    assert isinstance(captured["gateway_actor_config"].apply_chat_template_kwargs, dict)
+    assert captured["gateway_actor_config"].apply_chat_template_kwargs == expected_chat_template_kwargs
+
+
 class _FakeTransferQueue:
     def __init__(self):
         self.puts = []
@@ -131,6 +201,7 @@ class _FakeGatewayManager:
     def __init__(self, finalized_by_session_prefix: dict[str, list[Trajectory]]):
         self._finalized_by_prefix = finalized_by_session_prefix
         self.created_sessions = []
+        self.created_session_kwargs = []
         self.finalized_sessions = []
         self.aborted_sessions = []
 
@@ -142,6 +213,7 @@ def _lookup(self, session_id: str) -> list[Trajectory]:
 
     async def create_session(self, session_id: str, **kwargs):
         self.created_sessions.append(session_id)
+        self.created_session_kwargs.append(dict(kwargs))
         return SessionHandle(
             session_id=session_id,
             base_url=f"http://fake/{session_id}/v1",
@@ -156,20 +228,29 @@ async def abort_session(self, session_id: str) -> None:
         self.aborted_sessions.append(session_id)
 
 
-def _build_prompts(count: int = 2, *, global_steps: int = 7, validate: bool = False):
+def _build_prompts(
+    count: int = 2,
+    *,
+    global_steps: int = 7,
+    validate: bool = False,
+    do_sample: bool | None = None,
+):
     non_tensor_dict = {"global_steps": global_steps}
     if validate:
         non_tensor_dict["validate"] = True
+    tensor_dict = {
+        "raw_prompt": [[{"role": "user", "content": f"sample {i}"}] for i in range(count)],
+        "uid": [f"uid-{i}" for i in range(count)],
+        "data_source": ["deepeyes"] * count,
+        "reward_model": [{"ground_truth": f"answer-{i}"} for i in range(count)],
+        "extra_info": [{"index": i} for i in range(count)],
+        "tools_kwargs": [{"tool": i} for i in range(count)],
+        "agent_name": ["deepeyes"] * count,
+    }
+    if do_sample is not None:
+        tensor_dict["__do_sample__"] = [do_sample] * count
     return tu.get_tensordict(
-        tensor_dict={
-            "raw_prompt": [[{"role": "user", "content": f"sample {i}"}] for i in range(count)],
-            "uid": [f"uid-{i}" for i in range(count)],
-            "data_source": ["deepeyes"] * count,
-            "reward_model": [{"ground_truth": f"answer-{i}"} for i in range(count)],
-            "extra_info": [{"index": i} for i in range(count)],
-            "tools_kwargs": [{"tool": i} for i in range(count)],
-            "agent_name": ["deepeyes"] * count,
-        },
+        tensor_dict=tensor_dict,
         non_tensor_dict=non_tensor_dict,
     )
 
@@ -179,6 +260,7 @@ def _trajectory(
     prompt_ids: list[int] | None = None,
     response_ids: list[int] | None = None,
     response_logprobs: list[float] | None = None,
+    reward_info: dict[str, object] | None = None,
     num_turns: int = 2,
     extra_fields: dict[str, object] | None = None,
 ):
@@ -189,6 +271,7 @@ def _trajectory(
         response_ids=response_ids,
         response_mask=[1] * len(response_ids),
         response_logprobs=response_logprobs,
+        reward_info=dict(reward_info or {}),
         reward_score=None,
         num_turns=num_turns,
         multi_modal_data={"images": ["raw-image-should-not-be-written"]},
@@ -258,6 +341,68 @@ async def test_agent_runners_registry_materializes_runners_and_selects_by_agent_
     assert all("gateway_manager" not in call["kwargs"] for call in calls)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("validate", "do_sample", "expected_sampling_params"),
+    [
+        (
+            False,
+            None,
+            {
+                "temperature": 0.7,
+                "top_p": 0.8,
+                "top_k": 20,
+                "repetition_penalty": 1.0,
+                "logprobs": True,
+            },
+        ),
+        (
+            True,
+            None,
+            {
+                "temperature": 0,
+                "top_p": 0.95,
+                "top_k": -1,
+                "repetition_penalty": 1.0,
+                "logprobs": True,
+            },
+        ),
+        (
+            False,
+            False,
+            {
+                "temperature": 0,
+                "top_p": 1.0,
+                "top_k": -1,
+                "repetition_penalty": 1.0,
+                "logprobs": True,
+            },
+        ),
+    ],
+)
+async def test_framework_binds_sampling_defaults_to_gateway_sessions(
+    fake_tq,
+    validate,
+    do_sample,
+    expected_sampling_params,
+):
+    runtime = _FakeGatewayManager({"session-0-0": [_trajectory()]})
+    framework = await _build_framework_with_agent_runners(
+        agent_runners={"runner": _inline_runner_config(_async_noop_runner)},
+        gateway_manager=runtime,
+    )
+
+    await framework.generate_sequences(
+        _build_prompts(
+            count=1,
+            validate=validate,
+            do_sample=do_sample,
+        )
+    )
+
+    assert runtime.created_session_kwargs == [{"sampling_params": expected_sampling_params}]
+
+
 @pytest.mark.asyncio
 async def test_generate_sequences_writes_tq_schema_for_each_session(monkeypatch, fake_tq):
     """Full ``generate_sequences`` path writes one TQ batch per successful
@@ -333,6 +478,135 @@ async def test_generate_sequences_writes_tq_schema_for_each_session(monkeypatch,
     assert "multi_modal_data" not in fields.keys()
 
 
+@pytest.mark.asyncio
+async def test_generate_sequences_preserves_sorted_trajectory_order_and_rewards_final_target(fake_tq):
+    class _Tokenizer:
+        def apply_chat_template(self, messages, tokenize=True, add_generation_prompt=True, tools=None, **kwargs):
+            parts = []
+            for message in messages:
+                content = message.get("content", "")
+                if content is None:
+                    content = ""
+                parts.append(f"{message['role']}:{content}\n")
+            if add_generation_prompt:
+                parts.append("assistant:")
+            text = "".join(parts)
+            if tokenize:
+                return [ord(char) for char in text]
+            return text
+
+        def decode(self, token_ids, skip_special_tokens=True):
+            if hasattr(token_ids, "tolist"):
+                token_ids = token_ids.tolist()
+            return "".join(
+                chr(int(token_id.item() if hasattr(token_id, "item") else token_id)) for token_id in token_ids
+            )
+
+        def encode(self, text, add_special_tokens=False):
+            return [ord(char) for char in text]
+
+    class _Backend:
+        def __init__(self, steps):
+            self.steps = list(steps)
+
+        async def generate(self, request_id, *, prompt_ids, sampling_params, image_data=None, video_data=None):
+            del request_id, prompt_ids, sampling_params, image_data, video_data
+            text = self.steps.pop(0)
+            token_ids = [ord(char) for char in text]
+            return SimpleNamespace(
+                token_ids=token_ids,
+                log_probs=[-0.1] * len(token_ids),
+                stop_reason="completed",
+            )
+
+    class _ComputeScoreRemote:
+        def __init__(self):
+            self.calls = []
+
+        async def remote(self, data):
+            self.calls.append(data)
+            return {"reward_score": 0.5, "reward_extra_info": {"target": "final-main"}}
+
+    class _StubWorker:
+        def __init__(self):
+            self.compute_score = _ComputeScoreRemote()
+
+    worker = _StubWorker()
+    tokenizer = _Tokenizer()
+    real_session = GatewaySession(
+        SessionHandle(session_id="real-sorted-order"),
+        MessageCodec(tokenizer),
+        response_length=len("MAIN1") + 1,
+    )
+    backend = _Backend(["MAIN1", "SUB"])
+    main_first = [
+        {"role": "system", "content": "You are helpful."},
+        {"role": "user", "content": "main"},
+    ]
+    subagent = [
+        {"role": "system", "content": "You are a focused subagent."},
+        {"role": "user", "content": "sub"},
+    ]
+    main_too_long = [
+        {"role": "system", "content": "You are helpful."},
+        {"role": "user", "content": "main"},
+        {"role": "assistant", "content": "MAIN1"},
+        {"role": "user", "content": "too long"},
+    ]
+
+    await real_session.run_generation({"model": "dummy-model", "messages": main_first}, backend)
+    await real_session.run_generation({"model": "dummy-model", "messages": subagent}, backend)
+    outcome = await real_session.run_generation({"model": "dummy-model", "messages": main_too_long}, backend)
+    await real_session.set_reward_info({"branch": "main", "target": "final-main"})
+    trajectories = await real_session.finalize()
+
+    assert outcome.finish_reason == "length"
+    assert backend.steps == []
+    assert [tokenizer.decode(trajectory.response_ids) for trajectory in trajectories] == ["SUB", "MAIN1"]
+    assert trajectories[-1].extra_fields == {"finish_reason": "length"}
+    runtime = _FakeGatewayManager({"session-0-0": trajectories})
+    framework = await _build_framework_with_agent_runners(
+        agent_runners={"runner": _inline_runner_config(_async_noop_runner)},
+        gateway_manager=runtime,
+        reward_loop_worker_handles=[worker],
+        n=1,
+        val_n=1,
+    )
+
+    await framework.generate_sequences(_build_prompts(count=1, global_steps=12))
+
+    assert len(worker.compute_score.calls) == 1
+    data = worker.compute_score.calls[0]
+    final_trajectory = trajectories[-1]
+    assert data.batch["prompts"].tolist() == [final_trajectory.prompt_ids]
+    assert data.batch["responses"].tolist() == [final_trajectory.response_ids]
+    assert data.non_tensor_batch["extra_info"].tolist() == [{"index": 0, "branch": "main", "target": "final-main"}]
+    assert data.non_tensor_batch["__num_turns__"].tolist() == [final_trajectory.num_turns]
+
+    assert len(fake_tq.batch_puts) == 1
+    batch_put = fake_tq.batch_puts[0]
+    assert batch_put["keys"] == ["uid-0_0_0", "uid-0_0_1"]
+    assert [tag.get("finish_reason") for tag in batch_put["tags"]] == [None, "length"]
+    fields = batch_put["fields"]
+    assert [fields["prompts"][i].tolist() for i in range(len(trajectories))] == [
+        trajectory.prompt_ids for trajectory in trajectories
+    ]
+    assert [fields["responses"][i].tolist() for i in range(len(trajectories))] == [
+        trajectory.response_ids for trajectory in trajectories
+    ]
+    assert [tokenizer.decode(fields["responses"][i]) for i in range(len(trajectories))] == ["SUB", "MAIN1"]
+    for index, trajectory in enumerate(trajectories):
+        assert fields["rollout_log_probs"][index].tolist() == pytest.approx(trajectory.response_logprobs)
+    assert [fields["rm_scores"][i].tolist() for i in range(len(trajectories))] == [
+        [0.0] * (len(trajectory.response_ids) - 1) + [0.5] for trajectory in trajectories
+    ]
+    assert tu.get(fields, "reward_extra_info") == [
+        {"target": "final-main"},
+        {"target": "final-main"},
+    ]
+    assert fake_tq.puts == [{"key": "uid-0", "partition_id": "train", "tag": {"status": "finished"}}]
+
+
 @pytest.mark.asyncio
 async def test_generate_sequences_keeps_successful_sessions_when_one_session_fails(fake_tq):
     """A failed rollout session aborts only that session; other successful