Fix multi-turn truncation

tests/common/experience_extraction_test.py

@@ -1,15 +1,115 @@
+import asyncio
 import io
 from types import SimpleNamespace
 from unittest import TestCase
+from unittest.mock import AsyncMock, MagicMock
 
 import numpy as np
 import pybase64
 import torch
 
-from trinity.common.models.experience_extraction import convert_api_output_to_experience
+from trinity.common.models.experience_extraction import (
+    HistoryRecordingStream,
+    convert_api_output_to_experience,
+)
+from trinity.common.models.vllm_model import vLLMRolloutModel
 
 
 class TestExperienceExtraction(TestCase):
+    def test_vllm_generate_maps_finish_reason_per_output(self):
+        def logprob(value: float):
+            return {0: SimpleNamespace(logprob=value)}
+
+        model = vLLMRolloutModel.__new__(vLLMRolloutModel)
+        model.tokenizer = MagicMock()
+        model.tokenizer.decode.return_value = "prompt"
+        model._handle_prompt_truncation = MagicMock(return_value=([1, 2], True))
+        model._extract_routed_experts = MagicMock(return_value=None)
+        model._generate_internal = AsyncMock(
+            return_value=SimpleNamespace(
+                prompt_token_ids=[1, 2],
+                outputs=[
+                    SimpleNamespace(
+                        token_ids=[3],
+                        logprobs=[logprob(-0.1)],
+                        text="truncated",
+                        finish_reason="length",
+                    ),
+                    SimpleNamespace(
+                        token_ids=[4],
+                        logprobs=[logprob(-0.2)],
+                        text="complete",
+                        finish_reason="stop",
+                    ),
+                ],
+            )
+        )
+
+        experiences = asyncio.run(model.generate("prompt", n=2))
+
+        self.assertEqual(experiences[0].truncate_status, "response_truncated")
+        self.assertIsNone(experiences[1].truncate_status)
+
+    def test_convert_completion_output_maps_finish_reason(self):
+        output = SimpleNamespace(
+            prompt_token_ids=[1, 2],
+            choices=[
+                SimpleNamespace(
+                    token_ids=[3],
+                    logprobs=None,
+                    message=SimpleNamespace(content="truncated"),
+                    finish_reason="length",
+                    routed_experts=None,
+                ),
+                SimpleNamespace(
+                    token_ids=[4],
+                    logprobs=None,
+                    message=SimpleNamespace(content="complete"),
+                    finish_reason="stop",
+                    routed_experts=None,
+                ),
+            ],
+        )
+
+        experiences = convert_api_output_to_experience(output)
+
+        self.assertEqual(experiences[0].truncate_status, "response_truncated")
+        self.assertIsNone(experiences[1].truncate_status)
+
+    def test_stream_conversion_preserves_final_finish_reason(self):
+        chunks = [
+            SimpleNamespace(
+                prompt_token_ids=[1, 2],
+                choices=[
+                    SimpleNamespace(
+                        index=0,
+                        token_ids=[3],
+                        logprobs=None,
+                        delta=SimpleNamespace(content="go"),
+                        finish_reason=None,
+                    )
+                ],
+            ),
+            SimpleNamespace(
+                choices=[
+                    SimpleNamespace(
+                        index=0,
+                        token_ids=None,
+                        logprobs=None,
+                        delta=SimpleNamespace(content=None),
+                        finish_reason="length",
+                    )
+                ]
+            ),
+        ]
+        history = []
+
+        list(HistoryRecordingStream(iter(chunks), history))
+
+        self.assertEqual(len(history), 1)
+        self.assertEqual(history[0].response_text, "go")
+        self.assertEqual(history[0].truncate_status, "response_truncated")
+
     def test_convert_completion_output_extracts_sglang_routed_experts(self):
         routed_experts = torch.tensor(
             [

tests/explorer/step_wise_workflow_test.py

@@ -182,6 +182,36 @@ def test_workflows_stop_at_max_env_steps(self) -> None:
                 experiences = workflow.run()
             self.assertEqual(len(experiences), 3)
 
+    def test_workflows_stop_after_response_truncation(self) -> None:
+        for workflow_cls in _dummy_workflows:
+            call_count = 0
+
+            def next_experience():
+                nonlocal call_count
+                call_count += 1
+                return [
+                    Experience(
+                        tokens=Tensor([0, 0]),
+                        prompt_length=1,
+                        truncate_status=("response_truncated" if call_count == 2 else None),
+                    )
+                ]
+
+            self.model.extract_experience_from_history.side_effect = next_experience
+            task = Task(
+                workflow=workflow_cls,
+                repeat_times=self.taskset_config.repeat_times,
+                workflow_args={"max_env_steps": 10, "actual_steps": 10},
+            )
+            workflow = task.to_workflow(model=self.model)
+            if workflow.asynchronous:
+                experiences = asyncio.run(workflow.run_async())
+            else:
+                experiences = workflow.run()
+
+            self.assertEqual(len(experiences), 2)
+            self.assertEqual(call_count, 2)
+
     def test_workflows_raise_error(self) -> None:
         self.model.enable_history = False
         for workflow in _dummy_workflows:

tests/explorer/workflow_test.py

@@ -31,6 +31,7 @@
 from trinity.common.models.model import ModelWrapper
 from trinity.common.workflows import WORKFLOWS, Workflow
 from trinity.common.workflows.customized_math_workflows import MathBoxedWorkflow
+from trinity.common.workflows.envs.alfworld.alfworld_workflow import AlfworldWorkflow
 from trinity.common.workflows.eval_workflow import MathEvalWorkflow
 from trinity.common.workflows.workflow import MathWorkflow, MultiTurnWorkflow, Task
 from trinity.explorer.workflow_runner import WorkflowRunner
@@ -192,6 +193,32 @@ async def run_async(self):
 
 
 class WorkflowTest(unittest.TestCase):
+    def test_alfworld_stops_before_env_step_after_response_truncation(self) -> None:
+        workflow = AlfworldWorkflow.__new__(AlfworldWorkflow)
+        workflow.max_env_steps = 3
+        workflow.get_model_response = AsyncMock(
+            return_value=[
+                MockResponse(
+                    "go",
+                    truncate_status="response_truncated",
+                )
+            ]
+        )
+        final_experience = Experience(tokens=Tensor([0, 1]), prompt_length=1)
+        workflow.process_messages_to_experience_async = AsyncMock(return_value=final_experience)
+        env = MagicMock()
+        env.reset.return_value = ("initial observation", {})
+
+        experiences = asyncio.run(workflow.generate_env_inference_samples(env))
+
+        env.step.assert_not_called()
+        env.close.assert_called_once()
+        self.assertEqual(experiences, [final_experience])
+        self.assertEqual(
+            workflow.process_messages_to_experience_async.call_args.kwargs["truncate_status"],
+            "response_truncated",
+        )
+
     def test_math_workflow(self) -> None:
         model = MagicMock()
         model.chat.return_value = [

trinity/common/models/experience_extraction.py

@@ -172,6 +172,9 @@ def _convert_completion_output_to_experience(
             logprobs=extract_logprobs(choice),
             prompt_length=len(output.prompt_token_ids),
             response_text=getattr(choice.message, "content", None),
+            truncate_status=(
+                "response_truncated" if getattr(choice, "finish_reason", None) == "length" else None
+            ),
             routed_experts=_extract_completion_routed_experts(
                 output,
                 choice,
@@ -184,7 +187,7 @@ def _convert_completion_output_to_experience(
     ]
 
 
-def _convert_stream_chunks_to_experience(chunks: Sequence[Any]) -> List[Experience]:
+def _convert_stream_chunks_to_experience(chunks: Sequence[Any]) -> List[Experience]:  # noqa
     prompt_token_ids: Optional[List[int]] = None
     by_choice: Dict[int, Dict[str, Any]] = {}
 
@@ -201,9 +204,14 @@ def _convert_stream_chunks_to_experience(chunks: Sequence[Any]) -> List[Experien
                     "token_ids": [],
                     "logprobs": [],
                     "response_text_parts": [],
+                    "finish_reason": None,
                 }
             data = by_choice[idx]
 
+            finish_reason = getattr(choice, "finish_reason", None)
+            if finish_reason is not None:
+                data["finish_reason"] = finish_reason
+
             token_ids = getattr(choice, "token_ids", None)
             if token_ids is not None:
                 data["token_ids"].extend(token_ids)
@@ -240,6 +248,9 @@ def _convert_stream_chunks_to_experience(chunks: Sequence[Any]) -> List[Experien
                 logprobs=torch.tensor(data["logprobs"], dtype=torch.float32),
                 prompt_length=len(prompt_token_ids),
                 response_text=response_text,
+                truncate_status=(
+                    "response_truncated" if data["finish_reason"] == "length" else None
+                ),
             )
         )
     return exps

trinity/common/models/vllm_model.py

@@ -287,35 +287,39 @@ async def generate(
                 input_ids=output.prompt_token_ids,
                 multi_modal_data=prompt.get("multi_modal_data", {}),
             )
-        experiences = [
-            Experience(
-                tokens=torch.cat(
-                    (
-                        torch.tensor(output.prompt_token_ids, dtype=torch.int32),
-                        torch.tensor(output.outputs[i].token_ids, dtype=torch.int32),
-                    )
-                ),
-                logprobs=torch.cat(
-                    (
-                        torch.tensor(
-                            [
-                                list(logprob_dict.values())[0].logprob
-                                for logprob_dict in output.outputs[i].logprobs
-                            ],
-                            dtype=torch.float32,
-                        ),
-                    )
-                ),
-                prompt_length=len(output.prompt_token_ids),
-                prompt_text=self.tokenizer.decode(output.prompt_token_ids),
-                response_text=output.outputs[i].text,
-                multi_modal_inputs=combine_output_token_ids(
-                    output.outputs[i].token_ids, multi_modal_inputs
-                ),
-                routed_experts=self._extract_routed_experts(output, i),
+        experiences = []
+        for output_index, seq_output in enumerate(output.outputs):
+            experiences.append(
+                Experience(
+                    tokens=torch.cat(
+                        (
+                            torch.tensor(output.prompt_token_ids, dtype=torch.int32),
+                            torch.tensor(seq_output.token_ids, dtype=torch.int32),
+                        )
+                    ),
+                    logprobs=torch.cat(
+                        (
+                            torch.tensor(
+                                [
+                                    list(logprob_dict.values())[0].logprob
+                                    for logprob_dict in seq_output.logprobs
+                                ],
+                                dtype=torch.float32,
+                            ),
+                        )
+                    ),
+                    prompt_length=len(output.prompt_token_ids),
+                    prompt_text=self.tokenizer.decode(output.prompt_token_ids),
+                    response_text=seq_output.text,
+                    truncate_status=(
+                        "response_truncated" if seq_output.finish_reason == "length" else None
+                    ),
+                    multi_modal_inputs=combine_output_token_ids(
+                        seq_output.token_ids, multi_modal_inputs
+                    ),
+                    routed_experts=self._extract_routed_experts(output, output_index),
+                )
             )
-            for i in range(len(output.outputs))
-        ]
         return experiences
 
     async def logprobs(  # type: ignore [override]

trinity/common/workflows/envs/alfworld/alfworld_workflow.py

@@ -123,20 +123,30 @@ async def get_model_response_text(self, messages):
     async def generate_env_inference_samples(self, env) -> List[Experience]:
         observation, info = env.reset()
         final_reward = -0.1
+        done = False
+        response_truncate_status = None
+        r = -1
         memory = []
         memory.append({"role": "system", "content": AlfWORLD_SYSTEM_PROMPT})
         for r in range(self.max_env_steps):
             format_obs = format_observation(observation)
             memory = memory + [{"role": "user", "content": format_obs}]
-            response_text = await self.get_model_response_text(memory)
+            response = (await self.get_model_response(memory))[0]
+            response_text = response.response_text
+            response_truncate_status = response.truncate_status
             memory.append({"role": "assistant", "content": response_text})
+            if response_truncate_status == "response_truncated":
+                break
             action = parse_action(response_text)
             observation, reward, done, info = env.step(action)
             if done:
                 final_reward = reward
                 break
         experience = await self.process_messages_to_experience_async(
-            memory, final_reward, {"env_rounds": r, "env_done": 1 if done else 0}
+            memory,
+            final_reward,
+            {"env_rounds": r, "env_done": 1 if done else 0},
+            truncate_status=response_truncate_status,
         )
         # Close the env to save cpu memory
         env.close()

trinity/common/workflows/step_wise_workflow.py

@@ -38,7 +38,7 @@ def run(self) -> list[Experience]:
                 exp.eid.step = step
             # Store the step experiences
             experiences.extend(exps)
-            if not continue_run:
+            if _has_truncated_response(exps) or not continue_run:
                 break
 
         return experiences
@@ -89,7 +89,7 @@ async def run_async(self) -> list[Experience]:
                 exp.eid.step = step
             # Store the step experiences
             experiences.extend(exps)
-            if not continue_run:
+            if _has_truncated_response(exps) or not continue_run:
                 break
 
         return experiences
@@ -144,7 +144,7 @@ def run(self) -> list[Experience]:
                 exp.eid.step = step
             # Store the step experiences
             experiences.extend(exps)
-            if not continue_run:
+            if _has_truncated_response(exps) or not continue_run:
                 break
         reward = self.reward(experiences)
         for exp in experiences:
@@ -197,7 +197,7 @@ async def run_async(self) -> list[Experience]:
                 exp.eid.step = step
             # Store the step experiences
             experiences.extend(exps)
-            if not continue_run:
+            if _has_truncated_response(exps) or not continue_run:
                 break
         reward = await self.reward_async(experiences)
         for exp in experiences:
@@ -225,3 +225,8 @@ async def step_async(self, step_num: int) -> bool:
     async def reward_async(self, exps: list[Experience]) -> float:
         """Calculate the reward for the given experiences of the entire run asynchronously."""
         raise NotImplementedError
+
+
+def _has_truncated_response(exps: list[Experience]) -> bool:
+    """Return whether a model call ended because its response hit the length limit."""
+    return any(getattr(exp, "truncate_status", None) == "response_truncated" for exp in exps)