fix(scheduler): preserve worker context when resizing - #1044
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review. 📝 WalkthroughWalkthrough
ChangesWorker resize context
Shutdown response handling
Priority: ⬇️ Low Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to This change preserves worker context during scheduler resizing and tolerates missing shutdown results. No actionable merge-blocking risk remains. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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 `@src/executorlib/task_scheduler/interactive/blockallocation.py`:
- Around line 99-108: Update the resize path and _worker_kwargs so the first
newly added worker’s bootup event is explicitly released, while preserving
sequential handoff for later workers. Ensure concurrent startup cannot leave
added workers waiting on a statically captured next_bootup_event; use the
existing resize/startup synchronization mechanism or a resizable handoff, and
apply the same correction to the corresponding logic around the later referenced
block.
- Line 134: Protect the live-worker count increment in the resize logic with
self._alive_workers_lock, matching the locking used by _drain_dead_worker().
Keep the adjustment of self._alive_workers[0] atomic with respect to concurrent
worker-failure decrements.
In `@tests/unit/task_scheduler/interactive/test_blockallocation.py`:
- Around line 42-50: Update the test setup around scheduler._bootup_events to
use the valid bootup Event created by the scheduler instead of appending
FakeThread. Extend the worker kwargs assertions to verify the exact bootup
event, its signaled state, and the shared lock, while preserving the existing
worker_id and event-reference checks.
🪄 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: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: c7b1c3b1-83d1-4ab1-a842-0270f51a8f17
📒 Files selected for processing (2)
src/executorlib/task_scheduler/interactive/blockallocation.pytests/unit/task_scheduler/interactive/test_blockallocation.py
| scheduler._bootup_events.append(FakeThread) | ||
| scheduler.max_workers = 2 | ||
|
|
||
| worker = FakeThread.instances[-1] | ||
| self.assertEqual(worker.kwargs["worker_id"], 1) | ||
| self.assertIn("stop_function", worker.kwargs) | ||
| self.assertIn("bootup_event", worker.kwargs) | ||
| self.assertIn("next_bootup_event", worker.kwargs) | ||
| self.assertIs(worker.kwargs["alive_workers"], scheduler._alive_workers) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Use a valid bootup event and verify its state.
Line 42 appends the FakeThread class to _bootup_events, although production workers require an Event. The worker target is not executed, so this test passes even though the worker would fail when it calls bootup_event.wait().
Remove the invalid append. Assert the exact bootup event, its signaled state, and the shared lock.
Proposed fix
with patch(
"executorlib.task_scheduler.interactive.blockallocation.Thread",
FakeThread,
):
- scheduler._bootup_events.append(FakeThread)
scheduler.max_workers = 2
worker = FakeThread.instances[-1]
self.assertEqual(worker.kwargs["worker_id"], 1)
- self.assertIn("bootup_event", worker.kwargs)
- self.assertIn("next_bootup_event", worker.kwargs)
+ self.assertIs(worker.kwargs["bootup_event"], scheduler._bootup_events[1])
+ self.assertTrue(worker.kwargs["bootup_event"].is_set())
+ self.assertIsNone(worker.kwargs["next_bootup_event"])
self.assertIs(worker.kwargs["alive_workers"], scheduler._alive_workers)
+ self.assertIs(
+ worker.kwargs["alive_workers_lock"], scheduler._alive_workers_lock
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| scheduler._bootup_events.append(FakeThread) | |
| scheduler.max_workers = 2 | |
| worker = FakeThread.instances[-1] | |
| self.assertEqual(worker.kwargs["worker_id"], 1) | |
| self.assertIn("stop_function", worker.kwargs) | |
| self.assertIn("bootup_event", worker.kwargs) | |
| self.assertIn("next_bootup_event", worker.kwargs) | |
| self.assertIs(worker.kwargs["alive_workers"], scheduler._alive_workers) | |
| scheduler.max_workers = 2 | |
| worker = FakeThread.instances[-1] | |
| self.assertEqual(worker.kwargs["worker_id"], 1) | |
| self.assertIn("stop_function", worker.kwargs) | |
| self.assertIs(worker.kwargs["bootup_event"], scheduler._bootup_events[1]) | |
| self.assertTrue(worker.kwargs["bootup_event"].is_set()) | |
| self.assertIsNone(worker.kwargs["next_bootup_event"]) | |
| self.assertIs(worker.kwargs["alive_workers"], scheduler._alive_workers) | |
| self.assertIs( | |
| worker.kwargs["alive_workers_lock"], scheduler._alive_workers_lock | |
| ) |
🤖 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/unit/task_scheduler/interactive/test_blockallocation.py` around lines
42 - 50, Update the test setup around scheduler._bootup_events to use the valid
bootup Event created by the scheduler instead of appending FakeThread. Extend
the worker kwargs assertions to verify the exact bootup event, its signaled
state, and the shared lock, while preserving the existing worker_id and
event-reference checks.
- Fix lambda closure capturing self in _worker_kwargs causing reference
cycles that prevent __del__ and block worker threads on future_queue.get()
- Set bootup events for new workers when max_workers increases so threads
don't block forever on bootup_event.wait()
- Use .get('result') instead of ['result'] in communication.py for graceful
shutdown when spawned process is already dead
- Remove erroneous FakeThread class append in test blockallocation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1044 +/- ##
==========================================
+ Coverage 94.25% 94.28% +0.02%
==========================================
Files 39 39
Lines 2176 2186 +10
==========================================
+ Hits 2051 2061 +10
Misses 125 125 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Hi! Both failing checks on 0916e74 look infra-related rather than code-related. The benchmark miss is just ~2% over the timing threshold (mpi4py 53.39s vs. the 52.35s cutoff from process*1.15), and the same benchmark passed on the mpich job. unittest_flux_mpich died during UCX init ( Would a maintainer be able to re-run those two failed jobs? Happy to investigate further if anything looks genuinely connected to the change. |
|
@Muhtasim-Munif-Fahim I triggered the execution of the CI again and now all the tests passed. I have been on holidays for the last two weeks so I am currently working through a bit of a backlog and try to get back to your pull request until the end of the week. |
_alive_workers[0] is read and mutated from multiple worker threads via alive_workers_lock elsewhere; the resize path incremented it unguarded, racing with those threads. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2 participants
Problem
Summary by CodeRabbit
Bug Fixes
Tests