Skip to content

feat(runtime): add bounded distributed dispatch handles - #2270

Merged
YunjiQin merged 1 commit into
hw-native-sys:mainfrom
Crane-Liu:codex/worker-async-q1-dispatch-handle
Aug 12, 2026
Merged

YunjiQin merged 1 commit into
hw-native-sys:mainfrom
Crane-Liu:codex/worker-async-q1-dispatch-handle

Conversation

@Crane-Liu

@Crane-Liu Crane-Liu commented Aug 3, 2026 •

Copy link
Copy Markdown
Contributor

Summary

  • add DistributedWorker.submit() and public DistributedRunHandle while keeping run() and __call__() blocking
  • bound dispatch metadata to two reusable frames; a third submission drains the oldest handle before reuse
  • retain per-dispatch arguments, CallConfig, generated task metadata, and native handles through terminal completion
  • allocate HOST intermediates per frame so two in-flight dispatches do not share mutable scratch
  • serialize DFX submissions, isolate backpressure errors, and make close() safe for concurrent callers
  • submit persistent orchestration directly from the caller through Simpler Worker.submit(), retaining CommDomains without a Python dispatcher queue or background thread
  • publish provisional dispatch ownership before native submission and recover an accepted Simpler handle across interruption so its frame cannot be released early

Dependency

This is Q1 in the worker asynchronous pipeline stack. It is rebased on current PyPTO main (df3ac122) and uses the existing runtime gitlink at 3165cc89; this PR does not change the runtime submodule. Serving integration, end-to-end queue semantics, and cross-layer performance acceptance remain outside this PR.

Validation

  • rebuilt the worktree bindings and reinstalled the editable package
  • tests/ut/runtime/test_distributed_worker.py: 144 passed
  • tests/ut/runtime: 587 passed
  • full tests/ut: 9389 passed, 3 skipped
  • full pre-commit run --all-files: passed, including Ruff, Pyright, Markdown, C/C++, headers, English-only, docs en/zh parity/navigation/symbol coverage, broad-exception, and operator-name checks

@coderabbitai

coderabbitai Bot commented Aug 3, 2026 •

Copy link
Copy Markdown

Review Change Stack

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: f1501dcf-dd8d-4a27-9c0e-95fd469c7ab5

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

The runtime adds bounded asynchronous distributed dispatch through DistributedWorker.submit() and DistributedRunHandle. It manages frame reuse, buffer lifetimes, timeouts, errors, persistent domains, shutdown draining, and blocking compatibility through run().

Changes

Asynchronous distributed execution

Layer / File(s) Summary
Public submission and handle flow
python/pypto/runtime/distributed_runner.py, python/pypto/runtime/__init__.py, tests/ut/runtime/test_distributed_worker.py, docs/en/user/distributed/03-execution.md, docs/zh/user/distributed/03-execution.md, runtime
Adds DistributedWorker.submit() and the public DistributedRunHandle export. Dispatch metadata and arguments remain retained until completion. run() waits on asynchronous submission. Documentation and routing tests describe the new API.
Bounded frames and shutdown lifecycle
python/pypto/runtime/distributed_runner.py, tests/ut/runtime/test_distributed_worker.py, tests/st/distributed/test_l3_device_tensor.py
Adds two reusable dispatch frames, FIFO backpressure, timeout and error caching, frame retirement, input lifetime retention, and shutdown draining. System tests submit three dispatches with separate mutable buffers.
Persistent asynchronous dispatch
python/pypto/runtime/distributed_runner.py, tests/ut/runtime/test_distributed_worker.py
Persistent workers submit native handles and retain per-program domains and requests. Tests cover domain reuse, keepalive release, failure propagation, finalization ordering, and teardown errors.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

Poem

I’m a rabbit with handles to hop,
Two frames keep dispatch in a loop.
FIFO queues wait, then spring free,
Errors rest where their results should be.
Close drains the burrow with care—
Async carrots complete in the air.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 36.08% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description check ✅ Passed The description directly explains the bounded asynchronous dispatch changes and the validation results.
Title check ✅ Passed The title clearly and concisely identifies the new bounded distributed dispatch handles.

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread python/pypto/runtime/distributed_runner.py Outdated
Comment thread python/pypto/runtime/distributed_runner.py

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (3)
docs/en/user/distributed/03-execution.md (1)

76-81: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider documenting the handle API surface.

The section names DistributedRunHandle but does not describe its methods. result(timeout=...), wait(timeout=...), and done are 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 value

Remove the unused generation field.

_DispatchFrame.generation is only incremented when acquiring a dispatch frame. If stale frame references are not detected with generation, 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 value

Remove the two unused events and align the test double with the asserted failure.

request_finalizer_started and allow_request_finalizer_to_finish are created and set, but no code waits on them. They add no synchronization to this test.

worker_submit also discards fn, so failing_entry never runs. The asserted error comes only from native.complete(RuntimeError("persistent dispatch failed before cleanup")). The test therefore does not exercise a request-side entry failure, which its name implies. Either invoke fn inside worker_submit so failing_entry raises, 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

