Add: bounded asynchronous native run lane - #1650
Conversation
📝 WalkthroughWalkthroughThe PR adds bounded concurrent native-run preparation for eligible backends, asynchronous L2 submission with ChangesConcurrent native-run pipeline
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant Worker
participant ChipWorker
participant Runtime
participant RunHandle
Client->>Worker: submit L2 task
Worker->>ChipWorker: prepare native run
ChipWorker->>Runtime: reserve and prepare run identity
Worker->>ChipWorker: launch FIFO front run
ChipWorker->>Runtime: launch native run
Worker-->>Client: return RunHandle
Client->>RunHandle: wait
RunHandle->>Worker: poll and finalize
Worker->>ChipWorker: finalize native run
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 |
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (8)
src/common/platform/onboard/host/device_runner_base.cpp (1)
1356-1359: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueGuard the division by
cores_per_blockdim_.
block_dim_ = worker_count_ / cores_per_blockdim_introduces a division on a field that the previous code only multiplied by. If any arch leavescores_per_blockdim_at 0, this is undefined behaviour instead of a benign zero.🛡️ Proposed guard
void DeviceRunnerBase::activate_launch_shape(const Runtime &runtime) { worker_count_ = runtime.get_worker_count(); - block_dim_ = worker_count_ / cores_per_blockdim_; + block_dim_ = cores_per_blockdim_ > 0 ? worker_count_ / cores_per_blockdim_ : 0; }🤖 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 1356 - 1359, Update DeviceRunnerBase::activate_launch_shape to guard cores_per_blockdim_ before dividing worker_count_. Preserve the benign zero behavior by setting block_dim_ to zero when cores_per_blockdim_ is zero; otherwise retain the existing worker-count division.src/common/platform/onboard/host/c_api_shared.cpp (2)
711-713: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExplain why
runner_resources_ownedis set before provisioning.Line 711 sets the flag to
truebeforeprovision_native_run_resourcesruns at Line 712. On a provisioning failure,cleanup_failed_preparetherefore callsabandon_native_run_resourcesfor a slot that provisioning did not complete.That is the right choice for a partial provision, but the pre-set reads as a sequencing mistake. Add a one-line comment stating that a failed provision may still hold partial resources, so ownership is claimed before the attempt.
🤖 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/c_api_shared.cpp` around lines 711 - 713, In the preparation flow around runner_resources_owned and provision_native_run_resources, add a concise one-line comment explaining that provisioning failure may leave partial native resources held, so ownership must be claimed before the provisioning attempt for cleanup_failed_prepare to release them.
697-701: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueA rejected reservation drops its allocated trace invocation.
Line 676 allocates
trace_invand Line 677 stampstrace_start_ns. Whentry_reserve_native_runfails, Line 700 returns without callingemit_native_run_host_wall. Every other failure path routes throughcleanup_failed_prepare, which emits the span. Admission rejections therefore leave a gap in the host trace exactly where contention analysis needs a record.🤖 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/c_api_shared.cpp` around lines 697 - 701, Update the admission-rejection branch in the native-run preparation flow around try_reserve_native_run so it emits the allocated trace invocation via the existing cleanup_failed_prepare path or equivalent before destroying state and returning. Preserve the current error logging and -1 return while ensuring rejected reservations produce the same host-wall trace record as other failure paths.src/common/platform/onboard/host/device_runner_base.h (1)
132-135: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
pipeline_slot(),selected_arena_bank(), andarena_bank()can now throw.These accessors were plain field reads. They now route through the pthread TLS helper, which throws on key-creation or allocation failure.
arena_bank()at Line 1021 is on thesetup_static_arenaandacquire_pooled_*paths, andpipeline_slot()is read at the top ofDeviceRunner::run. Callers that previously treated these as infallible now have a new exception edge.Document the new throwing contract on the declarations so callers on
noexceptpaths do not adopt them by accident.Also applies to: 148-148, 1021-1021
🤖 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 132 - 135, Document on the declarations of pipeline_slot(), selected_arena_bank(), and arena_bank() that each accessor may throw due to pthread TLS key creation or allocation failure. Keep the existing signatures unchanged and make the contract visible to callers before they use these accessors in run, setup_static_arena, or acquire_pooled_* paths.src/common/worker/chip_worker.cpp (1)
684-688: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSimplify the occupancy check.
occupied != 0 && occupied != 1isoccupied > 1. The direct form states the intent: at most one predecessor may exist.♻️ Proposed simplification
- if (occupied != 0 && occupied != 1) { + if (occupied > 1) {🤖 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/worker/chip_worker.cpp` around lines 684 - 688, In the occupancy validation near the prepare_native_run ownership check, replace the `occupied != 0 && occupied != 1` condition with the equivalent `occupied > 1` check, preserving the existing runtime error and identity formatting.src/common/hierarchical/worker_manager.cpp (1)
318-325: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMake the non-queued mapping explicit in
dispatch.
dispatchpasses nostaged_run_id, soenqueue_dispatchcannot returnSTAGED_IDENTITY_CHANGEDtoday. Theelsebranch nevertheless reports any future non-STOPPINGresult as"endpoint capacity exceeded".dispatch_preparedalready switches on each enumerator. Match that shape so a new enumerator cannot be reported under the wrong message.♻️ Proposed change
if (result == EnqueueDispatchResult::STOPPING) { complete_unpublished(d, "WorkerThread::dispatch: worker is stopping"); - } else { + } else if (result == EnqueueDispatchResult::CAPACITY_EXCEEDED) { complete_unpublished(d, "WorkerThread::dispatch: endpoint capacity exceeded"); + } else { + complete_unpublished(d, "WorkerThread::dispatch: enqueue rejected the dispatch"); }🤖 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 318 - 325, Update dispatch to explicitly handle every non-queued EnqueueDispatchResult, matching dispatch_prepared’s switch structure. Preserve the STOPPING message, map the currently impossible STAGED_IDENTITY_CHANGED result explicitly, and ensure any newly added enumerator cannot fall through to “endpoint capacity exceeded.”python/simpler/worker.py (1)
2405-2414: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winShare one diagnostics predicate.
config_has_diagnosticshere andWorker._l2_config_has_diagnosticsat lines 8194-8202 contain the identical field list. Both must trackCallConfig::diagnostics_any(). Extract one module-level helper and call it from both sites, so a new diagnostic flag cannot be added to only one copy.♻️ Proposed shared helper
def _call_config_has_diagnostics(config: CallConfig) -> bool: # Mirrors CallConfig::diagnostics_any(). return bool( config.enable_l2_swimlane or config.enable_dump_args or config.enable_pmu or config.enable_dep_gen or config.enable_scope_stats )- def config_has_diagnostics(config: CallConfig) -> bool: - # Mirrors CallConfig::diagnostics_any(); these modes share native - # diagnostic state and therefore use the serial prepare fallback. - return bool( - config.enable_l2_swimlane - or config.enable_dump_args - or config.enable_pmu - or config.enable_dep_gen - or config.enable_scope_stats - ) + config_has_diagnostics = _call_config_has_diagnostics🤖 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` around lines 2405 - 2414, Extract the shared field-list logic from the local config_has_diagnostics function into a module-level _call_config_has_diagnostics(config: CallConfig) helper, preserving the fields that mirror CallConfig::diagnostics_any(). Update both config_has_diagnostics and Worker._l2_config_has_diagnostics to delegate to this helper so future diagnostic flags are maintained in one place.tests/ut/py/test_worker/test_host_worker.py (1)
405-432: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLet the harness vary the frame run id.
publishwrites a constantrun_idof 5 into every frame. The new concurrency tests therefore stage an active frame and a successor frame that share one run id and differ only bydispatch_id. Production stages a successor that belongs to a different run. The tests still prove dispatch-id ordering, so nothing is wrong today. Ifrun_two_frame_looplater gates preparation on run identity, these tests would pass without exercising that gate. Add arun_idparameter with a default of 5 and give the successor a distinct value in the concurrent-prepare tests.🤖 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 405 - 432, Update the test harness publish method to accept a run_id parameter defaulting to 5, and write that value into _OFF_FRAME_RUN_ID instead of the hardcoded constant. In the concurrent-prepare tests, pass a distinct run_id for the successor frame while preserving the existing default for other callers.
🤖 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/worker-manager.md`:
- Around line 241-247: Update the registry-control deferral sentence in the
worker-manager documentation to wait for both the active native run and any
prepared successor holding a native token. Keep the rule aligned with
run_two_frame_loop, where controls remain deferred while any staged frame owns a
native token.
In `@python/simpler/worker.py`:
- Around line 8570-8583: Update _wait_run_handle_accepted so it advances the L2
FIFO in bounded steps, rechecking the target run state after each
_l2_progress_locked call rather than allowing one call to run through terminal
completion. Return as soon as the target phase is no longer "prepared", while
preserving unknown-run and propagated-error handling; leave the orchestrator
path unchanged.
In `@src/common/platform/onboard/host/c_api_shared.cpp`:
- Around line 873-881: Update simpler_finalize_run to wrap
capture_native_run_thread_selection, select_pipeline_slot, and select_arena_bank
in exception handling that releases runner_claimed and runner_reserved, calls
destroy_native_run_state, and returns -1 on selection failure or thrown
exceptions. Apply the equivalent exception guard to simpler_launch_run around
its native thread selection calls, preserving its existing claim-release
behavior.
In `@src/common/platform/onboard/host/device_runner_base.cpp`:
- Around line 1659-1665: Protect the native_launch_signal_ read in
DeviceRunnerBase::publish_task_accepted with native_run_mu_, matching the
synchronization used by try_acquire_native_run and release_native_run. Hold the
mutex while copying or checking the pointer and notifying it, so the pointer
cannot race with updates or destruction; preserve the existing accepted_state
release store.
- Around line 72-107: Update the native-run selection flow around
run_selection(), select_pipeline_slot(), and select_arena_bank() so selection
state is isolated per DeviceRunnerBase instance rather than shared by thread
alone. Attach or copy NativeRunThreadSelection to the owning runner, or make the
existing TLS state resolve through a runner-specific key, ensuring contexts used
on the same host thread cannot reuse each other’s arena bank, retained buffers,
or run identity.
In `@src/common/platform/onboard/host/device_runner_base.h`:
- Around line 136-140: Remove noexcept from restore_native_run_thread_selection
in both its declaration and definition, allowing exceptions from run_selection
during create_thread’s initial TLS allocation to propagate safely instead of
terminating the process.
In `@src/common/worker/chip_worker.cpp`:
- Around line 814-826: Re-validate both lease_generation and run_epoch under
native_run_mu_ before writing completion state in poll_native_run and
wait_native_run. If the slot identity no longer matches, do not update the
successor state; otherwise preserve the existing writes to phase REAPED and, in
wait_native_run, wait_rc.
In
`@tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py`:
- Around line 196-199: Remove the timing-dependent torch.count_nonzero assertion
from the direct L2 submit test, while retaining the run_handle._terminal
assertion to verify submit returns a non-completed compatibility handle. Keep
the subsequent run_handle.wait(30.0) lifecycle check unchanged.
---
Nitpick comments:
In `@python/simpler/worker.py`:
- Around line 2405-2414: Extract the shared field-list logic from the local
config_has_diagnostics function into a module-level
_call_config_has_diagnostics(config: CallConfig) helper, preserving the fields
that mirror CallConfig::diagnostics_any(). Update both config_has_diagnostics
and Worker._l2_config_has_diagnostics to delegate to this helper so future
diagnostic flags are maintained in one place.
In `@src/common/hierarchical/worker_manager.cpp`:
- Around line 318-325: Update dispatch to explicitly handle every non-queued
EnqueueDispatchResult, matching dispatch_prepared’s switch structure. Preserve
the STOPPING message, map the currently impossible STAGED_IDENTITY_CHANGED
result explicitly, and ensure any newly added enumerator cannot fall through to
“endpoint capacity exceeded.”
In `@src/common/platform/onboard/host/c_api_shared.cpp`:
- Around line 711-713: In the preparation flow around runner_resources_owned and
provision_native_run_resources, add a concise one-line comment explaining that
provisioning failure may leave partial native resources held, so ownership must
be claimed before the provisioning attempt for cleanup_failed_prepare to release
them.
- Around line 697-701: Update the admission-rejection branch in the native-run
preparation flow around try_reserve_native_run so it emits the allocated trace
invocation via the existing cleanup_failed_prepare path or equivalent before
destroying state and returning. Preserve the current error logging and -1 return
while ensuring rejected reservations produce the same host-wall trace record as
other failure paths.
In `@src/common/platform/onboard/host/device_runner_base.cpp`:
- Around line 1356-1359: Update DeviceRunnerBase::activate_launch_shape to guard
cores_per_blockdim_ before dividing worker_count_. Preserve the benign zero
behavior by setting block_dim_ to zero when cores_per_blockdim_ is zero;
otherwise retain the existing worker-count division.
In `@src/common/platform/onboard/host/device_runner_base.h`:
- Around line 132-135: Document on the declarations of pipeline_slot(),
selected_arena_bank(), and arena_bank() that each accessor may throw due to
pthread TLS key creation or allocation failure. Keep the existing signatures
unchanged and make the contract visible to callers before they use these
accessors in run, setup_static_arena, or acquire_pooled_* paths.
In `@src/common/worker/chip_worker.cpp`:
- Around line 684-688: In the occupancy validation near the prepare_native_run
ownership check, replace the `occupied != 0 && occupied != 1` condition with the
equivalent `occupied > 1` check, preserving the existing runtime error and
identity formatting.
In `@tests/ut/py/test_worker/test_host_worker.py`:
- Around line 405-432: Update the test harness publish method to accept a run_id
parameter defaulting to 5, and write that value into _OFF_FRAME_RUN_ID instead
of the hardcoded constant. In the concurrent-prepare tests, pass a distinct
run_id for the successor frame while preserving the existing default for other
callers.
🪄 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: 6baf5a84-ef87-4133-9053-7b0416dc8098
📒 Files selected for processing (24)
docs/task-flow.mddocs/worker-manager.mdpython/bindings/task_interface.cpppython/simpler/task_interface.pypython/simpler/worker.pysrc/a2a3/platform/onboard/host/device_runner.cppsrc/a2a3/platform/onboard/host/device_runner.hsrc/common/hierarchical/scheduler.cppsrc/common/hierarchical/worker_manager.cppsrc/common/hierarchical/worker_manager.hsrc/common/log/include/common/strace.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/worker/chip_worker.cppsrc/common/worker/chip_worker.hsrc/common/worker/native_run_state.hsrc/common/worker/pto_runtime_c_api.htests/st/a2a3/host_build_graph/native_run_lifecycle/kernels/orchestration/long_vector_orch.cpptests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.pytests/ut/cpp/hierarchical/test_run_stream_slots.cpptests/ut/cpp/hierarchical/test_scheduler.cpptests/ut/py/test_worker/test_host_worker.pytests/ut/py/test_worker/test_startup_readiness.py
b160bcf to
6526943
Compare
|
Reviewed against 1. This reverts D1, which merged 4 hours ago in #1587The title says "align native prepared lane with v2 ownership", but the diff also undoes the uniform host-runtime pipeline ABI:
The commit message justifies this as "keep pipeline metadata optional for older runtimes". There are no older runtimes: If the v2 ownership rework genuinely needs a symbol to be optional, please name which one and why, rather than reverting the set. 2. Deleting the sim
|
774646d to
bc204de
Compare
bc204de to
8a70285
Compare
|
@ChaoWao Addressed the review summary:
Final validation includes all pre-commit hooks, 1093 Python unit tests, 78 C++ |
8a70285 to
e84ad57
Compare
- Return a live RunHandle from direct L2 submit and bound admission to one active run plus one eligible prepared successor. - Carry immutable per-run identity and resource selection through the required host-runtime ABI with context-bound HostApi callbacks, replacing thread-local runner and resource selection state. - Publish acceptance from the launch signal while preserving FIFO launch, serialized opaque runtime phases, and deterministic cleanup. - Keep simpler_init flat and carry run context in required per-run calls. - Keep A2/A3 HBG depth-two preparation capability-gated and A5 HBG depth-one until it owns per-slot execution resources. - Document the contracts and add isolation, acceptance, concurrency, and lifecycle regressions.
e84ad57 to
bfc79c5
Compare
|
Splitting out PR-A of this change as an independent, narrower PR: #1685 ("Refactor: bind platform host callbacks per run instead of per thread"). #1685 lands only the platform-layer prerequisite this PR bundles first:
It deliberately excludes the other two concerns still on this branch:
Implemented independently from this branch (not pushed to it). Once #1685 merges, the remaining two can rebase onto it — |
Split the platform HostApi into an immutable HostApiFunctions table plus a per-run HostApi value object that binds it to one runner and one run's pipeline slot / arena bank. Carry the run's resource selection and trace identity across the C ABI in a required NativeRunDescriptor, and delete the pthread-TLS resource-selection mechanism (NativeRunThreadSelection, both pthread keys, capture/restore) and the four setter exports it made necessary. Launch acceptance moves from a TLS setter into the per-run launch signal. Bump PTO_PIPELINE_CONTRACT_ABI_VERSION 1 -> 2 so a stale host_runtime.so (whose simpler_run/prepare_run still take the old flat argument list) is rejected at load rather than crashing at the first call. PR-A of the hw-native-sys#1650 three-way split; PR-B (direct-L2 async lane) and PR-C (HostTensorAccessScope RAII) follow. Co-Authored-By: Claude <noreply@anthropic.com>
Split the platform HostApi into an immutable HostApiFunctions table plus a per-run HostApi value object that binds it to one runner and one run's pipeline slot / arena bank. Carry the run's resource selection and trace identity across the C ABI in a required NativeRunDescriptor, and delete the pthread-TLS resource-selection mechanism (NativeRunThreadSelection, both pthread keys, capture/restore) and the four setter exports it made necessary. Launch acceptance moves from a TLS setter into the per-run launch signal. Bump PTO_PIPELINE_CONTRACT_ABI_VERSION 1 -> 2 so a stale host_runtime.so (whose simpler_run/prepare_run still take the old flat argument list) is rejected at load rather than crashing at the first call. PR-A of the hw-native-sys#1650 three-way split; PR-B (direct-L2 async lane) and PR-C (HostTensorAccessScope RAII) follow. Co-Authored-By: Claude <noreply@anthropic.com>
Summary
RunHandlefrom direct-L2 submit while keepingrun()assubmit(...).wait(); the lane and handle own launch, completion, finalization, FIFO handoff, and cleanup.NativeRunDescriptorcopied into opaque per-run state.HostApicontaining runner, slot, bank, and its function table; remove current-runner/resource-selection TLS and the four setter exports.simpler_initkeeps its original flat binary and prewarm arguments.Semantics
DeviceRunnerowns per-slot execution resources.HostApiis compiled together with every host runtime; no optional symbol or version/size descriptor layer is introduced.Validation
HostApiand launch-signal tests passed, including two-context routing and pre-marker failure semantics.The branch contains one commit rebased onto the latest
mainat push time.