Refactor: declare K=1 runtime pipeline contracts#1463
Conversation
📝 WalkthroughWalkthroughThe change replaces drain-based execution with explicit run lifecycle APIs, asynchronous Python ChangesRun-scoped execution lifecycle
Runtime pipeline contracts
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant Worker
participant RunHandle
participant Orchestrator
Caller->>Worker: submit(callable)
Worker->>Orchestrator: begin_run and build DAG
Worker-->>Caller: return RunHandle
Worker->>Orchestrator: close_run_submission(run_id)
Caller->>RunHandle: wait(timeout)
RunHandle->>Orchestrator: wait_run(run_id)
RunHandle->>Worker: perform cleanup
RunHandle->>Orchestrator: release_run(run_id)
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 |
4561329 to
9d1850f
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (3)
python/bindings/worker_bind.h (1)
371-371: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConsider releasing the GIL for
_release_run.
Orchestrator::release_runcan acquire the scheduler loop mutex and runallocator_->reset_to_empty()(src/common/hierarchical/orchestrator.cpplines 177-183). That blocks every other Python thread for the duration since the native scheduler never needs the GIL, there is no deadlock risk in releasing it.♻️ Proposed change
- .def("_release_run", &Orchestrator::release_run, nb::arg("run_id")); + .def( + "_release_run", &Orchestrator::release_run, nb::arg("run_id"), + nb::call_guard<nb::gil_scoped_release>() + );🤖 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/bindings/worker_bind.h` at line 371, Update the `_release_run` binding for `Orchestrator::release_run` to release the Python GIL while the native operation executes, using the binding library’s established call-guard mechanism. Preserve the existing `run_id` argument and release behavior.python/simpler/orchestrator.py (1)
143-150: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueStale docstring reference to
Worker.run().The module docstring above (lines 27-32) was updated to center on
Worker.submit()/RunHandle, but this class docstring still says "DAG builder. Valid only inside the orch function passed to Worker.run()." Update to mentionWorker.submit()(withWorker.run()as the blocking compatibility wrapper) for consistency with the rest of this file's updated terminology.📝 Proposed doc fix
class Orchestrator: - """DAG builder. Valid only inside the orch function passed to Worker.run(). + """DAG builder. Valid only inside the orch function passed to Worker.submit() + (or its blocking Worker.run() wrapper).🤖 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/orchestrator.py` around lines 143 - 150, Update the Orchestrator class docstring to state that it is valid only inside the orchestration function passed to Worker.submit(), while noting Worker.run() as the blocking compatibility wrapper. Keep the existing ownership and lifetime explanation unchanged.python/simpler/worker.py (1)
5738-5790: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
callableparameter shadows the Python builtin across the new submit API and its test mock. The root cause is the new publicWorker.submit/runsignatures choosingcallableas the parameter name; the test mock mirrors it only because it monkeypatches that exact signature.
python/simpler/worker.py#L5738-L5790: rename thecallableparameter (e.g. tocallbackororch_fn) insubmit,run,_submit_locked, and_submit_l3_locked.tests/ut/py/test_worker/test_host_worker.py#L1575-L1575: updatefake_submit's parameter name to match once the production signature is renamed.🤖 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 5738 - 5790, Rename the callable parameter to a non-shadowing name consistently in Worker.submit, Worker.run, _submit_locked, and _submit_l3_locked, updating their internal references while preserving behavior; also rename the corresponding parameter in tests/ut/py/test_worker/test_host_worker.py lines 1575-1575 in fake_submit to match the production signature.Source: Linters/SAST tools
🤖 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 381-383: Revise the documentation around ring.slot_state(id) and
reset_to_empty() to state that CONSUMED ends only the slot’s
logical/task-resource lifetime; the TaskSlotState entry remains in the deque
until reset_to_empty() resets the ring. Remove the implication that slots are
individually reclaimed while preserving the existing worker-level reset
condition.
In `@src/common/hierarchical/orchestrator.cpp`:
- Around line 197-207: Make report_task_error best-effort so late or reset run
IDs cannot throw from the scheduler thread. Add a tolerant try_get_run lookup
alongside get_run, update record_run_error to use it and silently drop errors
when no run exists, while preserving first-error recording for valid runs.
- Around line 164-184: The drain decision in Orchestrator::release_run is made
after releasing runs_mu_, allowing concurrent begin_run() activity between the
quiescence check and allocator reset. Move the allocator active_count() check
and the resulting reset decision into the runs_mu_ critical section, preserving
scheduler-loop mutex protection around reset_to_empty(); alternatively enforce
and document that begin_run() cannot run concurrently, but ensure the invariant
is explicit.
In `@src/common/hierarchical/scheduler.cpp`:
- Around line 75-78: Update the failure-callback flow in the scheduler around
is_failure, cfg_.on_task_failed_cb, and the aggregation path near lines 153-155
so each failed task or grouped failure is reported only once. Track whether the
task/group failure has already been reported across incoming completions,
suppressing the later aggregated callback while preserving reporting for
genuinely new failures.
In `@src/common/worker/chip_worker.cpp`:
- Around line 163-174: In the initialization flow, defer assigning
resolved_contract to pipeline_contract_ until create_device_context and
simpler_init complete successfully. Update the rollback paths in init() to avoid
exposing metadata from a failed runtime, while preserving contract validation
and handle cleanup.
In `@tests/ut/py/test_worker/test_startup_readiness.py`:
- Line 768: Rename the local variable callable in the ChipCallable.build test
setup to chip_callable, and update any references to it within the test.
---
Nitpick comments:
In `@python/bindings/worker_bind.h`:
- Line 371: Update the `_release_run` binding for `Orchestrator::release_run` to
release the Python GIL while the native operation executes, using the binding
library’s established call-guard mechanism. Preserve the existing `run_id`
argument and release behavior.
In `@python/simpler/orchestrator.py`:
- Around line 143-150: Update the Orchestrator class docstring to state that it
is valid only inside the orchestration function passed to Worker.submit(), while
noting Worker.run() as the blocking compatibility wrapper. Keep the existing
ownership and lifetime explanation unchanged.
In `@python/simpler/worker.py`:
- Around line 5738-5790: Rename the callable parameter to a non-shadowing name
consistently in Worker.submit, Worker.run, _submit_locked, and
_submit_l3_locked, updating their internal references while preserving behavior;
also rename the corresponding parameter in
tests/ut/py/test_worker/test_host_worker.py lines 1575-1575 in fake_submit to
match the production signature.
🪄 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: fc67cfa5-37eb-4853-a8ca-05ad6fe1cd2f
📒 Files selected for processing (37)
docs/comm-domain.mddocs/hierarchical_level_runtime.mddocs/orchestrator.mddocs/remote-l3-worker-design.mddocs/scheduler.mddocs/task-flow.mdpython/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/common/hierarchical/orchestrator.cppsrc/common/hierarchical/orchestrator.hsrc/common/hierarchical/ring.cppsrc/common/hierarchical/ring.hsrc/common/hierarchical/scheduler.cppsrc/common/hierarchical/scheduler.hsrc/common/hierarchical/tensormap.cppsrc/common/hierarchical/tensormap.hsrc/common/hierarchical/types.cppsrc/common/hierarchical/types.hsrc/common/hierarchical/worker.cppsrc/common/hierarchical/worker_manager.cppsrc/common/hierarchical/worker_manager.hsrc/common/worker/chip_worker.cppsrc/common/worker/chip_worker.hsrc/common/worker/pto_runtime_c_api.htests/ut/cpp/hierarchical/test_orchestrator.cpptests/ut/cpp/hierarchical/test_scheduler.cpptests/ut/cpp/hierarchical/test_tensormap.cpptests/ut/py/test_worker/test_error_propagation.pytests/ut/py/test_worker/test_host_worker.pytests/ut/py/test_worker/test_l3_l2_orch_comm.pytests/ut/py/test_worker/test_startup_readiness.py
9d1850f to
a9ec299
Compare
ChaoWao
changed the title
52e266c to
b3e89f4
Compare
Define the versioned PipelineContract C ABI and the HOST_PER_RUN / DEVICE_SCRATCH / EXEC_HANDLE resource classes, and let the A2/A3 host_build_graph and tensormap_and_ringbuffer runtimes declare the resources that cross KernelLaunch. A resource's copy count follows from its class: HOST_PER_RUN and EXEC_HANDLE need one instance per in-flight run, DEVICE_SCRATCH needs exactly one because it is not rewritten per run and device ops run one at a time. That leaves pipeline_depth as the only replication count, so a runtime that wants a small resource replicated and a large one shared says so by classifying them rather than by carrying a second count. host_build_graph materializes this run's own graph into the image it uploads, so every device region it declares is HOST_PER_RUN. tensormap_and_ringbuffer orchestrates on the device, so only its task args carry per-run content; its heap, SM and runtime image are built and uploaded once per callable sizing and then served from the prebuilt-arena cache, so they are DEVICE_SCRATCH. Zero is reserved as an unspecified kind. An entry a runtime never filled in stays zeroed, so a declaration that overstates resource_count is rejected whatever its filled entries are, and a kind may repeat for a runtime that needs several resources of one type. ChipWorker loads the declaration as optional runtime metadata and rejects one it cannot honor: wrong ABI version, too many resources, a depth other than 1, an unspecified or out-of-range kind, an out-of-range class, or a non-zero reserved bytes_per_copy. A runtime that exports no contract keeps the default depth-1 shape. The contract is published only once the runtime is up, so a failed init never reports the shape of a runtime this worker is not bound to. One runtime buffer and one pipeline slot are kept throughout. A depth above 1, slot selection, asynchronous acceptance and overlap remain later changes; the existing synchronous simpler_run -> launch_run -> reap_run path is unchanged.
Revision: kind sentinel replaces the repeated-kind rule;
|
2 participants
Summary
PipelineContractC ABI and theHOST_PER_RUN/DEVICE_SCRATCH/EXEC_HANDLEresource classeshost_build_graphandtensormap_and_ringbufferdeclare the resources that crossKernelLaunchpipeline_depthas the only replication count; a resource's copy count follows from its classresource_countis rejected whatever its filled entries arepipeline_contract.hbeside the contract, with unit testsinit()cannot report another runtime's shapesimpler_run -> launch_run -> reap_runpathWhy one replication count
HOST_PER_RUNandEXEC_HANDLEneed one instance per in-flight run;DEVICE_SCRATCHneeds exactly one, because it is not rewritten per run anddevice ops run one at a time. So the copy count is derivable from the class, and
a runtime that wants a small resource replicated while a large one stays shared
expresses that by classifying them — not by carrying a second global count whose
binding to particular resource kinds would live outside the contract.
That is exactly the difference between the two runtimes:
HOST_PER_RUNDEVICE_SCRATCHEXEC_HANDLEhost_build_graphGM_HEAP,GM_SM,RUNTIME_IMAGEtensormap_and_ringbufferTASK_ARGSGM_HEAP,GM_SM,RUNTIME_IMAGEhost_build_graphmaterializes this run's own graph into the image it uploads,so every device-resident region it declares carries per-run content.
tensormap_and_ringbufferorchestrates on the device, so only its task argscarry per-run content.
The class turns on rewritten per run or not, not on who writes the bytes. Its
heap, SM and runtime image are host-built and uploaded by
build_and_cache_prebuilt_arena, but only on a cache miss or eager prewarm —once per
(callable, sizing). A repeat run of the same callable takesbind_cached_runtime_image, which performs nocopy_to_device. One instance istherefore correct, and the header comment states that rather than the earlier
claim that the device builds these regions.
Admission
is_valid_depth1_pipeline_contractrejects a wrong ABI version, more resourcesthan fit, a depth other than 1, an unspecified or out-of-range kind, an
out-of-range class, and a non-zero reserved
bytes_per_copy. An entry a runtimenever filled in stays zeroed, so the unspecified-kind rule is what catches a
resource_countlarger than the entries actually filled in. A runtime thatexports no contract keeps the default depth-1 shape, so the declaration is
additive; a runtime that ships a contract this build cannot honor fails
init().Repeated kinds are legal: a kind names a resource type, not one instance of it,
so a runtime that needs two of a kind can say so.
Scope
One runtime buffer and one pipeline slot throughout. A depth above 1, slot
selection, asynchronous acceptance and overlap remain later changes. Nothing
consumes the declaration yet.
Series context
Part of the worker async-preparation series, following #1459 (per-run identity
and completion fence), #1460 (
Worker.submitand run handles) and #1462(launch/reap split), all merged. This declares what each runtime owns across the
KernelLaunchboundary that #1462 established.Validation
Rebased onto merged
main; the three superseded ancestor commits are dropped,so this is a single commit.
PipelineContractcasesbytes_per_copyrule or the unspecified-kind rule fails exactly the casescovering it, the latter failing both
RejectsAnUnspecifiedKindandRejectsAResourceCountPastTheFilledEntriesRejectsAResourceCountPastTheFilledEntriesloops over two different filledkinds, so it cannot pass by coincidence of which kind the declaration uses
libhost_runtime.sorebuilt for both runtimes, both exportingget_pipeline_contractsimulation —
vector_example(host_build_graph) andmixed_example(
tensormap_and_ringbuffer) each pass, which requiresChipWorker::inittoload and accept the runtime's contract, including the enlarged six-resource
tensormap_and_ringbufferdeclarationThe hardware runs quoted in earlier revisions were against a commit no longer in
this branch's history and have not been re-run; the onboard CI job is the
meaningful check.