Add: kernel-mode entry on Worker(level=2) and a pre-init capability probe

[sunkaixuan2018]

Targets feat/kernel-mode-integration-test.

What this is

PyPTO drives simpler through simpler.worker.Worker, but the kernel-mode entries on this branch exist only on ChipWorker. YunjiQin raised this on #2185 (the PyPTO integration target is Worker). This PR adds a kernel mode to Worker(level=2), so a framework process that already owns a device and a stream can use kernel mode through the public Worker.

worker = Worker(level=2, execution_mode="kernel", device_id=0,
                platform="a2a3", runtime="tensormap_and_ringbuffer")
assert worker.kernel_mode_supported          # answered before init, no device taken
worker.init(config=CallConfig())             # device 0 is already current on this thread
callable_id = worker.kernel_prepare_callable(chip_callable)
worker.kernel_launch(callable_id, args, caller_stream=stream)
...                                          # the caller synchronizes its own stream
worker.close()

API

Entry	Behaviour
Worker(level=2, execution_mode="kernel", ...)	Fixes the mode at construction. "program" is the default. A level other than 2 raises ValueError, and so does enable_sdma.
init(*, config)	Routes to ChipWorker.kernel_init on the device the calling thread already has current. config is required and keyword-only. prewarm_config is refused. Nothing is prewarmed or replayed.
kernel_mode_supported	Answers before init, through a new native probe (below). Cached. Returns False for a level other than 2. A missing runtime build raises.
kernel_prepare_callable(chip_callable) -> int	Returns the id the runtime minted. There is no dedup and no handle, and register() is not involved. Images stay referenced until a native teardown succeeds.
kernel_launch(callable_id, args, *, caller_stream) -> None	Enqueues and returns. There is no lease, no RunHandle and no wait. args is ChipStorageTaskArgs. An id this Worker did not mint, a non-integer id, or a zero stream is refused before the native call.
close()	Releases only the context: no stream sync, no device reset, no ACL finalize. A teardown failure raises, and a later close() retries.

Mode exclusion.

• A kernel-mode Worker refuses register/unregister/submit/run, malloc/free/copy_to/copy_from, create_buffer/make_tensor_arg/release_buffer and device_memory_info. The native memory query already refuses a kernel context.
• committed_device_memory stays available.
• A program-mode Worker refuses the kernel entries.

Decisions taken for this PR

1. Worker layer only, on this branch's ChipWorker. The Add: kernel-mode entry layer — give the four C entries callers (⑩a) #2185-head ChipWorker fixes (typed errors, teardown owed after a failed init) are not ported here. Where this base's ChipWorker falls short, the Worker compensates. ChipWorker.finalize swallows a device-teardown failure, so kernel close() checks initialized afterwards and raises to keep the journal entry. The Worker also keeps its own reference to every prepared image.

2. kernel_launch takes ChipStorageTaskArgs. The borrowed-tensor argument type is a follow-up; its public name is still undecided.

3. A real pre-init probe. This is the only C++ change. The new static ChipWorker::probe_kernel_mode_supported(host_lib_path, sim_context_path):

   • loads the host runtime;
   • asks simpler_kernel_mode_supported on a fresh device context that no init touches;
   • destroys that context before the handle is released.

   The C ABI already requires that entry to answer from the runtime build alone, so the probe takes no device. It is bound on _ChipWorker and wrapped as ChipWorker.probe_kernel_mode_supported(bins).

4. An unclosed kernel Worker leaks rather than finalizes.

   • The problem: ~ChipWorker finalizes on whatever thread destroys it, without the quiescence kernel teardown requires.
   • The behaviour: a kernel-mode Worker that is garbage-collected, or still alive at interpreter exit, emits a ResourceWarning and deliberately pins its ChipWorker for the rest of the process. The finalizer never references the Worker.
   • A side effect: nanobind reports that pinned instance as leaked at exit.

Concurrency, teardown, fork

• One lock linearizes prepare and launch against the native finalize.
  • A launch that overlaps a prepare, another launch or close() fails fast. The caller is expected to serialize these calls.
  • close() publishes CLOSED before it takes the lock, so a launch either sees CLOSED or finishes before finalize.
  • If a prepare or launch still holds the lock past the rollback grace period, close() raises TimeoutError and can be called again.
• Failed kernel_init rolls back to FAILED through the existing path. Nothing is pinned.
• Fork: the process that bound the context is recorded. A forked child that inherited the Worker cannot prepare, launch or tear the context down.

Validation

CI does not run for PRs into this branch (ci.yml triggers on PRs to main), so everything below was run by hand on a2a3 silicon. Hardware runs were wrapped in task-submit.

