tensorrt_llm/scaffolding/README_TREE_METHODS.md (1)

41-43: Nit: tighten wording and quote the stop token.

Apply:

-For long outputs, raise generation `max_tokens` and consider adding a stop sequence (e.g., `</think>`).
+For long outputs, increase generation `max_tokens` and consider adding a stop sequence (e.g., '</think>').

examples/scaffolding/run_tot_example.py (5)

11-25: CLI polish: remove unused arg and expose branch_factor.

jsonl_file is unused. Expose --branch_factor to match controller/documentation.

Apply:

-    parser.add_argument('--jsonl_file', type=str, default='./test.jsonl')
+    # Optional: dataset input can be added later
@@
     parser.add_argument('--selection_strategy',
                         type=str,
                         default="best",
                         choices=["best", "vote", "random"])
+    parser.add_argument('--branch_factor', type=int, default=2)

---

42-52: Make reward model optional.

Docs say PRM is optional; gate reward worker initialization.

Apply:

-    # Initialize reward worker if provided
-    reward_controller = None
-    reward_worker = TRTLLMWorker.init_with_new_llm(
-        args.reward_model_dir,
-        backend="pytorch",
-        max_batch_size=4,
-        max_num_tokens=8192,
-        kv_cache_free_gpu_memory_fraction=0.2,
-        disable_overlap_scheduler=True)
-    workers[PRMController.WorkerTag.REWARD] = reward_worker
+    # Initialize reward worker if provided
+    reward_controller = None
+    reward_worker = None
+    if args.reward_model_dir:
+        reward_worker = TRTLLMWorker.init_with_new_llm(
+            args.reward_model_dir,
+            backend="pytorch",
+            max_batch_size=4,
+            max_num_tokens=8192,
+            kv_cache_free_gpu_memory_fraction=0.2,
+            disable_overlap_scheduler=True)
+        workers[PRMController.WorkerTag.REWARD] = reward_worker

---

54-59: Construct PRMController only when reward_worker exists.

Apply:

-    reward_controller = PRMController(tokenizer=reward_worker.tokenizer)
+    if reward_worker:
+        reward_controller = PRMController(tokenizer=reward_worker.tokenizer)

---

61-67: Pass branch_factor from CLI.

Aligns example with documented parameter.

Apply:

-                               selection_strategy=args.selection_strategy,
-                               branch_factor=2)
+                               selection_strategy=args.selection_strategy,
+                               branch_factor=args.branch_factor)

---

74-77: Ensure worker shutdown on errors; print robustly.

Wrap run in try/finally and handle single vs batched returns.

Apply:

-    results = llm.generate(prompts)
-    print(results[0].outputs[0].text)
-    llm.shutdown(shutdown_workers=True)
-    print(f'main shut down done')
+    try:
+        results = llm.generate(prompts)
+        res = results if isinstance(results, list) else [results]
+        print(res[0].outputs[0].text)
+    finally:
+        llm.shutdown(shutdown_workers=True)
+        print("Main shutdown done")

examples/scaffolding/run_mcts_example.py (4)

18-24: Remove unused argument.

jsonl_file isn’t used.

Apply:

-    parser.add_argument('--jsonl_file', type=str, default='./test.jsonl')

---

39-48: Make reward model optional (mirrors TOT example).

Apply:

-    # Initialize reward worker if provided
-    reward_controller = None
-    reward_worker = TRTLLMWorker.init_with_new_llm(
-        args.reward_model_dir,
-        backend="pytorch",
-        max_batch_size=4,
-        max_num_tokens=8192,
-        kv_cache_free_gpu_memory_fraction=0.2,
-        disable_overlap_scheduler=True)
-    workers[PRMController.WorkerTag.REWARD] = reward_worker
+    # Initialize reward worker if provided
+    reward_controller = None
+    reward_worker = None
+    if args.reward_model_dir:
+        reward_worker = TRTLLMWorker.init_with_new_llm(
+            args.reward_model_dir,
+            backend="pytorch",
+            max_batch_size=4,
+            max_num_tokens=8192,
+            kv_cache_free_gpu_memory_fraction=0.2,
+            disable_overlap_scheduler=True)
+        workers[PRMController.WorkerTag.REWARD] = reward_worker

---

54-55: Construct PRMController only when reward_worker exists.

Apply:

-    reward_controller = PRMController(tokenizer=reward_worker.tokenizer)
+    if reward_worker:
+        reward_controller = PRMController(tokenizer=reward_worker.tokenizer)

---

70-73: Ensure worker shutdown on errors; print robustly.

Apply:

-    results = llm.generate(prompts)
-    print(results[0].outputs[0].text)
-    llm.shutdown(shutdown_workers=True)
-    print(f'main shut down done')
+    try:
+        results = llm.generate(prompts)
+        res = results if isinstance(results, list) else [results]
+        print(res[0].outputs[0].text)
+    finally:
+        llm.shutdown(shutdown_workers=True)
+        print("Main shutdown done")

tensorrt_llm/scaffolding/tree_controllers.py (9)

23-30: Polish method docstrings (punctuation, clarity).

Keeps D415 happy and improves readability.

 def is_leaf(self) -> bool:
-        """node has no children"""
+        """Return True if the node has no children."""
 ...
 def is_root(self) -> bool:
-        """node has no parent?"""
+        """Return True if the node has no parent."""

---

50-52: Remove or use untried_actions.

It’s declared but never referenced. Either wire it into expansion or drop it to avoid confusion.

---

110-117: Prefer explicit exception over assert for runtime validation.

assert can be stripped with optimization; raise ValueError instead.

-        assert len(
-            tasks) == 1, "MCTS Controller only supports single task processing"
+        if len(tasks) != 1:
+            raise ValueError("MCTSController only supports single task processing.")

---

118-121: Seed randomness for reproducibility (optional).

Honor task.seed/seed kwarg to make runs deterministic.

         initial_state = getattr(task, 'input_str', str(task)) or ""
+        # Optional reproducibility
+        seed = kwargs.get('seed', getattr(task, 'seed', None))
+        if seed is not None:
+            random.seed(seed)

---

253-276: Document and validate constructor params.

Add docstrings and basic validation (e.g., positive depths/iterations, branch_factor ≥ 1).

---

279-281: Prefer explicit exception over assert for runtime validation.

Same as MCTS: raise ValueError.

-        assert len(
-            tasks) == 1, "TOT Controller only supports single task processing"
+        if len(tasks) != 1:
+            raise ValueError("TOTController only supports single task processing.")

---

414-423: Narrow overly broad exception.

Catch specific types to avoid masking issues.

-        except Exception:
+        except (TypeError, AttributeError):
             final_task.stop = ['</think>']

---

449-456: Fix docstring style (D205, D415).

Add punctuation and a blank line after the summary.

-        def is_new_item(line: str) -> Optional[str]:
-            """Return the content of a new item header if the line starts a new approach/step item.
-            Supports:
+        def is_new_item(line: str) -> Optional[str]:
+            """Return the content of a new item header if the line starts a new approach/step item.
+
+            Supports:
             - 'Approach N: ...' or 'Step N: ...'
             - 'N. ...'
             - '- ...' or '* ...'
             """

---

150-166: Avoid brittle GenerationResult instantiation
Bypassing __init__ and poking _outputs is fragile—provide a small constructor/helper (or extend PRMController) to build a valid GenerationResult from raw text (or accept plain strings directly), instead of using GenerationResult.__new__ and private fields.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.