🐛 Bug: DSTrace.get_sibling_exps crashes with KeyError when the trace DAG has any node with more than one child

[Fr3ya]

🐛 Bug Description

DSTrace.get_sibling_exps in rdagent/scenarios/data_science/proposal/exp_gen/base.py#L114-L137 crashes with KeyError the moment the trace DAG contains any single node with more than one child:

# base.py:122-129 (verbatim)
touched_node_set = set()
for idx in range(len(self.dag_parent)):
    touched_node_set.add(idx)
    if self.dag_parent[idx] == self.NEW_ROOT:
        continue
    for parent in self.dag_parent[idx]:
        touched_node_set.remove(parent)   # ← KeyError when parent already removed

When node A has two children B and C:

idx	action	set state
A	add A; skip remove (NEW_ROOT)	{A}
B	add B; remove A	{B}
C	add C; remove A	KeyError — A already gone

The same pattern repeats at L132-L133 for exp_parent_idx removal across uncommitted_experiments — if two uncommitted experiments share a parent (a normal parallel-expansion case), it crashes there too.

The function is called by DSProposalV2ExpGen at proposal.py:1365:

sibling_exp = trace.get_sibling_exps() if trace.should_inject_diversity() else None

should_inject_diversity() returns True when DS_RD_SETTING.enable_cross_trace_diversity is enabled. Branching DAGs are the entire point of multi-trace exploration — so this crashes on the first sibling lookup after the first sibling expansion.

To Reproduce

Self-contained — no RD-Agent install needed:

class FakeExp:
    def __init__(self, label, parent=()):
        self.label = label
        self.local_selection = parent
    def __repr__(self):
        return f"Exp({self.label})"


class FakeTrace:
    NEW_ROOT = ()
    def __init__(self, dag_parent, hist, uncommitted=None, current=()):
        self.dag_parent = dag_parent
        self.hist = hist
        self.uncommitted_experiments = uncommitted or {}
        self._cur = current
    def get_current_selection(self):
        return self._cur


def get_sibling_exps(self, current_selection=None):
    """Verbatim port of DSTrace.get_sibling_exps."""
    if current_selection is None:
        current_selection = self.get_current_selection()
    ignore_leaf_idx = [current_selection[0]] if current_selection != self.NEW_ROOT else []
    sibling_exps = []
    touched_node_set = set()
    for idx in range(len(self.dag_parent)):
        touched_node_set.add(idx)
        if self.dag_parent[idx] == self.NEW_ROOT:
            continue
        for parent in self.dag_parent[idx]:
            touched_node_set.remove(parent)
    for loop_idx, exp in self.uncommitted_experiments.items():
        sibling_exps.append(exp)
        if (exp_parent_idx := exp.local_selection[0] if exp.local_selection != self.NEW_ROOT else None) is not None:
            touched_node_set.remove(exp_parent_idx)
    for idx in touched_node_set:
        if idx not in ignore_leaf_idx:
            sibling_exps.append(self.hist[idx][0])
    return sibling_exps


# Linear DAG  0 → 1 → 2: fine
t = FakeTrace(
    dag_parent=[FakeTrace.NEW_ROOT, (0,), (1,)],
    hist=[(FakeExp("A0"), None), (FakeExp("A1"), None), (FakeExp("A2"), None)],
    current=(2,),
)
print(f"Linear DAG ok: {get_sibling_exps(t)}")

# Branching DAG  0 → 1, 0 → 2: KeyError
t = FakeTrace(
    dag_parent=[FakeTrace.NEW_ROOT, (0,), (0,)],
    hist=[(FakeExp(f"B{i}"), None) for i in range(3)],
    current=(1,),
)
try:
    get_sibling_exps(t)
except KeyError as e:
    print(f"Branching DAG crashes: KeyError({e!r})")

Output:

Linear DAG ok: []
Branching DAG crashes: KeyError(0)

Expected Behavior

get_sibling_exps should compute the set of "leaf-ish" nodes (those not appearing as a parent of anything yet, plus uncommitted experiments) without crashing on shared parents. The intent is clear from the existing code structure; the bug is just set.remove raising on a key that was already removed.

Suggested fix

One character per line, on lines base.py:129 and base.py:133:

         for parent in self.dag_parent[idx]:
-            touched_node_set.remove(parent)
+            touched_node_set.discard(parent)
     for loop_idx, exp in self.uncommitted_experiments.items():
         sibling_exps.append(exp)
         if (exp_parent_idx := exp.local_selection[0] if exp.local_selection != self.NEW_ROOT else None) is not None:
-            touched_node_set.remove(exp_parent_idx)
+            touched_node_set.discard(exp_parent_idx)

set.discard is the idempotent counterpart to set.remove — no-op when the key is absent.

Verification of the fix against the branching DAG above:

Branching DAG (fixed): [Exp(B2)]    # node 2 — sibling of current selection 1, parent 0 not counted

Environment

• RD-Agent version: SHA 4f9ecb005881cddc08df0124a2e894c018007679 (HEAD at scan time)

Additional Notes

• Surfaced by a systematic scan for "set-based DAG bookkeeping with set.remove instead of set.discard."
• A unit test that builds a 3-node fan-out (dag_parent=[(), (0,), (0,)]) and calls get_sibling_exps would catch this. Worth adding alongside the fix.
• The trace_scheduler also has an unrelated leaf-selection bug worth filing separately (MCTSScheduler.select uses range(len(trace.hist)) instead of trace.get_leaves() — that scheduler will also surface odd behavior in the same branching topologies that trigger this KeyError).