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.