Add: eager prebuilt runtime-arena prewarm at worker init

[YunjiQin]

What

Populate the prebuilt runtime-arena cache eagerly at worker init from a
declared run config, so the first simpler_run with the same ring sizing
skips the ~800ms cold arena build.

On distributed (ep>1) runs this also removes the cross-card collective-wait
bubble: prebuilt-arena skew was converting a 1:1 dispatch into device_wall
idle. Measured on a distributed run: bind_end skew ~210ms → ~2.6ms, and
device_wall converges across cards.

How

• Extract the run-path build+upload+cache into a Runtime-free helper
  build_and_cache_prebuilt_arena(api, sizing). The lazy first-run miss path
  now builds+caches then re-binds via the cache-hit path, so run and
  prewarm share one code path (drops ensure_static_arenas' runtime arg +
  bind_launch_state).
• Add simpler_prewarm_config C-API + DeviceRunnerBase::prewarm_config. The
  runtime-specific prewarm_config_impl is a weak no-op default in the
  common onboard/sim device_runner_base, overridden by the strong trb a2a3
  impl — so hbg and archs without a prewarm impl still link (a5 untouched).
• Thread through ChipWorker (dlsym simpler_prewarm_config) and nanobind
  (_ChipWorker.prewarm_config, Worker.control_prewarm_config).
• Wire the L2/L3 control path: _CTRL_PREWARM_CONFIG mailbox command +
  chip-loop handler + worker-manager control_prewarm_config chain;
  Worker.init(prewarm_config=...) prewarms in-process (L2) or dispatches
  per chip (L3).

Scope / safety

• No-op for runtimes without a prebuilt arena (weak default returns success).
• a5 path untouched; a2a3 trb provides the only strong impl.
• Prebuilt-arena cache is single-slot per worker, so exactly one ring sizing
  is prewarmed per init.

Testing

Not yet verified on-device in this PR branch — flagging for CI / reviewer
onboard run.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

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

Run ID: 7565af93-573e-4fd1-a597-e8b537e2e1a5

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

📝 Walkthrough

Walkthrough

Adds configuration-driven runtime-arena prewarming. Python workers can accept a CallConfig, route it directly or through hierarchical mailboxes, and invoke runtime-specific APIs that build and cache the corresponding arena image.

Changes

Runtime Arena Prewarm

