Update: start L2 worker hierarchy eagerly in Worker.init() - #1381
Update: start L2 worker hierarchy eagerly in Worker.init()#1381YunjiQin wants to merge 1 commit into
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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughHierarchical workers now start eagerly during ChangesHierarchical startup timing
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Worker
participant Hierarchy
participant InnerWorker
participant ChildWorkers
Worker->>Hierarchy: Worker.init() calls _start_hierarchical()
Hierarchy->>InnerWorker: create and initialize L3 worker
InnerWorker->>ChildWorkers: eagerly fork sub/chip children
Worker->>InnerWorker: dispatch via inner_worker.run(...)
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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request changes the initialization behavior of hierarchical workers (level >= 3) to be eager rather than lazy. Specifically, _start_hierarchical(), which forks child processes and starts the C++ scheduler, is now called directly within init() instead of being deferred to the first run(). Consequently, various deferred startup calls and state checks across worker.py have been removed, and the documentation and unit tests have been updated to align with this eager initialization model. There are no review comments, so I have no feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
c628ed8 to
4925f72
Compare
4925f72 to
c8b094d
Compare
Move _start_hierarchical() into Worker.init() for L3+ workers, so the whole L2 hierarchy is live — chip/sub children forked, ChipWorker.init done, ChipCallables uploaded H2D, prewarm complete — by the time init() returns, instead of being deferred to the first run(). - Call _start_hierarchical() from init() for L3+, inside the try block so a startup failure is reaped by _cleanup_partial_init. - Remove the five now-redundant lazy trigger sites (run, create_host_buffer, remote require/register/unregister); each is reached only after init(), so the calls were idempotent no-ops. - New contract: input tensors must be allocated before init(), or via create_host_buffer after init(). SceneTestCase's L3 path rehosts generate_args() tensors onto create_host_buffer regions so the eagerly-forked chip children can see them. - Fix a latent create_host_buffer bug the rehosting exposed: _rewrite_blob_host_addrs rewrote EVERY tensor whose VA fell in a registered host-buffer range, including child_memory device pointers (an HCCL comm-window scratch slot). A collective's window pointer whose numeric VA coincidentally collided with a host-buffer parent VA got redirected into a create_host_buffer mapping, deadlocking the exchange. Skip child_memory tensors (byte 43 of the 128 B Tensor) — only host pointers are ever rewritten. - Registrations between init() and run() now take the post-start broadcast path; a childless L3 registering a Python fn after init raises eagerly instead of being silently snapshotted. - Update docstrings/comments and docs (task-flow, worker-manager, hierarchical_level_runtime) to describe eager startup. - Rework the three host-worker tests that asserted the removed "init-done-but-not-started" window; stub _start_hierarchical in the remote-id bookkeeping unit test that uses a fake C worker.
|
Closing as superseded by #1397 (merged as 483d97e). #1397 delivers the complete eager, transactional, recursive Thanks for the original direction. Closing to keep the eager-init work tracked in one place. |
2 participants
Summary
Move
_start_hierarchical()intoWorker.init()for L3+ workers, so thewhole L2 hierarchy is live — chip/sub children forked,
ChipWorker.initdone,
ChipCallables uploaded H2D, prewarm complete — by the timeinit()returns, instead of being deferred to the first
run()._start_hierarchical()frominit()(inside the try block, so astartup failure is reaped by
_cleanup_partial_init).run,create_host_buffer, remote require/register/unregister). Each is reachedonly after
init(), so the calls were idempotent no-ops.init()andrun()now takethe post-start broadcast path; a childless L3 registering a Python fn after
init()raises eagerly instead of being silently snapshotted.hierarchical_level_runtime) to describe eager startup.
test_host_workertests that asserted the removed"init-done-but-not-started" window; stub
_start_hierarchicalin theremote-id bookkeeping unit test that uses a fake C worker.
Testing
test_host_worker(63),test_callable_identity+test_remote_l3_lifecycle(84),test_l4_recursive+test_ensure_prepared+test_error_propagation(27).(320 cores,
OMP_NUM_THREADSunset) hangs fork-heavy L3 st tests due to apre-existing torch-OpenMP fork-safety issue that reproduces identically on
mainand is unrelated to this change (torchspawns a 320-thread poolat import, before any
Workerctor, so a forked chip child inherits acorrupt pool and deadlocks). Relying on CI for onboard/sim coverage.