📥 Commits

Reviewing files that changed from the base of the PR and between 4389165 and 15fbbe0.

📒 Files selected for processing (7)
  • docs/en/user/distributed/03-execution.md
  • docs/zh/user/distributed/03-execution.md
  • python/pypto/runtime/__init__.py
  • python/pypto/runtime/distributed_runner.py
  • runtime
  • tests/st/distributed/test_l3_device_tensor.py
  • tests/ut/runtime/test_distributed_worker.py

Comment thread python/pypto/runtime/distributed_runner.py Outdated
Comment thread python/pypto/runtime/distributed_runner.py
@lyfne123

lyfne123 commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

please fix ai comment

@Crane-Liu
Crane-Liu force-pushed the codex/worker-async-q1-dispatch-handle branch from 15fbbe0 to 03aec67 Compare August 10, 2026 09:13
@YunjiQin

Copy link
Copy Markdown
Collaborator

Proposal: remove the persistent-only Python dispatcher queue/thread and submit the persistent orchestration directly from DistributedWorker.submit().

The queue was necessary in the original persistent implementation (6ff376a3 / #2095): one long-lived outer Worker.run(self._persistent_orch) owned all requests, so the orchestration callback needed a queue to receive commands and had to run on a background thread. That architecture was removed by 340d6c56 / #2163, which changed persistent execution to one public Worker.run(run_request) fence per request. This PR changes that again to one Worker.submit(run_request) per request.

At this point the queue is only a synchronous RPC hop:

caller: queue.put(request) -> wait for submitted Event
thread: queue.get() -> Worker.submit(run_request) -> set Event

It provides no buffering or additional serialization: public DistributedWorker.submit() holds _dispatch_submit_mu until _submit_persistent() receives the native handle, while Simpler already serializes graph construction with Worker._submit_mu. Simpler supports submitting from a non-init-owner thread, establishes the callback run context on whichever thread calls submit(), and binds CommDomains to _RunResources, not to a Python thread. The native close/teardown owner-thread restriction is separate and is not solved by this dispatcher.

I suggest replacing the queue path with:

_submit_persistent
  -> build run_request with the persistent domain provider/reset logic
  -> self._w.submit(run_request)
  -> publish DistributedRunHandle

Keep _persistent_domains_by_program, domain validation/reset/detach/release, _dispatch_submit_mu, active-handle draining, and the persistent terminal-error policy. Close should stop admission, drain every published native handle, release retained domains when safe, and then close the underlying Worker. The queue/STOP sentinel, _PersistentRequest, background thread, ready/submitted events, and start/stop lifecycle can then be removed.

This still permits the intended overlap: construction of request N+1 happens synchronously in its Worker.submit() call while request N executes on device. The current queue does not make graph construction asynchronous because the caller waits for request.submitted before returning. If we eventually want submit() to return before graph construction finishes, that would require a real two-stage handle plus durable cancellation/ownership publication; the current Event proxy does not provide that contract.

Removing the queue also eliminates several cross-thread cancellation/error-transfer states implicated by the review findings. It does not by itself fix the accepted-native-run publication window: acceptance through installation in _active_dispatch_handles must still be interruption-safe before a frame may be released. I would add regression coverage for persistent direct submission, first-domain warm-up followed by two bounded handles, close/drain plus single domain release, synchronous graph failure, and interruption at the native-handle publication boundary, and update both persistent lifecycle guides accordingly.

@Crane-Liu
Crane-Liu force-pushed the codex/worker-async-q1-dispatch-handle branch from 03aec67 to 13d8d1d Compare August 11, 2026 11:17
@Crane-Liu

Copy link
Copy Markdown
Contributor Author

@YunjiQin Addressed in 13d8d1df:

  • Removed the persistent-only queue, STOP sentinel, _PersistentRequest, background thread, readiness/submission events, and start/stop lifecycle. Persistent requests now build their retained-domain callback and call Simpler Worker.submit() directly on the DistributedWorker.submit() caller thread.
  • Preserved domain validation, reset, detach/release, terminal-error handling, _dispatch_submit_mu, bounded active handles, close draining, and single domain release.
  • Closed the accepted-native-run publication window by installing provisional frame/handle ownership before native submission. If interruption lands before the returned handle is stored, the serialized submit recovers the sole newly accepted Simpler handle and keeps the frame active for normal drain/close; a frame is released immediately only when no native handle was accepted.
  • Added regression coverage for direct persistent submission, warm-up followed by two bounded handles and one close-time domain release, synchronous graph failure, and interruption at the native-handle publication boundary. Updated the English and Chinese persistent lifecycle guides.

The branch is also rebased onto current main. Focused runtime tests, full unit tests, and the complete pre-commit suite pass.

@Crane-Liu
Crane-Liu force-pushed the codex/worker-async-q1-dispatch-handle branch from 13d8d1d to 1370137 Compare August 12, 2026 02:40
@YunjiQin
YunjiQin merged commit a7baa03 into hw-native-sys:main Aug 12, 2026
15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

3 participants