Fix retryable worker teardown convergence - #1689
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthrough
ChangesWorker lifecycle
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Worker
participant CleanupJournal
participant ChildProcess
participant SharedMemory
Worker->>ChildProcess: Reap child process
ChildProcess-->>Worker: Exit status or ChildProcessError
Worker->>CleanupJournal: Retain surviving PID and mailbox
Worker->>SharedMemory: Reclaim namespace-scoped segments
CleanupJournal->>ChildProcess: Retry child reclamation
CleanupJournal->>SharedMemory: Retry retained cleanup
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
1403aa6 to
329340c
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/simpler/worker.py (1)
8193-8215: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winRename the misnamed loop key, and correct the now-stale docstring.
Two issues in this rewritten loop.
The loop unpacks
for token, entry in entries:, but_host_buf_registryis keyed bydata_ptr, not by token — see the insertion at line 8087 and the lookups in_free_host_buffer_lockedat lines 8144 and 8147. Theget(token)andpop(token)calls at lines 8213-8214 are therefore correct, because the value really is the data_ptr key. The name is wrong, andentry.tokenis a different number entirely (the_host_buf_token_countervalue assigned at line 8077). A later reader who trusts the name and "corrects" these calls to useentry.tokenwould introduce a registry leak.The docstring at lines 8196-8199 no longer matches the new control flow. An unmap failure now raises at line 8208, which skips
_close_host_shmand retains the registry entry for a journal retry. The claim that "every buffer's shm is closed even if its unmap broadcast (or a prior buffer) fails" is no longer true, and the retention is the point of the change.♻️ Proposed fix
def _release_all_host_buffers(self) -> None: """Unmap + free every still-registered host buffer (called from close()). - Per-buffer best-effort: every buffer's shm is closed even if its unmap - broadcast (or a prior buffer) fails, so one failure never strands the - rest; the first error is raised after all are attempted so close() - reports the leak rather than swallowing it to stderr.""" + Per-buffer best-effort: one buffer's failure never strands the rest, and + the first error is raised after all are attempted so close() reports the + leak rather than swallowing it to stderr. A buffer whose unmap broadcast + fails keeps its registry entry so the cleanup journal can retry it.""" with self._registry_lock: entries = list(self._host_buf_registry.items()) errors: list[BaseException] = [] - for token, entry in entries: + for data_ptr, entry in entries: try: if self._worker is not None: # resource presence, not lifecycle (see _close_host_shm) child_errors = self._broadcast_host_unmap(entry.token) if child_errors: raise RuntimeError(f"host buffer token={entry.token} unmap failed: {child_errors[0]}") # Tolerates a still-live view over a zero-copy buffer at close(): # unlinks the name regardless so the OS reclaims it once dropped. self._close_host_shm(entry) with self._registry_lock: - if self._host_buf_registry.get(token) is entry: - self._host_buf_registry.pop(token) + if self._host_buf_registry.get(data_ptr) is entry: + self._host_buf_registry.pop(data_ptr) self._rebuild_host_buf_snapshot()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/simpler/worker.py` around lines 8193 - 8215, Rename the loop key in _release_all_host_buffers from token to data_ptr to reflect that _host_buf_registry is keyed by data_ptr, while keeping the existing get(data_ptr) and pop(data_ptr) behavior unchanged. Update the method docstring to state that unmap failures retain the registry entry and skip _close_host_shm for journal retry, rather than claiming every buffer is closed after failures.
🧹 Nitpick comments (6)
python/simpler/task_interface.py (1)
1206-1211: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider storing the owner as a
Threadobject, not a raw ident.
threading.get_ident()values are recycled after a thread terminates. If the init-owner thread exits and a later thread receives the same ident, that later thread passes this owner check and callsself._impl.finalize()against a device context bound to the dead thread.
Workeralready avoids this by storingself._init_owner_thread: threading.Thread | Noneand comparing withthreading.current_thread()(seepython/simpler/worker.pyline 3867 and the owner check inclose()at line 8922). TheWorker.close()check masks the hole for the journaled_finalize_chippath, but a directChipWorkeruser — the usage documented in the class docstring at lines 1092-1099 — is not covered.♻️ Proposed change to a non-recyclable owner reference
- self._init_owner_thread_id: int | None = None + self._init_owner_thread: threading.Thread | None = None- self._init_owner_thread_id = threading.get_ident() + self._init_owner_thread = threading.current_thread()with self._lifecycle_lock: - owner = self._init_owner_thread_id - if owner is not None and owner != threading.get_ident(): + owner = self._init_owner_thread + if owner is not None and owner is not threading.current_thread(): raise RuntimeError("ChipWorker.finalize() must run on the thread that called ChipWorker.init()")
tests/ut/py/test_chip_worker.pyline 314 setsworker._init_owner_thread_iddirectly, so update that assignment toworker._init_owner_thread = threading.current_thread()as well.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/simpler/task_interface.py` around lines 1206 - 1211, Replace the raw ident ownership tracking used by ChipWorker.init/finalize with a threading.Thread reference, storing threading.current_thread() and comparing the current thread object in the owner check. Update all related state declarations and references, including the test assignment to _init_owner_thread, while preserving the existing finalize-in-progress validation.tests/ut/py/test_chip_worker.py (1)
302-329: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the second new
finalize()guard.This test covers the cross-thread rejection at
python/simpler/task_interface.pylines 1207-1209. The other new guard — rejectingfinalize()whileinit()is in progress, at lines 1210-1211 — has no test. That guard protects a concurrent-teardown window, so it is worth pinning.💚 Proposed additional test
def test_public_wrapper_rejects_finalize_during_init(self): from simpler.task_interface import ChipWorker # noqa: PLC0415 # pyright: ignore[reportAttributeAccessIssue] class FakeImpl: initialized = False device_id = 0 def finalize(self): raise AssertionError("finalize ran while init was in progress") worker = ChipWorker() worker._impl = FakeImpl() worker._init_in_progress = True with pytest.raises(RuntimeError, match="while ChipWorker.init\\(\\) is in progress"): worker.finalize()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ut/py/test_chip_worker.py` around lines 302 - 329, Add a unit test alongside test_public_wrapper_rejects_cross_thread_finalize named test_public_wrapper_rejects_finalize_during_init. Configure a fake implementation, set worker._init_in_progress to True, call worker.finalize(), and assert it raises RuntimeError with the message indicating ChipWorker.init() is in progress, ensuring FakeImpl.finalize is not reached.python/simpler/worker.py (3)
3092-3101: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMake the orphan-shm journal identity unique per segment.
Every orphan mailbox in one group is journaled under the same identity string,
f"{kind} mailbox without live child". This works today because the call at line 3101 usesjournal.add, which appends unconditionally.It becomes lossy if this identity is ever routed through
add_once(line 3148) or selected bydrive(only=...)(line 3165), because both key on the(kind, identity)pair. Then a second orphan mailbox in the same group would be silently dropped or skipped. Include the index to keep the key unique.♻️ Proposed change
- for shm in shms[len(pids) :]: + for orphan_index, shm in enumerate(shms[len(pids) :], start=len(pids)): def _cleanup_orphan_shm(_shm=shm): try: _shm.close() finally: with contextlib.suppress(FileNotFoundError): _shm.unlink() - journal.add("shm", f"{kind} mailbox without live child", _cleanup_orphan_shm) + journal.add("shm", f"{kind} mailbox {orphan_index} without live child", _cleanup_orphan_shm)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/simpler/worker.py` around lines 3092 - 3101, Update the orphan shared-memory journal entry in the loop around _cleanup_orphan_shm and journal.add so its identity includes the segment index, making each orphan mailbox unique within the group while preserving the existing kind and cleanup behavior.
8803-8804: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExpose the journal length instead of reading
_entriesdirectly.
CleanupJournalpublishesemptyanderrorsbut no size. This diagnostic reaches intoself._cleanup_journal._entriesto count items. A__len__onCleanupJournalkeeps the diagnostic on the public surface, and the test attests/ut/py/test_worker/test_cancellation_journal.pyline 376 also inspects_entries.♻️ Proposed change
`@property` def empty(self): return len(self._entries) == 0 + def __len__(self): + return len(self._entries) +if not self._cleanup_journal.empty: - parts.append(f"{len(self._cleanup_journal._entries)} cleanup journal item(s)") + parts.append(f"{len(self._cleanup_journal)} cleanup journal item(s)")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/simpler/worker.py` around lines 8803 - 8804, Update CleanupJournal to implement a public __len__ method returning its item count, then replace direct _entries access in the diagnostic with len(self._cleanup_journal). Update the cancellation journal test to use len(journal) instead of inspecting _entries directly.
6901-6904: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider recording the suppressed reclaim error.
contextlib.suppress(BaseException)discards every failure from_reclaim_child_groups, including theTimeoutErrornaming surviving pids that_reap_child_groupsraises at line 9164 and any mailbox close error.The debt itself is preserved —
_journal_child_survivorsruns before the raise, and_describe_live_resourcesreports journal entries at line 8804 — so this is an observability gap, not a leak. A rollback that leaves survivors currently reports nothing about why. Writing the error to stderr, as the other best-effort paths in this file do (for example line 8106), would keep the diagnostic.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/simpler/worker.py` around lines 6901 - 6904, Record failures from _reclaim_child_groups instead of silently discarding them in the contextlib.suppress block, writing the exception details to stderr consistent with other best-effort paths in worker.py. Preserve the existing suppression behavior so reclaim errors do not interrupt rollback or signal handling.tests/ut/py/test_worker/test_startup_readiness.py (1)
1615-1623: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd cases for the other new
add_workerrejections.This test covers the topology-parent assignment and the attached-child
init()guard. Three other new rejections inadd_workerhave no coverage:
- self-attachment at
python/simpler/worker.pylines 6110-6111;- an already-attached parent at lines 6119-6120;
- an already-attached child at lines 6123-6124.
The self-attachment case matters most.
_hierarchical_start_cvwraps a plain non-reentrantthreading.Lock(lines 3900-3901), so without that guardsorted(...)at line 6115 would yield the same Condition twice andwith first, secondwould deadlock the calling thread. A test pins that the guard stays ahead of the sort.💚 Proposed additional tests
def test_add_worker_rejects_self_attachment(self): # The self-check must precede the two-lock acquire: _hierarchical_start_cv # is non-reentrant, so acquiring it twice would deadlock. w = Worker(level=4, num_sub_workers=0) with pytest.raises(ValueError, match="cannot add a Worker to itself"): w.add_worker(w) def test_add_worker_rejects_already_attached_child(self): first_parent = Worker(level=4, num_sub_workers=0) second_parent = Worker(level=4, num_sub_workers=0) child = Worker(level=3, num_sub_workers=0) first_parent.add_worker(child) with pytest.raises(RuntimeError, match="already attached to another parent"): second_parent.add_worker(child) def test_add_worker_rejects_attached_parent(self): root = Worker(level=5, num_sub_workers=0) middle = Worker(level=4, num_sub_workers=0) leaf = Worker(level=3, num_sub_workers=0) root.add_worker(middle) with pytest.raises(RuntimeError, match="already attached as a child"): middle.add_worker(leaf)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ut/py/test_worker/test_startup_readiness.py` around lines 1615 - 1623, Add three tests alongside test_add_worker_freezes_child_before_topology_publication covering Worker.add_worker rejection paths: self-attachment must raise ValueError with the “cannot add a Worker to itself” message before lock acquisition, attaching an already-parented child must raise RuntimeError with “already attached to another parent,” and adding a child from an already-attached parent must raise RuntimeError with “already attached as a child.”
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/simpler/worker.py`:
- Around line 9267-9280: Restrict the initial self._cleanup_journal.drive() in
the teardown method to an only set containing the six pre-transport keys
registered in the preceding loop, preventing residual child or shm entries from
aborting transport shutdown. Preserve the later child-entry handling in
_reclaim_child_groups, and extend
test_both_teardown_paths_retry_each_resource_class with a case where a
pre-existing child journal entry survives into _teardown_worker_tree.
- Around line 6131-6145: Update the worker lifecycle/linkage guard used by
_linkage.add_worker() to reject additions when the worker or any ancestor is no
longer in the NEW state, thereby freezing the entire subtree during
_assign_shm_namespace(). Ensure the ancestor-state check covers recursive
descendants so no worker can be appended after the namespace traversal begins.
---
Outside diff comments:
In `@python/simpler/worker.py`:
- Around line 8193-8215: Rename the loop key in _release_all_host_buffers from
token to data_ptr to reflect that _host_buf_registry is keyed by data_ptr, while
keeping the existing get(data_ptr) and pop(data_ptr) behavior unchanged. Update
the method docstring to state that unmap failures retain the registry entry and
skip _close_host_shm for journal retry, rather than claiming every buffer is
closed after failures.
---
Nitpick comments:
In `@python/simpler/task_interface.py`:
- Around line 1206-1211: Replace the raw ident ownership tracking used by
ChipWorker.init/finalize with a threading.Thread reference, storing
threading.current_thread() and comparing the current thread object in the owner
check. Update all related state declarations and references, including the test
assignment to _init_owner_thread, while preserving the existing
finalize-in-progress validation.
In `@python/simpler/worker.py`:
- Around line 3092-3101: Update the orphan shared-memory journal entry in the
loop around _cleanup_orphan_shm and journal.add so its identity includes the
segment index, making each orphan mailbox unique within the group while
preserving the existing kind and cleanup behavior.
- Around line 8803-8804: Update CleanupJournal to implement a public __len__
method returning its item count, then replace direct _entries access in the
diagnostic with len(self._cleanup_journal). Update the cancellation journal test
to use len(journal) instead of inspecting _entries directly.
- Around line 6901-6904: Record failures from _reclaim_child_groups instead of
silently discarding them in the contextlib.suppress block, writing the exception
details to stderr consistent with other best-effort paths in worker.py. Preserve
the existing suppression behavior so reclaim errors do not interrupt rollback or
signal handling.
In `@tests/ut/py/test_chip_worker.py`:
- Around line 302-329: Add a unit test alongside
test_public_wrapper_rejects_cross_thread_finalize named
test_public_wrapper_rejects_finalize_during_init. Configure a fake
implementation, set worker._init_in_progress to True, call worker.finalize(),
and assert it raises RuntimeError with the message indicating ChipWorker.init()
is in progress, ensuring FakeImpl.finalize is not reached.
In `@tests/ut/py/test_worker/test_startup_readiness.py`:
- Around line 1615-1623: Add three tests alongside
test_add_worker_freezes_child_before_topology_publication covering
Worker.add_worker rejection paths: self-attachment must raise ValueError with
the “cannot add a Worker to itself” message before lock acquisition, attaching
an already-parented child must raise RuntimeError with “already attached to
another parent,” and adding a child from an already-attached parent must raise
RuntimeError with “already attached as a child.”
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 19b993b6-e9d7-4562-96dd-3eb2d8ff5bce
📒 Files selected for processing (5)
python/simpler/task_interface.pypython/simpler/worker.pytests/ut/py/test_chip_worker.pytests/ut/py/test_worker/test_cancellation_journal.pytests/ut/py/test_worker/test_startup_readiness.py
329340c to
fa7fcd1
Compare
Startup rollback and close share one journal-driven resource sequence. Unreaped child mailboxes stay owned until waitpid. The root reclaims nested SHM through an exact tree namespace. Parent-child topology freezes atomically. Public ChipWorker teardown enforces init-thread ownership.
|
@coderabbitai Addressed the review-body findings in fa7fcd1:
The branch is also rebased onto current main. Focused tests: 123 passed. Full device-free Worker/ChipWorker tests: 587 passed, 2 skipped. Ruff, pyright, and changed-file pre-commit hooks pass. |
|
Tip For best results, initiate chat on the files or code changes.
🐇 ✅ If you found this review helpful, would you consider giving us a shout-out on X? Thank you for using CodeRabbit! You are interacting with an AI system. |
1 participant
Summary
Contract
Testing