Layer / File(s)	Summary
Runtime arena cache builder src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp	Centralizes prebuilt arena construction and uses it for cache misses and the new prewarm_config_impl entry point.
Runtime prewarm API src/common/platform/onboard/host/*, src/common/platform/sim/host/*, src/common/worker/*	Adds runner and C APIs, weak runtime hooks, and ChipWorker lifecycle handling for configuration-only prewarming.
Hierarchical worker control src/common/hierarchical/*	Adds mailbox command CTRL_PREWARM_CONFIG and forwards packed CallConfig data to selected child workers.
Python prewarm initialization python/bindings/*, python/simpler/*	Exposes prewarm_config, accepts optional initialization configuration, and triggers direct or hierarchical prewarming during bootstrap.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PythonWorker
  participant WorkerManager
  participant ChipWorker
  participant simpler_prewarm_config
  participant prewarm_config_impl
  PythonWorker->>WorkerManager: control_prewarm_config(worker_id, CallConfig)
  WorkerManager->>ChipWorker: dispatch CTRL_PREWARM_CONFIG
  ChipWorker->>simpler_prewarm_config: prewarm ring sizing
  simpler_prewarm_config->>prewarm_config_impl: forward runtime_env ring parameters
  prewarm_config_impl-->>ChipWorker: cache prebuilt runtime image

Loading

Possibly related PRs

• hw-native-sys/simpler#1193: Introduces the prebuilt runtime-arena cache used by this prewarming path.
• hw-native-sys/simpler#1042: Adds the CallConfig.runtime_env ring-sizing data consumed by prewarming.
• hw-native-sys/simpler#1234: Modifies the same runtime-arena cache and bind behavior.

Poem

A rabbit found rings in a cache so bright,
Warmed arena paths before first flight.
Through mailboxes the configs hop,
Runtime images bloom on top.
“No cold start!” the bunny sings—
Prewarmed paths for speedy things!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 26.79% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: eager prewarming of the prebuilt runtime-arena during worker init.
Description check	✅ Passed	The description directly matches the changeset and explains the new prewarm path, control flow, and worker-init support.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

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.}

[gemini-code-assist]

Code Review

This pull request introduces an eager pre-warming mechanism (prewarm_config) to build and cache the prebuilt runtime-arena image based on the run configuration's ring sizing during initialization, thereby avoiding the cold build overhead (~800ms) on the first run. The changes span Python bindings, worker interfaces, C++ runtime makers, and hierarchical worker managers. Feedback focuses on validating the pre-warm configuration during initialization for all worker levels, adding defensive null checks and zero-initialization in control_prewarm_config to prevent undefined behavior, and using RAIIScopeGuard in the simulation implementation to ensure reliable cleanup of thread-local data.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[coderabbitai]

Actionable comments posted: 3

🤖 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 `@python/simpler/task_interface.py`:
- Around line 1164-1171: Correct the kwargs documentation in the method
associated with the shown CallConfig setup: remove the invalid ring_heap
override example and document setting config.runtime_env.ring_heap before
invoking the method, unless explicit kwargs mapping is implemented. Keep the
setattr loop consistent with only valid top-level CallConfig attributes.

In `@python/simpler/worker.py`:
- Around line 2865-2878: Ensure prewarming occurs during init() for hierarchical
workers: when prewarm_config is provided, initialize/start the hierarchy from
init() rather than deferring the existing prewarm control in
_start_hierarchical() until run(). Update the relevant
_init_hierarchical/_start_hierarchical flow so the configured arena is built
before init() returns, while preserving no-op behavior when prewarm_config is
None.

In `@src/common/hierarchical/worker_manager.cpp`:
- Around line 566-574: LocalMailboxEndpoint::control_prewarm_config currently
accepts undersized or null payloads and can leave stale configuration bytes.
Validate that config is non-null and config_size equals sizeof(CallConfig)
before acquiring the lock or writing the control command; reject invalid input
through the existing error-handling mechanism, then copy the full CallConfig
without clamping.

🪄 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

Run ID: 97d8560d-9a12-4ec3-962d-8e95b188617b

📥 Commits

Reviewing files that changed from the base of the PR and between cbbdfd4 and d7fb950.

📒 Files selected for processing (17)

• python/bindings/task_interface.cpp
• python/bindings/worker_bind.h
• python/simpler/task_interface.py
• python/simpler/worker.py
• src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp
• src/common/hierarchical/worker.h
• src/common/hierarchical/worker_manager.cpp
• src/common/hierarchical/worker_manager.h
• src/common/platform/onboard/host/c_api_shared.cpp
• src/common/platform/onboard/host/device_runner_base.cpp
• src/common/platform/onboard/host/device_runner_base.h
• src/common/platform/sim/host/c_api_shared.cpp
• src/common/platform/sim/host/device_runner_base.cpp
• src/common/platform/sim/host/device_runner_base.h
• src/common/worker/chip_worker.cpp
• src/common/worker/chip_worker.h
• src/common/worker/pto_runtime_c_api.h

[ChaoWao]

Design: can the arena prewarm ride the existing cw.init path instead of a new API vertical?

The functionality is sound and build_and_cache_prebuilt_arena is a nice refactor. My concern is the surface area: this threads a full new vertical (simpler_prewarm_config C-API ×2, DeviceRunnerBase::prewarm_config ×2, ChipWorker::prewarm_config + dlsym, and the whole CTRL_PREWARM_CONFIG chain — Worker/WorkerManager/WorkerThread/LocalMailboxEndpoint::control_prewarm_config + base virtual + Python handler) for something that is fork-time-constant.

The key observation: the arena depends only on config.runtime_env.ring_*, and self._prewarm_config is fixed in init() — strictly before _start_hierarchical runs. So the sizing is a known constant before any fork, which means no cross-process control command is needed to deliver it.

Concretely:

• L2 — _init_level2 already calls cw.init(dev, bins) and then separately cw.prewarm_config(...). Fold the sizing into cw.init(dev, bins, prewarm_config=...) and build inside init. Drops the separate ChipWorker/DeviceRunner/C-API prewarm methods.
• L3 — the chip child entry _chip_process_loop(buf, bins, dev_id, …) is os.fork() with plain COW arg-passing (the PR itself notes _l3_bins crosses with "no pickle barrier"). Pass prewarm_config=self._prewarm_config the same way and build inside the child's cw.init. This removes the entire CTRL_PREWARM_CONFIG mailbox chain (~6 forwarders + the Python handler).

That collapses ~12 of the new symbols; what stays is the part that earns its keep: the public knob Worker.init(prewarm_config=) and the build_and_cache_prebuilt_arena helper.

Two bonuses from building inside the child's cw.init:

1. Ordering fits better. simpler_init inside cw.init already brings the device up, so the build has a live context, and it lands under the existing _INIT_DONE barrier ([Bug] Idle scheduler thread independently latches PTO2_ERROR_SCHEDULER_TIMEOUT, causing fatal cascade in distributed runs #897) — whose comment already anticipates a "long init right tail" (e.g. large device_malloc). That's a more natural home than a mailbox send after dw.init().
2. The L4+ gap closes cheaply. prewarm_config is currently validated+stored at L4 but never propagated (_start_hierarchical gates on device_ids, and inner_worker.init() is called without it), so it silently no-ops on L4+ trees. Propagating one param down is simpler than a control command.

The only thing that would justify the mailbox route is dynamic, post-init re-prewarm of a different sizing (children already forked, can't re-fork) — but the PR defines this as single-slot, one sizing per init, so that case is out of scope. "Consistency with control_prepare" isn't quite it either: prepare is per-digest and parent-driven, whereas the arena sizing is one fork-constant value identical for every child, which is exactly what COW inheritance is for.

Not blocking the concept — just think the same behavior lands with roughly a third of the API surface if it goes through cw.init.