Feat: add bounded whole-run FIFO admission - #1541
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 Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe change introduces generation-safe pipeline leases, run-scoped FIFO admission, task-acceptance signaling, lease-aware mailbox dispatch, and AICore image-aware stream reuse across C++, Python bindings, device runners, documentation, and tests. ChangesPipeline admission and acceptance
Estimated code review effort: 5 (Critical) | ~120 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 9
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/orchestrator.cpp (1)
177-206: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winAvoid racing
fail_run_submissionagainst scheduler activation. The run phase is sampled without the completion/run lock, andcancel_unstarted_runis later allowed to mark slotsPENDING/READY/FREEasFAILED; if the closed run has already been activated toEXECUTING, the scheduler can dispatch the same slot asRUNNINGconcurrently, losing one terminal transition. Guard cancellation by the current run phase and use stateful atomic transitions for slot states.🤖 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/orchestrator.cpp` around lines 177 - 206, Update fail_run_submission and cancel_unstarted_run to synchronize phase inspection with the run/completion state, preventing cancellation from racing scheduler activation. Recheck the current phase while holding the appropriate lock before cancelling, and make each slot-state update use compare-and-exchange transitions that only convert valid non-running states; never mark an EXECUTING/RUNNING slot FAILED or overwrite a concurrent scheduler transition.
🧹 Nitpick comments (5)
src/common/platform/onboard/host/device_runner_base.h (1)
297-320: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDoc block missing the new
aicore_image_hashparam. The@paramlist enumerates every other argument; add a line for the new one so the staging contract stays self-describing (same forrecord_host_orch_callableat Line 332).🤖 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/platform/onboard/host/device_runner_base.h` around lines 297 - 320, Update the documentation for record_device_orch_callable to add an `@param` entry describing aicore_image_hash, matching its position and purpose in the function signature. Also add the corresponding parameter documentation to record_host_orch_callable, without changing implementation behavior.src/common/hierarchical/scheduler.cpp (1)
366-369: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: fold the
active_run_cb ? scoped : unscopedbranching into helpers. The same "resolve active run, bail on invalid, pick scoped vs unscoped queue op" shape is now repeated in all four dispatch paths plus the wait predicate; a pair of small private helpers (e.g.resolve_active_run()returningstd::optional<RunId>andpop_ready(...)) would keep future queue-API changes in one place.🤖 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 366 - 369, Consolidate the repeated active-run resolution and scoped/unscoped ready-queue branching across all four dispatch paths and the wait predicate in the scheduler class. Add small private helpers such as resolve_active_run() and pop_ready(...) that preserve invalid-run early returns while centralizing queue API selection, then update the affected paths to use them.python/simpler/task_interface.py (1)
1277-1284: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd one-line docstrings to match the neighbouring properties.
Every other property in this class documents its unit/semantics;
pipeline_depthandruntime_slot_countare the odd ones out.🤖 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/task_interface.py` around lines 1277 - 1284, Add concise one-line docstrings to the pipeline_depth and runtime_slot_count properties in the task interface, matching the neighboring properties’ style and documenting each value’s unit or semantics.tests/ut/cpp/hierarchical/test_orchestrator.cpp (1)
656-665: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winA failed
ASSERT_*between here andthird.get()turns the test into a hang.
ASSERT_TRUEreturns from the test body, and~futurefromstd::asyncblocks until the task completes — butbegin_run()only returns once a slot frees, which is exactly what the skipped lines would have done. Consider releasing the admission slot from a scope guard so an assertion failure fails fast instead of wedging the suite.🤖 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 656 - 665, Ensure the asynchronous third begin_run test always releases an admission slot before any ASSERT_* can exit the test, by adding a scope guard around the first run’s slot ownership and consumption cleanup. Preserve the existing successful-path behavior while making cleanup execute during assertion-induced early returns, so third.get() cannot leave the test hanging.src/common/platform/onboard/host/device_runner_base.cpp (1)
143-150: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winBank validation accepts
< PTO_PIPELINE_MAX_DEPTHbut only two banks exist.
select_arena_bankadmits any id belowPTO_PIPELINE_MAX_DEPTH, while every consumer is a binarybank == 0 ? bank0 : bank1. With today's depth of 2 this is benign, but raising the depth constant would silently alias banks 2..N onto bank1 and share arenas across concurrently admitted runs. Prefer indexing an array of banks, or validate against the actual bank count.♻️ Sketch
- if (bank_id >= PTO_PIPELINE_MAX_DEPTH) { - LOG_ERROR("arena bank %u is outside [0, %u)", bank_id, PTO_PIPELINE_MAX_DEPTH); + if (bank_id >= kArenaBankCount) { + LOG_ERROR("arena bank %u is outside [0, %u)", bank_id, kArenaBankCount); return -1; }with
gm_heap_arena_[kArenaBankCount]etc. so the selection isgm_heap_arena_[g_arena_bank].Also applies to: 193-211
🤖 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/platform/onboard/host/device_runner_base.cpp` around lines 143 - 150, Update select_arena_bank and the associated arena consumers to validate against the actual number of supported banks rather than PTO_PIPELINE_MAX_DEPTH. Define or reuse a two-bank count, reject IDs outside that range, and replace binary bank-selection logic with indexed access so each accepted bank maps to its own arena without aliasing.
🤖 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/orchestrator.md`:
- Around line 101-106: The documentation around the “later submit” behavior must
distinguish admission from acceptance: clarify that acceptance unblocks only an
already-admitted successor, while a third submission with bounded depth two
remains blocked in begin_run() until a terminal run releases its lease. Preserve
the existing endpoint acceptance and completion fallback details.
In `@docs/worker-manager.md`:
- Around line 175-181: Update the documentation around the mailbox polling loop
to state that on_accept() is a one-shot callback, invoked only on the first
observation of TASK_ACCEPTED. Document the acceptance_observed guard and
preserve TASK_DONE as the mailbox ownership boundary, so implementations do not
repeatedly decrement per-run acceptance accounting.
In `@python/simpler/worker.py`:
- Line 1544: In the comment near the worker implementation, replace the
ambiguous multiplication symbol “×” in “N×40B” with the ASCII character “x”,
preserving the rest of the comment unchanged.
- Around line 180-186: Update _PIPELINE_LEASE_FMT to use the C ABI’s packed
12-byte layout for PipelineSlotLease: uint32 slot_id, uint32 reserved, and
uint64 generation with no padding between fields. Ensure _OFF_ARGS resolves to
the established MAILBOX_OFF_ARGS offset and remains compatible with the C++
definition in pto_runtime_c_api.h.
In `@src/a2a3/platform/onboard/host/device_runner.cpp`:
- Around line 621-624: Make acceptance durable across the mailbox transition in
the dispatch flow around publish_task_accepted and
LocalMailboxEndpoint::run_with_accept: ensure a parent that observes TASK_DONE
without previously sampling TASK_ACCEPTED still invokes on_accept exactly once.
Use a latched acceptance flag/counter or equivalent state, and preserve the
existing behavior for parents that observe TASK_ACCEPTED first so
decrement_run_accepts is not duplicated.
In `@src/common/hierarchical/orchestrator.cpp`:
- Around line 360-371: Update Orchestrator::decrement_run_accepts to acquire
run->completion_mu before changing pending_accepts and before calling
completion_cv.notify_all(). Keep the decrement, underflow handling, error
recording, and notification behavior intact while ensuring the predicate
transition is synchronized with wait_run_accepted.
In `@src/common/hierarchical/types.cpp`:
- Around line 70-82: Update ReadyQueue’s unscoped access paths, including
try_pop(TaskSlot&) and try_front(TaskSlot&), so they no longer select buckets by
unordered_map iteration order; remove these overloads or add separate
run-insertion-order tracking that preserves FIFO, and update callers such as
try_pop_group(TaskSlot&) and try_pop_single(worker_id, TaskSlot&) accordingly.
In `@src/common/hierarchical/worker_manager.cpp`:
- Around line 397-404: Move the acceptance callback out of the mailbox-locked
section in the worker task flow around read_mailbox_state and
LocalMailboxEndpoint::run. Record that TASK_ACCEPTED was observed while polling,
then invoke on_accept_ only after the mailbox round-trip and mailbox_mu_ have
been released, preserving exactly-once callback behavior.
In `@tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py`:
- Line 114: After unpacking tensors into first_a, first_b, first_out, second_a,
second_b, and second_out, explicitly delete those references before
free_host_buffer runs; apply the same cleanup at the corresponding later
unpacking block, while guarding it for early-failure paths where the names may
not be bound.
---
Outside diff comments:
In `@src/common/hierarchical/orchestrator.cpp`:
- Around line 177-206: Update fail_run_submission and cancel_unstarted_run to
synchronize phase inspection with the run/completion state, preventing
cancellation from racing scheduler activation. Recheck the current phase while
holding the appropriate lock before cancelling, and make each slot-state update
use compare-and-exchange transitions that only convert valid non-running states;
never mark an EXECUTING/RUNNING slot FAILED or overwrite a concurrent scheduler
transition.
---
Nitpick comments:
In `@python/simpler/task_interface.py`:
- Around line 1277-1284: Add concise one-line docstrings to the pipeline_depth
and runtime_slot_count properties in the task interface, matching the
neighboring properties’ style and documenting each value’s unit or semantics.
In `@src/common/hierarchical/scheduler.cpp`:
- Around line 366-369: Consolidate the repeated active-run resolution and
scoped/unscoped ready-queue branching across all four dispatch paths and the
wait predicate in the scheduler class. Add small private helpers such as
resolve_active_run() and pop_ready(...) that preserve invalid-run early returns
while centralizing queue API selection, then update the affected paths to use
them.
In `@src/common/platform/onboard/host/device_runner_base.cpp`:
- Around line 143-150: Update select_arena_bank and the associated arena
consumers to validate against the actual number of supported banks rather than
PTO_PIPELINE_MAX_DEPTH. Define or reuse a two-bank count, reject IDs outside
that range, and replace binary bank-selection logic with indexed access so each
accepted bank maps to its own arena without aliasing.
In `@src/common/platform/onboard/host/device_runner_base.h`:
- Around line 297-320: Update the documentation for record_device_orch_callable
to add an `@param` entry describing aicore_image_hash, matching its position and
purpose in the function signature. Also add the corresponding parameter
documentation to record_host_orch_callable, without changing implementation
behavior.
In `@tests/ut/cpp/hierarchical/test_orchestrator.cpp`:
- Around line 656-665: Ensure the asynchronous third begin_run test always
releases an admission slot before any ASSERT_* can exit the test, by adding a
scope guard around the first run’s slot ownership and consumption cleanup.
Preserve the existing successful-path behavior while making cleanup execute
during assertion-induced early returns, so third.get() cannot leave the test
hanging.
🪄 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: 976e895c-87d2-41d9-ab42-f455679ef0bc
📒 Files selected for processing (45)
docs/orchestrator.mddocs/task-flow.mddocs/worker-manager.mdpython/bindings/task_interface.cpppython/bindings/worker_bind.hpython/simpler/orchestrator.pypython/simpler/task_interface.pypython/simpler/worker.pysrc/a2a3/platform/onboard/host/device_runner.cppsrc/a2a3/platform/onboard/host/device_runner.hsrc/a2a3/runtime/host_build_graph/host/runtime_maker.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cppsrc/a5/runtime/host_build_graph/host/runtime_maker.cppsrc/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cppsrc/common/hierarchical/orchestrator.cppsrc/common/hierarchical/orchestrator.hsrc/common/hierarchical/scheduler.cppsrc/common/hierarchical/scheduler.hsrc/common/hierarchical/types.cppsrc/common/hierarchical/types.hsrc/common/hierarchical/worker.cppsrc/common/hierarchical/worker.hsrc/common/hierarchical/worker_manager.cppsrc/common/hierarchical/worker_manager.hsrc/common/platform/onboard/host/c_api_shared.cppsrc/common/platform/onboard/host/device_runner_base.cppsrc/common/platform/onboard/host/device_runner_base.hsrc/common/platform/sim/host/c_api_shared.cppsrc/common/task_interface/chip_callable_layout.hsrc/common/task_interface/prepare_callable_common.hsrc/common/utils/fnv1a_64.hsrc/common/worker/chip_worker.cppsrc/common/worker/chip_worker.hsrc/common/worker/pipeline_contract.hsrc/common/worker/pipeline_slot_pool.hsrc/common/worker/pto_runtime_c_api.htests/st/a2a3/host_build_graph/run_stream_reuse/kernels/aiv/kernel_sub.cpptests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.pytests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.pytests/ut/cpp/hierarchical/test_orchestrator.cpptests/ut/cpp/hierarchical/test_pipeline_contract.cpptests/ut/cpp/hierarchical/test_scheduler.cpptests/ut/cpp/types/test_chip_callable_upload_immutable.cpptests/ut/py/test_callable_identity.pytests/ut/py/test_worker/test_host_worker.py
8fcabcf to
17d945b
Compare
|
B2 was rebuilt and force-with-lease updated as a single commit on the current main. Key correction: the FIFO head can execute while its graph callback is still open, so existing L3-L2 interactive callbacks no longer deadlock. A successor still cannot dispatch until the previous run is terminal. The busy-SUB requeue path now preserves the active run partition, and unscoped queue access preserves run insertion order. Validation is recorded in the PR body. Final hardware task |
255ed4b to
6b39042
Compare
6b39042 to
fe60768
Compare
|
Head is [P1] 取消与 dispatch 竞态两半都修了。前半：三条派发路径（SUB / NEXT_LEVEL group / single）的「读 READY … 写 RUNNING」合并成一次原子交换，抽成
确定性竞态测试用的是真实取消路径： [P1] 构图失败提前释放 RUNNING task 使用的 producer第一个循环现在记录本次成功 [P1] whole-run FIFO 未覆盖 callback 中的直接 device control
四个入口我全部当成 device effect，没有做 host-only / device-effect 二分——它们都跨 mailbox 在 child 上执行真实操作， 这条修复本身引入了一个我没预料的缺陷，是 onboard 回归测试抓出来的：那四个 nanobind 绑定不释放 GIL，于是一个可能无限期的 主线程连 [P2] API 与文档
Validation
原 onboard FIFO 用例的 |
fe60768 to
7c65ace
Compare
|
Head is [P1] poison 与 cancellation 争同一 slot 的清理权核实成立： [P1] group front→pop 窗口会终止进程核实成立，且 [P1] device-control admission 既漏拦又误拦三个子问题一个正解，按你的处方改了：
binding 释放 GIL —— 上一轮我在这上面栽过：持 GIL 做无界等待会冻住整个解释器，连让 active run 收尾的线程都跑不动。 [P2] 已终止 run 的 partition 可被重新创建
[P2] depth 文档矛盾
Validation
[P1-D] 采纳你的方案 1，并按你的拆分做成独立 PR我同意 control lane 那条路不成立 —— 它只排序 control、排不了 Scheduler task，而且按到达顺序反而会把 N 的 cleanup 挡在 N+1 之后；要补齐就得引入 run epoch、cleanup marker、task/control 共同 ownership 和多 mailbox 原子占用，等于把 run fence 重做一遍且更大。 所以我按你给的形状实现 B2a：automatic run finalization——两阶段 native fence（ 它会是独立 PR，本 PR 需要 rebase 到它之上，不先合。我不把 P1-D 塞进这个 diff：它改的是 run 终结的所有权（ 有一点我会在 B2a 里一并处理，因为它属于同一个契约：callback control 和 cleanup control 都显式携带 run_id / cleanup token，所有 admission wait 都在取 Python 锁之前。本 PR 已经把 callback 侧做成这样了，cleanup 侧要等 finalizer 存在才有地方挂。 |
7c65ace to
319ff8b
Compare
Reserve a generation-safe pipeline lease before graph construction, so a
callback learns there is no capacity before it builds a DAG rather than after:
`begin_run` blocks on the lease, and the depth it admits to is the minimum the
child backends published rather than a constant. A backend without a depth-two
contract therefore keeps serial behavior instead of being handed a slot-1 lease
it cannot serve.
Ready queues are partitioned by run and the scheduler pops only from the
partition of the run holding the FIFO head, which is what keeps two admitted
runs from interleaving device work. The head may execute while its own graph
callback is still open — existing callbacks submit device work and wait for L2
communication before returning — while successors stay gated until every prior
run is terminal.
Popping a slot is not owning it. A run whose callback throws fails and consumes
its own unstarted slots, and its fence can then release the run's lease, all
while the scheduler sits between reading READY and launching. Every dispatch
path now claims the slot with an atomic READY -> RUNNING exchange at the point
of launch, and a lost claim leaves the slot to whoever moved it. The
cancellation side retries its own exchange, because a producer completing
concurrently moves a consumer PENDING -> READY and a single attempt would leave
that slot dispatchable under a failing run. `dispatchable_run_id()` is what the
scheduler asks for: the head must also be EXECUTING and still own its lease.
Cancellation releases dependency references only for the slots it actually
failed. Releasing one held by a slot that is still RUNNING would let its
producer reach CONSUMED — and its HeapRing output be reclaimed — while the
device is still reading it, and the consumer's real completion would then
release the same reference twice.
`malloc` / `free` / `copy_to` / `copy_from` reach a child directly rather than
through a TaskSlot, so the ready-queue FIFO does not order them and the mailbox
mutex only serialises one command at a time. They now block until the calling
run holds the FIFO head, so a prepared successor cannot free or overwrite child
memory the active run is still reading. That wait can be long, so their
bindings release the GIL — holding it across the wait froze every other Python
thread, including the one that would have let the active run finish.
The lease reaches the runtime slot it names: `TaskSlot` carries
`{slot_id, generation}` through the local chip mailbox into `run_from_blob`, so
a production run executes under its own lease instead of unconditionally on
slot 0.
Both failure paths — a producer poisoning its consumer, and a run cancelling
its unstarted slots — claim a slot through the same exchange, and the winner is
the only one that writes the failure message. A plain store could put a
CONSUMED slot back to FAILED and then consume it, and release its dependency
references, a second time.
The group queue's head is observed before its target workers are checked, so a
run cancelling in that window can consume the slot and erase its partition.
That is a legal outcome, not a corrupt queue: throwing there would end the
process, because the scheduler thread has no handler.
A terminal run's ready-queue partitions stay erased. Both routes into
`enqueue_ready` can race the erase, and re-inserting would rebuild a partition
the scheduler will never read again.
Direct device control is ordered against the run that issued it, named
explicitly rather than read from `building_run_id_`: that field says a callback
is open somewhere, not that this thread is inside it, so a public
`Worker.copy_*` on another thread would otherwise be charged to whichever run
happens to be building. The wait happens before any Python lock is taken, since
a callback holding the provenance lock across it would block the paths that let
the active run finish, and its binding releases the GIL. Domain and region
creation are ordered the same way — they drive collective control on every
member chip.
A run whose own cleanup touches the device degrades itself to depth one. The
whole-run FIFO orders tasks; it cannot order cleanup, which happens after the
native fence and reaches a child through mailbox control rather than a
TaskSlot — N+1 allocating a domain while N is still releasing one can leave two
collectives each holding a different chip's mailbox and waiting for the other.
A run marks itself when it takes a CommDomain, an L3-L2 region or queue, a
remote slot reference, or calls direct device control, and the next submission
waits for that run's cleanup before it is admitted. Runs that only dispatch
tasks keep the full depth, so a control-free hot path is unaffected.
Graph callbacks are serialized, which is what makes that a fact rather than a
prediction: a predecessor's callback has always returned before the next
submission reaches the check, so whether it bears cleanup is already decided
and at most one such run can be outstanding. The wait deliberately does not
swallow its error the way pre-submission draining does — cleanup that failed
leaves collective device state this process can neither describe nor reclaim,
so the worker records it and refuses every later operation instead of admitting
work on top of it.
`Scheduler::Config::before_claim_cb` is a test seam, unset in production: the
instant between the pop and the claim is unreachable from outside, so the
losing side of the claim can only be exercised from there.
319ff8b to
5ccad44
Compare
2 participants
Summary
{slot_id, generation}through TaskSlot and the local chip mailbox into the B1 runtime slot.Rebase and compatibility correction
main@8e89f014; old Feat: add launch-acceptance flight fence #1467/B1 stack commits are no longer present.INVALID_RUN_ID.Validation
866 passed, 10 deselected.17d945bd: focused Python worker/L3-L2 UT184 passed.64/64 passed.task_20260728_235654_212932915886, exit 0.task_20260729_000230_235780427710, exit 0.l3_l2_message_queue, andl3_l2_orch_comm_stream.Non-goals
This PR establishes bounded whole-run admission and FIFO dispatch. The two-frame endpoint, HBG prepared epochs, and TMR prepared state remain B3, B4, and B5.