diff --git a/Dockerfile.quick-start b/Dockerfile.quick-start index 25791edc..5ab2e23b 100644 --- a/Dockerfile.quick-start +++ b/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. diff --git a/docs/dev-guide.md b/docs/dev-guide.md index 35009f5a..d34078e9 100644 --- a/docs/dev-guide.md +++ b/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 diff --git a/src/kvcr/guard.py b/src/kvcr/guard.py index 1f4d0611..8af4fc1d 100644 --- a/src/kvcr/guard.py +++ b/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() diff --git a/src/kvcr/kvcr_service.py b/src/kvcr/kvcr_service.py index 790355b4..a44d7edf 100644 --- a/src/kvcr/kvcr_service.py +++ b/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( diff --git a/src/kvcr/progress.py b/src/kvcr/progress.py index 9e772851..e27598ee 100644 --- a/src/kvcr/progress.py +++ b/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,12 +526,15 @@ 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 " @@ -538,19 +542,32 @@ def start(self) -> None: ) 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) diff --git a/tests/unit/test_guard.py b/tests/unit/test_guard.py index c8ba24f7..e241cc0b 100644 --- a/tests/unit/test_guard.py +++ b/tests/unit/test_guard.py @@ -10,6 +10,9 @@ import select import signal import socket +import threading +from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import TimeoutError as FutureTimeoutError from contextlib import contextmanager, nullcontext from types import SimpleNamespace from unittest.mock import Mock @@ -33,6 +36,7 @@ from kvcr.local_disk import _G3Residency from kvcr.local_dram import _LocalDramState from kvcr.memory import KVCRPoolSpec +from kvcr.progress import _KVCRProgress from kvcr.recovery_journal import ( _RECORD_BLOCK, _RECOVERY_ENCODER, @@ -42,7 +46,7 @@ _RecoveryBlock, _RecoveryMirror, ) -from kvcr.types import BlockKey +from kvcr.types import BlockKey, RegionDescriptor _TEST_SPEC = KVCRPoolSpec( pool_id="pool_0", @@ -335,11 +339,16 @@ def test_default_guard_waits_for_pidfd_without_heartbeats() -> None: assert guard._last_heartbeat is None -def test_a_serving_guard_reports_a_poll_failure_and_fences_its_core(caplog) -> None: +@pytest.mark.parametrize("serving", [True, False]) +def test_a_serving_guard_reports_a_poll_failure_and_fences_its_core( + caplog, serving +) -> None: """A mirror failure retires its journal, fences the core, and is reported.""" error = RecoveryMirrorError("recovery record is malformed") core = Mock() core.poll_completed.side_effect = error + # A standby polls its background preparation instead. + core._progress.raise_if_failed.side_effect = error core.close.side_effect = OSError("close failed") control = Mock() journal = Mock() @@ -348,10 +357,11 @@ def test_a_serving_guard_reports_a_poll_failure_and_fences_its_core(caplog) -> N guard._control = control guard._configured = _tier(16) guard._recovery._journal = journal - guard._serving = True + guard._serving = serving guard._core = core caplog.set_level(logging.ERROR, logger="kvcr.guard") + guard._poll() guard._poll() assert guard._failure is error @@ -480,7 +490,7 @@ def new_channel() -> Mock: assert journal.reset_called # Adoption registers the Guard transport without starting it. assert len(cores) == 1 - cores[0]._progress.prepare.assert_called_once_with() + cores[0]._progress.prepare.assert_called_once_with(wait=False) cores[0].start.assert_not_called() with pytest.raises(RecoveryMirrorError, match="another tier configuration"): guard._refuse_incompatible(_tier(16)) @@ -585,10 +595,115 @@ def new_channel() -> Mock: assert closed == ["core1", "control1", "core2", "control2", "attachment"] +def _registering_core(monkeypatch, step) -> Mock: + """A Guard core whose startup steps run through `step(name)`.""" + progress = _KVCRProgress( + lambda _: step("initialize"), + lambda _, __: ({}, False), + list, + lambda: None, + nixl_agent_name="guard", + memory_regions=({}, {"": RegionDescriptor(addr=128, size=128)}), + ) + progress._nixl_agent = Mock() + progress._nixl_agent.register_memory.side_effect = lambda *_, **__: step("register") + monkeypatch.setattr(progress, "_initialize_nixl", lambda: step("nixl")) + monkeypatch.setattr(progress, "_capture_agent_metadata", lambda: None) + core = Mock(_block_record_map={}, _progress=progress) + core.configure_mock( + start=progress.start, close=progress.close, is_quiescent=progress.is_quiescent + ) + return core + + +_PREPARED_ORDER = ["nixl", "register", "initialize"] + + +@pytest.mark.parametrize("promote", [True, False]) +@pytest.mark.parametrize("blocked", ["register", "nixl"]) +def test_only_promotion_waits_for_a_guard_registration( + monkeypatch, blocked, promote +) -> None: + """Promotion waits out a background preparation, in one fixed order, and a + release never does.""" + events, entered, release = [], threading.Event(), threading.Event() + + def step(name: str) -> None: + events.append(name) + if name == blocked: + entered.set() + assert release.wait(5) + + core = _registering_core(monkeypatch, step) + monkeypatch.setattr("kvcr.guard._KVCRCore", Mock(return_value=core)) + guard = _configurable_guard() + guard._recovery.attachment = _fake_attachment() + guard._recovery._journal = _Journal() + guard._adopt(Mock(), _tier(16)) + with ThreadPoolExecutor(1) as pool: + try: + assert entered.wait(1) + action = pool.submit(guard._promote if promote else guard._release) + # A release must not wait out the registration (or fail the service). + if promote: + with pytest.raises(FutureTimeoutError): + action.result(0.1) + else: + action.result(1) + assert action.running() is promote and not guard._serving + assert events == _PREPARED_ORDER[: _PREPARED_ORDER.index(blocked) + 1] + finally: + release.set() + action.result(5) + if promote: + assert guard._serving and events == _PREPARED_ORDER + core._progress.close() + else: + _wait_until(core._progress._prepare_ready.is_set) + guard._poll() + assert guard._orphans == [] and guard._failure is None + assert core.is_quiescent() and "initialize" not in events + + +def test_every_released_core_is_kept_until_it_has_registered(monkeypatch) -> None: + """Claims released while registering are all kept: closing the pool under + an unowned registration would fault the process.""" + entered, release = threading.Semaphore(0), threading.Event() + + def step(name: str) -> None: + if name == "register": + entered.release() + assert release.wait(5) + + cores = [_registering_core(monkeypatch, step) for _ in range(2)] + monkeypatch.setattr("kvcr.guard._KVCRCore", Mock(side_effect=cores)) + guard = _configurable_guard() + guard._recovery.attachment = _fake_attachment() + guard._recovery._journal = _Journal() + try: + for _ in range(2): + guard._adopt(Mock(), _tier(16)) + assert entered.acquire(timeout=1) + guard._release() + assert guard._orphans == cores + guard._poll() + assert guard._orphans == cores and guard._failure is None + # Once registered, the actor closes each orphan without a command. + guard._started = True + guard._thread.start() + release.set() + _wait_until(lambda: not guard._orphans, timeout=5) + finally: + release.set() + guard.close() + assert guard._failure is None + assert all(core.is_quiescent() for core in cores) + + def test_claim_preparation_failure_closes_transport(monkeypatch) -> None: core = Mock(_block_record_map={}) monkeypatch.setattr("kvcr.guard._KVCRCore", Mock(return_value=core)) - guard = _configurable_guard() + guard = _configurable_guard(startup_timeout_seconds=600) guard._recovery.attachment = _fake_attachment() guard._recovery._journal = _Journal() @@ -599,9 +714,11 @@ def test_claim_preparation_failure_closes_transport(monkeypatch) -> None: control.close.assert_called_once_with() core.close.assert_called_once_with() assert guard._failure is core._progress.prepare.side_effect - core._progress.prepare.assert_called_once_with() + core._progress.prepare.assert_called_once_with(wait=False) core.start.assert_not_called() assert guard._serving is False + # The service's budget reaches the Guard core; the default stays None. + assert guard_module._KVCRCore.call_args.args[0].startup_timeout_seconds == 600 @pytest.mark.parametrize( @@ -770,8 +887,9 @@ def test_recovery_close_error_stays_first_while_lease_cleanup_continues() -> Non owner.close.assert_called_once_with() +@pytest.mark.parametrize("orphaned", [False, True, "with_successor"]) def test_a_close_refused_by_a_moving_core_retains_the_pool_until_quiescent( - caplog, + caplog, orphaned ) -> None: """Close keeps everything while the core moves; once quiescent it contains.""" control = Mock() @@ -781,18 +899,31 @@ def test_a_close_refused_by_a_moving_core_retains_the_pool_until_quiescent( core.close.side_effect = error core.is_quiescent.return_value = False guard = _guard() - guard._control = control guard._configured = _tier(16) - guard._core = core guard._recovery.attachment = attachment caplog.set_level(logging.WARNING, logger="kvcr.guard") + if orphaned: + # A released core whose close failed in polling stays owned. (The + # failure fences the control channel, so none is attached here.) + guard._orphans.append(core) + guard._poll() + assert guard._failure is error and guard._orphans == [core] + caplog.clear() + else: + guard._control = control + guard._core = core + if orphaned == "with_successor": + # A newer core serving the pool while the released one still registers. + guard._core = successor = Mock() # Still moving bytes: nothing may be unmapped under it. with pytest.raises(RuntimeError) as raised: guard.close() assert raised.value is error - control.close.assert_not_called() attachment.close.assert_not_called() + if orphaned == "with_successor": + # The refused orphan must not leave the successor polling. + successor.close.assert_called_once_with() # Quiescent: the same failure is contained and everything goes back. core.is_quiescent.return_value = True @@ -800,8 +931,9 @@ def test_a_close_refused_by_a_moving_core_retains_the_pool_until_quiescent( assert caplog.messages == ["KVCR Guard core close failed after reaching quiescence"] assert caplog.records[0].exc_info is not None assert caplog.records[0].exc_info[1] is error - control.close.assert_called_once_with() attachment.close.assert_called_once_with() + if not orphaned: + control.close.assert_called_once_with() @pytest.mark.parametrize("refused_by", ["geometry", "handback", "hand-over"]) @@ -922,7 +1054,7 @@ def test_a_dropped_handback_still_leaves_the_new_lease_mirrored( # The pool went cold, not fatal: a fresh standby is ready but not serving. assert guard._serving is False assert guard._core is prepared - prepared._progress.prepare.assert_called_once_with() + prepared._progress.prepare.assert_called_once_with(wait=False) prepared.start.assert_not_called() # The claimant was told cold; the new lease is still mirrored, so this # primary's deposits survive its own death. diff --git a/tests/unit/test_kvcr_service.py b/tests/unit/test_kvcr_service.py index f07dd475..2e22dd6a 100644 --- a/tests/unit/test_kvcr_service.py +++ b/tests/unit/test_kvcr_service.py @@ -1528,6 +1528,18 @@ def test_service_heartbeat_timeout_is_opt_in() -> None: ) +def test_service_startup_timeout_defaults_to_sixty_seconds() -> None: + assert _parse_args(_service_args("1")).startup_timeout_seconds == 60 + with pytest.raises(SystemExit): + _parse_args(_service_args("1") + ["--startup-timeout-seconds", "0"]) + assert ( + _parse_args( + _service_args("1") + ["--startup-timeout-seconds", "600"] + ).startup_timeout_seconds + == 600 + ) + + @pytest.mark.parametrize( "value", [ diff --git a/tests/unit/test_progress.py b/tests/unit/test_progress.py index a4b13105..ed155f30 100644 --- a/tests/unit/test_progress.py +++ b/tests/unit/test_progress.py @@ -240,6 +240,8 @@ def test_prepare_registers_memory_before_backend_activation( ) assert agent.deregistered == [1] assert progress.is_quiescent() + with pytest.raises(RuntimeError, match="once"): + progress.start() # Stages are tracked either way; only DEBUG reads clocks and logs them. assert progress._startup_stage == "cleanup" stages = [m for m in caplog.messages if "progress_startup_stage " in m] @@ -249,6 +251,7 @@ def test_prepare_registers_memory_before_backend_activation( @pytest.mark.parametrize("prepared", [False, True]) def test_instance_startup_budget_remains_bounded(monkeypatch, prepared) -> None: """A longer default must not hide a blocked instance's shorter deadline.""" + join_timeout = progress_module._JOIN_TIMEOUT_SECONDS entered, release = threading.Event(), threading.Event() def initialize_agent(): @@ -267,12 +270,49 @@ def initialize_agent(): with pytest.raises(RuntimeError, match="timed out after 0.001s"): (progress.prepare if prepared else progress.start)() assert entered.wait(1) + # Close cannot cancel initialization, so nothing may be released yet. + monkeypatch.setattr(progress_module, "_JOIN_TIMEOUT_SECONDS", 0.01) + with pytest.raises(RuntimeError, match="did not stop"): + progress.close() + assert not progress.is_quiescent() finally: + monkeypatch.setattr(progress_module, "_JOIN_TIMEOUT_SECONDS", join_timeout) release.set() progress.close() assert progress.is_quiescent() +def test_start_bounds_its_wait_for_a_background_preparation(monkeypatch) -> None: + """A stalled registration fails start() instead of parking its caller.""" + entered, release = threading.Event(), threading.Event() + + def register() -> None: + entered.set() + assert release.wait(5) + + progress = _KVCRProgress( + lambda _: None, + lambda _, __: ({}, False), + list, + lambda: None, + startup_timeout_seconds=0.001, + ) + monkeypatch.setattr(progress, "_initialize_nixl", lambda: None) + monkeypatch.setattr(progress, "_register_memory_regions", register) + progress.prepare(wait=False) + assert entered.wait(1) + try: + with pytest.raises(RuntimeError, match="preparation timed out after 0.001s"): + progress.start() + # The registration is still running: nothing was activated or released. + assert not progress._activate.is_set() + assert not progress.is_quiescent() + finally: + release.set() + progress.close() + assert progress.is_quiescent() + + def test_close_drains_queued_submissions(monkeypatch) -> None: entered, resume, stopping = (threading.Event() for _ in range(3)) received = [] @@ -784,7 +824,8 @@ def close() -> None: assert lifecycle_threads[0] != main_thread -def test_startup_failure_is_reported_to_main() -> None: +@pytest.mark.parametrize("prepared", [False, True]) +def test_startup_failure_is_reported_to_main(monkeypatch, prepared) -> None: expected = RuntimeError("startup failed") def initialize(_progress: _KVCRProgress) -> None: @@ -796,6 +837,11 @@ def initialize(_progress: _KVCRProgress) -> None: lambda: [], lambda: None, ) + if prepared: + # A preparer that already died is not restarted. + monkeypatch.setattr(progress, "_initialize_nixl", Mock(side_effect=expected)) + progress.prepare(wait=False) + progress._thread.join(5) with pytest.raises(RuntimeError, match="startup failed") as exc_info: progress.start()