Guard preparation off the claim path; pinned registration by default

Dockerfile.quick-start

@@ -91,12 +91,6 @@ ENV DYN_REQUEST_PLANE=tcp \
     PYTHONHASHSEED=0 \
     VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES=1
 
-# On-demand paging keeps Guard preparation off the synchronous registration
-# path. Fallback stays on so hardware without host ODP still registers; set
-# UCX_REG_NONBLOCK_FALLBACK=n to prove ODP is the path actually taken.
-ENV UCX_REG_NONBLOCK_MEM_TYPES=host \
-    UCX_IB_ODP_PREFETCH=n
-
 # Fail the build if an incompatible source revision displaced any part of the
 # serving path. The validated NCCL restore intentionally differs from the
 # torch wheel's declared pin, so a global `pip check` is not a useful canary.

docs/dev-guide.md

@@ -179,9 +179,10 @@ The callback names above represent services implemented by the framework
 adapter; they are not provided by KVCR itself.
 
 `KVCRConfig.startup_timeout_seconds` accepts a positive integer to allow longer
-initialization waits. Omitting it or setting it to `None` uses 30 seconds.
+initialization waits. Omitting it or setting it to `None` uses 60 seconds.
 The timeout applies separately to the local preparation and startup waits;
-transfer deadlines are unchanged.
+transfer deadlines are unchanged. A Guard core uses the service's
+`--startup-timeout-seconds`.
 
 For `on_resilience_event` and optional framework-memory quarantine, see the
 [resilience contract](design_overview.md#failed-peers-and-dangling-operations).
@@ -269,7 +270,8 @@ uv run python -m kvcr.kvcr_service \
   --compatibility-digest example-model-layout
 ```
 
-All flags except `--log-level` and `--heartbeat-timeout-ms` are required.
+All flags except `--log-level`, `--heartbeat-timeout-ms` and
+`--startup-timeout-seconds` are required.
 
 | Flag | Meaning |
 | --- | --- |
@@ -279,6 +281,7 @@ All flags except `--log-level` and `--heartbeat-timeout-ms` are required.
 | `--pool-sizes-gb` | Comma-separated usable sizes of the ordered pools in every group |
 | `--compatibility-digest` | Exact digest every claimant must provide |
 | `--heartbeat-timeout-ms` | Opt in to Guard promotion after a missed heartbeat; omitted means pidfd-only detection. Workers send heartbeats every 50 ms; minimum timeout is 1000 ms. |
+| `--startup-timeout-seconds` | Guard core startup budget; defaults to 60 s. Raise it for large pools, whose registration takes longer. |
 | `--log-level` | Service and Guard logging threshold; defaults to `$KVCR_LOG_LEVEL` or `INFO` |
 
 Heartbeat promotion kills the worker it promotes over. On a missed deadline the
@@ -336,15 +339,10 @@ its recovered records; a claimant that cannot inherit the endpoint is refused.
 A clean release returns the Guard to standby and the group to claimable.
 Takeover and handback cost time linear in the number of recovered blocks.
 
-For large host pools, UCX on-demand paging (ODP) avoids synchronous page
-registration during Guard preparation. It needs host ODP in the NIC, driver and
-UCX backend, and shifts page faults to first access rather than speeding up
-backing-file allocation.
-
-- `UCX_REG_NONBLOCK_MEM_TYPES=host` and `UCX_IB_ODP_PREFETCH=n`: set by
-  `Dockerfile.quick-start`.
-- `UCX_REG_NONBLOCK_FALLBACK=n`: fails instead of silently reverting to blocking
-  registration. Use it to confirm ODP is the path actually taken.
+Registering a large host pool can take tens of seconds, so a Guard prepares
+its core in the background and a claim does not wait for it. A promotion waits
+for the preparation to finish, up to `--startup-timeout-seconds`; while it
+waits, the Guard processes no release, abort or close.
 
 `--log-level DEBUG` (or `KVCR_LOG_LEVEL=DEBUG`) logs source transfer refusals
 with their reason, startup and Guard promotion/handback stage timings, and

src/kvcr/guard.py

@@ -382,6 +382,7 @@ def __init__(
         owner: _KVCRPoolOwner | None = None,
         refusing: Callable[[], bool] = lambda: False,
         heartbeat_timeout_seconds: float | None = None,
+        startup_timeout_seconds: int | None = None,
     ) -> None:
         if heartbeat_timeout_seconds is not None and not (
             _MIN_HEARTBEAT_TIMEOUT_SECONDS <= heartbeat_timeout_seconds < float("inf")
@@ -400,6 +401,8 @@ def __init__(
         self._configured: _TierConfig | None = None
         self._recovery = _RecoveryState(spec, compatibility_digest, pool_sizes_bytes)
         self._core = None
+        # Released cores still registering; _poll closes each once registered.
+        self._orphans: list[_KVCRCore] = []
         self._commands: queue.Queue[_Command] = queue.Queue()
         self._ops = {
             "claim": self._claim,
@@ -421,6 +424,7 @@ def __init__(
         self._failure_callback = failure_callback or (lambda guard, error: None)
         self._last_heartbeat: float | None = None
         self._heartbeat_timeout = heartbeat_timeout_seconds
+        self._startup_timeout = startup_timeout_seconds
         # Identity of the lease whose fencing failure has been logged: the actor
         # retries every millisecond, and one warning per lease is enough.
         self._fence_failed: object | None = None
@@ -619,7 +623,11 @@ def _serve_commands(self) -> None:
                 else:
                     # Nothing publishes into a pool no primary holds, so an
                     # unclaimed Guard waits rather than waking on an empty journal.
-                    busy = self._serving or self._control is not None
+                    busy = (
+                        self._serving
+                        or self._control is not None
+                        or bool(self._orphans)
+                    )
                     timeout = _POLL_SECONDS if busy and not self._failure else None
                     command = self._commands.get(timeout=timeout)
             except queue.Empty:
@@ -919,9 +927,13 @@ def _release(self) -> None:
             # Re-adopt lets start_primary() retain or replace the mirror.
             self._recovery.mirror = None
         elif self._core is not None:
-            self._core.close()
+            if self._core._progress._prepare_ready.is_set():
+                self._core.close()
+                self._control = None
+            else:
+                # Closing now would wait out the registration.
+                self._orphans.append(self._core)
             self._core = None
-            self._control = None
         if self._recovery.mirror is not None:
             self._recovery.release()
         if self._control is not None:
@@ -945,9 +957,17 @@ def _poll(self) -> bool:
         if self._failure is not None:
             return False
         try:
+            for orphan in list(self._orphans):
+                if orphan._progress._prepare_ready.is_set():
+                    # Dropped only once closed; shutdown still sees a failed one.
+                    orphan.close()
+                    self._orphans.remove(orphan)
             if self._serving:
                 self._core.poll_completed()
             else:
+                if self._core is not None:
+                    # Preparation runs in the background; report its failure here.
+                    self._core._progress.raise_if_failed()
                 return self._recovery.poll()
         except RecoveryMirrorError as error:
             self._recovery.invalidate_journal()
@@ -1055,6 +1075,7 @@ def reject_pin(keys: object) -> int:
                 nixl_agent_name=f"KVCR-Guard-{uuid.uuid4()}",
                 pool_layouts=self._configured.pool_layouts,
                 nixl_listen_port=0,
+                startup_timeout_seconds=self._startup_timeout,
             ),
             KVCRBindings(
                 reject_pin,
@@ -1072,7 +1093,8 @@ def reject_pin(keys: object) -> int:
         )
         self._core = core
         self._log_promotion_stage("core_init_done")
-        core._progress.prepare()
+        # Registration can take long; promotion's start() waits for it.
+        core._progress.prepare(wait=False)
 
     def _hand_back(self) -> None:
         """Stop serving, leaving this pool group's state where the next primary looks.
@@ -1091,20 +1113,22 @@ def _hand_back(self) -> None:
         self._serving = False
 
     def _close_resources(self) -> None:
-        try:
-            if self._core is not None:
-                self._core.close()
-        except BaseException:
-            if not self._core.is_quiescent():
-                # Still moving bytes, so nothing below runs: unmapping the
-                # pool under a thread still writing into it faults the process.
-                raise
-            logger.warning(
-                "KVCR Guard core close failed after reaching quiescence",
-                exc_info=True,
-            )
-        # Every close runs regardless; the first failure is the raised one.
+        # Stop every core, released ones first; a refused close keeps the
+        # pool mapped below.
         failure: BaseException | None = None
+        for core in (*self._orphans, self._core):
+            if core is None:
+                continue
+            try:
+                self._close_core(core)
+            except BaseException as error:  # noqa: BLE001 - raised below
+                failure = failure or error
+            else:
+                if core in self._orphans:
+                    self._orphans.remove(core)
+        if failure is not None:
+            raise failure
+        # Every close runs regardless; the first failure is the raised one.
         for give_back in (
             self._close_control,
             self._recovery.close,
@@ -1117,6 +1141,19 @@ def _close_resources(self) -> None:
         if failure is not None:
             raise failure
 
+    @staticmethod
+    def _close_core(core: _KVCRCore) -> None:
+        try:
+            core.close()
+        except BaseException:
+            if not core.is_quiescent():
+                # Still registering or moving bytes: unmapping under it faults.
+                raise
+            logger.warning(
+                "KVCR Guard core close failed after reaching quiescence",
+                exc_info=True,
+            )
+
     def _close_control(self) -> None:
         if self._control is not None:
             self._control.close()

src/kvcr/kvcr_service.py

@@ -50,6 +50,7 @@
     _reclaim_pool_if_orphaned,
     _unlink_if_identity,
 )
+from .progress import _STARTUP_TIMEOUT_SECONDS
 
 logger = logging.getLogger(__name__)
 
@@ -79,6 +80,7 @@ def __init__(
         journal_bytes: int,
         compatibility_digest: str,
         heartbeat_timeout_seconds: float | None = None,
+        startup_timeout_seconds: int | None = None,
     ) -> None:
         self._pool_dir = Path(pool_dir).resolve()
         if not self._pool_dir.is_dir():
@@ -112,6 +114,7 @@ def __init__(
                         owner=owner,
                         refusing=self._refusing.is_set,
                         heartbeat_timeout_seconds=heartbeat_timeout_seconds,
+                        startup_timeout_seconds=startup_timeout_seconds,
                     )
                 except BaseException:
                     # Nothing has recorded this pool group yet, so the sweep
@@ -492,6 +495,7 @@ def __init__(
         compatibility_digest: str,
         journal_bytes: int = _DEFAULT_JOURNAL_BYTES,
         heartbeat_timeout_seconds: float | None = None,
+        startup_timeout_seconds: int | None = None,
     ) -> None:
         try:
             probe_pidfd = os.pidfd_open(os.getpid())
@@ -515,6 +519,7 @@ def __init__(
             journal_bytes,
             compatibility_digest,
             heartbeat_timeout_seconds,
+            startup_timeout_seconds,
         )
         try:
             self._server = _ThreadingUnixServer(
@@ -658,6 +663,12 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         type=float,
         help="opt in to early Guard promotion (minimum 1000 ms; default: pidfd-only)",
     )
+    parser.add_argument(
+        "--startup-timeout-seconds",
+        type=int,
+        default=int(_STARTUP_TIMEOUT_SECONDS),
+        help="Guard core preparation and startup budget (default: %(default)s s)",
+    )
     parser.add_argument(
         "--log-level",
         type=str.upper,
@@ -673,6 +684,8 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         < float("inf")
     ):
         parser.error("--heartbeat-timeout-ms must be at least 1000 and finite")
+    if args.startup_timeout_seconds is not None and args.startup_timeout_seconds <= 0:
+        parser.error("--startup-timeout-seconds must be positive")
     return args
 
 
@@ -698,6 +711,7 @@ def main() -> None:
                 if args.heartbeat_timeout_ms is not None
                 else None
             ),
+            startup_timeout_seconds=args.startup_timeout_seconds,
         )
         signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask)
         logger.info(

src/kvcr/progress.py

@@ -28,7 +28,8 @@
 _IDLE_WAIT_SECONDS = 0.001
 _OP_CLEANUP_TIMEOUT_SECONDS = 5.0
 _JOIN_TIMEOUT_SECONDS = 10.0
-_STARTUP_TIMEOUT_SECONDS = 30.0
+_STARTUP_TIMEOUT_SECONDS = 60.0
+_PREPARE_LOG_INTERVAL_SECONDS = 10.0
 _RELEASE_LOG_INTERVAL_SECONDS = 1.0
 _STOP = object()
 _OpId = tuple[str, Any]
@@ -525,32 +526,48 @@ def _release_transfer(self, transfer_id: int, state: _TransferState) -> bool:
         return True
 
     def start(self) -> None:
-        self._activate.set()
-        if not self._thread.is_alive():
+        # A failed preparer is not restarted; its error is raised below.
+        if self._thread.ident is None or not (self._thread.is_alive() or self._failure):
+            self._activate.set()
             self._thread.start()
-        timeout = self._startup_timeout_seconds
-        if timeout is None:
-            timeout = _STARTUP_TIMEOUT_SECONDS
+        else:
+            # Wait out a background preparation, then activate on its own budget.
+            self._wait_prepared()
+            self._activate.set()
+        timeout = self._startup_budget()
         if not self._ready.wait(timeout=timeout):
             raise RuntimeError(
                 "KVCR progress initialization timed out after "
                 f"{timeout:g}s (stage: {self._startup_stage})"
             )
         self.raise_if_failed()
 
-    def prepare(self) -> None:
+    def prepare(self, wait: bool = True) -> None:
         """Create the agent and register memory without starting backends."""
         self._thread.start()
-        timeout = self._startup_timeout_seconds
-        if timeout is None:
-            timeout = _STARTUP_TIMEOUT_SECONDS
-        if not self._prepare_ready.wait(timeout=timeout):
-            raise RuntimeError(
-                "KVCR progress preparation timed out after "
-                f"{timeout:g}s (stage: {self._startup_stage})"
-            )
+        if not wait:
+            return
+        self._wait_prepared()
         self.raise_if_failed()
 
+    def _wait_prepared(self) -> None:
+        timeout = self._startup_budget()
+        deadline = time.monotonic() + timeout
+        while not self._prepare_ready.wait(
+            min(_PREPARE_LOG_INTERVAL_SECONDS, deadline - time.monotonic())
+        ):
+            if time.monotonic() >= deadline:
+                raise RuntimeError(
+                    "KVCR progress preparation timed out after "
+                    f"{timeout:g}s (stage: {self._startup_stage})"
+                )
+            logger.info("KVCR waiting for preparation: %s", self._startup_stage)
+
+    def _startup_budget(self) -> float:
+        if self._startup_timeout_seconds is not None:
+            return self._startup_timeout_seconds
+        return _STARTUP_TIMEOUT_SECONDS
+
     def submit(self, item: object) -> None:
         self.raise_if_failed()
         self._submissions.put(item)