diff --git a/tests/common/experience_extraction_test.py b/tests/common/experience_extraction_test.py index dea984f6e48..329d99991a3 100644 --- a/tests/common/experience_extraction_test.py +++ b/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( [ diff --git a/tests/explorer/step_wise_workflow_test.py b/tests/explorer/step_wise_workflow_test.py index 389b8ec01cb..2127d603b49 100644 --- a/tests/explorer/step_wise_workflow_test.py +++ b/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: diff --git a/tests/explorer/workflow_test.py b/tests/explorer/workflow_test.py index 782824bc9a9..1c987590417 100644 --- a/tests/explorer/workflow_test.py +++ b/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 = [ diff --git a/trinity/common/models/experience_extraction.py b/trinity/common/models/experience_extraction.py index e70eb4dffdf..2662072e702 100644 --- a/trinity/common/models/experience_extraction.py +++ b/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 diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index b30e83e8287..1affc1e7621 100644 --- a/trinity/common/models/vllm_model.py +++ b/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] diff --git a/trinity/common/workflows/envs/alfworld/alfworld_workflow.py b/trinity/common/workflows/envs/alfworld/alfworld_workflow.py index 7fb29604a22..bbab895a910 100644 --- a/trinity/common/workflows/envs/alfworld/alfworld_workflow.py +++ b/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() diff --git a/trinity/common/workflows/step_wise_workflow.py b/trinity/common/workflows/step_wise_workflow.py index 1042866063a..b3231f2b0db 100644 --- a/trinity/common/workflows/step_wise_workflow.py +++ b/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)