Lane	Result
cpput, no hardware	168/168 passed
new UT test_worker_kernel_mode.py + test_chip_worker.py + test_startup_readiness.py	176/176 passed, no nanobind leak report
pyut, no hardware (full)	8 failed / 2372 passed / 11 errors. All 19 are the import torch set this box always hits (no torch installed); the new file adds none.
cpput, a2a3 hardware	2/2 passed, incl. test_kernel_mode_entry
pyut, a2a3 hardware	28/28 passed: the 22 kernel-mode and onboard cases the branch already had, plus the 6 new test_worker_kernel_mode_hw.py cases (H0 probe before init, H1 eager numerics, H2 prepare twice, H3 refusals then a valid launch, H4 second Worker on the same device, H5 HBG refused)
pre-commit	all 16 hooks pass on every changed file, no rewrites

Baseline on the untouched branch (4a5f28c9f), same box:

• cpput no-hardware: 168/168.
• cpput hardware: 2/2.
• pyut hardware: 21/21.
• The no-hardware failures are all import torch (this box has no torch).

Out of scope

• The borrowed tensor/scalar argument type, and the GIL-free native launch descriptor. Both still have open design questions.
• Kernel mode at L3+.
• The Add: kernel-mode entry layer — give the four C entries callers (⑩a) #2185-head ChipWorker fixes on this branch: typed error codes, finalize_device after a failed kernel_init, and teardown-failure propagation.

Known issues found on the way (not fixed here)

• ChipWorker::kernel_init destroys the context without finalize_device after a nonzero init rc.
• ChipWorker::finalize swallows a device-teardown failure, in program mode too. The Python ChipWorker.finalize wrapper then clears its registries.
• ~ChipWorker finalizes on the destroying thread.
• The chip_worker.h doc for kernel_prepare_callable still describes a caller_stream parameter.

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

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: Advanced

Run ID: 2efcea1a-78c0-4ec0-adbc-311d116c1d2e

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

---

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

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[YunjiQin]

Pushed a rebase onto b5a0ea0c9 (the head carrying the merged #2171 and #2172) plus one follow-up commit, using maintainer-edit. Summary of what changed and why, so you can check the resolutions.

Rebase: three conflicts, all resolved by keeping both sides

They all come from #2247 landing after this branch's base.

• src/common/worker/chip_worker.h — probe_kernel_mode_supported joins Add: adopt #2185's kernel entry hardening on the integration line #2247's finalize() doc, which now describes the throwing kernel teardown.
• tests/ut/py/test_chip_worker.py — TestChipWorkerKernelProbe joins TestChipWorkerKernelEntryLayer; the fake runtime keeps Add: adopt #2185's kernel entry hardening on the integration line #2247's live_handles tracking and finalize_failures / finalize_aborts alongside this branch's created_contexts counter.
• python/simpler/task_interface.py — the kernel_prepare_callable docstring takes this branch's account of the registration synchronize plus Add: adopt #2185's kernel entry hardening on the integration line #2247's sentence on how long the image stays referenced.

Follow-up commit: one compensation #2247 made unreachable

close() carried this:

self._chip_worker.finalize()
# ChipWorker.finalize returns without raising when the native device teardown
# fails, and the native worker then stays initialized. ...
if kernel and impl is not None and bool(getattr(impl, "initialized", False)):
    raise RuntimeError("kernel context teardown failed; ...")

After #2247 that comment is false and the branch is dead. ChipWorker::finalize now sets initialized_ = false and throws ChipWorkerError for a kernel context (chip_worker.cpp:676-682), so the raise is already what keeps the journal entry, the ChipWorker, its images and the GC pin for a retry — and initialized reads False on both the success and the failure path, so the check can never fire.

Removed it. The fake ChipWorker.finalize in test_worker_kernel_mode.py modelled the old behaviour (failed finalize returns, native worker stays initialized), so leaving it would have kept test_failed_teardown_keeps_the_context_for_retry asserting a contract the real one no longer has. The fake now clears initialized and raises ChipWorkerError, and the test expects that type with the status the entry reports.

Also now stale, left for you

Three of the four entries under "Known issues found on the way" are fixed on this head by #2247:

• kernel_init destroying the context without finalize_device after a nonzero init rc
• ChipWorker::finalize swallowing a device-teardown failure (kernel contexts; program mode still reports on stderr by D20)
• the chip_worker.h doc for kernel_prepare_callable still naming a caller_stream parameter

The fourth — ~ChipWorker finalizing on the destroying thread — still holds, so the GC pin's rationale is intact. #2247 only wrapped that destructor in a try/catch.

Verified on the rebased branch

• tests/ut -m "not requires_hardware" — 2511 passed (the 2448 branch baseline plus this PR's 63)
• tests/ut -m requires_hardware --platform a2a3 — 28/28 scheduled cases passed, matching this PR's own count
• pre-commit on every changed file — passed

Hardware ran through task-submit --device auto after the architecture precheck.

One note on the review thread that led here: I said earlier the hardware cases did not cover close invalidating the minted ids. That was wrong — I had read the summary rather than the code. _case_eager_values already asserts kernel_launch and kernel_prepare_callable are both refused after close(), and _case_prepare_twice covers two ids from one callable. Nothing to add there.

#2249 is the capture scene-test fix and is disjoint from this: it changes no Worker or ChipWorker code, and this changes no scene test, so they can land in either order.