Skip to content

Add: generation-safe two-frame local endpoint - #1574

Open
Crane-Liu wants to merge 2 commits into
hw-native-sys:mainfrom
Crane-Liu:codex/worker-async-b3-two-frame-endpoint
Open

Add: generation-safe two-frame local endpoint#1574
Crane-Liu wants to merge 2 commits into
hw-native-sys:mainfrom
Crane-Liu:codex/worker-async-b3-two-frame-endpoint

Conversation

@Crane-Liu

@Crane-Liu Crane-Liu commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Summary

  • split the local worker mailbox into one control frame and two generation-identified task frames
  • allow the parent to publish a second request while the child still executes the first, while preserving one sequential native executor
  • validate run, slot, generation, and dispatch identity across accept and completion states, and poison uncertain endpoints
  • retire every pre-publish dispatch sequence on failure and reject oversized argument blobs before publication
  • enable two-frame admission only for a2a3 host_build_graph workers with pipeline depth at least two; TMR, A5, SUB, remote, and recursive workers keep the single-frame fallback

Dependency

Depends on #1541 (B2 whole-run FIFO admission). Until #1541 merges, this PR intentionally includes that predecessor commit in its diff.

Why TMR remains single-frame

An initial B3 validation enabled the endpoint for TMR as well. Real-device l3_l2_message_queue and l3_l2_orch_comm_stream then reproduced an AIV UB out-of-bounds failure. B3 now gates the endpoint to host_build_graph. TMR prepared-state and RuntimeEnv compatibility belongs to B5; both TMR regressions pass again through the legacy single-frame path.

Validation

Tested commit: be51bfd

  • Ruff: modified Python files passed
  • Python non-hardware unit tests: 875 passed, 10 deselected
  • C++ non-hardware unit tests: 66 of 66 passed
  • final post-rebase targeted tests: 89 Python tests and 3 core C++ suites passed
  • queued architecture probe task_20260729_022217_132110221055: a2a3, exit 0
  • a2a3 real hardware task task_20260729_022237_13463413851: exit 0
    • worker_async_endpoint: passed; request B reached ACCEPTED while request A remained ACTIVE, with one native execution at a time
    • worker_async_fifo: passed
  • prior TMR fallback validation task_20260729_005744_12965644321: l3_l2_message_queue and l3_l2_orch_comm_stream passed

No simulation validation was run, per the agreed workflow.

@coderabbitai

coderabbitai Bot commented Jul 29, 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: 3b959e88-863a-4564-a11a-d9aa716d755e

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

Pipeline leases and run-scoped queues implement bounded whole-run FIFO admission. Framed mailboxes support depth-aware concurrent dispatch, while bindings, documentation, and tests expose and validate the new protocol.

Changes

Pipeline admission and dispatch

