diff --git a/tensorrt_llm/scaffolding/controller.py b/tensorrt_llm/scaffolding/controller.py index 4ea3aba87dc5..1cbe5de92640 100644 --- a/tensorrt_llm/scaffolding/controller.py +++ b/tensorrt_llm/scaffolding/controller.py @@ -302,7 +302,7 @@ def process(self, assert isinstance(majority_answer, str), "majority_vote failed" # The task returned by majority vote does not have output_tokens and logits. - tasks[0].result = tasks_list[majority_index][0].result + tasks[0].output_str = tasks_list[majority_index][0].output_str def majority_vote(self, candidates_tasks: List[List[Task]], **kwargs) -> Tuple[int, str]: diff --git a/tests/unittest/scaffolding/test_scaffolding.py b/tests/unittest/scaffolding/test_scaffolding.py index eaefc5386ad1..d6d322936583 100644 --- a/tests/unittest/scaffolding/test_scaffolding.py +++ b/tests/unittest/scaffolding/test_scaffolding.py @@ -1,6 +1,5 @@ # autoflake: skip_file -import pytest from scaffolding.test_worker import (create_trtllm_worker, deepseek_distill_7b_path, default_prompt) @@ -8,8 +7,6 @@ NativeGenerationController, ScaffoldingLlm) -_SCAFFOLDING_EXECUTOR_NVBUG = pytest.mark.skip(reason="https://nvbugs/6341070") - def create_scaffolding_llm_with_native_generation_controller( deepseek_distill_7b_path): @@ -48,7 +45,6 @@ def create_scaffolding_llm_with_majority_vote_controller( return llm -@_SCAFFOLDING_EXECUTOR_NVBUG def test_unbatched_scaffolding_sync(default_prompt, deepseek_distill_7b_path): scaffolding_llm = create_scaffolding_llm_with_native_generation_controller( deepseek_distill_7b_path) @@ -58,7 +54,6 @@ def test_unbatched_scaffolding_sync(default_prompt, deepseek_distill_7b_path): scaffolding_llm.shutdown(shutdown_workers=True) -@_SCAFFOLDING_EXECUTOR_NVBUG def test_batched_scaffolding_sync(default_prompt, deepseek_distill_7b_path): scaffolding_llm = create_scaffolding_llm_with_native_generation_controller( deepseek_distill_7b_path) @@ -72,7 +67,6 @@ def test_batched_scaffolding_sync(default_prompt, deepseek_distill_7b_path): scaffolding_llm.shutdown(shutdown_workers=True) -@_SCAFFOLDING_EXECUTOR_NVBUG def test_async_scaffolding_generation(default_prompt, deepseek_distill_7b_path): async def run_async_test(): @@ -88,7 +82,6 @@ async def run_async_test(): asyncio.run(run_async_test()) -@_SCAFFOLDING_EXECUTOR_NVBUG def test_majority_vote(default_prompt, deepseek_distill_7b_path): scaffolding_llm = create_scaffolding_llm_with_majority_vote_controller( deepseek_distill_7b_path, samples_num=3) diff --git a/tests/unittest/scaffolding/test_worker.py b/tests/unittest/scaffolding/test_worker.py index 14d22960d6d5..27d4353a1078 100644 --- a/tests/unittest/scaffolding/test_worker.py +++ b/tests/unittest/scaffolding/test_worker.py @@ -18,7 +18,6 @@ sys.path.append(os.path.join(os.path.dirname(__file__), '..')) from llmapi.test_llm import get_model_path -_SCAFFOLDING_EXECUTOR_NVBUG = pytest.mark.skip(reason="https://nvbugs/6341070") _SCAFFOLDING_MCP_NVBUG = pytest.mark.skip(reason="https://nvbugs/6337229") _TRTOAI_PROCESSPOOL_NVBUG = pytest.mark.skip(reason="https://nvbugs/6341072") @@ -98,7 +97,6 @@ def create_trtllm_worker(model_path): return TRTLLMWorker.init_with_new_llm(str(model_path), backend="pytorch") -@_SCAFFOLDING_EXECUTOR_NVBUG def test_trtllm_worker_generation(default_prompt, deepseek_distill_7b_path): worker = create_trtllm_worker(deepseek_distill_7b_path) task = GenerationTask.create_from_prompt(default_prompt)