Add: eager prebuilt runtime-arena prewarm at worker init

python/bindings/task_interface.cpp

@@ -873,6 +873,17 @@ NB_MODULE(_task_interface, m) {
             "Launch a callable_id from a TaskArgs (used for in-process callers). "
             "Returns None; timing is emitted as `[STRACE]` log markers."
         )
+        .def(
+            "prewarm_config",
+            [](ChipWorker &self, const CallConfig &config) {
+                self.prewarm_config(config);
+            },
+            nb::arg("config"),
+            "Eagerly build + cache the prebuilt runtime-arena image for config's ring "
+            "sizing so the first run() with the same sizing skips the (~800ms) cold "
+            "build. Config-only (no callable_id, no args); requires init() first. "
+            "A no-op for runtimes without a prebuilt arena (host_build_graph)."
+        )
         .def(
             "run_from_blob",
             [](ChipWorker &self, int32_t callable_id, uint64_t args_blob_ptr, size_t blob_capacity,

python/bindings/worker_bind.h

@@ -435,6 +435,16 @@ inline void bind_worker(nb::module_ &m) {
             },
             nb::arg("worker_id"), nb::arg("digest"), "Prewarm a NEXT_LEVEL child by callable digest."
         )
+        .def(
+            "control_prewarm_config",
+            [](Worker &self, int worker_id, const CallConfig &config) {
+                nb::gil_scoped_release release;
+                self.control_prewarm_config(worker_id, &config, sizeof(CallConfig));
+            },
+            nb::arg("worker_id"), nb::arg("config"),
+            "Prewarm a NEXT_LEVEL child's prebuilt runtime-arena cache from a "
+            "CallConfig so its first run skips the cold arena build."
+        )
         .def(
             "broadcast_register_all",
             [](Worker &self, uint64_t blob_ptr, uint64_t blob_size, nb::object digest) {

python/simpler/task_interface.py

@@ -1153,6 +1153,24 @@ def run(self, handle, args, config=None, **kwargs):
         state = self._resolve_handle(handle)
         self._run_slot(state.slot_id, args, config, **kwargs)
 
+    def prewarm_config(self, config=None, **kwargs) -> None:
+        """Eagerly build + cache the prebuilt runtime-arena image for ``config``.
+
+        Config-only warm-up: only the ``runtime_env`` ring sizing is read, so no
+        callable handle or args are needed. The first ``run`` with the same ring
+        sizing then hits the cache and skips the (~800ms) cold build. A no-op for
+        runtimes without a prebuilt arena (host_build_graph). Requires ``init``.
+
+        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)
+        self._impl.prewarm_config(config)
+
     def unregister_callable(self, handle) -> None:
         """Drop one live callable handle and release its private resources when final."""
         with self._registry_lock:

python/simpler/worker.py

@@ -219,6 +219,12 @@ def my_l4_orch(orch, args, config):
 _CTRL_MAP_HOST = 14
 _CTRL_UNMAP_HOST = 15
 
+# Eagerly build + cache the prebuilt runtime-arena image for a run config's ring
+# sizing so the first run() with the same sizing skips the (~800ms) cold build.
+# Sent from the parent at end of init() alongside CTRL_PREPARE. Carries the
+# packed CallConfig at _OFF_CONFIG (same slot a task dispatch uses).
+_CTRL_PREWARM_CONFIG = 16
+
 # MAP_HOST payload: token (u64), parent_va (u64), nbytes (u64), then the
 # NUL-free host-buffer shm name as the trailing bytes. UNMAP_HOST payload is the
 # token alone.
@@ -1226,6 +1232,12 @@ def _run_chip_main_loop(  # noqa: PLR0912, PLR0913, PLR0915 -- unified TASK_READ
                                 f"prepare chip={device_id}: callable hash {_format_digest(digest)} not registered"
                             )
                         _ensure_prepared(cw, registry, prepared, int(cid), device_id=device_id)
+                    elif sub_cmd == _CTRL_PREWARM_CONFIG:
+                        # Config-only: build + cache the prebuilt runtime arena for
+                        # this run config's ring sizing so the first run() hits the
+                        # cache. No callable digest — the arena is callable-agnostic.
+                        cfg = _read_config_from_mailbox(buf)
+                        cw.prewarm_config(cfg)
                     elif sub_cmd == _CTRL_REGISTER:
                         digest = _read_control_digest(buf)
                         payload_size = struct.unpack_from("Q", buf, _CTRL_OFF_ARG0)[0]
@@ -1566,6 +1578,11 @@ def __init__(
         self._hierarchical_start_mu = threading.Lock()
         self._hierarchical_start_cv = threading.Condition(self._hierarchical_start_mu)
 
+        # Optional CallConfig whose ring sizing is pre-warmed at init() so the
+        # first run() with the same sizing skips the (~800ms) cold prebuilt
+        # runtime-arena build. Set by init(prewarm_config=...); None = disabled.
+        self._prewarm_config: Any | None = None
+
         # Level-2 internals
         self._chip_worker: ChipWorker | None = None
 
@@ -3046,9 +3063,20 @@ def add_worker(self, worker: Worker) -> int:
     # init — auto-discovery
     # ------------------------------------------------------------------
 
-    def init(self) -> None:
+    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
 
         try:
             if self.level == 2:
@@ -3083,6 +3111,11 @@ def _init_level2(self) -> None:
             if isinstance(target, ChipCallable):
                 self._chip_worker._register_callable_at_slot(cid, target)
 
+        # Pre-warm the prebuilt runtime-arena cache for the declared config so the
+        # first run() with matching ring sizing skips the cold arena build.
+        if self._prewarm_config is not None:
+            self._chip_worker.prewarm_config(self._prewarm_config)
+
     def _init_hierarchical(self) -> None:
         device_ids = self._config.get("device_ids", [])
         n_sub = self._config.get("num_sub_workers", 0)
@@ -3311,6 +3344,14 @@ def _start_hierarchical(self) -> None:  # noqa: PLR0912 -- three parallel fork l
                         for worker_id in range(len(self._chip_shms)):
                             dw.control_prepare(worker_id, digest)
 
+            # Pre-warm each chip child's prebuilt runtime-arena cache for the
+            # declared config so the first dispatch with matching ring sizing
+            # skips the (~800ms) cold arena build. No-op when no prewarm config
+            # was passed to init() or for runtimes without a prebuilt arena.
+            if device_ids and self._prewarm_config is not None:
+                for worker_id in range(len(self._chip_shms)):
+                    dw.control_prewarm_config(worker_id, self._prewarm_config)
+
             self._hierarchical_started = True
             with self._hierarchical_start_cv:
                 self._hierarchical_start_state = "started"

src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp

@@ -622,8 +622,7 @@ static void apply_orch_sched_env_flags(Runtime *runtime) {
 // reserve sequence on a throwaway host arena. Idempotent across runs — the
 // pools are owned by DeviceRunner and freed in DeviceRunner::finalize().
 static bool ensure_static_arenas(
-    Runtime *runtime, const HostApi *api, const ArenaSizingConfig &sizing, const ArenaStaticSizes &sizes,
-    StaticArenaPtrs *out
+    const HostApi *api, const ArenaSizingConfig &sizing, const ArenaStaticSizes &sizes, StaticArenaPtrs *out
 ) {
     DeviceArena sizing_arena;  // discarded; only its computed arena_size is read
     PTO2RuntimeArenaLayout layout =
@@ -651,7 +650,6 @@ static bool ensure_static_arenas(
         LOG_ERROR("Failed to acquire pooled PTO2 shared memory");
         return false;
     }
-    runtime->set_gm_sm_ptr(out->gm_sm);
 
     out->runtime_arena_dev = api->acquire_pooled_runtime_arena();
     if (out->runtime_arena_dev == nullptr) {
@@ -678,8 +676,8 @@ static bool ensure_static_arenas(
 // The layout is stashed inside the image (rt->prebuilt_layout) so the AICPU can
 // recover every arena-internal offset after the rtMemcpy. Returns the layout
 // via `out_layout`; the runtime-arena device base travels separately on the
-// host Runtime (bind_launch_state), since the AICPU needs that pointer *before*
-// it can dereference the image.
+// host Runtime (set on the cache-hit path), since the AICPU needs that pointer
+// *before* it can dereference the image.
 static bool build_runtime_image(
     const ArenaSizingConfig &sizing, const ArenaStaticSizes &sizes, const StaticArenaPtrs &ptrs,
     DeviceArena *host_arena, PTO2RuntimeArenaLayout *out_layout
@@ -705,24 +703,6 @@ static bool build_runtime_image(
     return true;
 }
 
-// per-run: publish the launch state. Copy the staged args onto the runtime,
-// rtMemcpy the host image into the pooled runtime-arena region, and record the
-// device base + runtime offset the AICPU reads before dereferencing the image.
-static bool bind_launch_state(
-    Runtime *runtime, const HostApi *api, const StaticArenaPtrs &ptrs, const DeviceArena &host_arena,
-    const PTO2RuntimeArenaLayout &layout, const ChipStorageTaskArgs &device_args
-) {
-    runtime->set_orch_args(device_args);
-
-    int rc_upload = api->copy_to_device(ptrs.runtime_arena_dev, host_arena.base(), layout.offsets.arena_size);
-    if (rc_upload != 0) {
-        LOG_ERROR("Failed to rtMemcpy prebuilt runtime arena to device (rc=%d)", rc_upload);
-        return false;
-    }
-    runtime->set_prebuilt_arena(ptrs.runtime_arena_dev, layout.offsets.off_runtime);
-    return true;
-}
-
 static int bind_cached_runtime_image(
     Runtime *runtime, const HostApi *api, const PrebuiltRuntimeArenaCacheProbe &probe,
     const ChipStorageTaskArgs &device_args
@@ -754,7 +734,7 @@ static int bind_cached_runtime_image(
 }
 
 static void store_prebuilt_runtime_image(
-    Runtime *runtime, const HostApi *api, const PrebuiltRuntimeArenaCacheProbe &probe, const StaticArenaPtrs &ptrs,
+    const HostApi *api, const PrebuiltRuntimeArenaCacheProbe &probe, const StaticArenaPtrs &ptrs,
     const PTO2RuntimeArenaLayout &layout, const DeviceArena &host_arena
 ) {
     if (api->mark_prebuilt_runtime_arena_cached == nullptr) {
@@ -766,6 +746,40 @@ static void store_prebuilt_runtime_image(
     );
 }
 
+// Populate the DeviceRunnerBase prebuilt-arena cache for `sizing`: reserve the
+// pooled arenas, build the host image, rtMemcpy it to the pooled runtime-arena
+// region, and record it in the cache. Needs no Runtime and no per-run args —
+// the image is arg-independent (args are wired only on the cache-hit path via
+// bind_cached_runtime_image). Shared by the lazy first-run miss path and the
+// eager prewarm_config_impl entry, so both populate the cache identically.
+static bool build_and_cache_prebuilt_arena(const HostApi *api, const ArenaSizingConfig &sizing) {
+    ArenaStaticSizes sizes;
+    if (!derive_arena_static_sizes(sizing, &sizes)) {
+        return false;
+    }
+
+    StaticArenaPtrs ptrs;
+    if (!ensure_static_arenas(api, sizing, sizes, &ptrs)) {
+        return false;
+    }
+
+    DeviceArena host_arena;  // libc malloc backend; owns the image until upload
+    PTO2RuntimeArenaLayout layout;
+    if (!build_runtime_image(sizing, sizes, ptrs, &host_arena, &layout)) {
+        return false;
+    }
+
+    int rc_upload = api->copy_to_device(ptrs.runtime_arena_dev, host_arena.base(), layout.offsets.arena_size);
+    if (rc_upload != 0) {
+        LOG_ERROR("Failed to rtMemcpy prebuilt runtime arena to device (rc=%d)", rc_upload);
+        return false;
+    }
+
+    PrebuiltRuntimeArenaCacheProbe probe = make_prebuilt_runtime_arena_cache_probe(sizing);
+    store_prebuilt_runtime_image(api, probe, ptrs, layout, host_arena);
+    return true;
+}
+
 /**
  * Per-run binding: build device-side argument storage (tensor copy-out, GM
  * heap, PTO2 shared memory) and publish it to the runtime. Assumes the
@@ -777,9 +791,9 @@ static void store_prebuilt_runtime_image(
  * half runs only once per callable_id.
  *
  * Orchestrates the three lifecycles behind the bind: per-config arena sizing
- * (resolve_arena_sizing) + static pools (ensure_static_arenas) + host image
- * (build_runtime_image), and per-run args (stage_device_args) + launch publish
- * (bind_launch_state).
+ * (resolve_arena_sizing) + per-run args (stage_device_args) + the prebuilt
+ * runtime-arena image (build_and_cache_prebuilt_arena on a cache miss, then
+ * bind_cached_runtime_image wires the pointers onto the runtime).
  *
  * @param runtime    Pointer to pre-constructed Runtime
  * @param orch_args  Separated tensor/scalar arguments for this run
@@ -856,26 +870,18 @@ extern "C" int bind_callable_to_runtime_impl(
             return -1;
         }
         if (cache_rc != 0) {
-            ArenaStaticSizes sizes;
-            if (!derive_arena_static_sizes(sizing, &sizes)) {
+            // Miss: build + upload + cache the arena image, then re-run the
+            // cache lookup — which now hits and wires the runtime (args, gm/sm
+            // pointer, prebuilt-arena base). This is the same code the eager
+            // prewarm_config_impl path runs, so the cache contract is identical.
+            if (!build_and_cache_prebuilt_arena(api, sizing)) {
                 return -1;
             }
-
-            StaticArenaPtrs ptrs;
-            if (!ensure_static_arenas(runtime, api, sizing, sizes, &ptrs)) {
+            int rebind_rc = bind_cached_runtime_image(runtime, api, cache_probe, device_args);
+            if (rebind_rc != 0) {
+                LOG_ERROR("prebuilt runtime arena cache miss after eager build (rc=%d)", rebind_rc);
                 return -1;
             }
-
-            DeviceArena host_arena;  // libc malloc backend; owns the image until upload
-            PTO2RuntimeArenaLayout layout;
-            if (!build_runtime_image(sizing, sizes, ptrs, &host_arena, &layout)) {
-                return -1;
-            }
-
-            if (!bind_launch_state(runtime, api, ptrs, host_arena, layout, device_args)) {
-                return -1;
-            }
-            store_prebuilt_runtime_image(runtime, api, cache_probe, ptrs, layout, host_arena);
         }
     }
     int64_t t_prebuilt_end = _now_ms();
@@ -890,6 +896,32 @@ extern "C" int bind_callable_to_runtime_impl(
     return 0;
 }
 
+/**
+ * Eagerly populate the prebuilt runtime-arena cache for a run config, so the
+ * first bind_callable_to_runtime_impl with the same sizing hits the cache and
+ * skips the (~800ms) build + upload. Config-only: no callable, no per-run args
+ * — the arena image depends solely on the ring sizing. Requires the device to
+ * be initialized (pooled-arena device_malloc + rtMemcpy need a live context).
+ *
+ * @return 0 on success, -1 on failure
+ */
+extern "C" int prewarm_config_impl(
+    const HostApi *api, const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool
+) {
+    if (api == nullptr) {
+        LOG_ERROR("HostApi pointer is null");
+        return -1;
+    }
+
+    ArenaSizingConfig sizing;
+    if (!resolve_arena_sizing(ring_task_window, ring_heap, ring_dep_pool, &sizing)) {
+        return -1;
+    }
+
+    STRACE("simpler_prewarm.build");
+    return build_and_cache_prebuilt_arena(api, sizing) ? 0 : -1;
+}
+
 /**
  * Validate runtime results and cleanup.
  *

src/common/hierarchical/worker.h

@@ -100,6 +100,12 @@ class Worker {
     // used by the Python facade at end of _start_hierarchical).
     void control_prepare(int worker_id, const uint8_t *digest) { manager_.control_prepare(worker_id, digest); }
 
+    // Forward CTRL_PREWARM_CONFIG to a specific NEXT_LEVEL worker (config-arena
+    // prewarm path used by the Python facade at end of _start_hierarchical).
+    void control_prewarm_config(int worker_id, const void *config, size_t config_size) {
+        manager_.control_prewarm_config(worker_id, config, config_size);
+    }
+
     // Drive a single chip child through one CommDomain alloc / release.  The
     // Python orch facade is expected to call this on every participating chip
     // in parallel (one thread per chip) so the child-side `file_barrier` can

src/common/hierarchical/worker_manager.cpp

@@ -563,6 +563,17 @@ void LocalMailboxEndpoint::control_prepare(const uint8_t *digest) {
     run_control_command("control_prepare");
 }
 
+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_register(const char *shm_name, size_t blob_size, const uint8_t *digest) {
     std::lock_guard<std::mutex> lk(mailbox_mu_);
     // OFF_ERROR / OFF_ERROR_MSG are cleared by run_control_command — no
@@ -749,6 +760,11 @@ void WorkerThread::control_prepare(const uint8_t *digest) {
     endpoint_->control_prepare(digest);
 }
 
+void WorkerThread::control_prewarm_config(const void *config, size_t config_size) {
+    if (!endpoint_) throw std::runtime_error("control_prewarm_config: null endpoint");
+    endpoint_->control_prewarm_config(config, config_size);
+}
+
 void WorkerThread::control_register(const char *shm_name, size_t blob_size, const uint8_t *digest) {
     if (!endpoint_) throw std::runtime_error("control_register: null endpoint");
     endpoint_->control_register(shm_name, blob_size, digest);
@@ -977,6 +993,14 @@ void WorkerManager::control_prepare(int worker_id, const uint8_t *digest) {
     wt->control_prepare(digest);
 }
 
+void WorkerManager::control_prewarm_config(int worker_id, const void *config, size_t config_size) {
+    auto *wt = get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id);
+    if (wt == nullptr) {
+        throw std::runtime_error("control_prewarm_config: invalid worker_id " + std::to_string(worker_id));
+    }
+    wt->control_prewarm_config(config, config_size);
+}
+
 void WorkerManager::control_alloc_domain(int worker_id, const char *request_shm_name, const char *reply_shm_name) {
     auto *wt = get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id);
     if (wt == nullptr) {