Update: make whole-run FIFO admission failure-safe - #1541
Conversation
📝 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 存在才有地方挂。 |
76e877d to
6310890
Compare
bb43d90 to
5deb624
Compare
Enforce bounded whole-run FIFO ordering across scheduling, execution, and teardown. Make run, remote-buffer, CommDomain, and shared-memory ownership transactional under interruption, with exact-once cleanup publication. Expand lifecycle documentation plus native, Python fault-injection, and onboard coverage.
ChaoWao
changed the title
`_child_prov_lock` is held across the *native* half of malloc / free / copy_to / copy_from, so every device op of every next-level worker serializes on one lock belonging to the parent worker. The bindings already release the GIL there, which makes the effect easy to miss: unrelated Python threads keep running, but a `copy_to` on chip 0 blocks a `malloc` on chip 1 for its whole duration. Keep `_child_prov_lock` as the bookkeeping lock — it still makes each provenance mutation/read atomic, and the safety-first ordering is unchanged (record after a successful alloc, revoke before a native free) — and take a per-worker lock around the native call. Ops on the same worker stay mutually exclusive, so a copy can still never overlap that buffer's free; ops on different workers now overlap. The per-worker lock is always acquired before `_child_prov_lock` and never the reverse, so the pair cannot deadlock. This is necessary but, on this base, not sufficient: `_control_reservation` (hw-native-sys#1541) holds `_submit_mu` across the same native call at the same granularity, so an 8-chip upload still serializes there. Measured on 8 x 910B2 with this commit alone, uploading rank-stacked weights one thread per shard: threaded 6.5-8.0 GB/s vs serial 8.6-10.3 GB/s — no gain. The following commit addresses that lock; with both, the same harness reaches 19.2-34.0 GB/s threaded.
A control command that belongs to no run takes `_submit_mu` through `_control_reservation` and holds it across the native call, so on this base it — not `_child_prov_lock` — is what serializes an 8-chip upload: all 8 threads enter `Orchestrator.copy_to` within 14 us of each other, then each child starts its copy within ~1 ms of the previous one finishing, at full link speed. What such a command needs is "no run may be admitted while I run", which is a property of the worker; two commands on different chips can both have that at the same time. Make `_submit_mu` a shared/exclusive lock: run admission takes it exclusively, control takes it shared. Writer-preferring, so control traffic cannot starve a submit. The reservation's re-entrancy is untouched: it short-circuits on the thread-local set before reaching the lock. Measured on 8 x 910B2, one upload thread per shard, fresh shared-memory bands so every page is read cold exactly once: provenance locks only threaded 6.5-8.0 GB/s serial 8.6-10.3 + this commit threaded 19.2-34.0 GB/s serial 9.4-10.5 Draft: this is a change to the serializer hw-native-sys#1541 introduced, and the reviewer of that area may prefer a different shape. Sent as a separate commit so it can be dropped or replaced without touching the provenance-lock fix.
`_child_prov_lock` is held across the *native* half of malloc / free / copy_to / copy_from, so every device op of every next-level worker serializes on one lock belonging to the parent worker. The bindings already release the GIL there, which makes the effect easy to miss: unrelated Python threads keep running, but a `copy_to` on chip 0 blocks a `malloc` on chip 1 for its whole duration. Keep `_child_prov_lock` as the bookkeeping lock — it still makes each provenance mutation/read atomic, and the safety-first ordering is unchanged (record after a successful alloc, revoke before a native free) — and take a per-worker lock around the native call. Ops on the same worker stay mutually exclusive, so a copy can still never overlap that buffer's free; ops on different workers now overlap. The per-worker lock is always acquired before `_child_prov_lock` and never the reverse, so the pair cannot deadlock. This narrows an invariant an existing test pins down, so that test is updated rather than left passing by accident: `test_free_holds_lock_across_native_free` asserted that the *parent worker's* lock is held across the native free. It now asserts the narrower exclusion actually needed — that worker's own lock held across the native call, `_child_prov_lock` released, and the revoke committed first. Provenance is keyed by (worker_id, ptr) and revoked before the native free, so a concurrent dispatch reads the table under `_child_prov_lock` and finds the address already gone, or is about a different chip entirely. This is necessary but, on this base, not sufficient: `_control_reservation` (hw-native-sys#1541) holds `_submit_mu` across the same native call at the same granularity, so an 8-chip upload still serializes there. Measured on 8 x 910B2 with this commit alone, uploading rank-stacked weights one thread per shard: threaded 6.5-8.0 GB/s vs serial 8.6-10.3 GB/s — no gain. The following commit addresses that lock; with both, the same harness reaches 19.2-34.0 GB/s threaded.
A control command that belongs to no run takes `_submit_mu` through `_control_reservation` and holds it across the native call, so on this base it — not `_child_prov_lock` — is what serializes an 8-chip upload: all 8 threads enter `Orchestrator.copy_to` within 14 us of each other, then each child starts its copy within ~1 ms of the previous one finishing, at full link speed. What such a command needs is "no run may be admitted while I run", which is a property of the worker; two commands on different chips can both have that at the same time. Make `_submit_mu` a shared/exclusive lock: run admission takes it exclusively, control takes it shared. Writer-preferring, so control traffic cannot starve a submit. The reservation's re-entrancy is untouched: it short-circuits on the thread-local set before reaching the lock. Measured on 8 x 910B2, one upload thread per shard, fresh shared-memory bands so every page is read cold exactly once: provenance locks only threaded 6.5-8.0 GB/s serial 8.6-10.3 + this commit threaded 19.2-34.0 GB/s serial 9.4-10.5 Draft: this is a change to the serializer hw-native-sys#1541 introduced, and the reviewer of that area may prefer a different shape. Sent as a separate commit so it can be dropped or replaced without touching the provenance-lock fix.
2 participants
Summary
Validation