Add: eager prebuilt runtime-arena prewarm at worker init#1322
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds configuration-driven runtime-arena prewarming. Python workers can accept a ChangesRuntime Arena Prewarm
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
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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
| if self._initialized: | ||
| raise RuntimeError("Worker already initialized") | ||
| self._prewarm_config = prewarm_config |
There was a problem hiding this comment.
To ensure consistent and early error reporting, we should validate the prewarm_config CallConfig during init() for both L2 and L3+ workers. Currently, L2 workers validate it during init(), but L3+ workers only validate it during the first run() (when _start_hierarchical is called). Validating it immediately in init() prevents delayed failures and makes debugging easier.
if self._initialized:
raise RuntimeError("Worker already initialized")
if prewarm_config is not None:
prewarm_config.validate()
self._prewarm_config = prewarm_config| void LocalMailboxEndpoint::control_prewarm_config(const void *config, size_t config_size) { | ||
| std::lock_guard<std::mutex> lk(mailbox_mu_); | ||
| write_control_args(mbox(), CTRL_PREWARM_CONFIG); | ||
| // Stage the packed CallConfig at MAILBOX_OFF_CONFIG (same slot a task dispatch | ||
| // uses) so the child reads it with its normal config decode. Clamp to the | ||
| // CallConfig region so a mis-sized caller cannot overrun the args area. | ||
| size_t n = config_size < sizeof(CallConfig) ? config_size : sizeof(CallConfig); | ||
| std::memcpy(mbox() + MAILBOX_OFF_CONFIG, config, n); | ||
| run_control_command("control_prewarm_config"); | ||
| } |
There was a problem hiding this comment.
If config is nullptr, std::memcpy will cause a segmentation fault (undefined behavior). We should add a defensive null check for config before proceeding. Additionally, if config_size < sizeof(CallConfig), copying only config_size bytes leaves the remaining bytes of the CallConfig region in the mailbox uninitialized (or containing garbage from previous task dispatches). We should zero-initialize the destination CallConfig region in the mailbox before copying to ensure all fields are safely initialized.
void LocalMailboxEndpoint::control_prewarm_config(const void *config, size_t config_size) {
if (config == nullptr) {
throw std::invalid_argument("control_prewarm_config: config pointer is null");
}
std::lock_guard<std::mutex> lk(mailbox_mu_);
write_control_args(mbox(), CTRL_PREWARM_CONFIG);
// Stage the packed CallConfig at MAILBOX_OFF_CONFIG (same slot a task dispatch
// uses) so the child reads it with its normal config decode. Clamp to the
// CallConfig region so a mis-sized caller cannot overrun the args area.
size_t n = config_size < sizeof(CallConfig) ? config_size : sizeof(CallConfig);
std::memset(mbox() + MAILBOX_OFF_CONFIG, 0, sizeof(CallConfig));
std::memcpy(mbox() + MAILBOX_OFF_CONFIG, config, n);
run_control_command("control_prewarm_config");
}References
- Do not assume that allocated shared memory or device memory is zero-initialized. Always explicitly initialize all fields to prevent garbage values from causing undefined behavior.
- Defensively clamp counts or sizes read from shared memory to the maximum capacity of local buffers to prevent stack overflows caused by potential memory corruption.
| int simpler_prewarm_config(DeviceContextHandle ctx, const CallConfig *config) { | ||
| if (ctx == NULL || config == NULL) return -1; | ||
| SimDeviceRunnerBase *runner = static_cast<SimDeviceRunnerBase *>(ctx); | ||
|
|
||
| pthread_once(&g_runner_key_once, create_runner_key); | ||
| pthread_setspecific(g_runner_key, ctx); | ||
|
|
||
| STRACE_NEW_INV(); | ||
| STRACE("simpler_prewarm"); | ||
|
|
||
| try { | ||
| int rc = runner->prewarm_config( | ||
| &g_host_api, config->runtime_env.ring_task_window, config->runtime_env.ring_heap, | ||
| config->runtime_env.ring_dep_pool | ||
| ); | ||
| pthread_setspecific(g_runner_key, nullptr); | ||
| return rc; | ||
| } catch (...) { | ||
| pthread_setspecific(g_runner_key, nullptr); | ||
| return -1; | ||
| } | ||
| } |
There was a problem hiding this comment.
To ensure that the thread-local g_runner_key is reliably cleaned up on all exit paths (including exceptions or any future early returns), we should use RAIIScopeGuard to manage the key's lifetime. This matches the robust pattern used in the onboard implementation in src/common/platform/onboard/host/c_api_shared.cpp.
int simpler_prewarm_config(DeviceContextHandle ctx, const CallConfig *config) {
if (ctx == NULL || config == NULL) return -1;
SimDeviceRunnerBase *runner = static_cast<SimDeviceRunnerBase *>(ctx);
pthread_once(&g_runner_key_once, create_runner_key);
pthread_setspecific(g_runner_key, ctx);
auto tsd_guard = RAIIScopeGuard([]() {
pthread_setspecific(g_runner_key, nullptr);
});
STRACE_NEW_INV();
STRACE("simpler_prewarm");
try {
return runner->prewarm_config(
&g_host_api, config->runtime_env.ring_task_window, config->runtime_env.ring_heap,
config->runtime_env.ring_dep_pool
);
} catch (...) {
return -1;
}
}References
- Use RAII guards to manage resources like thread-specific data and device contexts, ensuring cleanup is automatically handled on all function exit paths.
- Clear thread-local data set via pthread_setspecific on all function exit paths, including exception handlers.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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
📒 Files selected for processing (17)
python/bindings/task_interface.cpppython/bindings/worker_bind.hpython/simpler/task_interface.pypython/simpler/worker.pysrc/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.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/platform/sim/host/device_runner_base.cppsrc/common/platform/sim/host/device_runner_base.hsrc/common/worker/chip_worker.cppsrc/common/worker/chip_worker.hsrc/common/worker/pto_runtime_c_api.h
| Args: | ||
| config: Optional CallConfig; if None a default is created. | ||
| **kwargs: Overrides applied to config (e.g. ``ring_heap=...``). | ||
| """ | ||
| if config is None: | ||
| config = CallConfig() | ||
| for k, v in kwargs.items(): | ||
| setattr(config, k, v) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the invalid ring-sizing kwargs example.
Line 1166 suggests ring_heap=..., but ring_heap belongs to CallConfig.runtime_env; setattr(config, "ring_heap", ...) will not configure the arena sizing. Document configuring config.runtime_env.ring_heap before calling this method, or map these kwargs explicitly.
🤖 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 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.
| def init(self, prewarm_config=None) -> None: | ||
| """Initialize the worker. | ||
|
|
||
| Args: | ||
| prewarm_config: Optional CallConfig. When given, its ring sizing | ||
| (``runtime_env.ring_task_window`` / ``ring_heap`` / | ||
| ``ring_dep_pool``) is eagerly built + cached at init so the first | ||
| ``run`` with the same sizing skips the (~800ms) cold prebuilt | ||
| runtime-arena build. A no-op for runtimes without a prebuilt arena | ||
| (host_build_graph). ``None`` (default) disables prewarm. | ||
| """ | ||
| if self._initialized: | ||
| raise RuntimeError("Worker already initialized") | ||
| self._prewarm_config = prewarm_config |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
Make L3+ prewarming happen during init().
For hierarchical workers, init() only calls _init_hierarchical; the prewarm control at Lines 3150-3152 is deferred until _start_hierarchical(), which is first called by run(). Thus init(prewarm_config=...) returns before any arena is built and the first run() still blocks on prewarming. Start the hierarchy when a prewarm config is supplied, or revise the API contract.
🤖 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 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.
| void LocalMailboxEndpoint::control_prewarm_config(const void *config, size_t config_size) { | ||
| std::lock_guard<std::mutex> lk(mailbox_mu_); | ||
| write_control_args(mbox(), CTRL_PREWARM_CONFIG); | ||
| // Stage the packed CallConfig at MAILBOX_OFF_CONFIG (same slot a task dispatch | ||
| // uses) so the child reads it with its normal config decode. Clamp to the | ||
| // CallConfig region so a mis-sized caller cannot overrun the args area. | ||
| size_t n = config_size < sizeof(CallConfig) ? config_size : sizeof(CallConfig); | ||
| std::memcpy(mbox() + MAILBOX_OFF_CONFIG, config, n); | ||
| run_control_command("control_prewarm_config"); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Reject partial CallConfig payloads.
Line 572 permits undersized input, leaving the mailbox config tail from a prior dispatch while the child decodes a full CallConfig. Require a non-null payload of exactly sizeof(CallConfig) before writing the command.
Proposed fix
void LocalMailboxEndpoint::control_prewarm_config(const void *config, size_t config_size) {
std::lock_guard<std::mutex> lk(mailbox_mu_);
+ if (config == nullptr || config_size != sizeof(CallConfig)) {
+ throw std::invalid_argument("control_prewarm_config requires a complete CallConfig");
+ }
write_control_args(mbox(), CTRL_PREWARM_CONFIG);
- size_t n = config_size < sizeof(CallConfig) ? config_size : sizeof(CallConfig);
- std::memcpy(mbox() + MAILBOX_OFF_CONFIG, config, n);
+ std::memcpy(mbox() + MAILBOX_OFF_CONFIG, config, sizeof(CallConfig));
run_control_command("control_prewarm_config");
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| void LocalMailboxEndpoint::control_prewarm_config(const void *config, size_t config_size) { | |
| std::lock_guard<std::mutex> lk(mailbox_mu_); | |
| write_control_args(mbox(), CTRL_PREWARM_CONFIG); | |
| // Stage the packed CallConfig at MAILBOX_OFF_CONFIG (same slot a task dispatch | |
| // uses) so the child reads it with its normal config decode. Clamp to the | |
| // CallConfig region so a mis-sized caller cannot overrun the args area. | |
| size_t n = config_size < sizeof(CallConfig) ? config_size : sizeof(CallConfig); | |
| std::memcpy(mbox() + MAILBOX_OFF_CONFIG, config, n); | |
| run_control_command("control_prewarm_config"); | |
| void LocalMailboxEndpoint::control_prewarm_config(const void *config, size_t config_size) { | |
| std::lock_guard<std::mutex> lk(mailbox_mu_); | |
| if (config == nullptr || config_size != sizeof(CallConfig)) { | |
| throw std::invalid_argument("control_prewarm_config requires a complete CallConfig"); | |
| } | |
| write_control_args(mbox(), CTRL_PREWARM_CONFIG); | |
| std::memcpy(mbox() + MAILBOX_OFF_CONFIG, config, sizeof(CallConfig)); | |
| run_control_command("control_prewarm_config"); | |
| } |
🤖 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 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.
YunjiQin
force-pushed
the
feat/prebuilt-arena-prewarm
branch
2 times, most recently
from
9f8b6ab to
843cb59
Compare
Populate the prebuilt runtime-arena cache eagerly at init from a declared run config, so the first simpler_run with the same ring sizing skips the ~800ms cold build. On distributed (ep>1) runs this also removes the cross-card collective-wait bubble: prebuilt skew was converting 1:1 into device_wall idle (bind_end skew ~210ms -> ~2.6ms, device_wall converges). - 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). a2a3 tensormap_and_ringbuffer only. Validated in a2a3sim (L2 build-at- init + cache-hit on run; L3 per-chip dispatch) and on-device ep=4.
1 participant
What
Populate the prebuilt runtime-arena cache eagerly at worker init from a
declared run config, so the first
simpler_runwith the same ring sizingskips 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_wallidle. Measured on a distributed run:
bind_endskew ~210ms → ~2.6ms, anddevice_wallconverges across cards.How
build_and_cache_prebuilt_arena(api, sizing). The lazy first-run miss pathnow builds+caches then re-binds via the cache-hit path, so
runandprewarmshare one code path (dropsensure_static_arenas' runtime arg +bind_launch_state).simpler_prewarm_configC-API +DeviceRunnerBase::prewarm_config. Theruntime-specific
prewarm_config_implis a weak no-op default in thecommon onboard/sim
device_runner_base, overridden by the strong trb a2a3impl — so
hbgand archs without a prewarm impl still link (a5 untouched).ChipWorker(dlsymsimpler_prewarm_config) and nanobind(
_ChipWorker.prewarm_config,Worker.control_prewarm_config)._CTRL_PREWARM_CONFIGmailbox command +chip-loop handler + worker-manager
control_prewarm_configchain;Worker.init(prewarm_config=...)prewarms in-process (L2) or dispatchesper chip (L3).
Scope / safety
is prewarmed per init.
Testing
Not yet verified on-device in this PR branch — flagging for CI / reviewer
onboard run.