Skip to content

fix: restrict MCTSScheduler leaf selection to available DAG leaves (#1429)#1437

Open
utsab345 wants to merge 1 commit into
microsoft:mainfrom
utsab345:fix/mcts-scheduler-leaf-selection
Open

fix: restrict MCTSScheduler leaf selection to available DAG leaves (#1429)#1437
utsab345 wants to merge 1 commit into
microsoft:mainfrom
utsab345:fix/mcts-scheduler-leaf-selection

Conversation

@utsab345

@utsab345 utsab345 commented Jun 30, 2026

Copy link
Copy Markdown
Contributor

Summary

Fixes #1429.

This PR updates MCTSScheduler.select() to follow the same leaf-selection semantics as the existing RoundRobinScheduler and ProbabilisticScheduler.

Changes

  • Replaced the candidate set built from all trace.hist indices with the actual DAG leaves returned by trace.get_leaves().
  • Filtered out leaves that are currently being expanded by another worker using uncommited_rec_status.
  • Preserved the existing MCTS scoring (PUCT) logic by only changing the candidate selection stage.

Why

Previously, MCTSScheduler.select() constructed its candidate set using every index in trace.hist, which caused two issues:

  • Internal (non-leaf) nodes could be selected for expansion, violating the scheduler contract and standard MCTS leaf-expansion semantics.
  • Leaves already being expanded by another worker were not filtered out, allowing duplicate work during parallel execution.

This implementation now mirrors the behavior of the sibling schedulers while keeping the remainder of the MCTS selection logic unchanged.

Testing

Verified the following:

  • Project compiles successfully.
  • Import test passes (299 subtests).
  • Only DAG leaf nodes are selected for expansion.
  • In-flight leaves tracked by uncommited_rec_status are never selected.
  • Existing MCTS selection behavior is preserved apart from the corrected candidate filtering.

This change is intentionally minimal and only affects the construction of the candidate leaf set.


📚 Documentation preview 📚: https://RDAgent--1437.org.readthedocs.build/en/1437/

…_status

The old code used list(set(range(len(trace.hist)))) as available leaves,
which includes internal (already-expanded) nodes and ignores
uncommited_rec_status. This caused two bugs:

1. MCTS expanded internal nodes instead of actual DAG leaves
2. Parallel workers could re-expand the same in-flight leaf

Fix mirrors the pattern used by RoundRobinScheduler and
ProbabilisticScheduler: trace.get_leaves() filtered by
uncommited_rec_status[leaf] == 0.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

🐛 Bug: MCTSScheduler.select uses all hist indices as "leaves" and ignores uncommited_rec_status — picks internal nodes and re-picks in-flight ones

1 participant