feat(runtime): add bounded distributed dispatch handles - #2270
Conversation
📝 WalkthroughWalkthroughThe runtime adds bounded asynchronous distributed dispatch through ChangesAsynchronous distributed execution
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 15fbbe0536
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| tensors = frame.tensors | ||
| tensors.clear() | ||
| tensors.update(state["base_tensors"]) |
There was a problem hiding this comment.
When submit() is used twice before either handle completes on a program whose HOST orchestrator emitted _alloc_intermediates, this copies state["base_tensors"] into each frame by reference. Those tensors are the pre-fork shared-memory scratch buffers created once in __init__, so both in-flight host_orch executions pass the same scratch tensors to chip tasks even if user inputs/outputs are distinct. That can silently cross-contaminate results; allocate separate intermediate tensors per dispatch frame during prepare or serialize such programs before publishing a second handle.
Useful? React with 👍 / 👎.
| if config is not None and config.enable_l2_swimlane: | ||
| _collect_l3_swimlane(compiled.output_dir, compiled.platform) | ||
| if _DfxOpts.from_run_config(config).any(): | ||
| _clear_dfx_dispatch_dirs(dfx_base) |
There was a problem hiding this comment.
When two DFX-enabled submissions are accepted back-to-back, this cleanup runs for the second submit while the first handle may still be writing or awaiting _collect_l3_swimlane(). Because _reset_dfx_dispatch_state restarts every run at rank*/d0, both handles use the same dfx_outputs paths; the second cleanup can delete the first dispatch's artifacts and both runs can clobber each other. Drain/serialize DFX submissions or give each handle a unique output prefix before clearing.
Useful? React with 👍 / 👎.
![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
🧹 Nitpick comments (3)
docs/en/user/distributed/03-execution.md (1)
76-81: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider documenting the handle API surface.
The section names
DistributedRunHandlebut does not describe its methods.result(timeout=...),wait(timeout=...), anddoneare public. A short list helps callers poll or bound their wait.🤖 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 `@docs/en/user/distributed/03-execution.md` around lines 76 - 81, Add a concise API summary in the distributed execution section for DistributedRunHandle, documenting the public result(timeout=...), wait(timeout=...), and done members, including that result returns the run outcome, wait supports bounded waiting, and done indicates completion.python/pypto/runtime/distributed_runner.py (1)
98-108: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused
generationfield.
_DispatchFrame.generationis only incremented when acquiring a dispatch frame. If stale frame references are not detected withgeneration, drop the field.🤖 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/pypto/runtime/distributed_runner.py` around lines 98 - 108, Remove the unused generation field from _DispatchFrame and delete any related initialization or increment logic when dispatch frames are acquired, while preserving the remaining frame metadata and reuse behavior.tests/ut/runtime/test_distributed_worker.py (1)
2503-2511: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the two unused events and align the test double with the asserted failure.
request_finalizer_startedandallow_request_finalizer_to_finishare created and set, but no code waits on them. They add no synchronization to this test.
worker_submitalso discardsfn, sofailing_entrynever runs. The asserted error comes only fromnative.complete(RuntimeError("persistent dispatch failed before cleanup")). The test therefore does not exercise a request-side entry failure, which its name implies. Either invokefninsideworker_submitsofailing_entryraises, or rename the test to describe native-handle failure propagation.♻️ Proposed cleanup of the unused events
m = patched_setup m["worker"]._live_domains = {} - request_finalizer_started = threading.Event() - allow_request_finalizer_to_finish = threading.Event() native = _ControlledNativeHandle()assert native.result_started.wait(timeout=2) - request_finalizer_started.set() # A failing request may already have submitted device work. Its caller # must not observe completion while the native handle is still finalizing. assert not caller_done.is_set() - allow_request_finalizer_to_finish.set() native.complete(RuntimeError("persistent dispatch failed before cleanup"))Also applies to: 2544-2551
🤖 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/runtime/test_distributed_worker.py` around lines 2503 - 2511, Remove the unused request_finalizer_started and allow_request_finalizer_to_finish events and their associated set/wait logic. Update worker_submit to invoke the submitted fn so failing_entry actually raises, preserving the test’s request-side failure scenario and existing assertions.
🤖 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/pypto/runtime/distributed_runner.py`:
- Around line 1634-1639: Update the backpressure drain in the dispatch
submission path around _active_dispatch_handles so draining oldest does not
propagate that handle’s cached failure to the new submitter. Consume the oldest
handle’s result while suppressing any exception, preserving the handle’s cached
outcome for its owner and allowing frame acquisition to continue.
- Around line 2381-2391: Make close() idempotent for concurrent callers by
adding and initializing a _closing flag alongside _accepting_dispatches, then
check and set _closing inside the existing _dispatch_submit_mu block before
entering teardown. Have subsequent close() calls return immediately, while
preserving the existing _closed guard and normal drain and release behavior.
---
Nitpick comments:
In `@docs/en/user/distributed/03-execution.md`:
- Around line 76-81: Add a concise API summary in the distributed execution
section for DistributedRunHandle, documenting the public result(timeout=...),
wait(timeout=...), and done members, including that result returns the run
outcome, wait supports bounded waiting, and done indicates completion.
In `@python/pypto/runtime/distributed_runner.py`:
- Around line 98-108: Remove the unused generation field from _DispatchFrame and
delete any related initialization or increment logic when dispatch frames are
acquired, while preserving the remaining frame metadata and reuse behavior.
In `@tests/ut/runtime/test_distributed_worker.py`:
- Around line 2503-2511: Remove the unused request_finalizer_started and
allow_request_finalizer_to_finish events and their associated set/wait logic.
Update worker_submit to invoke the submitted fn so failing_entry actually
raises, preserving the test’s request-side failure scenario and existing
assertions.
🪄 Autofix (Beta)
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: cfe77f61-a75c-4618-9fb1-ea0e7924b0f3
📒 Files selected for processing (7)
docs/en/user/distributed/03-execution.mddocs/zh/user/distributed/03-execution.mdpython/pypto/runtime/__init__.pypython/pypto/runtime/distributed_runner.pyruntimetests/st/distributed/test_l3_device_tensor.pytests/ut/runtime/test_distributed_worker.py
| if not self._active_dispatch_handles: | ||
| raise RuntimeError( | ||
| "DistributedWorker dispatch frames are occupied without owning handles" | ||
| ) | ||
| oldest = min(self._active_dispatch_handles, key=lambda handle: handle._dispatch_id) | ||
| oldest.result() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Backpressure re-raises another dispatch's failure to the new submitter.
oldest.result() re-raises the oldest dispatch's cached error. So submit() for a new, valid dispatch fails with the error of an earlier dispatch, and no frame is acquired. The owner of the old handle still sees that error on its own result() call, because result() caches it. Backpressure only needs the frame to be released, so drop the drained outcome here.
🐛 Proposed fix to isolate the drained failure
oldest = min(self._active_dispatch_handles, key=lambda handle: handle._dispatch_id)
- oldest.result()
+ try:
+ # Draining only needs the frame back. The drained handle caches
+ # its own outcome, so its owner still observes any failure.
+ oldest.result()
+ except BaseException: # noqa: BLE001 - reported to the drained handle's owner
+ pass📝 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.
| if not self._active_dispatch_handles: | |
| raise RuntimeError( | |
| "DistributedWorker dispatch frames are occupied without owning handles" | |
| ) | |
| oldest = min(self._active_dispatch_handles, key=lambda handle: handle._dispatch_id) | |
| oldest.result() | |
| if not self._active_dispatch_handles: | |
| raise RuntimeError( | |
| "DistributedWorker dispatch frames are occupied without owning handles" | |
| ) | |
| oldest = min(self._active_dispatch_handles, key=lambda handle: handle._dispatch_id) | |
| try: | |
| # Draining only needs the frame back. The drained handle caches | |
| # its own outcome, so its owner still observes any failure. | |
| oldest.result() | |
| except BaseException: # noqa: BLE001 - reported to the drained handle's owner | |
| pass |
🤖 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/pypto/runtime/distributed_runner.py` around lines 1634 - 1639, Update
the backpressure drain in the dispatch submission path around
_active_dispatch_handles so draining oldest does not propagate that handle’s
cached failure to the new submitter. Consume the oldest handle’s result while
suppressing any exception, preserving the handle’s cached outcome for its owner
and allowing frame acquisition to continue.
| def close(self) -> None: | ||
| """Release the Worker and comm rootinfo file. Idempotent.""" | ||
| if self._closed: | ||
| return | ||
| # Auto-free any DeviceTensors the caller forgot. Run BEFORE we set | ||
| # ``_closed`` so the per-op ``_require_open`` guard inside ``free`` | ||
| # still admits these calls, and BEFORE we tear down the underlying | ||
| # worker so the free path is still live. | ||
| self._close_owned_tensors() | ||
| # Serialize the admission transition with submit(). Once this block is | ||
| # entered, no partially published dispatch can race the drain below. | ||
| with self._dispatch_submit_mu: | ||
| if self._closed: | ||
| return | ||
| with self._dispatch_cv: | ||
| self._accepting_dispatches = False | ||
| self._dispatch_cv.notify_all() | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
The idempotence guard does not cover concurrent close() calls.
_closed becomes True only at line 2419, after the drain and the tensor release. Two threads that call close() at the same time both pass this guard and both run the full teardown, including self._w.close(). Gate on the admission flag inside the same mutex block so the second caller returns immediately.
🛡️ Proposed fix
with self._dispatch_submit_mu:
- if self._closed:
+ if self._closed or self._closing:
return
+ self._closing = True
with self._dispatch_cv:
self._accepting_dispatches = False
self._dispatch_cv.notify_all()Initialize self._closing = False next to self._accepting_dispatches in __init__.
🤖 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/pypto/runtime/distributed_runner.py` around lines 2381 - 2391, Make
close() idempotent for concurrent callers by adding and initializing a _closing
flag alongside _accepting_dispatches, then check and set _closing inside the
existing _dispatch_submit_mu block before entering teardown. Have subsequent
close() calls return immediately, while preserving the existing _closed guard
and normal drain and release behavior.
1 participant
Summary
DistributedWorker.submit()and publicDistributedRunHandlewhile keepingrun()and__call__()blockingCallConfig, generated task metadata, and native handles through terminal completionWorker.submit(), drain accepted work in FIFO order during close, and preserve cached failuresDependency chain
This is Q1 in the worker asynchronous pipeline stack and depends on:
The
runtimegitlink points to W4 commit999e70b73cf50bdffa7f9d5833d28e87523dc74d. Serving integration, end-to-end queue semantics, and cross-layer performance acceptance remain outside this PR.Validation
tests/ut/runtime/test_distributed_worker.py: 135 passedtests/ut/runtime: 542 passedtests/utwith 16 workers: 8620 passed, 2 skippedtask_20260803_064658_168142219291task_20260803_064735_169883129670, 1 passed on device 1No simulation run was used for this validation.