fix(scheduler): preserve worker context when resizing - #1044
fix(scheduler): preserve worker context when resizing#1044Muhtasim-Munif-Fahim wants to merge 3 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthrough
ChangesWorker resize context
Shutdown response handling
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 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
| def _worker_kwargs(self, worker_id: int) -> dict: | ||
| return self._process_kwargs | { | ||
| "worker_id": worker_id, | ||
| "stop_function": lambda: _interrupt_bootup_dict[self._self_id], | ||
| "bootup_event": self._bootup_events[worker_id], | ||
| "next_bootup_event": ( | ||
| self._bootup_events[worker_id + 1] | ||
| if worker_id + 1 < len(self._bootup_events) | ||
| else None | ||
| ), |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Release the first added worker from bootup.
_worker_kwargs() snapshots the original tail worker's next_bootup_event as None. After resize, the first added worker waits on the new unset event, but no existing worker can signal it. The added worker and later added workers remain blocked in bootup_event.wait().
Signal the first new bootup event when resizing. If strict worker-ID boot order must also apply during concurrent startup, use a resizable handoff instead of a static next_bootup_event.
Proposed fix
self._bootup_events.extend(
Event() for _ in range(max_workers - old_max_workers)
)
+ self._bootup_events[old_max_workers].set()
self._alive_workers[0] += max_workers - old_max_workersAlso applies to: 130-140
🤖 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 `@src/executorlib/task_scheduler/interactive/blockallocation.py` around lines
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.
| self._bootup_events.extend( | ||
| Event() for _ in range(max_workers - old_max_workers) | ||
| ) | ||
| self._alive_workers[0] += max_workers - old_max_workers |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Synchronize the live-worker count update.
_drain_dead_worker() decrements self._alive_workers[0] under self._alive_workers_lock. A worker can fail while this resize increments the same value. An unsynchronized increment can lose either update and make the scheduler treat live workers as dead.
Proposed fix
- self._alive_workers[0] += max_workers - old_max_workers
+ with self._alive_workers_lock:
+ self._alive_workers[0] += max_workers - old_max_workers📝 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.
| self._alive_workers[0] += max_workers - old_max_workers | |
| with self._alive_workers_lock: | |
| self._alive_workers[0] += max_workers - old_max_workers |
🤖 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 `@src/executorlib/task_scheduler/interactive/blockallocation.py` at 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.
| 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.19% 94.21% +0.01%
==========================================
Files 39 39
Lines 2137 2144 +7
==========================================
+ Hits 2013 2020 +7
Misses 124 124 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
2 participants
Problem
Summary by CodeRabbit
Bug Fixes
Tests