Layer / File(s) Summary
Run lifecycle and active-run scheduling
src/common/hierarchical/orchestrator.*, src/common/hierarchical/types.*, src/common/hierarchical/scheduler.*, src/common/worker/pipeline_slot_pool.h, src/common/hierarchical/worker.*
Runs acquire generation-safe leases, ready queues track RunId, and schedulers dispatch only the active FIFO run.
Framed mailbox execution
python/simpler/worker.py, src/common/hierarchical/worker_manager.*
Mailbox frames carry task identity and pipeline leases; two-frame execution, acceptance states, endpoint capacity, and failure signaling are added.
Public bindings and protocol documentation
python/bindings/*, python/simpler/task_interface.py, docs/task-flow.md
Bindings expose pipeline depth and frame constants, GIL release is added, and the lease/mailbox flow is documented.
Admission and dispatch validation
tests/ut/*, tests/st/a2a3/host_build_graph/*
Tests cover queue ordering, admission depth, bounded capacity, two-frame execution, and whole-run FIFO behavior.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related PRs

Poem

A rabbit hops through frames of two,
With leases green and queues in view.
FIFO paths now neatly flow,
While accepted tasks softly glow.
“Run on!” cheers Bunny, ears held high.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 9.93% 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.
Title check ✅ Passed The title is concise and accurately summarizes the main change: a generation-safe two-frame local endpoint.
Description check ✅ Passed The description is clearly related to the changeset and matches the documented two-frame endpoint, admission, and validation work.

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.

@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: 6

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
src/common/hierarchical/scheduler.cpp (1)

465-492: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win

break after one dispatch keeps the second task frame idle.

has_capacity() now allows two in-flight tasks per NEXT_LEVEL worker, but the loop still dispatches at most one slot per worker per dispatch_ready() pass. If both tasks of a run are already READY and enqueued, the second frame is only filled on the next wake (a completion or a fresh notify_ready), so the intended overlap is lost for graphs that enqueue both dispatches before the scheduler runs.

♻️ Proposed fix: keep dispatching while capacity remains
             s.state.store(TaskState::RUNNING, std::memory_order_release);
             worker->dispatch(WorkerDispatch{slot, 0});
-            break;
+            if (!worker->has_capacity()) break;
         }
🤖 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/common/hierarchical/scheduler.cpp` around lines 465 - 492, Update
Scheduler::dispatch_next_level_singles so it continues popping and dispatching
READY slots for each worker while worker->has_capacity() remains true, instead
of breaking after the first successful dispatch. Preserve the existing
validation, run filtering, and misrouting behavior, while ensuring the loop
stops when capacity is full or no eligible slot is available.
🧹 Nitpick comments (6)
tests/ut/py/test_worker/test_host_worker.py (1)

256-268: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Add a small sleep to the mailbox spin waits.

Both deadline loops poll with no sleep, so the test thread holds the GIL almost continuously while _run_chip_main_loop — a Python thread in the same interpreter — needs it to advance the frame state. On a loaded runner this is a real starvation source against the 5s deadline, and other waits in this file already use time.sleep(0.001).

♻️ Proposed change
         _mailbox_store_i32(frame1_addr + _OFF_STATE, worker_mod._TASK_READY)
         deadline = time.monotonic() + 5.0
         while _mailbox_load_i32(frame1_addr + _OFF_STATE) != worker_mod._TASK_ACCEPTED_STATE:
             assert time.monotonic() < deadline
+            time.sleep(0.001)
         assert calls == 1
 
         release[0].set()
         assert entered[1].wait(5.0)
         assert _mailbox_load_i32(frame0_addr + _OFF_STATE) == worker_mod._TASK_DONE
         assert _mailbox_load_i32(frame1_addr + _OFF_STATE) == worker_mod._TASK_ACTIVE
         release[1].set()
         deadline = time.monotonic() + 5.0
         while _mailbox_load_i32(frame1_addr + _OFF_STATE) != worker_mod._TASK_DONE:
             assert time.monotonic() < deadline
+            time.sleep(0.001)
🤖 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/py/test_worker/test_host_worker.py` around lines 256 - 268, Add
time.sleep(0.001) inside both deadline polling loops in the mailbox test around
_OFF_STATE checks, allowing _run_chip_main_loop to acquire the GIL while
preserving the existing 5-second timeout and state assertions.
tests/ut/cpp/hierarchical/test_orchestrator.cpp (1)

752-762: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

A regression here hangs the suite instead of failing it.

std::async(std::launch::async, ...) returns a future whose destructor blocks until the task completes. If admission stops being released (the exact regression this test guards), begin_run() never returns, ASSERT_EQ at Line 761 aborts the test body, and ~future then blocks forever — the run that would free the lease is never completed. Consider completing/aborting the blocking work in a scope guard (or driving the release before any fatal assertion) so a failure surfaces as a failed test rather than a stuck CI job. The same shape applies to the replacement future at Lines 793-797.

🤖 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/cpp/hierarchical/test_orchestrator.cpp` around lines 752 - 762, Make
the asynchronous admission checks in the current test, including the “third” and
“replacement” futures, failure-safe so a fatal assertion cannot destroy a
still-blocked std::future and hang the suite. Add scoped cleanup or perform the
lease-release/abort operation before any fatal assertion that depends on future
readiness, ensuring blocked begin_run work is always unblocked while preserving
the existing success-path assertions.
src/common/hierarchical/types.h (1)

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

Use TaskSlot for task_slots instead of a hard-coded int32_t.

Orchestrator::register_run_slot pushes TaskSlot values and cancel_unstarted_run copies them back into std::vector<TaskSlot>; spelling the element type as the alias keeps the two in step if TaskSlot's underlying type ever changes.

♻️ Proposed change
-    std::vector<int32_t> task_slots;
+    std::vector<TaskSlot> task_slots;
🤖 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/common/hierarchical/types.h` at line 154, Update the task_slots member
declaration to use the TaskSlot alias as its vector element type instead of
hard-coded int32_t, preserving the existing container and behavior.
src/common/hierarchical/worker_manager.h (1)

143-149: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a static assert guarding MAILBOX_ARGS_CAPACITY against underflow.

The capacity is an unsigned subtraction of two derived offsets; if the header region ever grows past MAILBOX_OFF_FRAME_PROTOCOL the result silently wraps to a huge value instead of failing the build. The neighbouring layout invariants already have asserts — this one is worth the same treatment.

🛡️ Proposed addition
 static constexpr size_t MAILBOX_ARGS_CAPACITY =
     static_cast<size_t>(MAILBOX_OFF_FRAME_PROTOCOL) - static_cast<size_t>(MAILBOX_OFF_TASK_ARGS_BLOB);
+static_assert(
+    MAILBOX_OFF_TASK_ARGS_BLOB < MAILBOX_OFF_FRAME_PROTOCOL, "task args blob region overlaps the frame identity region"
+);
🤖 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/common/hierarchical/worker_manager.h` around lines 143 - 149, Add a
compile-time assertion immediately after the MAILBOX_ARGS_CAPACITY definition to
verify that MAILBOX_OFF_FRAME_PROTOCOL is at least MAILBOX_OFF_TASK_ARGS_BLOB
before the unsigned subtraction can underflow. Keep the existing capacity
calculation unchanged and follow the neighboring mailbox layout invariant
assertion style.
python/simpler/worker.py (1)

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

Reusing the lease slot_id field as the depth-publication channel deserves a comment.

Startup packs cw.pipeline_depth into the lease's slot field and the parent reads it back at line 4747, while every later write to the same offset means "lease slot id". A one-line note here (and at the reader) would keep the dual meaning discoverable.

🤖 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/simpler/worker.py` at line 1912, Add concise comments at the
_PIPELINE_LEASE_FMT.pack_into write and its corresponding parent reader near
line 4747 documenting that the lease slot_id field temporarily carries
pipeline_depth during startup, while later writes use it as the lease slot ID.
src/common/hierarchical/worker_manager.cpp (1)

550-566: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Early returns inside the frame lock rely on endpoint_poisoned_ for publish-turn escape.

The poisoned check (553) and non-IDLE check (559) return before publishing; the first relies on an already-poisoned flag and the second sets it, so waiters escape. Worth an explicit comment stating that invariant, since it is what keeps the next_publish_id_ sequence from stalling here.

🤖 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/common/hierarchical/worker_manager.cpp` around lines 550 - 566, Add a
concise comment in the run_two_frame frame-lock early-return block explaining
that these publish-skipping exits rely on endpoint_poisoned_ being set, allowing
publish-turn waiters to escape without stalling next_publish_id_. Keep the
existing poisoned and non-IDLE checks and control flow unchanged.
🤖 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 `@docs/task-flow.md`:
- Around line 296-328: Remove or rewrite the dormant-capacity paragraph in
docs/task-flow.md lines 296-328 so it no longer describes single-frame, unleased
production behavior and remains consistent with the active depth-two lease/FIFO
contract; update the local mailbox diagram in docs/task-flow.md lines 430-431 to
represent PipelineSlotLease or its {slot_id, generation} fields.

In `@python/simpler/worker.py`:
- Around line 1748-1792: Add a bounded yield to admission_loop after each
polling sweep, including when no control or task frame is admitted, using a very
short sleep or time.sleep(0). Preserve immediate handling of shutdown, control
requests, and ready frames while preventing the admission thread from
continuously contending for the GIL with task execution.
- Line 1762: Update both zip calls in the worker logic— the
frame_bufs/frame_addrs loop and the _chip_shms/_chip_pids loop—to pass
strict=True, preserving the existing iteration while raising an error if the
paired collections differ in length.

In `@src/common/hierarchical/worker_manager.cpp`:
- Around line 528-548: Ensure every pre-publish exit in
LocalMailboxEndpoint::run_two_frame consumes or invalidates its dispatch
sequence: update next_publish_id_ and notify publish_cv_ before returning on the
no-free-task-frame path, and handle exceptions caught before the publish block
similarly rather than rethrowing with the sequence stranded. If the sequence
cannot be advanced safely, poison the endpoint so later publish waits do not
block indefinitely.
- Around line 582-600: Update the run_two_frame argument-blob handling to check
blob_bytes against MAILBOX_ARGS_CAPACITY before writing the header or payload.
When oversized, fail the completion with ENDPOINT_FAILURE and return before
publishing the frame; retain the existing serialization only for blobs that fit.

In
`@tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py`:
- Around line 125-132: Update the cleanup in the worker async endpoint test so
host buffers are freed only after run is confirmed terminal: after
run.wait(20.0), recheck run.done and skip or defer st_worker.free_host_buffer
when the run remains active, mirroring the sibling worker async FIFO test.
Preserve the existing exception suppression and buffer cleanup for completed
runs.

---

Outside diff comments:
In `@src/common/hierarchical/scheduler.cpp`:
- Around line 465-492: Update Scheduler::dispatch_next_level_singles so it
continues popping and dispatching READY slots for each worker while
worker->has_capacity() remains true, instead of breaking after the first
successful dispatch. Preserve the existing validation, run filtering, and
misrouting behavior, while ensuring the loop stops when capacity is full or no
eligible slot is available.

---

Nitpick comments:
In `@python/simpler/worker.py`:
- Line 1912: Add concise comments at the _PIPELINE_LEASE_FMT.pack_into write and
its corresponding parent reader near line 4747 documenting that the lease
slot_id field temporarily carries pipeline_depth during startup, while later
writes use it as the lease slot ID.

In `@src/common/hierarchical/types.h`:
- Line 154: Update the task_slots member declaration to use the TaskSlot alias
as its vector element type instead of hard-coded int32_t, preserving the
existing container and behavior.

In `@src/common/hierarchical/worker_manager.cpp`:
- Around line 550-566: Add a concise comment in the run_two_frame frame-lock
early-return block explaining that these publish-skipping exits rely on
endpoint_poisoned_ being set, allowing publish-turn waiters to escape without
stalling next_publish_id_. Keep the existing poisoned and non-IDLE checks and
control flow unchanged.

In `@src/common/hierarchical/worker_manager.h`:
- Around line 143-149: Add a compile-time assertion immediately after the
MAILBOX_ARGS_CAPACITY definition to verify that MAILBOX_OFF_FRAME_PROTOCOL is at
least MAILBOX_OFF_TASK_ARGS_BLOB before the unsigned subtraction can underflow.
Keep the existing capacity calculation unchanged and follow the neighboring
mailbox layout invariant assertion style.

In `@tests/ut/cpp/hierarchical/test_orchestrator.cpp`:
- Around line 752-762: Make the asynchronous admission checks in the current
test, including the “third” and “replacement” futures, failure-safe so a fatal
assertion cannot destroy a still-blocked std::future and hang the suite. Add
scoped cleanup or perform the lease-release/abort operation before any fatal
assertion that depends on future readiness, ensuring blocked begin_run work is
always unblocked while preserving the existing success-path assertions.

In `@tests/ut/py/test_worker/test_host_worker.py`:
- Around line 256-268: Add time.sleep(0.001) inside both deadline polling loops
in the mailbox test around _OFF_STATE checks, allowing _run_chip_main_loop to
acquire the GIL while preserving the existing 5-second timeout and state
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: 86fdd0f3-23b4-46cc-82fc-99e1b1d498fa

📥 Commits

Reviewing files that changed from the base of the PR and between 7e66f46 and 789f990.

📒 Files selected for processing (24)
  • docs/task-flow.md
  • python/bindings/task_interface.cpp
  • python/bindings/worker_bind.h
  • python/simpler/task_interface.py
  • python/simpler/worker.py
  • src/common/hierarchical/orchestrator.cpp
  • src/common/hierarchical/orchestrator.h
  • src/common/hierarchical/scheduler.cpp
  • src/common/hierarchical/scheduler.h
  • src/common/hierarchical/types.cpp
  • src/common/hierarchical/types.h
  • src/common/hierarchical/worker.cpp
  • src/common/hierarchical/worker.h
  • src/common/hierarchical/worker_manager.cpp
  • src/common/hierarchical/worker_manager.h
  • src/common/worker/pipeline_slot_pool.h
  • tests/st/a2a3/host_build_graph/worker_async_endpoint/kernels/orchestration/long_vector_orch.cpp
  • tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py
  • tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py
  • tests/ut/cpp/hierarchical/test_orchestrator.cpp
  • tests/ut/cpp/hierarchical/test_pipeline_contract.cpp
  • tests/ut/cpp/hierarchical/test_scheduler.cpp
  • tests/ut/py/test_callable_identity.py
  • tests/ut/py/test_worker/test_host_worker.py

Comment thread docs/task-flow.md
Comment thread python/simpler/worker.py
Comment thread python/simpler/worker.py Outdated
Comment thread src/common/hierarchical/worker_manager.cpp
Comment thread src/common/hierarchical/worker_manager.cpp
@Crane-Liu
Crane-Liu force-pushed the codex/worker-async-b3-two-frame-endpoint branch from 789f990 to be51bfd Compare July 29, 2026 09:22
@Crane-Liu Crane-Liu changed the title Feat: add generation-safe two-frame local endpoint Add: generation-safe two-frame local endpoint Jul 29, 2026
@Crane-Liu

Copy link
Copy Markdown
Contributor Author

Addressed the review summary and refreshed the branch on the latest main:

  • Scheduler now fills all available local endpoint frames in one pass.
  • Async admission tests always unblock their futures on failure.
  • TaskSlot typing, mailbox offset underflow guards, startup lease/depth comments, and GIL-friendly unit-test polling are updated.
  • The suggested production admission sleep was intentionally not added: .claude/rules/codestyle.md forbids sleeps/yields on dispatch paths, and docs/investigations/2026-07-host-dispatch-latency-budget.md records the measured latency cost. The existing poll remains the documented contract.
  • The early-return poison comment is obsolete after the implementation change: all pre-publish exits explicitly retire their sequence instead of relying on poison.

Validation on be51bfdf: 875 Python UT passed (10 deselected), 66 C++ UT passed, and queued a2a3 hardware task task_20260729_022237_13463413851 passed both endpoint and FIFO scenes. No simulation run was performed.

- Separate control traffic from two identified task frames
- Accept a queued frame while one sequential child executor is active
- Bound parent dispatch capacity and fill every available frame
- Retire publish ordering across every pre-publish failure path
- Reject oversized argument blobs before publishing mailbox state
- Preserve single-frame fallbacks on unsupported local backends
- Keep failure-path tests and host-buffer cleanup terminal-safe
- Release the GIL around native execution so admission remains live
@Crane-Liu
Crane-Liu force-pushed the codex/worker-async-b3-two-frame-endpoint branch from be51bfd to d6dfaa9 Compare July 29, 2026 09:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant