loadscope: forget a dead worker's collection so a still-collecting worker isn't mistaken for done

[dchaudhari7177]

Fixes #1362.

Problem

LoadScopeScheduling.remove_node() pops the dead worker from assigned_work but never removes its entry from registered_collections. Since collection_is_completed is len(registered_collections) >= numnodes, the stale entry keeps the collection looking "complete" while a later worker is still collecting. schedule() then runs and _assign_work_unit() raises KeyError indexing registered_collections for the worker that never registered — crashing the run with INTERNALERROR. This affects both --dist=loadscope and --dist=loadfile (the latter subclasses LoadScopeScheduling).

Fix

Drop the node from registered_collections in remove_node(), mirroring the assigned_work.pop(node) already there. A replacement worker then re-registers normally, and collection_is_completed correctly waits for every live worker.

This is safe for the shutdown path too: normal worker_workerfinished removals only happen after triggershutdown, and on a crash the replacement worker re-registers its collection.

Test

Added TestLoadScopeScheduling::test_remove_node_forgets_dead_worker_collection, which reproduces the scenario with MockNodes (3 expected workers; one dies mid-collection while another is still collecting). It fails on main (registered_collections still holds the dead node, collection_is_completed flips true early) and passes with the fix. Full testing/test_dsession.py is green; ruff check/format clean. Added changelog/1362.bugfix.rst.

[larsoner]

Is this related to / dup of #1363 ?

[dchaudhari7177]

No — different bug, different line, and they can both be wrong at once. Same file, which is why they look alike.

#1363 / #1313 is about mark_test_complete never pruning a finished work unit from assigned_work, so remove_node requeues it and a replacement worker gets an empty unit it can never complete. You closed it in favour of #1328, which fixes that class of hang for loadgroup.

This PR / #1362 is about remove_node popping the dead node from assigned_work but leaving it in registered_collections. Since

def collection_is_completed(self) -> bool:
    return len(self.registered_collections) >= self.numnodes

counts entries rather than live workers, the stale entry makes the collection look complete while a later worker is still collecting. schedule() then runs and _assign_work_unit raises KeyError indexing registered_collections for the worker that never registered. The failure mode is an INTERNALERROR traceback, not a hang.

The two diffs do not touch the same function — #1363 edits mark_test_complete (line ~252), this edits remove_node (line ~182) — and #1328, which landed instead, does not touch registered_collections at all. I re-read remove_node on master at e27d3bb to confirm: the registered_collections.pop is still absent, and testing/test_dsession.py::TestLoadScopeScheduling::test_remove_node_forgets_dead_worker_collection still fails there.

Happy to rebase and re-run if that helps.

[dchaudhari7177]

Not a duplicate — different bug, different line, and they don't overlap.

#1363 (and #1328, which you preferred) is about #1313: a crashed worker's completed work units get requeued, a replacement is handed an empty unit, and the session hangs. It changes mark_test_complete and _assign_work_unit.

This one is #1362: remove_node pops the dead worker from assigned_work but leaves its entry in registered_collections. collection_is_completed is len(registered_collections) >= numnodes, so the stale entry keeps collection looking complete while a later worker is still collecting; schedule() then runs and _assign_work_unit raises KeyError on the worker that never registered, so the run dies with INTERNALERROR rather than hanging. The fix is one line in remove_node, mirroring the assigned_work.pop(node) next to it.

So a session could hit either one independently: #1328 doesn't touch registered_collections, and this doesn't touch requeueing. Happy to rebase if that helps.

[davidheff]

Thanks for this. The stale entry is a real fault, and dropping it fixes the schedule() route. But I think a second route to the same KeyError remains with this change applied.

remove_node ends by calling _reschedule on every node in assigned_work. A replacement worker is in assigned_work as soon as add_node runs, so it's there before it has reported its collection. If a second worker crashes while the first replacement is still collecting, the second remove_node re-queues the dead worker's units and then reschedules the replacement. _reschedule sees nothing pending and calls _assign_work_unit, which raises KeyError looking up registered_collections for the replacement.

worker_errordown swallows KeyError from remove_node (except KeyError: pass), so this route doesn't take the run down the way the schedule() route does. But it skips handle_crashitem, so the test the second worker crashed on is never reported, and the loop stops before rescheduling any node after the replacement.

Reproduced against this branch (e6ef607) and against master with a real LoadScopeScheduling and mock nodes:

import types
from xdist.scheduler.loadscope import LoadScopeScheduling

class Node:
    def __init__(self, name):
        self.gateway = types.SimpleNamespace(id=name)
        self.shutting_down = False
        self.sent = []
    def send_runtest_some(self, indices):
        self.sent.extend(indices)
    def shutdown(self):
        self.shutting_down = True

config = types.SimpleNamespace(
    getvalue=lambda name: ["3*popen"],
    option=types.SimpleNamespace(loadscopereorder=False),
)
collection = [f"test_{m}.py::test_{i}" for m in "abcdef" for i in range(3)]
sched = LoadScopeScheduling(config)
a, b, c = Node("gw0"), Node("gw1"), Node("gw2")
for node in (a, b, c):
    sched.add_node(node)
for node in (a, b, c):
    sched.add_node_collection(node, collection)
    if sched.collection_is_completed:
        sched.schedule()

sched.remove_node(a)   # gw0 crashes
sched.add_node(Node("gw3"))   # its replacement starts collecting
sched.remove_node(b)   # gw1 crashes: KeyError, on this branch and on master

What closes both routes for us is the check from #1299: _reschedule returns early for a node that isn't in registered_collections yet. It's then scheduled normally when its own collection arrives. With that in place, dropping the stale entry is still correct, since it keeps collection_is_completed honest, but it's no longer what prevents the KeyError. Would you consider adding the guard here, or should the two PRs be combined?

For context, we hit this in a real run when two workers died in the same second. The INTERNALERROR came from the schedule() route and terminated the entire run before it completed.

[dchaudhari7177]

@davidheff thanks. That's a real second route, and your repro made it quick to pin down.

I added the guard here: _reschedule now returns early for a node that isn't in registered_collections yet. It's the same check as #1299, so whichever lands first, the other drops that line on rebase. I'm not trying to take over that PR; @bendrucker's test covers the mismatched-collection case, which this one doesn't.

New test, test_second_crash_while_replacement_collects, is your scenario on LoadScopeScheduling. It runs three workers, gw0 dies, a replacement joins and is still collecting, then gw1 dies. It asserts that the second remove_node doesn't raise and that the replacement is neither handed work nor shut down. Once the replacements have collected, schedule() gives both of them work. Without the guard it fails with KeyError in _assign_work_unit, the same as your repro.

I also updated the changelog fragment to cover both routes. testing/test_dsession.py gives 34 passed, 2 xfailed, and ruff 0.15.22 check and format are clean. I didn't finish the acceptance suite locally; it is very slow on Windows here, so CI is the check for that part.