diff --git a/examples/disaggregated/slurm/cache_transceiver_test/README.md b/examples/disaggregated/slurm/cache_transceiver_test/README.md index e99c57dc4fe3..25e2fd0b2668 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/README.md +++ b/examples/disaggregated/slurm/cache_transceiver_test/README.md @@ -55,7 +55,9 @@ itself flows over UCX/NIXL. - Every rank prints `UCX_TLS` and `UCX_NET_DEVICES`; `UCX_PROTO_INFO=used` is set so UCX >= 1.21 captures the selected GPU↔GPU transport in the stderr logs. - Failing transfers are reported per cell as `TRANSFER_ERROR` / `MISMATCH` / - `TIMEOUT` (the run continues — a single bad case never aborts the sweep). + `TIMEOUT`. A verification mismatch is safe to continue; a result that does + not prove transfer quiescence hard-aborts only the current UCX sweep so its + cache manager cannot recycle or deregister pages still owned by transport. ## Outputs (under `environment.work_dir`) @@ -112,10 +114,16 @@ within a `(combination)` across UCX sweeps. so no cross-node gather is needed. - A bad `UCX_NET_DEVICES`/`UCX_TLS` can hang a transfer. Hangs are handled at three layers so one stuck sweep never blocks the others: - 1. **`signal.alarm`** (per cell, `timeout_per_cell_s`) recovers Python-level - stalls and continues within the sweep. - 2. **Watchdog thread** (per cell) records `TIMEOUT` and `SIGKILL`s the process - for hangs inside a *GIL-released* native call (e.g. + 1. **The Python sender deadline plus `signal.alarm`** bounds Python-level + stalls. `timeout_per_cell_s` must be greater than five seconds. The Python + sender deadline is five seconds lower, leaving time for CTX and GEN to + exchange a final ownership decision. The alarm is the per-cell bound for + the C++ path and for later requests after the shared cell budget is spent. + A timeout or any other result without explicit completion retains the + pages and hard-aborts the current sweep; the harness never rebuilds a + cache manager over possibly active transfer memory. + 2. **Watchdog thread** (per cell) best-effort records `TIMEOUT` and `SIGKILL`s + the process for hangs inside a *GIL-released* native call (e.g. `check_*_transfer_status`); `srun --kill-on-bad-exit` then tears down the sweep. This is **best-effort**: it cannot fire if the hang is in a native call that *holds* the GIL (e.g. the UCX connection handshake inside diff --git a/examples/disaggregated/slurm/cache_transceiver_test/config.yaml b/examples/disaggregated/slurm/cache_transceiver_test/config.yaml index 5b108698e796..a3a58f4a68c2 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/config.yaml +++ b/examples/disaggregated/slurm/cache_transceiver_test/config.yaml @@ -98,9 +98,11 @@ ucx_env_sweep: run: - # Per (combination, request_length) cell hang threshold; a stuck transfer is reported as - # TIMEOUT and the harness moves on. The per-cell watchdog fires shortly after - # this and is always clamped below max_sweep_s. + # Per (combination, request_length) cell hard limit; must be >5 seconds. The + # Python sender deadline is set 5 seconds lower so it can return and complete + # the CTX/GEN ownership handshake first. The alarm remains the C++ path's + # per-cell bound. If quiescence is still unproven, the harness records the + # remaining cells and hard-aborts this sweep without reusing KV memory. timeout_per_cell_s: 60 # Hard wall-clock cap for ONE full sweep (all cases x request lengths). The # outer `timeout` around each srun enforces this, so a hung sweep can never diff --git a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index 3bfb64a8dcbc..9920dc740de7 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py @@ -40,8 +40,9 @@ import pickle import signal import sys +import time from dataclasses import dataclass, field -from typing import List, Optional +from typing import Any, Iterable, List, Optional, Sequence import torch import yaml @@ -69,6 +70,8 @@ # Must match report.py. RID_COMBINATION_STRIDE = 1_000_000 RID_REQLEN_STRIDE = 10_000 +TRANSFER_TIMEOUT_GRACE_SECONDS = 5 +ABORT_COORDINATION_TIMEOUT_SECONDS = 2 DTYPE_MAP = {"FP8": DataType.FP8, "HALF": DataType.HALF, "BF16": DataType.BF16} @@ -287,36 +290,203 @@ class _TransferError(Exception): pass -def _wait_ctx_complete(xcvr, rid, runtime): - """Block until this ctx request's send finishes (or errors). - - * C++ transceiver: check_context_transfer_status(None) is a true block-all; - a single call suffices (see the comment at the call site). - * PYTHON (V2) transceiver: even under block_all, each TxSession wait is - bounded by kv_transfer_sender_future_timeout_ms (default 1000 ms). A - slower peer handshake makes the call log "TxSession ... timed out" and - return with the request still DISAGG_CONTEXT_TRANS_IN_PROGRESS -- in the - real executor that is benign (the session stays open and is re-polled - every iteration), but returning here would free and refill the KV blocks - mid-flight, so the receiver reads the NEXT request's pattern (verify FAIL - on every request except the last). So poll until - this rid lands in the completed/failed lists. Collectively safe: every - ctx rank loops on the same rid and the per-call consensus makes all - ranks observe completion on the same iteration. The per-cell - signal.alarm and the hang detector bound the loop. - """ - if runtime != "PYTHON": - xcvr.check_context_transfer_status(None) +class _FatalTransferError(_TransferError): + """A transfer may still own KV pages, so this process must not reuse them.""" + + +def _request_ids(values: Iterable[Any]) -> set[int]: + return {value.py_request_id if hasattr(value, "py_request_id") else value for value in values} + + +def _context_completion_error( + rid: int, + completed: Iterable[Any], + failed: Iterable[Any], + state: Any, + error_state: Any, +) -> Optional[str]: + completed_rids = _request_ids(completed) + failed_rids = _request_ids(failed) + if rid in failed_rids: + return f"ctx transfer failed for rid={rid}" + if state == error_state: + return f"ctx transfer reported DISAGG_TRANS_ERROR for rid={rid}" + if rid not in completed_rids: + return ( + f"ctx block-all returned without completing rid={rid}: " + f"completed={sorted(completed_rids)} failed={sorted(failed_rids)}" + ) + return None + + +def _gen_completion_error( + rid: int, + completed: Iterable[Any], + failed: Iterable[Any], + cancelled: Iterable[Any], + state: Any, + complete_state: Any, + error_state: Any, +) -> Optional[str]: + completed_rids = _request_ids(completed) + failed_rids = _request_ids(failed) + cancelled_rids = _request_ids(cancelled) + if rid in failed_rids: + return f"gen transfer failed for rid={rid}" + if rid in cancelled_rids: + return f"gen transfer cancelled for rid={rid}" + if state == error_state: + return f"gen transfer reported DISAGG_TRANS_ERROR for rid={rid}" + if rid not in completed_rids: + return ( + f"gen block-all returned without completing rid={rid}: " + f"completed={sorted(completed_rids)} failed={sorted(failed_rids)} " + f"cancelled={sorted(cancelled_rids)}" + ) + if state != complete_state: + return f"gen transfer returned rid={rid} as complete with nonterminal state={state}" + return None + + +def _can_release_sequence(transfer_may_have_started: bool, transfer_completed: bool) -> bool: + """Return whether the transceiver has relinquished its ownership of the pages.""" + return not transfer_may_have_started or transfer_completed + + +def _release_sequence_if_safe( + mgr: Any, + req: Any, + kv_handle: Any, + use_v2: bool, + *, + transfer_may_have_started: bool, + transfer_completed: bool, +) -> bool: + if req is None or not _can_release_sequence(transfer_may_have_started, transfer_completed): + return False + free_sequence(mgr, req, kv_handle, use_v2) + return True + + +def _validate_context_completion(req: Any, status: Optional[Sequence[Any]]) -> None: + if status is None or len(status) != 2: + raise _FatalTransferError( + f"ctx block-all returned invalid status for rid={req.py_request_id}" + ) + completed, failed = status + error = _context_completion_error( + req.py_request_id, + completed, + failed, + req.state, + LlmRequestState.DISAGG_TRANS_ERROR, + ) + if error is not None: + raise _FatalTransferError(error) + + +def _validate_python_gen_completion(req: Any, status: Optional[Sequence[Any]]) -> None: + if status is None or len(status) != 3: + raise _FatalTransferError( + f"gen block-all returned invalid status for rid={req.py_request_id}" + ) + completed, failed, cancelled = status + error = _gen_completion_error( + req.py_request_id, + completed, + failed, + cancelled, + req.state, + LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE, + LlmRequestState.DISAGG_TRANS_ERROR, + ) + if error is not None: + raise _FatalTransferError(error) + + +def _first_reason(reasons: Iterable[Optional[str]], fallback: str) -> str: + return next((reason for reason in reasons if reason), fallback) + + +def _exchange_release_decision( + role: str, + comm: Any, + is_leader: bool, + zmq_sock: Any, + local_safe: bool, + local_reason: str = "", +) -> tuple[bool, str]: + """Agree across both MPI roles before either side releases KV pages.""" + role_safe = bool(comm.allreduce(1 if local_safe else 0, op=MPI.MIN)) + gathered_reasons = comm.gather(local_reason or None, root=0) + decision = None + if is_leader: + role_reason = ( + "" + if role_safe + else _first_reason( + gathered_reasons, + f"{role} peer rank did not prove transfer completion", + ) + ) + local_status = "COMPLETE" if role_safe else "FATAL" + try: + if role == "gen": + zmq_sock.send(pickle.dumps((local_status, role_reason))) + peer_status, peer_reason = pickle.loads(zmq_sock.recv()) + peer_role = "ctx" if role == "gen" else "gen" + if peer_status not in ("COMPLETE", "FATAL"): + combined_safe = False + combined_reason = f"invalid {peer_role} peer release status: {peer_status!r}" + elif role == "gen": + combined_safe = role_safe and peer_status == "COMPLETE" + combined_reason = "" if combined_safe else str(peer_reason or role_reason) + else: + combined_safe = role_safe and peer_status == "COMPLETE" + combined_reason = ( + "" + if combined_safe + else role_reason + if not role_safe + else str(peer_reason or "gen did not prove transfer completion") + ) + if role != "gen": + zmq_sock.send( + pickle.dumps(("COMPLETE" if combined_safe else "FATAL", combined_reason)) + ) + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - converted to process-fatal ownership outcome + combined_safe = False + combined_reason = f"final release handshake failed: {e!r}" + decision = (combined_safe, combined_reason) + return comm.bcast(decision, root=0) + + +def _hard_abort_process(comm: Any) -> None: + """Terminate without running transceiver/KV-manager finalizers.""" + try: + comm.Abort(137) + except Exception: # noqa: BLE001 - SIGKILL is the mandatory fallback + pass + os.kill(os.getpid(), signal.SIGKILL) + raise RuntimeError("SIGKILL unexpectedly returned") + + +def _coordinate_abort_after_leader_flush(comm: Any) -> None: + """Best-effort bounded rendezvous after the leader persists diagnostics.""" + try: + request = comm.Ibarrier() + deadline = time.monotonic() + ABORT_COORDINATION_TIMEOUT_SECONDS + while not request.Test(): + if time.monotonic() >= deadline: + return + time.sleep(0.01) + except Exception: # noqa: BLE001 - abort must remain the fallback return - while True: - completed, failed = xcvr.check_context_transfer_status(None) - if rid in failed: - raise _TransferError(f"ctx transfer failed for rid={rid}") - if rid in completed: - return -def _wait_gen_complete(xcvr, req, runtime): +def _wait_gen_complete(xcvr: Any, req: LlmRequest, runtime: str) -> None: """Block until this gen request's receive finishes (or errors). Block-all is the only safe wait here: returning while the receive is still @@ -333,7 +503,8 @@ def _wait_gen_complete(xcvr, req, runtime): detector bound this loop, so a genuinely stuck transfer is still caught. """ if runtime == "PYTHON": - xcvr.check_gen_transfer_status(None) # block_all + status = xcvr.check_gen_transfer_status(None) # block_all + _validate_python_gen_completion(req, status) return import time @@ -346,6 +517,10 @@ def _wait_gen_complete(xcvr, req, runtime): if req.state in terminal: break time.sleep(0.001) + if req.state == LlmRequestState.DISAGG_TRANS_ERROR: + raise _FatalTransferError( + f"gen transfer reported DISAGG_TRANS_ERROR for rid={req.py_request_id}" + ) def run_one_request( @@ -353,11 +528,10 @@ def run_one_request( ): """Transfer one request and verify it (gen side). - The per-request ZMQ handshake is lockstep-safe: the ctx leader ALWAYS sends - exactly one reply ("OK"+context_phase_params, or "ABORT"+reason) for each - gen "go", so a ctx-side error never leaves the gen side blocked or the - sockets out of sync. Errors (local exceptions or a transceiver-reported - DISAGG_TRANS_ERROR) are raised so the caller records TRANSFER_ERROR. + The leaders exchange setup and final ownership decisions over ZMQ. Neither + role releases its sequence until every local rank and the peer role report + successful completion. A result without quiescence proof raises + _FatalTransferError so the caller hard-aborts without running finalizers. Returns the gen-side verification result (True/False), or None on ctx. """ is_ctx = role == "ctx" @@ -366,6 +540,7 @@ def run_one_request( if is_ctx: local_err = None req = kv_handle = None + transfer_may_have_started = False try: req = make_request(True, rid, req_len, runtime) kv_handle = add_sequence(kvm, req, req_len, use_v2) @@ -373,53 +548,100 @@ def run_one_request( tensorrt_llm.logger.info( f"[ctx r{rank}] rid={rid} len={req_len}: transfer START (send)" ) + # The call can dispatch KV work before raising, so ownership becomes + # uncertain as soon as we enter it. + transfer_may_have_started = True xcvr.respond_and_send_async(req) - except Exception as e: # noqa: BLE001 - relay failure to gen, then raise + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - relayed to gen, then classified local_err = e - # Instance-wide consensus: every ctx rank must take the same branch at the - # collective check below, even if only some ranks failed (e.g. a UCX - # device error on a subset of GPUs). allreduce doubles as a barrier. + any_failed = comm.allreduce(1 if local_err is not None else 0, op=MPI.MAX) + any_started = comm.allreduce(1 if transfer_may_have_started else 0, op=MPI.MAX) + handshake_error = None if is_leader: - zmq_sock.recv() # gen leader's "go" - if not any_failed: - # context_phase_params is picklable and carries everything gen - # needs (endpoint, ctx_dp_rank, first_gen/draft tokens). - zmq_sock.send(pickle.dumps(("OK", req.context_phase_params))) - else: - reason = repr(local_err) if local_err is not None else "peer ctx rank failed" - zmq_sock.send(pickle.dumps(("ABORT", reason))) + try: + message = zmq_sock.recv() # gen leader's "go" + if message != b"go": + raise _TransferError(f"unexpected gen handshake message: {message!r}") + if not any_failed: + # context_phase_params carries the endpoint and generation metadata. + zmq_sock.send(pickle.dumps(("OK", req.context_phase_params))) + else: + reason = repr(local_err) if local_err is not None else "peer ctx rank failed" + status = "FATAL" if any_started else "ABORT" + zmq_sock.send(pickle.dumps((status, reason))) + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - broadcast before raising + handshake_error = repr(e) + handshake_error = comm.bcast(handshake_error, root=0) + if handshake_error is not None: + raise _FatalTransferError(f"initial ctx/gen handshake failed: {handshake_error}") if any_failed: - if req is not None: - try: - free_sequence(kvm, req, kv_handle, use_v2) - except Exception: # noqa: BLE001 - pass - raise local_err if local_err is not None else _TransferError("peer ctx rank failed") - # Always block_all (None): with a request count, C++ only waits up to - # kv_transfer_sender_future_timeout_ms (1000ms) and returns even if the - # transfer is still in progress. NIXL/UCX cold-start connection setup can - # exceed that, so the harness would free the request mid-transfer, leaving - # the gen side hung and the ctx sender thread asserting on a freed session. - # The PYTHON runtime additionally needs a poll loop on top of block_all; - # see _wait_ctx_complete. - _wait_ctx_complete(xcvr, rid, runtime) + reason = repr(local_err) if local_err is not None else "peer ctx rank failed" + if any_started: + raise _FatalTransferError(reason) + _release_sequence_if_safe( + kvm, + req, + kv_handle, + use_v2, + transfer_may_have_started=False, + transfer_completed=False, + ) + raise _TransferError(reason) + + transfer_error = None + try: + # block_all retries short polling slices until completion or the + # request-level KV-transfer deadline. Only the exact completed ID + # proves that the sender has stopped reading this request's pages. + status = xcvr.check_context_transfer_status(None) + _validate_context_completion(req, status) + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - becomes a cross-role fatal decision + transfer_error = e + + safe_to_release, reason = _exchange_release_decision( + role, + comm, + is_leader, + zmq_sock, + local_safe=transfer_error is None, + local_reason=repr(transfer_error) if transfer_error is not None else "", + ) + if not safe_to_release: + raise _FatalTransferError(reason) state = req.state tensorrt_llm.logger.info(f"[ctx r{rank}] rid={rid}: transfer DONE (send), state={state}") free_sequence(kvm, req, kv_handle, use_v2) - if state == LlmRequestState.DISAGG_TRANS_ERROR: - raise _TransferError("ctx transfer reported DISAGG_TRANS_ERROR") return None # gen side + initial_error = None if is_leader: - zmq_sock.send(b"go") - status, payload = pickle.loads(zmq_sock.recv()) + try: + zmq_sock.send(b"go") + status, payload = pickle.loads(zmq_sock.recv()) + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - ctx may already own in-flight pages + status, payload = "FATAL", f"initial ctx/gen handshake failed: {e!r}" + initial_error = repr(e) else: status, payload = None, None - status, payload = comm.bcast((status, payload), root=0) # syncs gen ranks + status, payload, initial_error = comm.bcast( + (status, payload, initial_error), root=0 + ) # syncs gen ranks + if initial_error is not None or status == "FATAL": + raise _FatalTransferError(str(payload)) if status == "ABORT": raise _TransferError(f"ctx aborted: {payload}") + if status != "OK": + raise _FatalTransferError(f"unexpected ctx handshake status: {status!r}") ctx_params = payload local_err = None @@ -428,36 +650,67 @@ def run_one_request( req = make_request(False, rid, req_len, runtime, ctx_params=ctx_params) kv_handle = add_sequence(kvm, req, req_len, use_v2) tensorrt_llm.logger.info(f"[gen r{rank}] rid={rid} len={req_len}: transfer START (recv)") + # The context sender is already live, and this call may partially + # dispatch before raising. Any setup error from here is process-fatal. xcvr.request_and_receive_async(req) + except _Timeout: + raise except Exception as e: # noqa: BLE001 local_err = e - # Instance-wide consensus so all gen ranks take the same branch at check_gen. + any_failed = comm.allreduce(1 if local_err is not None else 0, op=MPI.MAX) - if any_failed: - if req is not None: - try: - free_sequence(kvm, req, kv_handle, use_v2) - except Exception: # noqa: BLE001 - pass - raise local_err if local_err is not None else _TransferError("peer gen rank failed") - # Block until the receive actually completes (mirrors the ctx side) instead - # of returning on the sender timeout. See _wait_gen_complete for why the C++ - # path polls an int rather than passing None. - _wait_gen_complete(xcvr, req, runtime) - # The receive may land on a side CUDA stream; sync before reading. - torch.cuda.synchronize() + transfer_error = local_err + if not any_failed: + try: + # See _wait_gen_complete for why the C++ path polls an int rather + # than passing None. Python validates the exact returned request ID. + _wait_gen_complete(xcvr, req, runtime) + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - becomes a cross-role fatal decision + transfer_error = e + elif transfer_error is None: + transfer_error = _FatalTransferError("peer gen rank failed during receive setup") + + if transfer_error is None: + try: + # A terminal transceiver result is not sufficient ownership proof: + # wait for any side-stream work touching the received pages before + # either role is allowed to release its sequence. + torch.cuda.synchronize() + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - quiescence remains unproven + transfer_error = e + + safe_to_release, reason = _exchange_release_decision( + role, + comm, + is_leader, + zmq_sock, + local_safe=transfer_error is None, + local_reason=repr(transfer_error) if transfer_error is not None else "", + ) + if not safe_to_release: + raise _FatalTransferError(reason) + # Transfer completion and CUDA synchronization have released transport + # ownership. Verification errors are now safe mismatches: retain lockstep, + # free the sequence, and keep testing. + try: + local_ok = verify_request(kvm, req.py_request_id, rank, n_local_layers) + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - transfer is already proven quiescent + tensorrt_llm.logger.error(f"[gen r{rank}] rid={rid}: verification error: {e!r}") + local_ok = False + ok = bool(comm.allreduce(1 if local_ok else 0, op=MPI.MIN)) state = req.state - ok = state != LlmRequestState.DISAGG_TRANS_ERROR and verify_request( - kvm, req.py_request_id, rank, n_local_layers - ) tensorrt_llm.logger.info( f"[gen r{rank}] rid={rid}: transfer DONE (recv), state={state}, " f"verify={'PASS' if ok else 'FAIL'}" ) free_sequence(kvm, req, kv_handle, use_v2) - if state == LlmRequestState.DISAGG_TRANS_ERROR: - raise _TransferError("gen transfer reported DISAGG_TRANS_ERROR") return ok @@ -530,13 +783,18 @@ def main(): os.makedirs(status_dir, exist_ok=True) status_path = os.path.join(status_dir, f"sweep{sweep}_{role}.jsonl") status_f = open(status_path, "a") if is_leader else None + timeout_s = int(cfg["run"]["timeout_per_cell_s"]) + if timeout_s <= TRANSFER_TIMEOUT_GRACE_SECONDS: + raise ValueError( + "run.timeout_per_cell_s must be greater than " + f"{TRANSFER_TIMEOUT_GRACE_SECONDS}s to leave ownership-handshake headroom" + ) - # ZMQ leader channel. A REQ/REP handshake interrupted by a timeout leaves the - # socket unusable, so we reset it after any error to isolate the blast radius - # to a single (combination, request_length) cell. + # ZMQ leader channel. Transfer-ownership failures hard-abort this sweep, so + # an interrupted REQ/REP exchange is never reused in-process. zmq_ctx = None zmq_sock = None - rcv_timeout_ms = cfg["run"]["timeout_per_cell_s"] * 1000 + rcv_timeout_ms = timeout_s * 1000 def open_sock(): if not is_leader: @@ -569,16 +827,6 @@ def open_sock(): s.connect(f"tcp://{ctx_node}:{zmq_port}") return s - def reset_sock(): - nonlocal zmq_sock - if not is_leader: - return - try: - zmq_sock.close(linger=0) - except Exception: # noqa: BLE001 - pass - zmq_sock = open_sock() - zmq_sock = open_sock() cases = build_cases(cfg) @@ -611,7 +859,72 @@ def record(combination_idx, reqlen_idx, status, reason=""): ) status_f.flush() - timeout_s = cfg["run"]["timeout_per_cell_s"] + def record_remaining( + combination_idx: int, + reqlen_idx: Optional[int], + status: str, + reason: str, + ) -> None: + """Record every unrun cell before a fatal sweep abort.""" + for remaining_ci in range(combination_idx, len(cases)): + start_li = reqlen_idx if remaining_ci == combination_idx else 0 + if start_li is None: + start_li = 0 + for remaining_li in range(start_li, len(req_lens)): + cell_reason = ( + reason + if remaining_ci == combination_idx + and (reqlen_idx is None or remaining_li == reqlen_idx) + else f"skipped after fatal transfer outcome: {reason}" + ) + record(remaining_ci, remaining_li, status, cell_reason) + + def flush_status() -> None: + if status_f is None or status_f.closed: + return + status_f.flush() + try: + os.fsync(status_f.fileno()) + except OSError: + pass + status_f.close() + + def hard_abort_sweep( + combination_idx: int, + reqlen_idx: Optional[int], + status: str, + reason: str, + *, + coordinated: bool = False, + ) -> None: + """Record the failure and exit without deregistering live transfer memory.""" + try: + signal.alarm(0) + cancel_watchdog() + except Exception: # noqa: BLE001 - abort must still happen + pass + if is_leader: + try: + record_remaining(combination_idx, reqlen_idx, status, reason) + flush_status() + except Exception: # noqa: BLE001 - abort must still happen + pass + if coordinated: + # All ranks reach this path after an intra-role consensus/broadcast. + # The bounded collective prevents a nonleader from aborting before + # the leader's status file is durable, without trusting consensus + # in rank-local timeout/unknown-exception paths. + _coordinate_abort_after_leader_flush(comm) + elif not is_leader: + # Rank-local alarms/exceptions cannot safely enter a collective. + # Give the leader's independently armed deadline a short chance to + # persist the role verdict before this rank aborts the srun step. + time.sleep(ABORT_COORDINATION_TIMEOUT_SECONDS) + try: + sys.stdout.flush() + sys.stderr.flush() + finally: + _hard_abort_process(comm) # In-process hang detector (TensorRT-LLM's HangDetector): a side thread runs # an asyncio timer reset per cell; on expiry it dumps all thread stacks @@ -639,19 +952,27 @@ def record(combination_idx, reqlen_idx, status, reason=""): def _on_hang(): ci = hang_cell["ci"] - targets = range(len(req_lens)) if hang_cell["li"] is None else [hang_cell["li"]] - for li in targets: - record( - ci, - li, - "TIMEOUT", - f"hang detected during {hang_cell['what']} (>{watchdog_deadline}s)", + reason = f"hang detected during {hang_cell['what']} (>{watchdog_deadline}s)" + if is_leader: + try: + record_remaining(ci, hang_cell["li"], "TIMEOUT", reason) + flush_status() + except Exception: # noqa: BLE001 - SIGKILL must still happen + pass + else: + # All ranks arm the same cell deadline. Give the leader's watchdog + # a bounded opportunity to persist the role status before this + # rank's bad exit causes srun to terminate the whole step. + time.sleep(ABORT_COORDINATION_TIMEOUT_SECONDS) + try: + sys.stderr.write( + f"[{role} rank={rank}] WATCHDOG_KILL {hang_cell['what']} " + f"ci={ci} li={hang_cell['li']}\n" ) - sys.stderr.write( - f"[{role} rank={rank}] WATCHDOG_KILL {hang_cell['what']} ci={ci} li={hang_cell['li']}\n" - ) - sys.stderr.flush() - os.kill(os.getpid(), signal.SIGKILL) + sys.stdout.flush() + sys.stderr.flush() + finally: + os.kill(os.getpid(), signal.SIGKILL) hang_detector = HangDetector(timeout=watchdog_deadline, on_detected=_on_hang) hang_detector.start() @@ -686,6 +1007,10 @@ def arm_watchdog(combination_idx, reqlen_idx, what): # much slower host-staged tcp, so enable bounce for cross-node # transfers inside an NVLink domain. kv_cache_bounce_size_mb=int(cfg["kv_cache"].get("bounce_size_mb", 0)), + # For the Python sender, leave deterministic headroom for the + # signal handler and final CTX/GEN ownership handshake after its + # request deadline. The cell alarm remains the C++ path's bound. + kv_transfer_timeout_ms=(timeout_s - TRANSFER_TIMEOUT_GRACE_SECONDS) * 1000, ) # Build the cache manager + transceiver ONCE per case (the manager is @@ -701,6 +1026,17 @@ def arm_watchdog(combination_idx, reqlen_idx, what): ) signal.alarm(0) cancel_watchdog() + except _Timeout: + reason = f"setup exceeded {timeout_s}s with ownership state unknown" + try: + print( + f"[{role} rank={rank}] SETUP TIMEOUT {case['label']}", + file=sys.stderr, + flush=True, + ) + except Exception: # noqa: BLE001 - abort must still happen + pass + hard_abort_sweep(ci, None, "TIMEOUT", reason) except Exception as e: # noqa: BLE001 - setup failed for the whole case signal.alarm(0) cancel_watchdog() @@ -732,7 +1068,6 @@ def arm_watchdog(combination_idx, reqlen_idx, what): torch.cuda.empty_cache() continue - case_timed_out = False for li, req_len in enumerate(req_lens): try: signal.alarm(timeout_s) @@ -758,20 +1093,34 @@ def arm_watchdog(combination_idx, reqlen_idx, what): signal.alarm(0) cancel_watchdog() record(ci, li, "PASS" if (role != "gen" or all_ok) else "MISMATCH") - except _Timeout: - signal.alarm(0) - cancel_watchdog() - record(ci, li, "TIMEOUT", f"exceeded {timeout_s}s") - for remaining_li in range(li + 1, len(req_lens)): - record(ci, remaining_li, "TIMEOUT", "skipped after timeout in earlier req_len") - print( - f"[{role} rank={rank}] TIMEOUT {case['label']} req_len={req_len}", - file=sys.stderr, - flush=True, + except _FatalTransferError as e: + try: + print( + f"[{role} rank={rank}] FATAL {case['label']} req_len={req_len}: {e!r}", + file=sys.stderr, + flush=True, + ) + except Exception: # noqa: BLE001 - abort must still happen + pass + hard_abort_sweep( + ci, + li, + "TRANSFER_ERROR", + repr(e), + coordinated=True, ) - reset_sock() - case_timed_out = True - except Exception as e: # noqa: BLE001 - report any transceiver error + except _Timeout: + reason = f"exceeded {timeout_s}s with transfer quiescence unproven" + try: + print( + f"[{role} rank={rank}] TIMEOUT {case['label']} req_len={req_len}", + file=sys.stderr, + flush=True, + ) + except Exception: # noqa: BLE001 - abort must still happen + pass + hard_abort_sweep(ci, li, "TIMEOUT", reason) + except _TransferError as e: signal.alarm(0) cancel_watchdog() record(ci, li, "TRANSFER_ERROR", repr(e)) @@ -780,9 +1129,16 @@ def arm_watchdog(combination_idx, reqlen_idx, what): file=sys.stderr, flush=True, ) - any_timed_out = comm.allreduce(1 if case_timed_out else 0, op=MPI.MAX) - if any_timed_out: - break + except Exception as e: # noqa: BLE001 - unknown ownership state is fatal + try: + print( + f"[{role} rank={rank}] UNEXPECTED {case['label']} req_len={req_len}: {e!r}", + file=sys.stderr, + flush=True, + ) + except Exception: # noqa: BLE001 - abort must still happen + pass + hard_abort_sweep(ci, li, "TRANSFER_ERROR", repr(e)) # Tear down the case's transceiver and preserve its C++ CSVs. if hasattr(xcvr, "shutdown"): diff --git a/jenkins/scripts/perf/local/submit.py b/jenkins/scripts/perf/local/submit.py index 8e571eb2f890..e7a13536e203 100755 --- a/jenkins/scripts/perf/local/submit.py +++ b/jenkins/scripts/perf/local/submit.py @@ -957,6 +957,7 @@ def main(): # of the real worker steps; enable policy and timeouts come from # precheck_config (single owner). pcfg = _import_precheck_config(llm_src) + precheck_enabled = pcfg.precheck_enabled(config) script_prefix_lines.extend( pcfg.precheck_prefix_lines( config, @@ -967,16 +968,16 @@ def main(): hardware_config.get("gpus_per_ctx_server", 0) or 0, hardware_config.get("gpus_per_gen_server", 0) or 0, ), + llm_models_root=args.llm_models_root if precheck_enabled else None, ) ) # Add srun args for disagg srun_args_lines.extend( - [ - "--container-env=DISAGG_SERVING_TYPE", - "--container-env=pytestCommand", - ] + ["--container-env=DISAGG_SERVING_TYPE", "--container-env=pytestCommand"] ) + if precheck_enabled: + srun_args_lines.append("--container-env=LLM_MODELS_ROOT") else: worker_env_vars = ( f"TLLM_PROFILE_START_STOP='{tllm_profile_start_stop}' " diff --git a/jenkins/scripts/perf/submit.py b/jenkins/scripts/perf/submit.py index f1c54fb13690..21707e43e654 100755 --- a/jenkins/scripts/perf/submit.py +++ b/jenkins/scripts/perf/submit.py @@ -649,6 +649,57 @@ def get_test_output_dir(script_prefix_lines, test_case_name): return os.path.join(output_dir, test_case_name) if test_case_name else output_dir +class _PytestCommandEnvMissing(ValueError): + """A valid pytestCommand does not provide the requested leading variable.""" + + +def extract_pytest_command_env(script_prefix_lines, name): + """Read a leading environment assignment from the exported pytest command.""" + line = next((ln for ln in script_prefix_lines if "export pytestCommand=" in ln), None) + if line is None: + raise ValueError("launch prefix does not export pytestCommand") + try: + outer_tokens = shlex.split(line) + except ValueError as e: + raise ValueError(f"cannot parse exported pytestCommand: {e}") from e + command_assignment = next( + (token for token in outer_tokens if token.startswith("pytestCommand=")), None + ) + if command_assignment is None: + raise ValueError("launch prefix has a malformed pytestCommand export") + command = command_assignment.partition("=")[2] + try: + command_tokens = shlex.split(command) + except ValueError as e: + raise ValueError(f"cannot parse pytestCommand payload: {e}") from e + for token in command_tokens: + if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*=", token): + break + key, value = token.split("=", 1) + if key == name: + return value + raise _PytestCommandEnvMissing( + f"pytestCommand does not set leading environment variable {name}" + ) + + +def _resolve_llm_models_root(script_prefix_lines): + """Resolve the precheck model root from pytestCommand or the submitter env.""" + try: + return extract_pytest_command_env(script_prefix_lines, "LLM_MODELS_ROOT") + except _PytestCommandEnvMissing as e: + fallback = os.environ.get("LLM_MODELS_ROOT") + if fallback: + return fallback + # Fail closed when the precheck is enabled: without the model root it + # cannot reproduce serving's KV shape and model-specific defaults, so + # disabling the gate here would silently run an unvalidated workload. + raise ValueError( + f"{e}; LLM_MODELS_ROOT is also absent from the submitter environment " + "(pytestCommand is assembled by getPytestBaseCommandLine in L0_Test.groovy)" + ) from e + + def remove_whitespace_lines(lines): return [line.strip() for line in lines if line.strip()] @@ -852,6 +903,10 @@ def main(): # Enable/kill-switch policy and timeouts live in precheck_config # (single owner, shared with the local flow). pcfg = _import_precheck_config(args.llm_src) + precheck_enabled = pcfg.precheck_enabled(config) + llm_models_root = ( + _resolve_llm_models_root(script_prefix_lines) if precheck_enabled else None + ) script_prefix_lines.extend( pcfg.precheck_prefix_lines( config, @@ -863,14 +918,14 @@ def main(): hardware_config["gpus_per_gen_server"], ), stage_name=args.stage_name, + llm_models_root=llm_models_root, ) ) srun_args_lines.extend( - [ - "--container-env=DISAGG_SERVING_TYPE", - "--container-env=pytestCommand", - ] + ["--container-env=DISAGG_SERVING_TYPE", "--container-env=pytestCommand"] ) + if precheck_enabled: + srun_args_lines.append("--container-env=LLM_MODELS_ROOT") script_prefix_lines = remove_whitespace_lines(script_prefix_lines) script_prefix = "\n".join(script_prefix_lines) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 673d088e8443..516809e71034 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -83,6 +83,14 @@ # Number of worker threads for KV transfer queues (default: 1) KV_TRANSFER_NUM_THREADS = int(os.environ.get("TRTLLM_KV_TRANSFER_NUM_THREADS", "1")) +# Keep standalone TxSession waits responsive to cancellation even when callers +# do not configure a sender-future wait slice. +_FALLBACK_TX_WAIT_SLICE_S = 1.0 +# Keep direct/standalone TxSession block-all waits finite when the caller omits +# an overall deadline. KvCacheTransceiverV2 requires a configured transfer +# timeout before it creates either sender or receiver sessions. +_FALLBACK_TX_OVERALL_TIMEOUT_S = 60.0 + @dataclass class RecvReqInfo: @@ -1231,12 +1239,15 @@ def __init__( timeout_s: Optional[float] = None, prompt_len: Optional[int] = None, beam_width: int = 1, + overall_timeout_s: Optional[float] = None, ): super().__init__( sender, SessionArgsBase(params, prompt_len=prompt_len, beam_width=beam_width), ) self._timeout_s = timeout_s + self._overall_timeout_s = overall_timeout_s + self._deadline_monotonic_s: Optional[float] = None self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST self._sender: Sender # narrow base class type for Pylance self.request_id = request_id @@ -1287,6 +1298,11 @@ def send(self, slice: KVSlice) -> None: if self.transfer_start_time is None: self.transfer_start_time = tensorrt_llm.bindings.global_steady_clock_now() with self.lock: + if not self.kv_tasks: + overall_timeout_s = self._overall_timeout_s + if overall_timeout_s is None or overall_timeout_s <= 0: + overall_timeout_s = _FALLBACK_TX_OVERALL_TIMEOUT_S + self._deadline_monotonic_s = time.monotonic() + overall_timeout_s params = self._base_args.params slice_id = len(self.kv_tasks) task = KVSendTask( @@ -1363,11 +1379,15 @@ def has_transferring_tasks(self) -> bool: def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: """Poll or block until KV (and optionally aux) transfer finishes. - With blocking=True (default): waits up to _timeout_s for each task. + With blocking=True (default): retries bounded wait slices until the + transfer finishes or its overall deadline expires. Deadline expiry is + nonterminal: callers must retain the session and its KV pages because + peer writes may still be active. Errors and cancellation remain + terminal, but likewise do not prove that peer writes quiesced. With blocking=False: polls non-blockingly; returns None if any KV task or aux is not yet done. """ - if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED): + if self.has_failed(): return WaitResult.FAILED if not self.kv_tasks: return None @@ -1389,17 +1409,77 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: return None return WaitResult.COMPLETED + # send() normally anchors this once, at the first dispatched KV task. + # Keep direct/internal TxSession construction bounded as well, and never + # reset the deadline on a later block-all call. + if self._deadline_monotonic_s is None: + overall_timeout_s = self._overall_timeout_s + if overall_timeout_s is None or overall_timeout_s <= 0: + overall_timeout_s = _FALLBACK_TX_OVERALL_TIMEOUT_S + with self.lock: + if self._deadline_monotonic_s is None: + self._deadline_monotonic_s = time.monotonic() + overall_timeout_s + + # ``_timeout_s`` bounds one scheduler wait slice. The separate absolute + # deadline is shared by every KV task and aux; it is never reset by a + # later wait_complete() call. + wait_slice_s = self._timeout_s + if wait_slice_s is None or wait_slice_s <= 0: + wait_slice_s = _FALLBACK_TX_WAIT_SLICE_S + + def wait_for_task(task: SendTaskBase) -> WaitResult: + while True: + # A task/session terminal state observed at the deadline + # boundary takes precedence over TIMEOUT. + if self.has_failed(): + return WaitResult.FAILED + if task.status == TaskStatus.TRANSFERRED: + return WaitResult.COMPLETED + + remaining_s = None + if self._deadline_monotonic_s is not None: + remaining_s = self._deadline_monotonic_s - time.monotonic() + if remaining_s <= 0: + # The worker can publish terminal state between the + # checks above and the clock read. Preserve boundary + # precedence before classifying this as a timeout. + if self.has_failed(): + return WaitResult.FAILED + if task.status == TaskStatus.TRANSFERRED: + return WaitResult.COMPLETED + return WaitResult.TIMEOUT + timeout_s = wait_slice_s if remaining_s is None else min(wait_slice_s, remaining_s) + task.wait(timeout=timeout_s) + + # A bounded slice keeps cancellation and sibling failure observable. for task in self.kv_tasks: - if not task.wait(timeout=self._timeout_s): - return WaitResult.TIMEOUT - if task.status == TaskStatus.ERROR: - return WaitResult.FAILED - if self._need_aux and self.aux_task is not None: - if not self.aux_task.wait(timeout=self._timeout_s): - return WaitResult.TIMEOUT - if self.aux_task.status == TaskStatus.ERROR: + result = wait_for_task(task) + if result != WaitResult.COMPLETED: + return result + if self._need_aux: + if self.aux_task is None: + # _finalize_send() installs the aux task synchronously before + # publishing the request to _send_reqs. Once every KV task is + # terminal, a missing required aux task is an invariant error, + # not an asynchronously pending transfer. + with self.lock: + if self._terminal_status not in ( + SessionStatus.ERROR, + SessionStatus.CANCELLED, + ): + self._exception = RuntimeError( + "required auxiliary transfer was not dispatched" + ) + self._terminal_status = SessionStatus.ERROR return WaitResult.FAILED - return WaitResult.COMPLETED + result = wait_for_task(self.aux_task) + if result != WaitResult.COMPLETED: + return result + return ( + WaitResult.FAILED + if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED) + else WaitResult.COMPLETED + ) def set_exception(self, reason: str = ""): msg = f"TxSession {self.disagg_request_id} exception" @@ -2313,6 +2393,7 @@ class TransferWorkerConfig: tx_timeout_s: Optional[float] = None rx_timeout_s: Optional[float] = None bounce: Optional["Config"] = None + tx_overall_timeout_s: Optional[float] = None class TransferWorker: @@ -2347,6 +2428,7 @@ def create_tx_session(self, request: LlmRequest) -> TxSession: timeout_s=self._config.tx_timeout_s, prompt_len=request.prompt_len, beam_width=request.py_beam_width, + overall_timeout_s=self._config.tx_overall_timeout_s, ) def create_rx_session(self, request: LlmRequest) -> RxSession: diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 0d7e3429424b..fd8db9d0e9cc 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -85,10 +85,18 @@ def __init__( self._kv_cache_manager = kv_cache_manager self._mapping = mapping self.kv_transfer_timeout_ms = cache_transceiver_config.kv_transfer_timeout_ms + if self.kv_transfer_timeout_ms is None: + raise ValueError("KvCacheTransceiverV2 requires a finite kv_transfer_timeout_ms") self.kv_transfer_poll_interval_ms = cache_transceiver_config.kv_transfer_poll_interval_ms self._sender_future_timeout_ms = ( cache_transceiver_config.kv_transfer_sender_future_timeout_ms ) + transfer_timeout_s = self.kv_transfer_timeout_ms / 1000.0 + sender_wait_slice_s = ( + self._sender_future_timeout_ms / 1000.0 + if self._sender_future_timeout_ms is not None + else None + ) self._check_compatible() self._reuse_adapter: CacheReuseAdapter = create_cache_reuse_adapter(kv_cache_manager) @@ -115,8 +123,9 @@ def __init__( # can be in-flight simultaneously. AuxBuffer holds only small CPU metadata, so a # large multiplier is cheap. max_concurrent_sessions=max(1, int(kv_cache_manager.max_batch_size)) * 20000, - tx_timeout_s=self._sender_future_timeout_ms / 1000.0, - rx_timeout_s=self.kv_transfer_timeout_ms / 1000.0, + tx_timeout_s=sender_wait_slice_s, + tx_overall_timeout_s=transfer_timeout_s, + rx_timeout_s=transfer_timeout_s, # Size 0 turns bounce off; the per-transfer size gates are internal (tuned via # env: TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS for plain-KV payloads, # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES for recurrent-state payloads). @@ -490,8 +499,9 @@ def _gen_consensus_outcome(self, to_process, cancelled, failed, completed): to_process, cancelled, failed, completed, self._gen_allgather, self._gen_need_sync ) - def _ctx_consensus_outcome(self, to_process, cancelled, failed, completed, timed_out): - # TP first, then PP. timed_out is local-only (back-off signal). + def _ctx_consensus_outcome(self, to_process, cancelled, failed, completed): + # TP first, then PP. A local timeout remains nonterminal, so it is + # represented by the absence of that request from completed. c, f, d = self._consensus_outcome( to_process, cancelled, @@ -503,7 +513,7 @@ def _ctx_consensus_outcome(self, to_process, cancelled, failed, completed, timed if self._ctx_need_pp_sync: pp_allgather: Callable = getattr(self._dist, "pp_allgather") c, f, d = self._consensus_outcome(to_process, c, f, d, pp_allgather, True) - return c, f, d, timed_out + return c, f, d def _sync_transfer_timing(self, reqs: list): """Allgather timing for a batch of completed requests in one collective. @@ -722,7 +732,7 @@ def check_context_transfer_status( block_all, ) - completed, timed_out, failed, cancelled = [], [], [], [] + completed, failed, cancelled = [], [], [] for rid in to_process: session = self._send_sessions[rid] result = session.wait_complete(blocking=block_all) @@ -734,16 +744,17 @@ def check_context_transfer_status( continue elif result == WaitResult.TIMEOUT: logger.warning( - f"TxSession rid={session.disagg_request_id} timed out after {self._sender_future_timeout_ms}ms" + f"TxSession rid={session.disagg_request_id} exceeded " + f"kv_transfer_timeout_ms={self.kv_transfer_timeout_ms}ms; " + "keeping it in progress" ) - timed_out.append(rid) else: logger.warning(f"TxSession rid={session.disagg_request_id} failed") failed.append(rid) # All ranks must agree on per-rid outcome to avoid req.state divergence. - cancelled, failed, completed, timed_out = self._ctx_consensus_outcome( - to_process, cancelled, failed, completed, timed_out + cancelled, failed, completed = self._ctx_consensus_outcome( + to_process, cancelled, failed, completed ) for rid in cancelled: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 01696cdb6068..4379cc5172a9 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4228,14 +4228,18 @@ class CacheTransceiverConfig(StrictBaseModel, PybindMirror): kv_transfer_timeout_ms: Optional[PositiveInt] = Field( default=60000, description= - "Timeout in milliseconds for KV cache transfer. Requests exceeding this timeout will be cancelled." - ) + "KV cache transfer timeout in milliseconds. Blocking sender waits use " + "it as an absolute deadline; blocking receive task waits use it per " + "task. The Python V2 transceiver requires a finite value; None remains " + "available to other runtimes. It is distinct from the sender future " + "wait slice.") kv_transfer_sender_future_timeout_ms: Optional[PositiveInt] = Field( default=1000, description= - "Timeout in milliseconds to wait for the sender future to be ready when scheduled batch size is 0. This allows the request to be eventually cancelled by the user or because of kv_transfer_timeout_ms" - ) + "Duration in milliseconds of each bounded sender future wait slice " + "while polling KV transfer completion. It does not set the overall " + "transfer deadline.") kv_transfer_poll_interval_ms: Optional[PositiveInt] = Field( default=5000, diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md index 0722e0bf1618..027a4852774c 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md @@ -16,7 +16,7 @@ starts. | Same UCX env vars (incl. the `unset UCX_TLS` cases) | `jenkins/scripts/perf/submit.py` builds the precheck commands from the **same** `ucx_tls_cmd` + `$CTX/GEN_WORKER_ENV_VARS` strings as the worker steps; `slurm_precheck_run.sh` sources the same `slurm_env_setup.sh` (the `UCX_TLS=tcp` fixup) as `slurm_run.sh`. | | Same instance count / parallelism | One precheck `srun` per ctx/gen server with the same `-N/--ntasks/--ntasks-per-node/--mpi=pmix` and the same node slices (`-w`) as the real server steps (`slurm_launch_draft.sh`). TP/PP/CP/attention-DP come from the same `worker_config`. | | Same transceiver config | `CacheTransceiverConfig(**yaml["worker_config"][role]["cache_transceiver_config"])` — the yaml block is passed through verbatim (backend, `max_tokens_in_buffer`, timeouts, ...). | -| Same KV cache manager version + transceiver runtime | Explicit per-side `kv_cache_config.use_kv_cache_manager_v2` wins; absent means "auto" and resolves via `get_preferred_kv_cache_manager_version()`, and `transceiver_runtime: auto` resolves via `get_preferred_transceiver_runtime()` (NIXL-gated) — both through the same llm_utils code serving uses. V2 requires the Python transceiver (the C++ one only supports V1); a V2+CPP combination fails fast with INIT_ERROR. | +| Same KV cache manager version + transceiver runtime | The launch generator forwards the real test's `LLM_MODELS_ROOT` into every precheck process. Explicit per-side `kv_cache_config.use_kv_cache_manager_v2` wins; absent means "auto" and requires a registered model class, then resolves via `get_preferred_kv_cache_manager_version()`. `transceiver_runtime: auto` resolves via `get_preferred_transceiver_runtime()` (NIXL-gated) — both through the same llm_utils code serving uses. An unresolved manager-version `auto` setting fails with INIT_ERROR instead of silently assuming V1. V2 requires the Python transceiver (the C++ one only supports V1); a V2+CPP combination also fails fast. | Asymmetric layouts (ctx dep4 → gen dep16, ctx pp8 → gen tp32, ...) are supported: data is seeded per (request, **global** layer) and constant along @@ -58,6 +58,13 @@ so the precheck does too. Later reps run under the tight `wave_timeout_s`, which is what actually catches hangs. Set `PRECHECK_DEBUG=1` in the worker env to raise the C++/Python transceiver log levels when debugging a stall. +The Python sender's `kv_transfer_timeout_ms` is also a real request deadline. +If block-all returns without every expected request completed—or any rank +reports failure, cancellation, or unsynchronized receive work—the precheck +retains the KV pages, persists an ownership-fatal verdict, and hard-aborts the +instance without running transceiver/cache-manager finalizers. Pages are freed +only after exact completion on every rank (plus CUDA synchronization on gen). + ## Failure output The sbatch log gets a summary block: per-instance verdicts @@ -79,13 +86,13 @@ csv/ctx_/_.csv # Python transceiver per-task perf ```bash # Inspect what a yaml resolves to (no GPU needed): -python3 run_precheck.py --role gen --server-idx 0 --dry-run \ +LLM_MODELS_ROOT='' python3 run_precheck.py --role gen --server-idx 0 --dry-run \ --config ../disaggregated/.yaml --work-dir /tmp/ct --llm-src # On a SLURM allocation: one srun per instance, e.g. ctx dep4 + gen dep8: -srun -N1 --ntasks=4 --mpi=pmix python3 run_precheck.py --role ctx --server-idx 0 \ - --config --work-dir --llm-src & -srun -N2 --ntasks=8 --mpi=pmix python3 run_precheck.py --role gen --server-idx 0 \ +LLM_MODELS_ROOT='' srun -N1 --ntasks=4 --mpi=pmix python3 run_precheck.py \ + --role ctx --server-idx 0 --config --work-dir --llm-src & +LLM_MODELS_ROOT='' srun -N2 --ntasks=8 --mpi=pmix python3 run_precheck.py --role gen --server-idx 0 \ --config --work-dir --llm-src & wait ``` diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py index 53b79919a9cf..b5fb567daf43 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -26,6 +26,9 @@ import json import os +import shlex +from collections.abc import Mapping +from typing import Any # Optional per-yaml overrides live under a `cache_transceiver_precheck:` block. PRECHECK_DEFAULTS = { @@ -61,8 +64,9 @@ "verify_data": True, } -# Fallback KV shape when the model directory cannot be resolved: the precheck -# still exercises the exact network path, just with a synthetic cache shape. +# Fallback KV shape for dry-runs and explicitly selected manager versions when +# the model directory cannot be resolved. Manager-version "auto" resolution +# fails fast instead of silently pairing this shape with V1. FALLBACK_KV_SHAPE = { "num_layers": 32, "num_kv_heads": 8, @@ -127,17 +131,8 @@ def default_step_timeout_s(max_world): return 900 + wireup_timeout_s(max_world) -def precheck_prefix_lines( - cfg, benchmark_mode, config_path_expr, ucx_tls_cmd, max_world, stage_name="" -): - """Launch-script export lines wiring the precheck gate. - - Single owner of the enable/kill-switch policy, the step-timeout default, - and the export names the gate consumes — shared by - jenkins/scripts/perf/submit.py and jenkins/scripts/perf/local/submit.py. - `config_path_expr` and the env-var references are launch-script-side - expressions ($llmSrcNode etc.), expanded at sbatch runtime. - """ +def precheck_enabled(cfg): + """Resolve the shared yaml/environment enable policy.""" knobs = cfg.get("cache_transceiver_precheck", {}) or {} # Off by default until the gate is validated on the post-merge stages # (the precheck is a launch-script gate, not a pytest case, so it cannot @@ -160,6 +155,42 @@ def precheck_prefix_lines( ) else: enabled = bool(knobs.get("enabled", False)) + return enabled + + +def precheck_prefix_lines( + cfg: Mapping[str, Any], + benchmark_mode: str, + config_path_expr: str, + ucx_tls_cmd: str, + max_world: int, + stage_name: str = "", + llm_models_root: str | None = None, +) -> list[str]: + """Launch-script export lines wiring the precheck gate. + + Single owner of the enable/kill-switch policy, the step-timeout default, + and the export names the gate consumes — shared by + jenkins/scripts/perf/submit.py and jenkins/scripts/perf/local/submit.py. + + Args: + cfg: Parsed perf-sanity YAML configuration. + benchmark_mode: Precheck benchmark mode passed to the driver. + config_path_expr: Launch-script expression resolving the YAML path, + expanded at sbatch runtime along with its environment references. + ucx_tls_cmd: Shell prefix selecting the UCX transports. + max_world: Largest role world size, used to derive the step timeout. + stage_name: Optional stage name for the synthetic JUnit report. + llm_models_root: Model-root path exported when the precheck is enabled. + + Returns: + Generated launch-script export lines shared by both submit modules. + + Raises: + ValueError: If the precheck is enabled without a nonempty model root. + """ + knobs = cfg.get("cache_transceiver_precheck", {}) or {} + enabled = precheck_enabled(cfg) cmd = ( "python3 $llmSrcNode/tests/scripts/perf-sanity/cache_transceiver_precheck/" f"run_precheck.py --config {config_path_expr} " @@ -178,6 +209,12 @@ def precheck_prefix_lines( f'export pytestCommandCTXPrecheck="{ucx_tls_cmd} $CTX_WORKER_ENV_VARS {cmd} --role ctx"', f'export pytestCommandGENPrecheck="{ucx_tls_cmd} $GEN_WORKER_ENV_VARS {cmd} --role gen"', ] + if enabled: + if not llm_models_root: + raise ValueError("enabled cache-transceiver precheck requires LLM_MODELS_ROOT") + # Keep this as a top-level assignment. shlex.quote() is not safe when + # nested inside the double-quoted pytestCommand exports below. + lines.insert(0, f"export LLM_MODELS_ROOT={shlex.quote(llm_models_root)}") if stage_name: # Suite name for the synthetic junit xml the gate writes on failure # (absent -> the gate falls back to $SLURM_JOB_NAME). diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index e8dcea152fc9..8b9e1eda4846 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -85,6 +85,7 @@ # requests of a session; the peer stride only separates sessions, which talk # to distinct agents and therefore cannot alias tags with each other. RID_PEER_STRIDE = 1 << 24 +ABORT_COORDINATION_TIMEOUT_S = 2.0 class _Timeout(Exception): @@ -95,6 +96,10 @@ class _TransferError(Exception): pass +class _FatalTransferError(_TransferError): + """Transfer quiescence is unproven; finalizers must not release its memory.""" + + class _PeerAbort(Exception): pass @@ -103,6 +108,29 @@ def _alarm_handler(signum, frame): raise _Timeout() +def _coordinate_abort_after_leader_flush(comm): + """Bounded best-effort rendezvous after the leader writes its verdict.""" + try: + request = comm.Ibarrier() + deadline = time.monotonic() + ABORT_COORDINATION_TIMEOUT_S + while not request.Test(): + if time.monotonic() >= deadline: + return + time.sleep(0.01) + except Exception: # noqa: BLE001 - abort remains the mandatory fallback + return + + +def _hard_abort_process(comm): + """Terminate without running transceiver/KV-manager finalizers.""" + try: + comm.Abort(137) + except Exception: # noqa: BLE001 - SIGKILL is the mandatory fallback + pass + os.kill(os.getpid(), signal.SIGKILL) + raise RuntimeError("SIGKILL unexpectedly returned") + + def make_rid(ctx_idx, gen_idx, num_ctx, seq): """Unique rid: peer-session base + dense in-session sequence number.""" peer = gen_idx * num_ctx + ctx_idx @@ -214,6 +242,7 @@ def load_internal_apis(): CacheTransceiverConfig, KvCacheConfig, MTPDecodingConfig, + TorchLlmArgs, ) from tensorrt_llm.llmapi.llm_utils import ( _resolve_kv_cache_manager_v2_auto, @@ -242,6 +271,7 @@ def load_internal_apis(): CacheTransceiverConfig=CacheTransceiverConfig, KvCacheConfig=KvCacheConfig, MTPDecodingConfig=MTPDecodingConfig, + TorchLlmArgs=TorchLlmArgs, resolve_kv_cache_manager_v2_auto=_resolve_kv_cache_manager_v2_auto, resolve_transceiver_runtime_auto=_resolve_transceiver_runtime_auto, Mapping=Mapping, @@ -271,28 +301,39 @@ def _pattern_like(shape, dtype, device, seed): return rnd.to(dtype).to(device).expand(nb, kv, heads, tok, dim) -def _request_block_views(kvm, rid): - """Yield (global_layer, buffer, valid_block_indices) for this rank.""" +def _request_block_views(kvm, rid, prompt_len): + """Yield the prompt blocks transferred for this request on this rank.""" + num_prompt_blocks = (prompt_len + kvm.tokens_per_block - 1) // kvm.tokens_per_block for global_layer in kvm.pp_layers: blocks = kvm.get_batch_cache_indices([rid], layer_idx=global_layer)[0] + # V2 may reserve extra KV tokens for speculative decoding. At an exact + # block boundary those tokens allocate an additional page, but the + # transceiver intentionally trims its slice to prompt_len blocks. + # Verify the same payload range instead of the untransferred page. valid = [b for b in blocks if b >= 0] + if len(valid) < num_prompt_blocks: + raise _TransferError( + f"KV under-allocation for rid={rid} layer={global_layer}: " + f"required={num_prompt_blocks} available={len(valid)}" + ) + valid = valid[:num_prompt_blocks] if not valid: continue buf = kvm.get_buffers(global_layer, kv_layout="HND") yield global_layer, buf, valid -def fill_request(kvm, rid): - for global_layer, buf, valid in _request_block_views(kvm, rid): +def fill_request(kvm, rid, prompt_len): + for global_layer, buf, valid in _request_block_views(kvm, rid, prompt_len): shape = (len(valid), *buf.shape[1:]) buf[valid] = _pattern_like(shape, buf.dtype, buf.device, seed_for(rid, global_layer)) -def verify_request(kvm, rid): +def verify_request(kvm, rid, prompt_len): """Returns (ok, detail) comparing received blocks to the expected pattern.""" import torch - for global_layer, buf, valid in _request_block_views(kvm, rid): + for global_layer, buf, valid in _request_block_views(kvm, rid, prompt_len): recv = buf[valid] exp = _pattern_like(recv.shape, recv.dtype, recv.device, seed_for(rid, global_layer)) recv_f, exp_f = recv.float(), exp.float() # fp8 lacks direct compare ops @@ -320,8 +361,8 @@ def _lookup_model_cls(model_dir): def resolve_model_prefs(model_dir, side, cache_cfg): """Mirror serving's model-preference resolution (PR #15823 semantics). - - use_kv_cache_manager_v2 == "auto" (yaml absent): adopt the model - class's get_preferred_kv_cache_manager_version() value, falling back to V1 + - use_kv_cache_manager_v2 == "auto" (yaml absent): require the model + class and adopt its get_preferred_kv_cache_manager_version() value (llm_utils._resolve_kv_cache_manager_v2_auto). - cache_cfg.transceiver_runtime == "auto": adopt model_cls.get_preferred_transceiver_runtime(), NIXL-gated, via the @@ -333,6 +374,12 @@ def resolve_model_prefs(model_dir, side, cache_cfg): api = load_internal_apis() model_cls, hf_view = _lookup_model_cls(model_dir) + setting = side["use_kv_cache_manager_v2"] + if setting == "auto" and model_cls is None: + raise RuntimeError( + "use_kv_cache_manager_v2 is 'auto', but the precheck could not resolve " + f"a registered model class from model_dir={model_dir!r}; refusing to assume V1" + ) # Runtime BEFORE V2, like serving: the V2 resolver's disagg gating reads # cache_cfg.transceiver_runtime and treats an unresolved "auto" as non-PYTHON. @@ -340,31 +387,34 @@ def resolve_model_prefs(model_dir, side, cache_cfg): try: shim = types.SimpleNamespace(cache_transceiver_config=cache_cfg) api.resolve_transceiver_runtime_auto(shim, model_cls, hf_view) - except Exception as e: # noqa: BLE001 - fall back to the create() default (CPP) - print( - f"[precheck] WARNING: transceiver_runtime 'auto' resolution failed " - f"({e!r}); create_kv_cache_transceiver will fall back to CPP", - flush=True, - ) + except Exception as e: # noqa: BLE001 - resolver spans model extension hooks + raise RuntimeError( + "transceiver_runtime 'auto' resolution failed; refusing to validate " + "a runtime that may differ from serving" + ) from e - setting = side["use_kv_cache_manager_v2"] if setting == "auto": try: - # The REAL serving resolver, via the same shim pattern as the - # runtime resolution below -- one owner for the 'auto' semantics. - # cache_transceiver_config feeds the resolver's disagg gating - # (a V2 model preference requires the NIXL Python transceiver). - shim = types.SimpleNamespace( - kv_cache_config=types.SimpleNamespace(use_kv_cache_manager_v2="auto"), - cache_transceiver_config=cache_cfg, - speculative_config=None, - ) - use_v2 = bool(api.resolve_kv_cache_manager_v2_auto(shim, model_cls, hf_view)) - except Exception as e: # noqa: BLE001 - fall back like a missing model - print( - f"[precheck] WARNING: V2 'auto' resolution failed ({e!r}); assuming V1", flush=True - ) - use_v2 = False + parallel = side["parallel"] + llm_args_kwargs = { + "model": model_dir, + "tensor_parallel_size": parallel["tp"], + "pipeline_parallel_size": parallel["pp"], + "context_parallel_size": parallel["cp"], + "kv_cache_config": {"use_kv_cache_manager_v2": setting}, + "cache_transceiver_config": cache_cfg, + } + num_nextn = int(side.get("num_nextn_predict_layers", 0) or 0) + if num_nextn > 0: + llm_args_kwargs["speculative_config"] = api.MTPDecodingConfig( + num_nextn_predict_layers=num_nextn + ) + resolver_args = api.TorchLlmArgs(**llm_args_kwargs) + # Use the real serving arguments so the resolver sees the same + # disaggregation, parallelism, and speculative-decoding inputs. + use_v2 = bool(api.resolve_kv_cache_manager_v2_auto(resolver_args, model_cls, hf_view)) + except Exception as e: # noqa: BLE001 - resolver spans model extension hooks + raise RuntimeError("V2 'auto' resolution failed; refusing to assume V1") from e else: use_v2 = bool(setting) return use_v2 @@ -500,7 +550,8 @@ def free_sequence(kvm, req, use_v2): def _wait_gen_complete(xcvr, req, runtime, llm_request_state): """Block until this gen request's receive finishes (or errors). - PYTHON transceiver: check_gen_transfer_status(None) blocks for all. C++: + The Python transceiver is handled once per wave in gen_run_wave(), where + its returned request IDs can be checked before releasing KV pages. For C++, the int API can return before THIS request completes on a cold link, so poll for a terminal state (bounded by signal.alarm + hang detector). Logs periodic progress so a stalled transfer shows WHICH request is stuck @@ -508,8 +559,7 @@ def _wait_gen_complete(xcvr, req, runtime, llm_request_state): "RDMA write never completed"). """ if runtime == "PYTHON": - xcvr.check_gen_transfer_status(None) - return + raise ValueError("Python generation waves must be checked as a batch") terminal = ( llm_request_state.DISAGG_GENERATION_TRANS_COMPLETE, llm_request_state.DISAGG_TRANS_ERROR, @@ -916,47 +966,60 @@ def ctx_run_wave(self, peer_idx, li, req_len, rep, wave): rid = self._pair_rid(peer_idx, li, rep, pair) req = make_request(True, rid, req_len, self.runtime) add_sequence(self.kvm, req, req_len, self.use_v2) - fill_request(self.kvm, rid) + # Track ownership as soon as allocation succeeds. A later + # setup failure must retain the pages rather than free storage + # that an asynchronously dispatched sender may still read. + reqs[pair] = req + fill_request(self.kvm, rid, req_len) tensorrt_llm.logger.info( f"[ctx{self.server_idx} r{self.rank}] rid={rid} len={req_len}: send START" ) self.xcvr.respond_and_send_async(req) - reqs[pair] = req + except _Timeout: + raise except Exception as e: # noqa: BLE001 - relayed to gen, then raised local_err = e - reason = self._consensus_error(local_err) + try: + reason = self._consensus_error(local_err) - # Params for pair k come from its owning dp rank at pp stage 0 with - # attention DP; without DP every rank sends the same request, and the - # instance leader's params are the ones the real server would return. - if self.side["parallel"]["enable_attention_dp"]: - contributes = self.mapping.pp_rank == 0 - else: - contributes = self.is_leader - contrib = ( - {p: r.context_phase_params for p, r in reqs.items()} - if local_err is None and contributes - else {} - ) - gathered = self.comm.gather(contrib, root=0) - params_by_pair = {} - if self.is_leader: - for d in gathered: - params_by_pair.update(d or {}) - if reason is None: - missing = [p for p in wave if p not in params_by_pair] - if missing: - reason = f"missing context_phase_params for pairs {missing}" - # The missing-params check runs only on the leader (only it holds the - # gathered params). Broadcast the verdict so EVERY rank raises together: - # otherwise the leader raises here while the other ranks return and enter - # ctx_finish_wave's collective, the collective sequence diverges, and the - # step deadlocks until the watchdog SIGKILLs it (misreported as TIMEOUT). - reason = self.comm.bcast(reason, root=0) + # Params for pair k come from its owning dp rank at pp stage 0 with + # attention DP; without DP every rank sends the same request, and the + # instance leader's params are the ones the real server would return. + if self.side["parallel"]["enable_attention_dp"]: + contributes = self.mapping.pp_rank == 0 + else: + contributes = self.is_leader + contrib = ( + {p: r.context_phase_params for p, r in reqs.items()} + if local_err is None and contributes + else {} + ) + gathered = self.comm.gather(contrib, root=0) + params_by_pair = {} + if self.is_leader: + for d in gathered: + params_by_pair.update(d or {}) + if reason is None: + missing = [p for p in wave if p not in params_by_pair] + if missing: + reason = f"missing context_phase_params for pairs {missing}" + # The missing-params check runs only on the leader (only it holds the + # gathered params). Broadcast the verdict so EVERY rank raises together: + # otherwise the leader raises here while the other ranks return and enter + # ctx_finish_wave's collective, the collective sequence diverges, and the + # step deadlocks until the watchdog SIGKILLs it (misreported as TIMEOUT). + reason = self.comm.bcast(reason, root=0) + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - sends may already be in flight + raise _FatalTransferError(f"ctx send ownership consensus failed: {e!r}") from e if reason is not None: - self._free_all(reqs) - raise _TransferError(f"ctx send setup failed: {reason}") + # Send setup can fail asymmetrically after another rank dispatched + # work. A block-all collective is not safe from that state, and + # orderly teardown can deregister memory while a worker still owns + # it. Conservatively abort without running finalizers. + raise _FatalTransferError(f"ctx send setup failed: {reason}") return params_by_pair, reqs def ctx_finish_wave(self, reqs): @@ -966,24 +1029,48 @@ def ctx_finish_wave(self, reqs): t0 = time.monotonic() local_err = None try: - self.xcvr.check_context_transfer_status(None) # block-all + completed, failed = self.xcvr.check_context_transfer_status(None) # block-all + completed_rids = set(completed) + failed_rids = set(failed) + missing = [ + p + for p, req in reqs.items() + if req.py_request_id not in completed_rids | failed_rids + ] + if missing: + raise _TransferError( + f"block-all returned before terminal status for pairs {missing}" + ) + failed_pairs = [p for p, req in reqs.items() if req.py_request_id in failed_rids] + if failed_pairs: + raise _TransferError(f"ctx transfer failed for pairs {failed_pairs}") bad = [ p for p, r in reqs.items() if r.state == self.llm_request_state.DISAGG_TRANS_ERROR ] if bad: - local_err = _TransferError(f"ctx DISAGG_TRANS_ERROR on pairs {bad}") + raise _TransferError(f"ctx DISAGG_TRANS_ERROR on pairs {bad}") + except _Timeout: + raise except Exception as e: # noqa: BLE001 local_err = e - finally: - self._free_all(reqs) - states = {p: str(r.state) for p, r in reqs.items()} - tensorrt_llm.logger.info( - f"[ctx{self.server_idx} r{self.rank}] wave sends finished in " - f"{time.monotonic() - t0:.1f}s states={states}" - ) - reason = self._consensus_error(local_err) + try: + states = {p: str(r.state) for p, r in reqs.items()} + tensorrt_llm.logger.info( + f"[ctx{self.server_idx} r{self.rank}] wave sends finished in " + f"{time.monotonic() - t0:.1f}s states={states}" + ) + reason = self._consensus_error(local_err) + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - release proof is incomplete + raise _FatalTransferError(f"ctx transfer ownership consensus failed: {e!r}") from e if reason is not None: - raise _TransferError(f"ctx transfer failed: {reason}") + # Failed/cancelled/missing results do not prove that every NIXL + # reader has relinquished the source pages. Do not proceed to the + # ordinary shutdown path, which deregisters those pages. + raise _FatalTransferError(f"ctx transfer failed: {reason}") + # Every rank proved exact completion before any rank recycles pages. + self._free_all(reqs) def gen_run_wave(self, peer_idx, li, req_len, rep, wave, params_by_pair): """Receive + verify owned pairs. Returns (ok, mismatch_detail). @@ -1004,23 +1091,49 @@ def gen_run_wave(self, peer_idx, li, req_len, rep, wave, params_by_pair): False, rid, req_len, self.runtime, ctx_params=params_by_pair[pair] ) add_sequence(self.kvm, req, req_len, self.use_v2) + # Track every allocated sequence before receive dispatch. On + # setup failure, retain all pages and bypass normal teardown. + reqs[pair] = req tensorrt_llm.logger.info( f"[gen{self.server_idx} r{self.rank}] rid={rid} len={req_len}: recv START" ) self.xcvr.request_and_receive_async(req) - reqs[pair] = req + except _Timeout: + raise except Exception as e: # noqa: BLE001 local_err = e - reason = self._consensus_error(local_err) + try: + reason = self._consensus_error(local_err) + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - ctx sender is already live + raise _FatalTransferError(f"gen receive setup consensus failed: {e!r}") from e if reason is not None: - self._free_all(reqs) - raise _TransferError(f"gen receive setup failed: {reason}") + # The context sender is already live, and receive setup can fail + # after partially dispatching local work. Quiescence is unknown. + raise _FatalTransferError(f"gen receive setup failed: {reason}") mismatch = "" t0 = time.monotonic() + transfer_error = None try: - for pair, req in reqs.items(): - _wait_gen_complete(self.xcvr, req, self.runtime, self.llm_request_state) + if self.runtime == "PYTHON": + completed, failed, cancelled = self.xcvr.check_gen_transfer_status(None) + completed_rids = set(completed) + failed_rids = set(failed) + cancelled_rids = {req.py_request_id for req in cancelled} + expected_rids = {req.py_request_id for req in reqs.values()} + missing_rids = expected_rids - completed_rids - failed_rids - cancelled_rids + if failed_rids or cancelled_rids or missing_rids: + raise _TransferError( + "Python gen block-all did not complete every request: " + f"failed={sorted(failed_rids)} " + f"cancelled={sorted(cancelled_rids)} " + f"missing={sorted(missing_rids)}" + ) + else: + for req in reqs.values(): + _wait_gen_complete(self.xcvr, req, self.runtime, self.llm_request_state) if reqs: tensorrt_llm.logger.info( f"[gen{self.server_idx} r{self.rank}] wave recvs finished in " @@ -1031,20 +1144,48 @@ def gen_run_wave(self, peer_idx, li, req_len, rep, wave, params_by_pair): p for p, r in reqs.items() if r.state == self.llm_request_state.DISAGG_TRANS_ERROR ] if bad: - local_err = _TransferError(f"gen DISAGG_TRANS_ERROR on pairs {bad}") - elif self.plan["verify_data"] and rep >= self.plan["warmup_requests"]: + raise _TransferError(f"gen DISAGG_TRANS_ERROR on pairs {bad}") + incomplete = [ + p + for p, r in reqs.items() + if r.state != self.llm_request_state.DISAGG_GENERATION_TRANS_COMPLETE + ] + if incomplete: + raise _TransferError(f"gen requests not complete for pairs {incomplete}") + except _Timeout: + raise + except Exception as e: # noqa: BLE001 + transfer_error = e + try: + reason = self._consensus_error(transfer_error) + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - release proof is incomplete + raise _FatalTransferError(f"gen transfer ownership consensus failed: {e!r}") from e + if reason is not None: + # Exact completion plus CUDA stream synchronization is the release + # proof. Do not let a locally successful rank free while another + # rank still has transport-owned pages. + raise _FatalTransferError(f"gen transfer failed: {reason}") + + # Remote writes and local CUDA work are complete on every rank. Later + # byte-verification failures do not invalidate the ownership proof. + verification_error = None + try: + if self.plan["verify_data"] and rep >= self.plan["warmup_requests"]: for pair, req in reqs.items(): - ok, detail = verify_request(self.kvm, req.py_request_id) + ok, detail = verify_request(self.kvm, req.py_request_id, req_len) if not ok: mismatch = f"pair={pair} {detail}" break + except _Timeout: + raise except Exception as e: # noqa: BLE001 - local_err = e - finally: - self._free_all(reqs) - reason = self._consensus_error(local_err) + verification_error = e + self._free_all(reqs) + reason = self._consensus_error(verification_error) if reason is not None: - raise _TransferError(f"gen transfer failed: {reason}") + raise _TransferError(f"gen verification failed: {reason}") mismatches = [m for m in self.comm.allgather(mismatch) if m] return (not mismatches, "; ".join(mismatches[:4])) @@ -1071,13 +1212,19 @@ def _leader_send_recv(self, sock, obj, key): """REQ round-trip on the gen leader; broadcast the reply to all ranks.""" reply = None err = None + timed_out = False if self.is_leader: try: sock.send(pack_msg(obj, key)) reply = unpack_msg(sock.recv(), key) + except _Timeout as e: + timed_out = True + err = repr(e) except Exception as e: # noqa: BLE001 err = repr(e) - err, reply = self.comm.bcast((err, reply), root=0) + timed_out, err, reply = self.comm.bcast((timed_out, err, reply), root=0) + if timed_out: + raise _Timeout(err) if err: raise _TransferError(f"ZMQ control channel failed: {err}") return reply @@ -1143,13 +1290,18 @@ def ctx_serve_peer(runner, sock, peer_idx, arm, disarm, key): comm = runner.comm def leader_recv(): - msg, err = None, None + msg, err, timed_out = None, None, False if runner.is_leader: try: msg = unpack_msg(sock.recv(), key) + except _Timeout as e: + timed_out = True + err = repr(e) except Exception as e: # noqa: BLE001 err = repr(e) - err, msg = comm.bcast((err, msg), root=0) + timed_out, err, msg = comm.bcast((timed_out, err, msg), root=0) + if timed_out: + raise _Timeout(err) if err: raise _TransferError(f"ZMQ recv from gen_{peer_idx} failed: {err}") return msg @@ -1179,12 +1331,39 @@ def leader_reply(obj): ) try: params_by_pair, reqs = runner.ctx_run_wave(peer_idx, li, req_len, rep, wave) + except _FatalTransferError as e: + try: + leader_reply(("abort", str(e))) + except Exception: # noqa: BLE001 - preserve ownership-fatal outcome + pass + raise except _TransferError as e: leader_reply(("abort", str(e))) raise + except Exception as e: # noqa: BLE001 - dispatch may have started + fatal = _FatalTransferError(f"ctx send path failed with ownership unknown: {e!r}") + try: + leader_reply(("abort", str(fatal))) + except Exception: # noqa: BLE001 - preserve ownership-fatal outcome + pass + raise fatal from e # JSON object keys are strings; the gen side converts back to int. - leader_reply(("params", {str(p): params_to_wire(v) for p, v in params_by_pair.items()})) - runner.ctx_finish_wave(reqs) + try: + leader_reply(("params", {str(p): params_to_wire(v) for p, v in params_by_pair.items()})) + except _Timeout: + raise + except Exception as e: # noqa: BLE001 - sends may already be in flight + raise _FatalTransferError( + f"failed to publish params to gen_{peer_idx} after send dispatch: {e!r}" + ) from e + try: + runner.ctx_finish_wave(reqs) + except (_Timeout, _FatalTransferError, _TransferError): + raise + except Exception as e: # noqa: BLE001 - send quiescence is unknown + raise _FatalTransferError( + f"ctx completion proof failed with ownership unknown: {e!r}" + ) from e # The gen defers "done" until it has finished the schedules of ALL its # ctx peers, so every ctx instance stays alive for the whole precheck -- @@ -1276,13 +1455,33 @@ def gen_run_peer(runner, peer_idx, arm, disarm): case_ok = {} for li, req_len, rep, wave in _schedule(plan): arm(f"ctx_{peer_idx} len={req_len} rep={rep}", seconds=wave_timeout_s(plan, li, rep)) - reply = runner._leader_send_recv( - sock, ("go", {"li": li, "rep": rep, "wave": wave[0]}), key - ) - if reply[0] == "abort": - raise _TransferError(f"ctx_{peer_idx} aborted: {reply[1]}") - params_by_pair = {int(p): params_from_wire(v) for p, v in reply[1].items()} - ok, detail = runner.gen_run_wave(peer_idx, li, req_len, rep, wave, params_by_pair) + try: + reply = runner._leader_send_recv( + sock, ("go", {"li": li, "rep": rep, "wave": wave[0]}), key + ) + if reply[0] == "abort": + raise _FatalTransferError(f"ctx_{peer_idx} aborted: {reply[1]}") + if reply[0] != "params": + raise _FatalTransferError( + f"unexpected wave reply from ctx_{peer_idx}: {reply[:1]}" + ) + params_by_pair = {int(p): params_from_wire(v) for p, v in reply[1].items()} + except _Timeout: + raise + except _FatalTransferError: + raise + except Exception as e: # noqa: BLE001 - ctx may have dispatched sends + raise _FatalTransferError( + f"wave control failed after ctx_{peer_idx} may have dispatched sends: {e!r}" + ) from e + try: + ok, detail = runner.gen_run_wave(peer_idx, li, req_len, rep, wave, params_by_pair) + except (_Timeout, _FatalTransferError, _TransferError): + raise + except Exception as e: # noqa: BLE001 - receive quiescence is unknown + raise _FatalTransferError( + f"gen completion proof failed with ownership unknown: {e!r}" + ) from e if rep >= plan["warmup_requests"]: prev_ok, prev_detail = case_ok.get(req_len, (True, "")) case_ok[req_len] = (prev_ok and ok, prev_detail or detail) @@ -1378,14 +1577,26 @@ def _install_watchdog(runner, plan, rank): current_cell = {"what": "startup"} def _on_hang(): - runner.recorder.record("-", 0, "TIMEOUT", f"hang detected during {current_cell['what']}") - runner.recorder.finalize() - sys.stderr.write( - f"[precheck {runner.role}_{runner.server_idx} r{rank}] WATCHDOG_KILL " - f"{current_cell['what']}\n" - ) - sys.stderr.flush() - os.kill(os.getpid(), signal.SIGKILL) + if runner.is_leader: + try: + runner.recorder.record( + "-", 0, "TIMEOUT", f"hang detected during {current_cell['what']}" + ) + runner.recorder.finalize() + except Exception: # noqa: BLE001 - SIGKILL must still happen + pass + else: + # Let the leader's concurrently armed watchdog replace its status + # files before this rank's bad exit tears down the srun step. + time.sleep(ABORT_COORDINATION_TIMEOUT_S) + try: + sys.stderr.write( + f"[precheck {runner.role}_{runner.server_idx} r{rank}] WATCHDOG_KILL " + f"{current_cell['what']}\n" + ) + sys.stderr.flush() + finally: + os.kill(os.getpid(), signal.SIGKILL) # The detector must outlast the LONGEST legitimate wait (peer handshakes # are serialized across sessions); per-cell alarms are the tighter bound @@ -1434,6 +1645,43 @@ def record_peer_failure(peer, exc): return record_peer_failure +def _hard_abort_unquiesced(runner, current_cell, exc): + """Persist the ownership-fatal verdict, then exit without teardown.""" + try: + signal.alarm(0) + except Exception: # noqa: BLE001 - abort must still happen + pass + if isinstance(exc, _Timeout): + status = "TIMEOUT" + reason = f"exceeded the budget during {current_cell['what']}" + else: + status = "TRANSFER_ERROR" + reason = str(exc) + if runner.is_leader: + try: + peer = current_cell["what"] + runner.recorder.record(peer, 0, status, reason) + raise_abort_flag(runner.work_dir, f"{peer} {status}: {reason}") + runner.recorder.finalize( + extra={ + "kv_cache_manager": "V2" if runner.use_v2 else "V1", + "transceiver_runtime": runner.runtime, + } + ) + except Exception: # noqa: BLE001 - abort must still happen + pass + # Fatal transfer paths first reach instance-wide consensus. This bounded + # rendezvous lets the leader finish replacing the status files before a + # nonleader aborts MPI, but does not rely on coordination for rank-local + # signal timeouts. + _coordinate_abort_after_leader_flush(runner.comm) + try: + sys.stdout.flush() + sys.stderr.flush() + finally: + _hard_abort_process(runner.comm) + + def _consensus_abort_reason(runner): """Instance-wide agreed view of the fail-fast flag (leader reads, bcast). @@ -1475,6 +1723,8 @@ def _serve_gen_peers(runner, plan, arm, disarm, record_peer_failure): try: ctx_serve_peer(runner, socks.get(gj), gj, arm, disarm, keys.get(gj)) runner.recorder.record(f"gen_{gj}", 0, "PASS", "served all transfers") + except (_FatalTransferError, _Timeout): + raise except _PeerAbort as e: # A gen driver that failed elsewhere aborts our session as part of # fail-fast: record a (non-failing) SKIP, not our own failure -- @@ -1517,6 +1767,8 @@ def _drive_ctx_peers(runner, arm, disarm, record_peer_failure): try: sock, sess_key = gen_run_peer(runner, ci, arm, disarm) open_sessions.append((ci, sock, sess_key)) + except (_FatalTransferError, _Timeout): + raise except Exception as e: # noqa: BLE001 - failure sets the fail-fast flag record_peer_failure(f"ctx_{ci}", e) for ci, sock, sess_key in open_sessions: @@ -1620,6 +1872,16 @@ def main(argv=None): _serve_gen_peers(runner, plan, arm, disarm, record_peer_failure) else: _drive_ctx_peers(runner, arm, disarm, record_peer_failure) + except (_FatalTransferError, _Timeout) as e: + try: + print( + f"[precheck {runner.role}_{runner.server_idx} r{rank}] OWNERSHIP_FATAL: {e}", + file=sys.stderr, + flush=True, + ) + except Exception: # noqa: BLE001 - abort must still happen + pass + _hard_abort_unquiesced(runner, current_cell, e) finally: stop_watchdog() diff --git a/tests/unittest/disaggregated/test_cache_transceiver_harness.py b/tests/unittest/disaggregated/test_cache_transceiver_harness.py index cb37e89b4660..3b5aac687d40 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_harness.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_harness.py @@ -21,14 +21,19 @@ Requires: 1 GPU, mpirun, mpi4py, tensorrt_llm. """ +import ast import json import os +import pickle import shutil import signal import socket import subprocess import sys import time +import types +from pathlib import Path +from typing import Any, Iterable, Optional, Sequence from unittest.mock import MagicMock, call import pytest @@ -53,6 +58,118 @@ _TERMINATE_GRACE_SECONDS = 5 +def _load_driver_subset( + module_name: str, + selected_names: set[str], + stubs: dict[str, Any], +) -> types.ModuleType: + """Execute selected driver definitions with injected runtime stand-ins.""" + source = Path(DRIVER_SCRIPT).read_text() + tree = ast.parse(source, filename=DRIVER_SCRIPT) + selected = [ + node + for node in tree.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name in selected_names + ] + module = types.ModuleType(module_name) + module.__dict__.update(stubs) + exec( + compile(ast.Module(body=selected, type_ignores=[]), DRIVER_SCRIPT, "exec"), + module.__dict__, + ) + missing = selected_names - set(module.__dict__) + assert not missing, f"{module_name}: driver definitions not loaded: {sorted(missing)}" + return module + + +def _load_driver_ownership_helpers() -> types.ModuleType: + """Load pure ownership helpers without importing GPU/MPI runtime modules.""" + selected_names = { + "_Timeout", + "_TransferError", + "_FatalTransferError", + "_request_ids", + "_context_completion_error", + "_gen_completion_error", + "_can_release_sequence", + "_release_sequence_if_safe", + "_validate_context_completion", + "_validate_python_gen_completion", + "_first_reason", + "_exchange_release_decision", + "_hard_abort_process", + } + return _load_driver_subset( + "cache_transceiver_harness_ownership", + selected_names, + { + "Any": Any, + "Iterable": Iterable, + "Optional": Optional, + "Sequence": Sequence, + "MPI": types.SimpleNamespace(MIN="MIN"), + "LlmRequestState": types.SimpleNamespace( + DISAGG_GENERATION_TRANS_COMPLETE="gen_complete", + DISAGG_TRANS_ERROR="error", + ), + "free_sequence": MagicMock(), + "os": os, + "pickle": pickle, + "signal": signal, + }, + ) + + +OWNERSHIP = _load_driver_ownership_helpers() + + +def _load_driver_request_flow() -> types.ModuleType: + """Load the request flow with local stand-ins for GPU/MPI dependencies.""" + selected_names = { + "_Timeout", + "_TransferError", + "_FatalTransferError", + "_request_ids", + "_context_completion_error", + "_gen_completion_error", + "_can_release_sequence", + "_release_sequence_if_safe", + "_validate_context_completion", + "_validate_python_gen_completion", + "_first_reason", + "_exchange_release_decision", + "_wait_gen_complete", + "run_one_request", + } + return _load_driver_subset( + "cache_transceiver_harness_request_flow", + selected_names, + { + "Any": Any, + "Iterable": Iterable, + "Optional": Optional, + "Sequence": Sequence, + "MPI": types.SimpleNamespace(MAX="MAX", MIN="MIN"), + "LlmRequest": Any, + "LlmRequestState": types.SimpleNamespace( + DISAGG_GENERATION_TRANS_COMPLETE="gen_complete", + DISAGG_TRANS_ERROR="error", + ), + "add_sequence": MagicMock(), + "fill_request": MagicMock(), + "free_sequence": MagicMock(), + "make_request": MagicMock(), + "pickle": pickle, + "tensorrt_llm": types.SimpleNamespace(logger=MagicMock()), + "torch": types.SimpleNamespace(cuda=types.SimpleNamespace(synchronize=MagicMock())), + "verify_request": MagicMock(), + }, + ) + + +REQUEST_FLOW = _load_driver_request_flow() + + def _find_free_port(): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("", 0)) @@ -160,6 +277,380 @@ def test_terminate_process_groups_signals_group_after_leader_exit(self, monkeypa proc.wait.assert_not_called() +class _FakeDecisionComm: + def __init__(self, *, reduced: Optional[int] = None) -> None: + self.reduced = reduced + self.abort_calls = [] + + def allreduce(self, value, op=None): + del op + return value if self.reduced is None else self.reduced + + def gather(self, value, root=0): + del root + return [value] + + def bcast(self, value, root=0): + del root + return value + + def Abort(self, code): + self.abort_calls.append(code) + + +class _FakeDecisionSocket: + def __init__(self, *responses) -> None: + self.responses = list(responses) + self.sent = [] + + def send(self, payload) -> None: + self.sent.append(pickle.loads(payload)) + + def recv(self): + return pickle.dumps(self.responses.pop(0)) + + +class _FakeRunSocket: + def __init__(self, responses, events) -> None: + self.responses = list(responses) + self.events = events + self.sent = [] + self.recv_count = 0 + + def send(self, payload) -> None: + self.sent.append(payload if payload == b"go" else pickle.loads(payload)) + + def recv(self): + self.recv_count += 1 + if self.recv_count == 2: + self.events.append("peer_ack") + response = self.responses.pop(0) + if isinstance(response, BaseException): + raise response + if isinstance(response, bytes): + return response + return pickle.dumps(response) + + +class _FakeGenTransceiver: + def __init__(self, status, events, receive_error=None) -> None: + self.status = status + self.events = events + self.receive_error = receive_error + + def request_and_receive_async(self, req) -> None: + del req + if self.receive_error is not None: + raise self.receive_error + + def check_gen_transfer_status(self, timeout): + assert timeout is None + self.events.append("block_all") + if isinstance(self.status, BaseException): + raise self.status + return self.status + + +class _FakeContextTransceiver: + def __init__(self, send_error) -> None: + self.send_error = send_error + + def respond_and_send_async(self, req) -> None: + del req + raise self.send_error + + +def _run_gen_request( + status, + state, + *, + final_response=("COMPLETE", ""), + receive_error=None, + synchronize_error=None, +): + events = [] + req = types.SimpleNamespace(py_request_id=7, state=state) + comm = _FakeDecisionComm() + sock = _FakeRunSocket( + [("OK", types.SimpleNamespace()), final_response], + events, + ) + xcvr = _FakeGenTransceiver(status, events, receive_error=receive_error) + + REQUEST_FLOW.make_request.reset_mock(return_value=True, side_effect=True) + REQUEST_FLOW.make_request.return_value = req + REQUEST_FLOW.add_sequence.reset_mock(return_value=True, side_effect=True) + REQUEST_FLOW.add_sequence.return_value = "handle" + REQUEST_FLOW.fill_request.reset_mock(return_value=True, side_effect=True) + REQUEST_FLOW.free_sequence.reset_mock(return_value=True, side_effect=True) + REQUEST_FLOW.free_sequence.side_effect = lambda *args: events.append("free") + REQUEST_FLOW.verify_request.reset_mock(return_value=True, side_effect=True) + REQUEST_FLOW.verify_request.return_value = True + REQUEST_FLOW.torch.cuda.synchronize.reset_mock(return_value=True, side_effect=True) + + def synchronize(): + events.append("cuda_sync") + if synchronize_error is not None: + raise synchronize_error + + REQUEST_FLOW.torch.cuda.synchronize.side_effect = synchronize + result = REQUEST_FLOW.run_one_request( + "gen", + comm, + "manager", + xcvr, + "PYTHON", + True, + 2, + 7, + 16, + 0, + sock, + ) + return result, events, sock + + +def _run_ctx_request(send_error): + events = [] + req = types.SimpleNamespace( + py_request_id=7, + state="in_progress", + context_phase_params=types.SimpleNamespace(), + ) + comm = _FakeDecisionComm() + sock = _FakeRunSocket([b"go"], events) + xcvr = _FakeContextTransceiver(send_error) + + REQUEST_FLOW.make_request.reset_mock(return_value=True, side_effect=True) + REQUEST_FLOW.make_request.return_value = req + REQUEST_FLOW.add_sequence.reset_mock(return_value=True, side_effect=True) + REQUEST_FLOW.add_sequence.return_value = "handle" + REQUEST_FLOW.fill_request.reset_mock(return_value=True, side_effect=True) + REQUEST_FLOW.free_sequence.reset_mock(return_value=True, side_effect=True) + return REQUEST_FLOW.run_one_request( + "ctx", + comm, + "manager", + xcvr, + "PYTHON", + True, + 2, + 7, + 16, + 0, + sock, + ) + + +class TestRequestTransferOwnershipFlow: + @pytest.mark.parametrize( + ("status", "state"), + [ + (([], [], []), "gen_complete"), + (([], [7], []), "error"), + (([], [], [7]), "in_progress"), + (([7], [], []), "in_progress"), + ], + ids=("missing", "failed", "cancelled", "nonterminal"), + ) + def test_python_gen_unproven_completion_never_frees(self, status, state): + with pytest.raises(REQUEST_FLOW._FatalTransferError): + _run_gen_request(status, state, final_response=("FATAL", "gen unsafe")) + + REQUEST_FLOW.free_sequence.assert_not_called() + + def test_cuda_synchronize_error_before_release_handshake_never_frees(self): + with pytest.raises(REQUEST_FLOW._FatalTransferError): + _run_gen_request( + ([7], [], []), + "gen_complete", + final_response=("FATAL", "gen cuda synchronize failed"), + synchronize_error=RuntimeError("CUDA synchronize failed"), + ) + + REQUEST_FLOW.free_sequence.assert_not_called() + + def test_gen_partial_dispatch_setup_error_is_fatal_and_never_frees(self): + with pytest.raises(REQUEST_FLOW._FatalTransferError, match="partial gen dispatch"): + _run_gen_request( + ([7], [], []), + "gen_complete", + receive_error=RuntimeError("partial gen dispatch"), + ) + + REQUEST_FLOW.free_sequence.assert_not_called() + + def test_ctx_partial_dispatch_setup_error_is_fatal_and_never_frees(self): + with pytest.raises(REQUEST_FLOW._FatalTransferError, match="partial ctx dispatch"): + _run_ctx_request(RuntimeError("partial ctx dispatch")) + + REQUEST_FLOW.free_sequence.assert_not_called() + + @pytest.mark.parametrize("operation", ("receive", "block_all", "cuda_sync")) + def test_timeout_from_caught_driver_operation_propagates(self, operation): + timeout = REQUEST_FLOW._Timeout("cell deadline expired") + kwargs = {} + status = ([7], [], []) + if operation == "receive": + kwargs["receive_error"] = timeout + elif operation == "block_all": + status = timeout + else: + kwargs["synchronize_error"] = timeout + + with pytest.raises(REQUEST_FLOW._Timeout, match="cell deadline expired"): + _run_gen_request(status, "gen_complete", **kwargs) + + REQUEST_FLOW.free_sequence.assert_not_called() + + def test_gen_frees_only_after_completion_sync_and_both_role_ack(self): + result, events, sock = _run_gen_request(([7], [], []), "gen_complete") + + assert result is True + assert events == ["block_all", "cuda_sync", "peer_ack", "free"] + assert sock.sent == [b"go", ("COMPLETE", "")] + REQUEST_FLOW.free_sequence.assert_called_once_with( + "manager", + REQUEST_FLOW.make_request.return_value, + "handle", + True, + ) + + def test_context_fatal_ack_never_frees_completed_gen_sequence(self): + with pytest.raises(REQUEST_FLOW._FatalTransferError, match="ctx unsafe"): + _run_gen_request( + ([7], [], []), + "gen_complete", + final_response=("FATAL", "ctx unsafe"), + ) + + REQUEST_FLOW.free_sequence.assert_not_called() + + +class TestTransferOwnershipHelpers: + def test_context_requires_exact_completed_request(self): + assert OWNERSHIP._context_completion_error(7, [7], [], "in_progress", "error") is None + assert "failed" in OWNERSHIP._context_completion_error(7, [], [7], "error", "error") + assert "without completing" in OWNERSHIP._context_completion_error( + 7, [8], [], "in_progress", "error" + ) + + @pytest.mark.parametrize( + ("completed", "failed", "cancelled", "state", "expected"), + [ + ([7], [], [], "gen_complete", None), + ([], [7], [], "error", "failed"), + ([], [], [types.SimpleNamespace(py_request_id=7)], "in_progress", "cancelled"), + ([], [], [], "in_progress", "without completing"), + ([7], [], [], "in_progress", "nonterminal"), + ], + ) + def test_generation_requires_completed_terminal_request( + self, completed, failed, cancelled, state, expected + ): + error = OWNERSHIP._gen_completion_error( + 7, + completed, + failed, + cancelled, + state, + "gen_complete", + "error", + ) + if expected is None: + assert error is None + else: + assert expected in error + + @pytest.mark.parametrize( + ("started", "completed", "should_release"), + [ + (False, False, True), + (True, False, False), + (True, True, True), + ], + ) + def test_release_sequence_requires_quiescence(self, started, completed, should_release): + OWNERSHIP.free_sequence.reset_mock() + released = OWNERSHIP._release_sequence_if_safe( + "manager", + types.SimpleNamespace(py_request_id=7), + "handle", + True, + transfer_may_have_started=started, + transfer_completed=completed, + ) + + assert released is should_release + assert OWNERSHIP.free_sequence.call_count == int(should_release) + + def test_gen_release_handshake_requires_context_acknowledgement(self): + socket = _FakeDecisionSocket(("COMPLETE", "")) + + safe, reason = OWNERSHIP._exchange_release_decision( + "gen", _FakeDecisionComm(), True, socket, True + ) + + assert safe and reason == "" + assert socket.sent == [("COMPLETE", "")] + + def test_context_release_handshake_propagates_local_timeout(self): + socket = _FakeDecisionSocket(("COMPLETE", "")) + + safe, reason = OWNERSHIP._exchange_release_decision( + "ctx", + _FakeDecisionComm(reduced=0), + True, + socket, + False, + "sender deadline expired", + ) + + assert not safe + assert reason == "sender deadline expired" + assert socket.sent == [("FATAL", "sender deadline expired")] + + def test_release_handshake_rejects_invalid_peer_status(self): + socket = _FakeDecisionSocket(("UNKNOWN", "")) + + safe, reason = OWNERSHIP._exchange_release_decision( + "gen", _FakeDecisionComm(), True, socket, True + ) + + assert not safe + assert "invalid" in reason + + def test_context_release_handshake_sends_invalid_peer_verdict(self): + socket = _FakeDecisionSocket(("UNKNOWN", "")) + + safe, reason = OWNERSHIP._exchange_release_decision( + "ctx", _FakeDecisionComm(), True, socket, True + ) + + expected = "invalid gen peer release status: 'UNKNOWN'" + assert not safe + assert reason == expected + assert socket.sent == [("FATAL", expected)] + + def test_release_handshake_preserves_timeout(self): + socket = _FakeDecisionSocket() + socket.recv = MagicMock(side_effect=OWNERSHIP._Timeout("release deadline expired")) + + with pytest.raises(OWNERSHIP._Timeout, match="release deadline expired"): + OWNERSHIP._exchange_release_decision("gen", _FakeDecisionComm(), True, socket, True) + + def test_hard_abort_falls_back_to_sigkill(self, monkeypatch): + comm = _FakeDecisionComm() + kill = MagicMock() + monkeypatch.setattr(OWNERSHIP.os, "kill", kill) + + with pytest.raises(RuntimeError, match="SIGKILL unexpectedly returned"): + OWNERSHIP._hard_abort_process(comm) + + assert comm.abort_calls == [137] + kill.assert_called_once_with(os.getpid(), signal.SIGKILL) + + def _build_config(work_dir: str) -> dict: return { "hardware": {"gpus_per_node": 2}, diff --git a/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py b/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py index e8a92f99c99c..9b106a1d8d43 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py @@ -118,12 +118,12 @@ def _terminate_process_groups(processes): pass -def _disagg_yaml(num_ctx, num_gen, ctx_tp, gen_tp, request_lengths=(64,)): +def _disagg_yaml(num_ctx, num_gen, ctx_tp, gen_tp, request_lengths=(64,), mtp_draft_len=0): """Minimal disagg perf-sanity yaml shaped like the checked-in configs.""" tokens_per_block = 32 def side(tp): - return { + config = { "tensor_parallel_size": tp, "pipeline_parallel_size": 1, "kv_cache_config": { @@ -139,6 +139,12 @@ def side(tp): "max_tokens_in_buffer": 512, }, } + if mtp_draft_len: + config["speculative_config"] = { + "decoding_type": "MTP", + "max_draft_len": mtp_draft_len, + } + return config return { "metadata": {"model_dir_name": "tiny-llama"}, @@ -311,6 +317,17 @@ def test_precheck_passes(tmp_path, num_ctx, num_gen, ctx_tp, gen_tp): assert peers == {f"ctx_{ci}" for ci in range(num_ctx)} +@pytest.mark.timeout(300) +def test_precheck_passes_mtp_exact_block_boundary(tmp_path): + """Reserved MTP tokens must not expand the verified transfer payload.""" + pytest.importorskip("mpi4py") + cfg = _disagg_yaml(1, 1, 1, 1, request_lengths=(64,), mtp_draft_len=3) + config_path, models_root = _write_inputs(tmp_path, cfg) + work_dir, launched = _launch_instances(tmp_path, _jobs(cfg, config_path), models_root) + _wait_all(launched) + _assert_all_passed(work_dir, launched) + + @pytest.mark.timeout(300) def test_precheck_fails_fast_on_fingerprint_mismatch(tmp_path): """Mismatched ctx/gen yamls must produce FAIL verdicts, not a hang. diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py index 76eff6adc439..9185d9ef2283 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -16,14 +16,22 @@ from __future__ import annotations +import threading +from collections.abc import Callable from dataclasses import dataclass +from types import SimpleNamespace from typing import Optional from unittest.mock import Mock import pytest from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult -from tensorrt_llm._torch.disaggregation.native.transfer import TaskStatus, TxSession +from tensorrt_llm._torch.disaggregation.native.transfer import ( + TaskStatus, + TransferWorker, + TransferWorkerConfig, + TxSession, +) from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm.bindings import LlmRequestState @@ -58,6 +66,7 @@ def __init__( self._has_failed = has_failed self.blocking_calls: list[bool] = [] self.closed = False + self.aux_slot: Optional[int] = 0 @property def disagg_request_id(self) -> int: @@ -79,17 +88,41 @@ def has_failed(self) -> bool: def close(self) -> None: self.closed = True + self.aux_slot = None class _FakeTask: - def __init__(self, status: TaskStatus, wait_result: bool = True) -> None: + def __init__( + self, + status: TaskStatus, + wait_result: bool | list[bool] = True, + on_wait: Optional[Callable[[Optional[float]], None]] = None, + ) -> None: self.status = status - self._wait_result = wait_result + self._wait_results = list(wait_result) if isinstance(wait_result, list) else [wait_result] + self._on_wait = on_wait self.wait_calls: list[Optional[float]] = [] def wait(self, timeout: Optional[float] = None) -> bool: self.wait_calls.append(timeout) - return self._wait_result + if self._on_wait is not None: + self._on_wait(timeout) + result = self._wait_results.pop(0) if len(self._wait_results) > 1 else self._wait_results[0] + if result and self.status != TaskStatus.ERROR: + self.status = TaskStatus.TRANSFERRED + return result + + +class _FakeClock: + def __init__(self, now_s: float = 0.0) -> None: + self.now_s = now_s + + def monotonic(self) -> float: + return self.now_s + + def advance(self, elapsed_s: Optional[float]) -> None: + assert elapsed_s is not None + self.now_s += elapsed_s def _make_transceiver( @@ -100,19 +133,17 @@ def _make_transceiver( transceiver._send_sessions = sessions transceiver._send_reqs = reqs or {rid: _FakeRequest() for rid in sessions} transceiver._sender_future_timeout_ms = 123 + transceiver.kv_transfer_timeout_ms = 60_000 # Attributes read by check_context_transfer_status before it processes sessions. transceiver._ever_had_send_session = True transceiver._ctx_need_tp_sync = False transceiver._ctx_need_pp_sync = False transceiver._transfer_worker = _FakeTransferWorker() transceiver._ctx_consensus = lambda local_ids: list(local_ids) - transceiver._ctx_consensus_outcome = ( - lambda _to_process, cancelled, failed, completed, timed_out: ( - cancelled, - failed, - completed, - timed_out, - ) + transceiver._ctx_consensus_outcome = lambda _to_process, cancelled, failed, completed: ( + cancelled, + failed, + completed, ) return transceiver @@ -122,14 +153,20 @@ def _make_tx_session( *, need_aux: bool = False, aux_task: Optional[_FakeTask] = None, + timeout_s: Optional[float] = 0.25, + deadline_monotonic_s: Optional[float] = None, ) -> TxSession: session = object.__new__(TxSession) - session._timeout_s = 0.25 + session._timeout_s = timeout_s + session._overall_timeout_s = None + session._deadline_monotonic_s = deadline_monotonic_s session._need_aux = need_aux session._terminal_status = None + session._exception = None session.receiver_ready = True session.kv_tasks = kv_tasks session.aux_task = aux_task + session.lock = threading.Lock() session._closed = False session._aux_buffer = None session.aux_slot = None @@ -214,6 +251,37 @@ def test_context_transfer_status_block_all_uses_blocking_wait() -> None: assert 12 not in transceiver._send_reqs +def test_context_transfer_status_timeout_retains_session_and_request(monkeypatch) -> None: + session = _FakeSession(rid=16, wait_result=WaitResult.TIMEOUT) + req = _FakeRequest() + transceiver = _make_transceiver({16: session}, {16: req}) + transceiver.kv_transfer_timeout_ms = 60_000 + warning = Mock() + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.transceiver.logger.warning", + warning, + ) + + completed, failed = transceiver.check_context_transfer_status(None) + + completed_again, failed_again = transceiver.check_context_transfer_status(None) + + assert completed == [] + assert failed == [] + assert completed_again == [] + assert failed_again == [] + assert session.blocking_calls == [True, True] + assert not session.closed + assert session.aux_slot == 0 + assert transceiver._send_sessions == {16: session} + assert transceiver._send_reqs == {16: req} + assert warning.call_count == 2 + messages = [args[0] for args, _kwargs in warning.call_args_list] + assert all("rid=16" in message for message in messages) + assert all("kv_transfer_timeout_ms=60000ms" in message for message in messages) + assert all("keeping it in progress" in message for message in messages) + + def test_context_transfer_status_zero_budget_processes_task_level_failure() -> None: session = _FakeSession( rid=13, @@ -314,6 +382,38 @@ def fake_allgather(payload): assert new_completed == [7] # intersection only (8 is completed on the peer only) +def test_ctx_tp_consensus_does_not_complete_when_peer_times_out() -> None: + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._ctx_need_tp_sync = True + transceiver._ctx_need_pp_sync = False + transceiver._dist = SimpleNamespace( + tp_allgather=lambda payload: [payload, [[], [], []]], + ) + + cancelled, failed, completed = transceiver._ctx_consensus_outcome([21], [], [], [21]) + + assert cancelled == [] + assert failed == [] + assert completed == [] + + +def test_ctx_pp_consensus_does_not_complete_when_peer_times_out() -> None: + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._ctx_need_tp_sync = False + transceiver._ctx_need_pp_sync = True + transceiver._dist = SimpleNamespace( + tp_allgather=Mock(side_effect=AssertionError("TP allgather must be skipped")), + pp_allgather=lambda payload: [payload, [[], [], []]], + ) + + cancelled, failed, completed = transceiver._ctx_consensus_outcome([22], [], [], [22]) + + assert cancelled == [] + assert failed == [] + assert completed == [] + transceiver._dist.tp_allgather.assert_not_called() + + @pytest.mark.skip( reason="ctx idle fast-path was dropped from this branch. TODO: when the " "fast-path is reintroduced, its terminal-count reduction must mirror " @@ -340,7 +440,7 @@ def test_ctx_consensus_fastpath_skips_when_idle(monkeypatch) -> None: transceiver._dist.allreduce = Mock(return_value=0) transceiver._ctx_consensus = Mock(return_value=[]) transceiver._build_to_process = Mock(return_value=[]) - transceiver._ctx_consensus_outcome = Mock(return_value=([], [], [], [])) + transceiver._ctx_consensus_outcome = Mock(return_value=([], [], [])) transceiver._transfer_worker = _FakeTransferWorker() transceiver._close_failed_sessions = Mock() @@ -356,14 +456,506 @@ def test_ctx_consensus_fastpath_skips_when_idle(monkeypatch) -> None: transceiver._ctx_consensus.assert_called_once() -def test_tx_session_wait_complete_defaults_to_blocking() -> None: - task = _FakeTask(TaskStatus.INIT, wait_result=False) +def test_tx_session_blocking_wait_retries_wait_slices_until_complete() -> None: + task = _FakeTask(TaskStatus.INIT, wait_result=[False, True]) + session = _make_tx_session([task]) + + assert session.wait_complete() == WaitResult.COMPLETED + assert task.wait_calls == [0.25, 0.25] + + +def test_tx_session_blocking_wait_times_out_stalled_task_and_does_not_reset_deadline( + monkeypatch, +) -> None: + clock = _FakeClock() + task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False, on_wait=clock.advance) + session = _make_tx_session( + [task], + timeout_s=0.25, + deadline_monotonic_s=0.6, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.transfer.time.monotonic", + clock.monotonic, + ) + + assert session.wait_complete(blocking=True) == WaitResult.TIMEOUT + assert task.wait_calls == pytest.approx([0.25, 0.25, 0.1]) + assert not session._closed + assert not session.has_failed() + + wait_call_count = len(task.wait_calls) + assert session.wait_complete(blocking=True) == WaitResult.TIMEOUT + assert len(task.wait_calls) == wait_call_count + + +def test_tx_session_blocking_wait_uses_finite_overall_fallback_when_unset( + monkeypatch, +) -> None: + clock = _FakeClock() + task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False, on_wait=clock.advance) + session = _make_tx_session([task], timeout_s=0.25) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.transfer._FALLBACK_TX_OVERALL_TIMEOUT_S", + 0.6, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.transfer.time.monotonic", + clock.monotonic, + ) + + assert session.wait_complete(blocking=True) == WaitResult.TIMEOUT + assert task.wait_calls == pytest.approx([0.25, 0.25, 0.1]) + assert session._deadline_monotonic_s == pytest.approx(0.6) + + wait_call_count = len(task.wait_calls) + assert session.wait_complete(blocking=True) == WaitResult.TIMEOUT + assert len(task.wait_calls) == wait_call_count + + +def test_tx_session_completion_observed_at_deadline_wins_over_timeout(monkeypatch) -> None: + clock = _FakeClock() + task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False) + + def complete_at_deadline(timeout_s: Optional[float]) -> None: + clock.advance(timeout_s) + task.status = TaskStatus.TRANSFERRED + + task._on_wait = complete_at_deadline + session = _make_tx_session( + [task], + timeout_s=0.25, + deadline_monotonic_s=0.25, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.transfer.time.monotonic", + clock.monotonic, + ) + + assert session.wait_complete(blocking=True) == WaitResult.COMPLETED + assert task.wait_calls == [0.25] + + +@pytest.mark.parametrize( + ("terminal", "expected"), + [ + ("completed", WaitResult.COMPLETED), + ("failed", WaitResult.FAILED), + ("cancelled", WaitResult.FAILED), + ], +) +def test_tx_session_terminal_transition_during_deadline_read_wins_over_timeout( + monkeypatch, + terminal: str, + expected: WaitResult, +) -> None: + task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False) + session = _make_tx_session( + [task], + timeout_s=0.25, + deadline_monotonic_s=0.25, + ) + + def expire_after_terminal_transition() -> float: + if terminal == "completed": + task.status = TaskStatus.TRANSFERRED + elif terminal == "failed": + task.status = TaskStatus.ERROR + else: + session._terminal_status = SessionStatus.CANCELLED + return 0.25 + + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.transfer.time.monotonic", + expire_after_terminal_transition, + ) + + assert session.wait_complete(blocking=True) == expected + assert task.wait_calls == [] + + +def test_tx_session_failure_observed_at_deadline_wins_over_timeout(monkeypatch) -> None: + clock = _FakeClock() + task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False) + + def fail_at_deadline(timeout_s: Optional[float]) -> None: + clock.advance(timeout_s) + task.status = TaskStatus.ERROR + + task._on_wait = fail_at_deadline + session = _make_tx_session( + [task], + timeout_s=0.25, + deadline_monotonic_s=0.25, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.transfer.time.monotonic", + clock.monotonic, + ) + + assert session.wait_complete(blocking=True) == WaitResult.FAILED + assert task.wait_calls == [0.25] + + +def test_tx_session_cancellation_observed_at_deadline_wins_over_timeout(monkeypatch) -> None: + clock = _FakeClock() + task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False) + session = _make_tx_session( + [task], + timeout_s=0.25, + deadline_monotonic_s=0.25, + ) + + def cancel_at_deadline(timeout_s: Optional[float]) -> None: + clock.advance(timeout_s) + session._terminal_status = SessionStatus.CANCELLED + + task._on_wait = cancel_at_deadline + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.transfer.time.monotonic", + clock.monotonic, + ) + + assert session.wait_complete(blocking=True) == WaitResult.FAILED + assert task.wait_calls == [0.25] + + +def test_tx_session_tasks_and_aux_share_one_deadline(monkeypatch) -> None: + clock = _FakeClock() + first_wait_durations = iter([0.5, 0.1]) + + def advance_first_task(_timeout_s: Optional[float]) -> None: + clock.advance(next(first_wait_durations)) + + first_task = _FakeTask( + TaskStatus.TRANSFERRING, + wait_result=[False, True], + on_wait=advance_first_task, + ) + second_task = _FakeTask( + TaskStatus.TRANSFERRING, + wait_result=True, + on_wait=lambda _timeout_s: clock.advance(0.2), + ) + aux_task = _FakeTask( + TaskStatus.TRANSFERRING, + wait_result=False, + on_wait=clock.advance, + ) + session = _make_tx_session( + [first_task, second_task], + need_aux=True, + aux_task=aux_task, + timeout_s=0.5, + deadline_monotonic_s=1.0, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.transfer.time.monotonic", + clock.monotonic, + ) + + assert session.wait_complete(blocking=True) == WaitResult.TIMEOUT + assert first_task.wait_calls == pytest.approx([0.5, 0.5]) + assert second_task.wait_calls == pytest.approx([0.4]) + assert aux_task.wait_calls == pytest.approx([0.2]) + + +def test_tx_session_first_send_anchors_deadline_once(monkeypatch) -> None: + clock = _FakeClock(now_s=10.0) + sender = Mock() + sender._get_req_info.return_value = {} + params = SimpleNamespace( + schedule_style="CONTEXT_FIRST", + disagg_request_id=31, + ctx_request_id=None, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.transfer.time.monotonic", + clock.monotonic, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.transfer.tensorrt_llm.bindings.global_steady_clock_now", + Mock(return_value=123), + ) + session = TxSession( + request_id=31, + params=params, + sender=sender, + timeout_s=0.25, + overall_timeout_s=2.0, + ) + + assert session._deadline_monotonic_s is None + session.send(Mock()) + assert session._deadline_monotonic_s == 12.0 + + clock.advance(0.5) + session.send(Mock()) + assert session._deadline_monotonic_s == 12.0 + assert sender.dispatch_task.call_count == 2 + session.close() + + +@pytest.mark.parametrize( + ("transfer_timeout_ms", "sender_wait_ms", "expected_timeout_s", "expected_slice_s"), + [ + (60_000, 1_000, 60.0, 1.0), + (60_000, None, 60.0, None), + ], +) +def test_transceiver_wires_separate_sender_slice_and_overall_timeout( + monkeypatch, + transfer_timeout_ms: Optional[int], + sender_wait_ms: Optional[int], + expected_timeout_s: Optional[float], + expected_slice_s: Optional[float], +) -> None: + worker = SimpleNamespace(page_table=None) + worker_constructor = Mock(return_value=worker) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.transceiver.TransferWorker", + worker_constructor, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.transceiver.create_cache_reuse_adapter", + Mock(return_value=Mock()), + ) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.transceiver.bounce_config_from_size", + Mock(return_value=None), + ) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.transceiver.torch.cuda.current_device", + Mock(return_value=0), + ) + monkeypatch.setattr( + KvCacheTransceiverV2, + "_broadcast_instance_name", + lambda _self: "ctx", + ) + monkeypatch.setattr( + KvCacheTransceiverV2, + "_broadcast_context_endpoint", + lambda _self: "endpoint", + ) + monkeypatch.setattr(KvCacheTransceiverV2, "_init_sync_policy", lambda _self: None) + monkeypatch.setattr(KvCacheTransceiverV2, "_exchange_rank_info", lambda _self: None) + mapping = SimpleNamespace( + cp_size=1, + tp_rank=0, + tp_size=1, + enable_attention_dp=False, + ) + cache_config = SimpleNamespace( + kv_transfer_timeout_ms=transfer_timeout_ms, + kv_transfer_poll_interval_ms=5_000, + kv_transfer_sender_future_timeout_ms=sender_wait_ms, + kv_cache_bounce_size_mb=0, + ) + + KvCacheTransceiverV2( + mapping=mapping, + dist=Mock(), + kv_cache_manager=SimpleNamespace(max_batch_size=4), + cache_transceiver_config=cache_config, + ) + + worker_config = worker_constructor.call_args.args[0] + assert isinstance(worker_config, TransferWorkerConfig) + assert worker_config.tx_timeout_s == expected_slice_s + assert worker_config.tx_overall_timeout_s == expected_timeout_s + assert worker_config.rx_timeout_s == expected_timeout_s + + +def test_transceiver_rejects_unset_transfer_timeout() -> None: + cache_config = SimpleNamespace( + kv_transfer_timeout_ms=None, + kv_transfer_poll_interval_ms=5_000, + kv_transfer_sender_future_timeout_ms=1_000, + ) + + with pytest.raises( + ValueError, + match="KvCacheTransceiverV2 requires a finite kv_transfer_timeout_ms", + ): + KvCacheTransceiverV2( + mapping=Mock(), + dist=Mock(), + kv_cache_manager=Mock(), + cache_transceiver_config=cache_config, + ) + + +def test_transfer_worker_passes_overall_timeout_to_tx_session(monkeypatch) -> None: + session_constructor = Mock(return_value=Mock()) + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.transfer.TxSession", + session_constructor, + ) + worker = object.__new__(TransferWorker) + worker._config = TransferWorkerConfig( + kv_cache_manager=Mock(), + device_id=0, + instance_name="ctx", + tx_timeout_s=0.25, + tx_overall_timeout_s=60.0, + ) + worker._sender = Mock() + worker._aux_buffer = Mock() + request = SimpleNamespace( + py_disaggregated_params=Mock(), + py_request_id=41, + prompt_len=128, + py_beam_width=1, + ) + + worker.create_tx_session(request) + + session_constructor.assert_called_once_with( + request_id=41, + params=request.py_disaggregated_params, + sender=worker._sender, + aux_buffer=worker._aux_buffer, + timeout_s=0.25, + prompt_len=128, + beam_width=1, + overall_timeout_s=60.0, + ) + + +def test_context_transfer_status_block_all_drains_wait_slices_before_close() -> None: + task = _FakeTask(TaskStatus.INIT, wait_result=[False, True]) + session = _make_tx_session([task]) + transceiver = _make_transceiver({15: session}, {15: _FakeRequest()}) + + completed, failed = transceiver.check_context_transfer_status(None) + + assert completed == [15] + assert failed == [] + assert task.wait_calls == [0.25, 0.25] + assert session._closed + assert 15 not in transceiver._send_sessions + + +def test_tx_session_blocking_wait_treats_cancelled_session_as_terminal() -> None: + task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False) + session = _make_tx_session([task]) + session._terminal_status = SessionStatus.CANCELLED + + assert session.wait_complete(blocking=True) == WaitResult.FAILED + assert task.wait_calls == [] + + +def test_tx_session_blocking_wait_observes_cancellation_between_slices() -> None: + task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False) session = _make_tx_session([task]) + wait = task.wait + + def cancel_during_wait(timeout: Optional[float] = None) -> bool: + result = wait(timeout) + session._terminal_status = SessionStatus.CANCELLED + return result - assert session.wait_complete() == WaitResult.TIMEOUT + task.wait = cancel_during_wait + + assert session.wait_complete(blocking=True) == WaitResult.FAILED assert task.wait_calls == [0.25] +@pytest.mark.parametrize("timeout_s", [None, 0.0, -1.0]) +def test_tx_session_blocking_wait_uses_fallback_without_positive_timeout( + timeout_s: Optional[float], +) -> None: + task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False) + session = _make_tx_session([task], timeout_s=timeout_s) + wait = task.wait + + def cancel_during_wait(timeout: Optional[float] = None) -> bool: + result = wait(timeout) + session._terminal_status = SessionStatus.CANCELLED + return result + + task.wait = cancel_during_wait + + assert session.wait_complete(blocking=True) == WaitResult.FAILED + assert task.wait_calls == [1.0] + + +def test_tx_session_blocking_wait_treats_task_failure_as_terminal() -> None: + failed_task = _FakeTask(TaskStatus.ERROR) + pending_task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=[False, True]) + session = _make_tx_session([failed_task, pending_task]) + + assert session.wait_complete(blocking=True) == WaitResult.FAILED + assert failed_task.wait_calls == [] + # A failed task event does not prove sibling physical writers quiesced, so + # precheck callers retain the wave instead of treating failure as drained. + assert pending_task.wait_calls == [] + + +def test_tx_session_blocking_wait_detects_failed_sibling_behind_pending_task() -> None: + pending_task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False) + failed_task = _FakeTask(TaskStatus.TRANSFERRING) + session = _make_tx_session([pending_task, failed_task]) + + def fail_sibling(_timeout: Optional[float]) -> None: + failed_task.status = TaskStatus.ERROR + + pending_task._on_wait = fail_sibling + + assert session.wait_complete(blocking=True) == WaitResult.FAILED + assert pending_task.wait_calls == [0.25] + assert failed_task.wait_calls == [] + + +def test_tx_session_blocking_wait_retries_aux_wait_slices() -> None: + kv_task = _FakeTask(TaskStatus.TRANSFERRED) + aux_task = _FakeTask(TaskStatus.INIT, wait_result=[False, True]) + session = _make_tx_session([kv_task], need_aux=True, aux_task=aux_task) + + assert session.wait_complete(blocking=True) == WaitResult.COMPLETED + assert kv_task.wait_calls == [] + assert aux_task.wait_calls == [0.25, 0.25] + + +def test_tx_session_blocking_aux_wait_observes_cancellation_between_slices() -> None: + kv_task = _FakeTask(TaskStatus.TRANSFERRED) + aux_task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False) + session = _make_tx_session([kv_task], need_aux=True, aux_task=aux_task) + wait = aux_task.wait + + def cancel_during_wait(timeout: Optional[float] = None) -> bool: + result = wait(timeout) + session._terminal_status = SessionStatus.CANCELLED + return result + + aux_task.wait = cancel_during_wait + + assert session.wait_complete(blocking=True) == WaitResult.FAILED + assert kv_task.wait_calls == [] + assert aux_task.wait_calls == [0.25] + + +def test_tx_session_blocking_wait_fails_missing_required_aux() -> None: + task = _FakeTask(TaskStatus.TRANSFERRED) + session = _make_tx_session([task], need_aux=True) + + assert session.wait_complete(blocking=True) == WaitResult.FAILED + assert session.status == SessionStatus.ERROR + assert isinstance(session.exception, RuntimeError) + assert task.wait_calls == [] + + +def test_generation_first_tx_session_nonblocking_missing_aux_stays_pending() -> None: + task = _FakeTask(TaskStatus.TRANSFERRED) + session = _make_tx_session([task], need_aux=True) + + assert session.wait_complete(blocking=False) is None + assert session.status == SessionStatus.KV_TRANSFERRED + assert session.exception is None + assert task.wait_calls == [] + + def test_tx_session_wait_complete_nonblocking_returns_none_without_waiting() -> None: task = _FakeTask(TaskStatus.TRANSFERRING) session = _make_tx_session([task]) @@ -395,7 +987,7 @@ def test_check_context_runs_consensus_after_a_send() -> None: transceiver._ever_had_send_session = True transceiver._ctx_need_tp_sync = True transceiver._ctx_consensus = Mock(return_value=[]) - transceiver._ctx_consensus_outcome = Mock(return_value=([], [], [], [])) + transceiver._ctx_consensus_outcome = Mock(return_value=([], [], [])) transceiver.check_context_transfer_status(0) transceiver._ctx_consensus.assert_called_once() diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 252493d1dba2..55c9cb7a6580 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -19,7 +19,9 @@ import json import os +import subprocess import sys +import types import pytest @@ -298,6 +300,126 @@ def test_use_kv_cache_manager_v2_flags(): assert pcfg.side_plan(plan, "gen")["use_kv_cache_manager_v2"] is True +def test_resolve_model_prefs_auto_requires_registered_model(monkeypatch): + monkeypatch.setattr(rp, "load_internal_apis", lambda: types.SimpleNamespace()) + monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (None, None)) + cache_cfg = types.SimpleNamespace(transceiver_runtime="PYTHON") + + with pytest.raises(RuntimeError, match="refusing to assume V1"): + rp.resolve_model_prefs(None, {"use_kv_cache_manager_v2": "auto"}, cache_cfg) + + +def test_resolve_model_prefs_auto_propagates_model_preference_failure(monkeypatch): + class FailingModel: + @classmethod + def get_preferred_kv_cache_manager_version(cls, _pretrained_config): + raise RuntimeError("model hook failed") + + def resolve_v2(_shim, model_cls, pretrained_config): + return model_cls.get_preferred_kv_cache_manager_version(pretrained_config) + + api = types.SimpleNamespace( + TorchLlmArgs=lambda **kwargs: types.SimpleNamespace(**kwargs), + resolve_kv_cache_manager_v2_auto=resolve_v2, + ) + monkeypatch.setattr(rp, "load_internal_apis", lambda: api) + monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (FailingModel, object())) + cache_cfg = types.SimpleNamespace(transceiver_runtime="PYTHON") + side = { + "use_kv_cache_manager_v2": "auto", + "parallel": {"tp": 2, "pp": 1, "cp": 1}, + } + + with pytest.raises(RuntimeError, match="V2 'auto' resolution failed.*refusing to assume V1"): + rp.resolve_model_prefs("/model", side, cache_cfg) + + +def test_resolve_model_prefs_passes_model_metadata_to_resolver(monkeypatch): + captured = [] + + class Model: + pass + + def resolve_v2(shim, model_cls, pretrained_config): + captured.append((shim, model_cls, pretrained_config)) + return True + + api = types.SimpleNamespace( + TorchLlmArgs=lambda **kwargs: types.SimpleNamespace(**kwargs), + MTPDecodingConfig=lambda **kwargs: types.SimpleNamespace(**kwargs), + resolve_kv_cache_manager_v2_auto=resolve_v2, + ) + hf_view = object() + monkeypatch.setattr(rp, "load_internal_apis", lambda: api) + monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (Model, hf_view)) + cache_cfg = types.SimpleNamespace(transceiver_runtime="PYTHON") + side = { + "use_kv_cache_manager_v2": "auto", + "parallel": {"tp": 4, "pp": 2, "cp": 3}, + "num_nextn_predict_layers": 3, + } + + assert rp.resolve_model_prefs("/model", side, cache_cfg) + resolver_args, model_cls, pretrained_config = captured.pop() + assert resolver_args.model == "/model" + assert resolver_args.tensor_parallel_size == 4 + assert resolver_args.pipeline_parallel_size == 2 + assert resolver_args.context_parallel_size == 3 + assert resolver_args.kv_cache_config == {"use_kv_cache_manager_v2": "auto"} + assert resolver_args.cache_transceiver_config is cache_cfg + assert resolver_args.speculative_config.num_nextn_predict_layers == 3 + assert model_cls is Model + assert pretrained_config is hf_view + + +def test_resolve_model_prefs_auto_propagates_resolver_failure(monkeypatch): + class Model: + pass + + def fail_resolver(_shim, _model_cls, _pretrained_config): + raise RuntimeError("resolver failed") + + api = types.SimpleNamespace( + TorchLlmArgs=lambda **kwargs: types.SimpleNamespace(**kwargs), + resolve_kv_cache_manager_v2_auto=fail_resolver, + ) + monkeypatch.setattr(rp, "load_internal_apis", lambda: api) + monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (Model, object())) + cache_cfg = types.SimpleNamespace(transceiver_runtime="PYTHON") + side = { + "use_kv_cache_manager_v2": "auto", + "parallel": {"tp": 1, "pp": 1, "cp": 1}, + } + + with pytest.raises(RuntimeError, match="V2 'auto' resolution failed.*refusing to assume V1"): + rp.resolve_model_prefs("/model", side, cache_cfg) + + +def test_resolve_model_prefs_explicit_v1_does_not_require_model(monkeypatch): + monkeypatch.setattr(rp, "load_internal_apis", lambda: types.SimpleNamespace()) + monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (None, None)) + cache_cfg = types.SimpleNamespace(transceiver_runtime="PYTHON") + + assert not rp.resolve_model_prefs(None, {"use_kv_cache_manager_v2": False}, cache_cfg) + + +def test_resolve_model_prefs_runtime_auto_failure_is_not_silently_demoted(monkeypatch): + def fail_runtime_resolver(*_args): + raise RuntimeError("model runtime hook failed") + + api = types.SimpleNamespace(resolve_transceiver_runtime_auto=fail_runtime_resolver) + monkeypatch.setattr(rp, "load_internal_apis", lambda: api) + monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (type("Model", (), {}), None)) + cache_cfg = types.SimpleNamespace(transceiver_runtime="auto") + + with pytest.raises(RuntimeError, match="may differ from serving"): + rp.resolve_model_prefs( + "/model", + {"use_kv_cache_manager_v2": False}, + cache_cfg, + ) + + def test_model_kv_shape_vocab_size(tmp_path): model_dir = tmp_path / "m" model_dir.mkdir() @@ -360,10 +482,81 @@ def test_wireup_timeout_derivation(): def _enabled_line(cfg): - lines = pcfg.precheck_prefix_lines(cfg, "e2e", "$c", "unset &&", max_world=8) + lines = pcfg.precheck_prefix_lines( + cfg, + "e2e", + "$c", + "unset &&", + max_world=8, + llm_models_root="/models", + ) return next(x for x in lines if x.startswith("export ctPrecheckEnabled")) +@pytest.mark.parametrize( + "model_root", + ( + "/models with spaces", + "/models/it's", + "/models/$HOME/$(must-not-run)", + ), +) +def test_precheck_commands_export_model_root_safely(model_root, monkeypatch): + monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) + lines = pcfg.precheck_prefix_lines( + _disagg_yaml(cache_transceiver_precheck={"enabled": True}), + "e2e", + "$config", + "unset UCX_TLS &&", + max_world=8, + llm_models_root=model_root, + ) + + commands = [line for line in lines if "pytestCommand" in line] + assert len(commands) == 2 + assert all("python3" in line for line in commands) + + script = "\n".join(lines) + '\nprintf "%s" "$LLM_MODELS_ROOT"\n' + result = subprocess.run( + ["bash"], + input=script, + text=True, + capture_output=True, + check=True, + ) + assert result.stdout == model_root + + +def test_disabled_precheck_does_not_require_or_export_model_root(monkeypatch): + monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) + + lines = pcfg.precheck_prefix_lines( + _disagg_yaml(), + "e2e", + "$config", + "unset UCX_TLS &&", + max_world=8, + ) + + assert "export ctPrecheckEnabled=0" in lines + assert not any(line.startswith("export LLM_MODELS_ROOT=") for line in lines) + + +@pytest.mark.parametrize("llm_models_root", [None, ""]) +def test_enabled_precheck_requires_model_root(monkeypatch, llm_models_root): + monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) + + with pytest.raises(ValueError, match="requires LLM_MODELS_ROOT"): + pcfg.precheck_prefix_lines( + _disagg_yaml(cache_transceiver_precheck={"enabled": True}), + "e2e", + "$config", + "unset UCX_TLS &&", + max_world=8, + llm_models_root=llm_models_root, + ) + + def test_precheck_env_kill_switch_truthy(monkeypatch): """The TRTLLM_DISAGG_CT_PRECHECK kill switch parses the usual boolean spellings. @@ -475,7 +668,7 @@ class _FakeParams: ctx_dp_rank = 0 disagg_info_endpoint = None - def _mk_runner(self, role, server_idx, plan, work_dir, monkeypatch, fail_ctx=False): + def _mk_runner(self, role, server_idx, plan, work_dir, monkeypatch): import sys import types @@ -492,8 +685,6 @@ def _mk_runner(self, role, server_idx, plan, work_dir, monkeypatch, fail_ctx=Fal calls = {"waves": 0} def ctx_run_wave(peer_idx, li, req_len, rep, wave): - if fail_ctx: - raise rp._TransferError("injected ctx failure") calls["waves"] += 1 return {p: self._FakeParams() for p in wave}, {} @@ -503,7 +694,7 @@ def ctx_run_wave(peer_idx, li, req_len, rep, wave): runner._calls = calls return runner - def _run(self, tmp_path, monkeypatch, fail_ctx_idx=None): + def _run(self, tmp_path, monkeypatch, fail_peer_idx=None): import threading monkeypatch.setenv("SLURM_JOB_ID", "777") @@ -526,17 +717,27 @@ def _run(self, tmp_path, monkeypatch, fail_ctx_idx=None): noop = lambda *a, **k: None # noqa: E731 - signal.alarm needs main thread gen = self._mk_runner("gen", 0, plan, work, monkeypatch) - ctxs = [ - self._mk_runner( - "ctx", - i, - plan, - work, - monkeypatch, - fail_ctx=(fail_ctx_idx is not None and i == fail_ctx_idx), - ) - for i in range(2) - ] + ctxs = [self._mk_runner("ctx", i, plan, work, monkeypatch) for i in range(2)] + + if fail_peer_idx is not None: + real_gen_run_peer = rp.gen_run_peer + + def fail_before_transfer(runner, peer_idx, arm, disarm): + if peer_idx != fail_peer_idx: + return real_gen_run_peer(runner, peer_idx, arm, disarm) + sock, key = rp._gen_open_session(runner, peer_idx, arm) + reason = "injected pre-transfer failure" + # Publish fail-fast before the peer handles our abort, so both + # sides deterministically classify this as the same safe, + # pre-dispatch failure rather than an ownership-fatal abort. + rp.raise_abort_flag(runner.work_dir, f"ctx_{peer_idx} TRANSFER_ERROR: {reason}") + try: + runner._leader_send_recv(sock, ("abort", reason), key) + finally: + sock.close(linger=0) + raise rp._TransferError(reason) + + monkeypatch.setattr(rp, "gen_run_peer", fail_before_transfer) failures = [] @@ -551,12 +752,18 @@ def rec(peer, exc): ] for t in threads: t.start() - rp._drive_ctx_peers( - gen, noop, noop, rp._make_peer_failure_recorder(gen, noop, {"what": "test"}) - ) - for t in threads: - t.join(timeout=60) - assert not t.is_alive(), "ctx serve thread wedged" + try: + rp._drive_ctx_peers( + gen, noop, noop, rp._make_peer_failure_recorder(gen, noop, {"what": "test"}) + ) + finally: + # Always join every peer, even when the driver raises. Asserting + # inside the loop can itself strand later peers and trip CI's + # pytest-threadleak hook. + for thread in threads: + thread.join(timeout=5) + leaked = [thread.name for thread in threads if thread.is_alive()] + assert not leaked, f"ctx serve threads wedged: {leaked}" return plan, gen, ctxs, failures def test_two_ctx_full_pass(self, tmp_path, monkeypatch): @@ -577,22 +784,23 @@ def test_two_ctx_full_pass(self, tmp_path, monkeypatch): def test_ctx_failure_last_peer(self, tmp_path, monkeypatch): # The failing pair is driven LAST: the earlier healthy peer already # completed, so there is nothing left to fail-fast/skip. - plan, gen, ctxs, failures = self._run(tmp_path, monkeypatch, fail_ctx_idx=1) + plan, gen, ctxs, failures = self._run(tmp_path, monkeypatch, fail_peer_idx=1) # gen side: healthy peer unaffected, failing peer gets a clear verdict by_peer = {c["peer"]: c["status"] for c in gen.recorder.cases} assert by_peer == {"ctx_0": "PASS", "ctx_1": "TRANSFER_ERROR"} - # ctx_1's own serve loop surfaced the failure (its peer is gen_0) - assert ("gen_0", "_TransferError") in failures + assert not failures # ctx_0 served its full schedule and got the deferred done assert [c["status"] for c in ctxs[0].recorder.cases] == ["PASS"] + assert [c["status"] for c in ctxs[1].recorder.cases] == ["SKIP"] def test_fail_fast_skips_remaining(self, tmp_path, monkeypatch): # The FIRST-driven pair fails: the remaining pair must be skipped # (not tested against a fabric already known bad), and told to abort # so it tears down promptly instead of waiting out its handshake alarm. - plan, gen, ctxs, failures = self._run(tmp_path, monkeypatch, fail_ctx_idx=0) + plan, gen, ctxs, failures = self._run(tmp_path, monkeypatch, fail_peer_idx=0) by_peer = {c["peer"]: c["status"] for c in gen.recorder.cases} assert by_peer == {"ctx_0": "TRANSFER_ERROR", "ctx_1": "SKIP"} + assert not failures # ctx_1 never ran a single transfer wave: fail-fast reached it first. assert ctxs[1]._calls["waves"] == 0 # ctx_1 recorded a non-failing SKIP (its driver aborted the session). diff --git a/tests/unittest/others/test_cache_transceiver_precheck_run.py b/tests/unittest/others/test_cache_transceiver_precheck_run.py index d6b736a4d4bd..1520c9179e9e 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_run.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_run.py @@ -80,6 +80,55 @@ def test_seed_for_deterministic_and_distinct(): assert all(0 <= s <= 0x7FFFFFFF for s in seeds) +@pytest.mark.parametrize(("prompt_len", "expected_blocks"), ((1024, 8), (7408, 58))) +def test_request_block_views_excludes_untransferred_speculative_page(prompt_len, expected_blocks): + """V2's reserved MTP tokens must not expand the verified transfer range.""" + tokens_per_block = 128 + num_allocated = (prompt_len + 2 + tokens_per_block - 1) // tokens_per_block + allocated = [-1] + list(range(num_allocated)) + buffer = object() + + def get_batch_cache_indices(request_ids, layer_idx): + assert request_ids == [7] + assert layer_idx == 4 + return [allocated] + + def get_buffers(global_layer, kv_layout): + assert global_layer == 4 + assert kv_layout == "HND" + return buffer + + kvm = types.SimpleNamespace( + tokens_per_block=tokens_per_block, + pp_layers=[4], + get_batch_cache_indices=get_batch_cache_indices, + get_buffers=get_buffers, + ) + + views = list(rp._request_block_views(kvm, rid=7, prompt_len=prompt_len)) + + assert views == [(4, buffer, list(range(expected_blocks)))] + + +@pytest.mark.parametrize("available_blocks", [0, 1]) +def test_request_block_views_rejects_prompt_under_allocation(available_blocks): + allocated = [-1] + list(range(available_blocks)) + kvm = types.SimpleNamespace( + tokens_per_block=128, + pp_layers=[4], + get_batch_cache_indices=lambda _request_ids, layer_idx: [allocated], + get_buffers=lambda *_args, **_kwargs: pytest.fail( + "buffer lookup must not occur after under-allocation" + ), + ) + + with pytest.raises( + rp._TransferError, + match=rf"rid=7 layer=4: required=2 available={available_blocks}", + ): + list(rp._request_block_views(kvm, rid=7, prompt_len=256)) + + # --------------------------------------------------------------------------- # # HMAC control-channel wire format # --------------------------------------------------------------------------- # @@ -327,6 +376,579 @@ def test_timeout_budgets(): assert rp.wave_timeout_s(plan, 1, 0) == 180 +# --------------------------------------------------------------------------- # +# Model preference resolution +# --------------------------------------------------------------------------- # +def test_resolve_model_prefs_allows_registered_class_without_preference_hook(monkeypatch): + model_cls = type("ModelWithoutPreferenceHook", (), {}) + cache_cfg = types.SimpleNamespace(transceiver_runtime="CPP") + calls = [] + + def resolve_v2(shim, resolved_model_cls, pretrained_config): + calls.append((shim, resolved_model_cls, pretrained_config)) + return False + + hf_view = object() + monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (model_cls, hf_view)) + monkeypatch.setattr( + rp, + "load_internal_apis", + lambda: types.SimpleNamespace( + TorchLlmArgs=lambda **kwargs: types.SimpleNamespace(**kwargs), + resolve_kv_cache_manager_v2_auto=resolve_v2, + ), + ) + + use_v2 = rp.resolve_model_prefs( + "/models/example", + { + "use_kv_cache_manager_v2": "auto", + "parallel": {"tp": 1, "pp": 1, "cp": 1}, + }, + cache_cfg, + ) + + assert use_v2 is False + assert len(calls) == 1 + assert calls[0][1:] == (model_cls, hf_view) + + +# --------------------------------------------------------------------------- # +# Transfer ownership +# --------------------------------------------------------------------------- # +def _ctx_finish_runner(monkeypatch, check_status): + events = [] + monkeypatch.setitem( + sys.modules, + "tensorrt_llm", + types.SimpleNamespace(logger=types.SimpleNamespace(info=lambda *_args, **_kwargs: None)), + ) + runner = object.__new__(rp.PrecheckRunner) + runner.xcvr = types.SimpleNamespace(check_context_transfer_status=check_status) + runner.llm_request_state = types.SimpleNamespace(DISAGG_TRANS_ERROR="error") + runner.server_idx = 0 + runner.rank = 0 + runner._consensus_error = lambda err: None if err is None else repr(err) + runner._free_all = lambda reqs: events.append(("free", sorted(reqs))) + return runner, events + + +def test_ctx_finish_wave_frees_only_after_block_all_returns_every_request(monkeypatch): + events = [] + + def check_status(at_least_request_num): + events.append(("block_all", at_least_request_num)) + return [101, 102], [] + + runner, free_events = _ctx_finish_runner(monkeypatch, check_status) + reqs = { + 0: types.SimpleNamespace(py_request_id=101, state="in_progress"), + 1: types.SimpleNamespace(py_request_id=102, state="in_progress"), + } + runner._free_all = lambda owned: events.append(("free", sorted(owned))) + + runner.ctx_finish_wave(reqs) + + assert events == [("block_all", None), ("free", [0, 1])] + assert free_events == [] + + +def test_ctx_finish_wave_retains_pages_when_block_all_omits_request(monkeypatch): + runner, events = _ctx_finish_runner(monkeypatch, lambda _n: ([101], [])) + reqs = { + 0: types.SimpleNamespace(py_request_id=101, state="in_progress"), + 1: types.SimpleNamespace(py_request_id=102, state="in_progress"), + } + + with pytest.raises(rp._FatalTransferError, match="block-all returned before terminal"): + runner.ctx_finish_wave(reqs) + + assert events == [] + + +def test_ctx_finish_wave_retains_pages_when_block_all_raises(monkeypatch): + def check_status(_n): + raise RuntimeError("interrupted") + + runner, events = _ctx_finish_runner(monkeypatch, check_status) + reqs = {0: types.SimpleNamespace(py_request_id=101, state="in_progress")} + + with pytest.raises(rp._FatalTransferError, match="interrupted"): + runner.ctx_finish_wave(reqs) + + assert events == [] + + +def test_ctx_finish_wave_retains_pages_when_request_failed(monkeypatch): + runner, events = _ctx_finish_runner(monkeypatch, lambda _n: ([101], [102])) + reqs = { + 0: types.SimpleNamespace(py_request_id=101, state="in_progress"), + 1: types.SimpleNamespace(py_request_id=102, state="error"), + } + + with pytest.raises(rp._FatalTransferError, match=r"ctx transfer failed for pairs \[1\]"): + runner.ctx_finish_wave(reqs) + + assert events == [] + + +def test_ctx_finish_wave_does_not_free_when_peer_rank_is_unsafe(monkeypatch): + runner, events = _ctx_finish_runner(monkeypatch, lambda _n: ([101, 102], [])) + reqs = { + 0: types.SimpleNamespace(py_request_id=101, state="in_progress"), + 1: types.SimpleNamespace(py_request_id=102, state="in_progress"), + } + runner._consensus_error = lambda _err: "rank 1 did not prove completion" + + with pytest.raises(rp._FatalTransferError, match="rank 1 did not prove completion"): + runner.ctx_finish_wave(reqs) + + assert events == [] + + +def test_ctx_finish_wave_consensus_exception_is_fatal(monkeypatch): + runner, events = _ctx_finish_runner(monkeypatch, lambda _n: ([101], [])) + reqs = {0: types.SimpleNamespace(py_request_id=101, state="in_progress")} + + def fail_consensus(_error): + raise RuntimeError("MPI consensus failed") + + runner._consensus_error = fail_consensus + + with pytest.raises(rp._FatalTransferError, match="MPI consensus failed"): + runner.ctx_finish_wave(reqs) + + assert events == [] + + +def test_ctx_finish_wave_timeout_skips_consensus_and_free(monkeypatch): + def check_status(_n): + raise rp._Timeout("deadline") + + runner, events = _ctx_finish_runner(monkeypatch, check_status) + consensus_calls = [] + runner._consensus_error = lambda err: consensus_calls.append(err) + reqs = {0: types.SimpleNamespace(py_request_id=101, state="in_progress")} + + with pytest.raises(rp._Timeout, match="deadline"): + runner.ctx_finish_wave(reqs) + + assert consensus_calls == [] + assert events == [] + + +def test_ctx_run_wave_setup_error_retains_allocated_pages(monkeypatch): + monkeypatch.setitem( + sys.modules, + "tensorrt_llm", + types.SimpleNamespace(logger=types.SimpleNamespace(info=lambda *_args, **_kwargs: None)), + ) + monkeypatch.setattr( + rp, + "make_request", + lambda _is_ctx, rid, _req_len, _runtime: types.SimpleNamespace( + py_request_id=rid, context_phase_params=None + ), + ) + monkeypatch.setattr(rp, "add_sequence", lambda *_args: None) + monkeypatch.setattr(rp, "fill_request", lambda *_args: None) + + calls = {"send": 0, "free": 0} + + def respond(_req): + calls["send"] += 1 + if calls["send"] == 2: + raise RuntimeError("injected setup failure") + + runner = object.__new__(rp.PrecheckRunner) + runner.runtime = "PYTHON" + runner.kvm = object() + runner.use_v2 = True + runner.server_idx = 0 + runner.rank = 0 + runner.is_leader = True + runner.side = {"parallel": {"enable_attention_dp": False}} + runner.mapping = types.SimpleNamespace(pp_rank=0) + runner.xcvr = types.SimpleNamespace(respond_and_send_async=respond) + runner.comm = types.SimpleNamespace( + gather=lambda obj, root=0: [obj], + bcast=lambda obj, root=0: obj, + ) + runner._owned = lambda _wave: [0, 1] + runner._pair_rid = lambda _peer, _li, _rep, pair: 101 + pair + runner._consensus_error = lambda err: None if err is None else repr(err) + runner._free_all = lambda _reqs: calls.__setitem__("free", calls["free"] + 1) + + with pytest.raises(rp._FatalTransferError, match="injected setup failure"): + runner.ctx_run_wave(0, 0, 64, 0, [0, 1]) + + assert calls == {"send": 2, "free": 0} + + +def test_ctx_run_wave_post_dispatch_collective_error_is_fatal(monkeypatch): + monkeypatch.setitem( + sys.modules, + "tensorrt_llm", + types.SimpleNamespace(logger=types.SimpleNamespace(info=lambda *_args, **_kwargs: None)), + ) + monkeypatch.setattr( + rp, + "make_request", + lambda _is_ctx, rid, _req_len, _runtime: types.SimpleNamespace( + py_request_id=rid, context_phase_params=object() + ), + ) + monkeypatch.setattr(rp, "add_sequence", lambda *_args: None) + monkeypatch.setattr(rp, "fill_request", lambda *_args: None) + + runner = object.__new__(rp.PrecheckRunner) + runner.runtime = "PYTHON" + runner.kvm = object() + runner.use_v2 = True + runner.server_idx = 0 + runner.rank = 0 + runner.is_leader = True + runner.side = {"parallel": {"enable_attention_dp": False}} + runner.mapping = types.SimpleNamespace(pp_rank=0) + runner.xcvr = types.SimpleNamespace(respond_and_send_async=lambda _req: None) + runner.comm = types.SimpleNamespace( + gather=lambda _obj, root=0: (_ for _ in ()).throw(RuntimeError("MPI gather failed")), + bcast=lambda obj, root=0: obj, + ) + runner._owned = lambda _wave: [0] + runner._pair_rid = lambda *_args: 101 + runner._consensus_error = lambda _err: None + + with pytest.raises(rp._FatalTransferError, match="MPI gather failed"): + runner.ctx_run_wave(0, 0, 64, 0, [0]) + + +def _gen_run_wave_runner(monkeypatch, outcome): + requests = {} + events = [] + states = types.SimpleNamespace( + DISAGG_GENERATION_TRANS_COMPLETE="complete", + DISAGG_TRANS_ERROR="error", + ) + monkeypatch.setitem( + sys.modules, + "torch", + types.SimpleNamespace( + cuda=types.SimpleNamespace(synchronize=lambda: events.append("cuda_sync")) + ), + ) + monkeypatch.setitem( + sys.modules, + "tensorrt_llm", + types.SimpleNamespace(logger=types.SimpleNamespace(info=lambda *_args, **_kwargs: None)), + ) + + def make_request(_is_ctx, rid, _req_len, _runtime, ctx_params=None): + req = types.SimpleNamespace(py_request_id=rid, state="in_progress") + requests[rid] = req + return req + + monkeypatch.setattr(rp, "make_request", make_request) + monkeypatch.setattr(rp, "add_sequence", lambda *_args: None) + + def check_status(_at_least_request_num): + completed, failed, cancelled = outcome(requests) + for rid in completed: + requests[rid].state = states.DISAGG_GENERATION_TRANS_COMPLETE + for rid in failed: + requests[rid].state = states.DISAGG_TRANS_ERROR + return completed, failed, [requests[rid] for rid in cancelled] + + runner = object.__new__(rp.PrecheckRunner) + runner.runtime = "PYTHON" + runner.kvm = object() + runner.use_v2 = True + runner.server_idx = 0 + runner.rank = 0 + runner.side = {"parallel": {"enable_attention_dp": False}} + runner.mapping = types.SimpleNamespace(pp_rank=0) + runner.llm_request_state = states + runner.plan = {"verify_data": False, "warmup_requests": 0} + runner.xcvr = types.SimpleNamespace( + request_and_receive_async=lambda _req: None, + check_gen_transfer_status=check_status, + ) + runner.comm = types.SimpleNamespace(allgather=lambda value: [value]) + runner._owned = lambda _wave: [0, 1] + runner._pair_rid = lambda _peer, _li, _rep, pair: 101 + pair + runner._consensus_error = lambda err: None if err is None else repr(err) + runner._free_all = lambda owned: events.append(("free", sorted(owned))) + return runner, events + + +def test_gen_run_wave_frees_only_after_every_python_receive_completes(monkeypatch): + runner, events = _gen_run_wave_runner(monkeypatch, lambda _reqs: ([101, 102], [], [])) + + ok, detail = runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()}) + + assert ok and not detail + assert events == ["cuda_sync", ("free", [0, 1])] + + +def test_gen_run_wave_checks_python_status_on_empty_owner_rank(monkeypatch): + calls = [] + + def outcome(requests): + calls.append(dict(requests)) + return [], [], [] + + runner, events = _gen_run_wave_runner(monkeypatch, outcome) + runner._owned = lambda _wave: [] + + ok, detail = runner.gen_run_wave(0, 0, 64, 0, [0, 1], {}) + + assert ok and not detail + assert calls == [{}] + assert events == ["cuda_sync", ("free", [])] + + +def test_gen_run_wave_setup_error_retains_allocated_pages(monkeypatch): + runner, events = _gen_run_wave_runner(monkeypatch, lambda _reqs: ([], [], [])) + calls = 0 + + def receive(_req): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("injected setup failure") + + runner.xcvr.request_and_receive_async = receive + + with pytest.raises(rp._FatalTransferError, match="injected setup failure"): + runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()}) + + assert calls == 2 + assert events == [] + + +@pytest.mark.parametrize( + ("outcome", "message"), + ( + (lambda _reqs: ([101], [102], []), r"failed=\[102\]"), + (lambda _reqs: ([101], [], [102]), r"cancelled=\[102\]"), + (lambda _reqs: ([101], [], []), r"missing=\[102\]"), + ), +) +def test_gen_run_wave_retains_pages_without_all_successes(monkeypatch, outcome, message): + runner, events = _gen_run_wave_runner(monkeypatch, outcome) + + with pytest.raises(rp._FatalTransferError, match=message): + runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()}) + + assert events == [] + + +def test_gen_run_wave_does_not_free_when_peer_rank_is_unsafe(monkeypatch): + runner, events = _gen_run_wave_runner(monkeypatch, lambda _reqs: ([101, 102], [], [])) + consensus_calls = [] + + def consensus(error): + consensus_calls.append(error) + # First call covers receive setup; the second is the transfer-release + # proof and models another rank reporting a nonterminal request. + return None if len(consensus_calls) == 1 else "rank 1 did not prove completion" + + runner._consensus_error = consensus + + with pytest.raises(rp._FatalTransferError, match="rank 1 did not prove completion"): + runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()}) + + assert consensus_calls == [None, None] + assert events == ["cuda_sync"] + + +def test_gen_run_wave_transfer_consensus_exception_is_fatal(monkeypatch): + runner, events = _gen_run_wave_runner(monkeypatch, lambda _reqs: ([101, 102], [], [])) + consensus_calls = 0 + + def consensus(_error): + nonlocal consensus_calls + consensus_calls += 1 + if consensus_calls == 2: + raise RuntimeError("MPI transfer consensus failed") + return None + + runner._consensus_error = consensus + + with pytest.raises(rp._FatalTransferError, match="MPI transfer consensus failed"): + runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()}) + + assert events == ["cuda_sync"] + + +def test_gen_run_wave_timeout_skips_transfer_consensus_and_free(monkeypatch): + def timeout(_reqs): + raise rp._Timeout("deadline") + + runner, events = _gen_run_wave_runner(monkeypatch, timeout) + consensus_calls = [] + + def consensus(error): + consensus_calls.append(error) + return None + + runner._consensus_error = consensus + + with pytest.raises(rp._Timeout, match="deadline"): + runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()}) + + # Receive setup reaches consensus before block-all. The timeout itself + # must propagate directly into the process-fatal path. + assert consensus_calls == [None] + assert events == [] + + +def test_hard_abort_unquiesced_persists_verdict_before_abort(monkeypatch): + class AbortSentinel(Exception): + pass + + events = [] + recorder = types.SimpleNamespace( + record=lambda *args: events.append(("record", *args)), + finalize=lambda **kwargs: events.append(("finalize", kwargs)), + ) + runner = types.SimpleNamespace( + is_leader=True, + recorder=recorder, + work_dir="/tmp/precheck", + use_v2=True, + runtime="PYTHON", + comm=object(), + xcvr=types.SimpleNamespace(shutdown=lambda: events.append(("shutdown",))), + ) + monkeypatch.setattr(rp.signal, "alarm", lambda seconds: events.append(("alarm", seconds))) + monkeypatch.setattr( + rp, + "raise_abort_flag", + lambda work_dir, reason: events.append(("abort_flag", work_dir, reason)), + ) + monkeypatch.setattr( + rp, + "_coordinate_abort_after_leader_flush", + lambda comm: events.append(("coordinate", comm)), + ) + + def abort_process(comm): + events.append(("abort", comm)) + raise AbortSentinel + + monkeypatch.setattr(rp, "_hard_abort_process", abort_process) + + with pytest.raises(AbortSentinel): + rp._hard_abort_unquiesced( + runner, + {"what": "ctx wave"}, + rp._FatalTransferError("completion unproven"), + ) + + assert events == [ + ("alarm", 0), + ("record", "ctx wave", 0, "TRANSFER_ERROR", "completion unproven"), + ( + "abort_flag", + "/tmp/precheck", + "ctx wave TRANSFER_ERROR: completion unproven", + ), + ( + "finalize", + { + "extra": { + "kv_cache_manager": "V2", + "transceiver_runtime": "PYTHON", + } + }, + ), + ("coordinate", runner.comm), + ("abort", runner.comm), + ] + + +def test_hard_abort_unquiesced_still_aborts_when_status_write_fails(monkeypatch): + class AbortSentinel(Exception): + pass + + runner = types.SimpleNamespace( + is_leader=True, + recorder=types.SimpleNamespace( + record=lambda *_args: (_ for _ in ()).throw(OSError("disk full")), + ), + work_dir="/tmp/precheck", + use_v2=True, + runtime="PYTHON", + comm=object(), + ) + monkeypatch.setattr(rp.signal, "alarm", lambda _seconds: None) + monkeypatch.setattr(rp, "_coordinate_abort_after_leader_flush", lambda _comm: None) + monkeypatch.setattr( + rp, + "_hard_abort_process", + lambda _comm: (_ for _ in ()).throw(AbortSentinel()), + ) + + with pytest.raises(AbortSentinel): + rp._hard_abort_unquiesced( + runner, + {"what": "ctx wave"}, + rp._FatalTransferError("completion unproven"), + ) + + +def test_ctx_peer_loop_reraises_fatal_transfer(monkeypatch): + fatal = rp._FatalTransferError("not quiesced") + runner = types.SimpleNamespace( + side={"num_peers": 1}, + is_leader=False, + recorder=types.SimpleNamespace(record=lambda *_args: None), + ) + recorded = [] + + def serve(*_args, **_kwargs): + raise fatal + + monkeypatch.setattr(rp, "ctx_serve_peer", serve) + + with pytest.raises(rp._FatalTransferError, match="not quiesced"): + rp._serve_gen_peers( + runner, + plan={}, + arm=lambda *_args, **_kwargs: None, + disarm=lambda: None, + record_peer_failure=lambda *args: recorded.append(args), + ) + + assert recorded == [] + + +def test_gen_peer_loop_reraises_fatal_transfer(monkeypatch): + fatal = rp._FatalTransferError("not quiesced") + runner = types.SimpleNamespace( + side={"num_peers": 1}, + recorder=types.SimpleNamespace(record=lambda *_args: None), + ) + recorded = [] + monkeypatch.setattr(rp, "_consensus_abort_reason", lambda _runner: None) + + def run_peer(*_args, **_kwargs): + raise fatal + + monkeypatch.setattr(rp, "gen_run_peer", run_peer) + + with pytest.raises(rp._FatalTransferError, match="not quiesced"): + rp._drive_ctx_peers( + runner, + arm=lambda *_args, **_kwargs: None, + disarm=lambda: None, + record_peer_failure=lambda *args: recorded.append(args), + ) + + assert recorded == [] + + # --------------------------------------------------------------------------- # # Internal-API contract (imports tensorrt_llm; no GPU work) # --------------------------------------------------------------------------- # @@ -389,15 +1011,15 @@ def test_kv_cache_manager_ctor_kwargs(self, api, manager_attr): def test_serving_resolvers(self, api): import inspect - # The driver calls the serving resolver with a shim, model class, and - # pretrained config. Only the shim is required. + # The driver calls the serving resolver with real LLM args, the model + # class, and its pretrained config. v2 = inspect.signature(api.resolve_kv_cache_manager_v2_auto).parameters assert list(v2)[:3] == ["llm_args", "model_cls", "pretrained_config"] assert all(p.default is not inspect.Parameter.empty for p in list(v2.values())[1:]) rt = inspect.signature(api.resolve_transceiver_runtime_auto).parameters assert list(rt)[:1] == ["llm_args"] and len(rt) >= 3 - def test_model_preference_resolver_shim_supports_v2(self, api, monkeypatch): + def test_model_preference_resolver_supports_v2(self, api, monkeypatch): class _PreferV2: @classmethod def get_preferred_kv_cache_manager_version(cls, pretrained_config=None): @@ -413,7 +1035,10 @@ def get_preferred_kv_cache_manager_version(cls, pretrained_config=None): assert rp.resolve_model_prefs( "/tmp/dummy_model", - {"use_kv_cache_manager_v2": "auto"}, + { + "use_kv_cache_manager_v2": "auto", + "parallel": {"tp": 1, "pp": 1, "cp": 1}, + }, cache_cfg, ) @@ -442,6 +1067,8 @@ def test_hang_detector_surface(self, api): def test_config_constructors(self, api): cache_cfg = api.CacheTransceiverConfig(backend="UCX", max_tokens_in_buffer=1024) assert hasattr(cache_cfg, "transceiver_runtime") + llm_args = api.TorchLlmArgs(model="/tmp/model", tensor_parallel_size=2) + assert llm_args.tensor_parallel_size == 2 api.KvCacheConfigCpp(max_tokens=64, enable_block_reuse=False) api.MTPDecodingConfig(num_nextn_predict_layers=1) api.Mapping( diff --git a/tests/unittest/scripts/test_perf_submit.py b/tests/unittest/scripts/test_perf_submit.py index 3a824627abbd..de81570e14fd 100644 --- a/tests/unittest/scripts/test_perf_submit.py +++ b/tests/unittest/scripts/test_perf_submit.py @@ -23,8 +23,9 @@ from pytest_split.algorithms import LeastDurationAlgorithm REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent +CI_SUBMIT_PATH = REPO_ROOT / "jenkins" / "scripts" / "perf" / "submit.py" SUBMIT_PATHS = ( - REPO_ROOT / "jenkins" / "scripts" / "perf" / "submit.py", + CI_SUBMIT_PATH, REPO_ROOT / "jenkins" / "scripts" / "perf" / "local" / "submit.py", ) EXAMPLE_SUBMIT_PATH = REPO_ROOT / "examples" / "disaggregated" / "slurm" / "benchmark" / "submit.py" @@ -53,6 +54,11 @@ def submit_module(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatc return _load_module(request.param, monkeypatch) +@pytest.fixture +def ci_submit_module(monkeypatch: pytest.MonkeyPatch) -> ModuleType: + return _load_module(CI_SUBMIT_PATH, monkeypatch) + + @pytest.fixture def example_submit_module(monkeypatch: pytest.MonkeyPatch) -> ModuleType: return _load_module(EXAMPLE_SUBMIT_PATH, monkeypatch) @@ -72,11 +78,6 @@ def test_get_benchmark_config_accepts_positive_integer(submit_module: ModuleType assert benchmark_config["concurrency"] == int(concurrency) -@pytest.fixture -def ci_submit_module(monkeypatch: pytest.MonkeyPatch) -> ModuleType: - return _load_module(SUBMIT_PATHS[0], monkeypatch) - - def _select_ci_test_case_line( ci_submit_module: ModuleType, tmp_path: Path, @@ -273,3 +274,62 @@ def test_ci_submit_rejects_missing_pytest_split_durations( pytest_options="--splits 1 --group 1 --durations-path /remote/.test_durations", split_group=1, ) + + +@pytest.mark.parametrize( + ("assignment", "expected"), + ( + ("LLM_MODELS_ROOT=/models", "/models"), + ("LLM_MODELS_ROOT='/models with spaces'", "/models with spaces"), + ("LLM_MODELS_ROOT=/models/cache=production", "/models/cache=production"), + ), +) +def test_extract_pytest_command_env(ci_submit_module: ModuleType, assignment: str, expected: str): + lines = [f'export pytestCommand="LLM_ROOT=/src {assignment} COLUMNS=300 pytest -vv"'] + + assert ci_submit_module.extract_pytest_command_env(lines, "LLM_MODELS_ROOT") == expected + + +def test_extract_pytest_command_env_rejects_missing_leading_assignment( + ci_submit_module: ModuleType, +): + lines = ['export pytestCommand="LLM_ROOT=/src pytest LLM_MODELS_ROOT=/too-late"'] + + with pytest.raises(ValueError, match="does not set leading environment variable"): + ci_submit_module.extract_pytest_command_env(lines, "LLM_MODELS_ROOT") + + +def test_extract_pytest_command_env_rejects_malformed_export(ci_submit_module: ModuleType): + lines = ['export pytestCommand="LLM_ROOT=/src LLM_MODELS_ROOT=/models pytest'] + + with pytest.raises(ValueError, match="cannot parse exported pytestCommand"): + ci_submit_module.extract_pytest_command_env(lines, "LLM_MODELS_ROOT") + + +def test_resolve_llm_models_root_falls_back_to_submitter_env( + ci_submit_module: ModuleType, monkeypatch +): + monkeypatch.setenv("LLM_MODELS_ROOT", "/models/from-env") + lines = ['export pytestCommand="LLM_ROOT=/src pytest LLM_MODELS_ROOT=/too-late"'] + + assert ci_submit_module._resolve_llm_models_root(lines) == "/models/from-env" + + +def test_resolve_llm_models_root_explains_both_missing_sources( + ci_submit_module: ModuleType, monkeypatch +): + monkeypatch.delenv("LLM_MODELS_ROOT", raising=False) + lines = ['export pytestCommand="LLM_ROOT=/src pytest"'] + + with pytest.raises(ValueError, match="getPytestBaseCommandLine in L0_Test.groovy"): + ci_submit_module._resolve_llm_models_root(lines) + + +def test_resolve_llm_models_root_does_not_mask_malformed_command( + ci_submit_module: ModuleType, monkeypatch +): + monkeypatch.setenv("LLM_MODELS_ROOT", "/models/from-env") + lines = ['export pytestCommand="LLM_ROOT=/src pytest'] + + with pytest.raises(ValueError, match="cannot parse exported pytestCommand"): + ci_submit_module._resolve_llm_models_root(lines)