From 02126b0bbcb62b0e9402c9c8543af46cc679b6f3 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 22 Jun 2026 12:43:33 -0700 Subject: [PATCH 01/14] [None][feat] WideEP FT: add AlltoAll watchdog Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 366 ++++++++++++++++++ .../_torch/distributed/moe_alltoall.py | 87 ++++- .../communication/nvlink_one_sided.py | 78 +++- .../_torch/modules/test_alltoall_watchdog.py | 229 +++++++++++ 4 files changed, 756 insertions(+), 4 deletions(-) create mode 100644 tensorrt_llm/_torch/alltoall_watchdog.py create mode 100644 tests/unittest/_torch/modules/test_alltoall_watchdog.py diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py new file mode 100644 index 000000000000..7c35669e29ee --- /dev/null +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -0,0 +1,366 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Host-side watchdog for MoE AlltoAll completion flags. + +The NVLinkOneSided kernels signal each collective by writing the current +``flag_val`` into the rank-local completion flag table. A dead peer in the +silent-spin failure mode never writes its slot, so this watchdog polls the same +table from a CPU thread and reports peers whose flags do not reach the expected +generation before a bounded timeout. +""" + +from __future__ import annotations + +import threading +import time +from collections import deque +from dataclasses import dataclass +from typing import Callable, Deque, Mapping, Optional, Protocol, Sequence + +import torch + +from tensorrt_llm.logger import logger as tllm_logger + + +class CompletionFlagReader(Protocol): + """Reads one phase's rank-local completion flag row.""" + + def read_completion_flags(self, phase: str) -> Sequence[int]: + """Return ``ep_size`` flag values for ``phase``.""" + + +class EPGroupHealthLike(Protocol): + """Subset of EPGroupHealth used by the watchdog.""" + + def get_mask(self) -> int: + """Return the active-rank bitmask.""" + + def mark_failed(self, rank: int) -> bool: + """Mark ``rank`` failed and return whether state changed.""" + + +@dataclass(frozen=True) +class AlltoAllWatchdogTimeout: + """Details emitted when an AlltoAll phase times out.""" + + phase: str + expected_flag: int + observed_flags: tuple[int, ...] + missing_ranks: tuple[int, ...] + marked_failed_ranks: tuple[int, ...] + elapsed_s: float + + +@dataclass(frozen=True) +class _CollectiveWatch: + phase: str + expected_flag: int + active_mask: int + start_s: float + + +class _TorchCompletionFlagReader: + """Completion-flag reader backed by the MoE AlltoAll workspace tensor.""" + + def __init__( + self, + workspace: torch.Tensor, + ep_rank: int, + ep_size: int, + dispatch_completion_flags_offset: int, + combine_completion_flags_offset: int, + ) -> None: + if workspace.dim() != 2: + raise ValueError("workspace must be a 2D tensor [ep_size, size_per_rank]") + if not 0 <= ep_rank < ep_size: + raise ValueError(f"ep_rank must be in [0, {ep_size}), got {ep_rank}") + if workspace.size(0) != ep_size: + raise ValueError( + f"workspace first dimension must equal ep_size={ep_size}, got {workspace.size(0)}" + ) + self._workspace = workspace + self._ep_rank = ep_rank + self._ep_size = ep_size + self._offsets = { + "dispatch": int(dispatch_completion_flags_offset), + "combine": int(combine_completion_flags_offset), + } + + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + offset = self._offsets[phase] + end = offset + self._ep_size * 4 + flags = self._workspace[self._ep_rank, offset:end].view(torch.int32) + if flags.device.type != "cpu": + flags = flags.detach().cpu() + return tuple(int(v) for v in flags.tolist()) + + +class AlltoAllWatchdog: + """Background host thread that watches AlltoAll completion flags. + + The watchdog is intentionally opt-in. Callers queue phases with + :meth:`watch`; the thread polls them in FIFO order so a queued combine cannot + hide a still-spinning dispatch. + """ + + VALID_PHASES = frozenset({"dispatch", "combine"}) + + def __init__( + self, + *, + ep_size: int, + ep_rank: int, + completion_reader: CompletionFlagReader, + timeout_s: float, + poll_interval_s: float = 0.05, + health: Optional[EPGroupHealthLike] = None, + on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, + ) -> None: + if ep_size <= 0: + raise ValueError(f"ep_size must be > 0, got {ep_size}") + if not 0 <= ep_rank < ep_size: + raise ValueError(f"ep_rank must be in [0, {ep_size}), got {ep_rank}") + if timeout_s <= 0: + raise ValueError(f"timeout_s must be > 0, got {timeout_s}") + if poll_interval_s <= 0: + raise ValueError(f"poll_interval_s must be > 0, got {poll_interval_s}") + + self._ep_size = int(ep_size) + self._ep_rank = int(ep_rank) + self._completion_reader = completion_reader + self._timeout_s = float(timeout_s) + self._poll_interval_s = float(poll_interval_s) + self._health = health + self._on_timeout = on_timeout + + self._cv = threading.Condition() + self._queue: Deque[_CollectiveWatch] = deque() + self._stopping = False + self._thread: threading.Thread | None = None + self._last_error: BaseException | None = None + + @classmethod + def from_workspace( + cls, + *, + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + ep_size: int, + timeout_s: float, + poll_interval_s: float = 0.05, + health: Optional[EPGroupHealthLike] = None, + on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, + ) -> "AlltoAllWatchdog": + """Build a watchdog from the MoE AlltoAll workspace and metainfo.""" + dispatch_offset = int( + metainfo[metainfo_index["DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX"]].item() + ) + combine_offset = int( + metainfo[metainfo_index["COMBINE_COMPLETION_FLAGS_OFFSET_INDEX"]].item() + ) + reader = _TorchCompletionFlagReader( + workspace=workspace, + ep_rank=ep_rank, + ep_size=ep_size, + dispatch_completion_flags_offset=dispatch_offset, + combine_completion_flags_offset=combine_offset, + ) + return cls( + ep_size=ep_size, + ep_rank=ep_rank, + completion_reader=reader, + timeout_s=timeout_s, + poll_interval_s=poll_interval_s, + health=health, + on_timeout=on_timeout, + ) + + @property + def last_error(self) -> BaseException | None: + """Return the last polling-thread error, if any.""" + with self._cv: + return self._last_error + + def start(self) -> None: + """Start the background polling thread. Idempotent.""" + with self._cv: + if self._thread is not None and self._thread.is_alive(): + return + self._stopping = False + self._thread = threading.Thread( + target=self._run, + name=f"AlltoAllWatchdog-rank{self._ep_rank}", + daemon=True, + ) + self._thread.start() + + def stop(self, timeout_s: float | None = None) -> None: + """Stop the polling thread and wait for it to exit.""" + with self._cv: + self._stopping = True + self._queue.clear() + self._cv.notify_all() + thread = self._thread + if thread is not None: + thread.join(timeout=timeout_s) + + def watch( + self, + *, + phase: str, + expected_flag: int, + active_mask: int | None = None, + ) -> None: + """Queue a just-launched AlltoAll phase for watchdog polling.""" + if phase not in self.VALID_PHASES: + raise ValueError(f"phase must be one of {sorted(self.VALID_PHASES)}, got {phase!r}") + if expected_flag < 0: + raise ValueError(f"expected_flag must be non-negative, got {expected_flag}") + if active_mask is None: + if self._health is not None: + active_mask = self._health.get_mask() + else: + active_mask = (1 << self._ep_size) - 1 + if not (active_mask >> self._ep_rank) & 1: + raise ValueError("active_mask must include the local ep_rank") + + self.start() + with self._cv: + if self._stopping: + raise RuntimeError("cannot queue a stopped AlltoAllWatchdog") + self._queue.append( + _CollectiveWatch( + phase=phase, + expected_flag=int(expected_flag), + active_mask=int(active_mask), + start_s=time.monotonic(), + ) + ) + self._cv.notify_all() + + def wait_until_idle(self, timeout_s: float) -> bool: + """Wait until all queued phases complete or timeout handling clears them.""" + deadline = time.monotonic() + timeout_s + with self._cv: + while self._queue: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + self._cv.wait(timeout=remaining) + return True + + def __enter__(self) -> "AlltoAllWatchdog": + self.start() + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self.stop(timeout_s=1.0) + + def _active_ranks(self, active_mask: int) -> tuple[int, ...]: + return tuple(rank for rank in range(self._ep_size) if (active_mask >> rank) & 1) + + def _phase_complete(self, watch: _CollectiveWatch, observed_flags: tuple[int, ...]) -> bool: + return all( + observed_flags[rank] == watch.expected_flag + for rank in self._active_ranks(watch.active_mask) + ) + + def _missing_ranks( + self, watch: _CollectiveWatch, observed_flags: tuple[int, ...] + ) -> tuple[int, ...]: + return tuple( + rank + for rank in self._active_ranks(watch.active_mask) + if observed_flags[rank] != watch.expected_flag + ) + + def _handle_timeout(self, watch: _CollectiveWatch, observed_flags: tuple[int, ...]) -> None: + elapsed_s = time.monotonic() - watch.start_s + missing_ranks = self._missing_ranks(watch, observed_flags) + marked_failed: list[int] = [] + if self._health is not None: + for rank in missing_ranks: + if rank == self._ep_rank: + continue + if self._health.mark_failed(rank): + marked_failed.append(rank) + + event = AlltoAllWatchdogTimeout( + phase=watch.phase, + expected_flag=watch.expected_flag, + observed_flags=observed_flags, + missing_ranks=missing_ranks, + marked_failed_ranks=tuple(marked_failed), + elapsed_s=elapsed_s, + ) + tllm_logger.warning( + "AlltoAll watchdog timeout on rank %d during %s: expected flag %d, " + "missing ranks %s, observed flags %s", + self._ep_rank, + watch.phase, + watch.expected_flag, + list(missing_ranks), + list(observed_flags), + ) + if self._on_timeout is not None: + self._on_timeout(event) + + def _run(self) -> None: + while True: + with self._cv: + while not self._queue and not self._stopping: + self._cv.wait() + if self._stopping: + return + watch = self._queue[0] + + try: + observed_flags = tuple( + int(v) for v in self._completion_reader.read_completion_flags(watch.phase) + ) + if len(observed_flags) != self._ep_size: + raise RuntimeError( + f"completion reader returned {len(observed_flags)} flags; " + f"expected ep_size={self._ep_size}" + ) + except BaseException as exc: # noqa: BLE001 - keep watchdog failures visible. + with self._cv: + self._last_error = exc + self._queue.clear() + self._cv.notify_all() + tllm_logger.error("AlltoAll watchdog stopped after polling error: %s", exc) + return + + if self._phase_complete(watch, observed_flags): + with self._cv: + if self._queue and self._queue[0] is watch: + self._queue.popleft() + self._cv.notify_all() + continue + + if time.monotonic() - watch.start_s >= self._timeout_s: + self._handle_timeout(watch, observed_flags) + with self._cv: + # The GPU stream is no longer trustworthy once a collective + # times out. Drop queued follow-on phases so they do not + # produce duplicate or misleading reports. + self._queue.clear() + self._cv.notify_all() + continue + + with self._cv: + self._cv.wait(timeout=self._poll_interval_s) diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index ca3a50dcfd12..8284b2d1087a 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -9,11 +9,13 @@ import os from dataclasses import dataclass -from typing import Dict, Optional +from typing import Callable, Dict, Optional import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._torch.alltoall_watchdog import (AlltoAllWatchdog, + AlltoAllWatchdogTimeout) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -126,6 +128,11 @@ def __init__( num_slots: int, workspace_size_per_rank: int, num_experts: Optional[int] = None, + ep_group_health=None, + alltoall_watchdog_timeout_s: Optional[float] = None, + alltoall_watchdog_poll_interval_s: float = 0.05, + alltoall_watchdog_on_timeout: Optional[Callable[ + [AlltoAllWatchdogTimeout], None]] = None, ): """ Initialize MoeAlltoAll with workspace allocation. @@ -138,6 +145,12 @@ def __init__( Note: The terminology is mapped to `num_experts` in this class and the kernels. num_experts: (Optional) Number of experts for EPLB stats (must be <= num_slots). DO NOT provide this parameter if EPLB is not enabled. Note: The terminology is mapped to `eplb_stats_num_experts` in this class and the kernels. + ep_group_health: Optional EPGroupHealth-compatible object. When present, its mask is passed to the + CUDA kernels and used by the watchdog. + alltoall_watchdog_timeout_s: Optional timeout for the host-side AlltoAll watchdog. If None, the + watchdog is disabled. + alltoall_watchdog_poll_interval_s: Poll interval for the watchdog thread. + alltoall_watchdog_on_timeout: Optional callback invoked when the watchdog reports suspects. """ # Check for environment variable override workspace_mb_env = os.environ.get("TRTLLM_MOE_A2A_WORKSPACE_MB") @@ -214,6 +227,65 @@ def __init__( self.metainfo = self._WORKSPACE["metainfo"] # Internal state self._state: _A2AState = _A2AState() + self.ep_group_health = ep_group_health + self._watchdog_flag_generation = 0 + self._alltoall_watchdog: AlltoAllWatchdog | None = None + if alltoall_watchdog_timeout_s is not None: + self._watchdog_flag_generation = self._read_current_flag_val() + self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( + workspace=self.workspace, + metainfo=self.metainfo, + metainfo_index=self._METAINFO_INDEX, + ep_rank=self.ep_rank, + ep_size=self.ep_size, + timeout_s=alltoall_watchdog_timeout_s, + poll_interval_s=alltoall_watchdog_poll_interval_s, + health=self.ep_group_health, + on_timeout=alltoall_watchdog_on_timeout, + ) + + def _read_current_flag_val(self) -> int: + flag_val_offset = self.metainfo[ + self._METAINFO_INDEX["FLAG_VAL_OFFSET_INDEX"]].item() + flag_val = self.workspace[self.ep_rank, + flag_val_offset:flag_val_offset + 4].view( + torch.int32) + if flag_val.device.type != "cpu": + flag_val = flag_val.detach().cpu() + return int(flag_val.item()) + + def _get_active_rank_mask_tensor( + self, + active_rank_mask: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + if active_rank_mask is not None: + return active_rank_mask + if self.ep_group_health is None: + return None + return torch.tensor(self.ep_group_health.get_mask_words(), + dtype=torch.uint64, + device="cpu") + + def _active_mask_int( + self, active_rank_mask: Optional[torch.Tensor]) -> Optional[int]: + if active_rank_mask is not None: + mask_cpu = active_rank_mask.detach().cpu() + return sum( + int(word) << (64 * idx) + for idx, word in enumerate(mask_cpu.tolist())) + if self.ep_group_health is not None: + return self.ep_group_health.get_mask() + return None + + def _watch_collective(self, phase: str, + active_rank_mask: Optional[torch.Tensor]) -> None: + if self._alltoall_watchdog is None: + return + self._watchdog_flag_generation += 1 + self._alltoall_watchdog.watch( + phase=phase, + expected_flag=self._watchdog_flag_generation, + active_mask=self._active_mask_int(active_rank_mask), + ) def dispatch(self, token_selected_experts: torch.Tensor, @@ -221,7 +293,8 @@ def dispatch(self, runtime_max_tokens_per_rank: int, invalid_token_expert_id: Optional[int] = None, expert_id_payload_index: Optional[int] = None, - eplb_local_stats: Optional[torch.Tensor] = None): + eplb_local_stats: Optional[torch.Tensor] = None, + active_rank_mask: Optional[torch.Tensor] = None): """ Perform MoE all-to-all dispatch operation. @@ -232,6 +305,7 @@ def dispatch(self, invalid_token_expert_id: If not None, set the token_selected_experts of the invalid tokens to this expert id. This is used to notify the MoE to skip these tokens for GroupGEMM. expert_id_payload_index: The index of token_selected_experts in the input_payloads. Must be provided if invalid_token_expert_id is not None. eplb_local_stats: (Optional) [num_experts] tensor containing local statistics for EPLB + active_rank_mask: Optional uint64 CPU tensor overriding ep_group_health for this dispatch. Returns: recv_tensors: List of tensors received, each has shape [ep_size, max_tokens_per_rank, payload_num_elements_per_token] @@ -246,6 +320,7 @@ def dispatch(self, 0 ) == self.eplb_stats_num_experts, "eplb_local_stats size must match eplb_stats_num_experts" + active_rank_mask = self._get_active_rank_mask_tensor(active_rank_mask) recv_tensors, combine_payload_offset, eplb_gathered_stats = torch.ops.trtllm.moe_a2a_dispatch( token_selected_experts, input_payloads, @@ -257,7 +332,9 @@ def dispatch(self, self.top_k, self.num_experts, eplb_local_stats, + active_rank_mask, ) + self._watch_collective("dispatch", active_rank_mask) if eplb_gathered_stats.numel() == 0: eplb_gathered_stats = None @@ -287,6 +364,7 @@ def combine( runtime_max_tokens_per_rank: int, payload_in_workspace: bool = False, use_low_precision_combine: bool = False, + active_rank_mask: Optional[torch.Tensor] = None, ): """ Perform MoE all-to-all combine operation. @@ -296,6 +374,7 @@ def combine( runtime_max_tokens_per_rank: Maximum of the number of tokens of each DP rank's local batch. payload_in_workspace: If True, 'payload' is a view into 'workspace' at 'combine_payload_offset' and no staging copy is needed. If False, the op stages 'payload' into the workspace region before combining. use_low_precision_combine: If True, quantize the combine payload to FP8 for NVLink transfer (halves NVLink bandwidth usage, output precision is preserved). + active_rank_mask: Optional uint64 CPU tensor overriding ep_group_health for this combine. Returns: combined_output: [local_num_tokens, num_elements_per_token] tensor of combined results @@ -303,11 +382,13 @@ def combine( assert self._state.phase == "dispatched", "combine called before a successful dispatch" assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" + active_rank_mask = self._get_active_rank_mask_tensor(active_rank_mask) output = torch.ops.trtllm.moe_a2a_combine( payload, self._state.local_num_tokens, self.workspace, self.metainfo, runtime_max_tokens_per_rank, self.ep_rank, self.ep_size, self.top_k, self._state.combine_payload_offset, - payload_in_workspace, use_low_precision_combine) + payload_in_workspace, use_low_precision_combine, active_rank_mask) + self._watch_collective("combine", active_rank_mask) # Reset state for next round self.reset_state() diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index f9068f71d717..7b5ed16e8256 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py @@ -25,11 +25,12 @@ """ import os -from typing import Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Tuple import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._torch.alltoall_watchdog import AlltoAllWatchdog, AlltoAllWatchdogTimeout from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -151,6 +152,10 @@ def __init__( dtype: Optional[torch.dtype] = None, num_experts: Optional[int] = None, use_low_precision_combine: bool = False, + ep_group_health=None, + alltoall_watchdog_timeout_s: Optional[float] = None, + alltoall_watchdog_poll_interval_s: float = 0.05, + alltoall_watchdog_on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, ): """ Initialize NVLinkOneSided with workspace allocation. @@ -169,6 +174,12 @@ def __init__( use_low_precision_combine: If True, quantize the combine payload to FP8 for NVLink transfer (halves NVLink bandwidth usage, output precision is preserved). Corresponds to model_config.use_low_precision_moe_combine. + ep_group_health: Optional EPGroupHealth-compatible object. When present, its mask is passed to the + CUDA kernels and used by the watchdog. + alltoall_watchdog_timeout_s: Optional timeout for the host-side AlltoAll watchdog. If None, the + watchdog is disabled. + alltoall_watchdog_poll_interval_s: Poll interval for the watchdog thread. + alltoall_watchdog_on_timeout: Optional callback invoked when the watchdog reports suspects. """ super().__init__(mapping) @@ -300,6 +311,26 @@ def __init__( self.workspace = workspace_state["workspace"] self.moe_a2a_metainfo = workspace_state["metainfo"] self.max_num_tokens_per_rank = workspace_state["max_num_tokens_per_rank"] + self.ep_group_health = ep_group_health + self._watchdog_flag_generation = 0 + self._alltoall_watchdog: AlltoAllWatchdog | None = None + if alltoall_watchdog_timeout_s is not None: + self._watchdog_flag_generation = self._read_current_flag_val() + self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( + workspace=self.workspace, + metainfo=self.moe_a2a_metainfo, + metainfo_index={ + "FLAG_VAL_OFFSET_INDEX": self.FLAG_VAL_OFFSET_INDEX, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": self.DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": self.COMBINE_COMPLETION_FLAGS_OFFSET_INDEX, + }, + ep_rank=self.ep_rank, + ep_size=self.ep_size, + timeout_s=alltoall_watchdog_timeout_s, + poll_interval_s=alltoall_watchdog_poll_interval_s, + health=self.ep_group_health, + on_timeout=alltoall_watchdog_on_timeout, + ) # Initialize dispatch state self._dispatch_state = {"phase": "idle"} @@ -307,6 +338,42 @@ def __init__( # Invalid token expert ID (default to -1), the kernels in TRTLLM-gen is hard-code to support -1 only. self.invalid_token_expert_id: int = -1 + def _read_current_flag_val(self) -> int: + flag_val_offset = self.moe_a2a_metainfo[self.FLAG_VAL_OFFSET_INDEX].item() + flag_val = self.workspace[self.ep_rank, flag_val_offset : flag_val_offset + 4].view( + torch.int32 + ) + if flag_val.device.type != "cpu": + flag_val = flag_val.detach().cpu() + return int(flag_val.item()) + + def _get_active_rank_mask_tensor( + self, active_rank_mask: Optional[torch.Tensor] + ) -> Optional[torch.Tensor]: + if active_rank_mask is not None: + return active_rank_mask + if self.ep_group_health is None: + return None + return torch.tensor(self.ep_group_health.get_mask_words(), dtype=torch.uint64, device="cpu") + + def _active_mask_int(self, active_rank_mask: Optional[torch.Tensor]) -> Optional[int]: + if active_rank_mask is not None: + mask_cpu = active_rank_mask.detach().cpu() + return sum(int(word) << (64 * idx) for idx, word in enumerate(mask_cpu.tolist())) + if self.ep_group_health is not None: + return self.ep_group_health.get_mask() + return None + + def _watch_collective(self, phase: str, active_rank_mask: Optional[torch.Tensor]) -> None: + if self._alltoall_watchdog is None: + return + self._watchdog_flag_generation += 1 + self._alltoall_watchdog.watch( + phase=phase, + expected_flag=self._watchdog_flag_generation, + active_mask=self._active_mask_int(active_rank_mask), + ) + @staticmethod def is_platform_supported() -> bool: """ @@ -326,6 +393,9 @@ def destroy(self): return self._destroyed = True + if self._alltoall_watchdog is not None: + self._alltoall_watchdog.stop(timeout_s=1.0) + self._alltoall_watchdog = None workspace_key = getattr(self, "_workspace_key", None) if workspace_key is None: return @@ -409,6 +479,7 @@ def dispatch( assert eplb_local_stats.size(0) == self.eplb_stats_num_experts, ( "eplb_local_stats size must match eplb_stats_num_experts" ) + active_rank_mask = self._get_active_rank_mask_tensor(kwargs.get("active_rank_mask")) recv_buffers, combine_payload_offset, eplb_gathered_stats = ( torch.ops.trtllm.moe_a2a_dispatch( @@ -422,8 +493,10 @@ def dispatch( self.top_k, self.num_experts, eplb_local_stats, + active_rank_mask, ) ) + self._watch_collective("dispatch", active_rank_mask) if eplb_gathered_stats.numel() == 0: eplb_gathered_stats = None self._dispatch_state["eplb_gathered_stats"] = eplb_gathered_stats @@ -526,6 +599,7 @@ def combine( raise ValueError( f"final_hidden_states must be 2D or 3D, got {final_hidden_states.dim()}D" ) + active_rank_mask = self._get_active_rank_mask_tensor(kwargs.get("active_rank_mask")) output = torch.ops.trtllm.moe_a2a_combine( final_hidden_states, int(local_num_tokens), @@ -538,7 +612,9 @@ def combine( int(combine_payload_offset), bool(self.payload_in_workspace), bool(self.use_low_precision_combine), + active_rank_mask, ) + self._watch_collective("combine", active_rank_mask) # Reset state for next round self.reset_state() diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py new file mode 100644 index 000000000000..1168bb0f00fe --- /dev/null +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for AlltoAllWatchdog (WideEP fault tolerance, PR 1a.4).""" + +import threading +import time + +import pytest +import torch + +from tensorrt_llm._torch.alltoall_watchdog import AlltoAllWatchdog, AlltoAllWatchdogTimeout +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth + + +class FakeCompletionFlagReader: + """Thread-safe completion flag reader for pure-Python watchdog tests.""" + + def __init__(self, ep_size: int) -> None: + self._lock = threading.Lock() + self._flags = { + "dispatch": [0 for _ in range(ep_size)], + "combine": [0 for _ in range(ep_size)], + } + + def set_flags(self, phase: str, flags: list[int]) -> None: + with self._lock: + self._flags[phase] = list(flags) + + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + with self._lock: + return tuple(self._flags[phase]) + + +def _wait_for(predicate, timeout_s: float = 1.0) -> None: + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.005) + raise AssertionError("condition was not reached before timeout") + + +def test_watchdog_completes_when_all_active_flags_arrive() -> None: + health = EPGroupHealth(4) + reader = FakeCompletionFlagReader(ep_size=4) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.2, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + reader.set_flags("dispatch", [1, 1, 1, 1]) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + assert health.all_active() is True + + +def test_watchdog_timeout_reports_and_marks_missing_remote_ranks() -> None: + health = EPGroupHealth(4) + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("dispatch", [1, 0, 1, 0]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: len(events) == 1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + event = events[0] + assert event.phase == "dispatch" + assert event.expected_flag == 1 + assert event.observed_flags == (1, 0, 1, 0) + assert event.missing_ranks == (1, 3) + assert event.marked_failed_ranks == (1, 3) + assert health.get_failed_ranks() == frozenset({1, 3}) + + +def test_watchdog_ignores_ranks_already_failed_in_health_mask() -> None: + health = EPGroupHealth(4) + assert health.mark_failed(2) is True + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("dispatch", [1, 1, 0, 1]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.05, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + assert health.get_failed_ranks() == frozenset({2}) + + +def test_watchdog_reports_local_missing_but_does_not_mark_local_failed() -> None: + health = EPGroupHealth(4) + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("combine", [0, 2, 2, 2]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="combine", expected_flag=2) + _wait_for(lambda: len(events) == 1) + + event = events[0] + assert event.missing_ranks == (0,) + assert event.marked_failed_ranks == () + assert health.get_failed_ranks() == frozenset() + + +def test_watchdog_preserves_fifo_order_and_clears_followups_after_timeout() -> None: + health = EPGroupHealth(3) + reader = FakeCompletionFlagReader(ep_size=3) + reader.set_flags("dispatch", [1, 0, 1]) + reader.set_flags("combine", [0, 0, 0]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=3, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + watchdog.watch(phase="combine", expected_flag=2) + _wait_for(lambda: len(events) == 1) + assert watchdog.wait_until_idle(timeout_s=1.0) + time.sleep(0.05) + + assert len(events) == 1 + assert events[0].phase == "dispatch" + assert events[0].missing_ranks == (1,) + assert health.get_failed_ranks() == frozenset({1}) + + +def test_watchdog_from_workspace_reads_phase_specific_offsets() -> None: + ep_size = 3 + ep_rank = 1 + workspace = torch.zeros((ep_size, 64), dtype=torch.uint8) + metainfo = torch.zeros((10,), dtype=torch.int64) + metainfo_index = { + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 4, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 5, + } + metainfo[4] = 4 + metainfo[5] = 20 + workspace[ep_rank, 4:16].view(torch.int32).copy_(torch.tensor([7, 7, 7], dtype=torch.int32)) + workspace[ep_rank, 20:32].view(torch.int32).copy_(torch.tensor([0, 8, 8], dtype=torch.int32)) + health = EPGroupHealth(ep_size) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog.from_workspace( + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=ep_rank, + ep_size=ep_size, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=7) + assert watchdog.wait_until_idle(timeout_s=1.0) + + watchdog.watch(phase="combine", expected_flag=8) + _wait_for(lambda: len(events) == 1) + + assert events[0].phase == "combine" + assert events[0].missing_ranks == (0,) + assert events[0].marked_failed_ranks == (0,) + assert health.get_failed_ranks() == frozenset({0}) + + +def test_watchdog_rejects_active_mask_without_local_rank() -> None: + reader = FakeCompletionFlagReader(ep_size=4) + with AlltoAllWatchdog( + ep_size=4, + ep_rank=2, + completion_reader=reader, + timeout_s=0.1, + poll_interval_s=0.005, + ) as watchdog: + with pytest.raises(ValueError, match="local ep_rank"): + watchdog.watch(phase="dispatch", expected_flag=1, active_mask=0b1011) From ede407f3f6e319145229400f320f84090659c9f1 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 22 Jun 2026 15:04:38 -0700 Subject: [PATCH 02/14] [None][fix] Address AlltoAll watchdog review findings Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 116 +++++++++++++++--- .../_torch/distributed/moe_alltoall.py | 28 ++++- .../communication/communication_factory.py | 13 ++ .../communication/nvlink_one_sided.py | 11 +- .../modules/fused_moe/fused_moe_cutlass.py | 7 ++ .../modules/fused_moe/fused_moe_trtllm_gen.py | 9 +- .../_torch/modules/fused_moe/wide_ep_ft.py | 82 +++++++++++++ .../_torch/modules/test_alltoall_watchdog.py | 102 ++++++++++++++- 8 files changed, 345 insertions(+), 23 deletions(-) create mode 100644 tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 7c35669e29ee..79a6347b58e2 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -31,8 +31,13 @@ import torch +from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.logger import logger as tllm_logger +DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S = 5.0 +DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S = 0.1 +UNKNOWN_COMPLETION_FLAG = -(2**63) + class CompletionFlagReader(Protocol): """Reads one phase's rank-local completion flag row.""" @@ -51,6 +56,10 @@ def mark_failed(self, rank: int) -> bool: """Mark ``rank`` failed and return whether state changed.""" +class CompletionFlagReadTimeout(TimeoutError): + """Raised when the host watchdog cannot read completion flags in time.""" + + @dataclass(frozen=True) class AlltoAllWatchdogTimeout: """Details emitted when an AlltoAll phase times out.""" @@ -61,6 +70,7 @@ class AlltoAllWatchdogTimeout: missing_ranks: tuple[int, ...] marked_failed_ranks: tuple[int, ...] elapsed_s: float + poll_timed_out: bool = False @dataclass(frozen=True) @@ -81,6 +91,7 @@ def __init__( ep_size: int, dispatch_completion_flags_offset: int, combine_completion_flags_offset: int, + device_copy_timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, ) -> None: if workspace.dim() != 2: raise ValueError("workspace must be a 2D tensor [ep_size, size_per_rank]") @@ -97,11 +108,50 @@ def __init__( "dispatch": int(dispatch_completion_flags_offset), "combine": int(combine_completion_flags_offset), } + self._device_copy_timeout_s = float(device_copy_timeout_s) + self._copy_stream: torch.cuda.Stream | None = None + self._retired_copies: list[tuple[torch.Tensor, torch.cuda.Event]] = [] + if workspace.device.type == "cuda": + self._copy_stream = torch.cuda.Stream(device=workspace.device) + + def _prune_retired_copies(self) -> None: + self._retired_copies = [ + (host_flags, event) for host_flags, event in self._retired_copies if not event.query() + ] + + def _read_cuda_flags(self, flags: torch.Tensor) -> tuple[int, ...]: + assert self._copy_stream is not None + self._prune_retired_copies() + + host_flags = torch.empty( + (self._ep_size,), + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + event = torch.cuda.Event(blocking=False) + with torch.cuda.device(flags.device), torch.cuda.stream(self._copy_stream): + host_flags.copy_(flags.detach(), non_blocking=True) + event.record(self._copy_stream) + + deadline_s = time.monotonic() + self._device_copy_timeout_s + while not event.query(): + remaining_s = deadline_s - time.monotonic() + if remaining_s <= 0: + self._retired_copies.append((host_flags, event)) + raise CompletionFlagReadTimeout( + "timed out copying AlltoAll completion flags to host" + ) + time.sleep(min(remaining_s, 0.001)) + + return tuple(int(v) for v in host_flags.tolist()) def read_completion_flags(self, phase: str) -> tuple[int, ...]: offset = self._offsets[phase] end = offset + self._ep_size * 4 flags = self._workspace[self._ep_rank, offset:end].view(torch.int32) + if flags.device.type == "cuda": + return self._read_cuda_flags(flags) if flags.device.type != "cpu": flags = flags.detach().cpu() return tuple(int(v) for v in flags.tolist()) @@ -123,8 +173,8 @@ def __init__( ep_size: int, ep_rank: int, completion_reader: CompletionFlagReader, - timeout_s: float, - poll_interval_s: float = 0.05, + timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, health: Optional[EPGroupHealthLike] = None, on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, ) -> None: @@ -160,8 +210,8 @@ def from_workspace( metainfo_index: Mapping[str, int], ep_rank: int, ep_size: int, - timeout_s: float, - poll_interval_s: float = 0.05, + timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, health: Optional[EPGroupHealthLike] = None, on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, ) -> "AlltoAllWatchdog": @@ -178,6 +228,7 @@ def from_workspace( ep_size=ep_size, dispatch_completion_flags_offset=dispatch_offset, combine_completion_flags_offset=combine_offset, + device_copy_timeout_s=poll_interval_s, ) return cls( ep_size=ep_size, @@ -288,11 +339,18 @@ def _missing_ranks( if observed_flags[rank] != watch.expected_flag ) - def _handle_timeout(self, watch: _CollectiveWatch, observed_flags: tuple[int, ...]) -> None: + def _handle_timeout( + self, + watch: _CollectiveWatch, + observed_flags: tuple[int, ...], + *, + poll_timed_out: bool = False, + ) -> None: elapsed_s = time.monotonic() - watch.start_s missing_ranks = self._missing_ranks(watch, observed_flags) marked_failed: list[int] = [] - if self._health is not None: + has_known_flags = UNKNOWN_COMPLETION_FLAG not in observed_flags + if self._health is not None and (has_known_flags or not poll_timed_out): for rank in missing_ranks: if rank == self._ep_rank: continue @@ -306,20 +364,37 @@ def _handle_timeout(self, watch: _CollectiveWatch, observed_flags: tuple[int, .. missing_ranks=missing_ranks, marked_failed_ranks=tuple(marked_failed), elapsed_s=elapsed_s, + poll_timed_out=poll_timed_out, ) - tllm_logger.warning( - "AlltoAll watchdog timeout on rank %d during %s: expected flag %d, " - "missing ranks %s, observed flags %s", - self._ep_rank, - watch.phase, - watch.expected_flag, - list(missing_ranks), - list(observed_flags), - ) + if poll_timed_out: + tllm_logger.error( + "AlltoAll watchdog could not read completion flags on rank %d " + "during %s before timeout %.3fs; expected flag %d, active " + "ranks %s, observed flags %s, marked ranks %s", + self._ep_rank, + watch.phase, + elapsed_s, + watch.expected_flag, + list(self._active_ranks(watch.active_mask)), + list(observed_flags), + list(marked_failed), + ) + else: + tllm_logger.warning( + "AlltoAll watchdog timeout on rank %d during %s: expected flag %d, " + "missing ranks %s, observed flags %s", + self._ep_rank, + watch.phase, + watch.expected_flag, + list(missing_ranks), + list(observed_flags), + ) if self._on_timeout is not None: self._on_timeout(event) def _run(self) -> None: + last_observed_flags = tuple(UNKNOWN_COMPLETION_FLAG for _ in range(self._ep_size)) + poll_timed_out = False while True: with self._cv: while not self._queue and not self._stopping: @@ -337,6 +412,11 @@ def _run(self) -> None: f"completion reader returned {len(observed_flags)} flags; " f"expected ep_size={self._ep_size}" ) + last_observed_flags = observed_flags + poll_timed_out = False + except CompletionFlagReadTimeout: + observed_flags = last_observed_flags + poll_timed_out = True except BaseException as exc: # noqa: BLE001 - keep watchdog failures visible. with self._cv: self._last_error = exc @@ -350,16 +430,20 @@ def _run(self) -> None: if self._queue and self._queue[0] is watch: self._queue.popleft() self._cv.notify_all() + last_observed_flags = tuple(UNKNOWN_COMPLETION_FLAG for _ in range(self._ep_size)) + poll_timed_out = False continue if time.monotonic() - watch.start_s >= self._timeout_s: - self._handle_timeout(watch, observed_flags) + self._handle_timeout(watch, observed_flags, poll_timed_out=poll_timed_out) with self._cv: # The GPU stream is no longer trustworthy once a collective # times out. Drop queued follow-on phases so they do not # produce duplicate or misleading reports. self._queue.clear() self._cv.notify_all() + last_observed_flags = tuple(UNKNOWN_COMPLETION_FLAG for _ in range(self._ep_size)) + poll_timed_out = False continue with self._cv: diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index 8284b2d1087a..fe2c1020d201 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -8,14 +8,17 @@ # ruff: noqa: E501 import os +import sys from dataclasses import dataclass from typing import Callable, Dict, Optional import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory -from tensorrt_llm._torch.alltoall_watchdog import (AlltoAllWatchdog, - AlltoAllWatchdogTimeout) +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, AlltoAllWatchdog, + AlltoAllWatchdogTimeout) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -130,7 +133,8 @@ def __init__( num_experts: Optional[int] = None, ep_group_health=None, alltoall_watchdog_timeout_s: Optional[float] = None, - alltoall_watchdog_poll_interval_s: float = 0.05, + alltoall_watchdog_poll_interval_s: + float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, alltoall_watchdog_on_timeout: Optional[Callable[ [AlltoAllWatchdogTimeout], None]] = None, ): @@ -228,8 +232,12 @@ def __init__( # Internal state self._state: _A2AState = _A2AState() self.ep_group_health = ep_group_health + self._destroyed = False self._watchdog_flag_generation = 0 self._alltoall_watchdog: AlltoAllWatchdog | None = None + if (alltoall_watchdog_timeout_s is None + and self.ep_group_health is not None): + alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S if alltoall_watchdog_timeout_s is not None: self._watchdog_flag_generation = self._read_current_flag_val() self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( @@ -244,6 +252,20 @@ def __init__( on_timeout=alltoall_watchdog_on_timeout, ) + def destroy(self) -> None: + """Stop background watchdog resources owned by this wrapper.""" + if getattr(self, "_destroyed", False): + return + self._destroyed = True + watchdog = getattr(self, "_alltoall_watchdog", None) + if watchdog is not None: + watchdog.stop(timeout_s=1.0) + self._alltoall_watchdog = None + + def __del__(self) -> None: + if not sys.is_finalizing(): + self.destroy() + def _read_current_flag_val(self) -> int: flag_val_offset = self.metainfo[ self._METAINFO_INDEX["FLAG_VAL_OFFSET_INDEX"]].item() diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py index a4ec2ceefe44..17a67deee219 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py @@ -28,6 +28,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.logger import logger +from ..wide_ep_ft import get_wide_ep_ft_options from .allgather_reducescatter import AllGatherReduceScatter from .base import Communication from .deep_ep import DeepEP @@ -133,6 +134,9 @@ def create_strategy( try: enable_eplb = model_config.moe_load_balancer is not None + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = get_wide_ep_ft_options( + model_config + ) strategy = NVLinkOneSided( mapping, num_slots, @@ -143,6 +147,9 @@ def create_strategy( dtype=act_dtype, num_experts=num_experts if enable_eplb else None, use_low_precision_combine=use_low_precision_combine, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s=watchdog_poll_interval_s, ) logger.info("Selected communication strategy: NVLinkOneSided") return strategy @@ -285,6 +292,9 @@ def _create_forced_method( ) elif method in ["NVLINK_ONE_SIDED"]: enable_eplb = model_config.moe_load_balancer is not None + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = get_wide_ep_ft_options( + model_config + ) return NVLinkOneSided( mapping, num_slots, @@ -295,6 +305,9 @@ def _create_forced_method( dtype=act_dtype, num_experts=num_experts if enable_eplb else None, use_low_precision_combine=use_low_precision_combine, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s=watchdog_poll_interval_s, ) elif method == "DEEPEP": return DeepEP( diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index 7b5ed16e8256..9b3d72306775 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py @@ -30,7 +30,12 @@ import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory -from tensorrt_llm._torch.alltoall_watchdog import AlltoAllWatchdog, AlltoAllWatchdogTimeout +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + AlltoAllWatchdog, + AlltoAllWatchdogTimeout, +) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -154,7 +159,7 @@ def __init__( use_low_precision_combine: bool = False, ep_group_health=None, alltoall_watchdog_timeout_s: Optional[float] = None, - alltoall_watchdog_poll_interval_s: float = 0.05, + alltoall_watchdog_poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, alltoall_watchdog_on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, ): """ @@ -314,6 +319,8 @@ def __init__( self.ep_group_health = ep_group_health self._watchdog_flag_generation = 0 self._alltoall_watchdog: AlltoAllWatchdog | None = None + if alltoall_watchdog_timeout_s is None and self.ep_group_health is not None: + alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S if alltoall_watchdog_timeout_s is not None: self._watchdog_flag_generation = self._read_current_flag_val() self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index 3a89ec2fb35d..0b619bb31916 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -20,6 +20,7 @@ Fp4QuantizedTensor) from .interface import AlltoallMethodType, MoE from .quantization import UnquantizedFusedMoEMethod +from .wide_ep_ft import get_wide_ep_ft_options # isort: off from .quantization import ( @@ -331,6 +332,8 @@ def __init__( dtype, self.num_experts if self.layer_load_balancer else None, ) + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = ( + get_wide_ep_ft_options(model_config)) self.moe_a2a = MoeAlltoAll( mapping=self.mapping, @@ -340,6 +343,10 @@ def __init__( workspace_size_per_rank=workspace_size, num_experts=self.num_experts if self.layer_load_balancer else None, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s= + watchdog_poll_interval_s, ) elif self.alltoall_method_type == AlltoallMethodType.DeepEP or self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: raise NotImplementedError( diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 820d92a95e4a..6f59870c374b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -36,6 +36,7 @@ from ...utils import ActivationType, AuxStreamType, Fp4QuantizedTensor from .interface import AlltoallMethodType, MoE, MoEWeightLoadingMode from .moe_op_backend import MoEOpBackend, get_op_backend +from .wide_ep_ft import get_wide_ep_ft_options # isort: off from .quantization import ( @@ -291,6 +292,8 @@ def __init__( ep_size, self.routing_method.experts_per_token, max_num_tokens, hidden_size, dtype, self.num_experts if self.layer_load_balancer else None) + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = ( + get_wide_ep_ft_options(model_config)) self.moe_a2a = MoeAlltoAll( mapping=self.mapping, @@ -299,7 +302,11 @@ def __init__( num_slots=self.num_slots, workspace_size_per_rank=workspace_size, num_experts=self.num_experts - if self.layer_load_balancer else None) + if self.layer_load_balancer else None, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s= + watchdog_poll_interval_s) elif self.alltoall_method_type == AlltoallMethodType.DeepEP or self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: raise NotImplementedError( "DeepEP and DeepEPLowLatency are not supported for TRTLLMGenFusedMoE yet" diff --git a/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py new file mode 100644 index 000000000000..69e5c53f061e --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shared WideEP fault-tolerance options for MoE communication paths.""" + +from __future__ import annotations + +import os +from typing import Any, Optional + +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, +) + +from .ep_group_health import EPGroupHealth + +_ENABLE_ENV = "TRTLLM_ENABLE_WIDE_EP_FT" +_TIMEOUT_ENV = "TRTLLM_ALLTOALL_WATCHDOG_TIMEOUT_S" +_POLL_INTERVAL_ENV = "TRTLLM_ALLTOALL_WATCHDOG_POLL_INTERVAL_S" + +_HEALTH_KEY = "wide_ep_ft_ep_group_health" +_TIMEOUT_KEY = "alltoall_watchdog_timeout_s" +_POLL_INTERVAL_KEY = "alltoall_watchdog_poll_interval_s" + + +def _env_enabled() -> bool: + return os.environ.get(_ENABLE_ENV, "0").lower() in {"1", "true", "yes", "on"} + + +def _float_option(extra_attrs: dict, key: str, env_name: str, default: float) -> float: + if key in extra_attrs: + return float(extra_attrs[key]) + if env_name in os.environ: + return float(os.environ[env_name]) + return default + + +def get_wide_ep_ft_options( + model_config: Any, +) -> tuple[Optional[EPGroupHealth], Optional[float], float]: + """Return the shared EP health object and watchdog timing for a model. + + WideEP FT remains opt-in until the integration PR wires a public model + option. Callers can either inject ``wide_ep_ft_ep_group_health`` through + ``ModelConfig.extra_attrs`` or set ``TRTLLM_ENABLE_WIDE_EP_FT=1`` to create + one process-local health object shared by all MoE communication layers. + """ + + extra_attrs = getattr(model_config, "extra_attrs", {}) + health = extra_attrs.get(_HEALTH_KEY) or extra_attrs.get("ep_group_health") + if health is None and _env_enabled(): + health = EPGroupHealth(model_config.mapping.moe_ep_size) + extra_attrs[_HEALTH_KEY] = health + + poll_interval_s = _float_option( + extra_attrs, + _POLL_INTERVAL_KEY, + _POLL_INTERVAL_ENV, + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + ) + if health is None: + return None, None, poll_interval_s + + timeout_s = _float_option( + extra_attrs, + _TIMEOUT_KEY, + _TIMEOUT_ENV, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + ) + return health, timeout_s, poll_interval_s diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index 1168bb0f00fe..ba4b6774f412 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -16,12 +16,21 @@ import threading import time +from types import SimpleNamespace import pytest import torch -from tensorrt_llm._torch.alltoall_watchdog import AlltoAllWatchdog, AlltoAllWatchdogTimeout +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + UNKNOWN_COMPLETION_FLAG, + AlltoAllWatchdog, + AlltoAllWatchdogTimeout, + CompletionFlagReadTimeout, +) from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth +from tensorrt_llm._torch.modules.fused_moe.wide_ep_ft import get_wide_ep_ft_options class FakeCompletionFlagReader: @@ -43,6 +52,23 @@ def read_completion_flags(self, phase: str) -> tuple[int, ...]: return tuple(self._flags[phase]) +class TimeoutCompletionFlagReader: + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + raise CompletionFlagReadTimeout("blocked") + + +class OneGoodReadThenTimeoutReader: + def __init__(self, flags: tuple[int, ...]) -> None: + self._flags = flags + self._read_count = 0 + + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + self._read_count += 1 + if self._read_count == 1: + return self._flags + raise CompletionFlagReadTimeout("blocked") + + def _wait_for(predicate, timeout_s: float = 1.0) -> None: deadline = time.monotonic() + timeout_s while time.monotonic() < deadline: @@ -74,6 +100,32 @@ def test_watchdog_completes_when_all_active_flags_arrive() -> None: assert health.all_active() is True +def test_watchdog_defaults_match_design_doc() -> None: + reader = FakeCompletionFlagReader(ep_size=1) + watchdog = AlltoAllWatchdog(ep_size=1, ep_rank=0, completion_reader=reader) + + assert watchdog._timeout_s == DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + assert watchdog._poll_interval_s == DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S + + +def test_wide_ep_ft_options_create_shared_health_when_enabled(monkeypatch) -> None: + monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") + model_config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(moe_ep_size=4), + ) + + health, timeout_s, poll_interval_s = get_wide_ep_ft_options(model_config) + health_again, timeout_again_s, poll_again_s = get_wide_ep_ft_options(model_config) + + assert isinstance(health, EPGroupHealth) + assert health_again is health + assert timeout_s == DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + assert timeout_again_s == timeout_s + assert poll_interval_s == DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S + assert poll_again_s == poll_interval_s + + def test_watchdog_timeout_reports_and_marks_missing_remote_ranks() -> None: health = EPGroupHealth(4) reader = FakeCompletionFlagReader(ep_size=4) @@ -149,6 +201,54 @@ def test_watchdog_reports_local_missing_but_does_not_mark_local_failed() -> None assert health.get_failed_ranks() == frozenset() +def test_watchdog_poll_timeout_without_snapshot_fails_closed() -> None: + health = EPGroupHealth(3) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=3, + ep_rank=0, + completion_reader=TimeoutCompletionFlagReader(), + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: len(events) == 1) + + event = events[0] + assert event.poll_timed_out is True + assert event.observed_flags == (UNKNOWN_COMPLETION_FLAG,) * 3 + assert event.missing_ranks == (0, 1, 2) + assert event.marked_failed_ranks == () + assert health.all_active() is True + + +def test_watchdog_poll_timeout_with_prior_snapshot_marks_known_missing_rank() -> None: + health = EPGroupHealth(3) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=3, + ep_rank=0, + completion_reader=OneGoodReadThenTimeoutReader((1, 0, 1)), + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: len(events) == 1) + + event = events[0] + assert event.poll_timed_out is True + assert event.observed_flags == (1, 0, 1) + assert event.missing_ranks == (1,) + assert event.marked_failed_ranks == (1,) + assert health.get_failed_ranks() == frozenset({1}) + + def test_watchdog_preserves_fifo_order_and_clears_followups_after_timeout() -> None: health = EPGroupHealth(3) reader = FakeCompletionFlagReader(ep_size=3) From 003fe0cec868279a5f46d99ac2ca1462dfceb35a Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 22 Jun 2026 15:20:17 -0700 Subject: [PATCH 03/14] fix: make AlltoAll watchdog stop terminal Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 6 +++++- .../_torch/modules/test_alltoall_watchdog.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 79a6347b58e2..2688cfa0e2d7 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -197,6 +197,7 @@ def __init__( self._cv = threading.Condition() self._queue: Deque[_CollectiveWatch] = deque() + self._closed = False self._stopping = False self._thread: threading.Thread | None = None self._last_error: BaseException | None = None @@ -249,6 +250,8 @@ def last_error(self) -> BaseException | None: def start(self) -> None: """Start the background polling thread. Idempotent.""" with self._cv: + if self._closed: + raise RuntimeError("cannot start a stopped AlltoAllWatchdog") if self._thread is not None and self._thread.is_alive(): return self._stopping = False @@ -262,6 +265,7 @@ def start(self) -> None: def stop(self, timeout_s: float | None = None) -> None: """Stop the polling thread and wait for it to exit.""" with self._cv: + self._closed = True self._stopping = True self._queue.clear() self._cv.notify_all() @@ -291,7 +295,7 @@ def watch( self.start() with self._cv: - if self._stopping: + if self._closed: raise RuntimeError("cannot queue a stopped AlltoAllWatchdog") self._queue.append( _CollectiveWatch( diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index ba4b6774f412..cda4d813d146 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -108,6 +108,24 @@ def test_watchdog_defaults_match_design_doc() -> None: assert watchdog._poll_interval_s == DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S +def test_watchdog_stop_is_terminal() -> None: + reader = FakeCompletionFlagReader(ep_size=1) + watchdog = AlltoAllWatchdog( + ep_size=1, + ep_rank=0, + completion_reader=reader, + timeout_s=0.2, + poll_interval_s=0.005, + ) + watchdog.start() + watchdog.stop(timeout_s=1.0) + + with pytest.raises(RuntimeError, match="stopped AlltoAllWatchdog"): + watchdog.start() + with pytest.raises(RuntimeError, match="stopped AlltoAllWatchdog"): + watchdog.watch(phase="dispatch", expected_flag=1) + + def test_wide_ep_ft_options_create_shared_health_when_enabled(monkeypatch) -> None: monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") model_config = SimpleNamespace( From e3af911c239e9e6557f1f2d8dcd2145a790dc22b Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 23 Jun 2026 10:09:02 -0700 Subject: [PATCH 04/14] fix: address AlltoAll watchdog review comments Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../moeAlltoAllKernels.cu | 8 +- .../communicationKernels/moeAlltoAllKernels.h | 4 +- tensorrt_llm/_torch/alltoall_watchdog.py | 37 ++- .../_torch/distributed/moe_alltoall.py | 32 ++- .../communication/nvlink_one_sided.py | 30 +- .../_torch/modules/moe/test_moe_comm.py | 267 +++++++++++++++++- .../_torch/modules/test_alltoall_watchdog.py | 7 +- 7 files changed, 344 insertions(+), 41 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index 0c4431b8eaef..0f2d453c363a 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -741,7 +741,7 @@ __device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, i { int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k]; int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k]; - if (dst_idx < 0) + if (dst_idx < 0 || !is_rank_active(ptrs.active_rank_mask, target_rank)) { acc[k].fill(0.0f); continue; @@ -766,8 +766,12 @@ __device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, i #pragma unroll for (int k = 0; k < TOP_K; ++k) { - if (ptrs.topk_send_indices[local_token_idx * TOP_K + k] < 0) + int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k]; + int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k]; + if (dst_idx < 0 || !is_rank_active(ptrs.active_rank_mask, target_rank)) + { continue; // acc[k] already holds 0.0f from fill() above + } #pragma unroll for (int j = elems_per_vec - 1; j >= 0; --j) acc[k][j] = static_cast(reinterpret_cast(&acc[k])[j]); diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h index 9a6f3904c501..138ca92e71a8 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h @@ -93,8 +93,8 @@ struct CombineKernelPointers int const* topk_send_indices; // dst index per k, -1 for duplicates // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Combine skips flag - // writes/waits to/from masked peers; per-token accumulation uses topk_send_indices[k] < 0 - // (set by dispatch) to skip dead-targeted slots, so no explicit mask check is needed there. + // writes/waits to/from masked peers and also skips per-token accumulation for ranks that + // become inactive between dispatch and combine. uint64_t active_rank_mask[kRankMaskWords]; }; diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 2688cfa0e2d7..c41a3fb8fdb3 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -26,8 +26,9 @@ import threading import time from collections import deque +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass -from typing import Callable, Deque, Mapping, Optional, Protocol, Sequence +from typing import Protocol import torch @@ -110,6 +111,8 @@ def __init__( } self._device_copy_timeout_s = float(device_copy_timeout_s) self._copy_stream: torch.cuda.Stream | None = None + self._host_flags: torch.Tensor | None = None + self._copy_event: torch.cuda.Event | None = None self._retired_copies: list[tuple[torch.Tensor, torch.cuda.Event]] = [] if workspace.device.type == "cuda": self._copy_stream = torch.cuda.Stream(device=workspace.device) @@ -123,13 +126,17 @@ def _read_cuda_flags(self, flags: torch.Tensor) -> tuple[int, ...]: assert self._copy_stream is not None self._prune_retired_copies() - host_flags = torch.empty( - (self._ep_size,), - dtype=torch.int32, - device="cpu", - pin_memory=prefer_pinned(), - ) - event = torch.cuda.Event(blocking=False) + if self._host_flags is None: + self._host_flags = torch.empty( + (self._ep_size,), + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + if self._copy_event is None: + self._copy_event = torch.cuda.Event(blocking=False) + host_flags = self._host_flags + event = self._copy_event with torch.cuda.device(flags.device), torch.cuda.stream(self._copy_stream): host_flags.copy_(flags.detach(), non_blocking=True) event.record(self._copy_stream) @@ -139,6 +146,8 @@ def _read_cuda_flags(self, flags: torch.Tensor) -> tuple[int, ...]: remaining_s = deadline_s - time.monotonic() if remaining_s <= 0: self._retired_copies.append((host_flags, event)) + self._host_flags = None + self._copy_event = None raise CompletionFlagReadTimeout( "timed out copying AlltoAll completion flags to host" ) @@ -175,8 +184,8 @@ def __init__( completion_reader: CompletionFlagReader, timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, - health: Optional[EPGroupHealthLike] = None, - on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, + health: EPGroupHealthLike | None = None, + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None = None, ) -> None: if ep_size <= 0: raise ValueError(f"ep_size must be > 0, got {ep_size}") @@ -196,7 +205,7 @@ def __init__( self._on_timeout = on_timeout self._cv = threading.Condition() - self._queue: Deque[_CollectiveWatch] = deque() + self._queue: deque[_CollectiveWatch] = deque() self._closed = False self._stopping = False self._thread: threading.Thread | None = None @@ -213,8 +222,8 @@ def from_workspace( ep_size: int, timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, - health: Optional[EPGroupHealthLike] = None, - on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, + health: EPGroupHealthLike | None = None, + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None = None, ) -> "AlltoAllWatchdog": """Build a watchdog from the MoE AlltoAll workspace and metainfo.""" dispatch_offset = int( @@ -421,7 +430,7 @@ def _run(self) -> None: except CompletionFlagReadTimeout: observed_flags = last_observed_flags poll_timed_out = True - except BaseException as exc: # noqa: BLE001 - keep watchdog failures visible. + except Exception as exc: # noqa: BLE001 - keep watchdog failures visible. with self._cv: self._last_error = exc self._queue.clear() diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index fe2c1020d201..b16c2a25bc6d 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -9,6 +9,7 @@ import os import sys +import threading from dataclasses import dataclass from typing import Callable, Dict, Optional @@ -212,6 +213,8 @@ def __init__( "mnnvl_mem": mnnvl_mem, "workspace": workspace, "metainfo": metainfo, + "watchdog_flag_generation": 0, + "watchdog_flag_generation_lock": threading.Lock(), } else: assert self._WORKSPACE[ @@ -229,17 +232,20 @@ def __init__( self.mnnvl_mem = self._WORKSPACE["mnnvl_mem"] self.workspace = self._WORKSPACE["workspace"] self.metainfo = self._WORKSPACE["metainfo"] + if "watchdog_flag_generation_lock" not in self._WORKSPACE: + self._WORKSPACE["watchdog_flag_generation_lock"] = threading.Lock() + self._WORKSPACE[ + "watchdog_flag_generation"] = self._read_current_flag_val() # Internal state self._state: _A2AState = _A2AState() self.ep_group_health = ep_group_health self._destroyed = False - self._watchdog_flag_generation = 0 self._alltoall_watchdog: AlltoAllWatchdog | None = None if (alltoall_watchdog_timeout_s is None and self.ep_group_health is not None): alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S if alltoall_watchdog_timeout_s is not None: - self._watchdog_flag_generation = self._read_current_flag_val() + self._sync_watchdog_flag_generation() self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( workspace=self.workspace, metainfo=self.metainfo, @@ -276,6 +282,25 @@ def _read_current_flag_val(self) -> int: flag_val = flag_val.detach().cpu() return int(flag_val.item()) + def _sync_watchdog_flag_generation(self) -> None: + workspace_state = self._WORKSPACE + assert workspace_state is not None + lock = workspace_state["watchdog_flag_generation_lock"] + with lock: + workspace_state["watchdog_flag_generation"] = max( + int(workspace_state["watchdog_flag_generation"]), + self._read_current_flag_val(), + ) + + def _next_watchdog_flag_generation(self) -> int: + workspace_state = self._WORKSPACE + assert workspace_state is not None + lock = workspace_state["watchdog_flag_generation_lock"] + with lock: + workspace_state["watchdog_flag_generation"] = ( + int(workspace_state["watchdog_flag_generation"]) + 1) + return int(workspace_state["watchdog_flag_generation"]) + def _get_active_rank_mask_tensor( self, active_rank_mask: Optional[torch.Tensor]) -> Optional[torch.Tensor]: @@ -302,10 +327,9 @@ def _watch_collective(self, phase: str, active_rank_mask: Optional[torch.Tensor]) -> None: if self._alltoall_watchdog is None: return - self._watchdog_flag_generation += 1 self._alltoall_watchdog.watch( phase=phase, - expected_flag=self._watchdog_flag_generation, + expected_flag=self._next_watchdog_flag_generation(), active_mask=self._active_mask_int(active_rank_mask), ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index 9b3d72306775..db5966615356 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py @@ -25,6 +25,7 @@ """ import os +import threading from typing import Callable, Dict, List, Optional, Tuple import torch @@ -287,6 +288,8 @@ def __init__( "mnnvl_mem": mnnvl_mem, "workspace": workspace, "metainfo": metainfo, + "watchdog_flag_generation": 0, + "watchdog_flag_generation_lock": threading.Lock(), } NVLinkOneSided._WORKSPACES[self._workspace_key] = workspace_state else: @@ -312,17 +315,20 @@ def __init__( NVLinkOneSided._WORKSPACE_REFCOUNTS.get(self._workspace_key, 0) + 1 ) self._destroyed = False + self._workspace_state = workspace_state self.mnnvl_mem = workspace_state["mnnvl_mem"] self.workspace = workspace_state["workspace"] self.moe_a2a_metainfo = workspace_state["metainfo"] self.max_num_tokens_per_rank = workspace_state["max_num_tokens_per_rank"] + if "watchdog_flag_generation_lock" not in workspace_state: + workspace_state["watchdog_flag_generation_lock"] = threading.Lock() + workspace_state["watchdog_flag_generation"] = self._read_current_flag_val() self.ep_group_health = ep_group_health - self._watchdog_flag_generation = 0 self._alltoall_watchdog: AlltoAllWatchdog | None = None if alltoall_watchdog_timeout_s is None and self.ep_group_health is not None: alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S if alltoall_watchdog_timeout_s is not None: - self._watchdog_flag_generation = self._read_current_flag_val() + self._sync_watchdog_flag_generation() self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( workspace=self.workspace, metainfo=self.moe_a2a_metainfo, @@ -354,6 +360,22 @@ def _read_current_flag_val(self) -> int: flag_val = flag_val.detach().cpu() return int(flag_val.item()) + def _sync_watchdog_flag_generation(self) -> None: + lock = self._workspace_state["watchdog_flag_generation_lock"] + with lock: + self._workspace_state["watchdog_flag_generation"] = max( + int(self._workspace_state["watchdog_flag_generation"]), + self._read_current_flag_val(), + ) + + def _next_watchdog_flag_generation(self) -> int: + lock = self._workspace_state["watchdog_flag_generation_lock"] + with lock: + self._workspace_state["watchdog_flag_generation"] = ( + int(self._workspace_state["watchdog_flag_generation"]) + 1 + ) + return int(self._workspace_state["watchdog_flag_generation"]) + def _get_active_rank_mask_tensor( self, active_rank_mask: Optional[torch.Tensor] ) -> Optional[torch.Tensor]: @@ -374,10 +396,9 @@ def _active_mask_int(self, active_rank_mask: Optional[torch.Tensor]) -> Optional def _watch_collective(self, phase: str, active_rank_mask: Optional[torch.Tensor]) -> None: if self._alltoall_watchdog is None: return - self._watchdog_flag_generation += 1 self._alltoall_watchdog.watch( phase=phase, - expected_flag=self._watchdog_flag_generation, + expected_flag=self._next_watchdog_flag_generation(), active_mask=self._active_mask_int(active_rank_mask), ) @@ -424,6 +445,7 @@ def destroy(self): self.mnnvl_mem = None self.workspace = None self.moe_a2a_metainfo = None + self._workspace_state = None self._dispatch_state = {"phase": "destroyed"} def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: diff --git a/tests/unittest/_torch/modules/moe/test_moe_comm.py b/tests/unittest/_torch/modules/moe/test_moe_comm.py index 057f87a95f7a..ac75c895cf6d 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_comm.py +++ b/tests/unittest/_torch/modules/moe/test_moe_comm.py @@ -225,14 +225,31 @@ def _read_nvlink_topk_target_ranks( return raw.view(torch.int32).view(max_num_tokens, top_k).cpu() -def _run_nvlink_rank_mask_dispatch_combine( +def _read_nvlink_topk_send_indices( + comm: NVLinkOneSided, + max_num_tokens: int, + top_k: int, +) -> torch.Tensor: + """Read topk_send_indices[max_num_tokens, top_k] from NVLinkOneSided workspace.""" + from tensorrt_llm.bindings import internal as _tllm_internal + + offset_index = int(_tllm_internal.thop.MOE_A2A_TOPK_SEND_INDICES_OFFSET_INDEX) + offset = comm.moe_a2a_metainfo[offset_index].item() + raw = comm.workspace[ + comm.ep_rank, + offset : offset + max_num_tokens * top_k * 4, + ] + return raw.view(torch.int32).view(max_num_tokens, top_k).cpu() + + +def _run_nvlink_rank_mask_dispatch( comm: NVLinkOneSided, token_selected_experts: torch.Tensor, payload: torch.Tensor, runtime_max_tokens_per_rank: int, active_rank_mask: Optional[torch.Tensor], -) -> Tuple[torch.Tensor, torch.Tensor]: - """Run raw NVLink one-sided dispatch/combine with an optional active rank mask.""" +) -> Tuple[List[torch.Tensor], int, torch.Tensor, torch.Tensor]: + """Run raw NVLink one-sided dispatch with an optional active rank mask.""" recv_tensors, combine_payload_offset, _ = torch.ops.trtllm.moe_a2a_dispatch( token_selected_experts, [payload], @@ -252,24 +269,105 @@ def _run_nvlink_rank_mask_dispatch_combine( runtime_max_tokens_per_rank, comm.top_k, ) + topk_send_indices = _read_nvlink_topk_send_indices( + comm, + runtime_max_tokens_per_rank, + comm.top_k, + ) + return recv_tensors, int(combine_payload_offset), topk_target_ranks, topk_send_indices - combined = torch.ops.trtllm.moe_a2a_combine( - recv_tensors[0], - token_selected_experts.size(0), + +def _run_nvlink_rank_mask_combine( + comm: NVLinkOneSided, + combine_payload: torch.Tensor, + local_num_tokens: int, + runtime_max_tokens_per_rank: int, + combine_payload_offset: int, + active_rank_mask: Optional[torch.Tensor], +) -> torch.Tensor: + """Run raw NVLink one-sided combine with an optional active rank mask.""" + return torch.ops.trtllm.moe_a2a_combine( + combine_payload, + local_num_tokens, comm.workspace, comm.moe_a2a_metainfo, runtime_max_tokens_per_rank, comm.ep_rank, comm.ep_size, comm.top_k, - int(combine_payload_offset), + combine_payload_offset, False, # payload_in_workspace False, # use_low_precision active_rank_mask, ) + + +def _run_nvlink_rank_mask_dispatch_combine( + comm: NVLinkOneSided, + token_selected_experts: torch.Tensor, + payload: torch.Tensor, + runtime_max_tokens_per_rank: int, + active_rank_mask: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + """Run raw NVLink one-sided dispatch/combine with an optional active rank mask.""" + recv_tensors, combine_payload_offset, topk_target_ranks, _ = _run_nvlink_rank_mask_dispatch( + comm, + token_selected_experts, + payload, + runtime_max_tokens_per_rank, + active_rank_mask, + ) + combined = _run_nvlink_rank_mask_combine( + comm, + recv_tensors[0], + token_selected_experts.size(0), + runtime_max_tokens_per_rank, + combine_payload_offset, + active_rank_mask, + ) return combined.cpu(), topk_target_ranks +def _expected_nvlink_rank_mask_combine_output( + comm: NVLinkOneSided, + payload: torch.Tensor, + topk_target_ranks: torch.Tensor, + topk_send_indices: torch.Tensor, + local_num_tokens: int, + runtime_max_tokens_per_rank: int, + dead_ranks: Set[int], +) -> torch.Tensor: + """Compute combine output from dispatched workspace while skipping dead ranks.""" + from tensorrt_llm.bindings import internal as _tllm_internal + + hidden_size = payload.shape[-1] + expected = torch.zeros( + (local_num_tokens, hidden_size), + dtype=payload.dtype, + device=payload.device, + ) + payload_offset_index = int(_tllm_internal.thop.MOE_A2A_PAYLOAD_DATA_OFFSET_INDEX) + payload_offset = comm.moe_a2a_metainfo[payload_offset_index].item() + bytes_per_rank = ( + comm.ep_size * runtime_max_tokens_per_rank * hidden_size * payload.element_size() + ) + + for token_idx in range(local_num_tokens): + for k in range(comm.top_k): + target_rank = int(topk_target_ranks[token_idx, k].item()) + dst_idx = int(topk_send_indices[token_idx, k].item()) + if dst_idx < 0 or target_rank in dead_ranks: + continue + raw = comm.workspace[target_rank, payload_offset : payload_offset + bytes_per_rank] + recv_payload = raw.view(payload.dtype).view( + comm.ep_size, + runtime_max_tokens_per_rank, + hidden_size, + ) + expected[token_idx] += recv_payload[comm.ep_rank, dst_idx] + return expected.cpu() + + # ============================================================================ # Source Encoding Utilities # ============================================================================ @@ -992,6 +1090,7 @@ def _worker_rank_mask_all_active_matches_no_mask(config: CommTestConfig) -> dict ) return { + "rank": rank, "output_eq": torch.equal(out_no_mask, out_all_active), "topk_eq": torch.equal(topk_no_mask, topk_all_active), } @@ -1040,7 +1139,7 @@ def _worker_rank_mask_one_rank_masked( if rank == dead_rank: MPI.COMM_WORLD.barrier() - return {"status": "dead"} + return {"rank": rank, "status": "dead"} local_num_tokens = config.all_num_tokens[rank] torch.manual_seed(0xA2A + rank) @@ -1069,6 +1168,7 @@ def _worker_rank_mask_one_rank_masked( MPI.COMM_WORLD.barrier() return { + "rank": rank, "status": "alive", "combined": combined, "topk_target_ranks": topk_target_ranks, @@ -1082,6 +1182,83 @@ def _worker_rank_mask_one_rank_masked( comm.destroy() +def _worker_rank_mask_inactive_before_combine( + config: CommTestConfig, + dead_rank: int, +) -> dict: + """Dispatch with all ranks active, then omit one rank from combine's active mask.""" + rank = tllm.mpi_rank() + torch.cuda.set_device(rank) + + comm = None + try: + mapping = Mapping( + rank=rank, + tp_size=config.ep_size, + moe_ep_size=config.ep_size, + world_size=config.ep_size, + ) + comm = create_comm_object(config.comm_type, mapping, config) + + local_num_tokens = config.all_num_tokens[rank] + torch.manual_seed(0xA2A + rank) + token_selected_experts = torch.randint( + 0, + config.num_experts, + (local_num_tokens, config.top_k), + dtype=torch.int32, + device="cuda", + ) + payload = _make_rank_mask_payload(local_num_tokens, config.hidden_size, rank) + + recv_tensors, combine_payload_offset, topk_target_ranks, topk_send_indices = ( + _run_nvlink_rank_mask_dispatch( + comm, + token_selected_experts, + payload, + local_num_tokens, + active_rank_mask=_ep_mask_words(config.ep_size, dead_ranks=set()), + ) + ) + + if rank == dead_rank: + MPI.COMM_WORLD.barrier() + return {"rank": rank, "status": "dead"} + + dead_ranks = {dead_rank} + expected = _expected_nvlink_rank_mask_combine_output( + comm, + payload, + topk_target_ranks, + topk_send_indices, + local_num_tokens, + local_num_tokens, + dead_ranks, + ) + combined = _run_nvlink_rank_mask_combine( + comm, + recv_tensors[0], + local_num_tokens, + local_num_tokens, + combine_payload_offset, + active_rank_mask=_ep_mask_words(config.ep_size, dead_ranks=dead_ranks), + ).cpu() + + MPI.COMM_WORLD.barrier() + return { + "rank": rank, + "status": "alive", + "combined": combined, + "expected": expected, + } + except Exception: + traceback.print_exc() + raise + finally: + if comm is not None and hasattr(comm, "destroy"): + comm.destroy() + + # ============================================================================ # Verification Functions # ============================================================================ @@ -1807,7 +1984,7 @@ def _make_postquant_test_params(): @pytest.fixture(autouse=True) -def setup_test(): +def setup_test() -> None: torch.manual_seed(0x1234) tllm.logger.set_level("error") @@ -1886,7 +2063,8 @@ def _run_rank_mask_all_active_test( ) ) - for rank, result in enumerate(results): + for result in results: + rank = result["rank"] assert result["output_eq"], ( f"rank {rank}: combine output differs between no-mask and all-active mask" ) @@ -1915,7 +2093,8 @@ def _run_rank_mask_one_rank_masked_test( ) saw_dead = False - for rank, result in enumerate(results): + for result in results: + rank = result["rank"] if result["status"] == "dead": assert rank == dead_rank saw_dead = True @@ -1952,6 +2131,45 @@ def _run_rank_mask_one_rank_masked_test( assert saw_dead, f"dead rank {dead_rank} did not appear in results" +def _run_rank_mask_inactive_before_combine_test( + mpi_pool_executor, + dead_rank: int, + local_num_tokens: int, + top_k: int, +) -> None: + ep_size = mpi_pool_executor.num_workers + config = _make_rank_mask_config(ep_size, local_num_tokens, top_k) + _skip_if_rank_mask_config_unsupported(config) + assert 0 <= dead_rank < ep_size + + worker_args = [(config, dead_rank)] * config.ep_size + results = list( + mpi_pool_executor.map( + _worker_rank_mask_inactive_before_combine, + *zip(*worker_args), + ) + ) + + saw_dead = False + for result in results: + rank = result["rank"] + if result["status"] == "dead": + assert rank == dead_rank + saw_dead = True + continue + + assert result["status"] == "alive" + combined = result["combined"] + expected = result["expected"] + assert combined is not None + assert expected is not None + assert torch.equal(combined, expected), ( + f"rank {rank}: combine output included a rank masked inactive before combine" + ) + + assert saw_dead, f"dead rank {dead_rank} did not appear in results" + + # ============================================================================ # Test Class # ============================================================================ @@ -2019,7 +2237,7 @@ def test_moe_comm_rank_mask_all_active_matches_no_mask( mpi_pool_executor, local_num_tokens: int, top_k: int, - ): + ) -> None: """Verify all-active active_rank_mask matches omitted mask for NVLinkOneSided.""" _run_rank_mask_all_active_test(mpi_pool_executor, local_num_tokens, top_k) @@ -2039,7 +2257,7 @@ def test_moe_comm_rank_mask_one_rank_masked_completes( dead_rank: int, local_num_tokens: int, top_k: int, - ): + ) -> None: """Verify masked-dead rank is skipped by raw NVLinkOneSided moe_a2a ops.""" _run_rank_mask_one_rank_masked_test( mpi_pool_executor, @@ -2047,3 +2265,26 @@ def test_moe_comm_rank_mask_one_rank_masked_completes( local_num_tokens, top_k, ) + + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize( + "mpi_pool_executor,dead_rank,local_num_tokens,top_k", + [ + (4, 2, 16, 2), + ], + indirect=["mpi_pool_executor"], + ) + def test_moe_comm_rank_mask_inactive_before_combine_skips_stale_dispatch_slots( + self, + mpi_pool_executor, + dead_rank: int, + local_num_tokens: int, + top_k: int, + ) -> None: + """Verify combine skips slots from a rank masked inactive after dispatch.""" + _run_rank_mask_inactive_before_combine_test( + mpi_pool_executor, + dead_rank, + local_num_tokens, + top_k, + ) diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index cda4d813d146..b44140dbd21f 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -16,6 +16,7 @@ import threading import time +from collections.abc import Callable from types import SimpleNamespace import pytest @@ -69,7 +70,7 @@ def read_completion_flags(self, phase: str) -> tuple[int, ...]: raise CompletionFlagReadTimeout("blocked") -def _wait_for(predicate, timeout_s: float = 1.0) -> None: +def _wait_for(predicate: Callable[[], bool], timeout_s: float = 1.0) -> None: deadline = time.monotonic() + timeout_s while time.monotonic() < deadline: if predicate(): @@ -126,7 +127,9 @@ def test_watchdog_stop_is_terminal() -> None: watchdog.watch(phase="dispatch", expected_flag=1) -def test_wide_ep_ft_options_create_shared_health_when_enabled(monkeypatch) -> None: +def test_wide_ep_ft_options_create_shared_health_when_enabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") model_config = SimpleNamespace( extra_attrs={}, From a249e277bd3d42a50eccda5219991f344167718f Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 29 Jun 2026 11:40:25 -0700 Subject: [PATCH 05/14] fix: accept advanced watchdog completion flags Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 4 ++-- .../_torch/modules/test_alltoall_watchdog.py | 22 +++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index c41a3fb8fdb3..81d13438f10a 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -339,7 +339,7 @@ def _active_ranks(self, active_mask: int) -> tuple[int, ...]: def _phase_complete(self, watch: _CollectiveWatch, observed_flags: tuple[int, ...]) -> bool: return all( - observed_flags[rank] == watch.expected_flag + observed_flags[rank] >= watch.expected_flag for rank in self._active_ranks(watch.active_mask) ) @@ -349,7 +349,7 @@ def _missing_ranks( return tuple( rank for rank in self._active_ranks(watch.active_mask) - if observed_flags[rank] != watch.expected_flag + if observed_flags[rank] < watch.expected_flag ) def _handle_timeout( diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index b44140dbd21f..5d8a7826c70e 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -101,6 +101,28 @@ def test_watchdog_completes_when_all_active_flags_arrive() -> None: assert health.all_active() is True +def test_watchdog_completes_when_flags_advance_past_expected_generation() -> None: + health = EPGroupHealth(4) + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("dispatch", [2, 1, 5, 3]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + assert health.all_active() is True + + def test_watchdog_defaults_match_design_doc() -> None: reader = FakeCompletionFlagReader(ep_size=1) watchdog = AlltoAllWatchdog(ep_size=1, ep_rank=0, completion_reader=reader) From bcbb8c37fc3b99d69477ed3cac7e4735a114f46e Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 29 Jun 2026 12:38:33 -0700 Subject: [PATCH 06/14] fix: harden shared AlltoAll watchdog state Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 247 +++++++++++++++++- .../_torch/distributed/moe_alltoall.py | 109 ++------ .../communication/nvlink_one_sided.py | 105 +++----- .../_torch/modules/test_alltoall_watchdog.py | 178 ++++++++++++- 4 files changed, 464 insertions(+), 175 deletions(-) diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 81d13438f10a..1cb3fe88aae7 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -26,7 +26,7 @@ import threading import time from collections import deque -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Mapping, MutableMapping, Sequence from dataclasses import dataclass from typing import Protocol @@ -38,6 +38,23 @@ DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S = 5.0 DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S = 0.1 UNKNOWN_COMPLETION_FLAG = -(2**63) +_COMPLETION_FLAG_MASK = (1 << 32) - 1 +_COMPLETION_FLAG_HALF_RANGE = 1 << 31 +_WORKSPACE_WATCHDOG_STATE_KEY = "alltoall_watchdog_shared_state" +_WORKSPACE_WATCHDOG_STATE_INIT_LOCK = threading.Lock() + + +def _normalize_completion_flag(value: int) -> int: + return int(value) & _COMPLETION_FLAG_MASK + + +def _completion_flag_reached(observed: int, expected: int) -> bool: + if observed == UNKNOWN_COMPLETION_FLAG: + return False + # Counter values are ordered modulo uint32. A queued watch cannot lag by + # half the counter space within the bounded watchdog timeout. + distance = (_normalize_completion_flag(observed) - expected) & _COMPLETION_FLAG_MASK + return distance < _COMPLETION_FLAG_HALF_RANGE class CompletionFlagReader(Protocol): @@ -53,6 +70,9 @@ class EPGroupHealthLike(Protocol): def get_mask(self) -> int: """Return the active-rank bitmask.""" + def get_mask_words(self) -> tuple[int, ...]: + """Return the active-rank bitmask split into uint64 words.""" + def mark_failed(self, rank: int) -> bool: """Mark ``rank`` failed and return whether state changed.""" @@ -82,6 +102,25 @@ class _CollectiveWatch: start_s: float +@dataclass(frozen=True) +class _SharedWatchdogConfig: + ep_size: int + ep_rank: int + timeout_s: float + poll_interval_s: float + health: EPGroupHealthLike | None + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None + + +@dataclass +class _WorkspaceWatchdogState: + lock: threading.Lock + generation: int = 0 + watchdog: AlltoAllWatchdog | None = None + config: _SharedWatchdogConfig | None = None + ref_count: int = 0 + + class _TorchCompletionFlagReader: """Completion-flag reader backed by the MoE AlltoAll workspace tensor.""" @@ -166,6 +205,168 @@ def read_completion_flags(self, phase: str) -> tuple[int, ...]: return tuple(int(v) for v in flags.tolist()) +class AlltoAllWatchdogCoordinator: + """Shared watchdog plumbing for MoE AlltoAll frontends.""" + + def __init__( + self, + *, + workspace_state: MutableMapping[str, object], + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + health: EPGroupHealthLike | None = None, + ) -> None: + self._workspace_state = workspace_state + self._workspace = workspace + self._metainfo = metainfo + self._metainfo_index = metainfo_index + self._ep_rank = ep_rank + self._health = health + + def _find_shared_state(self) -> _WorkspaceWatchdogState | None: + state = self._workspace_state.get(_WORKSPACE_WATCHDOG_STATE_KEY) + if state is None: + return None + if not isinstance(state, _WorkspaceWatchdogState): + raise TypeError("invalid shared AlltoAll watchdog workspace state") + return state + + def _get_shared_state(self) -> _WorkspaceWatchdogState: + state = self._find_shared_state() + if state is not None: + return state + with _WORKSPACE_WATCHDOG_STATE_INIT_LOCK: + state = self._workspace_state.get(_WORKSPACE_WATCHDOG_STATE_KEY) + if state is None: + state = _WorkspaceWatchdogState(lock=threading.Lock()) + self._workspace_state[_WORKSPACE_WATCHDOG_STATE_KEY] = state + if not isinstance(state, _WorkspaceWatchdogState): + raise TypeError("invalid shared AlltoAll watchdog workspace state") + return state + + def read_current_flag_val(self) -> int: + flag_val_offset = int(self._metainfo[self._metainfo_index["FLAG_VAL_OFFSET_INDEX"]].item()) + flag_val = self._workspace[self._ep_rank, flag_val_offset : flag_val_offset + 4].view( + torch.int32 + ) + if flag_val.device.type != "cpu": + flag_val = flag_val.detach().cpu() + return _normalize_completion_flag(int(flag_val.item())) + + def active_rank_mask_tensor(self, active_rank_mask: torch.Tensor | None) -> torch.Tensor | None: + if active_rank_mask is not None: + return active_rank_mask + if self._health is None: + return None + return torch.tensor(self._health.get_mask_words(), dtype=torch.uint64, device="cpu") + + def active_mask_int(self, active_rank_mask: torch.Tensor | None) -> int | None: + if active_rank_mask is not None: + mask_cpu = active_rank_mask.detach().cpu() + return sum(int(word) << (64 * idx) for idx, word in enumerate(mask_cpu.tolist())) + if self._health is not None: + return self._health.get_mask() + return None + + def acquire_watchdog( + self, + *, + ep_size: int, + timeout_s: float, + poll_interval_s: float, + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None = None, + ) -> AlltoAllWatchdog: + config = _SharedWatchdogConfig( + ep_size=int(ep_size), + ep_rank=self._ep_rank, + timeout_s=float(timeout_s), + poll_interval_s=float(poll_interval_s), + health=self._health, + on_timeout=on_timeout, + ) + state = self._get_shared_state() + with state.lock: + if state.watchdog is None: + state.generation = self.read_current_flag_val() + state.watchdog = AlltoAllWatchdog.from_workspace( + workspace=self._workspace, + metainfo=self._metainfo, + metainfo_index=self._metainfo_index, + ep_rank=self._ep_rank, + ep_size=ep_size, + timeout_s=timeout_s, + poll_interval_s=poll_interval_s, + health=self._health, + on_timeout=on_timeout, + ) + state.config = config + elif not self._matching_config(state.config, config): + raise ValueError( + "AlltoAll wrappers sharing a workspace must use the same " + "watchdog configuration and EP health object" + ) + state.ref_count += 1 + watchdog = state.watchdog + assert watchdog is not None + return watchdog + + @staticmethod + def _matching_config( + existing: _SharedWatchdogConfig | None, + requested: _SharedWatchdogConfig, + ) -> bool: + return ( + existing is not None + and existing.ep_size == requested.ep_size + and existing.ep_rank == requested.ep_rank + and existing.timeout_s == requested.timeout_s + and existing.poll_interval_s == requested.poll_interval_s + and existing.health is requested.health + and existing.on_timeout is requested.on_timeout + ) + + def release_watchdog(self, watchdog: AlltoAllWatchdog) -> None: + state = self._get_shared_state() + watchdog_to_stop: AlltoAllWatchdog | None = None + with state.lock: + if state.watchdog is not watchdog or state.ref_count <= 0: + raise RuntimeError("attempted to release an unregistered AlltoAll watchdog") + state.ref_count -= 1 + if state.ref_count == 0: + watchdog_to_stop = state.watchdog + state.watchdog = None + state.config = None + if watchdog_to_stop is not None: + watchdog_to_stop.stop() + + def watch_collective( + self, + watchdog: AlltoAllWatchdog | None, + phase: str, + active_rank_mask: torch.Tensor | None, + ) -> None: + if watchdog is None: + state = self._find_shared_state() + if state is not None: + with state.lock: + if state.watchdog is not None: + state.generation = (state.generation + 1) & _COMPLETION_FLAG_MASK + return + active_mask = self.active_mask_int(active_rank_mask) + state = self._get_shared_state() + with state.lock: + if state.watchdog is not watchdog: + raise RuntimeError("AlltoAll watchdog is not registered for this workspace") + state.generation = (state.generation + 1) & _COMPLETION_FLAG_MASK + watchdog.watch( + phase=phase, + expected_flag=state.generation, + active_mask=active_mask, + ) + + class AlltoAllWatchdog: """Background host thread that watches AlltoAll completion flags. @@ -279,7 +480,7 @@ def stop(self, timeout_s: float | None = None) -> None: self._queue.clear() self._cv.notify_all() thread = self._thread - if thread is not None: + if thread is not None and thread is not threading.current_thread(): thread.join(timeout=timeout_s) def watch( @@ -292,8 +493,10 @@ def watch( """Queue a just-launched AlltoAll phase for watchdog polling.""" if phase not in self.VALID_PHASES: raise ValueError(f"phase must be one of {sorted(self.VALID_PHASES)}, got {phase!r}") - if expected_flag < 0: - raise ValueError(f"expected_flag must be non-negative, got {expected_flag}") + if not 0 <= expected_flag <= _COMPLETION_FLAG_MASK: + raise ValueError( + f"expected_flag must be in [0, {_COMPLETION_FLAG_MASK}], got {expected_flag}" + ) if active_mask is None: if self._health is not None: active_mask = self._health.get_mask() @@ -339,7 +542,7 @@ def _active_ranks(self, active_mask: int) -> tuple[int, ...]: def _phase_complete(self, watch: _CollectiveWatch, observed_flags: tuple[int, ...]) -> bool: return all( - observed_flags[rank] >= watch.expected_flag + _completion_flag_reached(observed_flags[rank], watch.expected_flag) for rank in self._active_ranks(watch.active_mask) ) @@ -349,7 +552,7 @@ def _missing_ranks( return tuple( rank for rank in self._active_ranks(watch.active_mask) - if observed_flags[rank] < watch.expected_flag + if not _completion_flag_reached(observed_flags[rank], watch.expected_flag) ) def _handle_timeout( @@ -362,8 +565,11 @@ def _handle_timeout( elapsed_s = time.monotonic() - watch.start_s missing_ranks = self._missing_ranks(watch, observed_flags) marked_failed: list[int] = [] - has_known_flags = UNKNOWN_COMPLETION_FLAG not in observed_flags - if self._health is not None and (has_known_flags or not poll_timed_out): + if ( + self._health is not None + and not poll_timed_out + and UNKNOWN_COMPLETION_FLAG not in observed_flags + ): for rank in missing_ranks: if rank == self._ep_rank: continue @@ -405,6 +611,14 @@ def _handle_timeout( if self._on_timeout is not None: self._on_timeout(event) + def _stop_after_error(self, exc: Exception) -> None: + with self._cv: + self._last_error = exc + self._closed = True + self._stopping = True + self._queue.clear() + self._cv.notify_all() + def _run(self) -> None: last_observed_flags = tuple(UNKNOWN_COMPLETION_FLAG for _ in range(self._ep_size)) poll_timed_out = False @@ -418,7 +632,8 @@ def _run(self) -> None: try: observed_flags = tuple( - int(v) for v in self._completion_reader.read_completion_flags(watch.phase) + _normalize_completion_flag(int(v)) + for v in self._completion_reader.read_completion_flags(watch.phase) ) if len(observed_flags) != self._ep_size: raise RuntimeError( @@ -431,10 +646,7 @@ def _run(self) -> None: observed_flags = last_observed_flags poll_timed_out = True except Exception as exc: # noqa: BLE001 - keep watchdog failures visible. - with self._cv: - self._last_error = exc - self._queue.clear() - self._cv.notify_all() + self._stop_after_error(exc) tllm_logger.error("AlltoAll watchdog stopped after polling error: %s", exc) return @@ -448,7 +660,14 @@ def _run(self) -> None: continue if time.monotonic() - watch.start_s >= self._timeout_s: - self._handle_timeout(watch, observed_flags, poll_timed_out=poll_timed_out) + try: + self._handle_timeout(watch, observed_flags, poll_timed_out=poll_timed_out) + except Exception as exc: # noqa: BLE001 - keep watchdog failures visible. + self._stop_after_error(exc) + tllm_logger.error( + "AlltoAll watchdog stopped after timeout handling error: %s", exc + ) + return with self._cv: # The GPU stream is no longer trustworthy once a collective # times out. Drop queued follow-on phases so they do not diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index b16c2a25bc6d..6e44510a8573 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -9,7 +9,6 @@ import os import sys -import threading from dataclasses import dataclass from typing import Callable, Dict, Optional @@ -19,7 +18,7 @@ from tensorrt_llm._torch.alltoall_watchdog import ( DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, AlltoAllWatchdog, - AlltoAllWatchdogTimeout) + AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, EPGroupHealthLike) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -132,13 +131,13 @@ def __init__( num_slots: int, workspace_size_per_rank: int, num_experts: Optional[int] = None, - ep_group_health=None, + ep_group_health: Optional[EPGroupHealthLike] = None, alltoall_watchdog_timeout_s: Optional[float] = None, alltoall_watchdog_poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, alltoall_watchdog_on_timeout: Optional[Callable[ [AlltoAllWatchdogTimeout], None]] = None, - ): + ) -> None: """ Initialize MoeAlltoAll with workspace allocation. @@ -213,8 +212,6 @@ def __init__( "mnnvl_mem": mnnvl_mem, "workspace": workspace, "metainfo": metainfo, - "watchdog_flag_generation": 0, - "watchdog_flag_generation_lock": threading.Lock(), } else: assert self._WORKSPACE[ @@ -232,29 +229,31 @@ def __init__( self.mnnvl_mem = self._WORKSPACE["mnnvl_mem"] self.workspace = self._WORKSPACE["workspace"] self.metainfo = self._WORKSPACE["metainfo"] - if "watchdog_flag_generation_lock" not in self._WORKSPACE: - self._WORKSPACE["watchdog_flag_generation_lock"] = threading.Lock() - self._WORKSPACE[ - "watchdog_flag_generation"] = self._read_current_flag_val() # Internal state self._state: _A2AState = _A2AState() self.ep_group_health = ep_group_health + workspace_state = self._WORKSPACE + assert workspace_state is not None + metainfo_index = self._METAINFO_INDEX + assert metainfo_index is not None + self._watchdog_coordinator = AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=self.workspace, + metainfo=self.metainfo, + metainfo_index=metainfo_index, + ep_rank=self.ep_rank, + health=self.ep_group_health, + ) self._destroyed = False self._alltoall_watchdog: AlltoAllWatchdog | None = None if (alltoall_watchdog_timeout_s is None and self.ep_group_health is not None): alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S if alltoall_watchdog_timeout_s is not None: - self._sync_watchdog_flag_generation() - self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( - workspace=self.workspace, - metainfo=self.metainfo, - metainfo_index=self._METAINFO_INDEX, - ep_rank=self.ep_rank, + self._alltoall_watchdog = self._watchdog_coordinator.acquire_watchdog( ep_size=self.ep_size, timeout_s=alltoall_watchdog_timeout_s, poll_interval_s=alltoall_watchdog_poll_interval_s, - health=self.ep_group_health, on_timeout=alltoall_watchdog_on_timeout, ) @@ -265,74 +264,13 @@ def destroy(self) -> None: self._destroyed = True watchdog = getattr(self, "_alltoall_watchdog", None) if watchdog is not None: - watchdog.stop(timeout_s=1.0) + self._watchdog_coordinator.release_watchdog(watchdog) self._alltoall_watchdog = None def __del__(self) -> None: if not sys.is_finalizing(): self.destroy() - def _read_current_flag_val(self) -> int: - flag_val_offset = self.metainfo[ - self._METAINFO_INDEX["FLAG_VAL_OFFSET_INDEX"]].item() - flag_val = self.workspace[self.ep_rank, - flag_val_offset:flag_val_offset + 4].view( - torch.int32) - if flag_val.device.type != "cpu": - flag_val = flag_val.detach().cpu() - return int(flag_val.item()) - - def _sync_watchdog_flag_generation(self) -> None: - workspace_state = self._WORKSPACE - assert workspace_state is not None - lock = workspace_state["watchdog_flag_generation_lock"] - with lock: - workspace_state["watchdog_flag_generation"] = max( - int(workspace_state["watchdog_flag_generation"]), - self._read_current_flag_val(), - ) - - def _next_watchdog_flag_generation(self) -> int: - workspace_state = self._WORKSPACE - assert workspace_state is not None - lock = workspace_state["watchdog_flag_generation_lock"] - with lock: - workspace_state["watchdog_flag_generation"] = ( - int(workspace_state["watchdog_flag_generation"]) + 1) - return int(workspace_state["watchdog_flag_generation"]) - - def _get_active_rank_mask_tensor( - self, - active_rank_mask: Optional[torch.Tensor]) -> Optional[torch.Tensor]: - if active_rank_mask is not None: - return active_rank_mask - if self.ep_group_health is None: - return None - return torch.tensor(self.ep_group_health.get_mask_words(), - dtype=torch.uint64, - device="cpu") - - def _active_mask_int( - self, active_rank_mask: Optional[torch.Tensor]) -> Optional[int]: - if active_rank_mask is not None: - mask_cpu = active_rank_mask.detach().cpu() - return sum( - int(word) << (64 * idx) - for idx, word in enumerate(mask_cpu.tolist())) - if self.ep_group_health is not None: - return self.ep_group_health.get_mask() - return None - - def _watch_collective(self, phase: str, - active_rank_mask: Optional[torch.Tensor]) -> None: - if self._alltoall_watchdog is None: - return - self._alltoall_watchdog.watch( - phase=phase, - expected_flag=self._next_watchdog_flag_generation(), - active_mask=self._active_mask_int(active_rank_mask), - ) - def dispatch(self, token_selected_experts: torch.Tensor, input_payloads: list[torch.Tensor], @@ -366,7 +304,8 @@ def dispatch(self, 0 ) == self.eplb_stats_num_experts, "eplb_local_stats size must match eplb_stats_num_experts" - active_rank_mask = self._get_active_rank_mask_tensor(active_rank_mask) + active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( + active_rank_mask) recv_tensors, combine_payload_offset, eplb_gathered_stats = torch.ops.trtllm.moe_a2a_dispatch( token_selected_experts, input_payloads, @@ -380,7 +319,9 @@ def dispatch(self, eplb_local_stats, active_rank_mask, ) - self._watch_collective("dispatch", active_rank_mask) + self._watchdog_coordinator.watch_collective(self._alltoall_watchdog, + "dispatch", + active_rank_mask) if eplb_gathered_stats.numel() == 0: eplb_gathered_stats = None @@ -428,13 +369,15 @@ def combine( assert self._state.phase == "dispatched", "combine called before a successful dispatch" assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" - active_rank_mask = self._get_active_rank_mask_tensor(active_rank_mask) + active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( + active_rank_mask) output = torch.ops.trtllm.moe_a2a_combine( payload, self._state.local_num_tokens, self.workspace, self.metainfo, runtime_max_tokens_per_rank, self.ep_rank, self.ep_size, self.top_k, self._state.combine_payload_offset, payload_in_workspace, use_low_precision_combine, active_rank_mask) - self._watch_collective("combine", active_rank_mask) + self._watchdog_coordinator.watch_collective(self._alltoall_watchdog, + "combine", active_rank_mask) # Reset state for next round self.reset_state() diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index db5966615356..0811abc9b5d0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py @@ -25,7 +25,6 @@ """ import os -import threading from typing import Callable, Dict, List, Optional, Tuple import torch @@ -35,7 +34,9 @@ DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, AlltoAllWatchdog, + AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, + EPGroupHealthLike, ) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger @@ -158,11 +159,11 @@ def __init__( dtype: Optional[torch.dtype] = None, num_experts: Optional[int] = None, use_low_precision_combine: bool = False, - ep_group_health=None, + ep_group_health: EPGroupHealthLike | None = None, alltoall_watchdog_timeout_s: Optional[float] = None, alltoall_watchdog_poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, alltoall_watchdog_on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, - ): + ) -> None: """ Initialize NVLinkOneSided with workspace allocation. @@ -288,8 +289,6 @@ def __init__( "mnnvl_mem": mnnvl_mem, "workspace": workspace, "metainfo": metainfo, - "watchdog_flag_generation": 0, - "watchdog_flag_generation_lock": threading.Lock(), } NVLinkOneSided._WORKSPACES[self._workspace_key] = workspace_state else: @@ -320,28 +319,27 @@ def __init__( self.workspace = workspace_state["workspace"] self.moe_a2a_metainfo = workspace_state["metainfo"] self.max_num_tokens_per_rank = workspace_state["max_num_tokens_per_rank"] - if "watchdog_flag_generation_lock" not in workspace_state: - workspace_state["watchdog_flag_generation_lock"] = threading.Lock() - workspace_state["watchdog_flag_generation"] = self._read_current_flag_val() self.ep_group_health = ep_group_health + self._watchdog_coordinator = AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=self.workspace, + metainfo=self.moe_a2a_metainfo, + metainfo_index={ + "FLAG_VAL_OFFSET_INDEX": self.FLAG_VAL_OFFSET_INDEX, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": self.DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": self.COMBINE_COMPLETION_FLAGS_OFFSET_INDEX, + }, + ep_rank=self.ep_rank, + health=self.ep_group_health, + ) self._alltoall_watchdog: AlltoAllWatchdog | None = None if alltoall_watchdog_timeout_s is None and self.ep_group_health is not None: alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S if alltoall_watchdog_timeout_s is not None: - self._sync_watchdog_flag_generation() - self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( - workspace=self.workspace, - metainfo=self.moe_a2a_metainfo, - metainfo_index={ - "FLAG_VAL_OFFSET_INDEX": self.FLAG_VAL_OFFSET_INDEX, - "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": self.DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX, - "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": self.COMBINE_COMPLETION_FLAGS_OFFSET_INDEX, - }, - ep_rank=self.ep_rank, + self._alltoall_watchdog = self._watchdog_coordinator.acquire_watchdog( ep_size=self.ep_size, timeout_s=alltoall_watchdog_timeout_s, poll_interval_s=alltoall_watchdog_poll_interval_s, - health=self.ep_group_health, on_timeout=alltoall_watchdog_on_timeout, ) @@ -351,57 +349,6 @@ def __init__( # Invalid token expert ID (default to -1), the kernels in TRTLLM-gen is hard-code to support -1 only. self.invalid_token_expert_id: int = -1 - def _read_current_flag_val(self) -> int: - flag_val_offset = self.moe_a2a_metainfo[self.FLAG_VAL_OFFSET_INDEX].item() - flag_val = self.workspace[self.ep_rank, flag_val_offset : flag_val_offset + 4].view( - torch.int32 - ) - if flag_val.device.type != "cpu": - flag_val = flag_val.detach().cpu() - return int(flag_val.item()) - - def _sync_watchdog_flag_generation(self) -> None: - lock = self._workspace_state["watchdog_flag_generation_lock"] - with lock: - self._workspace_state["watchdog_flag_generation"] = max( - int(self._workspace_state["watchdog_flag_generation"]), - self._read_current_flag_val(), - ) - - def _next_watchdog_flag_generation(self) -> int: - lock = self._workspace_state["watchdog_flag_generation_lock"] - with lock: - self._workspace_state["watchdog_flag_generation"] = ( - int(self._workspace_state["watchdog_flag_generation"]) + 1 - ) - return int(self._workspace_state["watchdog_flag_generation"]) - - def _get_active_rank_mask_tensor( - self, active_rank_mask: Optional[torch.Tensor] - ) -> Optional[torch.Tensor]: - if active_rank_mask is not None: - return active_rank_mask - if self.ep_group_health is None: - return None - return torch.tensor(self.ep_group_health.get_mask_words(), dtype=torch.uint64, device="cpu") - - def _active_mask_int(self, active_rank_mask: Optional[torch.Tensor]) -> Optional[int]: - if active_rank_mask is not None: - mask_cpu = active_rank_mask.detach().cpu() - return sum(int(word) << (64 * idx) for idx, word in enumerate(mask_cpu.tolist())) - if self.ep_group_health is not None: - return self.ep_group_health.get_mask() - return None - - def _watch_collective(self, phase: str, active_rank_mask: Optional[torch.Tensor]) -> None: - if self._alltoall_watchdog is None: - return - self._alltoall_watchdog.watch( - phase=phase, - expected_flag=self._next_watchdog_flag_generation(), - active_mask=self._active_mask_int(active_rank_mask), - ) - @staticmethod def is_platform_supported() -> bool: """ @@ -422,7 +369,7 @@ def destroy(self): self._destroyed = True if self._alltoall_watchdog is not None: - self._alltoall_watchdog.stop(timeout_s=1.0) + self._watchdog_coordinator.release_watchdog(self._alltoall_watchdog) self._alltoall_watchdog = None workspace_key = getattr(self, "_workspace_key", None) if workspace_key is None: @@ -508,7 +455,9 @@ def dispatch( assert eplb_local_stats.size(0) == self.eplb_stats_num_experts, ( "eplb_local_stats size must match eplb_stats_num_experts" ) - active_rank_mask = self._get_active_rank_mask_tensor(kwargs.get("active_rank_mask")) + active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( + kwargs.get("active_rank_mask") + ) recv_buffers, combine_payload_offset, eplb_gathered_stats = ( torch.ops.trtllm.moe_a2a_dispatch( @@ -525,7 +474,9 @@ def dispatch( active_rank_mask, ) ) - self._watch_collective("dispatch", active_rank_mask) + self._watchdog_coordinator.watch_collective( + self._alltoall_watchdog, "dispatch", active_rank_mask + ) if eplb_gathered_stats.numel() == 0: eplb_gathered_stats = None self._dispatch_state["eplb_gathered_stats"] = eplb_gathered_stats @@ -628,7 +579,9 @@ def combine( raise ValueError( f"final_hidden_states must be 2D or 3D, got {final_hidden_states.dim()}D" ) - active_rank_mask = self._get_active_rank_mask_tensor(kwargs.get("active_rank_mask")) + active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( + kwargs.get("active_rank_mask") + ) output = torch.ops.trtllm.moe_a2a_combine( final_hidden_states, int(local_num_tokens), @@ -643,7 +596,9 @@ def combine( bool(self.use_low_precision_combine), active_rank_mask, ) - self._watch_collective("combine", active_rank_mask) + self._watchdog_coordinator.watch_collective( + self._alltoall_watchdog, "combine", active_rank_mask + ) # Reset state for next round self.reset_state() diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index 5d8a7826c70e..0028f8a67685 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -27,6 +27,7 @@ DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, UNKNOWN_COMPLETION_FLAG, AlltoAllWatchdog, + AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, CompletionFlagReadTimeout, ) @@ -123,6 +124,33 @@ def test_watchdog_completes_when_flags_advance_past_expected_generation() -> Non assert health.all_active() is True +def test_watchdog_handles_signed_uint32_generation_boundaries() -> None: + reader = FakeCompletionFlagReader(ep_size=2) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=2, + ep_rank=0, + completion_reader=reader, + timeout_s=0.05, + poll_interval_s=0.005, + on_timeout=events.append, + ) as watchdog: + reader.set_flags("dispatch", [-(1 << 31), -(1 << 31)]) + watchdog.watch(phase="dispatch", expected_flag=1 << 31) + assert watchdog.wait_until_idle(timeout_s=1.0) + + reader.set_flags("dispatch", [-1, -1]) + watchdog.watch(phase="dispatch", expected_flag=(1 << 32) - 1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + reader.set_flags("dispatch", [0, 0]) + watchdog.watch(phase="dispatch", expected_flag=0) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + + def test_watchdog_defaults_match_design_doc() -> None: reader = FakeCompletionFlagReader(ep_size=1) watchdog = AlltoAllWatchdog(ep_size=1, ep_rank=0, completion_reader=reader) @@ -268,7 +296,7 @@ def test_watchdog_poll_timeout_without_snapshot_fails_closed() -> None: assert health.all_active() is True -def test_watchdog_poll_timeout_with_prior_snapshot_marks_known_missing_rank() -> None: +def test_watchdog_poll_timeout_with_prior_snapshot_does_not_mark_failed_rank() -> None: health = EPGroupHealth(3) events: list[AlltoAllWatchdogTimeout] = [] @@ -288,8 +316,33 @@ def test_watchdog_poll_timeout_with_prior_snapshot_marks_known_missing_rank() -> assert event.poll_timed_out is True assert event.observed_flags == (1, 0, 1) assert event.missing_ranks == (1,) - assert event.marked_failed_ranks == (1,) - assert health.get_failed_ranks() == frozenset({1}) + assert event.marked_failed_ranks == () + assert health.all_active() is True + + +def test_watchdog_callback_error_stops_and_clears_queue() -> None: + health = EPGroupHealth(2) + reader = FakeCompletionFlagReader(ep_size=2) + reader.set_flags("dispatch", [1, 0]) + + def raise_from_callback(event: AlltoAllWatchdogTimeout) -> None: + raise RuntimeError(f"callback failed for {event.phase}") + + with AlltoAllWatchdog( + ep_size=2, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=raise_from_callback, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: watchdog.last_error is not None) + assert watchdog.wait_until_idle(timeout_s=1.0) + assert isinstance(watchdog.last_error, RuntimeError) + with pytest.raises(RuntimeError, match="stopped AlltoAllWatchdog"): + watchdog.watch(phase="dispatch", expected_flag=2) def test_watchdog_preserves_fifo_order_and_clears_followups_after_timeout() -> None: @@ -359,6 +412,125 @@ def test_watchdog_from_workspace_reads_phase_specific_offsets() -> None: assert health.get_failed_ranks() == frozenset({0}) +def test_workspace_coordinators_share_fifo_watchdog() -> None: + ep_size = 3 + ep_rank = 0 + workspace_state: dict[str, object] = {} + workspace = torch.zeros((ep_size, 64), dtype=torch.uint8) + metainfo = torch.tensor([0, 4, 16], dtype=torch.int64) + metainfo_index = { + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 1, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 2, + } + workspace[ep_rank, 4:16].view(torch.int32).copy_(torch.tensor([1, 0, 1])) + health = EPGroupHealth(ep_size) + events: list[AlltoAllWatchdogTimeout] = [] + on_timeout = events.append + coordinators = [ + AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=ep_rank, + health=health, + ) + for _ in range(2) + ] + watchdogs = [ + coordinator.acquire_watchdog( + ep_size=ep_size, + timeout_s=0.02, + poll_interval_s=0.005, + on_timeout=on_timeout, + ) + for coordinator in coordinators + ] + + try: + assert watchdogs[0] is watchdogs[1] + coordinators[0].watch_collective(watchdogs[0], "dispatch", None) + coordinators[1].watch_collective(watchdogs[1], "combine", None) + _wait_for(lambda: len(events) == 1) + assert watchdogs[0].wait_until_idle(timeout_s=1.0) + time.sleep(0.05) + + assert len(events) == 1 + assert events[0].phase == "dispatch" + assert events[0].missing_ranks == (1,) + assert health.get_failed_ranks() == frozenset({1}) + finally: + for coordinator, watchdog in zip(coordinators, watchdogs): + coordinator.release_watchdog(watchdog) + + +def test_workspace_coordinator_wraps_shared_generation() -> None: + workspace_state: dict[str, object] = {} + workspace = torch.zeros((1, 32), dtype=torch.uint8) + workspace[0, 0:4].view(torch.int32).fill_(-1) + metainfo = torch.tensor([0, 4, 8], dtype=torch.int64) + metainfo_index = { + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 1, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 2, + } + coordinator = AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=0, + ) + watchdog = coordinator.acquire_watchdog( + ep_size=1, + timeout_s=0.05, + poll_interval_s=0.005, + ) + + try: + coordinator.watch_collective(watchdog, "dispatch", None) + assert watchdog.wait_until_idle(timeout_s=1.0) + finally: + coordinator.release_watchdog(watchdog) + + +def test_unmonitored_coordinator_advances_shared_generation() -> None: + workspace_state: dict[str, object] = {} + workspace = torch.zeros((2, 32), dtype=torch.uint8) + metainfo = torch.tensor([0, 4, 12], dtype=torch.int64) + metainfo_index = { + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 1, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 2, + } + coordinators = [ + AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=0, + ) + for _ in range(2) + ] + events: list[AlltoAllWatchdogTimeout] = [] + watchdog = coordinators[0].acquire_watchdog( + ep_size=2, + timeout_s=0.02, + poll_interval_s=0.005, + on_timeout=events.append, + ) + + try: + coordinators[1].watch_collective(None, "dispatch", None) + coordinators[0].watch_collective(watchdog, "dispatch", None) + _wait_for(lambda: len(events) == 1) + assert events[0].expected_flag == 2 + finally: + coordinators[0].release_watchdog(watchdog) + + def test_watchdog_rejects_active_mask_without_local_rank() -> None: reader = FakeCompletionFlagReader(ep_size=4) with AlltoAllWatchdog( From 1c11081c9e856e30af89e01357e06c1642006fb3 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 12:43:40 -0700 Subject: [PATCH 07/14] fix: make AlltoAll watchdog detection-only Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 72 ++++++++++++------- .../_torch/distributed/moe_alltoall.py | 16 +++-- .../communication/nvlink_one_sided.py | 15 ++-- .../modules/fused_moe/ep_group_health.py | 13 ++-- .../_torch/modules/fused_moe/wide_ep_ft.py | 15 ++-- .../_torch/modules/test_alltoall_watchdog.py | 62 ++++++++++++---- 6 files changed, 135 insertions(+), 58 deletions(-) diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 1cb3fe88aae7..166d4507e926 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -19,6 +19,12 @@ silent-spin failure mode never writes its slot, so this watchdog polls the same table from a CPU thread and reports peers whose flags do not reach the expected generation before a bounded timeout. + +Timeouts are detection events, not committed membership changes. The optional +EP health object supplies the already-committed active peers to watch, and is +read-only here. Higher-layer recovery coordination consumes ``on_timeout`` +events and commits a new membership only after the data plane and expert +placement are ready for the same generation. """ from __future__ import annotations @@ -65,7 +71,7 @@ def read_completion_flags(self, phase: str) -> Sequence[int]: class EPGroupHealthLike(Protocol): - """Subset of EPGroupHealth used by the watchdog.""" + """Read-only committed EP membership used by AlltoAll frontends.""" def get_mask(self) -> int: """Return the active-rank bitmask.""" @@ -73,9 +79,6 @@ def get_mask(self) -> int: def get_mask_words(self) -> tuple[int, ...]: """Return the active-rank bitmask split into uint64 words.""" - def mark_failed(self, rank: int) -> bool: - """Mark ``rank`` failed and return whether state changed.""" - class CompletionFlagReadTimeout(TimeoutError): """Raised when the host watchdog cannot read completion flags in time.""" @@ -83,13 +86,17 @@ class CompletionFlagReadTimeout(TimeoutError): @dataclass(frozen=True) class AlltoAllWatchdogTimeout: - """Details emitted when an AlltoAll phase times out.""" + """Detection details emitted when an AlltoAll phase times out. + + ``missing_ranks`` contains suspects relative to the committed active mask + captured for this collective. Receiving this event does not publish a new + active mask. + """ phase: str expected_flag: int observed_flags: tuple[int, ...] missing_ranks: tuple[int, ...] - marked_failed_ranks: tuple[int, ...] elapsed_s: float poll_timed_out: bool = False @@ -206,7 +213,11 @@ def read_completion_flags(self, phase: str) -> tuple[int, ...]: class AlltoAllWatchdogCoordinator: - """Shared watchdog plumbing for MoE AlltoAll frontends.""" + """Shared watchdog plumbing for MoE AlltoAll frontends. + + ``health`` is the committed membership snapshot source. This class reads + it when a dispatch starts but never publishes detected failures into it. + """ def __init__( self, @@ -256,12 +267,38 @@ def read_current_flag_val(self) -> int: return _normalize_completion_flag(int(flag_val.item())) def active_rank_mask_tensor(self, active_rank_mask: torch.Tensor | None) -> torch.Tensor | None: + """Capture a caller override or the current committed mask. + + The returned tensor does not alias a caller-owned override, so it can + safely represent the mask used by one dispatch/combine pair. + """ if active_rank_mask is not None: - return active_rank_mask + return active_rank_mask.detach().clone() if self._health is None: return None return torch.tensor(self._health.get_mask_words(), dtype=torch.uint64, device="cpu") + def active_rank_mask_for_combine( + self, + dispatch_active_rank_mask: torch.Tensor | None, + requested_active_rank_mask: torch.Tensor | None, + ) -> torch.Tensor | None: + """Reuse the dispatch mask and reject a conflicting combine override. + + This guarantees one local mask snapshot across a dispatch/combine pair. + Atomic membership across layers and iterations remains the recovery + coordinator's responsibility. + """ + if requested_active_rank_mask is None: + return dispatch_active_rank_mask + requested_mask = self.active_rank_mask_tensor(requested_active_rank_mask) + assert requested_mask is not None + if dispatch_active_rank_mask is None or not torch.equal( + dispatch_active_rank_mask, requested_mask + ): + raise ValueError("active_rank_mask must match the mask captured at dispatch") + return dispatch_active_rank_mask + def active_mask_int(self, active_rank_mask: torch.Tensor | None) -> int | None: if active_rank_mask is not None: mask_cpu = active_rank_mask.detach().cpu() @@ -372,7 +409,8 @@ class AlltoAllWatchdog: The watchdog is intentionally opt-in. Callers queue phases with :meth:`watch`; the thread polls them in FIFO order so a queued combine cannot - hide a still-spinning dispatch. + hide a still-spinning dispatch. A timeout is reported through + ``on_timeout`` without mutating the committed EP membership. """ VALID_PHASES = frozenset({"dispatch", "combine"}) @@ -564,24 +602,11 @@ def _handle_timeout( ) -> None: elapsed_s = time.monotonic() - watch.start_s missing_ranks = self._missing_ranks(watch, observed_flags) - marked_failed: list[int] = [] - if ( - self._health is not None - and not poll_timed_out - and UNKNOWN_COMPLETION_FLAG not in observed_flags - ): - for rank in missing_ranks: - if rank == self._ep_rank: - continue - if self._health.mark_failed(rank): - marked_failed.append(rank) - event = AlltoAllWatchdogTimeout( phase=watch.phase, expected_flag=watch.expected_flag, observed_flags=observed_flags, missing_ranks=missing_ranks, - marked_failed_ranks=tuple(marked_failed), elapsed_s=elapsed_s, poll_timed_out=poll_timed_out, ) @@ -589,14 +614,13 @@ def _handle_timeout( tllm_logger.error( "AlltoAll watchdog could not read completion flags on rank %d " "during %s before timeout %.3fs; expected flag %d, active " - "ranks %s, observed flags %s, marked ranks %s", + "ranks %s, observed flags %s; reporting detection event only", self._ep_rank, watch.phase, elapsed_s, watch.expected_flag, list(self._active_ranks(watch.active_mask)), list(observed_flags), - list(marked_failed), ) else: tllm_logger.warning( diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index 6e44510a8573..f1692d2473fc 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -31,6 +31,7 @@ class _A2AState: local_num_tokens: int | None = None combine_payload_offset: int | None = None eplb_gathered_stats: torch.Tensor | None = None + active_rank_mask: torch.Tensor | None = None class MoeAlltoAll: @@ -149,8 +150,8 @@ def __init__( Note: The terminology is mapped to `num_experts` in this class and the kernels. num_experts: (Optional) Number of experts for EPLB stats (must be <= num_slots). DO NOT provide this parameter if EPLB is not enabled. Note: The terminology is mapped to `eplb_stats_num_experts` in this class and the kernels. - ep_group_health: Optional EPGroupHealth-compatible object. When present, its mask is passed to the - CUDA kernels and used by the watchdog. + ep_group_health: Optional read-only committed EP membership. When present, its mask is passed to the + CUDA kernels and defines the peers expected by the watchdog. Timeout detection never mutates it. alltoall_watchdog_timeout_s: Optional timeout for the host-side AlltoAll watchdog. If None, the watchdog is disabled. alltoall_watchdog_poll_interval_s: Poll interval for the watchdog thread. @@ -289,7 +290,8 @@ def dispatch(self, invalid_token_expert_id: If not None, set the token_selected_experts of the invalid tokens to this expert id. This is used to notify the MoE to skip these tokens for GroupGEMM. expert_id_payload_index: The index of token_selected_experts in the input_payloads. Must be provided if invalid_token_expert_id is not None. eplb_local_stats: (Optional) [num_experts] tensor containing local statistics for EPLB - active_rank_mask: Optional uint64 CPU tensor overriding ep_group_health for this dispatch. + active_rank_mask: Optional uint64 CPU tensor overriding committed membership for this dispatch. The + captured value is reused by combine. Returns: recv_tensors: List of tensors received, each has shape [ep_size, max_tokens_per_rank, payload_num_elements_per_token] @@ -329,6 +331,7 @@ def dispatch(self, self._state.local_num_tokens = token_selected_experts.size(0) self._state.combine_payload_offset = combine_payload_offset self._state.eplb_gathered_stats = eplb_gathered_stats + self._state.active_rank_mask = active_rank_mask self._state.phase = "dispatched" if invalid_token_expert_id is not None: @@ -361,7 +364,8 @@ def combine( runtime_max_tokens_per_rank: Maximum of the number of tokens of each DP rank's local batch. payload_in_workspace: If True, 'payload' is a view into 'workspace' at 'combine_payload_offset' and no staging copy is needed. If False, the op stages 'payload' into the workspace region before combining. use_low_precision_combine: If True, quantize the combine payload to FP8 for NVLink transfer (halves NVLink bandwidth usage, output precision is preserved). - active_rank_mask: Optional uint64 CPU tensor overriding ep_group_health for this combine. + active_rank_mask: Optional uint64 CPU tensor. If supplied, it must match the mask captured by dispatch + for this collective. Returns: combined_output: [local_num_tokens, num_elements_per_token] tensor of combined results @@ -369,8 +373,8 @@ def combine( assert self._state.phase == "dispatched", "combine called before a successful dispatch" assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" - active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( - active_rank_mask) + active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( + self._state.active_rank_mask, active_rank_mask) output = torch.ops.trtllm.moe_a2a_combine( payload, self._state.local_num_tokens, self.workspace, self.metainfo, runtime_max_tokens_per_rank, self.ep_rank, diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index 0811abc9b5d0..076e1f5e6b81 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py @@ -181,8 +181,8 @@ def __init__( use_low_precision_combine: If True, quantize the combine payload to FP8 for NVLink transfer (halves NVLink bandwidth usage, output precision is preserved). Corresponds to model_config.use_low_precision_moe_combine. - ep_group_health: Optional EPGroupHealth-compatible object. When present, its mask is passed to the - CUDA kernels and used by the watchdog. + ep_group_health: Optional read-only committed EP membership. When present, its mask is passed to + the CUDA kernels and defines the peers expected by the watchdog. Timeout detection never mutates it. alltoall_watchdog_timeout_s: Optional timeout for the host-side AlltoAll watchdog. If None, the watchdog is disabled. alltoall_watchdog_poll_interval_s: Poll interval for the watchdog thread. @@ -424,7 +424,8 @@ def dispatch( token_final_scales: Router weights [local_num_tokens, top_k] all_rank_num_tokens: Token counts per rank [ep_size] use_dp_padding: Whether to use DP padding (optional) - **kwargs: Strategy-specific arguments (unused) + **kwargs: Strategy-specific arguments. ``active_rank_mask`` may override the committed membership + for dispatch; the captured value is reused by combine. Returns: Tuple of (hidden_states, hidden_states_sf, token_selected_slots, token_final_scales) @@ -483,6 +484,7 @@ def dispatch( self._dispatch_state["combine_payload_offset"] = int(combine_payload_offset) self._dispatch_state["local_num_tokens"] = token_selected_slots.size(0) self._dispatch_state["runtime_max_tokens_per_rank"] = runtime_max_tokens_per_rank + self._dispatch_state["active_rank_mask"] = active_rank_mask self._dispatch_state["phase"] = "dispatched" # Extract results from recv_buffers @@ -545,6 +547,8 @@ def combine( final_hidden_states: Output from MoE computation Shape: [ep_size, max_tokens_per_rank, hidden_size] or [ep_size * max_tokens_per_rank, hidden_size] (will be reshaped) + **kwargs: Strategy-specific arguments. If ``active_rank_mask`` is supplied, it must match the mask + captured by dispatch for this collective. Returns: Combined output tensor [local_num_tokens, hidden_size] @@ -579,8 +583,9 @@ def combine( raise ValueError( f"final_hidden_states must be 2D or 3D, got {final_hidden_states.dim()}D" ) - active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( - kwargs.get("active_rank_mask") + active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( + self._dispatch_state.get("active_rank_mask"), + kwargs.get("active_rank_mask"), ) output = torch.ops.trtllm.moe_a2a_combine( final_hidden_states, diff --git a/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py b/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py index 1864d4ba825d..b1aba97c818c 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py +++ b/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py @@ -16,17 +16,18 @@ """EP group health tracking for WideEP fault tolerance. This module provides :class:`EPGroupHealth`, a process-local, thread-safe data -structure that records which ranks in an Expert Parallel (EP) group are currently -alive vs. failed. It is the single source of truth for EP rank health within one -process and is consumed by: +structure that records the committed active membership of an Expert Parallel +(EP) group. It is the single source of truth for the data-plane rank mask within +one process and is consumed by: * AlltoAll communication backends (rank masking on dispatch / combine) - * The host-side AlltoAll watchdog (failure detection) + * The host-side AlltoAll watchdog (read-only expected-peer snapshot) * The MoE load balancer (emergency-mask reconfiguration) * The model engine and PyExecutor (degraded health reporting) -Cross-process consensus on which ranks are dead (failure broadcast across the EP -group) is the responsibility of higher-layer coordination components and is not +Failure detection does not mutate this object directly. Cross-process consensus, +expert-placement preparation, and atomic publication of a new committed mask are +the responsibility of higher-layer recovery coordination components and are not performed here. """ diff --git a/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py index 69e5c53f061e..fb1586ae1054 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py +++ b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py @@ -26,17 +26,19 @@ from .ep_group_health import EPGroupHealth -_ENABLE_ENV = "TRTLLM_ENABLE_WIDE_EP_FT" +_ENABLE_ENV = "TLLM_FAULT_TOLERANCE_MODE" _TIMEOUT_ENV = "TRTLLM_ALLTOALL_WATCHDOG_TIMEOUT_S" _POLL_INTERVAL_ENV = "TRTLLM_ALLTOALL_WATCHDOG_POLL_INTERVAL_S" +# This object contains membership committed by higher-layer recovery +# coordination. Detection threads must treat it as read-only. _HEALTH_KEY = "wide_ep_ft_ep_group_health" _TIMEOUT_KEY = "alltoall_watchdog_timeout_s" _POLL_INTERVAL_KEY = "alltoall_watchdog_poll_interval_s" def _env_enabled() -> bool: - return os.environ.get(_ENABLE_ENV, "0").lower() in {"1", "true", "yes", "on"} + return os.environ.get(_ENABLE_ENV) == "1" def _float_option(extra_attrs: dict, key: str, env_name: str, default: float) -> float: @@ -50,12 +52,15 @@ def _float_option(extra_attrs: dict, key: str, env_name: str, default: float) -> def get_wide_ep_ft_options( model_config: Any, ) -> tuple[Optional[EPGroupHealth], Optional[float], float]: - """Return the shared EP health object and watchdog timing for a model. + """Return committed EP membership and watchdog timing for a model. WideEP FT remains opt-in until the integration PR wires a public model option. Callers can either inject ``wide_ep_ft_ep_group_health`` through - ``ModelConfig.extra_attrs`` or set ``TRTLLM_ENABLE_WIDE_EP_FT=1`` to create - one process-local health object shared by all MoE communication layers. + ``ModelConfig.extra_attrs`` or set ``TLLM_FAULT_TOLERANCE_MODE=1`` to create + one process-local membership object shared by all MoE communication layers. + The AlltoAll watchdog reads this object to determine expected peers and + reports suspects through its ``on_timeout`` seam; it never mutates the + committed membership directly. """ extra_attrs = getattr(model_config, "extra_attrs", {}) diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index 0028f8a67685..931b0dd5e5a6 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -180,7 +180,7 @@ def test_watchdog_stop_is_terminal() -> None: def test_wide_ep_ft_options_create_shared_health_when_enabled( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") + monkeypatch.setenv("TLLM_FAULT_TOLERANCE_MODE", "1") model_config = SimpleNamespace( extra_attrs={}, mapping=SimpleNamespace(moe_ep_size=4), @@ -197,8 +197,50 @@ def test_wide_ep_ft_options_create_shared_health_when_enabled( assert poll_again_s == poll_interval_s -def test_watchdog_timeout_reports_and_marks_missing_remote_ranks() -> None: +def test_wide_ep_ft_options_ignore_legacy_enable_flag(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("TLLM_FAULT_TOLERANCE_MODE", raising=False) + monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") + model_config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(moe_ep_size=4), + ) + + health, timeout_s, _ = get_wide_ep_ft_options(model_config) + + assert health is None + assert timeout_s is None + + +def test_watchdog_coordinator_reuses_dispatch_mask_for_combine() -> None: health = EPGroupHealth(4) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((4, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + dispatch_mask = coordinator.active_rank_mask_tensor(None) + assert dispatch_mask is not None + + # Simulate a higher-layer commit at an invalid mid-collective point. The + # local dispatch/combine pair must keep using its dispatch snapshot. + health.mark_failed(2) + combine_mask = coordinator.active_rank_mask_for_combine(dispatch_mask, None) + + assert combine_mask is dispatch_mask + assert combine_mask.tolist() == [0b1111, 0] + with pytest.raises(ValueError, match="mask captured at dispatch"): + coordinator.active_rank_mask_for_combine( + dispatch_mask, + torch.tensor(health.get_mask_words(), dtype=torch.uint64), + ) + + +def test_watchdog_no_detected_failure_publication_to_committed_health() -> None: + health = EPGroupHealth(4) + committed_before = health.snapshot() reader = FakeCompletionFlagReader(ep_size=4) reader.set_flags("dispatch", [1, 0, 1, 0]) events: list[AlltoAllWatchdogTimeout] = [] @@ -221,8 +263,8 @@ def test_watchdog_timeout_reports_and_marks_missing_remote_ranks() -> None: assert event.expected_flag == 1 assert event.observed_flags == (1, 0, 1, 0) assert event.missing_ranks == (1, 3) - assert event.marked_failed_ranks == (1, 3) - assert health.get_failed_ranks() == frozenset({1, 3}) + assert not hasattr(event, "marked_failed_ranks") + assert health.snapshot() == committed_before def test_watchdog_ignores_ranks_already_failed_in_health_mask() -> None: @@ -248,7 +290,7 @@ def test_watchdog_ignores_ranks_already_failed_in_health_mask() -> None: assert health.get_failed_ranks() == frozenset({2}) -def test_watchdog_reports_local_missing_but_does_not_mark_local_failed() -> None: +def test_watchdog_reports_local_missing_without_changing_committed_health() -> None: health = EPGroupHealth(4) reader = FakeCompletionFlagReader(ep_size=4) reader.set_flags("combine", [0, 2, 2, 2]) @@ -268,7 +310,6 @@ def test_watchdog_reports_local_missing_but_does_not_mark_local_failed() -> None event = events[0] assert event.missing_ranks == (0,) - assert event.marked_failed_ranks == () assert health.get_failed_ranks() == frozenset() @@ -292,7 +333,6 @@ def test_watchdog_poll_timeout_without_snapshot_fails_closed() -> None: assert event.poll_timed_out is True assert event.observed_flags == (UNKNOWN_COMPLETION_FLAG,) * 3 assert event.missing_ranks == (0, 1, 2) - assert event.marked_failed_ranks == () assert health.all_active() is True @@ -316,7 +356,6 @@ def test_watchdog_poll_timeout_with_prior_snapshot_does_not_mark_failed_rank() - assert event.poll_timed_out is True assert event.observed_flags == (1, 0, 1) assert event.missing_ranks == (1,) - assert event.marked_failed_ranks == () assert health.all_active() is True @@ -370,7 +409,7 @@ def test_watchdog_preserves_fifo_order_and_clears_followups_after_timeout() -> N assert len(events) == 1 assert events[0].phase == "dispatch" assert events[0].missing_ranks == (1,) - assert health.get_failed_ranks() == frozenset({1}) + assert health.get_failed_ranks() == frozenset() def test_watchdog_from_workspace_reads_phase_specific_offsets() -> None: @@ -408,8 +447,7 @@ def test_watchdog_from_workspace_reads_phase_specific_offsets() -> None: assert events[0].phase == "combine" assert events[0].missing_ranks == (0,) - assert events[0].marked_failed_ranks == (0,) - assert health.get_failed_ranks() == frozenset({0}) + assert health.get_failed_ranks() == frozenset() def test_workspace_coordinators_share_fifo_watchdog() -> None: @@ -459,7 +497,7 @@ def test_workspace_coordinators_share_fifo_watchdog() -> None: assert len(events) == 1 assert events[0].phase == "dispatch" assert events[0].missing_ranks == (1,) - assert health.get_failed_ranks() == frozenset({1}) + assert health.get_failed_ranks() == frozenset() finally: for coordinator, watchdog in zip(coordinators, watchdogs): coordinator.release_watchdog(watchdog) From f6bec2d41f9bca4b0ce138db93fe69d4a00450bd Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 26 Jun 2026 15:26:34 -0700 Subject: [PATCH 08/14] Add WideEP error classification patterns Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/pyexecutor/error_classification.py | 33 +++++- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../pyexecutor/test_error_classification.py | 104 ++++++++++++++++++ 3 files changed, 137 insertions(+), 1 deletion(-) create mode 100644 tests/unittest/_torch/pyexecutor/test_error_classification.py diff --git a/tensorrt_llm/_torch/pyexecutor/error_classification.py b/tensorrt_llm/_torch/pyexecutor/error_classification.py index c356e6ca3a62..be10331f6a9e 100644 --- a/tensorrt_llm/_torch/pyexecutor/error_classification.py +++ b/tensorrt_llm/_torch/pyexecutor/error_classification.py @@ -26,9 +26,26 @@ IMMEDIATE_FATAL_PATTERNS: list[str] = [ "cudaerrorillegaladdress", "cudaerrorlaunchfailure", + "cudaerrorcontextisdestroyed", + "cudaerrornodevice", "illegal memory access", "device-side assert", - "unrecoverable", + "cuda context destroyed", + "cuda context is destroyed", + "cuda context was destroyed", + "cuda context is unrecoverable", + "unrecoverable error in the engine", + "cuda error no device", + "gpu has fallen off the bus", + "xid 79", + "nccl communicator abort", + "nccl communicator was aborted", + "ncclcommabort", + "nvshmem peer unreachable", + "mpi rank terminated", + "mpi worker terminated", + "mpi rank exited unexpectedly", + "mpi worker exited unexpectedly", ] # Patterns that are serious but may be transient (e.g. a single OOM @@ -38,6 +55,20 @@ "cuda out of memory", "cuda error", "nccl error", + "nccl timeout", + "nccl operation timed out", + "alltoall timeout", + "all-to-all timeout", + "alltoall watchdog", + "completion_flags timeout", + "deep_ep buffer barrier hang", + "deepep buffer barrier hang", + "intranode::barrier", + "symmetric memory access violation", + "rdma timeout", + "nixl transfer failed", + "nixl transfer entered error state", + "nixl transfer wait timed out", ] diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 6c603e43a381..a0bbd6ab4f93 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -35,6 +35,7 @@ l0_a10: - unittest/_torch/executor/test_kv_pool_rebalance.py - unittest/_torch/executor/test_disagg_index_mapper_early_release.py - unittest/_torch/pyexecutor/test_kv_cache_compression_manager.py + - unittest/_torch/pyexecutor/test_error_classification.py - unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py - unittest/_torch/modules/dwdp/test_dwdp_manager.py - unittest/_torch/modules/dwdp/test_dwdp_mapping.py diff --git a/tests/unittest/_torch/pyexecutor/test_error_classification.py b/tests/unittest/_torch/pyexecutor/test_error_classification.py new file mode 100644 index 000000000000..9e90b2b55493 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_error_classification.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +import sys +from pathlib import Path + +import pytest + +_MODULE_PATH = ( + Path(__file__).resolve().parents[4] / "tensorrt_llm/_torch/pyexecutor/error_classification.py" +) +_SPEC = importlib.util.spec_from_file_location("error_classification_under_test", _MODULE_PATH) +assert _SPEC is not None +assert _SPEC.loader is not None +_ERROR_CLASSIFICATION = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _ERROR_CLASSIFICATION +_SPEC.loader.exec_module(_ERROR_CLASSIFICATION) + +ErrorBudget = _ERROR_CLASSIFICATION.ErrorBudget +classify_error = _ERROR_CLASSIFICATION.classify_error + + +@pytest.mark.parametrize( + "message", + [ + "CUDA context destroyed", + "cudaErrorContextIsDestroyed during collective setup", + "CUDA error: cudaErrorNoDevice on rank 37", + "NVRM: Xid 79, GPU has fallen off the bus", + "CUDA context is destroyed after device reset", + "CUDA context was destroyed after device reset", + "CUDA context is unrecoverable after device reset", + "Unrecoverable error in the engine", + "ncclCommAbort returned during communicator teardown", + "NCCL communicator was aborted during PP send", + "NVSHMEM peer unreachable for EP rank 12", + "MPI rank terminated with signal 9", + "MPI worker exited unexpectedly", + ], +) +def test_wide_ep_immediate_fatal_patterns(message): + assert classify_error(message) == "immediate_fatal" + + +@pytest.mark.parametrize( + "message", + [ + "NCCL timeout", + "AlltoAll timeout waiting for completion_flags[3][7]", + "AlltoAll watchdog timed out waiting for peer rank 7", + "completion_flags timeout", + "NCCL operation timed out during all_reduce", + "deep_ep buffer barrier hang", + "deepep buffer barrier hang", + "DeepEP Buffer.__del__ hung in intranode::barrier", + "Symmetric memory access violation reading dead peer", + "RDMA timeout on cross-node transfer", + "NIXL transfer failed: remote peer closed", + "NIXL transfer entered error state", + "NIXL transfer wait timed out after 5000 ms", + ], +) +def test_wide_ep_severe_patterns(message): + assert classify_error(message) == "severe" + + +@pytest.mark.parametrize( + "message", + [ + "AlltoAll slow path observed for rank 3", + "NCCL retry scheduled after transient transport hiccup", + "ECC correctable error threshold warning", + "Application marked the request unrecoverable but retryable", + ], +) +def test_wide_ep_nonfatal_signals_remain_transient(message): + assert classify_error(message) == "transient" + + +def test_error_budget_charges_wide_ep_severe_patterns(): + budget = ErrorBudget(recovery_rate=0.0) + + assert not budget.consume("NIXL transfer failed: remote peer closed") + assert budget.consume("AlltoAll timeout waiting for completion_flags[0][2]") + + +def test_error_budget_bypasses_budget_for_wide_ep_fatal_patterns(): + budget = ErrorBudget(recovery_rate=0.0) + + assert budget.consume("MPI rank terminated with signal 9") + assert budget.budget == 1.0 From e94306a46729cecaa825f2b9bf7572ec2f9be1f0 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 09:15:38 -0700 Subject: [PATCH 09/14] [TRTLLM-13548][feat] Add MPI FT subcommunicator and broadcast thread Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/distributed/communicator.py | 459 ++++ .../_torch/pyexecutor/ep_failure_broadcast.py | 1564 ++++++++++++ .../integration/test_lists/test-db/l0_a10.yml | 6 + .../_torch/distributed/test_communicator.py | 681 ++++++ .../_ep_failure_broadcast_mpi_worker.py | 622 +++++ .../pyexecutor/test_ep_failure_broadcast.py | 2107 +++++++++++++++++ .../test_ep_failure_broadcast_mpi.py | 687 ++++++ 7 files changed, 6126 insertions(+) create mode 100644 tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py create mode 100644 tests/unittest/_torch/distributed/test_communicator.py create mode 100644 tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py create mode 100644 tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py create mode 100644 tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast_mpi.py diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 6fbbf16a19df..7e78bf36ee66 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -1,6 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import math import pickle # nosec B403 +import threading from abc import ABC, abstractmethod +from dataclasses import dataclass from enum import IntEnum from functools import lru_cache, wraps from typing import Any, Callable, List, Optional, Tuple @@ -73,6 +90,448 @@ def reduce_op_to_mpi(op: ReduceOp) -> MPI.Op: return _reduce_op_to_mpi_dict[op] +@dataclass(frozen=True) +class MpiFtSubcommSetup: + """Collectively validated WideEP FT communicator startup state.""" + + comm: Any + local_rank: int + ep_size: int + ulfm_available: bool + + +_MPI_FT_PROCESS_LIFETIME_REFS: list[Any] = [] +_MPI_FT_PROCESS_LIFETIME_REFS_LOCK = threading.Lock() + + +def _retain_mpi_ft_comm(comm: Any) -> None: + """Keep a communicator with unsafe cleanup reachable until process exit.""" + with _MPI_FT_PROCESS_LIFETIME_REFS_LOCK: + _MPI_FT_PROCESS_LIFETIME_REFS.append(comm) + + +def _is_mpi_ft_int(value: Any) -> bool: + """Reject bool/float values that compare equal to MPI integer fields.""" + return isinstance(value, int) and not isinstance(value, bool) + + +def _mpi_ft_local_topology(mapping: Mapping, parent_rank: int, parent_size: int, + provided_thread_level: int, + health_size: Optional[int]) -> dict[str, Any]: + """Build a serializable startup record without raising locally.""" + world_size = None + mapping_rank = None + ep_group = None + ep_size = None + ep_rank = None + mapping_error = None + try: + world_size = mapping.world_size + mapping_rank = mapping.rank + ep_group = tuple(mapping.moe_ep_group) + ep_size = mapping.moe_ep_size + ep_rank = mapping.moe_ep_rank + except Exception as error: + mapping_error = ("invalid WideEP FT mapping topology: " + f"{type(error).__name__}: {error}") + + error_kind = None + error_message = None + if provided_thread_level < MPI.THREAD_MULTIPLE: + error_kind = "runtime" + error_message = ( + "WideEP FT requires MPI.THREAD_MULTIPLE because its control-plane " + "thread overlaps other MPI traffic " + f"(provided={provided_thread_level}, required={MPI.THREAD_MULTIPLE})" + ) + elif mapping_error is not None: + error_kind = "value" + error_message = mapping_error + elif not ep_group: + error_kind = "value" + error_message = "mapping.moe_ep_group must not be empty" + elif len(ep_group) != ep_size: + error_kind = "value" + error_message = ( + "mapping.moe_ep_group size must match mapping.moe_ep_size, " + f"got {len(ep_group)} and {ep_size}") + else: + try: + has_duplicate_ranks = len(set(ep_group)) != len(ep_group) + has_invalid_rank = any( + not isinstance(rank, int) or rank < 0 or rank >= parent_size + for rank in ep_group) + except TypeError as error: + error_kind = "value" + error_message = f"mapping.moe_ep_group contains invalid ranks: {error}" + else: + if has_duplicate_ranks: + error_kind = "value" + error_message = ( + "mapping.moe_ep_group contains duplicate ranks: " + f"{ep_group}") + elif parent_size != world_size: + error_kind = "runtime" + error_message = ( + "WideEP FT parent communicator size must match " + f"mapping.world_size, got {parent_size} and {world_size}") + elif parent_rank != mapping_rank: + error_kind = "runtime" + error_message = ( + "WideEP FT parent communicator rank must match mapping.rank, " + f"got {parent_rank} and {mapping_rank}") + elif has_invalid_rank: + error_kind = "value" + error_message = ( + "mapping.moe_ep_group contains a rank outside the parent " + f"communicator: {ep_group}") + elif ep_size != world_size or set(ep_group) != set( + range(parent_size)): + error_kind = "value" + error_message = ( + "WideEP FT MVP requires one MoE EP group spanning the " + "full MPI world; " + f"got world_size={world_size}, moe_ep_group={ep_group}") + elif health_size is not None and health_size != ep_size: + error_kind = "value" + error_message = ( + "EPGroupHealth size must match mapping.moe_ep_size, " + f"got {health_size} and {ep_size}") + elif parent_rank not in ep_group: + error_kind = "value" + error_message = ( + f"local rank {parent_rank} is not in its MoE EP group " + f"{ep_group}") + elif ep_group.index(parent_rank) != ep_rank: + error_kind = "value" + error_message = ( + "mapping.moe_ep_group order must match mapping.moe_ep_rank, " + f"got group index {ep_group.index(parent_rank)} and EP rank " + f"{ep_rank}") + + return { + "parent_rank": parent_rank, + "parent_size": parent_size, + "world_size": world_size, + "mapping_rank": mapping_rank, + "ep_group": ep_group, + "ep_size": ep_size, + "ep_rank": ep_rank, + "health_size": health_size, + "error_kind": error_kind, + "error_message": error_message, + } + + +def _validate_mpi_ft_topologies(topologies: List[Any], + parent_size: int) -> None: + """Raise the same validation error on every parent-communicator rank.""" + if len(topologies) != parent_size: + raise RuntimeError( + "WideEP FT startup allgather returned an unexpected number of " + f"topologies: got {len(topologies)}, expected {parent_size}") + + for gather_rank, topology in enumerate(topologies): + if not isinstance(topology, dict): + raise RuntimeError( + "WideEP FT startup allgather returned an invalid topology for " + f"parent rank {gather_rank}: {topology!r}") + + # Object allgather returns records in communicator-rank order. Selecting the + # first error makes every rank raise the same exception before Split. + for gather_rank, topology in enumerate(topologies): + error_message = topology.get("error_message") + if error_message is None: + continue + reporting_rank = topology.get("parent_rank", gather_rank) + message = ("WideEP FT startup validation failed on parent rank " + f"{reporting_rank}: {error_message}") + if topology.get("error_kind") == "value": + raise ValueError(message) + raise RuntimeError(message) + + for gather_rank, topology in enumerate(topologies): + reported_parent_rank = topology.get("parent_rank") + if (not _is_mpi_ft_int(reported_parent_rank) + or reported_parent_rank != gather_rank): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: allgather slot " + f"{gather_rank} reports parent rank " + f"{reported_parent_rank!r}") + reported_parent_size = topology.get("parent_size") + if (not _is_mpi_ft_int(reported_parent_size) + or reported_parent_size != parent_size): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports communicator size " + f"{reported_parent_size!r}, expected {parent_size}") + + ep_group = topology.get("ep_group") + if not isinstance(ep_group, (list, tuple)): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports invalid EP group {ep_group!r}") + if len(ep_group) != parent_size: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports EP group size {len(ep_group)}, " + f"expected {parent_size}") + if any( + isinstance(peer_rank, bool) or not _is_mpi_ft_int(peer_rank) + or peer_rank < 0 or peer_rank >= parent_size + for peer_rank in ep_group): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports invalid EP group ranks {ep_group!r}") + if len(set(ep_group)) != parent_size: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports duplicate EP group ranks {ep_group!r}") + world_size = topology.get("world_size") + if not _is_mpi_ft_int(world_size) or world_size != parent_size: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports world size " + f"{world_size!r}, expected {parent_size}") + mapping_rank = topology.get("mapping_rank") + if not _is_mpi_ft_int(mapping_rank) or mapping_rank != gather_rank: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports mapping rank " + f"{mapping_rank!r}") + ep_size = topology.get("ep_size") + if not _is_mpi_ft_int(ep_size) or ep_size != parent_size: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports EP size {ep_size!r}, " + f"expected {parent_size}") + ep_rank = topology.get("ep_rank") + if (not _is_mpi_ft_int(ep_rank) or ep_rank < 0 or ep_rank >= parent_size + or ep_group[ep_rank] != gather_rank): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports invalid EP rank {ep_rank!r} for " + f"group {ep_group!r}") + health_size = topology.get("health_size") + if health_size is not None and (not _is_mpi_ft_int(health_size) + or health_size != parent_size): + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports health size {health_size}, " + f"expected {parent_size}") + # Only dereference peer groups after every gathered record has passed the + # schema/range checks above. This keeps malformed records rank-attributed and + # prevents a valid earlier record from indexing into unchecked peer data. + for gather_rank, topology in enumerate(topologies): + ep_group = topology["ep_group"] + for peer_rank in ep_group: + peer_group = topologies[peer_rank].get("ep_group") + if peer_group != ep_group: + raise RuntimeError( + "WideEP FT startup topology is inconsistent: parent rank " + f"{gather_rank} reports EP group {ep_group}, but member " + f"rank {peer_rank} reports {peer_group}") + + +def _mpi_ft_post_split_status(ft_comm: Any, parent_rank: int, ep_rank: int, + ep_size: int) -> dict[str, Any]: + """Configure a derived communicator without raising before reconciliation.""" + error_message = None + ulfm_available = False + try: + # Install the non-fatal handler before any other operation on the new + # communicator. In particular, do not inspect a possibly broken handle + # while it still inherits the parent's fatal error policy. + ft_comm.Set_errhandler(MPI.ERRORS_RETURN) + actual_rank = ft_comm.Get_rank() + actual_size = ft_comm.Get_size() + if actual_size != ep_size or actual_rank != ep_rank: + error_message = ( + "MPI_Comm_split created an unexpected WideEP FT communicator: " + f"rank={actual_rank}, size={actual_size}, " + f"expected rank={ep_rank}, size={ep_size}") + except Exception as error: + error_message = ("WideEP FT communicator setup raised " + f"{type(error).__name__}: {error}") + + if error_message is None and callable(getattr( + ft_comm, "Is_revoked", None)) and callable( + getattr(ft_comm, "Revoke", None)): + try: + already_revoked = ft_comm.Is_revoked() + except Exception as error: + if not _is_unsupported_mpi_ulfm_error(error): + error_message = ("WideEP FT ULFM probe raised " + f"{type(error).__name__}: {error}") + else: + if already_revoked: + error_message = ( + "WideEP FT communicator is already revoked at construction") + else: + ulfm_available = True + return { + "parent_rank": parent_rank, + "error_message": error_message, + "ulfm_available": ulfm_available, + } + + +def _is_unsupported_mpi_ulfm_error(error: BaseException) -> bool: + if isinstance(error, NotImplementedError): + return True + mpi_exception = getattr(MPI, "Exception", None) + if not isinstance(mpi_exception, type) or not isinstance( + error, mpi_exception): + return False + get_error_class = getattr(error, "Get_error_class", None) + if not callable(get_error_class): + return False + unsupported_classes = { + value + for value in ( + getattr(MPI, "ERR_NOT_SUPPORTED", None), + getattr(MPI, "ERR_UNSUPPORTED_OPERATION", None), + ) if value is not None + } + try: + error_class = get_error_class() + except Exception: + # This helper runs while building a rank-local status that must reach + # the parent allgather. A broken exception classifier is not evidence of + # an unsupported operation, but it must not strand peer ranks either. + return False + return error_class in unsupported_classes + + +def _mpi_ft_post_split_result( + statuses: List[Any], + parent_size: int, +) -> tuple[Optional[str], bool]: + """Reconcile setup errors and the global ULFM capability.""" + if len(statuses) != parent_size: + return ( + "WideEP FT post-split allgather returned an unexpected number of " + f"statuses: got {len(statuses)}, expected {parent_size}", + False, + ) + for gather_rank, status in enumerate(statuses): + if not isinstance(status, dict): + return ( + "WideEP FT post-split allgather returned an invalid status for " + f"parent rank {gather_rank}: {status!r}", + False, + ) + if status.get("parent_rank") != gather_rank: + return ( + "WideEP FT post-split status is inconsistent: allgather slot " + f"{gather_rank} reports parent rank " + f"{status.get('parent_rank')}", + False, + ) + if not isinstance(status.get("ulfm_available"), bool): + return ( + "WideEP FT post-split allgather returned incomplete capability " + f"state for parent rank {gather_rank}: {status!r}", + False, + ) + + ulfm_available = all(status["ulfm_available"] for status in statuses) + for gather_rank, status in enumerate(statuses): + error_message = status.get("error_message") + if error_message is not None: + return ( + "WideEP FT communicator setup failed on parent rank " + f"{gather_rank}: {error_message}", + False, + ) + return None, ulfm_available + + +def create_mpi_ft_subcomm( + mapping: Mapping, + parent_comm: Optional[Any] = None, + health_size: Optional[int] = None) -> MpiFtSubcommSetup: + """Create the dedicated MPI control-plane communicator for WideEP FT. + + This operation is collective across ``parent_comm`` and must run on every + rank during startup, before any background failure-broadcast threads are + launched. The MVP requires one MoE EP group spanning the parent world, + ordered by the EP-local rank used by :class:`EPGroupHealth`. + + Before splitting, ranks exchange their local validation outcome and EP + topology on the healthy parent communicator. This prevents one rank from + raising locally while its peers block forever inside ``MPI_Comm_split``. + + Args: + mapping: Distributed topology for the local rank. + parent_comm: Parent MPI communicator. Defaults to TRT-LLM's active + ``mpi_comm()``, which wraps ``MPI.COMM_WORLD`` in the standard + launch path and preserves custom communicator sessions. + health_size: Optional local ``EPGroupHealth`` size to validate + collectively before splitting. + + Returns: + The world-spanning EP MPI communicator and its collectively validated + rank, size, and ULFM capability. + + Raises: + RuntimeError: If MPI is unavailable, does not provide + ``MPI.THREAD_MULTIPLE``, or creates an unexpected communicator. + ValueError: If the mapping's EP group is inconsistent. + """ + if MPI is None: + raise RuntimeError( + "mpi4py is required to create the WideEP FT communicator") + + parent_comm = mpi_comm() if parent_comm is None else parent_comm + parent_rank = parent_comm.Get_rank() + parent_size = parent_comm.Get_size() + local_topology = _mpi_ft_local_topology(mapping, parent_rank, parent_size, + MPI.Query_thread(), health_size) + topologies = parent_comm.allgather(local_topology) + _validate_mpi_ft_topologies(topologies, parent_size) + + ep_group = local_topology["ep_group"] + ep_rank = local_topology["ep_rank"] + ep_size = local_topology["ep_size"] + + # The MVP topology gate above makes the world-spanning group's first rank + # a stable color shared by every process. + ft_comm = parent_comm.Split(color=ep_group[0], key=ep_rank) + try: + setup_status = _mpi_ft_post_split_status(ft_comm, parent_rank, ep_rank, + ep_size) + setup_statuses = parent_comm.allgather(setup_status) + setup_error, ulfm_available = _mpi_ft_post_split_result( + setup_statuses, parent_size) + if setup_error is None: + setup = MpiFtSubcommSetup( + comm=ft_comm, + local_rank=ep_rank, + ep_size=ep_size, + ulfm_available=ulfm_available, + ) + # Enforce process-lifetime ownership at the creation boundary. A + # direct or future caller must not be able to drop the last Python + # reference and trigger rank-local collective MPI_Comm_free. + _retain_mpi_ft_comm(ft_comm) + return setup + except Exception: + # Once Split succeeds, never let an unexpected post-Split exception + # drop the last reference and trigger rank-local communicator cleanup. + _retain_mpi_ft_comm(ft_comm) + raise + + # MPI_Comm_free is collective. Once setup has reported any communicator + # invariant failure, attempting collective cleanup on that same handle is + # not provably bounded, even with MPI_ERRORS_RETURN installed. Keep the + # handle reachable and let MPI reclaim it at process teardown. + _retain_mpi_ft_comm(ft_comm) + logger.warning("WideEP FT retained a communicator after startup failure: " + f"{setup_error}") + raise RuntimeError(setup_error) + + class Distributed(ABC): def __init__(self, mapping: Mapping): diff --git a/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py b/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py new file mode 100644 index 000000000000..fd3cd1018c51 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py @@ -0,0 +1,1564 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Out-of-band MPI failure propagation for WideEP fault tolerance. + +The model-forward thread can be stuck in an AlltoAll kernel when a peer fails, +so it cannot participate in recovery consensus. :class:`MpiFtSubcomm` owns a +dedicated EP-scoped MPI communicator and progresses only nonblocking +point-to-point requests on an independent CPU thread. + +The component implements authoritative single-failure propagation. Multi-rank +suspect/confirm consensus and communicator reconstruction are intentionally +left to the later WideEP FT phases. + +Every survivor echoes a newly observed failure. Receiving one echo from every +active survivor proves that all of them updated their local health. This remains +an asynchronous protocol, not an iteration barrier: the model-engine hook polls +:meth:`health_is_reconciled` at safe iteration boundaries, and the later +barrier-piggyback phase integrates the same state directly into the iteration +barrier. ULFM revoke is reserved for terminal control-plane aborts because MPI +does not carry enough information to distinguish a successful commit revoke +from an emergency abort at remote ranks. Without ULFM, terminal state is echoed +to every survivor; failure to confirm that echo within a bounded interval uses +``MPI_Abort`` on the world-spanning FT communicator as a fail-stop fallback. +""" + +from __future__ import annotations + +import math +import queue +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass +from enum import Enum, IntEnum, auto +from typing import TYPE_CHECKING, Protocol, cast + +import numpy as np + +try: + from mpi4py import MPI +except ImportError: + MPI = None + +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import ( + EPGroupHealth, + EPGroupHealthSnapshot, +) +from tensorrt_llm.logger import logger + +if TYPE_CHECKING: + from tensorrt_llm.mapping import Mapping + + +_FAILURE_TAG = 0 +_MESSAGE_WORDS = 2 + + +class _MessageKind(IntEnum): + FAILURE = 1 + ABORT = 2 + + +_PROCESS_LIFETIME_REFS: list[object] = [] +_PROCESS_LIFETIME_REFS_LOCK = threading.Lock() + + +def _retain_for_process_lifetime(*references: object) -> None: + """Keep active MPI objects reachable until interpreter teardown.""" + with _PROCESS_LIFETIME_REFS_LOCK: + _PROCESS_LIFETIME_REFS.extend(references) + + +class _MpiRequest(Protocol): + """Subset of mpi4py request operations used by the progress loop.""" + + def Cancel(self) -> None: ... + + def Test(self) -> bool: ... + + +class _MpiComm(Protocol): + """Subset of mpi4py communicator operations used by this component.""" + + def Get_rank(self) -> int: ... + + def Get_size(self) -> int: ... + + def Irecv(self, buffer: np.ndarray, source: int, tag: int) -> _MpiRequest: ... + + def Isend(self, buffer: np.ndarray, dest: int, tag: int) -> _MpiRequest: ... + + def Is_revoked(self) -> bool: ... + + def Revoke(self) -> None: ... + + def Abort(self, errorcode: int) -> None: ... + + def Set_errhandler(self, errhandler: object) -> None: ... + + +class _MpiModule(Protocol): + """Subset of the mpi4py ``MPI`` module used by this component.""" + + ERRORS_RETURN: object + ERR_PROC_FAILED: int + ERR_UNKNOWN: int + THREAD_MULTIPLE: int + Exception: type[BaseException] + + def Query_thread(self) -> int: ... + + +FailureReceivedCallback = Callable[[int, int, float], None] + + +@dataclass(frozen=True) +class MpiFtSubcommConfig: + """Runtime-independent tuning for the FT progress thread. + + Args: + poll_interval_sec: Maximum delay between MPI progress passes. The + default leaves margin below the 100 ms agreement budget. + startup_timeout_sec: Maximum time to wait for receive requests to be + posted when :meth:`MpiFtSubcomm.start` is called. + stop_timeout_sec: Default bounded join timeout used by + :meth:`MpiFtSubcomm.stop`. + reconcile_timeout_sec: Maximum time for every active survivor to + announce the same failure before the control plane fails closed. + unattributed_error_timeout_sec: Maximum time to wait for the host + watchdog to identify a rank after non-ULFM MPI reports a generic + transport error. The default leaves margin for the design's + five-second watchdog bound while keeping terminal escalation + inside the ten-second recovery budget. + abort_timeout_sec: Maximum time to reconcile a relayed terminal ABORT + before falling back to communicator-wide ``MPI_Abort`` when ULFM + is unavailable. + """ + + poll_interval_sec: float = 0.01 + startup_timeout_sec: float = 2.0 + stop_timeout_sec: float = 2.0 + reconcile_timeout_sec: float = 1.0 + unattributed_error_timeout_sec: float = 6.0 + abort_timeout_sec: float = 0.1 + + def __post_init__(self) -> None: + for name, value in ( + ("poll_interval_sec", self.poll_interval_sec), + ("startup_timeout_sec", self.startup_timeout_sec), + ("stop_timeout_sec", self.stop_timeout_sec), + ("reconcile_timeout_sec", self.reconcile_timeout_sec), + ("unattributed_error_timeout_sec", self.unattributed_error_timeout_sec), + ("abort_timeout_sec", self.abort_timeout_sec), + ): + if not math.isfinite(value) or value <= 0: + raise ValueError(f"{name} must be finite and > 0, got {value}") + + +@dataclass +class _PendingSend: + request: _MpiRequest + buffer: np.ndarray + destination: int + failed_rank: int + kind: _MessageKind + + +@dataclass +class _PendingReceive: + request: _MpiRequest + buffer: np.ndarray + source: int + + +@dataclass(frozen=True) +class _OutboundMessage: + kind: _MessageKind + failed_rank: int + + +class _Lifecycle(Enum): + CREATED = auto() + STARTING = auto() + RUNNING = auto() + STOPPING = auto() + STOPPED = auto() + FAILED = auto() + + +class MpiFtSubcomm: + """Broadcast EP-rank failures on a dedicated MPI control plane. + + Construction creates the FT communicator collectively on the caller's + thread. Call :meth:`start` only after every rank has constructed the + component. During normal operation, the progress thread owns MPI calls on + this communicator; :meth:`broadcast_failure` only updates local health and + enqueues a report, so it never waits for a peer or MPI progress. Because + ``MPI.THREAD_MULTIPLE`` is required, a caller whose start/stop deadline + proves that thread is wedged may issue the terminal Revoke/Abort fallback. + + Args: + mapping: Distributed mapping whose ``moe_ep_group`` defines the FT + communicator and whose ``moe_ep_rank`` defines the local rank. + health: Process-local EP health state updated by sent and received + failure reports. + config: Progress-loop timing configuration. + on_failure_received: Optional callback invoked as + ``(failed_rank, source_rank, monotonic_time)`` only when a received + report changes local health. Delivery runs on a dedicated daemon + thread so callback latency cannot block MPI progress. + comm: Pre-created FT communicator. Intended for tests; production + callers should let the component create it from ``mapping``. + mpi_module: MPI module paired with ``comm``. Intended for tests. + + Raises: + RuntimeError: If MPI support or thread support is insufficient, or if + the communicator does not match ``mapping``. + ValueError: If ``health`` and ``mapping`` describe different EP groups. + """ + + def __init__( + self, + mapping: Mapping, + health: EPGroupHealth, + config: MpiFtSubcommConfig | None = None, + on_failure_received: FailureReceivedCallback | None = None, + *, + comm: _MpiComm | None = None, + mpi_module: _MpiModule | None = None, + ) -> None: + if mpi_module is None: + if MPI is None: + raise RuntimeError("mpi4py is required for WideEP failure propagation") + mpi_module = cast(_MpiModule, MPI) + self._mpi = mpi_module + + self._config = config or MpiFtSubcommConfig() + comm_created_collectively = comm is None + collective_setup = None + if comm_created_collectively: + # Keep the collective startup validation and Split on the + # constructing thread. In particular, do not raise a rank-local + # mapping error before peers have entered the same collective + # validation path. + from tensorrt_llm._torch.distributed.communicator import create_mpi_ft_subcomm + + collective_setup = create_mpi_ft_subcomm( + mapping, + health_size=health.moe_world_size, + ) + comm = cast(_MpiComm, collective_setup.comm) + + assert comm is not None + self._comm = comm + if comm_created_collectively: + # create_mpi_ft_subcomm already reconciled thread support, mapping, + # health size, communicator rank/size, and ERRORS_RETURN across the + # parent communicator. Repeating any of those MPI operations here + # would reintroduce a rank-local failure after collective startup. + assert collective_setup is not None + self._local_rank = collective_setup.local_rank + self._ep_size = collective_setup.ep_size + else: + if self._mpi.Query_thread() < self._mpi.THREAD_MULTIPLE: + raise RuntimeError( + "WideEP FT requires MPI.THREAD_MULTIPLE because its control-plane " + "thread overlaps other MPI traffic" + ) + + ep_group = tuple(mapping.moe_ep_group) + if len(ep_group) != mapping.moe_ep_size: + raise ValueError( + "mapping.moe_ep_group size must match mapping.moe_ep_size, " + f"got {len(ep_group)} and {mapping.moe_ep_size}" + ) + expected_world = set(range(mapping.world_size)) + if mapping.moe_ep_size != mapping.world_size or set(ep_group) != expected_world: + raise ValueError( + "WideEP FT MVP requires one MoE EP group spanning the full MPI world; " + f"got world_size={mapping.world_size}, moe_ep_group={ep_group}" + ) + if health.moe_world_size != mapping.moe_ep_size: + raise ValueError( + "EPGroupHealth size must match mapping.moe_ep_size, " + f"got {health.moe_world_size} and {mapping.moe_ep_size}" + ) + self._local_rank = mapping.moe_ep_rank + self._ep_size = mapping.moe_ep_size + # Install the non-fatal handler before inspecting the injected + # communicator so even validation errors are returned to Python. + self._comm.Set_errhandler(self._mpi.ERRORS_RETURN) + actual_rank = self._comm.Get_rank() + actual_size = self._comm.Get_size() + if actual_rank != self._local_rank or actual_size != self._ep_size: + raise RuntimeError( + "WideEP FT communicator does not match the mapping: " + f"rank={actual_rank}, size={actual_size}, " + f"expected rank={self._local_rank}, size={self._ep_size}" + ) + + self._health = health + self._on_failure_received = on_failure_received + self._outbound_reports: queue.SimpleQueue[_OutboundMessage] = queue.SimpleQueue() + self._announced_failures: set[int] = set() + self._failure_reporters: dict[int, set[int]] = {} + self._reconcile_deadlines: dict[int, float] = {} + self._unattributed_error_deadlines: dict[int, float] = {} + self._failure_fanout_posted: set[int] = set() + self._failure_fanout_complete: set[int] = set() + self._abort_reporters: set[int] = set() + self._abort_expected_reporters: frozenset[int] = frozenset() + self._abort_deadline: float | None = None + self._abort_payload = -1 + self._abort_announced = False + self._abort_fanout_posted = False + self._abort_fanout_complete = False + self._protocol_lock = threading.Lock() + self._failure_claim_lock = threading.Lock() + self._accepted_failed_rank: int | None = None + self._accepted_failure_generation: int | None = None + self._stop_event = threading.Event() + self._abort_requested = threading.Event() + self._wake_event = threading.Event() + self._ready_event = threading.Event() + self._deadline_monitor_stop_event = threading.Event() + self._deadline_monitor_wake_event = threading.Event() + self._deadline_monitor_ready_event = threading.Event() + self._transport_poisoned = threading.Event() + self._progress_failed = threading.Event() + self._lifecycle = _Lifecycle.CREATED + self._lifecycle_lock = threading.Lock() + self._thread: threading.Thread | None = None + self._deadline_monitor_thread: threading.Thread | None = None + self._last_error: BaseException | None = None + self._error_lock = threading.Lock() + self._retained_requests: list[_PendingSend | _PendingReceive] = [] + self._process_lifetime_retained = comm_created_collectively + if comm_created_collectively: + # The FT communicator is a process-lifetime control-plane + # resource. Letting mpi4py destroy it when one broadcaster is + # garbage-collected would call a collective Free at rank-local + # and nondeterministic times. + _retain_for_process_lifetime(self._comm) + self._shutdown_deadline: float | None = None + self._callback_queue: queue.SimpleQueue[tuple[int, int, float] | None] = queue.SimpleQueue() + self._callback_stop_event = threading.Event() + self._callback_thread: threading.Thread | None = None + self._terminal_action_lock = threading.Lock() + self._revoke_attempted = False + self._world_abort_attempted = False + self._ulfm_available = ( + collective_setup.ulfm_available if collective_setup is not None else self._probe_ulfm() + ) + + @property + def ulfm_available(self) -> bool: + """Whether this MPI build implements the ULFM communicator methods.""" + with self._terminal_action_lock: + return self._ulfm_available + + @property + def last_error(self) -> BaseException | None: + """First terminal control-plane error, if any.""" + with self._error_lock: + return self._last_error + + def start(self) -> None: + """Start the dedicated MPI progress thread. + + This method is not restartable. A second call, including one after + :meth:`stop`, raises ``RuntimeError``. + """ + thread_start_error: BaseException | None = None + with self._lifecycle_lock: + if self._lifecycle is not _Lifecycle.CREATED: + raise RuntimeError( + f"WideEP FT broadcaster cannot start from {self._lifecycle.name.lower()} state" + ) + self._lifecycle = _Lifecycle.STARTING + try: + if self._on_failure_received is not None: + callback_thread = threading.Thread( + target=self._callback_loop, + name="wide-ep-ft-callback", + daemon=True, + ) + callback_thread.start() + # Store only a successfully started thread. stop() may join + # every stored thread and joining an unstarted Thread raises. + self._callback_thread = callback_thread + deadline_monitor_thread = threading.Thread( + target=self._deadline_monitor_loop, + name="wide-ep-ft-deadline-monitor", + daemon=True, + ) + deadline_monitor_thread.start() + self._deadline_monitor_thread = deadline_monitor_thread + progress_thread = threading.Thread( + target=self._progress_loop, + name="wide-ep-ft-broadcast", + daemon=True, + ) + progress_thread.start() + self._thread = progress_thread + except Exception as error: + # Publish the error before FAILED so concurrent observers never + # see a terminal lifecycle without its cause. + self._record_error(error) + self._progress_failed.set() + self._lifecycle = _Lifecycle.FAILED + thread_start_error = error + + if thread_start_error is not None: + self._request_deadline_monitor_stop() + self._request_callback_stop() + self._retain_requests([]) + # There is no progress thread to relay ABORT, so fail closed on the + # caller thread. This is the only thread that can own the FT comm + # after progress-thread creation itself failed. + self._fail_closed_immediately(thread_start_error) + callback_thread = self._callback_thread + if callback_thread is not None and callback_thread.is_alive(): + callback_thread.join(self._config.stop_timeout_sec) + deadline_monitor_thread = self._deadline_monitor_thread + if deadline_monitor_thread is not None and deadline_monitor_thread.is_alive(): + deadline_monitor_thread.join(self._config.stop_timeout_sec) + raise RuntimeError( + "WideEP FT control-plane thread failed to start" + ) from thread_start_error + + startup_deadline = time.monotonic() + self._config.startup_timeout_sec + progress_ready = self._ready_event.wait(max(0.0, startup_deadline - time.monotonic())) + deadline_monitor_ready = self._deadline_monitor_ready_event.wait( + max(0.0, startup_deadline - time.monotonic()) + ) + if not progress_ready or not deadline_monitor_ready: + timeout_error = TimeoutError( + "WideEP FT control-plane threads did not start before the timeout" + ) + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + # Stop telemetry before the terminal MPI action. Real MPI_Abort + # does not return, while test doubles and Revoke do; either way no + # callback worker should be left waiting for progress-thread cleanup. + self._request_deadline_monitor_stop() + self._request_callback_stop() + # The progress thread may be wedged inside the MPI call whose + # startup deadline just expired, so it cannot execute the terminal + # action itself. MPI.THREAD_MULTIPLE is a construction invariant; + # use the caller thread to issue the terminal communicator action. + self._fail_closed_immediately(timeout_error) + # The progress thread may be stuck inside the first MPI Irecv and + # therefore unable to reach its ``finally`` block. Do not leave the + # independent telemetry worker alive after start() has failed. + callback_thread = self._callback_thread + if callback_thread is not None and callback_thread is not threading.current_thread(): + callback_thread.join(self._config.stop_timeout_sec) + deadline_monitor_thread = self._deadline_monitor_thread + if ( + deadline_monitor_thread is not None + and deadline_monitor_thread is not threading.current_thread() + ): + deadline_monitor_thread.join(self._config.stop_timeout_sec) + raise timeout_error + with self._lifecycle_lock: + startup_error = self.last_error + lifecycle = self._lifecycle + if startup_error is not None: + self._lifecycle = _Lifecycle.FAILED + elif lifecycle is not _Lifecycle.RUNNING: + raise RuntimeError( + "WideEP FT broadcaster did not reach running state during startup " + f"(state={lifecycle.name.lower()})" + ) + if startup_error is not None: + self._request_deadline_monitor_stop() + deadline_monitor_thread = self._deadline_monitor_thread + if ( + deadline_monitor_thread is not None + and deadline_monitor_thread is not threading.current_thread() + ): + deadline_monitor_thread.join(self._config.stop_timeout_sec) + raise RuntimeError("WideEP FT progress thread failed during startup") from startup_error + + def stop(self, timeout: float | None = None) -> None: + """Stop progress without blocking on MPI requests or freeing the comm. + + Healthy requests are cancelled and polled within the same bounded + timeout. After a peer or transport failure, requests and their NumPy + buffers are retained until process teardown because cancelling, + waiting for, or freeing them can hang inside MPI. + + If a known failure is not yet reconciled, the progress thread keeps + running until agreement or the existing reconciliation deadline drives + the terminal fail-stop path. + + Args: + timeout: Maximum join time. Defaults to + :attr:`MpiFtSubcommConfig.stop_timeout_sec`. + + Raises: + ValueError: If ``timeout`` is not finite and positive. + TimeoutError: If the progress thread does not exit in time. + """ + timeout = self._config.stop_timeout_sec if timeout is None else timeout + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError(f"timeout must be finite and > 0, got {timeout}") + + stop_deadline = time.monotonic() + timeout + with self._lifecycle_lock: + if self._lifecycle is _Lifecycle.CREATED: + self._lifecycle = _Lifecycle.STOPPED + self._request_deadline_monitor_stop() + self._request_callback_stop() + return + callback_thread = self._callback_thread + deadline_monitor_thread = self._deadline_monitor_thread + if ( + self._lifecycle is _Lifecycle.STOPPED + and (callback_thread is None or not callback_thread.is_alive()) + and (deadline_monitor_thread is None or not deadline_monitor_thread.is_alive()) + ): + return + thread = self._thread + if self._lifecycle in (_Lifecycle.STARTING, _Lifecycle.RUNNING): + self._lifecycle = _Lifecycle.STOPPING + join_grace = min(0.01, timeout / 2) + self._shutdown_deadline = stop_deadline - join_grace + self._stop_event.set() + self._wake_event.set() + self._request_callback_stop() + + if thread is not None: + thread.join(max(0.0, stop_deadline - time.monotonic())) + if thread.is_alive(): + timeout_error = TimeoutError( + "WideEP FT progress thread did not stop before the timeout" + ) + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + # Do not enqueue fail-stop work onto the thread that this + # timeout proves is not making MPI progress. + self._request_deadline_monitor_stop() + self._fail_closed_immediately(timeout_error) + raise timeout_error + + self._request_deadline_monitor_stop() + if ( + deadline_monitor_thread is not None + and deadline_monitor_thread is not threading.current_thread() + ): + deadline_monitor_thread.join(max(0.0, stop_deadline - time.monotonic())) + if deadline_monitor_thread.is_alive(): + timeout_error = TimeoutError( + "WideEP FT deadline monitor did not stop before the timeout" + ) + self._fail_closed_immediately(timeout_error) + raise timeout_error + + if callback_thread is not None and callback_thread is not threading.current_thread(): + callback_thread.join(max(0.0, stop_deadline - time.monotonic())) + if callback_thread.is_alive(): + timeout_error = TimeoutError( + "WideEP FT callback thread did not stop before the timeout" + ) + # Callbacks are telemetry only. A stuck callback may make this + # stop call time out, but it must not poison MPI state or make + # one rank take the poisoned-world process-exit path. + raise timeout_error + + with self._lifecycle_lock: + if self._lifecycle is not _Lifecycle.FAILED: + self._lifecycle = _Lifecycle.STOPPED + + def pre_failover(self, failed_rank: int) -> bool: + """Mark and asynchronously announce one failed EP rank. + + This is the nonblocking seam used by failure detectors. The report is + queued even when ``health`` was already updated by the caller, because + the host watchdog normally marks local health before invoking the + broadcaster. + + Args: + failed_rank: EP-local rank in ``[0, ep_size)``. + + Returns: + ``True`` if this call changed local health, else ``False``. + """ + return self.broadcast_failure(failed_rank) + + def broadcast_failure(self, failed_rank: int) -> bool: + """Mark and asynchronously broadcast one failed EP rank. + + This method performs no MPI calls and never waits for the progress + thread. Repeated calls for the same rank are coalesced after the first + local announcement. + """ + self._validate_rank(failed_rank) + with self._lifecycle_lock: + if self._lifecycle is not _Lifecycle.RUNNING or self._progress_failed.is_set(): + raise RuntimeError( + f"WideEP FT broadcaster is not running (state={self._lifecycle.name.lower()})" + ) + start_time = time.monotonic() + try: + self._claim_failure(failed_rank) + except RuntimeError as error: + self._lifecycle = _Lifecycle.FAILED + self._request_terminal_abort(error, failed_rank, self._local_rank) + raise + changed = self._health.mark_failed(failed_rank) + enqueued = self._observe_failure(failed_rank, self._local_rank) + logger.warning( + f"WideEP FT local failure rank={failed_rank} changed={changed} " + f"enqueued={enqueued} elapsed_ms={(time.monotonic() - start_time) * 1000.0:.3f}" + ) + return changed + + def world_is_poisoned(self) -> bool: + """Return whether this communicator epoch has ever observed failure.""" + return ( + self._transport_poisoned.is_set() + or self._accepted_failure_rank() is not None + or not self._health.all_active() + ) + + def failure_is_reconciled(self, failed_rank: int) -> bool: + """Return whether every active survivor has announced ``failed_rank``. + + This method performs no MPI calls and is safe to poll from an iteration + boundary. A terminal progress error always fails closed. + """ + self._validate_rank(failed_rank) + with self._lifecycle_lock: + if self._lifecycle is not _Lifecycle.RUNNING: + return False + if self.last_error is not None or self._progress_failed.is_set(): + return False + snapshot = self._health.snapshot() + if failed_rank not in snapshot.failed_ranks: + return False + accepted_failure, accepted_generation = self._accepted_failure_state() + with self._protocol_lock: + if not self._reconciliation_state_is_valid_locked( + snapshot, accepted_failure, accepted_generation + ): + return False + return self._failure_is_reconciled_locked(failed_rank, snapshot.mask) + + def health_is_reconciled(self) -> bool: + """Return whether all locally known failures are agreed by survivors.""" + with self._lifecycle_lock: + if self._lifecycle is not _Lifecycle.RUNNING: + return False + if self.last_error is not None or self._progress_failed.is_set(): + return False + snapshot = self._health.snapshot() + accepted_failure, accepted_generation = self._accepted_failure_state() + if not snapshot.failed_ranks: + return accepted_failure is None and not self._transport_poisoned.is_set() + with self._protocol_lock: + if not self._reconciliation_state_is_valid_locked( + snapshot, accepted_failure, accepted_generation + ): + return False + return all( + self._failure_is_reconciled_locked(failed_rank, snapshot.mask) + for failed_rank in snapshot.failed_ranks + ) + + def _validate_rank(self, rank: int) -> None: + if not 0 <= rank < self._ep_size: + raise ValueError(f"failed_rank must be in [0, {self._ep_size}), got {rank}") + + def _claim_failure(self, failed_rank: int) -> None: + """Atomically enforce the single-distinct-failure MVP contract.""" + with self._failure_claim_lock: + snapshot = self._health.snapshot() + known_failures = snapshot.failed_ranks + conflicting_health = known_failures - {failed_rank} + accepted = self._accepted_failed_rank + if (accepted is not None and accepted != failed_rank) or conflicting_health: + first_failure = accepted + if first_failure is None: + first_failure = min(conflicting_health) + raise RuntimeError( + "WideEP FT MVP supports exactly one distinct failed rank; " + f"already accepted rank {first_failure}, received rank {failed_rank}" + ) + if accepted is None: + self._accepted_failed_rank = failed_rank + # The watchdog may have marked the rank before calling us. If + # not, exactly one effective mark_failed() must be the next + # health transition. Capturing that expected generation here + # prevents an independent mutation in the claim-to-observe + # window from becoming the accepted epoch baseline. + self._accepted_failure_generation = snapshot.generation + ( + 0 if failed_rank in known_failures else 1 + ) + + def _accepted_failure_rank(self) -> int | None: + with self._failure_claim_lock: + return self._accepted_failed_rank + + def _accepted_failure_state(self) -> tuple[int | None, int | None]: + with self._failure_claim_lock: + return self._accepted_failed_rank, self._accepted_failure_generation + + def _observe_failure(self, failed_rank: int, reporter: int) -> bool: + """Record one authoritative report and enqueue this rank's echo once.""" + now = time.monotonic() + snapshot = self._health.snapshot() + message: _OutboundMessage | None = None + with self._protocol_lock: + unattributed_deadline = self._unattributed_error_deadlines.get(failed_rank) + if unattributed_deadline != math.inf: + self._unattributed_error_deadlines.pop(failed_rank, None) + reporters = self._failure_reporters.setdefault(failed_rank, set()) + reporters.update((reporter, self._local_rank)) + self._reconcile_deadlines.setdefault( + failed_rank, now + self._config.reconcile_timeout_sec + ) + failure_enqueued = failed_rank not in self._announced_failures + if failure_enqueued: + self._announced_failures.add(failed_rank) + message = _OutboundMessage(_MessageKind.FAILURE, failed_rank) + if self._failure_is_reconciled_locked(failed_rank, snapshot.mask): + self._reconcile_deadlines.pop(failed_rank, None) + if message is not None: + self._outbound_reports.put(message) + self._wake_event.set() + self._deadline_monitor_wake_event.set() + return failure_enqueued + + def _reconciliation_state_is_valid_locked( + self, + snapshot: EPGroupHealthSnapshot, + accepted_failure: int | None, + accepted_generation: int | None, + ) -> bool: + """Reject unresolved transport errors and same-epoch rank restoration.""" + if self._unattributed_error_deadlines: + return False + if accepted_failure is None: + return not snapshot.failed_ranks and accepted_generation is None + return ( + snapshot.failed_ranks == frozenset({accepted_failure}) + and accepted_generation == snapshot.generation + ) + + def _request_terminal_abort( + self, + error: BaseException, + payload: int, + reporter: int, + ) -> None: + """Relay terminal state before the progress thread fails closed. + + The relay is best effort because it uses the same FT communicator. If + survivor echoes do not converge before the bounded deadline, the + progress thread escalates to ULFM revoke or communicator-wide + ``MPI_Abort``. This method performs no MPI calls. + """ + self._record_error(error) + snapshot = self._health.snapshot() + enqueue_abort = False + now = time.monotonic() + with self._protocol_lock: + if self._abort_deadline is None: + self._abort_payload = payload + self._abort_expected_reporters = frozenset( + rank for rank in range(self._ep_size) if snapshot.mask & (1 << rank) + ) + self._abort_deadline = now + self._config.abort_timeout_sec + self._abort_reporters.update((reporter, self._local_rank)) + if not self._abort_announced: + self._abort_announced = True + enqueue_abort = True + self._abort_requested.set() + if enqueue_abort: + self._outbound_reports.put(_OutboundMessage(_MessageKind.ABORT, self._abort_payload)) + self._wake_event.set() + self._deadline_monitor_wake_event.set() + + def _failure_is_reconciled_locked(self, failed_rank: int, active_mask: int) -> bool: + required_reporters = {rank for rank in range(self._ep_size) if active_mask & (1 << rank)} + if not required_reporters: + # There is nobody left who can safely execute the next collective. + # Do not let the usual ``empty set is a subset`` rule open the + # iteration gate in this terminal boundary case. + return False + reporters = self._failure_reporters.get(failed_rank, set()) + return failed_rank in self._failure_fanout_complete and required_reporters.issubset( + reporters + ) + + def _probe_ulfm(self) -> bool: + if not callable(getattr(self._comm, "Is_revoked", None)) or not callable( + getattr(self._comm, "Revoke", None) + ): + return False + try: + already_revoked = self._comm.Is_revoked() + except Exception as error: + if isinstance(error, NotImplementedError) or self._is_unsupported_ulfm_error(error): + logger.debug(f"WideEP FT ULFM is unavailable: {error}") + return False + self._record_error(error) + self._retain_requests([]) + raise + if already_revoked: + error = RuntimeError("WideEP FT communicator is already revoked at construction") + self._record_error(error) + self._retain_requests([]) + raise error + return True + + def _is_unsupported_ulfm_error(self, error: BaseException) -> bool: + if not isinstance(error, self._mpi.Exception): + return False + get_error_class = getattr(error, "Get_error_class", None) + if not callable(get_error_class): + return False + unsupported_classes = { + value + for value in ( + getattr(self._mpi, "ERR_NOT_SUPPORTED", None), + getattr(self._mpi, "ERR_UNSUPPORTED_OPERATION", None), + ) + if value is not None + } + try: + error_class = get_error_class() + except Exception: + # Error inspection is advisory. Preserve the original probe error + # as the construction failure when the runtime cannot classify it. + return False + return error_class in unsupported_classes + + def _progress_loop(self) -> None: + pending_sends: list[_PendingSend] = [] + pending_receives: dict[int, _PendingReceive] = {} + try: + for peer in range(self._ep_size): + if peer != self._local_rank: + pending_receives[peer] = self._post_receive(peer) + with self._lifecycle_lock: + if self._lifecycle is _Lifecycle.STARTING: + self._lifecycle = _Lifecycle.RUNNING + self._ready_event.set() + + while True: + if self._progress_failed.is_set(): + terminal_error = self.last_error + if terminal_error is None: + terminal_error = RuntimeError("WideEP FT progress was asked to fail closed") + raise terminal_error + self._wake_event.clear() + self._drain_outbound_reports(pending_sends) + self._progress_sends(pending_sends) + self._progress_receives(pending_receives) + self._check_accepted_failure_epoch() + self._check_unattributed_transport_deadlines() + self._check_reconciliation_deadlines() + with self._lifecycle_lock: + terminal_pending = ( + self._abort_requested.is_set() or self._lifecycle is _Lifecycle.FAILED + ) + stop_requested = self._stop_event.is_set() and not terminal_pending + if terminal_pending: + if self._progress_terminal_abort(): + break + self._wake_event.wait(self._config.poll_interval_sec) + continue + if stop_requested: + safe_to_stop, stop_error = self._stop_reconciliation_state() + if stop_error is not None: + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + self._request_terminal_abort( + stop_error, + -1, + self._local_rank, + ) + elif safe_to_stop: + break + self._wake_event.wait(self._config.poll_interval_sec) + except Exception as error: + self._fail_closed_immediately(error) + logger.error(f"WideEP FT progress thread failed: {error}") + finally: + self._ready_event.set() + self._request_deadline_monitor_stop() + self._request_callback_stop() + try: + self._cleanup_requests(pending_sends, pending_receives) + except Exception as error: + self._retain_requests([*pending_sends, *pending_receives.values()]) + logger.error(f"WideEP FT request cleanup failed: {error}") + self._fail_closed_after_cleanup(error) + with self._lifecycle_lock: + if self._last_error is not None: + self._lifecycle = _Lifecycle.FAILED + elif self._stop_event.is_set(): + self._lifecycle = _Lifecycle.STOPPED + + def _post_receive(self, source: int) -> _PendingReceive: + buffer = np.empty(_MESSAGE_WORDS, dtype=np.int64) + request = self._comm.Irecv(buffer, source=source, tag=_FAILURE_TAG) + return _PendingReceive(request=request, buffer=buffer, source=source) + + def _drain_outbound_reports(self, pending_sends: list[_PendingSend]) -> None: + while True: + try: + message = self._outbound_reports.get_nowait() + except queue.Empty: + return + + failed_rank = message.failed_rank + snapshot = self._health.snapshot() + if message.kind is _MessageKind.ABORT: + # Terminal state must reach every rank still believed active. + # If one of them is actually the second failed rank, send + # failure falls through to the fail-stop MPI_Abort fallback. + peers = [ + peer + for peer in range(self._ep_size) + if peer != self._local_rank and snapshot.mask & (1 << peer) + ] + else: + peers = [ + peer + for peer in range(self._ep_size) + if peer != self._local_rank + and peer != failed_rank + and snapshot.mask & (1 << peer) + ] + if not peers: + with self._protocol_lock: + if message.kind is _MessageKind.ABORT: + self._abort_fanout_posted = True + self._abort_fanout_complete = True + else: + self._failure_fanout_posted.add(failed_rank) + self._failure_fanout_complete.add(failed_rank) + continue + for peer in peers: + buffer = np.asarray([int(message.kind), failed_rank], dtype=np.int64) + request = self._comm.Isend(buffer, dest=peer, tag=_FAILURE_TAG) + pending_sends.append( + _PendingSend( + request=request, + buffer=buffer, + destination=peer, + failed_rank=failed_rank, + kind=message.kind, + ) + ) + with self._protocol_lock: + if message.kind is _MessageKind.ABORT: + # Posting is not completion: rendezvous sends can still + # require this thread to call Test before peers can receive + # the relay. + self._abort_fanout_posted = True + else: + self._failure_fanout_posted.add(failed_rank) + + def _progress_sends(self, pending_sends: list[_PendingSend]) -> None: + incomplete: list[_PendingSend] = [] + for pending in pending_sends: + completed = pending.request.Test() + if not completed: + incomplete.append(pending) + pending_sends[:] = incomplete + incomplete_failure_ranks = { + pending.failed_rank for pending in incomplete if pending.kind is _MessageKind.FAILURE + } + abort_incomplete = any(pending.kind is _MessageKind.ABORT for pending in incomplete) + with self._protocol_lock: + self._failure_fanout_complete.update( + self._failure_fanout_posted - incomplete_failure_ranks + ) + if self._abort_fanout_posted and not abort_incomplete: + self._abort_fanout_complete = True + + def _progress_receives(self, pending_receives: dict[int, _PendingReceive]) -> None: + for source, pending in list(pending_receives.items()): + try: + completed = pending.request.Test() + except Exception as error: + if not isinstance(error, self._mpi.Exception): + raise + self._transport_poisoned.set() + # A request that raised may still own its receive buffer. Keep + # both alive until process teardown instead of assuming the + # implementation completed it while returning the error. + self._retain_requests([pending]) + if self._is_failed_peer_error(error): + del pending_receives[source] + self._handle_failed_peer(source, source) + continue + if not self.ulfm_available: + # Without ULFM, MPI implementations may surface the dead + # source as a generic error. Do not guess the failed rank; + # drop only this fixed-source receive and keep polling the + # other survivors for an authoritative watchdog report. + del pending_receives[source] + self._track_unattributed_transport_error(source) + logger.warning( + f"WideEP FT receive from EP rank {source} failed without ULFM: {error}" + ) + continue + raise + if not completed: + continue + + del pending_receives[source] + message_kind_value = int(pending.buffer[0]) + failed_rank = int(pending.buffer[1]) + try: + message_kind = _MessageKind(message_kind_value) + except ValueError: + logger.warning( + f"WideEP FT ignored invalid message kind {message_kind_value} " + f"from EP rank {source}" + ) + else: + if message_kind is _MessageKind.ABORT and -1 <= failed_rank < self._ep_size: + self._handle_received_abort(failed_rank, source) + elif message_kind is _MessageKind.FAILURE and 0 <= failed_rank < self._ep_size: + self._handle_received_report(failed_rank, source) + else: + logger.warning( + f"WideEP FT ignored invalid {message_kind.name.lower()} payload " + f"{failed_rank} from EP rank {source}" + ) + + if self._health.is_active(source) and not self._stop_event.is_set(): + try: + pending_receives[source] = self._post_receive(source) + except Exception as error: + if isinstance(error, self._mpi.Exception) and not self.ulfm_available: + self._transport_poisoned.set() + self._track_unattributed_transport_error(source) + logger.warning( + f"WideEP FT could not repost receive from EP rank {source} " + f"without ULFM: {error}" + ) + else: + raise + + def _handle_received_report(self, failed_rank: int, source: int) -> None: + start_time = time.monotonic() + changed = self._accept_received_failure( + failed_rank, + reporter=source, + transport_poisoned=False, + ) + if changed is None: + return + if not changed: + return + received_time = time.monotonic() + logger.warning( + f"WideEP FT received failure rank={failed_rank} source={source} " + f"elapsed_ms={(received_time - start_time) * 1000.0:.3f}" + ) + self._invoke_callback(failed_rank, source, received_time) + + def _handle_failed_peer(self, failed_rank: int, source: int) -> None: + changed = self._accept_received_failure( + failed_rank, + reporter=self._local_rank, + transport_poisoned=True, + ) + if changed is None: + return + if changed: + detected_time = time.monotonic() + logger.warning( + f"WideEP FT observed failed EP rank {failed_rank} through the FT communicator" + ) + self._invoke_callback(failed_rank, source, detected_time) + + def _accept_received_failure( + self, + failed_rank: int, + *, + reporter: int, + transport_poisoned: bool, + ) -> bool | None: + """Atomically claim, mark, and observe one progress-thread failure.""" + claim_error: RuntimeError | None = None + with self._lifecycle_lock: + try: + self._claim_failure(failed_rank) + except RuntimeError as error: + self._lifecycle = _Lifecycle.FAILED + claim_error = error + else: + changed = self._health.mark_failed(failed_rank) + if transport_poisoned: + self._transport_poisoned.set() + self._observe_failure(failed_rank, reporter) + return changed + + assert claim_error is not None + self._request_terminal_abort( + claim_error, + failed_rank, + self._local_rank, + ) + return None + + def _handle_received_abort(self, payload: int, source: int) -> None: + error = RuntimeError( + f"WideEP FT received terminal ABORT from EP rank {source} (payload={payload})" + ) + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + self._request_terminal_abort(error, payload, source) + + def _invoke_callback(self, failed_rank: int, source: int, event_time: float) -> None: + if self._on_failure_received is None or self._callback_stop_event.is_set(): + return + self._callback_queue.put((failed_rank, source, event_time)) + + def _callback_loop(self) -> None: + """Deliver telemetry callbacks without blocking MPI progress.""" + while True: + event = self._callback_queue.get() + if event is None: + return + failed_rank, source, event_time = event + try: + callback = self._on_failure_received + if callback is not None: + callback(failed_rank, source, event_time) + except Exception as error: + # Telemetry must never terminate either control-plane thread. + logger.warning(f"WideEP FT failure callback raised: {error}") + + def _request_callback_stop(self) -> None: + if self._callback_stop_event.is_set(): + return + self._callback_stop_event.set() + self._callback_queue.put(None) + + def _request_deadline_monitor_stop(self) -> None: + self._deadline_monitor_stop_event.set() + self._deadline_monitor_wake_event.set() + + def _deadline_monitor_loop(self) -> None: + """Enforce Python-side deadlines even when an MPI call is wedged. + + Normal communicator traffic remains owned by the progress thread. This + monitor only observes protected protocol state. If a terminal relay + itself reaches its deadline, ``MPI.THREAD_MULTIPLE`` permits this thread + to issue the existing one-shot Revoke/Abort fail-stop action. + """ + self._deadline_monitor_ready_event.set() + try: + while True: + # Clear before checking stop/deadline state so a concurrent + # setter cannot be erased between the check and wait below. + self._deadline_monitor_wake_event.clear() + if self._deadline_monitor_stop_event.is_set(): + return + if self._progress_failed.is_set(): + return + + self._check_accepted_failure_epoch() + self._check_unattributed_transport_deadlines() + self._check_reconciliation_deadlines() + + if self._progress_failed.is_set(): + return + now = time.monotonic() + with self._protocol_lock: + abort_deadline = self._abort_deadline + abort_reconciled = ( + self._abort_fanout_complete + and self._abort_expected_reporters.issubset(self._abort_reporters) + ) + if self._abort_requested.is_set() and abort_deadline is not None: + if abort_reconciled: + # The relay reached every expected survivor. The progress + # thread will break cleanly when it resumes. Do not mark + # progress failed here: its exception path would otherwise + # escalate a reconciled non-ULFM terminal relay to + # MPI_Abort. + self._wake_event.set() + return + if now >= abort_deadline: + self._fail_closed_immediately( + TimeoutError( + "WideEP FT terminal relay did not complete before " + "the abort deadline" + ) + ) + self._wake_event.set() + return + + self._deadline_monitor_wake_event.wait(self._config.poll_interval_sec) + except Exception as error: + self._fail_closed_immediately(error) + self._wake_event.set() + finally: + # start() also uses this event to distinguish a scheduled monitor + # from a thread that failed before its loop became observable. + self._deadline_monitor_ready_event.set() + + def _is_failed_peer_error(self, error: BaseException) -> bool: + get_error_class = getattr(error, "Get_error_class", None) + if not callable(get_error_class): + return False + try: + error_class = get_error_class() + except Exception: + # Preserve the original transport failure as the terminal cause (or + # let the non-ULFM watchdog identify the source) when an MPI runtime + # cannot even classify its own exception. + return False + unknown = getattr(self._mpi, "ERR_UNKNOWN", None) + failure_classes = { + value + for value in (getattr(self._mpi, "ERR_PROC_FAILED", None),) + if value is not None and value != unknown + } + return error_class in failure_classes + + def _track_unattributed_transport_error(self, source: int) -> None: + """Bound how long a generic non-ULFM peer error can remain unresolved.""" + # The watchdog report and the MPI error race independently. If this + # source is already the accepted failed rank, the transport error is + # explained and must not arm a stale deadline after _observe_failure() + # has already cleared the error-first ordering. + deadline = time.monotonic() + self._config.unattributed_error_timeout_sec + with self._protocol_lock: + if self._accepted_failure_rank() == source: + return + self._unattributed_error_deadlines.setdefault(source, deadline) + self._deadline_monitor_wake_event.set() + + def _check_accepted_failure_epoch(self) -> None: + """Fail closed when health leaves the accepted single-failure epoch.""" + error: RuntimeError | None = None + accepted_failure: int | None = None + with self._lifecycle_lock: + if self._lifecycle not in (_Lifecycle.RUNNING, _Lifecycle.STOPPING): + return + snapshot = self._health.snapshot() + accepted_failure, accepted_generation = self._accepted_failure_state() + # A watchdog may mark health immediately before calling + # pre_failover(). Until that call claims a rank, do not mistake the + # documented pre-mark seam for an epoch violation. + if accepted_failure is None: + return + if ( + snapshot.failed_ranks == frozenset({accepted_failure}) + and snapshot.generation == accepted_generation + ): + return + error = RuntimeError( + "WideEP FT accepted failure epoch changed before communicator " + "reconstruction " + f"(accepted_rank={accepted_failure}, " + f"expected_generation={accepted_generation}, snapshot={snapshot})" + ) + self._lifecycle = _Lifecycle.FAILED + + assert error is not None + assert accepted_failure is not None + self._request_terminal_abort(error, accepted_failure, self._local_rank) + + def _check_unattributed_transport_deadlines(self) -> None: + now = time.monotonic() + # Match the public reconciliation gate's lifecycle -> protocol lock + # order. Publish FAILED while both locks exclude a concurrent gate or + # local broadcast, then latch the expired poison before either can run. + with self._lifecycle_lock: + if self._lifecycle not in (_Lifecycle.RUNNING, _Lifecycle.STOPPING): + return + with self._protocol_lock: + expired_sources = [ + source + for source, deadline in self._unattributed_error_deadlines.items() + if now >= deadline + ] + if expired_sources: + self._lifecycle = _Lifecycle.FAILED + # Retaining these few per-source entries for this + # communicator epoch is harmless and permanently keeps the + # protocol gate closed. The infinity latch also prevents + # repeated terminal requests on later progress passes. + for source in expired_sources: + self._unattributed_error_deadlines[source] = math.inf + if not expired_sources: + return + error = TimeoutError( + "WideEP FT could not attribute non-ULFM transport errors before " + f"the watchdog deadline: {expired_sources}" + ) + self._request_terminal_abort(error, -1, self._local_rank) + + def _check_reconciliation_deadlines(self) -> None: + now = time.monotonic() + expired: list[int] = [] + with self._lifecycle_lock: + if self._lifecycle not in (_Lifecycle.RUNNING, _Lifecycle.STOPPING): + return + snapshot = self._health.snapshot() + with self._protocol_lock: + for failed_rank, deadline in list(self._reconcile_deadlines.items()): + if self._failure_is_reconciled_locked(failed_rank, snapshot.mask): + del self._reconcile_deadlines[failed_rank] + elif now >= deadline: + expired.append(failed_rank) + if expired: + self._lifecycle = _Lifecycle.FAILED + if expired: + error = TimeoutError( + f"WideEP FT did not reconcile failed EP ranks before the timeout: {expired}" + ) + self._request_terminal_abort(error, expired[0], self._local_rank) + + def _stop_reconciliation_state(self) -> tuple[bool, BaseException | None]: + """Return whether shutdown can stop MPI progress without splitting ranks.""" + snapshot = self._health.snapshot() + failed_ranks = snapshot.failed_ranks + accepted_failure, accepted_generation = self._accepted_failure_state() + with self._protocol_lock: + reconciliation_state_is_valid = self._reconciliation_state_is_valid_locked( + snapshot, + accepted_failure, + accepted_generation, + ) + if not reconciliation_state_is_valid: + return False, RuntimeError( + "WideEP FT cannot stop after the accepted failure epoch changed " + f"(accepted_rank={accepted_failure}, " + f"expected_generation={accepted_generation}, " + f"snapshot={snapshot})" + ) + if not failed_ranks: + if self._transport_poisoned.is_set(): + return False, RuntimeError( + "WideEP FT cannot stop safely after an unattributed transport error" + ) + return True, None + + with self._protocol_lock: + unreconciled = [ + failed_rank + for failed_rank in failed_ranks + if not self._failure_is_reconciled_locked(failed_rank, snapshot.mask) + ] + untracked = [ + failed_rank + for failed_rank in unreconciled + if failed_rank not in self._reconcile_deadlines + ] + if not unreconciled: + return True, None + if untracked: + return False, RuntimeError( + "WideEP FT cannot stop with locally failed ranks that were never announced: " + f"{untracked}" + ) + # Their existing reconciliation deadlines bound how long stop keeps + # progressing. Expiry enters the normal terminal ABORT path. + return False, None + + def _progress_terminal_abort(self) -> bool: + """Return ``True`` after terminal state is globally safe to stop.""" + if self.ulfm_available: + # Revoke is itself a communicator-wide terminal notification; it + # does not depend on rendezvous Isends reaching completion. + self._progress_failed.set() + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + try: + self._try_revoke() + except Exception as error: + logger.warning(f"WideEP FT terminal revoke failed: {error}") + self._try_world_abort() + return True + if self.ulfm_available: + return True + # A runtime unsupported result disables ULFM and falls through to + # the relayed-ABORT protocol. Clear the flag so the next loop pass + # can continue progressing those requests. + self._progress_failed.clear() + + now = time.monotonic() + with self._protocol_lock: + reconciled = self._abort_fanout_complete and self._abort_expected_reporters.issubset( + self._abort_reporters + ) + deadline = self._abort_deadline + expired = deadline is not None and now >= deadline + if not reconciled and not expired: + return False + + self._progress_failed.set() + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + # A final echo can race the deadline check. Once every expected + # survivor has acknowledged ABORT, fail-stop escalation is no longer + # needed even if the wall-clock deadline has just elapsed. + if expired and not reconciled: + self._try_world_abort() + return True + + def _fail_closed_immediately(self, error: BaseException) -> None: + """Abort globally after an error makes bounded relay unsafe.""" + with self._lifecycle_lock: + self._record_error(error) + self._progress_failed.set() + self._lifecycle = _Lifecycle.FAILED + if self.ulfm_available: + try: + self._try_revoke() + if self.ulfm_available: + return + except Exception as revoke_error: + logger.warning(f"WideEP FT could not revoke after terminal error: {revoke_error}") + self._try_world_abort() + + def _fail_closed_after_cleanup(self, error: BaseException) -> None: + """Terminate the world when healthy MPI requests cannot be quiesced. + + At this point the progress loop has already stopped, so a relayed ABORT + cannot be made reliable and ULFM revoke alone would not update peers' + process-local shutdown state. Communicator-wide abort is the only + deterministic way to prevent rank-asymmetric MPI_Finalize behavior. + """ + self._record_error(error) + self._progress_failed.set() + with self._lifecycle_lock: + self._lifecycle = _Lifecycle.FAILED + self._try_world_abort() + + def _try_world_abort(self) -> None: + """Last-resort fail-stop when terminal state cannot be reconciled.""" + with self._terminal_action_lock: + if self._world_abort_attempted: + return + # MPI_Abort normally never returns. Mark the attempt before calling + # it so test doubles, broken runtimes, or a later-resuming progress + # thread cannot issue the communicator-wide action twice. + self._world_abort_attempted = True + abort = getattr(self._comm, "Abort", None) + if not callable(abort): + logger.error( + "WideEP FT communicator has no MPI_Abort fallback; " + "survivors must terminate through the poisoned-world shutdown hook" + ) + return + try: + abort(1) + except Exception as error: + # Real MPI_Abort does not return. Test doubles and broken runtimes + # may do so or raise; the local process remains terminal either way. + logger.error(f"WideEP FT communicator-wide abort failed: {error}") + + def _retain_requests(self, requests: list[_PendingSend | _PendingReceive]) -> None: + self._retained_requests.extend(requests) + references: list[object] = list(requests) + if not self._process_lifetime_retained: + self._process_lifetime_retained = True + references.append(self._comm) + _retain_for_process_lifetime(*references) + + def _cleanup_requests( + self, + pending_sends: list[_PendingSend], + pending_receives: dict[int, _PendingReceive], + ) -> None: + requests: list[_PendingSend | _PendingReceive] = [ + *pending_sends, + *pending_receives.values(), + ] + # Never cancel, wait for, or free MPI operations after a peer or the + # communicator has failed. Keeping each request together with its + # NumPy buffer is the only bounded shutdown behavior in that state; + # the later model-engine shutdown hook skips MPI_Finalize entirely. + if self.last_error is not None or self.world_is_poisoned(): + self._retain_requests(requests) + return + + if not requests: + return + + # Healthy shutdown can locally cancel the preposted receives/sends. + # Poll Test rather than calling Wait so stop() remains bounded even on + # an MPI implementation with surprising cancellation behavior. + for pending in requests: + try: + pending.request.Cancel() + except Exception as error: + if not isinstance(error, self._mpi.Exception): + raise + logger.warning(f"WideEP FT could not cancel a healthy MPI request: {error}") + + remaining = requests + deadline = self._shutdown_deadline or time.monotonic() + while remaining and time.monotonic() < deadline: + next_remaining: list[_PendingSend | _PendingReceive] = [] + for pending in remaining: + try: + if not pending.request.Test(): + next_remaining.append(pending) + except Exception as error: + if not isinstance(error, self._mpi.Exception): + raise + next_remaining.append(pending) + remaining = next_remaining + if remaining: + time.sleep(min(0.001, max(0.0, deadline - time.monotonic()))) + + if remaining: + raise TimeoutError( + f"{len(remaining)} WideEP FT MPI requests remained active after " + "the healthy shutdown timeout" + ) + + def _record_error(self, error: BaseException) -> None: + self._transport_poisoned.set() + with self._error_lock: + if self._last_error is None: + self._last_error = error + + def _try_revoke(self) -> None: + # Claim the one permitted Revoke attempt, but never hold a Python lock + # across MPI. A caller-thread timeout and a later-resuming progress + # thread may enter terminal handling concurrently; the claimant owns + # any unsupported/error fallback. + with self._terminal_action_lock: + if not self._ulfm_available or self._revoke_attempted: + return + self._revoke_attempted = True + try: + self._comm.Revoke() + except Exception as error: + if isinstance(error, NotImplementedError) or self._is_unsupported_ulfm_error(error): + with self._terminal_action_lock: + self._ulfm_available = False + logger.warning(f"WideEP FT revoke is unavailable at runtime: {error}") + return + raise diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index a0bbd6ab4f93..b8a0823591ca 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -34,6 +34,12 @@ l0_a10: - unittest/_torch/executor/test_kv_cache_budget_split.py - unittest/_torch/executor/test_kv_pool_rebalance.py - unittest/_torch/executor/test_disagg_index_mapper_early_release.py + - unittest/_torch/distributed/test_communicator.py + - unittest/_torch/pyexecutor/test_ep_failure_broadcast.py + # MPI prerequisites are required by default. Local environments can opt out + # with TLLM_ALLOW_MPI_FT_SMOKE_SKIP=1. The three-minute outer timeout covers + # four 30-second launcher modes plus bounded cleanup. + - unittest/_torch/pyexecutor/test_ep_failure_broadcast_mpi.py TIMEOUT (3) ISOLATION - unittest/_torch/pyexecutor/test_kv_cache_compression_manager.py - unittest/_torch/pyexecutor/test_error_classification.py - unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py diff --git a/tests/unittest/_torch/distributed/test_communicator.py b/tests/unittest/_torch/distributed/test_communicator.py new file mode 100644 index 000000000000..0ed029297a00 --- /dev/null +++ b/tests/unittest/_torch/distributed/test_communicator.py @@ -0,0 +1,681 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.distributed import communicator + + +def _mapping(**overrides: object) -> SimpleNamespace: + values = { + "world_size": 8, + "rank": 5, + "moe_ep_group": list(range(8)), + "moe_ep_size": 8, + "moe_ep_rank": 5, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _valid_topologies(parent_size: int = 8) -> list[dict[str, object]]: + topologies = [] + for rank in range(parent_size): + ep_group = tuple(range(parent_size)) + topologies.append( + { + "parent_rank": rank, + "parent_size": parent_size, + "world_size": parent_size, + "mapping_rank": rank, + "ep_group": ep_group, + "ep_size": len(ep_group), + "ep_rank": ep_group.index(rank), + "health_size": None, + "error_kind": None, + "error_message": None, + } + ) + return topologies + + +def _valid_setup_statuses(parent_size: int = 8) -> list[dict[str, object]]: + return [ + { + "parent_rank": rank, + "error_message": None, + "ulfm_available": True, + } + for rank in range(parent_size) + ] + + +def _communicators(*, parent_rank: int = 5, parent_size: int = 8) -> tuple[Mock, Mock]: + ft_comm = Mock() + ft_comm.Get_rank.return_value = 5 + ft_comm.Get_size.return_value = 8 + ft_comm.Is_revoked.return_value = False + + parent_comm = Mock() + parent_comm.Get_rank.return_value = parent_rank + parent_comm.Get_size.return_value = parent_size + parent_comm.Split.return_value = ft_comm + + def _allgather(local_record: dict[str, object]) -> list[dict[str, object]]: + if "ep_group" in local_record: + records = _valid_topologies(parent_size) + else: + records = _valid_setup_statuses(parent_size) + records[parent_rank] = local_record + return records + + parent_comm.allgather.side_effect = _allgather + return parent_comm, ft_comm + + +def _fake_mpi(parent_comm: Mock, *, thread_level: int = 3) -> SimpleNamespace: + return SimpleNamespace( + COMM_WORLD=parent_comm, + ERRORS_RETURN=object(), + THREAD_MULTIPLE=3, + Exception=Exception, + Query_thread=Mock(return_value=thread_level), + ) + + +def test_create_mpi_ft_subcomm_splits_ep_group_and_sets_error_handler( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + fake_mpi = _fake_mpi(parent_comm) + monkeypatch.setattr(communicator, "MPI", fake_mpi) + active_comm = Mock(return_value=parent_comm) + monkeypatch.setattr(communicator, "mpi_comm", active_comm) + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + try: + result = communicator.create_mpi_ft_subcomm(_mapping()) + + assert result.comm is ft_comm + assert result.local_rank == 5 + assert result.ep_size == 8 + assert result.ulfm_available is True + active_comm.assert_called_once_with() + assert parent_comm.allgather.call_count == 2 + parent_comm.Split.assert_called_once_with(color=0, key=5) + ft_comm.Set_errhandler.assert_called_once_with(fake_mpi.ERRORS_RETURN) + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_requires_thread_multiple( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, _ = _communicators() + fake_mpi = _fake_mpi(parent_comm, thread_level=2) + monkeypatch.setattr(communicator, "MPI", fake_mpi) + + with pytest.raises(RuntimeError, match=r"requires MPI\.THREAD_MULTIPLE"): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +@pytest.mark.parametrize( + ("mapping_overrides", "error_match"), + [ + ({"moe_ep_group": []}, "must not be empty"), + ({"moe_ep_group": list(range(7))}, "size must match mapping.moe_ep_size"), + ({"moe_ep_group": [0, 1, 2, 3, 4, 5, 6, 6]}, "contains duplicate ranks"), + ({"moe_ep_group": [0, 1, 2, 3, 4, 5, 6, 8]}, "outside the parent communicator"), + ( + {"moe_ep_group": [0, 1, 2, 3], "moe_ep_size": 4}, + "spanning the full MPI world", + ), + ( + {"moe_ep_group": [5, 0, 1, 2, 3, 4, 6, 7]}, + "group order must match mapping.moe_ep_rank", + ), + ], +) +def test_create_mpi_ft_subcomm_validates_ep_group( + monkeypatch: pytest.MonkeyPatch, + mapping_overrides: dict[str, object], + error_match: str, +) -> None: + parent_comm, _ = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(ValueError, match=error_match): + communicator.create_mpi_ft_subcomm(_mapping(**mapping_overrides), parent_comm) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +@pytest.mark.parametrize( + ("parent_rank", "parent_size", "error_match"), + [ + (5, 7, "size must match mapping.world_size"), + (6, 8, "rank must match mapping.rank"), + ], +) +def test_create_mpi_ft_subcomm_validates_parent_mapping( + monkeypatch: pytest.MonkeyPatch, + parent_rank: int, + parent_size: int, + error_match: str, +) -> None: + parent_comm, _ = _communicators(parent_rank=parent_rank, parent_size=parent_size) + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(RuntimeError, match=error_match): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +@pytest.mark.parametrize( + ("mapping_error", "error_match"), + [ + (IndexError("missing EP group"), "IndexError: missing EP group"), + (ValueError("invalid EP group"), "ValueError: invalid EP group"), + ], + ids=["index-error", "value-error"], +) +def test_create_mpi_ft_subcomm_collectively_reports_mapping_accessor_error( + monkeypatch: pytest.MonkeyPatch, + mapping_error: Exception, + error_match: str, +) -> None: + class FailingMapping: + world_size = 8 + rank = 5 + moe_ep_size = 8 + moe_ep_rank = 5 + + @property + def moe_ep_group(self) -> list[int]: + raise mapping_error + + parent_comm, _ = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(ValueError, match=error_match): + communicator.create_mpi_ft_subcomm(FailingMapping(), parent_comm) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +def test_create_mpi_ft_subcomm_uses_collectively_validated_ep_size( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class SingleReadMapping: + world_size = 8 + rank = 5 + moe_ep_group = list(range(8)) + + def __init__(self) -> None: + self.ep_size_reads = 0 + self.ep_rank_reads = 0 + + @property + def moe_ep_size(self) -> int: + self.ep_size_reads += 1 + if self.ep_size_reads > 1: + raise ValueError("moe_ep_size was read after collective validation") + return 8 + + @property + def moe_ep_rank(self) -> int: + self.ep_rank_reads += 1 + if self.ep_rank_reads > 1: + raise ValueError("moe_ep_rank was read after collective validation") + return 5 + + parent_comm, ft_comm = _communicators() + fake_mpi = _fake_mpi(parent_comm) + monkeypatch.setattr(communicator, "MPI", fake_mpi) + mapping = SingleReadMapping() + + result = communicator.create_mpi_ft_subcomm(mapping, parent_comm) + + assert result.comm is ft_comm + assert result.local_rank == 5 + assert result.ep_size == 8 + assert result.ulfm_available is True + assert mapping.ep_size_reads == 1 + assert mapping.ep_rank_reads == 1 + assert parent_comm.allgather.call_count == 2 + ft_comm.Set_errhandler.assert_called_once_with(fake_mpi.ERRORS_RETURN) + + +@pytest.mark.parametrize( + ("failure_kind", "error_match"), + [ + ("rank-mismatch", "unexpected WideEP FT communicator"), + ("size-mismatch", "unexpected WideEP FT communicator"), + ("get-rank-error", "RuntimeError: cannot inspect rank"), + ("get-size-error", "RuntimeError: cannot inspect size"), + ("errhandler", "RuntimeError: cannot set error handler"), + ], +) +def test_create_mpi_ft_subcomm_retains_comm_after_post_split_failure( + monkeypatch: pytest.MonkeyPatch, + failure_kind: str, + error_match: str, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + if failure_kind == "rank-mismatch": + ft_comm.Get_rank.return_value = 4 + elif failure_kind == "size-mismatch": + ft_comm.Get_size.return_value = 7 + elif failure_kind == "get-rank-error": + ft_comm.Get_rank.side_effect = RuntimeError("cannot inspect rank") + elif failure_kind == "get-size-error": + ft_comm.Get_size.side_effect = RuntimeError("cannot inspect size") + else: + ft_comm.Set_errhandler.side_effect = RuntimeError("cannot set error handler") + + try: + with pytest.raises(RuntimeError, match=error_match): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert parent_comm.allgather.call_count == 2 + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_installs_error_handler_before_inspection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + fake_mpi = _fake_mpi(parent_comm) + monkeypatch.setattr(communicator, "MPI", fake_mpi) + calls = [] + ft_comm.Set_errhandler.side_effect = lambda _handler: calls.append("errhandler") + ft_comm.Get_rank.side_effect = lambda: calls.append("rank") or 5 + ft_comm.Get_size.side_effect = lambda: calls.append("size") or 8 + + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert calls[:3] == ["errhandler", "rank", "size"] + + +def test_create_mpi_ft_subcomm_collectively_disables_ulfm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + default_allgather = parent_comm.allgather.side_effect + + def _allgather(local_record: dict[str, object]) -> list[dict[str, object]]: + records = default_allgather(local_record) + if "ep_group" not in local_record: + records[2]["ulfm_available"] = False + return records + + parent_comm.allgather.side_effect = _allgather + + result = communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert result.comm is ft_comm + assert result.ulfm_available is False + + +def test_create_mpi_ft_subcomm_collectively_reports_ulfm_probe_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + ft_comm.Is_revoked.side_effect = RuntimeError("probe failed") + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + try: + with pytest.raises( + RuntimeError, + match="setup failed on parent rank 5.*ULFM probe raised.*probe failed", + ): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_reconciles_ulfm_error_classifier_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class ClassifierFailureMpiException(Exception): + def Get_error_class(self) -> int: + raise RuntimeError("cannot classify MPI error") + + parent_comm, ft_comm = _communicators() + fake_mpi = _fake_mpi(parent_comm) + fake_mpi.Exception = ClassifierFailureMpiException + monkeypatch.setattr(communicator, "MPI", fake_mpi) + ft_comm.Is_revoked.side_effect = ClassifierFailureMpiException("probe failed") + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + try: + with pytest.raises( + RuntimeError, + match="setup failed on parent rank 5.*ULFM probe raised.*probe failed", + ): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert parent_comm.allgather.call_count == 2 + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_treats_unsupported_ulfm_probe_as_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + ft_comm.Is_revoked.side_effect = NotImplementedError("ULFM is not compiled") + + result = communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert result.ulfm_available is False + ft_comm.Free.assert_not_called() + + +def test_create_mpi_ft_subcomm_collectively_rejects_revoked_comm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + ft_comm.Is_revoked.return_value = True + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + try: + with pytest.raises(RuntimeError, match="already revoked at construction"): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_retains_comm_after_remote_post_split_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + default_allgather = parent_comm.allgather.side_effect + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + def _allgather(local_record: dict[str, object]) -> list[dict[str, object]]: + records = default_allgather(local_record) + if "ep_group" not in local_record: + records[2]["error_message"] = "remote communicator setup failed" + return records + + parent_comm.allgather.side_effect = _allgather + + try: + with pytest.raises( + RuntimeError, + match="setup failed on parent rank 2: remote communicator setup failed", + ): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + ft_comm.Set_errhandler.assert_called_once() + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_retains_comm_when_post_split_allgather_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + default_allgather = parent_comm.allgather.side_effect + allgather_error = RuntimeError("post-split allgather failed") + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + def _allgather(local_record: dict[str, object]) -> list[dict[str, object]]: + if "ep_group" in local_record: + return default_allgather(local_record) + raise allgather_error + + parent_comm.allgather.side_effect = _allgather + + try: + with pytest.raises(RuntimeError, match="post-split allgather failed") as raised: + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert raised.value is allgather_error + assert parent_comm.allgather.call_count == 2 + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_retains_comm_after_malformed_post_split_status( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, ft_comm = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + default_allgather = parent_comm.allgather.side_effect + initial_refs = len(communicator._MPI_FT_PROCESS_LIFETIME_REFS) + + def _allgather(local_record: dict[str, object]) -> list[object]: + records = default_allgather(local_record) + if "ep_group" not in local_record: + records[2] = None + return records + + parent_comm.allgather.side_effect = _allgather + + try: + with pytest.raises( + RuntimeError, + match="invalid status for parent rank 2: None", + ): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + ft_comm.Free.assert_not_called() + assert communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] == [ft_comm] + finally: + del communicator._MPI_FT_PROCESS_LIFETIME_REFS[initial_refs:] + + +def test_create_mpi_ft_subcomm_propagates_remote_validation_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, _ = _communicators() + topologies = _valid_topologies() + topologies[2]["error_kind"] = "value" + topologies[2]["error_message"] = "mapping.moe_ep_group must not be empty" + parent_comm.allgather.side_effect = None + parent_comm.allgather.return_value = topologies + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises( + ValueError, match="startup validation failed on parent rank 2.*must not be empty" + ): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + parent_comm.Split.assert_not_called() + + +def test_create_mpi_ft_subcomm_collectively_validates_health_size( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, _ = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(ValueError, match="EPGroupHealth size must match"): + communicator.create_mpi_ft_subcomm( + _mapping(), + parent_comm, + health_size=7, + ) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +def test_create_mpi_ft_subcomm_rejects_inconsistent_remote_ep_group( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, _ = _communicators() + topologies = _valid_topologies() + topologies[6]["ep_group"] = (6, 0, 1, 2, 3, 4, 5, 7) + parent_comm.allgather.side_effect = None + parent_comm.allgather.return_value = topologies + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(RuntimeError, match="startup topology is inconsistent"): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + parent_comm.Split.assert_not_called() + + +@pytest.mark.parametrize( + ("field", "value", "error_match"), + [ + ("ep_group", None, "invalid EP group"), + ("ep_group", "01234567", "invalid EP group"), + ("ep_group", (0, 1, True, 3, 4, 5, 6, 7), "invalid EP group ranks"), + ("ep_group", (0, 1, 2, 3, 4, 5, 6, 8), "invalid EP group ranks"), + ("ep_group", (0, 1, 2, 3, 4, 5, 6, 6), "duplicate EP group ranks"), + ("ep_size", None, "reports EP size"), + ("parent_rank", 2.0, "reports parent rank"), + ("parent_size", 8.0, "reports communicator size"), + ("world_size", 8.0, "reports world size"), + ("mapping_rank", 2.0, "reports mapping rank"), + ("ep_size", 8.0, "reports EP size"), + ("ep_rank", 2.0, "reports invalid EP rank"), + ("health_size", 8.0, "reports health size"), + ], + ids=[ + "missing-group", + "string-group", + "bool-rank", + "out-of-range-rank", + "duplicate-rank", + "missing-ep-size", + "float-parent-rank", + "float-parent-size", + "float-world-size", + "float-mapping-rank", + "float-ep-size", + "float-ep-rank", + "float-health-size", + ], +) +def test_create_mpi_ft_subcomm_rejects_malformed_remote_topology_schema( + monkeypatch: pytest.MonkeyPatch, + field: str, + value: object, + error_match: str, +) -> None: + parent_comm, _ = _communicators() + topologies = _valid_topologies() + topologies[2][field] = value + parent_comm.allgather.side_effect = None + parent_comm.allgather.return_value = topologies + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(RuntimeError, match=error_match): + communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + parent_comm.Split.assert_not_called() + + +def test_create_mpi_ft_subcomm_rejects_float_local_ep_size_before_split( + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent_comm, _ = _communicators() + monkeypatch.setattr(communicator, "MPI", _fake_mpi(parent_comm)) + + with pytest.raises(RuntimeError, match=r"reports EP size 8\.0"): + communicator.create_mpi_ft_subcomm( + _mapping(moe_ep_size=8.0), + parent_comm, + ) + + parent_comm.allgather.assert_called_once() + parent_comm.Split.assert_not_called() + + +def test_create_mpi_ft_subcomm_uses_pkl5_compatible_object_allgather( + monkeypatch: pytest.MonkeyPatch, +) -> None: + ft_comm = Mock() + ft_comm.Get_rank.return_value = 5 + ft_comm.Get_size.return_value = 8 + ft_comm.Is_revoked.return_value = False + + class Pkl5LikeParent: + def __init__(self) -> None: + self.gathered_records = [] + self.split_args = None + + def Get_rank(self) -> int: + return 5 + + def Get_size(self) -> int: + return 8 + + def allgather(self, local_record: dict[str, object]) -> list[dict[str, object]]: + self.gathered_records.append(local_record) + if "ep_group" in local_record: + records = _valid_topologies() + else: + records = _valid_setup_statuses() + records[5] = local_record + return records + + def Split(self, *, color: int, key: int) -> Mock: + self.split_args = (color, key) + return ft_comm + + parent_comm = Pkl5LikeParent() + fake_mpi = _fake_mpi(Mock()) + monkeypatch.setattr(communicator, "MPI", fake_mpi) + + result = communicator.create_mpi_ft_subcomm(_mapping(), parent_comm) + + assert result.comm is ft_comm + assert result.local_rank == 5 + assert result.ep_size == 8 + assert result.ulfm_available is True + assert len(parent_comm.gathered_records) == 2 + assert parent_comm.split_args == (0, 5) + ft_comm.Set_errhandler.assert_called_once_with(fake_mpi.ERRORS_RETURN) diff --git a/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py b/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py new file mode 100644 index 000000000000..42b1d119607c --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py @@ -0,0 +1,622 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Logical WideEP failure-injection worker using a real MPI transport. + +The smoke keeps every process alive so ``COMM_WORLD`` can coordinate portable +results. Actual peer death and MPI runtime error classification are outside its +scope because common launchers terminate the whole job when a worker exits +without ``MPI_Finalize``. +""" + +import atexit +import os +import sys +import time +import traceback +from types import SimpleNamespace + +import numpy as np +from mpi4py import MPI + +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth +from tensorrt_llm._torch.pyexecutor.ep_failure_broadcast import MpiFtSubcomm, MpiFtSubcommConfig +from tensorrt_llm.mapping import MpiTopology + +_SKIP_MARKER = "WIDEEP_MPI_SMOKE_SKIP:" +_PROPAGATION_MARKER = "WIDEEP_MPI_PROPAGATION" +_INVALID_TOPOLOGY_MARKER = "WIDEEP_MPI_INVALID_TOPOLOGY_OK" +_HEALTHY_LIFECYCLE_MARKER = "WIDEEP_MPI_HEALTHY_LIFECYCLE_OK" +_ABORT_WORLD_MARKER = "WIDEEP_MPI_ABORT_WORLD_OK" +_TERMINAL_READY_MARKER = "WIDEEP_MPI_TERMINAL_READY" +_TERMINAL_COMPLETE_MARKER = "WIDEEP_MPI_TERMINAL_COMPLETE" +_TERMINAL_DONE_MARKER = "WIDEEP_MPI_TERMINAL_DONE" +_TERMINAL_ERROR_MARKER = "WIDEEP_MPI_TERMINAL_ERROR" +_TERMINAL_ATEXIT_MARKER = "WIDEEP_MPI_TERMINAL_ATEXIT_RAN" +_PROPAGATION_TARGET_SEC = 0.1 +_CONVERGENCE_TIMEOUT_SEC = 2.0 +_ALLOW_SKIP_ENV = "TLLM_ALLOW_MPI_FT_SMOKE_SKIP" +_HEALTHY_MODE = "healthy" +_TERMINAL_MODE = "terminal" +_WORKER_SKIPPED = False + + +class _NonUlfmComm: + """Typed MPI communicator view that makes smoke behavior independent of ULFM.""" + + def __init__(self, comm: MPI.Intracomm) -> None: + self._comm = comm + self.abort_calls = 0 + + def Get_rank(self) -> int: + return self._comm.Get_rank() + + def Get_size(self) -> int: + return self._comm.Get_size() + + def Set_errhandler(self, errhandler: MPI.Errhandler) -> None: + self._comm.Set_errhandler(errhandler) + + def Irecv(self, buffer: np.ndarray, source: int, tag: int) -> MPI.Request: + return self._comm.Irecv(buffer, source=source, tag=tag) + + def Isend(self, buffer: np.ndarray, dest: int, tag: int) -> MPI.Request: + return self._comm.Isend(buffer, dest=dest, tag=tag) + + def Abort(self, errorcode: int) -> None: + self.abort_calls += 1 + self._comm.Abort(errorcode) + + +def _mapping(world_size: int, rank: int) -> MpiTopology: + return MpiTopology( + world_size=world_size, + rank=rank, + tp_size=world_size, + moe_tp_size=1, + moe_ep_size=world_size, + ) + + +def _unsupported_environment(message: str, *, rank: int) -> int: + """Skip optional local runs, but fail a designated required smoke run.""" + global _WORKER_SKIPPED + _WORKER_SKIPPED = True + required = os.environ.get(_ALLOW_SKIP_ENV) != "1" + if rank == 0: + if required: + print(f"WideEP MPI FT smoke requirement failed: {message}", file=sys.stderr) + else: + print(f"{_SKIP_MARKER} {message}") + return 2 if required else 0 + + +def _synchronize_phase_error( + world: MPI.Intracomm, + phase: str, + local_error: str | None, +) -> str | None: + """Return the same rank-attributed phase failure before the next collective.""" + errors = world.allgather(local_error) + failures = [(rank, error) for rank, error in enumerate(errors) if error is not None] + if not failures: + return None + details = "\n".join(f"rank {rank}:\n{error}" for rank, error in failures) + return f"{phase} failed before the next MPI phase:\n{details}" + + +def _run_invalid_topology_smoke(world: MPI.Intracomm) -> str | None: + """Prove a rank-local topology error is reconciled before ``Split``.""" + rank = world.Get_rank() + world_size = world.Get_size() + ep_group = tuple(range(world_size)) + if rank == 0: + ep_group = ep_group[:-1] + mapping = SimpleNamespace( + world_size=world_size, + rank=rank, + moe_ep_size=world_size, + moe_ep_rank=rank, + moe_ep_group=ep_group, + ) + caught_error: str | None = None + unexpected_broadcaster: MpiFtSubcomm | None = None + try: + unexpected_broadcaster = MpiFtSubcomm( + mapping, + EPGroupHealth(world_size), + MpiFtSubcommConfig(startup_timeout_sec=5.0, stop_timeout_sec=5.0), + ) + except Exception as error: + caught_error = f"{type(error).__name__}: {error}" + + caught_errors = world.allgather(caught_error) + expected_fragment = "startup validation failed on parent rank 0" + if any(error is None for error in caught_errors): + return f"invalid topology unexpectedly constructed on some ranks: {caught_errors}" + if len(set(caught_errors)) != 1 or expected_fragment not in caught_errors[0]: + return f"invalid topology errors were not identical across ranks: {caught_errors}" + + # This barrier is intentionally after the failed constructor. It proves the + # parent communicator did not strand any rank in MPI_Comm_split. + world.Barrier() + if rank == 0: + print(_INVALID_TOPOLOGY_MARKER, flush=True) + if unexpected_broadcaster is not None: + return "invalid topology unexpectedly returned a broadcaster" + return None + + +def _run_healthy_lifecycle_smoke(world: MPI.Intracomm) -> str | None: + """Exercise production construction, progress startup, and clean shutdown.""" + rank = world.Get_rank() + world_size = world.Get_size() + mapping = _mapping(world_size, rank) + broadcaster: MpiFtSubcomm | None = None + error: str | None = None + + try: + config = MpiFtSubcommConfig( + startup_timeout_sec=5.0, + stop_timeout_sec=5.0, + reconcile_timeout_sec=5.0, + ) + # This path obtains TRT-LLM's pkl5 parent through mpi_comm(), performs + # collective validation and Split, and owns MPI progress directly. + broadcaster = MpiFtSubcomm( + mapping, + EPGroupHealth(world_size), + config, + ) + ft_comm = broadcaster._comm + if ft_comm.Get_rank() != rank or ft_comm.Get_size() != world_size: + raise AssertionError( + "MPI_Comm_split returned an unexpected FT communicator: " + f"rank={ft_comm.Get_rank()}, size={ft_comm.Get_size()}" + ) + if ft_comm.Get_errhandler() != MPI.ERRORS_RETURN: + raise AssertionError("FT communicator does not use MPI.ERRORS_RETURN") + broadcaster.start() + broadcaster.stop() + if broadcaster.last_error is not None: + raise AssertionError( + "Production-construction progress failed during healthy shutdown: " + f"{broadcaster.last_error}" + ) + except Exception: + error = traceback.format_exc() + finally: + if broadcaster is not None: + try: + broadcaster.stop(timeout=5.0) + if broadcaster.last_error is not None: + raise AssertionError( + "MPI progress failed during healthy lifecycle smoke: " + f"{broadcaster.last_error}" + ) + except Exception: + cleanup_error = traceback.format_exc() + error = f"{error or ''}\ncleanup failure:\n{cleanup_error}" + # Production-created FT communicators are intentionally retained for + # process lifetime. Freeing one here would violate the same ownership + # contract this smoke is meant to exercise. + + return error + + +def _run_failure_propagation_smoke(world: MPI.Intracomm) -> str | None: + """Exercise logical failure fanout and sticky state over real MPI traffic.""" + rank = world.Get_rank() + world_size = world.Get_size() + mapping = _mapping(world_size, rank) + health = EPGroupHealth(world_size) + owner: MpiFtSubcomm | None = None + broadcaster: MpiFtSubcomm | None = None + detector_rank = 0 + failure_rank = world_size - 1 + error: str | None = None + propagation_converged = False + + setup_error: str | None = None + try: + config = MpiFtSubcommConfig( + startup_timeout_sec=5.0, + stop_timeout_sec=5.0, + reconcile_timeout_sec=5.0, + ) + # Create the real FT communicator collectively, then inject a non-ULFM + # view so this smoke deterministically exercises typed Isend/Irecv+Test. + owner = MpiFtSubcomm( + mapping, + EPGroupHealth(world_size), + config, + ) + ft_comm = owner._comm + owner.stop() + broadcaster = MpiFtSubcomm( + mapping, + health, + config, + comm=_NonUlfmComm(ft_comm), + ) + broadcaster.start() + except Exception: + setup_error = traceback.format_exc() + error = _synchronize_phase_error(world, "failure-propagation setup", setup_error) + + detector_elapsed_sec: float | None = None + local_converged = False + if error is None: + phase_errors: list[str] = [] + try: + # Every broadcaster is running before rank 0 starts the latency + # clock. This synchronization is deliberately outside the measured + # interval. + world.Barrier() + except Exception: + phase_errors.append(f"readiness:\n{traceback.format_exc()}") + + trigger_error: str | None = None + propagation_start: float | None = None + if not phase_errors and rank == detector_rank: + propagation_start = time.monotonic() + try: + assert broadcaster is not None + broadcaster.pre_failover(failure_rank) + except Exception: + trigger_error = traceback.format_exc() + + observation_error: str | None = None + try: + # The victim process remains alive only so the smoke can coordinate + # its result on COMM_WORLD. Survivors exclude its logical rank and + # exercise the same failure/reconciliation fanout as production. + assert broadcaster is not None + deadline = time.monotonic() + _CONVERGENCE_TIMEOUT_SEC + while not phase_errors and rank != failure_rank and time.monotonic() < deadline: + failure_observed = not health.is_active(failure_rank) + failure_reconciled = broadcaster.failure_is_reconciled(failure_rank) + if failure_observed and failure_reconciled: + if rank == detector_rank: + assert propagation_start is not None + detector_elapsed_sec = time.monotonic() - propagation_start + break + time.sleep(0.005) + local_converged = rank == failure_rank or ( + not health.is_active(failure_rank) + and broadcaster.failure_is_reconciled(failure_rank) + and broadcaster.health_is_reconciled() + and broadcaster.world_is_poisoned() + ) + except Exception: + observation_error = traceback.format_exc() + + if trigger_error is not None: + phase_errors.append(f"trigger:\n{trigger_error}") + if observation_error is not None: + phase_errors.append(f"observation:\n{observation_error}") + error = _synchronize_phase_error( + world, + "failure-propagation trigger/observation", + "\n".join(phase_errors) or None, + ) + + if error is None: + all_converged = world.allreduce(local_converged, op=MPI.LAND) + # Share only rank 0's completed duration after the timed interval; raw + # monotonic timestamps never cross rank or host boundaries. + detector_elapsed_sec = world.bcast(detector_elapsed_sec, root=detector_rank) + if all_converged and detector_elapsed_sec is not None: + propagation_converged = True + else: + error = ( + f"failure rank {failure_rank} did not converge within " + f"{_CONVERGENCE_TIMEOUT_SEC:.0f}s" + ) + if rank == detector_rank and detector_elapsed_sec is not None: + elapsed_ms = detector_elapsed_sec * 1000.0 + target_ms = _PROPAGATION_TARGET_SEC * 1000.0 + print( + f"{_PROPAGATION_MARKER} world_size={world_size} " + f"elapsed_ms={elapsed_ms!r} target_ms={target_ms!r} " + f"target_met={elapsed_ms < target_ms}", + flush=True, + ) + + stop_succeeded = False + if broadcaster is not None: + # Deliberately do not reactivate failure_rank. Survivors must take the + # poisoned shutdown path and retain their unmatched receive requests. + try: + broadcaster.stop(timeout=5.0) + if broadcaster.last_error is not None: + raise AssertionError( + f"MPI progress failed during propagation smoke: {broadcaster.last_error}" + ) + stop_succeeded = True + except Exception: + cleanup_error = traceback.format_exc() + error = f"{error or ''}\ncleanup failure:\n{cleanup_error}" + if propagation_converged and stop_succeeded: + try: + if rank == failure_rank: + if broadcaster.world_is_poisoned(): + raise AssertionError("logical victim unexpectedly entered poisoned state") + if broadcaster._retained_requests: + raise AssertionError("logical victim retained healthy MPI requests") + else: + if not broadcaster.world_is_poisoned(): + raise AssertionError("survivor lost sticky poisoned state during stop") + if not broadcaster._retained_requests: + raise AssertionError("survivor did not retain poisoned MPI requests") + except Exception: + sticky_error = traceback.format_exc() + error = f"{error or ''}\nsticky-poison failure:\n{sticky_error}" + if owner is not None: + try: + owner.stop(timeout=5.0) + except Exception: + cleanup_error = traceback.format_exc() + error = f"{error or ''}\nowner cleanup failure:\n{cleanup_error}" + + return error + + +def _write_terminal_manifest(lines: list[str]) -> None: + """Emit canonical rank evidence from the single rank-0 output stream.""" + sys.stdout.write("".join(f"{line}\n" for line in lines)) + sys.stdout.flush() + + +def _terminal_atexit_sentinel() -> None: + """Negative evidence: this must not run in intentional no-Finalize mode.""" + marker = f"{_TERMINAL_ATEXIT_MARKER} pid={os.getpid()}\n".encode() + try: + os.write(sys.stderr.fileno(), marker) + except BaseException: + pass + + +def _run_terminal_abort_smoke(world: MPI.Intracomm) -> str | None: + """Exercise terminal ABORT relay while keeping ``COMM_WORLD`` healthy.""" + rank = world.Get_rank() + world_size = world.Get_size() + mapping = _mapping(world_size, rank) + broadcaster: MpiFtSubcomm | None = None + wrapped_comm: _NonUlfmComm | None = None + error: str | None = None + + setup_error: str | None = None + try: + # Construct a fresh production communicator, then hide ULFM so this + # scenario specifically validates the bounded ABORT echo fallback. + owner = MpiFtSubcomm( + mapping, + EPGroupHealth(world_size), + MpiFtSubcommConfig(startup_timeout_sec=5.0, stop_timeout_sec=5.0), + ) + ft_comm = owner._comm + owner.stop() + wrapped_comm = _NonUlfmComm(ft_comm) + broadcaster = MpiFtSubcomm( + mapping, + EPGroupHealth(world_size), + MpiFtSubcommConfig( + startup_timeout_sec=5.0, + stop_timeout_sec=5.0, + reconcile_timeout_sec=5.0, + abort_timeout_sec=3.0, + ), + comm=wrapped_comm, + ) + broadcaster.start() + except Exception: + setup_error = traceback.format_exc() + error = _synchronize_phase_error(world, "terminal-ABORT setup", setup_error) + + if error is None: + trigger_error: str | None = None + try: + world.Barrier() + if rank == 0: + # A transport-independent terminal request is the protocol seam + # used by local second-failure and timeout paths. + assert broadcaster is not None + broadcaster._request_terminal_abort( + RuntimeError("real-MPI terminal smoke"), -1, rank + ) + except Exception: + trigger_error = traceback.format_exc() + error = _synchronize_phase_error( + world, + "terminal-ABORT trigger", + trigger_error, + ) + + local_converged = False + if error is None: + observation_error: str | None = None + try: + assert broadcaster is not None + assert wrapped_comm is not None + deadline = time.monotonic() + _CONVERGENCE_TIMEOUT_SEC + while time.monotonic() < deadline: + if broadcaster._progress_failed.is_set(): + break + time.sleep(0.005) + local_converged = ( + broadcaster._progress_failed.is_set() + and broadcaster.last_error is not None + and broadcaster.world_is_poisoned() + and wrapped_comm.abort_calls == 0 + ) + except Exception: + observation_error = traceback.format_exc() + error = _synchronize_phase_error( + world, + "terminal-ABORT observation", + observation_error, + ) + + if error is None: + if not world.allreduce(local_converged, op=MPI.LAND): + error = "terminal ABORT did not converge on every MPI rank" + else: + # The ABORT protocol is confined to the FT subcommunicator. A + # parent barrier after convergence proves COMM_WORLD remains usable. + world.Barrier() + if rank == 0: + print(_ABORT_WORLD_MARKER, flush=True) + if broadcaster is not None: + try: + broadcaster.stop(timeout=5.0) + if not broadcaster.world_is_poisoned(): + raise AssertionError("terminal FT state lost its poisoned-world latch") + except Exception: + cleanup_error = traceback.format_exc() + error = f"{error or ''}\nterminal cleanup failure:\n{cleanup_error}" + + # Do not free the terminal communicator. Its component intentionally + # retains potentially active requests until process teardown. + return error + + +def _report_scenario_errors( + world: MPI.Intracomm, + scenario: str, + local_error: str | None, +) -> bool: + rank = world.Get_rank() + errors = world.gather(local_error, root=0) + failed = None + if rank == 0: + failed = any(error is not None for error in errors) + if failed: + for error_rank, error in enumerate(errors): + if error is not None: + print( + f"{scenario} failed on rank {error_rank}:\n{error}", + file=sys.stderr, + ) + return world.bcast(failed, root=0) + + +def main() -> int: + world = MPI.COMM_WORLD + expected_world_size = int(sys.argv[1]) if len(sys.argv) > 1 else 0 + mode = sys.argv[2] if len(sys.argv) > 2 else "" + world_size = world.Get_size() + rank = world.Get_rank() + if world_size < 2 or world_size != expected_world_size: + return _unsupported_environment( + f"expected {expected_world_size} MPI ranks, got {world_size}", + rank=rank, + ) + if MPI.Query_thread() < MPI.THREAD_MULTIPLE: + return _unsupported_environment( + f"MPI.THREAD_MULTIPLE is unavailable (provided={MPI.Query_thread()})", + rank=rank, + ) + + if mode == _HEALTHY_MODE: + scenarios = ( + ("invalid-topology", _run_invalid_topology_smoke), + ("healthy-lifecycle", _run_healthy_lifecycle_smoke), + ) + elif mode == _TERMINAL_MODE: + scenarios = ( + ("failure-broadcast", _run_failure_propagation_smoke), + ("terminal-abort", _run_terminal_abort_smoke), + ) + else: + if rank == 0: + print(f"unknown WideEP MPI smoke mode: {mode!r}", file=sys.stderr) + return 2 + for scenario, run_scenario in scenarios: + if _report_scenario_errors(world, scenario, run_scenario(world)): + return 1 + if mode == _HEALTHY_MODE and rank == 0: + print(_HEALTHY_LIFECYCLE_MARKER, flush=True) + return 0 + + +if __name__ == "__main__": + worker_mode = sys.argv[2] if len(sys.argv) > 2 else "" + if worker_mode == _TERMINAL_MODE: + # A normal interpreter shutdown would run mpi4py's Finalize hook. The + # launcher rejects this sentinel, so the smoke independently proves that + # the terminal path below bypassed Python atexit via os._exit(). + atexit.register(_terminal_atexit_sentinel) + exit_code = main() + if worker_mode == _TERMINAL_MODE and not _WORKER_SKIPPED: + # This mode deliberately leaves the FT control plane terminal and + # retains its requests. Skip mpi4py's collective MPI_Finalize just as + # the production poisoned-world shutdown hook will do. + sys.stdout.flush() + sys.stderr.flush() + if exit_code == 0: + world = MPI.COMM_WORLD + rank = world.Get_rank() + world_size = world.Get_size() + try: + # Gather evidence first, then let rank 0 emit a canonical + # manifest. MPI launchers forward separate per-rank pipes and do + # not preserve cross-rank write atomicity in the merged stream. + ready_ranks = world.gather(rank, root=0) + if rank == 0: + if sorted(ready_ranks) != list(range(world_size)): + raise AssertionError(f"invalid terminal READY ranks: {ready_ranks}") + _write_terminal_manifest( + [ + f"{_TERMINAL_READY_MARKER} rank={ready_rank} world_size={world_size}" + for ready_rank in ready_ranks + ] + ) + world.Barrier() + done_ranks = world.gather(rank, root=0) + if rank == 0: + if sorted(done_ranks) != list(range(world_size)): + raise AssertionError(f"invalid terminal DONE ranks: {done_ranks}") + _write_terminal_manifest( + [ + f"{_TERMINAL_COMPLETE_MARKER} world_size={world_size}", + *( + f"{_TERMINAL_DONE_MARKER} rank={done_rank} world_size={world_size}" + for done_rank in done_ranks + ), + ] + ) + # No rank can trigger launcher fail-fast until rank 0 has copied + # the complete gathered manifest to stdout. + world.Barrier() + except BaseException: + try: + _write_terminal_manifest([f"{_TERMINAL_ERROR_MARKER} rank={rank} exit_code=1"]) + except BaseException: + pass + traceback.print_exc() + sys.stderr.flush() + exit_code = 1 + else: + rank = MPI.COMM_WORLD.Get_rank() + try: + _write_terminal_manifest( + [f"{_TERMINAL_ERROR_MARKER} rank={rank} exit_code={exit_code}"] + ) + except BaseException: + pass + os._exit(exit_code) + # Healthy mode returns normally so mpi4py's MPI_Finalize path remains part + # of the smoke test. + sys.exit(exit_code) diff --git a/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py new file mode 100644 index 000000000000..f71227103da3 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py @@ -0,0 +1,2107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the WideEP MPI failure-broadcast control plane. + +The tests use in-memory MPI, communicator, and request doubles. They exercise +the real progress thread without requiring mpi4py, GPUs, or a multi-rank +launcher. +""" + +from __future__ import annotations + +import gc +import sys +import threading +import time +import types +import weakref +from collections import defaultdict +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np +import pytest + +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth +from tensorrt_llm._torch.pyexecutor import ep_failure_broadcast +from tensorrt_llm._torch.pyexecutor.ep_failure_broadcast import MpiFtSubcomm, MpiFtSubcommConfig + +_TEST_CONFIG = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=2.0, + unattributed_error_timeout_sec=2.0, + abort_timeout_sec=2.0, +) + +_FAILURE_MESSAGE = int(ep_failure_broadcast._MessageKind.FAILURE) +_ABORT_MESSAGE = int(ep_failure_broadcast._MessageKind.ABORT) + + +@pytest.mark.parametrize("value", [0.0, -1.0, float("nan"), float("inf"), float("-inf")]) +def test_config_requires_positive_finite_timeouts(value: float) -> None: + with pytest.raises(ValueError, match="must be finite and > 0"): + MpiFtSubcommConfig(poll_interval_sec=value) + with pytest.raises(ValueError, match="must be finite and > 0"): + MpiFtSubcommConfig(unattributed_error_timeout_sec=value) + with pytest.raises(ValueError, match="must be finite and > 0"): + MpiFtSubcommConfig(abort_timeout_sec=value) + + +def test_default_poll_interval_leaves_margin_below_agreement_budget() -> None: + assert MpiFtSubcommConfig().poll_interval_sec == 0.01 + + +@dataclass(frozen=True) +class _FakeMapping: + world_size: int + rank: int + moe_ep_size: int + moe_ep_rank: int + moe_ep_group: tuple[int, ...] + + @classmethod + def ep_world(cls, size: int, rank: int = 0) -> "_FakeMapping": + return cls( + world_size=size, + rank=rank, + moe_ep_size=size, + moe_ep_rank=rank, + moe_ep_group=tuple(range(size)), + ) + + +class _FakeMpiException(RuntimeError): + def __init__(self, error_class: int, message: str = "fake MPI failure") -> None: + super().__init__(message) + self._error_class = error_class + + def Get_error_class(self) -> int: + return self._error_class + + +class _FakeMPI: + ERRORS_RETURN = object() + ERR_PROC_FAILED = 101 + ERR_REVOKED = 102 + ERR_UNSUPPORTED_OPERATION = 103 + ERR_UNKNOWN = -1 + THREAD_MULTIPLE = 3 + Exception = _FakeMpiException + + def __init__(self, thread_level: int = THREAD_MULTIPLE) -> None: + self._thread_level = thread_level + + def Query_thread(self) -> int: + return self._thread_level + + +class _FakeRequest: + """Controllable nonblocking request that never owns its MPI buffer.""" + + def __init__(self, *, completed: bool = False) -> None: + self._completed = threading.Event() + if completed: + self._completed.set() + self._error: BaseException | None = None + self._lock = threading.Lock() + self.cancel_calls = 0 + self.test_calls = 0 + + def Cancel(self) -> None: + with self._lock: + self.cancel_calls += 1 + self._completed.set() + + def Test(self) -> bool: + with self._lock: + self.test_calls += 1 + error = self._error + if error is not None: + raise error + return self._completed.is_set() + + def complete(self) -> None: + self._completed.set() + + def fail(self, error: BaseException) -> None: + with self._lock: + self._error = error + + +class _BlockingRequest(_FakeRequest): + """Request whose Test call can hold the progress thread for stop tests.""" + + def __init__(self) -> None: + super().__init__() + self.entered_test = threading.Event() + self.release_test = threading.Event() + + def Test(self) -> bool: + self.entered_test.set() + self.release_test.wait() + return super().Test() + + +class _NonCompletingCancelRequest(_FakeRequest): + """Cancellation request that remains incomplete until explicitly released.""" + + def Cancel(self) -> None: + with self._lock: + self.cancel_calls += 1 + + +@dataclass(frozen=True) +class _ReceiveRecord: + request: _FakeRequest + buffer: np.ndarray + thread_id: int + + +@dataclass(frozen=True) +class _SendRecord: + request: _FakeRequest + buffer_ref: weakref.ReferenceType[np.ndarray] + destination: int + message_kind: int + payload: int + tag: int + thread_id: int + + +class _FakeComm: + def __init__( + self, + *, + size: int, + rank: int = 0, + receive_factory: Callable[[int], _FakeRequest] | None = None, + send_factory: Callable[[int], _FakeRequest] | None = None, + ulfm_probe_error: BaseException | None = None, + ulfm_revoked: bool = False, + revoke_error: BaseException | None = None, + abort_error: BaseException | None = None, + ) -> None: + self._size = size + self._rank = rank + self._receive_factory = receive_factory or (lambda _source: _FakeRequest()) + self._send_factory = send_factory or (lambda _destination: _FakeRequest()) + self._ulfm_probe_error = ulfm_probe_error + self._ulfm_revoked = ulfm_revoked + self._revoke_error = revoke_error + self._abort_error = abort_error + self._lock = threading.Lock() + self._receives: dict[int, list[_ReceiveRecord]] = defaultdict(list) + self._sends: list[_SendRecord] = [] + self.errhandler: object | None = None + self.errhandler_calls = 0 + self.is_revoked_calls = 0 + self.revoke_calls = 0 + self.abort_calls = 0 + + def Get_rank(self) -> int: + return self._rank + + def Get_size(self) -> int: + return self._size + + def Set_errhandler(self, errhandler: object) -> None: + self.errhandler_calls += 1 + self.errhandler = errhandler + + def Is_revoked(self) -> bool: + self.is_revoked_calls += 1 + if self._ulfm_probe_error is not None: + raise self._ulfm_probe_error + return self._ulfm_revoked + + def Revoke(self) -> None: + self.revoke_calls += 1 + if self._revoke_error is not None: + raise self._revoke_error + + def Abort(self, _errorcode: int) -> None: + self.abort_calls += 1 + if self._abort_error is not None: + raise self._abort_error + + def Irecv(self, buffer: np.ndarray, source: int, tag: int) -> _FakeRequest: + del tag + request = self._receive_factory(source) + record = _ReceiveRecord( + request=request, + buffer=buffer, + thread_id=threading.get_ident(), + ) + with self._lock: + self._receives[source].append(record) + return request + + def Isend(self, buffer: np.ndarray, dest: int, tag: int) -> _FakeRequest: + request = self._send_factory(dest) + record = _SendRecord( + request=request, + buffer_ref=weakref.ref(buffer), + destination=dest, + message_kind=int(buffer[0]), + payload=int(buffer[1]), + tag=tag, + thread_id=threading.get_ident(), + ) + with self._lock: + self._sends.append(record) + return request + + def receive_count(self, source: int) -> int: + with self._lock: + return len(self._receives[source]) + + def latest_receive(self, source: int) -> _ReceiveRecord: + with self._lock: + return self._receives[source][-1] + + def deliver( + self, + source: int, + failed_rank: int, + *, + message_kind: int = _FAILURE_MESSAGE, + ) -> None: + record = self.latest_receive(source) + record.buffer[:] = (message_kind, failed_rank) + record.request.complete() + + @property + def sends(self) -> tuple[_SendRecord, ...]: + with self._lock: + return tuple(self._sends) + + +def _wait_until( + predicate: Callable[[], bool], + *, + timeout: float = 1.0, + description: str = "condition", +) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.001) + raise AssertionError(f"timed out waiting for {description}") + + +def _make_broadcaster( + *, + size: int = 4, + rank: int = 0, + health: EPGroupHealth | None = None, + comm: _FakeComm | None = None, + mpi: _FakeMPI | None = None, + callback: Callable[[int, int, float], None] | None = None, +) -> tuple[MpiFtSubcomm, EPGroupHealth, _FakeComm, _FakeMPI]: + health = health or EPGroupHealth(size) + comm = comm or _FakeComm(size=size, rank=rank) + mpi = mpi or _FakeMPI() + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(size, rank), + health, + _TEST_CONFIG, + callback, + comm=comm, + mpi_module=mpi, + ) + return broadcaster, health, comm, mpi + + +def test_nonblocking_fanout_announces_externally_premarked_rank() -> None: + """The watchdog's prior local mark must not suppress cross-rank fanout.""" + health = EPGroupHealth(4) + assert health.mark_failed(2) is True + broadcaster, _, comm, _ = _make_broadcaster(health=health) + broadcaster.start() + caller_thread = threading.get_ident() + try: + assert broadcaster.broadcast_failure(2) is False + _wait_until(lambda: len(comm.sends) == 2, description="failure fanout") + _wait_until( + lambda: all(record.request.test_calls > 0 for record in comm.sends), + description="send request polling", + ) + + assert {record.destination for record in comm.sends} == {1, 3} + assert {record.payload for record in comm.sends} == {2} + assert all(record.thread_id != caller_thread for record in comm.sends) + assert all(record.request.test_calls > 0 for record in comm.sends) + assert all(record.buffer_ref() is not None for record in comm.sends) + + # A repeated detector report is locally coalesced. + assert broadcaster.pre_failover(2) is False + time.sleep(2 * _TEST_CONFIG.poll_interval_sec) + assert len(comm.sends) == 2 + + for send in comm.sends: + send.request.complete() + comm.deliver(source=1, failed_rank=2) + comm.deliver(source=3, failed_rank=2) + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + finally: + broadcaster.stop() + + +def test_survivor_echoes_reconcile_failure_before_next_collective() -> None: + broadcaster, _, comm, _ = _make_broadcaster(size=4) + broadcaster.start() + try: + assert broadcaster.pre_failover(2) is True + assert broadcaster.failure_is_reconciled(2) is False + _wait_until(lambda: len(comm.sends) == 2, description="initial failure fanout") + + comm.deliver(source=1, failed_rank=2) + _wait_until( + lambda: comm.receive_count(1) == 2, + description="source 1 receive repost", + ) + assert broadcaster.failure_is_reconciled(2) is False + + comm.deliver(source=3, failed_rank=2) + for send in comm.sends: + send.request.complete() + _wait_until( + lambda: broadcaster.failure_is_reconciled(2), + description="all survivor echoes", + ) + time.sleep(2 * _TEST_CONFIG.poll_interval_sec) + assert len(comm.sends) == 2 + assert {send.destination for send in comm.sends} == {1, 3} + assert all(send.message_kind == _FAILURE_MESSAGE for send in comm.sends) + assert broadcaster.health_is_reconciled() is True + assert comm.revoke_calls == 0 + finally: + broadcaster.stop() + + +def test_rank_reactivation_cannot_reopen_reconciliation_for_same_comm_epoch() -> None: + broadcaster, health, comm, _ = _make_broadcaster(size=3) + broadcaster.start() + try: + assert broadcaster.pre_failover(2) is True + _wait_until(lambda: len(comm.sends) == 1, description="failure fanout") + comm.sends[0].request.complete() + comm.deliver(source=1, failed_rank=2) + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + + assert health.mark_active(2) is True + _wait_until( + lambda: isinstance(broadcaster.last_error, RuntimeError), + description="reactivated-rank terminal handling", + ) + _wait_until(lambda: comm.revoke_calls == 1, description="epoch-violation revoke") + assert broadcaster.world_is_poisoned() is True + assert broadcaster.failure_is_reconciled(2) is False + assert broadcaster.health_is_reconciled() is False + + # Even restoring the failed bit cannot restore the old consensus: the + # intervening generation change belongs to a future communicator epoch. + assert health.mark_failed(2) is True + assert broadcaster.world_is_poisoned() is True + assert broadcaster.failure_is_reconciled(2) is False + assert broadcaster.health_is_reconciled() is False + finally: + broadcaster.stop() + + +def test_reactivation_between_failure_claim_and_observe_cannot_become_epoch_baseline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + health = EPGroupHealth(2) + assert health.mark_failed(1) is True + broadcaster, _, _, _ = _make_broadcaster(size=2, health=health) + original_claim = broadcaster._claim_failure + mutated = False + + def claim_then_turn_over_epoch(failed_rank: int) -> None: + nonlocal mutated + original_claim(failed_rank) + if not mutated: + mutated = True + assert health.mark_active(failed_rank) is True + assert health.mark_failed(failed_rank) is True + + monkeypatch.setattr(broadcaster, "_claim_failure", claim_then_turn_over_epoch) + broadcaster.start() + try: + assert broadcaster.broadcast_failure(1) is False + _wait_until( + lambda: 1 in broadcaster._failure_fanout_complete, + description="single-survivor fanout completion", + ) + assert health.generation == 3 + _wait_until( + lambda: isinstance(broadcaster.last_error, RuntimeError), + description="turnover terminal handling", + ) + assert broadcaster.failure_is_reconciled(1) is False + assert broadcaster.health_is_reconciled() is False + finally: + broadcaster.stop() + + +def test_second_failure_between_claim_and_observe_cannot_open_single_failure_gate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + health = EPGroupHealth(3) + broadcaster, _, _, _ = _make_broadcaster(size=3, health=health) + original_claim = broadcaster._claim_failure + mutated = False + + def claim_then_add_second_failure(failed_rank: int) -> None: + nonlocal mutated + original_claim(failed_rank) + if not mutated: + mutated = True + assert health.mark_failed(1) is True + + monkeypatch.setattr(broadcaster, "_claim_failure", claim_then_add_second_failure) + broadcaster.start() + try: + assert broadcaster.broadcast_failure(2) is True + _wait_until( + lambda: 2 in broadcaster._failure_fanout_complete, + description="single-survivor fanout completion", + ) + assert health.get_failed_ranks() == frozenset({1, 2}) + _wait_until( + lambda: isinstance(broadcaster.last_error, RuntimeError), + description="second-health-failure terminal handling", + ) + assert broadcaster.failure_is_reconciled(2) is False + assert broadcaster.health_is_reconciled() is False + finally: + broadcaster.stop() + + +def test_reconciliation_timeout_fails_closed_while_mpi_test_is_wedged() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.02, + ) + blocking_request = _BlockingRequest() + comm = _FakeComm( + size=3, + receive_factory=lambda source: blocking_request if source == 1 else _FakeRequest(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(3), + EPGroupHealth(3), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + _wait_until(blocking_request.entered_test.is_set, description="blocked MPI progress") + assert broadcaster.pre_failover(2) is True + assert blocking_request.release_test.is_set() is False + assert broadcaster._outbound_reports.empty() is False + + _wait_until( + lambda: isinstance(broadcaster.last_error, TimeoutError), + description="terminal reconciliation timeout", + ) + _wait_until(lambda: comm.revoke_calls == 1, description="deadline-monitor revoke") + assert blocking_request.release_test.is_set() is False + assert broadcaster.health_is_reconciled() is False + assert broadcaster.world_is_poisoned() is True + with pytest.raises(RuntimeError, match="not running"): + broadcaster.pre_failover(1) + blocking_request.release_test.set() + broadcaster.stop() + + +def test_reconciliation_timeout_world_aborts_when_mpi_test_is_wedged_without_ulfm() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.01, + abort_timeout_sec=0.01, + ) + blocking_request = _BlockingRequest() + comm = _FakeComm( + size=3, + receive_factory=lambda source: blocking_request if source == 1 else _FakeRequest(), + ulfm_probe_error=NotImplementedError(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(3), + EPGroupHealth(3), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + _wait_until(blocking_request.entered_test.is_set, description="blocked MPI progress") + assert broadcaster.pre_failover(2) is True + + _wait_until(lambda: comm.abort_calls == 1, description="deadline-monitor MPI_Abort") + assert blocking_request.release_test.is_set() is False + assert isinstance(broadcaster.last_error, TimeoutError) + assert broadcaster.world_is_poisoned() is True + + blocking_request.release_test.set() + broadcaster.stop() + + +def test_second_distinct_local_failure_fails_closed_before_health_mutation() -> None: + broadcaster, health, comm, _ = _make_broadcaster(size=4) + broadcaster.start() + assert broadcaster.pre_failover(3) is True + + with pytest.raises(RuntimeError, match="exactly one distinct failed rank") as raised: + broadcaster.pre_failover(2) + + assert health.get_failed_ranks() == frozenset({3}) + _wait_until(lambda: broadcaster.last_error is raised.value, description="terminal failure") + _wait_until(lambda: comm.revoke_calls == 1, description="progress-thread abort revoke") + assert broadcaster.health_is_reconciled() is False + with pytest.raises(RuntimeError, match="not running"): + broadcaster.pre_failover(1) + broadcaster.stop() + + +def test_no_ulfm_terminal_abort_relays_and_converges_without_world_abort() -> None: + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, health, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + with pytest.raises(RuntimeError, match="exactly one distinct failed rank"): + broadcaster.pre_failover(1) + + _wait_until( + lambda: any(send.message_kind == _ABORT_MESSAGE for send in comm.sends), + description="terminal ABORT relay", + ) + for send in comm.sends: + if send.message_kind == _ABORT_MESSAGE: + send.request.complete() + comm.deliver(source=1, failed_rank=1, message_kind=_ABORT_MESSAGE) + _wait_until(broadcaster._progress_failed.is_set, description="ABORT convergence") + + assert health.get_failed_ranks() == frozenset({2}) + assert comm.abort_calls == 0 + assert comm.revoke_calls == 0 + broadcaster.stop() + + +def test_remote_conflicting_failure_relays_abort_before_world_abort() -> None: + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, health, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + _wait_until( + lambda: any(send.message_kind == _FAILURE_MESSAGE for send in comm.sends), + description="first failure fanout", + ) + _wait_until(lambda: comm.receive_count(1) == 1, description="conflicting report receive") + comm.deliver(source=1, failed_rank=1) + + _wait_until( + lambda: any(send.message_kind == _ABORT_MESSAGE for send in comm.sends), + description="remote-conflict ABORT relay", + ) + assert comm.abort_calls == 0 + for send in comm.sends: + if send.message_kind == _ABORT_MESSAGE: + send.request.complete() + _wait_until(lambda: comm.receive_count(1) == 2, description="survivor ABORT receive") + comm.deliver(source=1, failed_rank=1, message_kind=_ABORT_MESSAGE) + _wait_until(broadcaster._progress_failed.is_set, description="remote-conflict convergence") + + assert health.get_failed_ranks() == frozenset({2}) + assert isinstance(broadcaster.last_error, RuntimeError) + assert "exactly one distinct failed rank" in str(broadcaster.last_error) + assert comm.abort_calls == 0 + broadcaster.stop() + + +def test_reconciled_terminal_abort_at_deadline_does_not_world_abort() -> None: + comm = _FakeComm(size=1, ulfm_probe_error=NotImplementedError()) + broadcaster, _, _, _ = _make_broadcaster(size=1, comm=comm) + broadcaster._request_terminal_abort(RuntimeError("terminal"), -1, 0) + broadcaster._drain_outbound_reports([]) + with broadcaster._protocol_lock: + broadcaster._abort_deadline = time.monotonic() - 1.0 + + assert broadcaster._progress_terminal_abort() is True + assert comm.abort_calls == 0 + broadcaster.stop() + + +def test_monitor_does_not_world_abort_reconciled_relay_when_mpi_test_resumes() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + abort_timeout_sec=0.2, + ) + blocking_request = _BlockingRequest() + comm = _FakeComm( + size=2, + receive_factory=lambda _source: blocking_request, + ulfm_probe_error=NotImplementedError(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(2), + EPGroupHealth(2), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + _wait_until(blocking_request.entered_test.is_set, description="blocked MPI progress") + broadcaster._request_terminal_abort(RuntimeError("terminal"), -1, 0) + with broadcaster._protocol_lock: + broadcaster._abort_reporters.add(1) + broadcaster._abort_fanout_posted = True + broadcaster._abort_fanout_complete = True + broadcaster._deadline_monitor_wake_event.set() + + _wait_until( + lambda: broadcaster._deadline_monitor_thread is not None + and not broadcaster._deadline_monitor_thread.is_alive(), + description="monitor-observed terminal reconciliation", + ) + assert broadcaster._progress_failed.is_set() is False + assert comm.abort_calls == 0 + + blocking_request.release_test.set() + _wait_until(broadcaster._progress_failed.is_set, description="clean terminal progress stop") + broadcaster.stop() + assert comm.abort_calls == 0 + + +def test_terminal_abort_waits_until_its_relay_completes() -> None: + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster._handle_received_abort(-1, 1) + broadcaster._handle_received_abort(-1, 2) + + assert broadcaster._progress_terminal_abort() is False + pending_sends: list[ep_failure_broadcast._PendingSend] = [] + broadcaster._drain_outbound_reports(pending_sends) + assert {send.destination for send in comm.sends} == {1, 2} + broadcaster._progress_sends(pending_sends) + assert broadcaster._progress_terminal_abort() is False + + for send in comm.sends: + send.request.complete() + broadcaster._progress_sends(pending_sends) + assert pending_sends == [] + assert broadcaster._progress_terminal_abort() is True + assert comm.abort_calls == 0 + broadcaster.stop() + + +def test_received_terminal_abort_is_relayed_without_mutating_health() -> None: + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, health, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="source 1 receive") + comm.deliver(source=1, failed_rank=2, message_kind=_ABORT_MESSAGE) + + _wait_until( + lambda: len([send for send in comm.sends if send.message_kind == _ABORT_MESSAGE]) == 2, + description="relayed terminal ABORT", + ) + for send in comm.sends: + if send.message_kind == _ABORT_MESSAGE: + send.request.complete() + _wait_until(lambda: comm.receive_count(2) == 1, description="source 2 receive") + comm.deliver(source=2, failed_rank=2, message_kind=_ABORT_MESSAGE) + _wait_until(broadcaster._progress_failed.is_set, description="peer ABORT convergence") + + assert health.all_active() is True + assert isinstance(broadcaster.last_error, RuntimeError) + assert broadcaster.world_is_poisoned() is True + assert comm.abort_calls == 0 + broadcaster.stop() + assert broadcaster.world_is_poisoned() is True + + +def test_no_ulfm_abort_timeout_uses_world_abort_before_stop_completes() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + abort_timeout_sec=0.02, + ) + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(3), + EPGroupHealth(3), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + with pytest.raises(RuntimeError, match="exactly one distinct failed rank"): + broadcaster.pre_failover(1) + + # stop_event must not let the progress loop skip the pending ABORT fallback. + broadcaster.stop(timeout=0.5) + assert comm.abort_calls == 1 + assert comm.revoke_calls == 0 + + +def test_stop_cannot_bypass_abort_requested_during_blocked_drain() -> None: + send_entered = threading.Event() + release_send = threading.Event() + + def blocking_send(_destination: int) -> _FakeRequest: + send_entered.set() + release_send.wait() + return _FakeRequest() + + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + abort_timeout_sec=0.02, + ) + comm = _FakeComm( + size=3, + send_factory=blocking_send, + ulfm_probe_error=NotImplementedError(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(3), + EPGroupHealth(3), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="terminal receive") + comm.deliver(source=1, failed_rank=-1, message_kind=_ABORT_MESSAGE) + _wait_until(send_entered.is_set, description="blocked ABORT drain") + + stop_errors: list[BaseException] = [] + + def stop_broadcaster() -> None: + try: + broadcaster.stop(timeout=0.5) + except BaseException as error: + stop_errors.append(error) + + stop_thread = threading.Thread(target=stop_broadcaster) + stop_thread.start() + _wait_until(broadcaster._stop_event.is_set, description="stop during ABORT drain") + release_send.set() + stop_thread.join(timeout=1.0) + + assert stop_thread.is_alive() is False + assert stop_errors == [] + assert comm.abort_calls == 1 + assert broadcaster._progress_failed.is_set() + + +def test_concurrent_local_and_remote_distinct_failures_accept_only_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + broadcaster, health, comm, _ = _make_broadcaster(size=4) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="remote receive") + + original_claim = broadcaster._claim_failure + local_claimed = threading.Event() + release_local_claim = threading.Event() + + def racing_claim(failed_rank: int) -> None: + original_claim(failed_rank) + if failed_rank == 2: + local_claimed.set() + release_local_claim.wait() + + monkeypatch.setattr(broadcaster, "_claim_failure", racing_claim) + local_errors: list[BaseException] = [] + + def report_local_failure() -> None: + try: + broadcaster.pre_failover(2) + except BaseException as error: + local_errors.append(error) + + local_thread = threading.Thread(target=report_local_failure) + local_thread.start() + _wait_until(local_claimed.is_set, description="local failure claim") + comm.deliver(source=1, failed_rank=3) + # The remote handler must wait for the local lifecycle transaction to commit + # its accepted rank and health generation atomically. + time.sleep(2 * _TEST_CONFIG.poll_interval_sec) + assert health.all_active() is True + release_local_claim.set() + local_thread.join(timeout=1.0) + assert local_thread.is_alive() is False + + _wait_until(lambda: broadcaster.last_error is not None, description="second-failure rejection") + _wait_until(lambda: comm.revoke_calls == 1, description="terminal abort revoke") + assert len(health.get_failed_ranks()) == 1 + assert health.get_failed_ranks() == frozenset({2}) + assert "exactly one distinct failed rank" in str(broadcaster.last_error) + assert local_errors == [] + broadcaster.stop() + + +def test_isend_post_error_to_survivor_is_terminal_without_ulfm() -> None: + error = _FakeMpiException(999, "send post failed") + + def fail_send(_destination: int) -> _FakeRequest: + raise error + + comm = _FakeComm( + size=3, + send_factory=fail_send, + ulfm_probe_error=NotImplementedError(), + ) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + + _wait_until(lambda: broadcaster.last_error is error, description="terminal Isend error") + assert broadcaster.health_is_reconciled() is False + assert comm.revoke_calls == 0 + assert comm.sends == () + broadcaster.stop() + + +def test_send_test_error_is_terminal_and_revokes_with_ulfm() -> None: + error = _FakeMpiException(999, "send Test failed") + failed_request = _FakeRequest() + failed_request.fail(error) + comm = _FakeComm(size=3, send_factory=lambda _destination: failed_request) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + + _wait_until(lambda: broadcaster.last_error is error, description="terminal send Test error") + _wait_until(lambda: comm.revoke_calls == 1, description="ULFM emergency revoke") + assert broadcaster.health_is_reconciled() is False + assert failed_request.cancel_calls == 0 + assert broadcaster._retained_requests + broadcaster.stop() + + +def test_received_report_is_idempotent_and_callback_runs_once() -> None: + callbacks: list[tuple[int, int, float]] = [] + broadcaster, health, comm, _ = _make_broadcaster( + size=3, + callback=lambda failed, source, when: callbacks.append((failed, source, when)), + ) + broadcaster.start() + try: + _wait_until(lambda: comm.receive_count(1) == 1, description="initial receive") + comm.deliver(source=1, failed_rank=2) + _wait_until(lambda: health.is_active(2) is False, description="received failure") + _wait_until(lambda: comm.receive_count(1) == 2, description="reposted receive") + + comm.deliver(source=1, failed_rank=2) + _wait_until(lambda: comm.receive_count(1) == 3, description="duplicate receive progress") + _wait_until(lambda: len(callbacks) == 1, description="failure callback") + + assert health.generation == 1 + assert len(callbacks) == 1 + failed_rank, source, event_time = callbacks[0] + assert (failed_rank, source) == (2, 1) + assert isinstance(event_time, float) + _wait_until(lambda: len(comm.sends) == 1, description="failure relay") + comm.sends[0].request.complete() + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + finally: + broadcaster.stop() + + +def test_monitor_cannot_observe_remote_claim_before_health_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + broadcaster, health, comm, _ = _make_broadcaster(size=3) + original_claim = broadcaster._claim_failure + claim_entered = threading.Event() + release_claim = threading.Event() + + def claim_then_pause(failed_rank: int) -> None: + original_claim(failed_rank) + claim_entered.set() + release_claim.wait() + + monkeypatch.setattr(broadcaster, "_claim_failure", claim_then_pause) + broadcaster.start() + try: + comm.deliver(source=1, failed_rank=2) + _wait_until(claim_entered.is_set, description="paused remote failure claim") + time.sleep(3 * _TEST_CONFIG.poll_interval_sec) + + assert health.is_active(2) is True + assert broadcaster.last_error is None + assert comm.revoke_calls == 0 + + release_claim.set() + _wait_until(lambda: not health.is_active(2), description="committed remote failure") + _wait_until(lambda: len(comm.sends) == 1, description="remote failure echo") + comm.sends[0].request.complete() + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + finally: + release_claim.set() + broadcaster.stop() + + +def test_blocking_callback_does_not_block_mpi_progress_and_stop_is_bounded() -> None: + callback_entered = threading.Event() + release_callback = threading.Event() + + def blocking_callback(_failed: int, _source: int, _when: float) -> None: + callback_entered.set() + release_callback.wait() + + broadcaster, health, comm, _ = _make_broadcaster( + size=3, + callback=blocking_callback, + ) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="initial receive") + comm.deliver(source=1, failed_rank=2) + _wait_until(callback_entered.is_set, description="blocking callback") + _wait_until(lambda: len(comm.sends) == 1, description="failure relay") + comm.sends[0].request.complete() + + # The MPI thread must repost and continue polling while user telemetry is blocked. + _wait_until(lambda: comm.receive_count(1) == 2, description="receive repost") + comm.deliver(source=1, failed_rank=2) + _wait_until(lambda: comm.receive_count(1) == 3, description="duplicate receive progress") + assert health.get_failed_ranks() == frozenset({2}) + + start = time.monotonic() + with pytest.raises(TimeoutError, match="callback thread did not stop"): + broadcaster.stop(timeout=0.02) + assert time.monotonic() - start < 0.2 + assert broadcaster.last_error is None + assert broadcaster.world_is_poisoned() is True + + release_callback.set() + broadcaster.stop(timeout=0.5) + + +def test_invalid_received_rank_is_ignored_and_receive_is_reposted() -> None: + callbacks: list[tuple[int, int, float]] = [] + broadcaster, health, comm, _ = _make_broadcaster( + size=3, + callback=lambda failed, source, when: callbacks.append((failed, source, when)), + ) + broadcaster.start() + try: + _wait_until(lambda: comm.receive_count(1) == 1, description="initial receive") + comm.deliver(source=1, failed_rank=99) + _wait_until(lambda: comm.receive_count(1) == 2, description="receive after invalid payload") + + assert health.all_active() is True + assert callbacks == [] + assert broadcaster.last_error is None + finally: + broadcaster.stop() + + +def test_single_rank_group_has_no_peer_requests_or_sends() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.02, + abort_timeout_sec=0.02, + ) + comm = _FakeComm(size=1, ulfm_probe_error=NotImplementedError()) + health = EPGroupHealth(1) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(1), + health, + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + try: + assert broadcaster.broadcast_failure(0) is True + _wait_until(broadcaster._outbound_reports.empty, description="empty single-rank report") + assert health.get_failed_ranks() == frozenset({0}) + assert broadcaster.world_is_poisoned() is True + assert broadcaster.failure_is_reconciled(0) is False + assert broadcaster.health_is_reconciled() is False + assert comm.sends == () + assert comm.receive_count(0) == 0 + _wait_until( + lambda: isinstance(broadcaster.last_error, TimeoutError), + description="empty-survivor terminal handling", + ) + finally: + broadcaster.stop() + + +def test_lifecycle_is_not_restartable_and_stop_before_start_is_terminal() -> None: + stopped, _, _, _ = _make_broadcaster(size=1) + stopped.stop() + stopped.stop() + with pytest.raises(RuntimeError, match="cannot start from stopped"): + stopped.start() + + running, _, _, _ = _make_broadcaster(size=1) + running.start() + try: + with pytest.raises(RuntimeError, match="cannot start from running"): + running.start() + finally: + running.stop() + with pytest.raises(RuntimeError, match="cannot start from stopped"): + running.start() + + +@pytest.mark.parametrize( + "failing_thread_name", + [ + "wide-ep-ft-callback", + "wide-ep-ft-deadline-monitor", + "wide-ep-ft-broadcast", + ], + ids=["callback", "deadline-monitor", "progress"], +) +def test_thread_start_failure_is_terminal_and_stoppable( + monkeypatch: pytest.MonkeyPatch, + failing_thread_name: str, +) -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + original_start = threading.Thread.start + start_error = OSError(f"cannot start {failing_thread_name}") + + def failing_start(thread: threading.Thread) -> None: + if thread.name == failing_thread_name: + raise start_error + original_start(thread) + + monkeypatch.setattr(threading.Thread, "start", failing_start) + comm = _FakeComm(size=2) + broadcaster, _, _, _ = _make_broadcaster( + size=2, + comm=comm, + callback=lambda _failed, _source, _when: None, + ) + + try: + with pytest.raises(RuntimeError, match="thread failed to start") as raised: + broadcaster.start() + + assert raised.value.__cause__ is start_error + assert broadcaster.last_error is start_error + assert broadcaster._progress_failed.is_set() + assert broadcaster._thread is None + assert comm.revoke_calls == 1 + if failing_thread_name == "wide-ep-ft-callback": + assert broadcaster._callback_thread is None + assert broadcaster._deadline_monitor_thread is None + elif failing_thread_name == "wide-ep-ft-deadline-monitor": + assert broadcaster._callback_thread is not None + assert broadcaster._callback_thread.is_alive() is False + assert broadcaster._deadline_monitor_thread is None + else: + assert broadcaster._callback_thread is not None + assert broadcaster._callback_thread.is_alive() is False + assert broadcaster._deadline_monitor_thread is not None + assert broadcaster._deadline_monitor_thread.is_alive() is False + + # A failed start must not leave an unstarted Thread object that stop() + # later attempts to join. + broadcaster.stop() + finally: + del registry[initial_registry_size:] + + +def test_concurrent_stop_prevents_start_from_publishing_running() -> None: + receive_entered = threading.Event() + release_receive = threading.Event() + + def blocking_receive(_source: int) -> _FakeRequest: + receive_entered.set() + release_receive.wait() + return _FakeRequest() + + comm = _FakeComm(size=2, receive_factory=blocking_receive) + broadcaster, _, _, _ = _make_broadcaster(size=2, comm=comm) + start_errors: list[BaseException] = [] + stop_errors: list[BaseException] = [] + + def start_broadcaster() -> None: + try: + broadcaster.start() + except BaseException as error: + start_errors.append(error) + + def stop_broadcaster() -> None: + try: + broadcaster.stop(timeout=0.5) + except BaseException as error: + stop_errors.append(error) + + start_thread = threading.Thread(target=start_broadcaster) + start_thread.start() + _wait_until(receive_entered.is_set, description="blocked startup receive") + stop_thread = threading.Thread(target=stop_broadcaster) + stop_thread.start() + _wait_until(broadcaster._stop_event.is_set, description="concurrent stop") + release_receive.set() + + start_thread.join(timeout=1.0) + stop_thread.join(timeout=1.0) + assert start_thread.is_alive() is False + assert stop_thread.is_alive() is False + assert len(start_errors) == 1 + assert "did not reach running state" in str(start_errors[0]) + assert stop_errors == [] + assert broadcaster.health_is_reconciled() is False + + +def test_startup_failure_publishes_error_before_failed_flag( + monkeypatch: pytest.MonkeyPatch, +) -> None: + comm = _FakeComm(size=1) + broadcaster, _, _, _ = _make_broadcaster(size=1, comm=comm) + error = _FakeMpiException(999, "forced startup failure") + record_entered = threading.Event() + release_record = threading.Event() + original_record_error = broadcaster._record_error + + def blocked_record_error(recorded_error: BaseException) -> None: + record_entered.set() + release_record.wait() + original_record_error(recorded_error) + + def failing_progress_loop() -> None: + broadcaster._fail_closed_immediately(error) + broadcaster._ready_event.set() + + monkeypatch.setattr(broadcaster, "_record_error", blocked_record_error) + monkeypatch.setattr(broadcaster, "_progress_loop", failing_progress_loop) + start_errors: list[BaseException] = [] + + def start_broadcaster() -> None: + try: + broadcaster.start() + except BaseException as start_error: + start_errors.append(start_error) + + start_thread = threading.Thread(target=start_broadcaster) + start_thread.start() + _wait_until(record_entered.is_set, description="blocked terminal publication") + assert broadcaster._progress_failed.is_set() is False + assert broadcaster.last_error is None + release_record.set() + start_thread.join(timeout=1.0) + + assert start_thread.is_alive() is False + assert len(start_errors) == 1 + assert "failed during startup" in str(start_errors[0]) + assert broadcaster.last_error is error + assert broadcaster._progress_failed.is_set() + assert comm.revoke_calls == 1 + + +def test_startup_timeout_fails_closed_before_progress_resumes() -> None: + receive_entered = threading.Event() + release_receive = threading.Event() + + def blocking_receive(_source: int) -> _FakeRequest: + receive_entered.set() + release_receive.wait() + return _FakeRequest() + + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.02, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + abort_timeout_sec=0.02, + ) + comm = _FakeComm( + size=2, + receive_factory=blocking_receive, + ulfm_probe_error=NotImplementedError(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(2), + EPGroupHealth(2), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + start_errors: list[BaseException] = [] + + def start_broadcaster() -> None: + try: + broadcaster.start() + except BaseException as error: + start_errors.append(error) + + start_thread = threading.Thread(target=start_broadcaster) + start_thread.start() + _wait_until(receive_entered.is_set, description="blocked startup") + start_thread.join(timeout=1.0) + + assert start_thread.is_alive() is False + assert len(start_errors) == 1 + assert isinstance(start_errors[0], TimeoutError) + assert "did not start" in str(start_errors[0]) + assert broadcaster._abort_requested.is_set() is False + assert isinstance(broadcaster.last_error, TimeoutError) + # The progress thread is still blocked in Irecv, so the caller thread must + # execute the no-ULFM fail-stop fallback before start() returns. + assert release_receive.is_set() is False + assert comm.abort_calls == 1 + release_receive.set() + broadcaster.stop(timeout=0.5) + assert comm.abort_calls == 1 + + +def test_startup_timeout_stops_callback_before_blocked_receive_is_released() -> None: + receive_entered = threading.Event() + release_receive = threading.Event() + + def blocking_receive(_source: int) -> _FakeRequest: + receive_entered.set() + release_receive.wait() + return _FakeRequest() + + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.02, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + unattributed_error_timeout_sec=0.2, + abort_timeout_sec=0.02, + ) + comm = _FakeComm( + size=2, + receive_factory=blocking_receive, + ulfm_probe_error=NotImplementedError(), + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(2), + EPGroupHealth(2), + config, + lambda _failed, _source, _when: None, + comm=comm, + mpi_module=_FakeMPI(), + ) + + try: + with pytest.raises(TimeoutError, match="did not start"): + broadcaster.start() + assert receive_entered.is_set() + assert release_receive.is_set() is False + assert broadcaster._callback_thread is not None + assert broadcaster._callback_thread.is_alive() is False + assert comm.abort_calls == 1 + finally: + release_receive.set() + broadcaster.stop(timeout=0.5) + assert comm.abort_calls == 1 + + +def test_reconciliation_gates_fail_closed_outside_running_lifecycle() -> None: + broadcaster, _, _, _ = _make_broadcaster(size=2) + assert broadcaster.health_is_reconciled() is False + assert broadcaster.failure_is_reconciled(1) is False + + broadcaster.start() + assert broadcaster.pre_failover(1) is True + _wait_until( + lambda: broadcaster.failure_is_reconciled(1), + description="single-survivor reconciliation", + ) + assert broadcaster.health_is_reconciled() is True + + broadcaster.stop() + assert broadcaster.failure_is_reconciled(1) is False + assert broadcaster.health_is_reconciled() is False + + +@pytest.mark.parametrize("timeout", [0.0, -1.0, float("nan"), float("inf"), float("-inf")]) +def test_stop_requires_positive_finite_timeout(timeout: float) -> None: + broadcaster, _, _, _ = _make_broadcaster(size=1) + with pytest.raises(ValueError, match="finite and > 0"): + broadcaster.stop(timeout) + broadcaster.start() + broadcaster.stop() + + +def test_stop_is_bounded_when_mpi_test_is_stuck() -> None: + blocking_request = _BlockingRequest() + comm = _FakeComm(size=2, receive_factory=lambda _source: blocking_request) + broadcaster, _, _, _ = _make_broadcaster(size=2, comm=comm) + broadcaster.start() + _wait_until(blocking_request.entered_test.is_set, description="blocked MPI Test") + + start = time.monotonic() + with pytest.raises(TimeoutError, match="did not stop"): + broadcaster.stop(timeout=0.01) + assert time.monotonic() - start < 0.2 + assert blocking_request.release_test.is_set() is False + assert comm.revoke_calls == 1 + assert comm.abort_calls == 0 + + blocking_request.release_test.set() + broadcaster.stop(timeout=0.5) + assert comm.revoke_calls == 1 + assert comm.abort_calls == 0 + + +def test_broadcast_rejects_report_after_stop_begins() -> None: + blocking_request = _BlockingRequest() + comm = _FakeComm(size=2, receive_factory=lambda _source: blocking_request) + broadcaster, health, _, _ = _make_broadcaster(size=2, comm=comm) + broadcaster.start() + _wait_until(blocking_request.entered_test.is_set, description="blocked MPI Test") + + stop_errors: list[BaseException] = [] + + def stop_broadcaster() -> None: + try: + broadcaster.stop(timeout=0.5) + except BaseException as error: + stop_errors.append(error) + + stop_thread = threading.Thread(target=stop_broadcaster) + stop_thread.start() + _wait_until(broadcaster._stop_event.is_set, description="stop transition") + + with pytest.raises(RuntimeError, match="state=stopping"): + broadcaster.broadcast_failure(1) + assert health.all_active() is True + + blocking_request.release_test.set() + stop_thread.join(timeout=1.0) + assert stop_thread.is_alive() is False + assert stop_errors == [] + + +def test_healthy_stop_cancels_preposted_receives() -> None: + broadcaster, health, comm, _ = _make_broadcaster(size=3) + broadcaster.start() + _wait_until( + lambda: comm.receive_count(1) == 1 and comm.receive_count(2) == 1, + description="preposted receives", + ) + receives = [comm.latest_receive(source).request for source in (1, 2)] + + assert health.all_active() is True + broadcaster.stop() + + assert [request.cancel_calls for request in receives] == [1, 1] + assert all(request.Test() is True for request in receives) + assert broadcaster._retained_requests == [] + + +def test_healthy_stop_cleanup_timeout_aborts_world_and_retains_request() -> None: + request = _NonCompletingCancelRequest() + comm = _FakeComm(size=2, receive_factory=lambda _source: request) + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.05, + reconcile_timeout_sec=0.2, + ) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(2), + EPGroupHealth(2), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="preposted receive") + + start = time.monotonic() + broadcaster.stop() + assert time.monotonic() - start < 0.2 + assert request.cancel_calls == 1 + assert isinstance(broadcaster.last_error, TimeoutError) + assert broadcaster._retained_requests + assert comm.abort_calls == 1 + assert comm.revoke_calls == 0 + + +def test_poisoned_stop_keeps_progressing_until_failure_is_reconciled() -> None: + broadcaster, health, comm, _ = _make_broadcaster(size=3) + broadcaster.start() + _wait_until( + lambda: comm.receive_count(1) == 1 and comm.receive_count(2) == 1, + description="preposted receives", + ) + receives = [comm.latest_receive(source).request for source in (1, 2)] + assert broadcaster.broadcast_failure(2) is True + _wait_until(lambda: len(comm.sends) == 1, description="pending failure send") + send = comm.sends[0] + + assert health.get_failed_ranks() == frozenset({2}) + initial_test_calls = send.request.test_calls + stop_errors: list[BaseException] = [] + + def stop_broadcaster() -> None: + try: + broadcaster.stop() + except BaseException as error: + stop_errors.append(error) + + stop_thread = threading.Thread(target=stop_broadcaster) + stop_thread.start() + _wait_until(broadcaster._stop_event.is_set, description="stop request") + _wait_until( + lambda: send.request.test_calls > initial_test_calls, + description="continued send progress during stop", + ) + assert stop_thread.is_alive() is True + + send.request.complete() + comm.deliver(source=1, failed_rank=2) + stop_thread.join(timeout=1.0) + + assert stop_thread.is_alive() is False + assert stop_errors == [] + assert [request.cancel_calls for request in receives] == [0, 0] + assert send.request.cancel_calls == 0 + assert broadcaster.last_error is None + assert broadcaster._retained_requests + + +def test_poisoned_resources_outlive_broadcaster_object() -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + + def retain_then_drop_owner() -> tuple[ + weakref.ReferenceType[_FakeRequest], + weakref.ReferenceType[np.ndarray], + weakref.ReferenceType[_FakeComm], + ]: + broadcaster, _, comm, _ = _make_broadcaster(size=3) + broadcaster.start() + _wait_until( + lambda: comm.receive_count(1) == 1 and comm.receive_count(2) == 1, + description="preposted receives", + ) + assert broadcaster.pre_failover(2) is True + _wait_until(lambda: len(comm.sends) == 1, description="pending send") + send = comm.sends[0] + send.request.complete() + comm.deliver(source=1, failed_rank=2) + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + retained_receive = comm.latest_receive(2) + references = ( + weakref.ref(retained_receive.request), + weakref.ref(retained_receive.buffer), + weakref.ref(comm), + ) + broadcaster.stop() + # The fake communicator owns test bookkeeping references that a real + # MPI communicator does not. Drop those so only the process-lifetime + # registry can keep the poisoned request and buffer reachable. + comm._receives.clear() + return references + + request_ref, buffer_ref, comm_ref = retain_then_drop_owner() + gc.collect() + try: + assert request_ref() is not None + assert buffer_ref() is not None + assert comm_ref() is not None + assert len(registry) > initial_registry_size + finally: + del registry[initial_registry_size:] + + +def test_startup_failure_retains_poisoned_comm_without_pending_requests() -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + + def fail_first_receive(_source: int) -> _FakeRequest: + raise _FakeMpiException(999, "first Irecv failed") + + def start_then_drop_owner() -> weakref.ReferenceType[_FakeComm]: + comm = _FakeComm(size=2, receive_factory=fail_first_receive) + comm_ref = weakref.ref(comm) + broadcaster, _, _, _ = _make_broadcaster(size=2, comm=comm) + with pytest.raises(RuntimeError, match="failed during startup"): + broadcaster.start() + assert broadcaster._retained_requests == [] + assert isinstance(broadcaster.last_error, _FakeMpiException) + return comm_ref + + comm_ref = start_then_drop_owner() + gc.collect() + try: + assert comm_ref() is not None + assert len(registry) > initial_registry_size + assert any(reference is comm_ref() for reference in registry[initial_registry_size:]) + finally: + del registry[initial_registry_size:] + + +def test_constructor_requires_mpi_thread_multiple() -> None: + mpi = _FakeMPI(thread_level=_FakeMPI.THREAD_MULTIPLE - 1) + comm = _FakeComm(size=2) + with pytest.raises(RuntimeError, match="MPI.THREAD_MULTIPLE"): + MpiFtSubcomm( + _FakeMapping.ep_world(2), + EPGroupHealth(2), + _TEST_CONFIG, + comm=comm, + mpi_module=mpi, + ) + assert comm.errhandler is None + + +def test_collectively_created_comm_uses_frozen_setup_and_lives_for_process( + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + comm = _FakeComm(size=2) + mpi = _FakeMPI() + + class SingleReadMapping: + def __init__(self) -> None: + self.rank_reads = 0 + + @property + def moe_ep_rank(self) -> int: + self.rank_reads += 1 + if self.rank_reads > 1: + raise RuntimeError("mapping rank was read after collective validation") + return 0 + + mapping = SingleReadMapping() + + def create_mpi_ft_subcomm( + setup_mapping: SingleReadMapping, + *, + health_size: int, + ) -> types.SimpleNamespace: + assert health_size == 2 + local_rank = setup_mapping.moe_ep_rank + comm.Set_errhandler(mpi.ERRORS_RETURN) + return types.SimpleNamespace( + comm=comm, + local_rank=local_rank, + ep_size=2, + ulfm_available=False, + ) + + distributed = types.ModuleType("tensorrt_llm._torch.distributed") + distributed.__path__ = [] + communicator = types.ModuleType("tensorrt_llm._torch.distributed.communicator") + communicator.create_mpi_ft_subcomm = create_mpi_ft_subcomm + distributed.communicator = communicator + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.distributed", distributed) + monkeypatch.setitem( + sys.modules, + "tensorrt_llm._torch.distributed.communicator", + communicator, + ) + + try: + broadcaster = MpiFtSubcomm( + mapping, + EPGroupHealth(2), + _TEST_CONFIG, + mpi_module=mpi, + ) + + assert broadcaster._comm is comm + assert broadcaster._local_rank == 0 + assert broadcaster._ep_size == 2 + assert mapping.rank_reads == 1 + assert comm.errhandler_calls == 1 + assert comm.is_revoked_calls == 0 + assert any(reference is comm for reference in registry[initial_registry_size:]) + broadcaster.stop() + del broadcaster + gc.collect() + assert any(reference is comm for reference in registry[initial_registry_size:]) + finally: + del registry[initial_registry_size:] + + +def test_constructor_rejects_non_world_spanning_ep_group_for_mvp() -> None: + mapping = _FakeMapping( + world_size=4, + rank=0, + moe_ep_size=2, + moe_ep_rank=0, + moe_ep_group=(0, 1), + ) + comm = _FakeComm(size=2) + with pytest.raises(ValueError, match="spanning the full MPI world"): + MpiFtSubcomm( + mapping, + EPGroupHealth(2), + _TEST_CONFIG, + comm=comm, + mpi_module=_FakeMPI(), + ) + assert comm.errhandler is None + + +@pytest.mark.parametrize( + "probe_error", + [ + NotImplementedError("ULFM not compiled"), + _FakeMpiException(_FakeMPI.ERR_UNSUPPORTED_OPERATION), + ], + ids=["not-implemented", "unsupported-operation"], +) +def test_ulfm_probe_falls_back_when_runtime_does_not_support_it( + probe_error: BaseException, +) -> None: + comm = _FakeComm(size=2, ulfm_probe_error=probe_error) + broadcaster, _, _, _ = _make_broadcaster(size=2, comm=comm) + assert broadcaster.ulfm_available is False + assert comm.is_revoked_calls == 1 + broadcaster.stop() + + +def test_ulfm_probe_propagates_unrelated_mpi_error() -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + error = _FakeMpiException(_FakeMPI.ERR_UNKNOWN, "corrupt communicator") + comm = _FakeComm(size=2, ulfm_probe_error=error) + try: + with pytest.raises(_FakeMpiException, match="corrupt communicator") as raised: + _make_broadcaster(size=2, comm=comm) + assert raised.value is error + assert any(reference is comm for reference in registry[initial_registry_size:]) + finally: + del registry[initial_registry_size:] + + +def test_ulfm_probe_retains_original_error_when_error_class_inspection_raises() -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + + class UnclassifiableMpiError(_FakeMpiException): + def Get_error_class(self) -> int: + raise RuntimeError("error-class lookup failed") + + probe_error = UnclassifiableMpiError( + _FakeMPI.ERR_UNSUPPORTED_OPERATION, + "unclassifiable ULFM probe failure", + ) + comm = _FakeComm(size=2, ulfm_probe_error=probe_error) + try: + with pytest.raises(UnclassifiableMpiError, match="unclassifiable") as raised: + _make_broadcaster(size=2, comm=comm) + assert raised.value is probe_error + assert any(reference is comm for reference in registry[initial_registry_size:]) + finally: + del registry[initial_registry_size:] + + +def test_constructor_rejects_and_retains_already_revoked_comm() -> None: + registry = ep_failure_broadcast._PROCESS_LIFETIME_REFS + initial_registry_size = len(registry) + + def construct_then_drop_owner() -> weakref.ReferenceType[_FakeComm]: + comm = _FakeComm(size=2, ulfm_revoked=True) + comm_ref = weakref.ref(comm) + with pytest.raises(RuntimeError, match="already revoked"): + _make_broadcaster(size=2, comm=comm) + return comm_ref + + comm_ref = construct_then_drop_owner() + gc.collect() + try: + assert comm_ref() is not None + assert len(registry) > initial_registry_size + finally: + del registry[initial_registry_size:] + + +def test_peer_failure_reconciles_without_revoking_control_plane() -> None: + callbacks: list[tuple[int, int, float]] = [] + broadcaster, health, comm, _ = _make_broadcaster( + size=3, + callback=lambda failed, source, when: callbacks.append((failed, source, when)), + ) + assert broadcaster.ulfm_available is True + broadcaster.start() + try: + _wait_until(lambda: comm.receive_count(1) == 1, description="peer receive") + comm.latest_receive(1).request.fail(_FakeMpiException(_FakeMPI.ERR_PROC_FAILED)) + + _wait_until(lambda: health.is_active(1) is False, description="peer failure classification") + _wait_until(lambda: len(comm.sends) == 1, description="peer failure relay") + + assert health.get_failed_ranks() == frozenset({1}) + _wait_until(lambda: len(callbacks) == 1, description="failed-peer callback") + assert [(failed, source) for failed, source, _ in callbacks] == [(1, 1)] + assert comm.sends[0].destination == 2 + assert comm.sends[0].message_kind == _FAILURE_MESSAGE + assert comm.sends[0].payload == 1 + assert comm.revoke_calls == 0 + assert broadcaster.last_error is None + + comm.sends[0].request.complete() + comm.deliver(source=2, failed_rank=1) + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + _wait_until(lambda: comm.receive_count(2) == 2, description="reposted receive") + assert broadcaster.health_is_reconciled() is True + assert broadcaster.last_error is None + assert comm.revoke_calls == 0 + finally: + broadcaster.stop() + + +def test_remote_revoke_after_local_reconciliation_still_fails_closed() -> None: + error = _FakeMpiException(_FakeMPI.ERR_REVOKED, "remote revoke") + comm = _FakeComm(size=3) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + _wait_until(lambda: len(comm.sends) == 1, description="failure relay") + comm.sends[0].request.complete() + _wait_until(lambda: comm.receive_count(1) == 1, description="initial peer receive") + comm.deliver(source=1, failed_rank=2) + _wait_until(broadcaster.health_is_reconciled, description="local reconciliation") + _wait_until( + lambda: comm.receive_count(1) == 2, + description="receive after failure report", + ) + + comm.latest_receive(1).request.fail(error) + _wait_until(lambda: broadcaster.last_error is error, description="terminal remote revoke") + assert broadcaster.health_is_reconciled() is False + assert comm.revoke_calls == 1 + with pytest.raises(RuntimeError, match="not running"): + broadcaster.pre_failover(1) + assert broadcaster._health.get_failed_ranks() == frozenset({2}) + broadcaster.stop() + + +def test_premature_remote_revoke_fails_closed() -> None: + error = _FakeMpiException(_FakeMPI.ERR_REVOKED, "premature revoke") + comm = _FakeComm(size=3) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + assert broadcaster.pre_failover(2) is True + _wait_until(lambda: comm.receive_count(1) == 1, description="initial peer receive") + + comm.latest_receive(1).request.fail(error) + _wait_until(lambda: broadcaster.last_error is error, description="terminal revoke") + assert broadcaster.health_is_reconciled() is False + assert comm.revoke_calls == 1 + broadcaster.stop() + + +def test_generic_mpi_error_is_terminal_and_revokes_once() -> None: + comm = _FakeComm(size=2) + broadcaster, health, _, _ = _make_broadcaster(size=2, comm=comm) + broadcaster.start() + try: + _wait_until(lambda: comm.receive_count(1) == 1, description="peer receive") + error = _FakeMpiException(999, "generic transport failure") + comm.latest_receive(1).request.fail(error) + + _wait_until(lambda: broadcaster.last_error is error, description="terminal MPI error") + _wait_until(lambda: comm.revoke_calls == 1, description="communicator revoke") + assert health.all_active() is True + with pytest.raises(RuntimeError, match="broadcaster is not running"): + broadcaster.broadcast_failure(1) + finally: + broadcaster.stop() + assert comm.revoke_calls == 1 + + +def test_proc_failed_is_classified_without_full_ulfm_support() -> None: + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, health, _, _ = _make_broadcaster(size=3, comm=comm) + assert broadcaster.ulfm_available is False + broadcaster.start() + try: + _wait_until(lambda: comm.receive_count(1) == 1, description="peer receive") + comm.latest_receive(1).request.fail(_FakeMpiException(_FakeMPI.ERR_PROC_FAILED)) + _wait_until(lambda: health.is_active(1) is False, description="peer failure classification") + _wait_until(lambda: len(comm.sends) == 1, description="peer failure relay") + + assert health.get_failed_ranks() == frozenset({1}) + assert comm.sends[0].destination == 2 + assert comm.sends[0].payload == 1 + assert broadcaster.last_error is None + assert comm.revoke_calls == 0 + comm.sends[0].request.complete() + comm.deliver(source=2, failed_rank=1) + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + finally: + broadcaster.stop() + + +def test_non_ulfm_generic_receive_error_keeps_other_sources_live() -> None: + callbacks: list[tuple[int, int, float]] = [] + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster, health, _, _ = _make_broadcaster( + size=3, + comm=comm, + callback=lambda failed, source, when: callbacks.append((failed, source, when)), + ) + broadcaster.start() + try: + _wait_until( + lambda: comm.receive_count(1) == 1 and comm.receive_count(2) == 1, + description="fixed-source receives", + ) + failed_receive = comm.latest_receive(1).request + failed_receive.fail(_FakeMpiException(999, "implementation-specific peer error")) + _wait_until( + lambda: len(broadcaster._retained_requests) == 1, + description="retained failed receive", + ) + + assert health.all_active() is True + assert broadcaster.last_error is None + assert failed_receive.cancel_calls == 0 + + # Another survivor's watchdog report remains authoritative even when + # the dead source's fixed receive failed with a generic error. + comm.deliver(source=2, failed_rank=1) + _wait_until(lambda: health.is_active(1) is False, description="authoritative report") + _wait_until(lambda: len(comm.sends) == 1, description="failure relay") + comm.sends[0].request.complete() + _wait_until(lambda: comm.receive_count(2) == 2, description="live-source receive repost") + _wait_until(lambda: len(callbacks) == 1, description="authoritative callback") + + assert [(failed, source) for failed, source, _ in callbacks] == [(1, 2)] + assert broadcaster.last_error is None + assert comm.revoke_calls == 0 + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + finally: + broadcaster.stop() + + +def test_non_ulfm_generic_error_after_watchdog_report_does_not_arm_stale_deadline() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + unattributed_error_timeout_sec=0.01, + abort_timeout_sec=0.02, + ) + comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(3), + EPGroupHealth(3), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + try: + _wait_until( + lambda: comm.receive_count(1) == 1 and comm.receive_count(2) == 1, + description="fixed-source receives", + ) + comm.deliver(source=2, failed_rank=1) + _wait_until(lambda: len(comm.sends) == 1, description="failure relay") + + # The transport notification can arrive after the watchdog report. It + # is already attributed to rank 1 and must not create a deadline that + # aborts an otherwise reconciled survivor set later. + failed_receive = comm.latest_receive(1).request + failed_receive.fail(_FakeMpiException(999, "late generic peer error")) + _wait_until( + lambda: failed_receive + in [pending.request for pending in broadcaster._retained_requests], + description="retained explained receive", + ) + comm.sends[0].request.complete() + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + time.sleep(3 * config.unattributed_error_timeout_sec) + + assert broadcaster.last_error is None + assert broadcaster.health_is_reconciled() is True + assert comm.abort_calls == 0 + finally: + broadcaster.stop() + + +def test_unattributed_error_from_other_source_keeps_reconciliation_gate_closed() -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.5, + unattributed_error_timeout_sec=0.05, + abort_timeout_sec=0.02, + ) + comm = _FakeComm(size=4, ulfm_probe_error=NotImplementedError()) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(4), + EPGroupHealth(4), + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + broadcaster.start() + try: + _wait_until( + lambda: all(comm.receive_count(source) == 1 for source in (1, 2, 3)), + description="fixed-source receives", + ) + comm.deliver(source=1, failed_rank=3) + _wait_until(lambda: comm.receive_count(1) == 2, description="source 1 receive repost") + _wait_until(lambda: len(comm.sends) == 2, description="failure fanout") + + # Source 1 already reported rank 3, but a later generic error from + # source 1 is not explained by that different failed rank. + comm.latest_receive(1).request.fail(_FakeMpiException(999, "second transport error")) + _wait_until( + lambda: bool(broadcaster._unattributed_error_deadlines), + description="unattributed deadline", + ) + comm.deliver(source=2, failed_rank=3) + for send in comm.sends: + send.request.complete() + _wait_until(lambda: comm.receive_count(2) == 2, description="source 2 report") + _wait_until( + lambda: 3 in broadcaster._failure_fanout_complete, + description="completed failure fanout", + ) + + assert broadcaster.failure_is_reconciled(3) is False + assert broadcaster.health_is_reconciled() is False + _wait_until( + lambda: isinstance(broadcaster.last_error, TimeoutError), + description="unattributed transport deadline", + ) + _wait_until(lambda: comm.abort_calls == 1, description="terminal fallback") + finally: + broadcaster.stop() + + +def test_non_ulfm_unattributed_receive_error_aborts_after_bounded_deadlines( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = MpiFtSubcommConfig( + poll_interval_sec=0.001, + startup_timeout_sec=0.5, + stop_timeout_sec=0.5, + reconcile_timeout_sec=0.2, + unattributed_error_timeout_sec=0.01, + abort_timeout_sec=0.05, + ) + comm = _FakeComm(size=2, ulfm_probe_error=NotImplementedError()) + health = EPGroupHealth(2) + broadcaster = MpiFtSubcomm( + _FakeMapping.ep_world(2), + health, + config, + comm=comm, + mpi_module=_FakeMPI(), + ) + terminal_request_entered = threading.Event() + release_terminal_request = threading.Event() + original_request_terminal_abort = broadcaster._request_terminal_abort + + def blocking_terminal_request( + error: BaseException, + payload: int, + reporter: int, + ) -> None: + terminal_request_entered.set() + assert release_terminal_request.wait(timeout=1.0) + original_request_terminal_abort(error, payload, reporter) + + monkeypatch.setattr(broadcaster, "_request_terminal_abort", blocking_terminal_request) + broadcaster.start() + _wait_until(lambda: comm.receive_count(1) == 1, description="fixed-source receive") + failed_receive = comm.latest_receive(1).request + failed_receive.fail(_FakeMpiException(999, "unattributed peer error")) + + _wait_until(terminal_request_entered.is_set, description="blocked terminal publication") + # The request is deliberately blocked before last_error is recorded. The + # lifecycle and protocol poison must already make the public gate fail + # closed, with no lock held that prevents the iteration thread from polling. + try: + assert broadcaster.last_error is None + assert broadcaster._lifecycle is ep_failure_broadcast._Lifecycle.FAILED + assert broadcaster._unattributed_error_deadlines[1] == float("inf") + assert broadcaster.health_is_reconciled() is False + finally: + release_terminal_request.set() + + _wait_until( + lambda: isinstance(broadcaster.last_error, TimeoutError), + description="unattributed transport deadline", + ) + assert "could not attribute" in str(broadcaster.last_error) + # Expiry remains protocol-visible for the rest of this communicator epoch. + assert 1 in broadcaster._unattributed_error_deadlines + assert broadcaster.health_is_reconciled() is False + _wait_until( + lambda: any(send.message_kind == _ABORT_MESSAGE for send in comm.sends), + description="terminal ABORT relay", + ) + _wait_until(lambda: comm.abort_calls == 1, description="bounded MPI_Abort fallback") + + assert health.all_active() is True + assert failed_receive.cancel_calls == 0 + assert comm.revoke_calls == 0 + broadcaster.stop() + + +def test_send_buffer_lives_until_request_test_completes() -> None: + comm = _FakeComm(size=3) + broadcaster, _, _, _ = _make_broadcaster(size=3, comm=comm) + broadcaster.start() + try: + assert broadcaster.broadcast_failure(2) is True + _wait_until(lambda: len(comm.sends) == 1, description="pending send") + record = comm.sends[0] + _wait_until(lambda: record.request.test_calls > 0, description="send Test polling") + + gc.collect() + assert record.buffer_ref() is not None + + record.request.complete() + + def buffer_released() -> bool: + gc.collect() + return record.buffer_ref() is None + + _wait_until(buffer_released, description="completed send buffer release") + comm.deliver(source=1, failed_rank=2) + _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + finally: + broadcaster.stop() diff --git a/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast_mpi.py b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast_mpi.py new file mode 100644 index 000000000000..d158fc8f7b29 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast_mpi.py @@ -0,0 +1,687 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Launcher for logical WideEP failure injection over a real MPI transport.""" + +import importlib.util +import math +import os +import re +import shutil +import signal +import subprocess # nosec B404 +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Protocol + +import pytest + +_SKIP_MARKER = "WIDEEP_MPI_SMOKE_SKIP:" +_PROPAGATION_MARKER = "WIDEEP_MPI_PROPAGATION" +_INVALID_TOPOLOGY_MARKER = "WIDEEP_MPI_INVALID_TOPOLOGY_OK" +_HEALTHY_LIFECYCLE_MARKER = "WIDEEP_MPI_HEALTHY_LIFECYCLE_OK" +_ABORT_WORLD_MARKER = "WIDEEP_MPI_ABORT_WORLD_OK" +_TERMINAL_READY_MARKER = "WIDEEP_MPI_TERMINAL_READY" +_TERMINAL_COMPLETE_MARKER = "WIDEEP_MPI_TERMINAL_COMPLETE" +_TERMINAL_DONE_MARKER = "WIDEEP_MPI_TERMINAL_DONE" +_TERMINAL_ERROR_MARKER = "WIDEEP_MPI_TERMINAL_ERROR" +_TERMINAL_ATEXIT_MARKER = "WIDEEP_MPI_TERMINAL_ATEXIT_RAN" +_WORKER_TIMEOUT_SEC = 30.0 +_REAP_TIMEOUT_SEC = 5.0 +_ALLOW_SKIP_ENV = "TLLM_ALLOW_MPI_FT_SMOKE_SKIP" +_HEALTHY_MODE = "healthy" +_TERMINAL_MODE = "terminal" +_PROPAGATION_PATTERN = re.compile( + rf"^{_PROPAGATION_MARKER} world_size=(\d+) " + r"elapsed_ms=(\S+) target_ms=(\S+) target_met=(True|False)$", + re.MULTILINE, +) +_TERMINAL_READY_PATTERN = re.compile(rf"^{_TERMINAL_READY_MARKER} rank=(\d+) world_size=(\d+)$") +_TERMINAL_COMPLETE_PATTERN = re.compile(rf"^{_TERMINAL_COMPLETE_MARKER} world_size=(\d+)$") +_TERMINAL_DONE_PATTERN = re.compile(rf"^{_TERMINAL_DONE_MARKER} rank=(\d+) world_size=(\d+)$") + + +class _ReapableProcess(Protocol): + def kill(self) -> None: ... + + def wait(self, timeout: float | None = None) -> int: ... + + +def _mpi_launcher() -> str | None: + return shutil.which("mpiexec") or shutil.which("mpirun") + + +def _smoke_is_required(environment: Mapping[str, str] | None = None) -> bool: + """Return whether missing MPI prerequisites must fail this invocation.""" + environment = os.environ if environment is None else environment + return environment.get(_ALLOW_SKIP_ENV) != "1" + + +def _handle_missing_prerequisite(reason: str, *, required: bool) -> None: + if required: + pytest.fail( + f"WideEP MPI FT smoke is required but cannot run: {reason}", + pytrace=False, + ) + pytest.skip(reason) + + +def _subprocess_output_text(output: str | bytes | None) -> str: + if output is None: + return "" + if isinstance(output, bytes): + return output.decode("utf-8", errors="replace") + return output + + +def _merge_subprocess_output(*outputs: str | bytes | None) -> str: + """Preserve timeout output without duplicating communicate's cumulative data.""" + merged = "" + for output in outputs: + text = _subprocess_output_text(output) + if not text or text in merged: + continue + if merged and merged in text: + merged = text + else: + merged += text + return merged + + +def _validate_rank_marker_set( + output: str, + expected_world_size: int, + marker: str, + pattern: re.Pattern[str], +) -> None: + """Require one exact structured marker from every expected rank.""" + marker_lines = [line for line in output.splitlines() if marker in line] + # MPI launchers may prefix a forwarded rank-0 stdout line. The worker emits + # all success evidence from one canonical stream, so strip only the prefix + # and retain strict validation of the marker payload itself. + marker_payloads = [line[line.index(marker) :] for line in marker_lines] + matches = [pattern.fullmatch(payload) for payload in marker_payloads] + malformed = [line for line, match in zip(marker_lines, matches) if match is None] + reported = [ + (int(match.group(1)), int(match.group(2))) for match in matches if match is not None + ] + expected_ranks = list(range(expected_world_size)) + reported_ranks = sorted(rank for rank, _world_size in reported) + invalid_world_sizes = sorted( + {world_size for _rank, world_size in reported if world_size != expected_world_size} + ) + if malformed or invalid_world_sizes or reported_ranks != expected_ranks: + pytest.fail( + f"expected exactly one {marker} from ranks " + f"{expected_ranks} with world_size={expected_world_size}, got " + f"markers={reported}, malformed={malformed}, " + f"invalid_world_sizes={invalid_world_sizes}:\n{output}", + pytrace=False, + ) + + +def _validate_terminal_completion(output: str, expected_world_size: int) -> None: + """Require exact READY/DONE sets around one global completion marker.""" + if _TERMINAL_ERROR_MARKER in output: + pytest.fail(f"terminal MPI worker reported an error:\n{output}", pytrace=False) + if _TERMINAL_ATEXIT_MARKER in output: + pytest.fail( + "terminal MPI worker entered Python atexit/mpi4py Finalize instead " + f"of the intentional os._exit path:\n{output}", + pytrace=False, + ) + + _validate_rank_marker_set( + output, + expected_world_size, + _TERMINAL_READY_MARKER, + _TERMINAL_READY_PATTERN, + ) + + completion_lines = [line for line in output.splitlines() if _TERMINAL_COMPLETE_MARKER in line] + completion_payloads = [ + line[line.index(_TERMINAL_COMPLETE_MARKER) :] for line in completion_lines + ] + completion_matches = [ + _TERMINAL_COMPLETE_PATTERN.fullmatch(payload) for payload in completion_payloads + ] + if ( + len(completion_matches) != 1 + or completion_matches[0] is None + or int(completion_matches[0].group(1)) != expected_world_size + ): + pytest.fail( + f"expected exactly one valid {_TERMINAL_COMPLETE_MARKER} with " + f"world_size={expected_world_size}, got {completion_lines}:\n{output}", + pytrace=False, + ) + _validate_rank_marker_set( + output, + expected_world_size, + _TERMINAL_DONE_MARKER, + _TERMINAL_DONE_PATTERN, + ) + + +def _validate_worker_result( + returncode: int, + output: str, + *, + required: bool, + mode: str = _HEALTHY_MODE, + expected_world_size: int | None = None, +) -> None: + """Validate normal-finalize and intentional no-finalize launcher results.""" + # A skip marker never excuses a launcher failure. Required workers return a + # nonzero status instead of printing this marker, but retain this ordering + # so a malformed optional worker cannot hide a launch error. + if returncode != 0 and (mode != _TERMINAL_MODE or _SKIP_MARKER in output): + pytest.fail(output or f"MPI smoke launcher exited with status {returncode}", pytrace=False) + if _SKIP_MARKER not in output: + if mode == _TERMINAL_MODE: + if expected_world_size is None: + raise ValueError("terminal MPI smoke validation requires expected_world_size") + # Open MPI reports a nonzero launcher status when workers + # intentionally bypass MPI_Finalize via os._exit(). Structured + # rank readiness plus global completion distinguishes that expected + # status from a real worker failure or MPI_Abort. + _validate_terminal_completion(output, expected_world_size) + return + if required: + pytest.fail( + f"WideEP MPI FT smoke is required but the worker skipped:\n{output.strip()}", + pytrace=False, + ) + pytest.skip(output.strip()) + + +def _parse_propagation_metric(output: str, expected_world_size: int) -> tuple[float, bool]: + """Validate and return the functional smoke's structured latency metric.""" + metric_lines = [line for line in output.splitlines() if _PROPAGATION_MARKER in line] + metric_payloads = [line[line.index(_PROPAGATION_MARKER) :] for line in metric_lines] + matches = [_PROPAGATION_PATTERN.fullmatch(payload) for payload in metric_payloads] + if len(matches) != 1 or matches[0] is None: + pytest.fail( + f"expected exactly one valid {_PROPAGATION_MARKER} line, got {metric_lines}:\n{output}", + pytrace=False, + ) + match = matches[0] + assert match is not None + world_size = int(match.group(1)) + try: + elapsed_ms = float(match.group(2)) + target_ms = float(match.group(3)) + except ValueError: + pytest.fail(f"invalid propagation metric:\n{match.group(0)}", pytrace=False) + target_met = match.group(4) == "True" + if world_size != expected_world_size: + pytest.fail( + f"propagation metric reported world_size={world_size}, expected {expected_world_size}", + pytrace=False, + ) + if not math.isfinite(elapsed_ms) or elapsed_ms < 0: + pytest.fail(f"invalid elapsed_ms={elapsed_ms}", pytrace=False) + if not math.isfinite(target_ms) or target_ms != 100.0: + pytest.fail(f"unexpected target_ms={target_ms}, expected 100", pytrace=False) + if target_met is not (elapsed_ms < target_ms): + pytest.fail( + f"target_met={target_met} is inconsistent with " + f"elapsed_ms={elapsed_ms} and target_ms={target_ms}", + pytrace=False, + ) + return elapsed_ms, target_met + + +def _validate_propagation_budget(output: str, expected_world_size: int) -> float: + """Require the real-MPI logical-failure fanout to meet its 100 ms budget.""" + elapsed_ms, target_met = _parse_propagation_metric(output, expected_world_size) + if not target_met: + pytest.fail( + f"WideEP MPI propagation exceeded the <100 ms budget: elapsed_ms={elapsed_ms}", + pytrace=False, + ) + return elapsed_ms + + +def _kill_and_reap_launcher(process: _ReapableProcess, timeout: float) -> bool: + """Kill the launcher and make one bounded attempt to reap it.""" + try: + process.kill() + except ProcessLookupError: + pass + try: + process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + return False + return True + + +@pytest.mark.parametrize( + ("environment", "expected"), + [ + ({}, True), + ({_ALLOW_SKIP_ENV: "0"}, True), + ({_ALLOW_SKIP_ENV: "1"}, False), + ], +) +def test_mpi_ft_smoke_required_mode(environment: dict[str, str], expected: bool) -> None: + assert _smoke_is_required(environment) is expected + + +def test_mpi_ft_smoke_missing_prerequisite_is_optional_locally() -> None: + with pytest.raises(pytest.skip.Exception, match="missing launcher"): + _handle_missing_prerequisite("missing launcher", required=False) + + +def test_mpi_ft_smoke_missing_prerequisite_fails_when_required() -> None: + with pytest.raises(pytest.fail.Exception, match="required.*missing launcher"): + _handle_missing_prerequisite("missing launcher", required=True) + + +def test_mpi_ft_smoke_launcher_failure_precedes_skip_marker() -> None: + with pytest.raises(pytest.fail.Exception, match="launcher failed"): + _validate_worker_result( + 1, + f"launcher failed\n{_SKIP_MARKER} unsupported thread level", + required=False, + ) + + +def test_mpi_ft_smoke_worker_skip_fails_when_required() -> None: + with pytest.raises(pytest.fail.Exception, match="worker skipped"): + _validate_worker_result( + 0, + f"{_SKIP_MARKER} unsupported thread level", + required=True, + ) + + +def test_mpi_ft_smoke_worker_skip_is_optional_when_opted_out() -> None: + with pytest.raises(pytest.skip.Exception, match="unsupported thread level"): + _validate_worker_result( + 0, + f"{_SKIP_MARKER} unsupported thread level", + required=False, + ) + + +def _terminal_completion_output(world_size: int) -> str: + ready = [ + f"{_TERMINAL_READY_MARKER} rank={rank} world_size={world_size}" + for rank in range(world_size) + ] + done = [ + f"{_TERMINAL_DONE_MARKER} rank={rank} world_size={world_size}" for rank in range(world_size) + ] + return "\n".join([*ready, f"{_TERMINAL_COMPLETE_MARKER} world_size={world_size}", *done]) + + +@pytest.mark.parametrize("returncode", [0, 1, 137], ids=["zero", "open-mpi", "signal-style"]) +def test_mpi_ft_terminal_result_accepts_complete_evidence(returncode: int) -> None: + _validate_worker_result( + returncode, + _terminal_completion_output(4), + required=True, + mode=_TERMINAL_MODE, + expected_world_size=4, + ) + + +@pytest.mark.parametrize( + "output", + [ + "", + "\n".join( + [ + f"{_TERMINAL_READY_MARKER} rank=0 world_size=2", + f"{_TERMINAL_COMPLETE_MARKER} world_size=2", + ] + ), + "\n".join( + [ + f"{_TERMINAL_READY_MARKER} rank=0 world_size=2", + f"{_TERMINAL_READY_MARKER} rank=0 world_size=2", + f"{_TERMINAL_READY_MARKER} rank=1 world_size=2", + f"{_TERMINAL_COMPLETE_MARKER} world_size=2", + ] + ), + "\n".join( + [ + f"{_TERMINAL_READY_MARKER} rank=0 world_size=3", + f"{_TERMINAL_READY_MARKER} rank=1 world_size=3", + f"{_TERMINAL_COMPLETE_MARKER} world_size=3", + ] + ), + ], + ids=["missing", "incomplete", "duplicate", "wrong-world-size"], +) +def test_mpi_ft_terminal_result_rejects_incomplete_rank_set(output: str) -> None: + with pytest.raises(pytest.fail.Exception, match="expected exactly one"): + _validate_worker_result( + 1, + output, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +@pytest.mark.parametrize( + "completion_suffix", + [ + "", + "\n".join( + [ + f"{_TERMINAL_COMPLETE_MARKER} world_size=2", + f"{_TERMINAL_COMPLETE_MARKER} world_size=2", + ] + ), + f"{_TERMINAL_COMPLETE_MARKER} world_size=3", + f"{_TERMINAL_COMPLETE_MARKER} world_size=2 trailing-garbage", + f"{_TERMINAL_COMPLETE_MARKER} world_size=2\n{_TERMINAL_ERROR_MARKER} rank=0", + ], + ids=["missing", "duplicate", "wrong-world-size", "malformed", "worker-error"], +) +def test_mpi_ft_terminal_result_rejects_invalid_global_completion( + completion_suffix: str, +) -> None: + ready = "\n".join(f"{_TERMINAL_READY_MARKER} rank={rank} world_size=2" for rank in range(2)) + done = "\n".join(f"{_TERMINAL_DONE_MARKER} rank={rank} world_size=2" for rank in range(2)) + output = "\n".join(part for part in (ready, completion_suffix, done) if part) + with pytest.raises(pytest.fail.Exception): + _validate_worker_result( + 1, + output, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +def test_mpi_ft_terminal_result_rejects_crash_before_done() -> None: + ready_and_complete = "\n".join( + [ + *(f"{_TERMINAL_READY_MARKER} rank={rank} world_size=2" for rank in range(2)), + f"{_TERMINAL_COMPLETE_MARKER} world_size=2", + ] + ) + with pytest.raises(pytest.fail.Exception, match=_TERMINAL_DONE_MARKER): + _validate_worker_result( + 137, + ready_and_complete, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +@pytest.mark.parametrize( + "done_suffix", + [ + f"{_TERMINAL_DONE_MARKER} rank=0 world_size=2", + "\n".join( + [ + f"{_TERMINAL_DONE_MARKER} rank=0 world_size=2", + f"{_TERMINAL_DONE_MARKER} rank=0 world_size=2", + f"{_TERMINAL_DONE_MARKER} rank=1 world_size=2", + ] + ), + "\n".join( + [ + f"{_TERMINAL_DONE_MARKER} rank=0 world_size=3", + f"{_TERMINAL_DONE_MARKER} rank=1 world_size=3", + ] + ), + "\n".join( + [ + f"{_TERMINAL_DONE_MARKER} rank=0 world_size=2", + f"{_TERMINAL_DONE_MARKER} rank=1 world_size=2 trailing-garbage", + ] + ), + ], + ids=["incomplete", "duplicate", "wrong-world-size", "malformed"], +) +def test_mpi_ft_terminal_result_rejects_invalid_done_set(done_suffix: str) -> None: + ready = "\n".join(f"{_TERMINAL_READY_MARKER} rank={rank} world_size=2" for rank in range(2)) + output = "\n".join([ready, f"{_TERMINAL_COMPLETE_MARKER} world_size=2", done_suffix]) + with pytest.raises(pytest.fail.Exception, match=_TERMINAL_DONE_MARKER): + _validate_worker_result( + 1, + output, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +def test_mpi_ft_terminal_result_accepts_launcher_prefixes_and_noise() -> None: + output = "\n".join( + [ + "unrelated launcher diagnostic", + *(f"[rank0]: {line}" for line in _terminal_completion_output(2).splitlines()), + "another unrelated diagnostic", + ] + ) + + _validate_worker_result( + 1, + output, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +def test_mpi_ft_terminal_result_rejects_python_atexit_evidence() -> None: + output = "\n".join( + [ + _terminal_completion_output(2), + f"{_TERMINAL_ATEXIT_MARKER} pid=123", + ] + ) + + with pytest.raises(pytest.fail.Exception, match="atexit/mpi4py Finalize"): + _validate_worker_result( + 0, + output, + required=True, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +def test_mpi_ft_terminal_launcher_failure_precedes_skip_marker() -> None: + with pytest.raises(pytest.fail.Exception, match="launcher failed"): + _validate_worker_result( + 1, + f"launcher failed\n{_SKIP_MARKER} unsupported thread level", + required=False, + mode=_TERMINAL_MODE, + expected_world_size=2, + ) + + +@pytest.mark.parametrize( + ("outputs", "expected"), + [ + ((None, "complete"), "complete"), + ((b"partial ", b"partial complete"), "partial complete"), + (("first\n", "second\n"), "first\nsecond\n"), + ((b"bad-\xff",), "bad-\ufffd"), + ], + ids=["none", "cumulative", "incremental", "bytes-replacement"], +) +def test_merge_subprocess_output_preserves_timeout_diagnostics( + outputs: tuple[str | bytes | None, ...], expected: str +) -> None: + assert _merge_subprocess_output(*outputs) == expected + + +def test_parse_propagation_metric_validates_target_without_requiring_it() -> None: + elapsed_ms, target_met = _parse_propagation_metric( + f"{_PROPAGATION_MARKER} world_size=4 elapsed_ms=125.5 target_ms=100.0 target_met=False", + 4, + ) + + assert elapsed_ms == 125.5 + assert target_met is False + + +def test_parse_propagation_metric_accepts_launcher_prefix_and_noise() -> None: + metric = f"{_PROPAGATION_MARKER} world_size=4 elapsed_ms=25.5 target_ms=100.0 target_met=True" + output = f"launcher diagnostic\n[rank0]: {metric}\nmore noise" + + assert _parse_propagation_metric(output, 4) == (25.5, True) + + +def test_validate_propagation_budget_rejects_semantically_valid_miss() -> None: + metric = f"{_PROPAGATION_MARKER} world_size=4 elapsed_ms=125.5 target_ms=100.0 target_met=False" + assert _parse_propagation_metric(metric, 4) == (125.5, False) + + with pytest.raises(pytest.fail.Exception, match="exceeded the <100 ms budget"): + _validate_propagation_budget(metric, 4) + + +@pytest.mark.parametrize("wait_times_out", [False, True], ids=["reaped", "still-running"]) +def test_kill_and_reap_launcher_is_bounded(wait_times_out: bool) -> None: + class FakeProcess: + def __init__(self) -> None: + self.kill_calls = 0 + self.wait_timeouts: list[float] = [] + + def kill(self) -> None: + self.kill_calls += 1 + + def wait(self, timeout: float | None = None) -> int: + assert timeout is not None + self.wait_timeouts.append(timeout) + if wait_times_out: + raise subprocess.TimeoutExpired("mpiexec", timeout) + return -signal.SIGKILL + + process = FakeProcess() + assert _kill_and_reap_launcher(process, 0.25) is (not wait_times_out) + assert process.kill_calls == 1 + assert process.wait_timeouts == [0.25] + + +@pytest.mark.parametrize( + "metric", + [ + "", + f"{_PROPAGATION_MARKER} world_size=2 elapsed_ms=nan target_ms=100 target_met=False", + f"{_PROPAGATION_MARKER} world_size=2 elapsed_ms=-1 target_ms=100 target_met=True", + f"{_PROPAGATION_MARKER} world_size=2 elapsed_ms=1 target_ms=99 target_met=True", + f"{_PROPAGATION_MARKER} world_size=2 elapsed_ms=101 target_ms=100 target_met=True", + f"{_PROPAGATION_MARKER} world_size=2 elapsed_ms=1 target_ms=100 target_met=True trailing", + ], +) +def test_parse_propagation_metric_rejects_malformed_semantics(metric: str) -> None: + with pytest.raises(pytest.fail.Exception): + _parse_propagation_metric(metric, 2) + + +@pytest.mark.parametrize("world_size", [2, 4], ids=["2-ranks", "4-ranks"]) +def test_ep_failure_broadcast_real_mpi(world_size: int) -> None: + """Exercise logical failure fanout and shutdown over a real MPI transport. + + The victim remains alive for portable ``COMM_WORLD`` result coordination; + actual process death and MPI error classification are outside this smoke. + """ + required = _smoke_is_required() + if importlib.util.find_spec("mpi4py") is None: + _handle_missing_prerequisite("mpi4py is not installed", required=required) + launcher = _mpi_launcher() + if launcher is None: + _handle_missing_prerequisite("mpiexec/mpirun is not installed", required=required) + + worker = Path(__file__).with_name("_ep_failure_broadcast_mpi_worker.py") + environment = os.environ.copy() + # Open MPI requires explicit opt-in when CI runs inside a root container. + # Other MPI implementations ignore these variables. + environment["OMPI_ALLOW_RUN_AS_ROOT"] = "1" + environment["OMPI_ALLOW_RUN_AS_ROOT_CONFIRM"] = "1" + environment["OMPI_MCA_rmaps_base_oversubscribe"] = "1" + environment["PYTHONUNBUFFERED"] = "1" + if not required: + environment[_ALLOW_SKIP_ENV] = "1" + + outputs: dict[str, str] = {} + for mode in (_HEALTHY_MODE, _TERMINAL_MODE): + command = [ + launcher, + "-n", + str(world_size), + sys.executable, + str(worker), + str(world_size), + mode, + ] + process = subprocess.Popen( # nosec B603 + command, + env=environment, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + start_new_session=True, + text=True, + ) + try: + output, _ = process.communicate(timeout=_WORKER_TIMEOUT_SEC) + except subprocess.TimeoutExpired as timeout_error: + partial_output = _subprocess_output_text(timeout_error.output) + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + # mpiexec may exit between communicate() timing out and cleanup. + pass + try: + output, _ = process.communicate(timeout=_REAP_TIMEOUT_SEC) + except subprocess.TimeoutExpired as reap_error: + output = _merge_subprocess_output(partial_output, reap_error.output) + reaped_after_kill = _kill_and_reap_launcher(process, _REAP_TIMEOUT_SEC) + if process.stdout is not None: + try: + process.stdout.close() + except OSError: + pass + reap_status = ( + "launcher was reaped after direct SIGKILL" + if reaped_after_kill + else "launcher still could not be reaped after direct SIGKILL" + ) + pytest.fail( + f"{world_size}-rank WideEP MPI {mode} smoke timed out after " + f"{_WORKER_TIMEOUT_SEC:.0f}s and could not be reaped within " + f"{_REAP_TIMEOUT_SEC:.0f}s; {reap_status}\n{output}" + ) + output = _merge_subprocess_output(partial_output, output) + pytest.fail( + f"{world_size}-rank WideEP MPI {mode} smoke timed out after " + f"{_WORKER_TIMEOUT_SEC:.0f}s\n{output}" + ) + + _validate_worker_result( + process.returncode, + output, + required=required, + mode=mode, + expected_world_size=world_size, + ) + outputs[mode] = output + + _validate_propagation_budget(outputs[_TERMINAL_MODE], world_size) + assert _INVALID_TOPOLOGY_MARKER in outputs[_HEALTHY_MODE] + assert _HEALTHY_LIFECYCLE_MARKER in outputs[_HEALTHY_MODE] + assert _ABORT_WORLD_MARKER in outputs[_TERMINAL_MODE] From a18840c87c4045595b75abc81a5e8d36ccaa5f2c Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 12:43:43 -0700 Subject: [PATCH 10/14] fix: separate detected and committed EP health Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/pyexecutor/ep_failure_broadcast.py | 287 +++++++++++------- .../_ep_failure_broadcast_mpi_worker.py | 11 +- .../pyexecutor/test_ep_failure_broadcast.py | 126 ++++++-- 3 files changed, 270 insertions(+), 154 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py b/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py index fd3cd1018c51..3ac2a74b85e9 100644 --- a/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py +++ b/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py @@ -13,26 +13,29 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Out-of-band MPI failure propagation for WideEP fault tolerance. +"""Out-of-band MPI failure-detection propagation for WideEP fault tolerance. The model-forward thread can be stuck in an AlltoAll kernel when a peer fails, so it cannot participate in recovery consensus. :class:`MpiFtSubcomm` owns a dedicated EP-scoped MPI communicator and progresses only nonblocking point-to-point requests on an independent CPU thread. -The component implements authoritative single-failure propagation. Multi-rank -suspect/confirm consensus and communicator reconstruction are intentionally -left to the later WideEP FT phases. +The component propagates one rank-failure detection and updates only the +process-local *detected* :class:`EPGroupHealth`. It never updates the committed +execution mask. Multi-rank suspect/confirm consensus, commit authorization, +and communicator reconstruction are intentionally left to later WideEP FT +phases. Every survivor echoes a newly observed failure. Receiving one echo from every -active survivor proves that all of them updated their local health. This remains -an asynchronous protocol, not an iteration barrier: the model-engine hook polls -:meth:`health_is_reconciled` at safe iteration boundaries, and the later -barrier-piggyback phase integrates the same state directly into the iteration -barrier. ULFM revoke is reserved for terminal control-plane aborts because MPI -does not carry enough information to distinguish a successful commit revoke -from an emergency abort at remote ranks. Without ULFM, terminal state is echoed -to every survivor; failure to confirm that echo within a bounded interval uses +active survivor proves only that they propagated the same detection. Detection +reconciliation does **not** authorize request resume, EPLB changes, NCCL +reconfiguration or use, or committed-mask mutation. The future 1c.4b recovery +coordinator consumes the detection callback and owns that authorization. + +ULFM revoke is reserved for terminal control-plane aborts because MPI does not +carry enough information to distinguish a successful commit revoke from an +emergency abort at remote ranks. Without ULFM, terminal state is echoed to every +survivor; failure to confirm that echo within a bounded interval uses ``MPI_Abort`` on the world-spanning FT communicator as a fail-stop fallback. """ @@ -123,7 +126,11 @@ class _MpiModule(Protocol): def Query_thread(self) -> int: ... -FailureReceivedCallback = Callable[[int, int, float], None] +FailureDetectedCallback = Callable[[int, int, float], None] +# Compatibility name for callers of the original draft API. New integrations +# should use ``FailureDetectedCallback`` because this callback carries detected, +# never committed, state. +FailureReceivedCallback = FailureDetectedCallback @dataclass(frozen=True) @@ -201,25 +208,30 @@ class _Lifecycle(Enum): class MpiFtSubcomm: - """Broadcast EP-rank failures on a dedicated MPI control plane. + """Broadcast EP-rank failure detections on a dedicated MPI control plane. Construction creates the FT communicator collectively on the caller's thread. Call :meth:`start` only after every rank has constructed the component. During normal operation, the progress thread owns MPI calls on - this communicator; :meth:`broadcast_failure` only updates local health and - enqueues a report, so it never waits for a peer or MPI progress. Because - ``MPI.THREAD_MULTIPLE`` is required, a caller whose start/stop deadline - proves that thread is wedged may issue the terminal Revoke/Abort fallback. + this communicator; :meth:`report_detected_failure` updates only local + detected health and enqueues a report, so it never waits for a peer or MPI + progress. The watchdog must call this method instead of pre-marking health. + Because ``MPI.THREAD_MULTIPLE`` is required, a caller whose start/stop + deadline proves that thread is wedged may issue the terminal Revoke/Abort + fallback. Args: mapping: Distributed mapping whose ``moe_ep_group`` defines the FT communicator and whose ``moe_ep_rank`` defines the local rank. - health: Process-local EP health state updated by sent and received - failure reports. + detected_health: Process-local EP detection state updated by sent and + received failure reports. This must be separate from the committed + execution-mask state owned by the recovery coordinator. config: Progress-loop timing configuration. - on_failure_received: Optional callback invoked as - ``(failed_rank, source_rank, monotonic_time)`` only when a received - report changes local health. Delivery runs on a dedicated daemon + on_failure_detected: Optional callback invoked as + ``(failed_rank, source_rank, monotonic_time)`` when a local or + received report first changes detected health. This is the handoff + to the future 1c.4b recovery coordinator; it is detection evidence, + not commit authorization. Delivery runs on a dedicated daemon thread so callback latency cannot block MPI progress. comm: Pre-created FT communicator. Intended for tests; production callers should let the component create it from ``mapping``. @@ -228,19 +240,27 @@ class MpiFtSubcomm: Raises: RuntimeError: If MPI support or thread support is insufficient, or if the communicator does not match ``mapping``. - ValueError: If ``health`` and ``mapping`` describe different EP groups. + ValueError: If ``detected_health`` and ``mapping`` describe different + EP groups, or if both callback keyword names are provided. """ def __init__( self, mapping: Mapping, - health: EPGroupHealth, + detected_health: EPGroupHealth, config: MpiFtSubcommConfig | None = None, - on_failure_received: FailureReceivedCallback | None = None, + on_failure_detected: FailureDetectedCallback | None = None, *, comm: _MpiComm | None = None, mpi_module: _MpiModule | None = None, + on_failure_received: FailureReceivedCallback | None = None, ) -> None: + if on_failure_detected is not None and on_failure_received is not None: + raise ValueError( + "Specify only on_failure_detected; on_failure_received is a compatibility alias" + ) + if on_failure_detected is None: + on_failure_detected = on_failure_received if mpi_module is None: if MPI is None: raise RuntimeError("mpi4py is required for WideEP failure propagation") @@ -259,7 +279,7 @@ def __init__( collective_setup = create_mpi_ft_subcomm( mapping, - health_size=health.moe_world_size, + health_size=detected_health.moe_world_size, ) comm = cast(_MpiComm, collective_setup.comm) @@ -292,10 +312,10 @@ def __init__( "WideEP FT MVP requires one MoE EP group spanning the full MPI world; " f"got world_size={mapping.world_size}, moe_ep_group={ep_group}" ) - if health.moe_world_size != mapping.moe_ep_size: + if detected_health.moe_world_size != mapping.moe_ep_size: raise ValueError( "EPGroupHealth size must match mapping.moe_ep_size, " - f"got {health.moe_world_size} and {mapping.moe_ep_size}" + f"got {detected_health.moe_world_size} and {mapping.moe_ep_size}" ) self._local_rank = mapping.moe_ep_rank self._ep_size = mapping.moe_ep_size @@ -311,8 +331,8 @@ def __init__( f"expected rank={self._local_rank}, size={self._ep_size}" ) - self._health = health - self._on_failure_received = on_failure_received + self._detected_health = detected_health + self._on_failure_detected = on_failure_detected self._outbound_reports: queue.SimpleQueue[_OutboundMessage] = queue.SimpleQueue() self._announced_failures: set[int] = set() self._failure_reporters: dict[int, set[int]] = {} @@ -391,7 +411,7 @@ def start(self) -> None: ) self._lifecycle = _Lifecycle.STARTING try: - if self._on_failure_received is not None: + if self._on_failure_detected is not None: callback_thread = threading.Thread( target=self._callback_loop, name="wide-ep-ft-callback", @@ -452,7 +472,7 @@ def start(self) -> None: ) with self._lifecycle_lock: self._lifecycle = _Lifecycle.FAILED - # Stop telemetry before the terminal MPI action. Real MPI_Abort + # Stop callback delivery before the terminal MPI action. Real MPI_Abort # does not return, while test doubles and Revoke do; either way no # callback worker should be left waiting for progress-thread cleanup. self._request_deadline_monitor_stop() @@ -464,7 +484,7 @@ def start(self) -> None: self._fail_closed_immediately(timeout_error) # The progress thread may be stuck inside the first MPI Irecv and # therefore unable to reach its ``finally`` block. Do not leave the - # independent telemetry worker alive after start() has failed. + # independent callback worker alive after start() has failed. callback_thread = self._callback_thread if callback_thread is not None and callback_thread is not threading.current_thread(): callback_thread.join(self._config.stop_timeout_sec) @@ -576,37 +596,30 @@ def stop(self, timeout: float | None = None) -> None: timeout_error = TimeoutError( "WideEP FT callback thread did not stop before the timeout" ) - # Callbacks are telemetry only. A stuck callback may make this - # stop call time out, but it must not poison MPI state or make - # one rank take the poisoned-world process-exit path. + # The detection handoff does not own MPI state. A stuck future + # coordinator callback may make this stop call time out, but it + # must not make one rank take the poisoned-world process-exit + # path. raise timeout_error with self._lifecycle_lock: if self._lifecycle is not _Lifecycle.FAILED: self._lifecycle = _Lifecycle.STOPPED - def pre_failover(self, failed_rank: int) -> bool: - """Mark and asynchronously announce one failed EP rank. + def report_detected_failure(self, failed_rank: int) -> bool: + """Record and asynchronously announce one detected EP-rank failure. - This is the nonblocking seam used by failure detectors. The report is - queued even when ``health`` was already updated by the caller, because - the host watchdog normally marks local health before invoking the - broadcaster. + This is the nonblocking seam invoked by the host watchdog. The + broadcaster owns the detected-health transition; callers must not mark + ``detected_health`` before invoking it. This method performs no MPI + calls and never waits for the progress thread. Repeated reports for the + same rank are coalesced after the first local announcement. Args: failed_rank: EP-local rank in ``[0, ep_size)``. Returns: - ``True`` if this call changed local health, else ``False``. - """ - return self.broadcast_failure(failed_rank) - - def broadcast_failure(self, failed_rank: int) -> bool: - """Mark and asynchronously broadcast one failed EP rank. - - This method performs no MPI calls and never waits for the progress - thread. Repeated calls for the same rank are coalesced after the first - local announcement. + ``True`` if this call changed local detected health, else ``False``. """ self._validate_rank(failed_rank) with self._lifecycle_lock: @@ -621,27 +634,45 @@ def broadcast_failure(self, failed_rank: int) -> bool: self._lifecycle = _Lifecycle.FAILED self._request_terminal_abort(error, failed_rank, self._local_rank) raise - changed = self._health.mark_failed(failed_rank) + changed = self._detected_health.mark_failed(failed_rank) enqueued = self._observe_failure(failed_rank, self._local_rank) + detected_time = time.monotonic() logger.warning( f"WideEP FT local failure rank={failed_rank} changed={changed} " - f"enqueued={enqueued} elapsed_ms={(time.monotonic() - start_time) * 1000.0:.3f}" + f"enqueued={enqueued} elapsed_ms={(detected_time - start_time) * 1000.0:.3f}" ) + if changed: + self._invoke_detection_callback(failed_rank, self._local_rank, detected_time) return changed + def pre_failover(self, failed_rank: int) -> bool: + """Compatibility alias for :meth:`report_detected_failure`. + + Despite the historical name, this method does not commit failover and + callers must not pre-mark detected health. + """ + return self.report_detected_failure(failed_rank) + + def broadcast_failure(self, failed_rank: int) -> bool: + """Compatibility alias for :meth:`report_detected_failure`.""" + return self.report_detected_failure(failed_rank) + def world_is_poisoned(self) -> bool: """Return whether this communicator epoch has ever observed failure.""" return ( self._transport_poisoned.is_set() or self._accepted_failure_rank() is not None - or not self._health.all_active() + or not self._detected_health.all_active() ) - def failure_is_reconciled(self, failed_rank: int) -> bool: - """Return whether every active survivor has announced ``failed_rank``. + def failure_detection_is_reconciled(self, failed_rank: int) -> bool: + """Return whether every active survivor announced the same detection. - This method performs no MPI calls and is safe to poll from an iteration - boundary. A terminal progress error always fails closed. + This is detection-plane evidence only. A ``True`` result does not + authorize request resume, EPLB changes, NCCL reconfiguration or use, or + committed-mask mutation. The future 1c.4b recovery coordinator owns + those decisions. This method performs no MPI calls, and a terminal + progress error always fails closed. """ self._validate_rank(failed_rank) with self._lifecycle_lock: @@ -649,38 +680,51 @@ def failure_is_reconciled(self, failed_rank: int) -> bool: return False if self.last_error is not None or self._progress_failed.is_set(): return False - snapshot = self._health.snapshot() + snapshot = self._detected_health.snapshot() if failed_rank not in snapshot.failed_ranks: return False accepted_failure, accepted_generation = self._accepted_failure_state() with self._protocol_lock: - if not self._reconciliation_state_is_valid_locked( + if not self._detection_reconciliation_state_is_valid_locked( snapshot, accepted_failure, accepted_generation ): return False - return self._failure_is_reconciled_locked(failed_rank, snapshot.mask) + return self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask) - def health_is_reconciled(self) -> bool: - """Return whether all locally known failures are agreed by survivors.""" + def detected_health_is_reconciled(self) -> bool: + """Return whether survivors propagated every local failure detection. + + This is not a commit gate. A ``True`` result does not authorize request + resume, EPLB changes, NCCL reconfiguration or use, or committed-mask + mutation. The future 1c.4b recovery coordinator owns those decisions. + """ with self._lifecycle_lock: if self._lifecycle is not _Lifecycle.RUNNING: return False if self.last_error is not None or self._progress_failed.is_set(): return False - snapshot = self._health.snapshot() + snapshot = self._detected_health.snapshot() accepted_failure, accepted_generation = self._accepted_failure_state() if not snapshot.failed_ranks: return accepted_failure is None and not self._transport_poisoned.is_set() with self._protocol_lock: - if not self._reconciliation_state_is_valid_locked( + if not self._detection_reconciliation_state_is_valid_locked( snapshot, accepted_failure, accepted_generation ): return False return all( - self._failure_is_reconciled_locked(failed_rank, snapshot.mask) + self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask) for failed_rank in snapshot.failed_ranks ) + def failure_is_reconciled(self, failed_rank: int) -> bool: + """Compatibility alias for :meth:`failure_detection_is_reconciled`.""" + return self.failure_detection_is_reconciled(failed_rank) + + def health_is_reconciled(self) -> bool: + """Compatibility alias for :meth:`detected_health_is_reconciled`.""" + return self.detected_health_is_reconciled() + def _validate_rank(self, rank: int) -> None: if not 0 <= rank < self._ep_size: raise ValueError(f"failed_rank must be in [0, {self._ep_size}), got {rank}") @@ -688,7 +732,7 @@ def _validate_rank(self, rank: int) -> None: def _claim_failure(self, failed_rank: int) -> None: """Atomically enforce the single-distinct-failure MVP contract.""" with self._failure_claim_lock: - snapshot = self._health.snapshot() + snapshot = self._detected_health.snapshot() known_failures = snapshot.failed_ranks conflicting_health = known_failures - {failed_rank} accepted = self._accepted_failed_rank @@ -701,15 +745,18 @@ def _claim_failure(self, failed_rank: int) -> None: f"already accepted rank {first_failure}, received rank {failed_rank}" ) if accepted is None: + if failed_rank in known_failures: + raise RuntimeError( + "WideEP FT detected health was mutated before the failure " + "broadcaster accepted the report; the watchdog must call " + "report_detected_failure() without pre-marking health" + ) self._accepted_failed_rank = failed_rank - # The watchdog may have marked the rank before calling us. If - # not, exactly one effective mark_failed() must be the next - # health transition. Capturing that expected generation here - # prevents an independent mutation in the claim-to-observe - # window from becoming the accepted epoch baseline. - self._accepted_failure_generation = snapshot.generation + ( - 0 if failed_rank in known_failures else 1 - ) + # Exactly one effective mark_failed() must be the next detected + # health transition. Capturing that generation here prevents an + # independent mutation in the claim-to-observe window from + # becoming the accepted epoch baseline. + self._accepted_failure_generation = snapshot.generation + 1 def _accepted_failure_rank(self) -> int | None: with self._failure_claim_lock: @@ -720,9 +767,9 @@ def _accepted_failure_state(self) -> tuple[int | None, int | None]: return self._accepted_failed_rank, self._accepted_failure_generation def _observe_failure(self, failed_rank: int, reporter: int) -> bool: - """Record one authoritative report and enqueue this rank's echo once.""" + """Record one detection report and enqueue this rank's echo once.""" now = time.monotonic() - snapshot = self._health.snapshot() + snapshot = self._detected_health.snapshot() message: _OutboundMessage | None = None with self._protocol_lock: unattributed_deadline = self._unattributed_error_deadlines.get(failed_rank) @@ -737,7 +784,7 @@ def _observe_failure(self, failed_rank: int, reporter: int) -> bool: if failure_enqueued: self._announced_failures.add(failed_rank) message = _OutboundMessage(_MessageKind.FAILURE, failed_rank) - if self._failure_is_reconciled_locked(failed_rank, snapshot.mask): + if self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask): self._reconcile_deadlines.pop(failed_rank, None) if message is not None: self._outbound_reports.put(message) @@ -745,7 +792,7 @@ def _observe_failure(self, failed_rank: int, reporter: int) -> bool: self._deadline_monitor_wake_event.set() return failure_enqueued - def _reconciliation_state_is_valid_locked( + def _detection_reconciliation_state_is_valid_locked( self, snapshot: EPGroupHealthSnapshot, accepted_failure: int | None, @@ -775,7 +822,7 @@ def _request_terminal_abort( ``MPI_Abort``. This method performs no MPI calls. """ self._record_error(error) - snapshot = self._health.snapshot() + snapshot = self._detected_health.snapshot() enqueue_abort = False now = time.monotonic() with self._protocol_lock: @@ -795,12 +842,12 @@ def _request_terminal_abort( self._wake_event.set() self._deadline_monitor_wake_event.set() - def _failure_is_reconciled_locked(self, failed_rank: int, active_mask: int) -> bool: + def _failure_detection_is_reconciled_locked(self, failed_rank: int, active_mask: int) -> bool: required_reporters = {rank for rank in range(self._ep_size) if active_mask & (1 << rank)} if not required_reporters: - # There is nobody left who can safely execute the next collective. - # Do not let the usual ``empty set is a subset`` rule open the - # iteration gate in this terminal boundary case. + # There is no survivor that can participate in later recovery. + # Do not let the usual ``empty set is a subset`` rule claim + # detection propagation completed in this terminal case. return False reporters = self._failure_reporters.get(failed_rank, set()) return failed_rank in self._failure_fanout_complete and required_reporters.issubset( @@ -930,7 +977,7 @@ def _drain_outbound_reports(self, pending_sends: list[_PendingSend]) -> None: return failed_rank = message.failed_rank - snapshot = self._health.snapshot() + snapshot = self._detected_health.snapshot() if message.kind is _MessageKind.ABORT: # Terminal state must reach every rank still believed active. # If one of them is actually the second failed rank, send @@ -1016,7 +1063,7 @@ def _progress_receives(self, pending_receives: dict[int, _PendingReceive]) -> No # Without ULFM, MPI implementations may surface the dead # source as a generic error. Do not guess the failed rank; # drop only this fixed-source receive and keep polling the - # other survivors for an authoritative watchdog report. + # other survivors for a rank-attributed detection report. del pending_receives[source] self._track_unattributed_transport_error(source) logger.warning( @@ -1048,7 +1095,7 @@ def _progress_receives(self, pending_receives: dict[int, _PendingReceive]) -> No f"{failed_rank} from EP rank {source}" ) - if self._health.is_active(source) and not self._stop_event.is_set(): + if self._detected_health.is_active(source) and not self._stop_event.is_set(): try: pending_receives[source] = self._post_receive(source) except Exception as error: @@ -1078,7 +1125,7 @@ def _handle_received_report(self, failed_rank: int, source: int) -> None: f"WideEP FT received failure rank={failed_rank} source={source} " f"elapsed_ms={(received_time - start_time) * 1000.0:.3f}" ) - self._invoke_callback(failed_rank, source, received_time) + self._invoke_detection_callback(failed_rank, source, received_time) def _handle_failed_peer(self, failed_rank: int, source: int) -> None: changed = self._accept_received_failure( @@ -1093,7 +1140,7 @@ def _handle_failed_peer(self, failed_rank: int, source: int) -> None: logger.warning( f"WideEP FT observed failed EP rank {failed_rank} through the FT communicator" ) - self._invoke_callback(failed_rank, source, detected_time) + self._invoke_detection_callback(failed_rank, source, detected_time) def _accept_received_failure( self, @@ -1111,7 +1158,7 @@ def _accept_received_failure( self._lifecycle = _Lifecycle.FAILED claim_error = error else: - changed = self._health.mark_failed(failed_rank) + changed = self._detected_health.mark_failed(failed_rank) if transport_poisoned: self._transport_poisoned.set() self._observe_failure(failed_rank, reporter) @@ -1133,24 +1180,25 @@ def _handle_received_abort(self, payload: int, source: int) -> None: self._lifecycle = _Lifecycle.FAILED self._request_terminal_abort(error, payload, source) - def _invoke_callback(self, failed_rank: int, source: int, event_time: float) -> None: - if self._on_failure_received is None or self._callback_stop_event.is_set(): + def _invoke_detection_callback(self, failed_rank: int, source: int, event_time: float) -> None: + if self._on_failure_detected is None or self._callback_stop_event.is_set(): return self._callback_queue.put((failed_rank, source, event_time)) def _callback_loop(self) -> None: - """Deliver telemetry callbacks without blocking MPI progress.""" + """Deliver detection handoffs without blocking MPI progress.""" while True: event = self._callback_queue.get() if event is None: return failed_rank, source, event_time = event try: - callback = self._on_failure_received + callback = self._on_failure_detected if callback is not None: callback(failed_rank, source, event_time) except Exception as error: - # Telemetry must never terminate either control-plane thread. + # A coordinator callback must never terminate either + # control-plane thread. logger.warning(f"WideEP FT failure callback raised: {error}") def _request_callback_stop(self) -> None: @@ -1256,35 +1304,40 @@ def _track_unattributed_transport_error(self, source: int) -> None: self._deadline_monitor_wake_event.set() def _check_accepted_failure_epoch(self) -> None: - """Fail closed when health leaves the accepted single-failure epoch.""" + """Fail closed when detected health leaves the accepted failure epoch.""" error: RuntimeError | None = None accepted_failure: int | None = None with self._lifecycle_lock: if self._lifecycle not in (_Lifecycle.RUNNING, _Lifecycle.STOPPING): return - snapshot = self._health.snapshot() + snapshot = self._detected_health.snapshot() accepted_failure, accepted_generation = self._accepted_failure_state() - # A watchdog may mark health immediately before calling - # pre_failover(). Until that call claims a rank, do not mistake the - # documented pre-mark seam for an epoch violation. if accepted_failure is None: - return - if ( + if not snapshot.failed_ranks: + return + error = RuntimeError( + "WideEP FT detected health changed outside the failure " + "broadcaster; the watchdog must call " + "report_detected_failure() without pre-marking health " + f"(snapshot={snapshot})" + ) + elif ( snapshot.failed_ranks == frozenset({accepted_failure}) and snapshot.generation == accepted_generation ): return - error = RuntimeError( - "WideEP FT accepted failure epoch changed before communicator " - "reconstruction " - f"(accepted_rank={accepted_failure}, " - f"expected_generation={accepted_generation}, snapshot={snapshot})" - ) + else: + error = RuntimeError( + "WideEP FT accepted failure epoch changed before communicator " + "reconstruction " + f"(accepted_rank={accepted_failure}, " + f"expected_generation={accepted_generation}, snapshot={snapshot})" + ) self._lifecycle = _Lifecycle.FAILED assert error is not None - assert accepted_failure is not None - self._request_terminal_abort(error, accepted_failure, self._local_rank) + payload = accepted_failure if accepted_failure is not None else -1 + self._request_terminal_abort(error, payload, self._local_rank) def _check_unattributed_transport_deadlines(self) -> None: now = time.monotonic() @@ -1322,10 +1375,10 @@ def _check_reconciliation_deadlines(self) -> None: with self._lifecycle_lock: if self._lifecycle not in (_Lifecycle.RUNNING, _Lifecycle.STOPPING): return - snapshot = self._health.snapshot() + snapshot = self._detected_health.snapshot() with self._protocol_lock: for failed_rank, deadline in list(self._reconcile_deadlines.items()): - if self._failure_is_reconciled_locked(failed_rank, snapshot.mask): + if self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask): del self._reconcile_deadlines[failed_rank] elif now >= deadline: expired.append(failed_rank) @@ -1339,11 +1392,11 @@ def _check_reconciliation_deadlines(self) -> None: def _stop_reconciliation_state(self) -> tuple[bool, BaseException | None]: """Return whether shutdown can stop MPI progress without splitting ranks.""" - snapshot = self._health.snapshot() + snapshot = self._detected_health.snapshot() failed_ranks = snapshot.failed_ranks accepted_failure, accepted_generation = self._accepted_failure_state() with self._protocol_lock: - reconciliation_state_is_valid = self._reconciliation_state_is_valid_locked( + reconciliation_state_is_valid = self._detection_reconciliation_state_is_valid_locked( snapshot, accepted_failure, accepted_generation, @@ -1366,7 +1419,7 @@ def _stop_reconciliation_state(self) -> tuple[bool, BaseException | None]: unreconciled = [ failed_rank for failed_rank in failed_ranks - if not self._failure_is_reconciled_locked(failed_rank, snapshot.mask) + if not self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask) ] untracked = [ failed_rank diff --git a/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py b/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py index 42b1d119607c..ccbff841fbcd 100644 --- a/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py +++ b/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py @@ -273,7 +273,7 @@ def _run_failure_propagation_smoke(world: MPI.Intracomm) -> str | None: propagation_start = time.monotonic() try: assert broadcaster is not None - broadcaster.pre_failover(failure_rank) + broadcaster.report_detected_failure(failure_rank) except Exception: trigger_error = traceback.format_exc() @@ -281,12 +281,13 @@ def _run_failure_propagation_smoke(world: MPI.Intracomm) -> str | None: try: # The victim process remains alive only so the smoke can coordinate # its result on COMM_WORLD. Survivors exclude its logical rank and - # exercise the same failure/reconciliation fanout as production. + # exercise the same detection/reconciliation fanout as production. + # This smoke does not model recovery commit or request resumption. assert broadcaster is not None deadline = time.monotonic() + _CONVERGENCE_TIMEOUT_SEC while not phase_errors and rank != failure_rank and time.monotonic() < deadline: failure_observed = not health.is_active(failure_rank) - failure_reconciled = broadcaster.failure_is_reconciled(failure_rank) + failure_reconciled = broadcaster.failure_detection_is_reconciled(failure_rank) if failure_observed and failure_reconciled: if rank == detector_rank: assert propagation_start is not None @@ -295,8 +296,8 @@ def _run_failure_propagation_smoke(world: MPI.Intracomm) -> str | None: time.sleep(0.005) local_converged = rank == failure_rank or ( not health.is_active(failure_rank) - and broadcaster.failure_is_reconciled(failure_rank) - and broadcaster.health_is_reconciled() + and broadcaster.failure_detection_is_reconciled(failure_rank) + and broadcaster.detected_health_is_reconciled() and broadcaster.world_is_poisoned() ) except Exception: diff --git a/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py index f71227103da3..8ab5975cd934 100644 --- a/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py +++ b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py @@ -309,34 +309,34 @@ def _make_broadcaster( *, size: int = 4, rank: int = 0, - health: EPGroupHealth | None = None, + detected_health: EPGroupHealth | None = None, comm: _FakeComm | None = None, mpi: _FakeMPI | None = None, callback: Callable[[int, int, float], None] | None = None, ) -> tuple[MpiFtSubcomm, EPGroupHealth, _FakeComm, _FakeMPI]: - health = health or EPGroupHealth(size) + detected_health = detected_health or EPGroupHealth(size) comm = comm or _FakeComm(size=size, rank=rank) mpi = mpi or _FakeMPI() broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(size, rank), - health, + detected_health, _TEST_CONFIG, callback, comm=comm, mpi_module=mpi, ) - return broadcaster, health, comm, mpi + return broadcaster, detected_health, comm, mpi -def test_nonblocking_fanout_announces_externally_premarked_rank() -> None: - """The watchdog's prior local mark must not suppress cross-rank fanout.""" - health = EPGroupHealth(4) - assert health.mark_failed(2) is True - broadcaster, _, comm, _ = _make_broadcaster(health=health) +def test_watchdog_callback_records_and_fans_out_detected_failure() -> None: + detected_health = EPGroupHealth(4) + broadcaster, _, comm, _ = _make_broadcaster(detected_health=detected_health) broadcaster.start() caller_thread = threading.get_ident() try: - assert broadcaster.broadcast_failure(2) is False + watchdog_callback = broadcaster.report_detected_failure + assert watchdog_callback(2) is True + assert detected_health.get_failed_ranks() == frozenset({2}) _wait_until(lambda: len(comm.sends) == 2, description="failure fanout") _wait_until( lambda: all(record.request.test_calls > 0 for record in comm.sends), @@ -350,7 +350,7 @@ def test_nonblocking_fanout_announces_externally_premarked_rank() -> None: assert all(record.buffer_ref() is not None for record in comm.sends) # A repeated detector report is locally coalesced. - assert broadcaster.pre_failover(2) is False + assert watchdog_callback(2) is False time.sleep(2 * _TEST_CONFIG.poll_interval_sec) assert len(comm.sends) == 2 @@ -358,17 +358,76 @@ def test_nonblocking_fanout_announces_externally_premarked_rank() -> None: send.request.complete() comm.deliver(source=1, failed_rank=2) comm.deliver(source=3, failed_rank=2) - _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") + _wait_until( + broadcaster.detected_health_is_reconciled, + description="detection reconciliation", + ) + finally: + broadcaster.stop() + + +def test_externally_premarked_detected_health_fails_closed() -> None: + detected_health = EPGroupHealth(2) + assert detected_health.mark_failed(1) is True + broadcaster, _, comm, _ = _make_broadcaster( + size=2, + detected_health=detected_health, + ) + + try: + try: + broadcaster.start() + except RuntimeError: + # The progress thread can detect the invalid initial state before + # start() observes RUNNING, or immediately after start() returns. + pass + _wait_until( + lambda: isinstance(broadcaster.last_error, RuntimeError), + description="pre-marked detected-health rejection", + ) + assert "without pre-marking health" in str(broadcaster.last_error) + _wait_until(lambda: comm.revoke_calls == 1, description="pre-mark fail-closed revoke") + finally: + broadcaster.stop() + + +def test_detection_reconciliation_does_not_mutate_committed_health() -> None: + callbacks: list[tuple[int, int, float]] = [] + detected_health = EPGroupHealth(2) + committed_health = EPGroupHealth(2) + initial_committed_snapshot = committed_health.snapshot() + broadcaster, _, _, _ = _make_broadcaster( + size=2, + detected_health=detected_health, + callback=lambda failed, source, when: callbacks.append((failed, source, when)), + ) + broadcaster.start() + try: + assert broadcaster.report_detected_failure(1) is True + _wait_until( + lambda: len(callbacks) == 1, + description="local detection coordinator handoff", + ) + _wait_until( + broadcaster.detected_health_is_reconciled, + description="detection reconciliation", + ) + + assert detected_health.get_failed_ranks() == frozenset({1}) + assert detected_health.generation == 1 + assert committed_health.snapshot() == initial_committed_snapshot + assert committed_health.all_active() is True + assert callbacks[0][:2] == (1, 0) finally: broadcaster.stop() -def test_survivor_echoes_reconcile_failure_before_next_collective() -> None: +def test_survivor_echoes_reconcile_failure_detection() -> None: broadcaster, _, comm, _ = _make_broadcaster(size=4) broadcaster.start() try: - assert broadcaster.pre_failover(2) is True - assert broadcaster.failure_is_reconciled(2) is False + assert broadcaster.report_detected_failure(2) is True + assert broadcaster.failure_detection_is_reconciled(2) is False _wait_until(lambda: len(comm.sends) == 2, description="initial failure fanout") comm.deliver(source=1, failed_rank=2) @@ -376,20 +435,20 @@ def test_survivor_echoes_reconcile_failure_before_next_collective() -> None: lambda: comm.receive_count(1) == 2, description="source 1 receive repost", ) - assert broadcaster.failure_is_reconciled(2) is False + assert broadcaster.failure_detection_is_reconciled(2) is False comm.deliver(source=3, failed_rank=2) for send in comm.sends: send.request.complete() _wait_until( - lambda: broadcaster.failure_is_reconciled(2), + lambda: broadcaster.failure_detection_is_reconciled(2), description="all survivor echoes", ) time.sleep(2 * _TEST_CONFIG.poll_interval_sec) assert len(comm.sends) == 2 assert {send.destination for send in comm.sends} == {1, 3} assert all(send.message_kind == _FAILURE_MESSAGE for send in comm.sends) - assert broadcaster.health_is_reconciled() is True + assert broadcaster.detected_health_is_reconciled() is True assert comm.revoke_calls == 0 finally: broadcaster.stop() @@ -429,8 +488,7 @@ def test_reactivation_between_failure_claim_and_observe_cannot_become_epoch_base monkeypatch: pytest.MonkeyPatch, ) -> None: health = EPGroupHealth(2) - assert health.mark_failed(1) is True - broadcaster, _, _, _ = _make_broadcaster(size=2, health=health) + broadcaster, _, _, _ = _make_broadcaster(size=2, detected_health=health) original_claim = broadcaster._claim_failure mutated = False @@ -439,6 +497,7 @@ def claim_then_turn_over_epoch(failed_rank: int) -> None: original_claim(failed_rank) if not mutated: mutated = True + assert health.mark_failed(failed_rank) is True assert health.mark_active(failed_rank) is True assert health.mark_failed(failed_rank) is True @@ -465,7 +524,7 @@ def test_second_failure_between_claim_and_observe_cannot_open_single_failure_gat monkeypatch: pytest.MonkeyPatch, ) -> None: health = EPGroupHealth(3) - broadcaster, _, _, _ = _make_broadcaster(size=3, health=health) + broadcaster, _, _, _ = _make_broadcaster(size=3, detected_health=health) original_claim = broadcaster._claim_failure mutated = False @@ -684,8 +743,10 @@ def test_monitor_does_not_world_abort_reconciled_relay_when_mpi_test_resumes() - broadcaster._deadline_monitor_wake_event.set() _wait_until( - lambda: broadcaster._deadline_monitor_thread is not None - and not broadcaster._deadline_monitor_thread.is_alive(), + lambda: ( + broadcaster._deadline_monitor_thread is not None + and not broadcaster._deadline_monitor_thread.is_alive() + ), description="monitor-observed terminal reconciliation", ) assert broadcaster._progress_failed.is_set() is False @@ -940,7 +1001,7 @@ def test_received_report_is_idempotent_and_callback_runs_once() -> None: broadcaster.stop() -def test_monitor_cannot_observe_remote_claim_before_health_commit( +def test_monitor_cannot_observe_remote_claim_before_detected_health_update( monkeypatch: pytest.MonkeyPatch, ) -> None: broadcaster, health, comm, _ = _make_broadcaster(size=3) @@ -965,7 +1026,7 @@ def claim_then_pause(failed_rank: int) -> None: assert comm.revoke_calls == 0 release_claim.set() - _wait_until(lambda: not health.is_active(2), description="committed remote failure") + _wait_until(lambda: not health.is_active(2), description="recorded remote detection") _wait_until(lambda: len(comm.sends) == 1, description="remote failure echo") comm.sends[0].request.complete() _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") @@ -1806,7 +1867,7 @@ def test_remote_revoke_after_local_reconciliation_still_fails_closed() -> None: assert comm.revoke_calls == 1 with pytest.raises(RuntimeError, match="not running"): broadcaster.pre_failover(1) - assert broadcaster._health.get_failed_ranks() == frozenset({2}) + assert broadcaster._detected_health.get_failed_ranks() == frozenset({2}) broadcaster.stop() @@ -1892,14 +1953,14 @@ def test_non_ulfm_generic_receive_error_keeps_other_sources_live() -> None: assert broadcaster.last_error is None assert failed_receive.cancel_calls == 0 - # Another survivor's watchdog report remains authoritative even when - # the dead source's fixed receive failed with a generic error. + # Another survivor's rank-attributed detection remains usable even + # when the dead source's fixed receive failed with a generic error. comm.deliver(source=2, failed_rank=1) - _wait_until(lambda: health.is_active(1) is False, description="authoritative report") + _wait_until(lambda: health.is_active(1) is False, description="detection report") _wait_until(lambda: len(comm.sends) == 1, description="failure relay") comm.sends[0].request.complete() _wait_until(lambda: comm.receive_count(2) == 2, description="live-source receive repost") - _wait_until(lambda: len(callbacks) == 1, description="authoritative callback") + _wait_until(lambda: len(callbacks) == 1, description="detection callback") assert [(failed, source) for failed, source, _ in callbacks] == [(1, 2)] assert broadcaster.last_error is None @@ -1941,8 +2002,9 @@ def test_non_ulfm_generic_error_after_watchdog_report_does_not_arm_stale_deadlin failed_receive = comm.latest_receive(1).request failed_receive.fail(_FakeMpiException(999, "late generic peer error")) _wait_until( - lambda: failed_receive - in [pending.request for pending in broadcaster._retained_requests], + lambda: ( + failed_receive in [pending.request for pending in broadcaster._retained_requests] + ), description="retained explained receive", ) comm.sends[0].request.complete() From 988912571d82db5398a01fe903060f9f4aba72ad Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 10:17:14 -0700 Subject: [PATCH 11/14] Add NCCL fault-tolerance wrapper for WideEP Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../tensorrt_llm/runtime/utils/mpiTags.h | 20 +- .../tensorrt_llm/runtime/utils/ncclHostApi.h | 52 + .../runtime/utils/ncclUniqueIdRendezvous.h | 102 ++ cpp/tensorrt_llm/common/attentionOp.cpp | 141 +- cpp/tensorrt_llm/common/ncclUtils.cpp | 541 +++++- cpp/tensorrt_llm/common/ncclUtils.h | 79 +- cpp/tensorrt_llm/common/opUtils.cpp | 1460 ++++++++++++++++- cpp/tensorrt_llm/common/opUtils.h | 97 +- cpp/tensorrt_llm/runtime/CMakeLists.txt | 3 +- cpp/tensorrt_llm/runtime/ncclCommunicator.cpp | 841 +++++++++- cpp/tensorrt_llm/runtime/ncclCommunicator.h | 112 +- .../runtime/utils/ncclUniqueIdRendezvous.cpp | 787 +++++++++ cpp/tensorrt_llm/thop/allgatherOp.cpp | 66 +- cpp/tensorrt_llm/thop/allreduceOp.cpp | 70 +- cpp/tensorrt_llm/thop/alltoallOp.cpp | 34 +- cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp | 173 +- cpp/tensorrt_llm/thop/ncclCommunicatorOp.h | 9 +- cpp/tensorrt_llm/thop/outputTensor.h | 2 +- cpp/tensorrt_llm/thop/reducescatterOp.cpp | 87 +- cpp/tests/unit_tests/multi_gpu/CMakeLists.txt | 129 +- .../multi_gpu/ncclUniqueIdRendezvousTest.cpp | 971 +++++++++++ .../unit_tests/multi_gpu/ncclUtilsTest.cpp | 247 ++- tensorrt_llm/_mnnvl_utils.py | 35 +- .../_torch/custom_ops/torch_custom_ops.py | 36 +- .../_torch/distributed/communicator.py | 312 +++- .../distributed/nccl_fault_tolerance.py | 229 +++ tensorrt_llm/_torch/distributed/ops.py | 135 +- .../communication/allgather_reducescatter.py | 174 +- tensorrt_llm/_torch/modules/linear.py | 38 +- .../_torch/pyexecutor/model_engine.py | 64 +- .../integration/test_lists/test-db/l0_a10.yml | 2 + .../moe/test_allgather_reducescatter_ft.py | 861 ++++++++++ .../modules/test_pp_nccl_communicator.py | 782 +++++++++ 33 files changed, 8257 insertions(+), 434 deletions(-) create mode 100644 cpp/include/tensorrt_llm/runtime/utils/ncclHostApi.h create mode 100644 cpp/include/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h create mode 100644 cpp/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.cpp create mode 100644 cpp/tests/unit_tests/multi_gpu/ncclUniqueIdRendezvousTest.cpp create mode 100644 tensorrt_llm/_torch/distributed/nccl_fault_tolerance.py create mode 100644 tests/unittest/_torch/modules/moe/test_allgather_reducescatter_ft.py create mode 100644 tests/unittest/_torch/modules/test_pp_nccl_communicator.py diff --git a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h index 32c086c84ee9..caa0651d2a2f 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h +++ b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -71,7 +71,23 @@ enum class MpiTag : int // KvCacheEventManager kKvCacheEventSize = 1026, - kKvCacheEvent = 1027 + kKvCacheEvent = 1027, + + // Fault-tolerant NCCL communicator rendezvous. Each ownership path and + // protocol phase has a distinct tag so stale or simultaneous messages + // cannot be decoded as another phase's payload. + kNcclCommReady = 1028, + kNcclCommUniqueId = 1029, + kNcclCommAck = 1030, + kNcclPpReady = 1031, + kNcclPpUniqueId = 1032, + kNcclPpAck = 1033, + + // Seeds used with the canonical group to derive MPI_Comm_create_group + // tags for dedicated, pre-failure NCCL rendezvous channels. These are + // separate from message tags and distinguish raw-op and PP ownership. + kNcclCommControl = 1034, + kNcclPpControl = 1035 }; } // namespace tensorrt_llm::mpi diff --git a/cpp/include/tensorrt_llm/runtime/utils/ncclHostApi.h b/cpp/include/tensorrt_llm/runtime/utils/ncclHostApi.h new file mode 100644 index 000000000000..a477547902b0 --- /dev/null +++ b/cpp/include/tensorrt_llm/runtime/utils/ncclHostApi.h @@ -0,0 +1,52 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace tensorrt_llm::runtime +{ + +// NCCL host APIs must not run concurrently from different threads for +// communicators associated with the same CUDA device. TRT-LLM can own both a +// PP communicator and several cached raw-op communicators, each with its own +// watchdog, so a per-communicator mutex is insufficient. A process-wide gate +// is deliberately stricter than a per-device gate and avoids relying on CUDA +// thread-local device state in watchdog threads. +// +// Recursive locking lets low-level error paths abort a communicator while the +// operation wrapper already owns the gate. A grouped operation keeps one lock +// alive from ncclGroupStart through ncclGroupEnd. +inline std::recursive_mutex& getNcclHostApiMutex() +{ + // Communicator registries and watchdogs are also process-lifetime + // singletons, with cross-translation-unit destruction order unspecified. + // Keep the gate alive until process exit so late communicator teardown can + // never touch an already-destroyed mutex. + static auto* mutex = new std::recursive_mutex; + return *mutex; +} + +using NcclHostApiLock = std::unique_lock; + +inline NcclHostApiLock acquireNcclHostApiLock() +{ + return NcclHostApiLock{getNcclHostApiMutex()}; +} + +} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h b/cpp/include/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h new file mode 100644 index 000000000000..67b76aa9df3c --- /dev/null +++ b/cpp/include/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h @@ -0,0 +1,102 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/runtime/utils/mpiUtils.h" + +#include +#include +#include +#include + +#if ENABLE_MULTI_DEVICE +#include +#endif + +namespace tensorrt_llm::runtime +{ + +#if ENABLE_MULTI_DEVICE + +struct NcclUniqueIdRendezvousTags +{ + mpi::MpiTag ready; + mpi::MpiTag id; + mpi::MpiTag ack; +}; + +//! An MPI control channel dedicated to NCCL unique-ID rendezvous. +//! +//! The channel is created collectively by the initial NCCL group while every +//! member is healthy. It retains a process-lifetime MPI communicator whose +//! error handler is MPI_ERRORS_RETURN, leaving the parent/session communicator unchanged. The +//! original parent-communicator ranks are retained so a later survivor subset +//! can keep using stable world-rank IDs even though MPI ranks in this channel +//! are compact. +//! +//! This isolates rendezvous traffic and error handling, but it is not a ULFM +//! communicator repair. Post-failure point-to-point progress still requires an +//! MPI implementation and launcher configured to let survivors continue. +class NcclUniqueIdRendezvousComm +{ +public: + NcclUniqueIdRendezvousComm(mpi::MpiComm comm, std::vector worldRanks, int worldRank); + + NcclUniqueIdRendezvousComm(NcclUniqueIdRendezvousComm const&) = delete; + NcclUniqueIdRendezvousComm& operator=(NcclUniqueIdRendezvousComm const&) = delete; + + [[nodiscard]] mpi::MpiComm const& mpiComm() const noexcept; + [[nodiscard]] std::vector const& worldRanks() const noexcept; + [[nodiscard]] int worldRank() const noexcept; + [[nodiscard]] int commRank(int worldRank) const; + [[nodiscard]] int worldRank(int commRank) const; + +private: + mpi::MpiComm mComm; + std::vector mWorldRanks; + int mWorldRank; +}; + +//! Create a dedicated control channel over initialRanks. +//! +//! Every member of initialRanks must call this before any process failure. +//! Non-members must not call. creationTagSeed and the canonical group derive a +//! bounded, deterministic MPI creation tag so concurrently-created overlapping +//! groups do not normally cross-pair. +std::shared_ptr createNcclUniqueIdRendezvousComm( + std::vector const& initialRanks, int worldRank, mpi::MpiComm const& parentComm, int creationTagSeed); + +// Exchange a fresh NCCL unique ID among only activeRanks. READY/ID/ACK tokens +// prevent delayed eager MPI messages from an earlier timed-out attempt from +// pairing different IDs across ranks. rendezvousId is a coordinator-provided +// logical attempt identity: all ranks in one attempt must pass the same value, +// and retries for the same logical communicator must use strictly increasing +// values. The protocol discards older same-communicator messages after seeing +// a later value while retaining messages for future values and other groups. +// All ranks use original/global rank IDs. +// Callers must keep one logical communicator per (control channel, tag set, +// active-rank set), or construct multiple instances in the same order on every +// rank. Recovery reuses the pre-failure control channel and performs no MPI +// collective. +ncclUniqueId exchangeNcclUniqueId(std::vector const& activeRanks, NcclUniqueIdRendezvousComm const& controlComm, + NcclUniqueIdRendezvousTags tags, std::uint64_t rendezvousId, std::chrono::steady_clock::time_point deadline); + +#endif // ENABLE_MULTI_DEVICE + +} // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 0ea70da0d73b..c9c571bdd0b4 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -33,6 +33,7 @@ #include "tensorrt_llm/runtime/utils/debugUtils.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include +#include #include using namespace tensorrt_llm::kernels; @@ -363,20 +364,26 @@ int AttentionOp::ulyssesContextPreprocess(T const* input, T* output, T* buffer, // Do all to all #if ENABLE_MULTI_DEVICE - ncclGroupStart(); - for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (cpIdx != mCpRank) + auto commLease = acquireComm(mCpNcclComm); + auto const comm = commLease.get(); + auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses context preprocess)"); + commLease.groupStart("ncclGroupStart(ulysses context preprocess)"); + auto checkNcclEnqueue = [&](ncclResult_t result) + { commLease.checkEnqueue(result, "NCCL send/recv(ulysses context preprocess)"); }; + for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - NCCLCHECK(ncclSend(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), - (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, - stream)); - NCCLCHECK(ncclRecv(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), - (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, - stream)); + if (cpIdx != mCpRank) + { + checkNcclEnqueue(ncclSend(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), + (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream)); + checkNcclEnqueue(ncclRecv(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), + (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream)); + } } + commLease.groupEnd("ncclGroupEnd(ulysses context preprocess)"); + commLease.track(watchdogToken, stream); } - ncclGroupEnd(); sync_check_cuda_error(stream); #endif // ENABLE_MULTI_DEVICE @@ -422,28 +429,36 @@ int AttentionOp::ulyssesContextPostprocess(T* input, T* output, T* buffer, Enque // all-to-all #if ENABLE_MULTI_DEVICE size_t const elementNum = partialTokenNum * getHeadSize() * mNumAttnHeads; - ncclGroupStart(); - for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (cpIdx != mCpRank) + auto commLease = acquireComm(mCpNcclComm); + auto const comm = commLease.get(); + auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses context postprocess)"); + commLease.groupStart("ncclGroupStart(ulysses context postprocess)"); + auto checkNcclEnqueue = [&](ncclResult_t result) + { commLease.checkEnqueue(result, "NCCL send/recv(ulysses context postprocess)"); }; + for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (mFP8AttenOutput) - { - NCCLCHECK(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum, ncclInt8, - cpIdx, *mCpNcclComm, stream)); - NCCLCHECK(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum, ncclInt8, - cpIdx, *mCpNcclComm, stream)); - } - else + if (cpIdx != mCpRank) { - NCCLCHECK(ncclSend( - buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream)); - NCCLCHECK(ncclRecv( - input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream)); + if (mFP8AttenOutput) + { + checkNcclEnqueue(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum, + ncclInt8, cpIdx, comm, stream)); + checkNcclEnqueue(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum, + ncclInt8, cpIdx, comm, stream)); + } + else + { + checkNcclEnqueue(ncclSend( + buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream)); + checkNcclEnqueue( + ncclRecv(input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream)); + } } } + commLease.groupEnd("ncclGroupEnd(ulysses context postprocess)"); + commLease.track(watchdogToken, stream); } - ncclGroupEnd(); #endif // ENABLE_MULTI_DEVICE // transpose_1_reverse + view @@ -489,20 +504,26 @@ int AttentionOp::ulyssesGenerationPreprocess( #if ENABLE_MULTI_DEVICE auto const partialHeads = mNumAttnHeads + 2 * mNumAttnKVHeads; - ncclGroupStart(); - for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (cpIdx != mCpRank) + auto commLease = acquireComm(mCpNcclComm); + auto const comm = commLease.get(); + auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses generation preprocess)"); + commLease.groupStart("ncclGroupStart(ulysses generation preprocess)"); + auto checkNcclEnqueue = [&](ncclResult_t result) + { commLease.checkEnqueue(result, "NCCL send/recv(ulysses generation preprocess)"); }; + for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - NCCLCHECK(ncclSend(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), - (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, - stream)); - NCCLCHECK(ncclRecv(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), - (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, - stream)); + if (cpIdx != mCpRank) + { + checkNcclEnqueue(ncclSend(buffer + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), + (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream)); + checkNcclEnqueue(ncclRecv(output + cpIdx * (partialTokenNum * getHeadSize() * partialHeads), + (partialTokenNum * getHeadSize() * partialHeads), (*getDtypeMap())[mType], cpIdx, comm, stream)); + } } + commLease.groupEnd("ncclGroupEnd(ulysses generation preprocess)"); + commLease.track(watchdogToken, stream); } - ncclGroupEnd(); sync_check_cuda_error(stream); #endif // ENABLE_MULTI_DEVICE return 0; @@ -525,28 +546,36 @@ int AttentionOp::ulyssesGenerationPostprocess(T* input, T* output, T* buffer, in // do all-to-all #if ENABLE_MULTI_DEVICE size_t const elementNum = partialTokenNum * getHeadSize() * mNumAttnHeads; - ncclGroupStart(); - for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (cpIdx != mCpRank) + auto commLease = acquireComm(mCpNcclComm); + auto const comm = commLease.get(); + auto const watchdogToken = commLease.begin(stream, "NCCL send/recv(ulysses generation postprocess)"); + commLease.groupStart("ncclGroupStart(ulysses generation postprocess)"); + auto checkNcclEnqueue = [&](ncclResult_t result) + { commLease.checkEnqueue(result, "NCCL send/recv(ulysses generation postprocess)"); }; + for (int cpIdx = 0; cpIdx < mCpSize; cpIdx++) { - if (mFP8AttenOutput) + if (cpIdx != mCpRank) { - NCCLCHECK(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum, ncclInt8, - cpIdx, *mCpNcclComm, stream)); - NCCLCHECK(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum, ncclInt8, - cpIdx, *mCpNcclComm, stream)); - } - else - { - NCCLCHECK(ncclSend( - input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream)); - NCCLCHECK(ncclRecv( - buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, *mCpNcclComm, stream)); + if (mFP8AttenOutput) + { + checkNcclEnqueue(ncclSend(reinterpret_cast<__nv_fp8_e4m3*>(input) + cpIdx * elementNum, elementNum, + ncclInt8, cpIdx, comm, stream)); + checkNcclEnqueue(ncclRecv(reinterpret_cast<__nv_fp8_e4m3*>(buffer) + cpIdx * elementNum, elementNum, + ncclInt8, cpIdx, comm, stream)); + } + else + { + checkNcclEnqueue( + ncclSend(input + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream)); + checkNcclEnqueue(ncclRecv( + buffer + cpIdx * elementNum, elementNum, (*getDtypeMap())[mType], cpIdx, comm, stream)); + } } } + commLease.groupEnd("ncclGroupEnd(ulysses generation postprocess)"); + commLease.track(watchdogToken, stream); } - ncclGroupEnd(); #endif // ENABLE_MULTI_DEVICE // do transpose_1_reverse @@ -3198,11 +3227,13 @@ int AttentionOp::initialize() noexcept return 0; } #if ENABLE_MULTI_DEVICE - if (mCpSize > 1 && COMM_SESSION.getSize() > 1) + if (mCpSize > 1) { - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_CHECK_WITH_INFO(mCpGroup.size() == static_cast(mCpSize), + "Context-parallel group size %zu does not match configured CP size %d", mCpGroup.size(), mCpSize); + TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); mCpNcclComm = getComm(mCpGroup); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } #endif // ENABLE_MULTI_DEVICE return 0; diff --git a/cpp/tensorrt_llm/common/ncclUtils.cpp b/cpp/tensorrt_llm/common/ncclUtils.cpp index e36cfddcd404..980fc502fd92 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.cpp +++ b/cpp/tensorrt_llm/common/ncclUtils.cpp @@ -21,8 +21,13 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" +#include +#include #include #include +#include namespace { @@ -70,6 +75,7 @@ struct NcclMemGuard { if (ptr) { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); TLLM_NCCL_CHECK_WARN(ncclMemFree(ptr)); } } @@ -97,6 +103,62 @@ namespace constexpr int kNcclWindowMinRuntimeVersion = NCCL_VERSION(2, 28, 0); constexpr int kNcclGb10WindowFixedVersion = NCCL_VERSION(2, 30, 4); constexpr int kGb10RealSmVersion = 121; +constexpr auto kRawNcclPollInterval = std::chrono::milliseconds{1}; + +struct RawNcclCallCompletion +{ + ncclResult_t result; + bool completed; +}; + +std::chrono::milliseconds getRawNcclCallTimeout() +{ + constexpr int64_t defaultTimeoutMs = 5000; + auto const* value = std::getenv("TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS"); + if (value != nullptr) + { + try + { + auto const parsed = std::stoll(value); + if (parsed > 0) + { + return std::chrono::milliseconds{parsed}; + } + } + catch (...) + { + } + } + return std::chrono::milliseconds{defaultTimeoutMs}; +} + +// The caller must hold the process-wide NCCL host-API gate from the initial +// call through this poll. NCCL permits only ncclCommGetAsyncError after a +// nonblocking API returns ncclInProgress. +RawNcclCallCompletion completeRawNcclCall( + ncclComm_t comm, ncclResult_t initialResult, std::chrono::steady_clock::time_point deadline) noexcept +{ + if (initialResult != ncclInProgress) + { + return {initialResult, true}; + } + + while (std::chrono::steady_clock::now() < deadline) + { + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(comm, &asyncResult); + if (queryResult == ncclSuccess && asyncResult != ncclInProgress) + { + return {asyncResult, true}; + } + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + return {queryResult, true}; + } + std::this_thread::sleep_for(kRawNcclPollInterval); + } + return {ncclInProgress, false}; +} bool isGb10Platform(int realSmVersion, bool isIntegrated) { @@ -107,6 +169,7 @@ bool isGb10Platform(int realSmVersion, bool isIntegrated) bool queryNcclWindowSupported() { #if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0) + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); int version = 0; if (ncclGetVersion(&version) != ncclSuccess) { @@ -211,8 +274,11 @@ bool isNcclWindowSupported() NcclCommResourceManager& NcclCommResourceManager::getInstance() noexcept { - static NcclCommResourceManager instance; - return instance; + // Communicator registries and watchdogs are process-lifetime. Keep their + // resource owner alive until process exit instead of relying on undefined + // cross-translation-unit singleton destruction order. + static auto* instance = new NcclCommResourceManager; + return *instance; } NcclCommResourceManager::~NcclCommResourceManager() @@ -271,18 +337,44 @@ void NcclCommResourceManager::registerResource(ncclComm_t comm, ResourceCleanupF debugName ? debugName : "unnamed", static_cast(comm), resources.size()); } -void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept +void NcclCommResourceManager::beginAbortCleanup(ncclComm_t comm) noexcept { - if (!comm) + if (comm == nullptr) { return; } + std::lock_guard lock(mMutex); + mAbortCleanups.insert(comm); +} + +void NcclCommResourceManager::waitForAbortCleanup(ncclComm_t comm) noexcept +{ + if (comm == nullptr) + { + return; + } + std::unique_lock lock(mMutex); + mAbortCleanupComplete.wait(lock, [this, comm]() { return mAbortCleanups.find(comm) == mAbortCleanups.end(); }); +} + +bool NcclCommResourceManager::cleanupResources(ncclComm_t comm, bool communicatorAborted) noexcept +{ + if (!comm) + { + return true; + } // Check if we're in the process of being destroyed // If so, skip cleanup - the destructor will handle it proactively if (mIsDestroying.load(std::memory_order_acquire)) { - return; + if (communicatorAborted) + { + std::lock_guard lock(mMutex); + mAbortCleanups.erase(comm); + mAbortCleanupComplete.notify_all(); + } + return false; } std::vector resourcesToClean; @@ -297,20 +389,42 @@ void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept // Double-check after acquiring lock (destruction may have started) if (mIsDestroying.load(std::memory_order_acquire)) { - return; + if (communicatorAborted) + { + mAbortCleanups.erase(comm); + mAbortCleanupComplete.notify_all(); + } + return false; + } + + if (communicatorAborted) + { + mAbortCleanups.insert(comm); + } + if (mCleanupsInProgress.find(comm) != mCleanupsInProgress.end()) + { + // Another thread already owns the callbacks. Its completion + // path rechecks the retirement marker, including one that was + // installed after that owner started, before waking handle + // reusers. + return true; } auto it = mCommResources.find(comm); if (it == mCommResources.end()) { // Nothing registered for this comm, nothing to clean up - return; + if (mAbortCleanups.erase(comm) != 0) + { + mAbortCleanupComplete.notify_all(); + } + return true; } // Move resources out (preserves order) and remove from map resourcesToClean = std::move(it->second); mCommResources.erase(it); - + mCleanupsInProgress.insert(comm); // Logging may fail during static destruction, so wrap in try-catch try { @@ -326,12 +440,13 @@ void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept { // If mutex access fails during static destruction, just return. // This prevents segfaults when the singleton is being destroyed. - return; + return false; } } // Clean up outside the lock to avoid deadlocks if cleanup functions try to access the manager // Order is preserved: resources are cleaned up in registration order + bool cleanupSucceeded = true; for (auto& [cleanup, name] : resourcesToClean) { try @@ -346,10 +461,11 @@ void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept { // Ignore logging failures during static destruction } - cleanup(); + cleanupSucceeded = cleanup() && cleanupSucceeded; } catch (std::exception const& e) { + cleanupSucceeded = false; try { TLLM_LOG_ERROR("[NCCLUtil] Exception during cleanup of resource '%s' for NCCL comm %p: %s", @@ -362,6 +478,7 @@ void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept } catch (...) { + cleanupSucceeded = false; try { TLLM_LOG_ERROR("[NCCLUtil] Unknown exception during cleanup of resource '%s' for NCCL comm %p", @@ -373,6 +490,22 @@ void NcclCommResourceManager::cleanupResources(ncclComm_t comm) noexcept } } } + + { + std::lock_guard lock(mMutex); + mCleanupsInProgress.erase(comm); + if (mAbortCleanups.erase(comm) != 0) + { + mAbortCleanupComplete.notify_all(); + } + } + return cleanupSucceeded; +} + +bool NcclCommResourceManager::isAbortCleanup(ncclComm_t comm) const noexcept +{ + std::lock_guard lock(mMutex); + return mAbortCleanups.find(comm) != mAbortCleanups.end(); } bool NcclCommResourceManager::hasResources(ncclComm_t comm) const noexcept @@ -396,93 +529,171 @@ size_t NcclCommResourceManager::getResourceCount(ncclComm_t comm) const noexcept NCCLWindowAllocator& NCCLWindowAllocator::getInstance() { - static NCCLWindowAllocator instance; - return instance; + static auto* instance = new NCCLWindowAllocator; + return *instance; } NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size) +{ + return requestBufferImpl(comm, size, nullptr, nullptr); +} + +NCCLWindowBuffer NCCLWindowAllocator::requestBuffer( + std::shared_ptr const& comm, size_t size, ncclComm_t* associatedComm) +{ + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null"); + if (associatedComm != nullptr) + { + *associatedComm = nullptr; + } + return requestBufferImpl(nullptr, size, &comm, associatedComm); +} + +NCCLWindowBuffer NCCLWindowAllocator::requestBufferImpl( + ncclComm_t comm, size_t size, std::shared_ptr const* managedComm, ncclComm_t* associatedComm) { if (!isNcclWindowSupported()) { return NCCLWindowBuffer(); } - TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null"); TLLM_CHECK_WITH_INFO(size > 0, "Buffer size must be greater than 0"); - std::lock_guard lock(mMutex); + bool const managedFaultTolerance = managedComm != nullptr && isNcclFaultToleranceEnabled(); + if (managedComm == nullptr) + { + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null"); + NcclCommResourceManager::getInstance().waitForAbortCleanup(comm); + } + // Pool hits need only mMutex. On a miss, release all communicator/pool + // locks, serialize allocation attempts, then recheck: another thread may + // have populated the pool while this thread waited. This preserves + // collective ordering without adding a second lock to the steady-state + // pooled-buffer path or inverting the host-gate -> pool-lock cleanup order. + std::unique_lock allocationLock(mAllocationMutex, std::defer_lock); + int handle = 0; + while (true) + { + std::unique_ptr validationLease; + if (managedComm != nullptr) + { + if (managedFaultTolerance) + { + // Keep lock ordering consistent with abort cleanup: + // communicator state / NCCL host gate precede the pool mutex. + validationLease = std::make_unique(acquireComm(*managedComm)); + comm = validationLease->get(); + } + else + { + TLLM_CHECK_WITH_INFO(*managedComm != nullptr, "NCCL communicator cannot be null"); + comm = **managedComm; + } + } + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null"); + if (associatedComm != nullptr) + { + *associatedComm = comm; + } + std::unique_lock poolLock(mMutex); - // Register cleanup callback for this communicator if not already registered - // This is cheap even if no buffers exist yet - cleanup will just return early - registerBufferCleanup(comm); + // Register cleanup callback for this communicator if not already registered. + registerBufferCleanup(comm); + // Keep validationLease until poolLock releases. Lease destruction can + // drain abort cleanup, which re-enters this non-recursive pool mutex. - // Check if we have an available buffer of at least the requested size for this communicator - // Use best-fit: find the smallest buffer that's >= requested size - auto& commBuffers = mBufferPool[comm]; - auto bestFit = commBuffers.end(); - size_t bestFitSize = std::numeric_limits::max(); + auto& commBuffers = mBufferPool[comm]; + auto bestFit = commBuffers.end(); + size_t bestFitSize = std::numeric_limits::max(); - for (auto it = commBuffers.begin(); it != commBuffers.end(); ++it) - { - if (!it->inUse && it->buffer.size >= size && it->buffer.size < bestFitSize) + for (auto it = commBuffers.begin(); it != commBuffers.end(); ++it) { - bestFit = it; - bestFitSize = it->buffer.size; + if (!it->inUse && it->buffer.size >= size && it->buffer.size < bestFitSize) + { + bestFit = it; + bestFitSize = it->buffer.size; + } } - } - if (bestFit != commBuffers.end()) - { - bestFit->inUse = true; - TLLM_LOG_TRACE( - "[NCCLUtil] Reusing NCCL window buffer for comm %p: handle=%d, ptr=%p, size=%zu (requested: %zu)", - static_cast(comm), bestFit->buffer.handle, bestFit->buffer.ptr, bestFit->buffer.size, size); - return bestFit->buffer; - } + if (bestFit != commBuffers.end()) + { + bestFit->inUse = true; + TLLM_LOG_TRACE( + "[NCCLUtil] Reusing NCCL window buffer for comm %p: handle=%d, ptr=%p, size=%zu (requested: %zu)", + static_cast(comm), bestFit->buffer.handle, bestFit->buffer.ptr, bestFit->buffer.size, size); + return bestFit->buffer; + } - // If a previous allocateAndRegisterBuffer call collectively failed for this comm at a size - // no larger than this request, do not retry the known-failing new allocation path. Smaller - // requests and already-pooled buffers can still use NCCL windows. - auto const failureIt = mMinSymmetricFailureSize.find(comm); - if (failureIt != mMinSymmetricFailureSize.end() && size >= failureIt->second) - { - TLLM_LOG_DEBUG("[NCCLUtil] Skipping NCCL window allocation for comm %p, size=%zu; known failure threshold=%zu", - static_cast(comm), size, failureIt->second); - return NCCLWindowBuffer(); - } + auto const failureIt = mMinSymmetricFailureSize.find(comm); + if (failureIt != mMinSymmetricFailureSize.end() && size >= failureIt->second) + { + TLLM_LOG_DEBUG( + "[NCCLUtil] Skipping NCCL window allocation for comm %p, size=%zu; known failure threshold=%zu", + static_cast(comm), size, failureIt->second); + return NCCLWindowBuffer(); + } - // No available buffer found, avoid registration during CUDA graph capture - auto stream = at::cuda::getCurrentCUDAStream(); - cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; - auto capture_err = cudaStreamIsCapturing(stream, &capture_status); - if (capture_err != cudaSuccess) - { - TLLM_LOG_DEBUG("[NCCLUtil] cudaStreamIsCapturing failed: %s", cudaGetErrorString(capture_err)); - } - if (capture_err == cudaSuccess && capture_status != cudaStreamCaptureStatusNone) - { - TLLM_LOG_DEBUG("[NCCLUtil] Skipping NCCL window allocation during capture for comm %p (requested: %zu)", - static_cast(comm), size); - return NCCLWindowBuffer(); + auto stream = at::cuda::getCurrentCUDAStream(); + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + auto const captureResult = cudaStreamIsCapturing(stream, &captureStatus); + if (captureResult != cudaSuccess) + { + TLLM_LOG_DEBUG("[NCCLUtil] cudaStreamIsCapturing failed: %s", cudaGetErrorString(captureResult)); + } + if (captureResult == cudaSuccess && captureStatus != cudaStreamCaptureStatusNone) + { + TLLM_LOG_DEBUG("[NCCLUtil] Skipping NCCL window allocation during capture for comm %p (requested: %zu)", + static_cast(comm), size); + return NCCLWindowBuffer(); + } + + if (!allocationLock.owns_lock()) + { + poolLock.unlock(); + validationLease.reset(); + allocationLock.lock(); + continue; + } + + TLLM_LOG_TRACE( + "[NCCLUtil] Allocating new NCCL window buffer for comm %p, size=%zu", static_cast(comm), size); + handle = static_cast(commBuffers.size()); + poolLock.unlock(); + validationLease.reset(); + break; } - // No available buffer found, allocate a new one - TLLM_LOG_TRACE( - "[NCCLUtil] Allocating new NCCL window buffer for comm %p, size=%zu", static_cast(comm), size); - int handle = static_cast(commBuffers.size()); - NCCLWindowBuffer buffer = allocateAndRegisterBuffer(comm, size, handle); + // CUDA allocation must not hold the communicator mutex or process-wide + // NCCL gate: a watchdog needs both in order to abort stalled work. + std::unique_ptr completionLease; + NCCLWindowBuffer buffer = allocateAndRegisterBuffer(comm, size, handle, managedComm, &completionLease); // Only cache valid buffers. allocateAndRegisterBuffer returns an empty buffer when any rank // failed ncclMemAlloc (collective fallback to plain allreduce); caching it would leak a // permanently "in use" empty entry per request because releaseBuffer is a no-op for nullptr. if (buffer.isValid()) { - commBuffers.push_back({buffer, true}); + // The managed registration lease remains live until the buffer is in + // the pool. That closes the gap where the watchdog could abort and run + // cleanup after registration but before the resource became visible. + { + std::lock_guard poolLock(mMutex); + mBufferPool[comm].push_back({buffer, true}); + } + completionLease.reset(); } else { + std::unique_ptr healthyLease; + if (managedFaultTolerance) + { + healthyLease = std::make_unique(acquireComm(*managedComm)); + TLLM_CHECK_WITH_INFO( + healthyLease->get() == comm, "NCCL communicator changed while recording a window-allocation fallback"); + } // The collective allreduce inside allocateAndRegisterBuffer agreed that this request // cannot use symmetric memory on at least one rank. Remember the smallest failing // request size so repeated too-large autotuner probes do not keep stressing this path. + std::lock_guard poolLock(mMutex); recordSymmetricFailureLocked(comm, size); } @@ -607,7 +818,8 @@ bool NCCLWindowAllocator::isCommValid(ncclComm_t comm) const noexcept return comm != nullptr; } -NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle) +NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle, + std::shared_ptr const* managedComm, std::unique_ptr* completionLease) { // Step 1: Pre-allocate the rank-sync flag before ncclMemAlloc. ncclMemAlloc can fail // asymmetrically with ncclUnhandledCudaError on configurations where the symmetric/VMM path @@ -626,7 +838,10 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, // the stream-ordered flag copy and collective fallback so the failing rank still reaches // ncclAllReduce with the other ranks. void* ncclPtr = nullptr; - TLLM_NCCL_CHECK_WARN(ncclMemAlloc(&ncclPtr, size)); + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + TLLM_NCCL_CHECK_WARN(ncclMemAlloc(&ncclPtr, size)); + } int const localAllocOk = (ncclPtr != nullptr) ? 1 : 0; NcclMemGuard ncclGuard{ncclPtr}; // frees ncclPtr on any early return or exception clearCudaErrorIfSymmetricAllocationFailed(localAllocOk); @@ -640,8 +855,66 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, TLLM_CUDA_CHECK_WARN( cudaMemcpyAsync(rankSyncFlag, &localAllocOk, sizeof(localAllocOk), cudaMemcpyHostToDevice, stream)); } - TLLM_NCCL_CHECK(ncclAllReduce(rankSyncFlag, rankSyncFlag, 1, ncclInt32, ncclMin, comm, stream)); - TLLM_CUDA_CHECK_WARN(cudaStreamSynchronize(stream)); + if (managedComm != nullptr) + { + try + { + std::uint64_t watchdogToken = 0; + { + auto commLease = acquireComm(*managedComm); + auto const currentComm = commLease.get(); + TLLM_CHECK_WITH_INFO(currentComm == comm, + "NCCL communicator changed while allocating a window buffer; retry with the active communicator"); + watchdogToken = commLease.begin(stream, "NCCL window allocation agreement"); + auto const allocAgreementResult + = ncclAllReduce(rankSyncFlag, rankSyncFlag, 1, ncclInt32, ncclMin, currentComm, stream); + commLease.check(allocAgreementResult, "ncclAllReduce(window allocation agreement)"); + commLease.track(watchdogToken, stream); + } + if (isNcclFaultToleranceEnabled()) + { + waitCommOperation(*managedComm, watchdogToken, "NCCL window allocation agreement"); + } + else + { + TLLM_CUDA_CHECK_WARN(cudaStreamSynchronize(stream)); + } + } + catch (...) + { + // A failed NCCL kernel may still hold the temporary flag even + // after communicator abort. Quarantine four bytes instead of + // letting cudaFree introduce a device-wide wait on recovery. + static_cast(flagGuard.release()); + throw; + } + } + else + { + RawNcclCallCompletion completion{ncclSuccess, true}; + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + auto const allocAgreementResult + = ncclAllReduce(rankSyncFlag, rankSyncFlag, 1, ncclInt32, ncclMin, comm, stream); + completion = completeRawNcclCall( + comm, allocAgreementResult, std::chrono::steady_clock::now() + getRawNcclCallTimeout()); + } + if (!completion.completed) + { + // The background enqueue may still acquire the flag after the + // local deadline. Keep its storage alive rather than racing it. + static_cast(flagGuard.release()); + TLLM_THROW("NCCL error: window allocation agreement timed out while the communicator was in progress"); + } + if (completion.result != ncclSuccess) + { + // NCCL documents device work as indeterminate after an async + // communicator error. Quarantine the flag before surfacing it. + static_cast(flagGuard.release()); + TLLM_NCCL_CHECK(completion.result); + } + TLLM_CUDA_CHECK_WARN(cudaStreamSynchronize(stream)); + } int allAllocOk = 0; TLLM_CUDA_CHECK(cudaMemcpy(&allAllocOk, rankSyncFlag, sizeof(int), cudaMemcpyDeviceToHost)); @@ -663,10 +936,60 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, // Failure here is non-fatal: warn and fall back to regular allreduce. // ncclGuard frees ncclPtr on return. ncclWindow_t window = nullptr; - ncclResult_t const regResult = ncclCommWindowRegister(comm, ncclPtr, size, &window, NCCL_WIN_COLL_SYMMETRIC); - TLLM_NCCL_CHECK_WARN(regResult); + ncclResult_t regResult = ncclSuccess; + if (managedComm != nullptr) + { + auto commLease = acquireComm(*managedComm); + auto const currentComm = commLease.get(); + TLLM_CHECK_WITH_INFO(currentComm == comm, + "NCCL communicator changed while registering a window buffer; retry with the active communicator"); + regResult = ncclCommWindowRegister(currentComm, ncclPtr, size, &window, NCCL_WIN_COLL_SYMMETRIC); + try + { + regResult = commLease.checkOptional(regResult, "ncclCommWindowRegister"); + } + catch (...) + { + // Registration may still reference the allocation after an + // asynchronous communicator failure. Quarantine it instead + // of calling ncclMemFree on the recovery path. + static_cast(ncclGuard.release()); + throw; + } + if (regResult == ncclSuccess && completionLease != nullptr && isNcclFaultToleranceEnabled()) + { + *completionLease = std::make_unique(std::move(commLease)); + } + } + else + { + ncclResult_t initialResult = ncclSuccess; + RawNcclCallCompletion completion{ncclSuccess, true}; + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + initialResult = ncclCommWindowRegister(comm, ncclPtr, size, &window, NCCL_WIN_COLL_SYMMETRIC); + completion + = completeRawNcclCall(comm, initialResult, std::chrono::steady_clock::now() + getRawNcclCallTimeout()); + if (!completion.completed) + { + // Registration can still complete in the background and retain + // ncclPtr. The unmanaged caller owns communicator recovery, so + // quarantine the allocation locally and surface the timeout. + static_cast(ncclGuard.release()); + TLLM_THROW("NCCL error: window registration timed out while the communicator was in progress"); + } + regResult = completion.result; + if (initialResult == ncclInProgress && regResult != ncclSuccess) + { + // Once an asynchronous registration fails, NCCL does not promise + // that the allocation is no longer referenced. + static_cast(ncclGuard.release()); + TLLM_THROW("NCCL error: window registration failed asynchronously: %s", ncclGetErrorString(regResult)); + } + } if (regResult != ncclSuccess) { + TLLM_LOG_WARNING("NCCL window registration failed for comm %p; falling back to regular allreduce: %s", + static_cast(comm), ncclGetErrorString(regResult)); return NCCLWindowBuffer{}; } @@ -709,24 +1032,38 @@ void NCCLWindowAllocator::registerBufferCleanup(ncclComm_t comm) // Register cleanup with the resource manager NcclCommResourceManager::getInstance().registerResource( - comm, [this, comm]() { this->cleanupBuffersForComm(comm); }, "NCCLWindowAllocator"); + comm, [this, comm]() { return this->cleanupBuffersForComm(comm); }, "NCCLWindowAllocator"); } -void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept +bool NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept { if (!comm) { - return; + return true; } - // Synchronize CUDA to ensure all operations using these buffers are complete - // before we deregister windows and free memory - cudaError_t cudaErr = cudaDeviceSynchronize(); - if (cudaErr != cudaSuccess) + bool const communicatorAborted = NcclCommResourceManager::getInstance().isAbortCleanup(comm); + bool cudaSynchronized = true; + + // Never hold the process-wide NCCL host gate across a CUDA synchronization: + // work from another failed communicator may be what the synchronization is + // waiting for, and that communicator's watchdog needs the gate to abort it. + tensorrt_llm::runtime::NcclHostApiLock hostApiLock; + if (!communicatorAborted) { - TLLM_LOG_WARNING("[NCCLUtil] cudaDeviceSynchronize failed with error: %d before cleanup for comm %p", cudaErr, - static_cast(comm)); - // Continue anyway - the sync failure might be from a previous error + cudaError_t cudaErr = cudaDeviceSynchronize(); + if (cudaErr != cudaSuccess) + { + cudaSynchronized = false; + TLLM_LOG_WARNING("[NCCLUtil] cudaDeviceSynchronize failed with error: %d before cleanup for comm %p", + cudaErr, static_cast(comm)); + } + if (cudaSynchronized) + { + // Healthy cleanup uses NCCL deregistration/free below. Acquire the gate + // before mMutex to preserve the communicator -> allocator lock order. + hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + } } std::lock_guard lock(mMutex); @@ -735,7 +1072,7 @@ void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept if (mRegisteredComms.find(comm) == mRegisteredComms.end()) { // Already cleaned up or never registered - return; + return true; } auto commIt = mBufferPool.find(comm); @@ -744,7 +1081,7 @@ void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept // No buffers to clean up, but mark as cleaned mRegisteredComms.erase(comm); mMinSymmetricFailureSize.erase(comm); - return; + return true; } TLLM_LOG_TRACE( @@ -772,25 +1109,50 @@ void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept TLLM_LOG_DEBUG("[NCCLUtil] NCCL window allocator teardown for comm %p: %zu buffers, %zu bytes total", static_cast(comm), commIt->second.size(), totalBytes); + if (communicatorAborted || !cudaSynchronized) + { + // Tensor storages or non-NCCL CUDA work may still reference these + // allocations. There is no safe stream fence once the communicator + // has failed, so quarantine them for the remaining process lifetime + // instead of risking a UAF or another blocking CUDA call. Rank failure + // is rare and process restart ultimately reclaims the memory. + TLLM_LOG_WARNING("[NCCLUtil] Quarantining %zu NCCL window buffers (%zu bytes) after %s", commIt->second.size(), + totalBytes, communicatorAborted ? "communicator abort" : "failed CUDA synchronization"); + mBufferPool.erase(commIt); + mRegisteredComms.erase(comm); + mMinSymmetricFailureSize.erase(comm); + return communicatorAborted; + } + + bool quarantineRemainingBuffers = false; for (auto& entry : commIt->second) { if (entry.buffer.isValid()) { + if (quarantineRemainingBuffers) + { + continue; + } // Deregister the window - the communicator is still valid at this point // (cleanup happens before ncclCommDestroy), but we need to be careful // if buffers are still in use by active operations - if (entry.buffer.window && comm) + if (entry.buffer.window && comm && !communicatorAborted) { // Note: Even if buffer is marked inUse, we must deregister since // the communicator is being destroyed. The communicator is valid, // but we should handle potential errors gracefully. - ncclResult_t result = ncclCommWindowDeregister(comm, entry.buffer.window); - if (result != ncclSuccess) + auto const initialResult = ncclCommWindowDeregister(comm, entry.buffer.window); + auto const completion = completeRawNcclCall( + comm, initialResult, std::chrono::steady_clock::now() + getRawNcclCallTimeout()); + if (!completion.completed || completion.result != ncclSuccess) { TLLM_LOG_WARNING( - "[NCCLUtil] ncclCommWindowDeregister failed with error: %d for comm %p, " - "window %p (buffer inUse: %d)", - result, static_cast(comm), static_cast(entry.buffer.window), entry.inUse); + "[NCCLUtil] ncclCommWindowDeregister %s for comm %p, window %p " + "(result: %s, buffer inUse: %d); quarantining this and remaining buffers", + completion.completed ? "failed" : "timed out", static_cast(comm), + static_cast(entry.buffer.window), ncclGetErrorString(completion.result), entry.inUse); + quarantineRemainingBuffers = true; + continue; } } @@ -820,6 +1182,7 @@ void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept mBufferPool.erase(commIt); mRegisteredComms.erase(comm); mMinSymmetricFailureSize.erase(comm); + return !quarantineRemainingBuffers; } #endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0) diff --git a/cpp/tensorrt_llm/common/ncclUtils.h b/cpp/tensorrt_llm/common/ncclUtils.h index fdc88f1aaaba..2b629cf72212 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.h +++ b/cpp/tensorrt_llm/common/ncclUtils.h @@ -30,11 +30,13 @@ #include #include +#include #include #include #include #include #include +#include #include #include #include @@ -61,6 +63,10 @@ TRTLLM_NAMESPACE_BEGIN +class NcclCommLease; +bool isNcclFaultToleranceEnabled(); +std::optional getCommAsyncError(std::shared_ptr const& comm); + namespace common::nccl_util { @@ -68,8 +74,10 @@ namespace common::nccl_util // NCCL Resource Management //============================================================================== -// Resource cleanup function type. Called before the NCCL communicator is destroyed. -using ResourceCleanupFunc = std::function; +// Resource cleanup function type. Called before the NCCL communicator is +// destroyed. False tells the owner that graceful communicator destruction is +// unsafe and it must take the abort/quarantine path instead. +using ResourceCleanupFunc = std::function; // Manages resources associated with NCCL communicators. Thread-safe singleton that maintains // a pool of resources per NCCL comm. Resources are automatically cleaned up when the @@ -88,7 +96,18 @@ class NcclCommResourceManager // the shared_ptr deleter before ncclCommDestroy. // Thread-safe: Uses global mutex to serialize cleanup operations. // Order-preserving: Resources are cleaned up in registration order. - void cleanupResources(ncclComm_t comm) noexcept; + [[nodiscard]] bool cleanupResources(ncclComm_t comm, bool communicatorAborted = false) noexcept; + + // Mark a successfully-aborted native handle as retired before releasing + // the NCCL host gate. A later communicator whose native pointer is reused + // waits for the old pointer-keyed cleanup before it can be published. + void beginAbortCleanup(ncclComm_t comm) noexcept; + void waitForAbortCleanup(ncclComm_t comm) noexcept; + + // True only while abort-specific callbacks for comm are running. Resource + // owners use this to skip operations that require a live communicator or + // a device-wide synchronization (both can hang on a failed communicator). + bool isAbortCleanup(ncclComm_t comm) const noexcept; // Check if a communicator has registered resources. bool hasResources(ncclComm_t comm) const noexcept; @@ -108,7 +127,13 @@ class NcclCommResourceManager using ResourceEntry = std::pair; mutable std::mutex mMutex; + std::condition_variable mAbortCleanupComplete; std::unordered_map> mCommResources; + std::unordered_set mAbortCleanups; + // Track every callback owner, not just callers explicitly marked as + // aborted. A normal cleanup may already own the callbacks when another + // thread retires the same native handle. + std::unordered_set mCleanupsInProgress; std::atomic mIsDestroying{false}; }; @@ -134,6 +159,7 @@ class NcclCommResource { mCleanup(mResource); } + return true; }, debugName); } @@ -223,8 +249,18 @@ class NCCLWindowAllocator // If an unused buffer of at least the requested size exists for this communicator, it will be reused. // Uses best-fit strategy: selects the smallest available buffer that meets the size requirement. // Otherwise, a new buffer is allocated and registered. + // Caller-owned/raw communicator overload. It supports a nonblocking call + // result by bounded polling, but it cannot coordinate with the TRT-LLM + // registry's abort/replacement lock. Registry-managed communicators must + // use the shared_ptr overload below. NCCLWindowBuffer requestBuffer(ncclComm_t comm, size_t size); + // Managed-communicator overload. CUDA allocation and tensor construction + // happen without an operation lease; the allocator acquires short leases + // only around NCCL calls so a watchdog can abort stalled work. + NCCLWindowBuffer requestBuffer( + std::shared_ptr const& comm, size_t size, ncclComm_t* associatedComm = nullptr); + // Search for a buffer by pointer. Returns an invalid buffer if not found. // This matches the UBManager.search_buffer() interface. NCCLWindowBuffer searchBuffer(ncclComm_t comm, void* ptr) const; @@ -264,7 +300,10 @@ class NCCLWindowAllocator ~NCCLWindowAllocator() = default; // Allocate a new buffer and register it with NCCL as a window - NCCLWindowBuffer allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle); + NCCLWindowBuffer requestBufferImpl( + ncclComm_t comm, size_t size, std::shared_ptr const* managedComm, ncclComm_t* associatedComm); + NCCLWindowBuffer allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle, + std::shared_ptr const* managedComm, std::unique_ptr* completionLease); // Record a failed new symmetric allocation (assumes mMutex is already locked). void recordSymmetricFailureLocked(ncclComm_t comm, size_t size); @@ -282,7 +321,7 @@ class NCCLWindowAllocator void registerBufferCleanup(ncclComm_t comm); // Cleanup all buffers for a specific communicator - void cleanupBuffersForComm(ncclComm_t comm) noexcept; + [[nodiscard]] bool cleanupBuffersForComm(ncclComm_t comm) noexcept; struct BufferEntry { @@ -291,6 +330,7 @@ class NCCLWindowAllocator }; mutable std::mutex mMutex; + std::mutex mAllocationMutex; std::unordered_map> mBufferPool; std::unordered_set mRegisteredComms; // Smallest request size that is known to fail collectively for each communicator. @@ -305,11 +345,12 @@ class ScopedNCCLWindowBuffer public: ScopedNCCLWindowBuffer(std::shared_ptr comm, size_t size) : mComm(std::move(comm)) + , mRawComm(nullptr) , mBuffer{} { - if (mComm && *mComm) + if (mComm) { - mBuffer = NCCLWindowAllocator::getInstance().requestBuffer(*mComm, size); + mBuffer = NCCLWindowAllocator::getInstance().requestBuffer(mComm, size, &mRawComm); } } @@ -317,7 +358,7 @@ class ScopedNCCLWindowBuffer { if (mBuffer.isValid()) { - NCCLWindowAllocator::getInstance().releaseBuffer(*mComm, mBuffer.ptr); + NCCLWindowAllocator::getInstance().releaseBuffer(mRawComm, mBuffer.ptr); } } @@ -348,6 +389,7 @@ class ScopedNCCLWindowBuffer private: std::shared_ptr mComm; + ncclComm_t mRawComm; NCCLWindowBuffer mBuffer; }; @@ -357,6 +399,7 @@ class ScopedNCCLWindowBuffer inline std::pair createNCCLWindowTensor( std::shared_ptr comm, at::IntArrayRef shape, torch::ScalarType dtype) { + ncclComm_t rawComm = nullptr; // Calculate buffer size int64_t buffer_size = std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies()) * torch::elementSize(dtype); @@ -376,7 +419,7 @@ inline std::pair createNCCLWindowTensor( auto& allocator = NCCLWindowAllocator::getInstance(); NCCLWindowBuffer buffer; - if (!comm || !*comm) + if (!comm || (!isNcclFaultToleranceEnabled() && !*comm)) { TLLM_LOG_DEBUG("[createNCCLWindowTensor] null comm; returning invalid buffer"); return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); @@ -384,10 +427,17 @@ inline std::pair createNCCLWindowTensor( try { - buffer = allocator.requestBuffer(*comm, buffer_size); + buffer = allocator.requestBuffer(comm, buffer_size, &rawComm); } catch (std::exception const& e) { + // Local allocation failures may fall back to a regular CUDA tensor. + // Never hide a watchdog error: continuing would reuse an aborted + // communicator on the fallback collective. + if (isNcclFaultToleranceEnabled() && getCommAsyncError(comm).has_value()) + { + throw; + } TLLM_LOG_DEBUG("[createNCCLWindowTensor] requestBuffer failed; returning invalid buffer: %s", e.what()); return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); } @@ -398,9 +448,16 @@ inline std::pair createNCCLWindowTensor( TLLM_LOG_DEBUG("[createNCCLWindowTensor] invalid buffer returned from requestBuffer; returning invalid buffer"); return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); } + TLLM_CHECK_WITH_INFO(rawComm != nullptr, "NCCL window allocator returned a buffer without its communicator"); // Create custom deleter that releases the buffer - auto deleter = [comm, ptr = buffer.ptr](void*) { NCCLWindowAllocator::getInstance().releaseBuffer(*comm, ptr); }; + auto deleter = [comm = std::move(comm), rawComm, ptr = buffer.ptr](void*) + { + // Keep a generic caller-owned communicator alive until its window + // tensor releases the buffer; rawComm remains the stable pool key. + static_cast(comm); + NCCLWindowAllocator::getInstance().releaseBuffer(rawComm, ptr); + }; // Create tensor from the buffer auto tensor = torch::from_blob(buffer.ptr, shape, strides_vec, deleter, torch::dtype(dtype).device(torch::kCUDA)); diff --git a/cpp/tensorrt_llm/common/opUtils.cpp b/cpp/tensorrt_llm/common/opUtils.cpp index 35236bd1c8b4..900c2abffe67 100644 --- a/cpp/tensorrt_llm/common/opUtils.cpp +++ b/cpp/tensorrt_llm/common/opUtils.cpp @@ -20,15 +20,26 @@ #include "tensorrt_llm/runtime/ipcNvlsMemory.h" #include "tensorrt_llm/runtime/utils/mpiTags.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" +#include "tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h" #include "cuda.h" #include #include #include +#include +#include +#include +#include #include +#include +#include #include +#include #include +#include +#include TRTLLM_NAMESPACE_BEGIN #if ENABLE_MULTI_DEVICE @@ -52,16 +63,55 @@ std::unordered_map* getDtypeMap() namespace { -// Get NCCL unique ID for a group of ranks. -ncclUniqueId getUniqueId(std::set const& group) +using Clock = std::chrono::steady_clock; + +std::chrono::milliseconds getPositiveDurationFromEnv(char const* name, int64_t defaultMs) +{ + auto const* value = std::getenv(name); + if (value == nullptr) + { + return std::chrono::milliseconds{defaultMs}; + } + try + { + auto const parsed = std::stoll(value); + if (parsed > 0) + { + return std::chrono::milliseconds{parsed}; + } + } + catch (...) + { + } + TLLM_LOG_WARNING("Ignoring invalid %s=%s; expected a positive integer in milliseconds", name, value); + return std::chrono::milliseconds{defaultMs}; +} + +std::string groupToString(std::set const& group) +{ + std::ostringstream oss; + for (auto it = group.begin(); it != group.end(); ++it) + { + if (it != group.begin()) + { + oss << ','; + } + oss << *it; + } + return oss.str(); +} + +ncclUniqueId getLegacyUniqueId(std::set const& group) { auto const rank = COMM_SESSION.getRank(); - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, rank); - ncclUniqueId id; + ncclUniqueId id{}; if (rank == *group.begin()) { - NCCLCHECK_THROW(ncclGetUniqueId(&id)); - for (auto it = std::next(std::begin(group), 1); it != group.end(); ++it) + { + auto hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + NCCLCHECK_THROW(ncclGetUniqueId(&id)); + } + for (auto it = std::next(group.begin()); it != group.end(); ++it) { COMM_SESSION.sendValue(id, *it, tensorrt_llm::mpi::MpiTag::kDefault); } @@ -70,115 +120,1357 @@ ncclUniqueId getUniqueId(std::set const& group) { COMM_SESSION.recvValue(id, *group.begin(), tensorrt_llm::mpi::MpiTag::kDefault); } - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, rank); return id; } -} // namespace -std::shared_ptr getComm(std::set const& group) +void configureNcclEnvironment() { - auto const rank = COMM_SESSION.getRank(); - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, rank); - static std::map, std::shared_ptr> commMap; - static std::mutex mutex; - std::lock_guard lock(mutex); - std::ostringstream oss; - int index = 0; - for (auto const& rank : group) +#if defined(_WIN32) + // Need static connection initialization for accurate KV cache size estimation + if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) + _putenv_s("NCCL_RUNTIME_CONNECT", "0"); + // Disable graph register to avoid startup hangs + if (getenv("NCCL_GRAPH_REGISTER") == nullptr) + _putenv_s("NCCL_GRAPH_REGISTER", "0"); +#else + setenv("NCCL_RUNTIME_CONNECT", "0", 0); + setenv("NCCL_GRAPH_REGISTER", "0", 0); + // NCCL aborts during init if it tries NVLS multicast but the fabric/IMEX + // plane can't bind it. Disable NVLS when it isn't actually usable so NCCL + // falls back to NVLink P2P. No-overwrite preserves an explicit user setting. + if (!tensorrt_llm::runtime::ipcNvlsSupported()) + { + setenv("NCCL_NVLS_ENABLE", "0", 0); + } +#endif // _WIN32 +} + +class NcclCommState +{ +public: + explicit NcclCommState(std::set group, + std::shared_ptr controlComm = nullptr, bool recovery = false, + std::uint64_t rendezvousId = 0) + : mGroup(std::move(group)) + , mControlComm(std::move(controlComm)) + , mWorldRank(mControlComm != nullptr ? mControlComm->worldRank() : COMM_SESSION.getRank()) + , mFaultToleranceEnabled(tensorrt_llm::isNcclFaultToleranceEnabled()) + , mPollInterval(mFaultToleranceEnabled + ? getPositiveDurationFromEnv("TRTLLM_NCCL_WATCHDOG_POLL_INTERVAL_MS", 100) + : std::chrono::milliseconds{100}) + , mOperationTimeout(mFaultToleranceEnabled + ? getPositiveDurationFromEnv("TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS", 5000) + : std::chrono::milliseconds{5000}) + , mInitTimeout(mFaultToleranceEnabled + ? (recovery ? getPositiveDurationFromEnv("TRTLLM_NCCL_RECOVERY_TIMEOUT_MS", 5000) + : getPositiveDurationFromEnv("TRTLLM_NCCL_INIT_TIMEOUT_MS", 120000)) + : std::chrono::milliseconds{120000}) { - if (index != 0) + TLLM_CHECK_WITH_INFO(!recovery || mFaultToleranceEnabled, + "NCCL error: recovery communicator construction requires TLLM_FAULT_TOLERANCE_MODE=1"); + if (mFaultToleranceEnabled && mControlComm == nullptr) { - oss << ","; + std::vector const initialRanks(mGroup.begin(), mGroup.end()); + mControlComm = tensorrt_llm::runtime::createNcclUniqueIdRendezvousComm( + initialRanks, mWorldRank, COMM_SESSION, static_cast(tensorrt_llm::mpi::MpiTag::kNcclCommControl)); + } + initialize(rendezvousId); + if (!mFaultToleranceEnabled) + { + return; + } + try + { + mWatchdog = std::thread([this]() { watchdogLoop(); }); + } + catch (...) + { + ncclComm_t abortedComm = nullptr; + if (mComm != nullptr) + { + auto const comm = mComm; + ncclResult_t abortResult; + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + abortResult = ncclCommAbort(comm); + if (abortResult == ncclSuccess) + { + common::nccl_util::NcclCommResourceManager::getInstance().beginAbortCleanup(comm); + } + } + if (abortResult == ncclSuccess) + { + mComm = nullptr; + abortedComm = comm; + } + else + { + TLLM_LOG_ERROR("NCCL communicator leaked after watchdog thread creation and abort failed: %s", + ncclGetErrorString(abortResult)); + } + } + if (abortedComm != nullptr) + { + static_cast( + common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(abortedComm, true)); + } + throw; } - oss << rank; - index++; } - auto groupStr = oss.str(); - auto it = commMap.find(group); - if (it != commMap.end()) + + ~NcclCommState() { - auto ncclComm = it->second; - TLLM_LOG_TRACE("NCCL comm for group(%s) is cached for rank %d", groupStr.c_str(), rank); - return ncclComm; + { + std::lock_guard lock(mMutex); + mStop = true; + } + mWakeup.notify_all(); + if (mWatchdog.joinable()) + { + mWatchdog.join(); + } + + cleanupAbortedResources(); + + std::lock_guard lock(mMutex); + if (mComm != nullptr) + { + auto const comm = mComm; + mComm = nullptr; + if (mError.has_value() || mFaultToleranceEnabled) + { + ncclResult_t abortResult; + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + abortResult = ncclCommAbort(comm); + if (abortResult == ncclSuccess) + { + common::nccl_util::NcclCommResourceManager::getInstance().beginAbortCleanup(comm); + } + } + if (abortResult == ncclSuccess) + { + static_cast( + common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(comm, true)); + quarantinePendingOperationsLocked(); + } + else + { + // Do not release registered buffers while a failed abort + // may have left kernels referencing them. + try + { + TLLM_LOG_WARNING("Leaking failed NCCL communicator for group (%s) after abort error: %s", + groupToString(mGroup).c_str(), ncclGetErrorString(abortResult)); + } + catch (...) + { + } + } + if (abortResult != ncclSuccess) + { + quarantinePendingOperationsLocked(); + } + destroyPooledEventsLocked(); + return; + } + bool const resourcesCleaned + = common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(comm); + if (!resourcesCleaned) + { + ncclResult_t abortResult; + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + abortResult = ncclCommAbort(comm); + } + if (abortResult != ncclSuccess) + { + try + { + TLLM_LOG_WARNING( + "NCCL resource cleanup did not complete and ncclCommAbort failed for group (%s): %s", + groupToString(mGroup).c_str(), ncclGetErrorString(abortResult)); + } + catch (...) + { + } + } + quarantinePendingOperationsLocked(); + destroyPooledEventsLocked(); + return; + } + ncclResult_t result; + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(comm, &asyncResult); + if (queryResult != ncclSuccess || asyncResult != ncclSuccess) + { + // The watcher is already stopped. Re-sample immediately + // before destroy so a late asynchronous error cannot turn + // ncclCommDestroy into an unbounded blocking call. + result = ncclCommAbort(comm); + if (result != ncclSuccess) + { + try + { + TLLM_LOG_WARNING("Final NCCL health check failed and abort failed for group (%s): %s", + groupToString(mGroup).c_str(), ncclGetErrorString(result)); + } + catch (...) + { + } + } + quarantinePendingOperationsLocked(); + destroyPooledEventsLocked(); + return; + } + result = ncclCommDestroy(comm); + } + if (result != ncclSuccess) + { + try + { + TLLM_LOG_WARNING("ncclCommDestroy failed for group (%s): %s", groupToString(mGroup).c_str(), + ncclGetErrorString(result)); + } + catch (...) + { + } + } + if (result == ncclSuccess) + { + destroyCompletedWatchdogEventsLocked(); + } + else + { + quarantinePendingOperationsLocked(); + destroyPooledEventsLocked(); + } + return; + } + destroyPooledEventsLocked(); } - TLLM_LOG_TRACE("Init NCCL comm for group(%s) for rank %d", groupStr.c_str(), rank); - ncclUniqueId id = getUniqueId(group); - int groupRank = 0; - for (auto const& currentRank : group) + NcclCommState(NcclCommState const&) = delete; + NcclCommState& operator=(NcclCommState const&) = delete; + + ncclComm_t requireCommLocked() const { - if (rank == currentRank) - break; - ++groupRank; + if (mError.has_value()) + { + TLLM_THROW("NCCL error: communicator was aborted: %s", mError->c_str()); + } + TLLM_CHECK_WITH_INFO(mComm != nullptr, "NCCL error: communicator was aborted"); + return mComm; } - TLLM_CHECK(static_cast(groupRank) < group.size()); - std::shared_ptr ncclComm(new ncclComm_t, - [](ncclComm_t* comm) + + void checkResultLocked(ncclResult_t result, char const* operation) + { + if (result == ncclSuccess) { - if (!comm) + return; + } + + auto const deadline = Clock::now() + mOperationTimeout; + while (result == ncclInProgress && Clock::now() < deadline) + { + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(requireCommLocked(), &asyncResult); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + result = queryResult; + break; + } + if (queryResult == ncclSuccess && asyncResult == ncclSuccess) { return; } + if (queryResult == ncclSuccess && asyncResult != ncclInProgress) + { + result = asyncResult; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } + + std::string message; + if (result == ncclInProgress) + { + message = std::string(operation) + " timed out while the NCCL communicator was in progress"; + } + else + { + message = std::string(operation) + " failed with NCCL error: " + ncclGetErrorString(result); + } + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + + ncclResult_t checkOptionalResultLocked(ncclResult_t result, char const* operation) + { + if (result == ncclSuccess || result == ncclInvalidArgument) + { + return result; + } + if (result != ncclInProgress) + { + checkResultLocked(result, operation); + } - // STEP 1: Clean up resources and destroy NCCL communicator if it's valid - if (*comm) + auto const deadline = Clock::now() + mOperationTimeout; + while (Clock::now() < deadline) + { + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(requireCommLocked(), &asyncResult); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + checkResultLocked(queryResult, "ncclCommGetAsyncError(optional NCCL operation)"); + } + if (queryResult == ncclSuccess) { - // Clean up all registered resources FIRST - // The cleanupResources function uses a destruction guard to safely handle - // static destruction order issues - it will return early if the singleton - // is being destroyed (in which case the destructor handles cleanup proactively) - tensorrt_llm::common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(*comm); + if (asyncResult == ncclSuccess) + { + return asyncResult; + } + if (asyncResult != ncclInProgress) + { + // ncclInvalidArgument is a benign capability fallback only + // when the optional API returns it synchronously. Once the + // call reports ncclInProgress, every eventual non-success + // result is an asynchronous communicator failure. + checkResultLocked(asyncResult, operation); + } + } + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } + + auto const message = std::string(operation) + " timed out while the optional NCCL operation was in progress"; + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } - // Now destroy the NCCL communicator - ncclResult_t result = ncclCommDestroy(*comm); - if (result != ncclSuccess) + void checkEnqueueResultLocked(ncclResult_t result, char const* operation) + { + if (result == ncclSuccess || result == ncclInProgress) + { + return; + } + auto const message = std::string(operation) + " failed with NCCL error: " + ncclGetErrorString(result); + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + + void waitOperation(std::uint64_t token, char const* operation) + { + TLLM_CHECK_WITH_INFO(token != 0, "NCCL error: watchdog operation token must not be zero in FT mode"); + + while (true) + { + std::string pendingError; + { + std::lock_guard lock(mMutex); + if (mError.has_value()) { - // Logging may fail during static destruction, so wrap in try-catch - try + pendingError = *mError; + } + else + { + TLLM_CHECK_WITH_INFO(token < mNextOperationToken, + "NCCL error: unknown watchdog operation token %llu", static_cast(token)); + auto const pending = std::find_if(mPendingOperations.begin(), mPendingOperations.end(), + [token](PendingOperation const& candidate) { return candidate.token == token; }); + // The watchdog may have observed and retired a fast + // operation between track() releasing its lease and this + // waiter acquiring the state mutex. + if (pending == mPendingOperations.end()) { - TLLM_LOG_WARNING("ncclCommDestroy failed with error: %d", result); + return; } - catch (...) + TLLM_CHECK_WITH_INFO(pending->completion != nullptr, + "NCCL error: watchdog operation token %llu has not been tracked", + static_cast(token)); + + auto const now = Clock::now(); + if (pending->start != nullptr) + { + auto const startResult = cudaEventQuery(pending->start); + if (startResult == cudaSuccess) + { + recycleEventLocked(pending->start); + pending->start = nullptr; + pending->deadline = now + mOperationTimeout; + pending->armed = true; + } + else if (startResult != cudaErrorNotReady) + { + pendingError = pending->name + " start marker failed: " + cudaGetErrorString(startResult); + } + } + + bool completionReady = false; + if (pendingError.empty() && pending->start == nullptr) + { + auto const completionResult = cudaEventQuery(pending->completion); + if (completionResult == cudaSuccess) + { + completionReady = true; + } + else if (completionResult != cudaErrorNotReady) + { + pendingError + = pending->name + " failed with CUDA error: " + cudaGetErrorString(completionResult); + } + else if (pending->armed && now >= pending->deadline) + { + pendingError = pending->name + " timed out before its CUDA completion marker"; + } + } + + bool communicatorReady = false; + if (pendingError.empty()) + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(requireCommLocked(), &asyncResult); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + pendingError = std::string(operation) + + " failed while querying NCCL: " + ncclGetErrorString(queryResult); + } + else if (queryResult == ncclSuccess && asyncResult != ncclSuccess + && asyncResult != ncclInProgress) + { + pendingError = std::string(operation) + + " failed with NCCL error: " + ncclGetErrorString(asyncResult); + } + else + { + communicatorReady = queryResult == ncclSuccess && asyncResult == ncclSuccess; + } + } + + if (pendingError.empty() && completionReady && communicatorReady) { - // Ignore logging failures during static destruction + recycleEventLocked(pending->completion); + mPendingOperations.erase(pending); + return; + } + if (!pendingError.empty()) + { + setErrorAndAbortLocked(pendingError); + pendingError = mError.value_or(pendingError); } } + } - // Clear the communicator value before freeing the pointer - *comm = nullptr; + if (!pendingError.empty()) + { + cleanupAbortedResources(); + TLLM_THROW("NCCL error: communicator was aborted: %s", pendingError.c_str()); } + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } + } - // STEP 2: Always free the pointer memory (regardless of whether *comm was valid) - delete comm; - }); -#if defined(_WIN32) - // Need static connection initialization for accurate KV cache size estimation - if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) - _putenv_s("NCCL_RUNTIME_CONNECT", "0"); - // Disable graph register to avoid startup hangs - if (getenv("NCCL_GRAPH_REGISTER") == nullptr) - _putenv_s("NCCL_GRAPH_REGISTER", "0"); -#else - setenv("NCCL_RUNTIME_CONNECT", "0", 0); - setenv("NCCL_GRAPH_REGISTER", "0", 0); - // NCCL aborts during init if it tries NVLS multicast but the fabric/IMEX - // plane can't bind it. Disable NVLS when it isn't actually usable so NCCL - // falls back to NVLink P2P. No-overwrite preserves an explicit user setting. - if (!tensorrt_llm::runtime::ipcNvlsSupported()) + uint64_t beginOperationLocked(cudaStream_t stream, char const* operation) { - setenv("NCCL_NVLS_ENABLE", "0", 0); + requireCommLocked(); + if (!mFaultToleranceEnabled) + { + return 0; + } + + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + auto const captureResult = cudaStreamIsCapturing(stream, &captureStatus); + if (captureResult != cudaSuccess) + { + auto const message = std::string(operation) + + " failed to query CUDA stream capture state: " + cudaGetErrorString(captureResult); + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + if (captureStatus != cudaStreamCaptureStatusNone) + { + TLLM_CHECK_WITH_INFO(!mFaultToleranceEnabled, + "NCCL error: CUDA graph capture of managed NCCL operations is disabled while " + "TLLM_FAULT_TOLERANCE_MODE=1; captured graph nodes retain the old communicator across recovery"); + // An event recorded while capturing belongs to the graph and may + // not execute until long after capture ends. Arming a wall-clock + // deadline here would therefore report a false timeout before the + // graph is launched. Non-FT graph launches retain legacy async + // error monitoring; FT mode above requires eager launches so every + // operation receives the completion deadline below. + TLLM_LOG_DEBUG("Skipping eager NCCL completion deadline for %s during CUDA graph capture", operation); + return 0; + } + + if (mPendingOperations.size() >= kMaxPendingOperations) + { + auto const message = std::string(operation) + " exceeded the bounded NCCL watchdog operation queue"; + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + + auto const token = mNextOperationToken++; + cudaEvent_t start = acquireEventLocked(); + mPendingOperations.push_back({token, start, nullptr, {}, false, operation}); + auto const result = cudaEventRecord(start, stream); + if (result != cudaSuccess) + { + auto const message + = std::string(operation) + " failed to record its start marker: " + cudaGetErrorString(result); + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + + mWakeup.notify_all(); + return token; } -#endif // _WIN32 -#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) - ncclConfig_t config = NCCL_CONFIG_INITIALIZER; - config.graphUsageMode = 1; - NCCLCHECK_THROW(ncclCommInitRankConfig(ncclComm.get(), group.size(), id, groupRank, &config)); -#else - NCCLCHECK_THROW(ncclCommInitRank(ncclComm.get(), group.size(), id, groupRank)); -#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) - commMap[group] = ncclComm; - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, rank); - return ncclComm; + + void finishOperationLocked(uint64_t token, cudaStream_t stream) + { + requireCommLocked(); + if (token == 0) + { + return; + } + + auto const operation = std::find_if(mPendingOperations.begin(), mPendingOperations.end(), + [token](PendingOperation const& pending) { return pending.token == token; }); + TLLM_CHECK_WITH_INFO(operation != mPendingOperations.end(), "NCCL error: unknown watchdog operation token %llu", + static_cast(token)); + TLLM_CHECK_WITH_INFO(operation->completion == nullptr, + "NCCL error: watchdog operation token %llu was already completed", static_cast(token)); + + operation->completion = acquireEventLocked(); + auto const result = cudaEventRecord(operation->completion, stream); + if (result != cudaSuccess) + { + auto const message + = operation->name + " failed to record its completion marker: " + cudaGetErrorString(result); + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + mWakeup.notify_all(); + } + + void abort(std::string const& reason) + { + { + std::lock_guard lock(mMutex); + auto hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + setErrorAndAbortLocked(reason); + TLLM_CHECK_WITH_INFO(mComm == nullptr, + "NCCL error: communicator abort failed; refusing to initialize a replacement while the old handle " + "is live"); + } + cleanupAbortedResources(); + } + + std::optional getError() const + { + std::lock_guard lock(mMutex); + return mError; + } + + std::shared_ptr const& controlComm() const noexcept + { + return mControlComm; + } + + int worldRank() const noexcept + { + return mWorldRank; + } + + ncclComm_t* slot() noexcept + { + return &mComm; + } + + std::mutex& mutex() noexcept + { + return mMutex; + } + + void cleanupAbortedResources() noexcept + { + std::vector abortedComms; + { + std::lock_guard lock(mMutex); + abortedComms.swap(mPendingAbortCleanup); + } + for (auto const comm : abortedComms) + { + static_cast(common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(comm, true)); + } + } + +private: + void initialize(std::uint64_t rendezvousId) + { + configureNcclEnvironment(); + auto const rank = mWorldRank; + TLLM_CHECK_WITH_INFO(!mGroup.empty(), "NCCL communicator group must not be empty"); + auto const rankIt = mGroup.find(rank); + TLLM_CHECK_WITH_INFO(rankIt != mGroup.end(), "Global rank %d is not in NCCL communicator group (%s)", rank, + groupToString(mGroup).c_str()); + + auto const groupRank = static_cast(std::distance(mGroup.begin(), rankIt)); + if (!mFaultToleranceEnabled) + { + auto const id = getLegacyUniqueId(mGroup); + auto hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + ncclComm_t comm = nullptr; +#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) + ncclConfig_t config = NCCL_CONFIG_INITIALIZER; + config.graphUsageMode = 1; + auto const result = ncclCommInitRankConfig(&comm, static_cast(mGroup.size()), id, groupRank, &config); +#else + auto const result = ncclCommInitRank(&comm, static_cast(mGroup.size()), id, groupRank); +#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) + if (result != ncclSuccess) + { + if (comm != nullptr) + { + static_cast(ncclCommAbort(comm)); + } + TLLM_THROW("NCCL communicator initialization failed for group (%s) with NCCL error: %s", + groupToString(mGroup).c_str(), ncclGetErrorString(result)); + } + hostApiLock.unlock(); + common::nccl_util::NcclCommResourceManager::getInstance().waitForAbortCleanup(comm); + mComm = comm; + return; + } + + TLLM_CHECK_WITH_INFO(mControlComm != nullptr, "NCCL rendezvous control channel is not initialized"); + TLLM_CHECK_WITH_INFO(rank == mControlComm->worldRank(), + "NCCL error: MPI rank %d does not match rendezvous world rank %d", rank, mControlComm->worldRank()); + auto const deadline = Clock::now() + mInitTimeout; + std::vector const group(mGroup.begin(), mGroup.end()); + auto const id = tensorrt_llm::runtime::exchangeNcclUniqueId(group, *mControlComm, + {tensorrt_llm::mpi::MpiTag::kNcclCommReady, tensorrt_llm::mpi::MpiTag::kNcclCommUniqueId, + tensorrt_llm::mpi::MpiTag::kNcclCommAck}, + rendezvousId, deadline); + ncclConfig_t config = NCCL_CONFIG_INITIALIZER; + config.blocking = 0; +#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) + config.graphUsageMode = 1; +#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) + + auto hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + ncclComm_t comm = nullptr; + auto result = ncclCommInitRankConfig(&comm, static_cast(mGroup.size()), id, groupRank, &config); + while (result == ncclInProgress && Clock::now() < deadline) + { + ncclResult_t asyncResult = ncclInProgress; + auto const queryResult = ncclCommGetAsyncError(comm, &asyncResult); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + result = queryResult; + break; + } + if (queryResult == ncclSuccess && asyncResult == ncclSuccess) + { + result = ncclSuccess; + break; + } + if (queryResult == ncclSuccess && asyncResult != ncclInProgress) + { + result = asyncResult; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } + + if (result != ncclSuccess) + { + if (comm != nullptr) + { + ncclCommAbort(comm); + } + if (result == ncclInProgress) + { + TLLM_THROW( + "NCCL error: communicator initialization timed out for group (%s)", groupToString(mGroup).c_str()); + } + TLLM_THROW("NCCL communicator initialization failed for group (%s) with NCCL error: %s", + groupToString(mGroup).c_str(), ncclGetErrorString(result)); + } + hostApiLock.unlock(); + common::nccl_util::NcclCommResourceManager::getInstance().waitForAbortCleanup(comm); + mComm = comm; + } + + void setErrorAndAbortLocked(std::string const& reason) + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + if (!mError.has_value()) + { + mError = reason; + } + if (mComm == nullptr) + { + return; + } + + auto const comm = mComm; + // Abort first: graceful window cleanup synchronizes the entire device + // and can hang forever while a failed NCCL kernel is still resident. + auto const abortResult = ncclCommAbort(comm); + if (abortResult == ncclSuccess) + { + common::nccl_util::NcclCommResourceManager::getInstance().beginAbortCleanup(comm); + mComm = nullptr; + quarantinePendingOperationsLocked(); + // Cleanup is deferred until the operation lease and any allocator + // lock have unwound. Running callbacks here can recursively enter + // NCCLWindowAllocator and deadlock its non-recursive mutex. + mPendingAbortCleanup.push_back(comm); + } + else + { + // Keep both the handle and registered buffers alive. Releasing + // them while NCCL kernels may still reference them is unsafe; a + // later control-path call can retry the idempotent abort. + TLLM_LOG_WARNING("ncclCommAbort failed for group (%s): %s", groupToString(mGroup).c_str(), + ncclGetErrorString(abortResult)); + } + } + + void quarantinePendingOperationsLocked() noexcept + { + if (!mPendingOperations.empty()) + { + // cudaEventDestroy is asynchronous, but the event may still be + // referenced by work from the failed communicator. Keep the + // handles alive for the rest of the process instead of making + // assumptions about CUDA progress after abort. + TLLM_LOG_WARNING( + "Quarantining %zu NCCL completion events after communicator abort", mPendingOperations.size()); + } + mPendingOperations.clear(); + } + + void destroyCompletedWatchdogEventsLocked() noexcept + { + size_t quarantined = 0; + for (auto const& operation : mPendingOperations) + { + for (auto const event : {operation.start, operation.completion}) + { + if (event == nullptr) + { + continue; + } + if (cudaEventQuery(event) == cudaSuccess) + { + static_cast(cudaEventDestroy(event)); + } + else + { + ++quarantined; + } + } + } + mPendingOperations.clear(); + if (quarantined != 0) + { + TLLM_LOG_WARNING("Quarantining %zu incomplete NCCL watchdog events during healthy teardown", quarantined); + } + for (auto const event : mEventPool) + { + static_cast(cudaEventDestroy(event)); + } + mEventPool.clear(); + } + + void destroyPooledEventsLocked() noexcept + { + for (auto const event : mEventPool) + { + static_cast(cudaEventDestroy(event)); + } + mEventPool.clear(); + } + + cudaEvent_t acquireEventLocked() + { + if (!mEventPool.empty()) + { + auto const event = mEventPool.back(); + mEventPool.pop_back(); + return event; + } + + cudaEvent_t event = nullptr; + auto const result = cudaEventCreateWithFlags(&event, cudaEventDisableTiming); + if (result != cudaSuccess) + { + auto const message + = std::string("failed to allocate an NCCL watchdog event: ") + cudaGetErrorString(result); + setErrorAndAbortLocked(message); + TLLM_THROW("NCCL error: communicator was aborted: %s", message.c_str()); + } + return event; + } + + void recycleEventLocked(cudaEvent_t event) noexcept + { + if (event == nullptr) + { + return; + } + if (mEventPool.size() < kMaxPooledEvents) + { + mEventPool.push_back(event); + } + else + { + // The event has completed, so releasing pool overflow cannot race + // device work. Keeping the pool bounded avoids unbounded steady- + // state driver resources under bursty workloads. + static_cast(cudaEventDestroy(event)); + } + } + + void watchdogLoop() noexcept + { + while (true) + { + std::unique_lock lock(mMutex); + if (mWakeup.wait_for(lock, mPollInterval, [this]() { return mStop; })) + { + return; + } + if (mComm == nullptr) + { + if (mError.has_value()) + { + lock.unlock(); + cleanupAbortedResources(); + return; + } + continue; + } + if (mError.has_value()) + { + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + setErrorAndAbortLocked(*mError); + } + lock.unlock(); + cleanupAbortedResources(); + continue; + } + + std::string pendingError; + auto const now = Clock::now(); + for (auto operation = mPendingOperations.begin(); operation != mPendingOperations.end();) + { + if (operation->start != nullptr) + { + auto const startResult = cudaEventQuery(operation->start); + if (startResult == cudaSuccess) + { + recycleEventLocked(operation->start); + operation->start = nullptr; + operation->deadline = now + mOperationTimeout; + operation->armed = true; + } + else if (startResult != cudaErrorNotReady) + { + pendingError = operation->name + " start marker failed: " + cudaGetErrorString(startResult); + break; + } + else + { + ++operation; + continue; + } + } + + if (operation->completion == nullptr) + { + if (operation->armed && now >= operation->deadline) + { + pendingError = operation->name + " timed out before its completion marker was recorded"; + break; + } + ++operation; + continue; + } + auto const eventResult = cudaEventQuery(operation->completion); + if (eventResult == cudaSuccess) + { + recycleEventLocked(operation->completion); + operation = mPendingOperations.erase(operation); + continue; + } + if (eventResult != cudaErrorNotReady) + { + pendingError = operation->name + " failed with CUDA error: " + cudaGetErrorString(eventResult); + break; + } + if (operation->armed && now >= operation->deadline) + { + pendingError = operation->name + " timed out before its CUDA completion marker"; + break; + } + ++operation; + } + + { + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + if (!pendingError.empty()) + { + setErrorAndAbortLocked(pendingError); + } + else + { + ncclResult_t asyncResult = ncclSuccess; + auto const queryResult = ncclCommGetAsyncError(mComm, &asyncResult); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + setErrorAndAbortLocked(std::string("ncclCommGetAsyncError failed with NCCL error: ") + + ncclGetErrorString(queryResult)); + } + else if (asyncResult != ncclSuccess && asyncResult != ncclInProgress) + { + setErrorAndAbortLocked(std::string("NCCL watchdog observed asynchronous NCCL error: ") + + ncclGetErrorString(asyncResult)); + } + } + } + lock.unlock(); + cleanupAbortedResources(); + } + } + + std::set mGroup; + std::shared_ptr mControlComm; + int const mWorldRank; + bool mFaultToleranceEnabled; + ncclComm_t mComm{nullptr}; + mutable std::mutex mMutex; + std::condition_variable mWakeup; + bool mStop{false}; + std::optional mError; + std::vector mPendingAbortCleanup; + + struct PendingOperation + { + uint64_t token; + cudaEvent_t start; + cudaEvent_t completion; + Clock::time_point deadline; + bool armed; + std::string name; + }; + + std::vector mPendingOperations; + std::vector mEventPool; + uint64_t mNextOperationToken{1}; + static constexpr size_t kMaxPendingOperations = 4096; + static constexpr size_t kMaxPooledEvents = 256; + std::chrono::milliseconds mPollInterval; + std::chrono::milliseconds mOperationTimeout; + std::chrono::milliseconds mInitTimeout; + std::thread mWatchdog; +}; + +class NcclCommRegistry +{ +public: + static NcclCommRegistry& instance() + { + // Raw communicators, watchdogs, and their CUDA/MPI dependencies are + // process-lifetime resources. Avoid cross-translation-unit static + // destruction where the resource manager or allocator could disappear + // before a live watchdog/state attempts final cleanup. + static auto* registry = new NcclCommRegistry; + return *registry; + } + + std::shared_ptr get(std::set const& group) + { + std::lock_guard lock(mMutex); + auto it = mComms.find(group); + if (it == mComms.end()) + { + auto state = std::make_shared(group); + mSlots[state->slot()] = state; + it = mComms.emplace(group, std::move(state)).first; + } + return std::shared_ptr{it->second, it->second->slot()}; + } + + std::shared_ptr find(std::shared_ptr const& comm) + { + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator holder is null"); + std::lock_guard lock(mMutex); + auto const it = mSlots.find(comm.get()); + TLLM_CHECK_WITH_INFO(it != mSlots.end(), "NCCL communicator is not managed by the TRT-LLM registry"); + auto state = it->second.lock(); + TLLM_CHECK_WITH_INFO(state != nullptr, "NCCL communicator registry entry has expired"); + return state; + } + + std::shared_ptr find(std::set const& group) + { + std::lock_guard lock(mMutex); + auto const it = mComms.find(group); + TLLM_CHECK_WITH_INFO( + it != mComms.end(), "No cached NCCL communicator exists for group (%s)", groupToString(group).c_str()); + return it->second; + } + + void abortAndReinit(std::set const& oldGroup, std::set const& activeGroup, std::uint64_t rendezvousId) + { + std::lock_guard const recoveryLock(mRecoveryMutex); + TLLM_CHECK_WITH_INFO(rendezvousId != 0, + "NCCL error: recovery rendezvous ID must be nonzero; zero is reserved for initial bootstrap"); + TLLM_CHECK_WITH_INFO(!activeGroup.empty(), "active NCCL rank group must not be empty"); + TLLM_CHECK_WITH_INFO(std::includes(oldGroup.begin(), oldGroup.end(), activeGroup.begin(), activeGroup.end()), + "active NCCL rank group (%s) must be a subset of old group (%s)", groupToString(activeGroup).c_str(), + groupToString(oldGroup).c_str()); + TLLM_CHECK_WITH_INFO(tensorrt_llm::isNcclFaultToleranceEnabled(), + "NCCL error: communicator reinitialization requires TLLM_FAULT_TOLERANCE_MODE=1; otherwise captured " + "CUDA graphs may retain the old communicator"); + std::shared_ptr oldState; + std::shared_ptr existingTargetState; + { + std::lock_guard lock(mMutex); + auto const old = mComms.find(oldGroup); + if (old != mComms.end()) + { + oldState = old->second; + } + auto const existingTarget = mComms.find(activeGroup); + if (existingTarget != mComms.end()) + { + existingTargetState = existingTarget->second; + } + } + TLLM_CHECK_WITH_INFO( + oldState != nullptr, "No cached NCCL communicator exists for group (%s)", groupToString(oldGroup).c_str()); + auto const& controlComm = oldState->controlComm(); + TLLM_CHECK_WITH_INFO(controlComm != nullptr, "NCCL rendezvous control channel is not initialized"); + auto const worldRank = controlComm->worldRank(); + TLLM_CHECK_WITH_INFO( + activeGroup.count(worldRank) != 0, "Failed rank %d must not call NCCL abort_and_reinit", worldRank); + oldState->abort("NCCL error: communicator was aborted for survivor reinitialization"); + if (existingTargetState != nullptr && existingTargetState != oldState) + { + // Explicit recovery is a distributed generation change. Never let + // one rank reuse a locally healthy cached target communicator + // while another rank enters the fresh-ID rendezvous. + existingTargetState->abort("NCCL error: cached target communicator was retired for reinitialization"); + } + + // Do not mutate oldState: existing operations retain it and must fail + // immediately instead of silently switching rank numbering underneath + // an in-flight call. Subsequent operations request activeGroup. + std::lock_guard lock(mMutex); + auto state = std::make_shared(activeGroup, controlComm, true, rendezvousId); + mSlots[state->slot()] = state; + mComms[activeGroup] = std::move(state); + } + + std::optional getError(std::set const& group) + { + std::shared_ptr state; + { + std::lock_guard lock(mMutex); + auto const it = mComms.find(group); + if (it == mComms.end()) + { + return std::nullopt; + } + state = it->second; + } + return state->getError(); + } + +private: + std::mutex mRecoveryMutex; + std::mutex mMutex; + std::map, std::shared_ptr> mComms; + std::unordered_map> mSlots; +}; + +} // namespace + +struct NcclCommLease::Impl +{ + explicit Impl(std::shared_ptr state_) + : state(std::move(state_)) + , lock(state->mutex()) + { + // Complete the call_once-protected capability probe before taking the + // host gate. The probe itself briefly takes that gate for + // ncclGetVersion; reversing the order can deadlock with a concurrent + // first-time capability query. + static_cast(common::nccl_util::isNcclWindowSupported()); + hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + state->requireCommLocked(); + } + + ~Impl() + { + if (groupActive) + { + auto const result = ncclGroupEnd(); + groupActive = false; + try + { + state->checkResultLocked(result, "ncclGroupEnd(exception cleanup)"); + } + catch (...) + { + // Destructors cannot propagate. checkResultLocked has already + // latched the classifier-friendly error and aborted the comm. + } + } + if (hostApiLock.owns_lock()) + { + hostApiLock.unlock(); + } + if (lock.owns_lock()) + { + lock.unlock(); + } + state->cleanupAbortedResources(); + } + + std::shared_ptr state; + std::unique_lock lock; + tensorrt_llm::runtime::NcclHostApiLock hostApiLock; + bool groupActive{false}; +}; + +NcclCommLease::NcclCommLease(std::unique_ptr impl) + : mImpl(std::move(impl)) +{ +} + +NcclCommLease::NcclCommLease(ncclComm_t legacyComm) + : mLegacyComm(legacyComm) +{ +} + +NcclCommLease::~NcclCommLease() +{ + if (mLegacyGroupActive) + { + auto const result = ncclGroupEnd(); + mLegacyGroupActive = false; + if (result != ncclSuccess) + { + TLLM_LOG_WARNING( + "Failed to close legacy NCCL group during exception cleanup: %s", ncclGetErrorString(result)); + } + } +} + +NcclCommLease::NcclCommLease(NcclCommLease&& other) noexcept + : mImpl(std::move(other.mImpl)) + , mLegacyComm(std::exchange(other.mLegacyComm, nullptr)) + , mLegacyGroupActive(std::exchange(other.mLegacyGroupActive, false)) +{ +} + +NcclCommLease& NcclCommLease::operator=(NcclCommLease&& other) noexcept +{ + if (this != &other) + { + if (mLegacyGroupActive) + { + static_cast(ncclGroupEnd()); + } + mImpl = std::move(other.mImpl); + mLegacyComm = std::exchange(other.mLegacyComm, nullptr); + mLegacyGroupActive = std::exchange(other.mLegacyGroupActive, false); + } + return *this; +} + +ncclComm_t NcclCommLease::get() const +{ + if (mImpl == nullptr) + { + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + return mLegacyComm; + } + return mImpl->state->requireCommLocked(); +} + +void NcclCommLease::check(ncclResult_t result, char const* operation) const +{ + if (mImpl == nullptr) + { + (void) operation; + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + NCCLCHECK_THROW(result); + return; + } + mImpl->state->checkResultLocked(result, operation); +} + +ncclResult_t NcclCommLease::checkOptional(ncclResult_t result, char const* operation) const +{ + if (mImpl == nullptr) + { + (void) operation; + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + if (result != ncclInvalidArgument) + { + NCCLCHECK_THROW(result); + } + return result; + } + return mImpl->state->checkOptionalResultLocked(result, operation); +} + +void NcclCommLease::checkEnqueue(ncclResult_t result, char const* operation) const +{ + if (mImpl == nullptr) + { + check(result, operation); + return; + } + if (result != ncclSuccess && result != ncclInProgress && mImpl->groupActive) + { + // Close the thread-local group before aborting. Otherwise the first + // collective on a replacement communicator can be nested inside the + // failed group. + static_cast(ncclGroupEnd()); + mImpl->groupActive = false; + } + mImpl->state->checkEnqueueResultLocked(result, operation); +} + +uint64_t NcclCommLease::begin(cudaStream_t stream, char const* operation) const +{ + if (mImpl == nullptr) + { + (void) stream; + (void) operation; + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + return 0; + } + return mImpl->state->beginOperationLocked(stream, operation); +} + +void NcclCommLease::track(uint64_t token, cudaStream_t stream) const +{ + if (mImpl == nullptr) + { + (void) token; + (void) stream; + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + return; + } + mImpl->state->finishOperationLocked(token, stream); +} + +void NcclCommLease::groupStart(char const* operation) const +{ + if (mImpl == nullptr) + { + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + TLLM_CHECK_WITH_INFO(!mLegacyGroupActive, "NCCL error: nested NCCL groups are not supported by this lease"); + (void) operation; + NCCLCHECK_THROW(ncclGroupStart()); + mLegacyGroupActive = true; + return; + } + TLLM_CHECK_WITH_INFO(!mImpl->groupActive, "NCCL error: nested NCCL groups are not supported by this lease"); + mImpl->state->checkResultLocked(ncclGroupStart(), operation); + mImpl->groupActive = true; +} + +void NcclCommLease::groupEnd(char const* operation) const +{ + if (mImpl == nullptr) + { + TLLM_CHECK_WITH_INFO(mLegacyComm != nullptr, "NCCL communicator lease was moved from"); + TLLM_CHECK_WITH_INFO(mLegacyGroupActive, "NCCL error: no NCCL group is active on this lease"); + auto const result = ncclGroupEnd(); + mLegacyGroupActive = false; + (void) operation; + NCCLCHECK_THROW(result); + return; + } + TLLM_CHECK_WITH_INFO(mImpl->groupActive, "NCCL error: no NCCL group is active on this lease"); + auto const result = ncclGroupEnd(); + mImpl->groupActive = false; + mImpl->state->checkResultLocked(result, operation); +} + +std::shared_ptr getComm(std::set const& group) +{ + auto const groupString = groupToString(group); + // Do not query COMM_SESSION here: after survivor reinitialization that + // communicator may include a dead process and retain its fatal handler. + TLLM_LOG_TRACE("Get NCCL comm for group(%s)", groupString.c_str()); + return NcclCommRegistry::instance().get(group); +} + +int getCommWorldRank(std::shared_ptr const& comm) +{ + return NcclCommRegistry::instance().find(comm)->worldRank(); +} + +bool isNcclFaultToleranceEnabled() +{ + // Communicator mode cannot change after the first managed communicator is + // used: its blocking configuration and watchdog ownership are fixed. + static bool const enabled = tensorrt_llm::mpi::isFaultToleranceModeEnabled(); + return enabled; +} + +NcclCommLease acquireComm(std::shared_ptr const& comm) +{ + // FT mode is a startup-only process setting. Keep the default-off path + // equivalent to the legacy blocking communicator: no registry lookup, + // heap allocation, communicator mutex, or process-wide NCCL gate per op. + if (!isNcclFaultToleranceEnabled()) + { + TLLM_CHECK_WITH_INFO(comm != nullptr && *comm != nullptr, "NCCL communicator is null"); + return NcclCommLease{*comm}; + } + auto state = NcclCommRegistry::instance().find(comm); + return NcclCommLease{std::make_unique(std::move(state))}; +} + +void waitCommOperation(std::shared_ptr const& comm, std::uint64_t token, char const* operation) +{ + if (!isNcclFaultToleranceEnabled()) + { + TLLM_CHECK_WITH_INFO(token == 0, "NCCL error: legacy communicator received a watchdog operation token"); + return; + } + auto state = NcclCommRegistry::instance().find(comm); + state->waitOperation(token, operation); +} + +void abortAndReinitComm(std::set const& oldGroup, std::set const& activeGroup, std::uint64_t rendezvousId) +{ + NcclCommRegistry::instance().abortAndReinit(oldGroup, activeGroup, rendezvousId); +} + +std::optional getCommAsyncError(std::set const& group) +{ + return NcclCommRegistry::instance().getError(group); +} + +std::optional getCommAsyncError(std::shared_ptr const& comm) +{ + auto state = NcclCommRegistry::instance().find(comm); + return state->getError(); } #endif // ENABLE_MULTI_DEVICE diff --git a/cpp/tensorrt_llm/common/opUtils.h b/cpp/tensorrt_llm/common/opUtils.h index 72e5a5ea3e09..306c9cffd2e1 100644 --- a/cpp/tensorrt_llm/common/opUtils.h +++ b/cpp/tensorrt_llm/common/opUtils.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -29,6 +29,8 @@ #include #endif // ENABLE_MULTI_DEVICE +#include +#include #include #include #include @@ -37,6 +39,7 @@ #include #include #include +#include #include "tensorrt_llm/common/nvmlWrapper.h" @@ -216,8 +219,100 @@ inline bool isBuilding() std::unordered_map* getDtypeMap(); +/// Access to a cached NCCL communicator. +/// +/// NCCL communicators used by the raw TRT-LLM collectives are configured as +/// nonblocking so that the fault-tolerance watchdog can abort them. NCCL does +/// not allow host API calls and ncclCommAbort to race on the same communicator, +/// so every FT raw collective must hold this serialized lease from its first +/// communicator call through ncclGroupEnd (when a group is used). When FT is +/// disabled the lease is an inline, non-locking wrapper around the legacy +/// blocking communicator. +class NcclCommLease +{ +public: + ~NcclCommLease(); + + NcclCommLease(NcclCommLease const&) = delete; + NcclCommLease& operator=(NcclCommLease const&) = delete; + NcclCommLease(NcclCommLease&&) noexcept; + NcclCommLease& operator=(NcclCommLease&&) noexcept; + + /// Return the native communicator, or throw a classifier-friendly error + /// when its watchdog has aborted it. + [[nodiscard]] ncclComm_t get() const; + + /// Complete a nonblocking NCCL host call. ncclInProgress is polled via + /// ncclCommGetAsyncError before another host API is issued. + void check(ncclResult_t result, char const* operation) const; + + /// Complete an optional NCCL capability call. An immediate + /// ncclInvalidArgument is returned as a documented, nonfatal unsupported + /// result. Once a call returns ncclInProgress, every eventual non-success + /// result is latched and aborted as an asynchronous communicator failure. + [[nodiscard]] ncclResult_t checkOptional(ncclResult_t result, char const* operation) const; + + /// Validate an operation queued inside ncclGroupStart/ncclGroupEnd. + /// ncclInProgress is completed by the enclosing ncclGroupEnd call. + void checkEnqueue(ncclResult_t result, char const* operation) const; + + /// Record a stream marker immediately before an asynchronous NCCL + /// operation. The timeout is armed only when this marker completes, so + /// unrelated work already queued on the stream cannot cause a false + /// timeout. + [[nodiscard]] uint64_t begin(cudaStream_t stream, char const* operation) const; + + /// Record the matching completion marker after the NCCL operation has + /// been enqueued. The watchdog aborts if it does not complete by the + /// deadline armed by begin(). A zero token denotes the default-off legacy + /// path. FT mode rejects capture so a recovered communicator can never be + /// bypassed by stale graph nodes. + void track(uint64_t token, cudaStream_t stream) const; + + /// Start/end an NCCL group tracked by this lease. If any intervening C++ + /// operation throws, the lease destructor closes the group before + /// releasing the shared NCCL host-API gate. + void groupStart(char const* operation) const; + void groupEnd(char const* operation) const; + +private: + struct Impl; + explicit NcclCommLease(std::unique_ptr impl); + explicit NcclCommLease(ncclComm_t legacyComm); + std::unique_ptr mImpl; + ncclComm_t mLegacyComm{nullptr}; + mutable bool mLegacyGroupActive{false}; + + friend NcclCommLease acquireComm(std::shared_ptr const& comm); +}; + std::shared_ptr getComm(std::set const& group); +/// Return the immutable original/global MPI rank associated with a managed communicator. +int getCommWorldRank(std::shared_ptr const& comm); + +/// Return the startup-cached raw-NCCL fault-tolerance mode. +bool isNcclFaultToleranceEnabled(); + +/// Acquire exclusive host-API access to a cached NCCL communicator. +NcclCommLease acquireComm(std::shared_ptr const& comm); + +/// Abort the communicator for oldGroup and initialize a fresh communicator +/// containing only activeGroup. Both groups contain global MPI ranks. Only +/// ranks in activeGroup may call this function. +void abortAndReinitComm(std::set const& oldGroup, std::set const& activeGroup, std::uint64_t rendezvousId); + +/// Wait for a begin()/track() operation without retaining its lease. The +/// operation timeout is armed only after its start marker reaches the head of +/// the stream, and other communicators remain free to enqueue or abort. +void waitCommOperation(std::shared_ptr const& comm, std::uint64_t token, char const* operation); + +/// Return the stable watchdog error for a cached communicator, if any. +std::optional getCommAsyncError(std::set const& group); + +/// Return the stable watchdog error for the communicator referenced by comm. +std::optional getCommAsyncError(std::shared_ptr const& comm); + #endif // ENABLE_MULTI_DEVICE //! To save GPU memory, all the plugins share the same cublas and cublasLt handle globally. diff --git a/cpp/tensorrt_llm/runtime/CMakeLists.txt b/cpp/tensorrt_llm/runtime/CMakeLists.txt index ca81fbb0f6cd..9c0c0f388b05 100644 --- a/cpp/tensorrt_llm/runtime/CMakeLists.txt +++ b/cpp/tensorrt_llm/runtime/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & # AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); you may not @@ -16,6 +16,7 @@ include(FetchContent) set(SRCS utils/mpiUtils.cpp + utils/ncclUniqueIdRendezvous.cpp utils/numpyUtils.cpp utils/runtimeUtils.cpp utils/debugUtils.cu diff --git a/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp b/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp index dc79fd5f48e2..51014d56a0a5 100644 --- a/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp +++ b/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -18,6 +18,18 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/runtime/utils/multiDeviceUtils.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" +#include "tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include #if ENABLE_MULTI_DEVICE #include @@ -27,6 +39,15 @@ using namespace tensorrt_llm::runtime; namespace { +using namespace std::chrono_literals; + +constexpr auto kInitialReadyTimeout = 120s; +#if ENABLE_MULTI_DEVICE +constexpr auto kReadyPollInterval = 1ms; +constexpr auto kRecoveryReadyTimeout = 5s; +constexpr auto kWatcherPollInterval = 100ms; +#endif + #if ENABLE_MULTI_DEVICE ncclDataType_t toNcclType(nvinfer1::DataType dataType) @@ -43,19 +64,154 @@ ncclDataType_t toNcclType(nvinfer1::DataType dataType) #if ENABLE_BF16 case nvinfer1::DataType::kBF16: return ncclBfloat16; #endif // ENABLE_BF16 - default: TLLM_THROW("Unsupported data type: %d", static_cast(dataType)); + default: TLLM_THROW("NCCL error: unsupported data type: %d", static_cast(dataType)); } } + +std::string ncclErrorMessage(char const* operation, ncclResult_t result) +{ + std::ostringstream message; + message << "NCCL error during " << operation << ": " << ncclGetErrorString(result); + return message.str(); +} + +std::chrono::milliseconds getOperationTimeout() +{ + static auto const timeout = [] + { + constexpr int64_t defaultTimeoutMs = 5000; + auto const* value = std::getenv("TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS"); + if (value != nullptr) + { + try + { + auto const parsed = std::stoll(value); + if (parsed > 0) + { + return std::chrono::milliseconds{parsed}; + } + } + catch (...) + { + } + TLLM_LOG_WARNING( + "Ignoring invalid TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS=%s; expected positive milliseconds", value); + } + return std::chrono::milliseconds{defaultTimeoutMs}; + }(); + return timeout; +} + #endif // ENABLE_MULTI_DEVICE } // namespace +NcclCommunicator::NcclCommunicator(ncclComm_t comm) + : mFaultToleranceEnabled{mpi::isFaultToleranceModeEnabled()} + , mComm{comm} +{ +#if ENABLE_MULTI_DEVICE + TLLM_CHECK_WITH_INFO(mComm != nullptr, "NCCL error: cannot wrap a null communicator"); + + auto const hostApiLock = acquireNcclHostApiLock(); + try + { + if (mFaultToleranceEnabled) + { + waitUntilReady(mComm, "wrapping communicator", std::chrono::steady_clock::now() + kInitialReadyTimeout); + } + TLLM_NCCL_CHECK(ncclCommCount(mComm, &mInitialWorldSize)); + TLLM_NCCL_CHECK(ncclCommUserRank(mComm, &mWorldRank)); + mActiveRanks.resize(mInitialWorldSize); + std::iota(mActiveRanks.begin(), mActiveRanks.end(), 0); + } + catch (...) + { + ncclCommAbort(mComm); + mComm = nullptr; + throw; + } + startWatcher(); +#else + (void) comm; +#endif // ENABLE_MULTI_DEVICE +} + +NcclCommunicator::NcclCommunicator(int worldSize, int rank, mpi::MpiComm const& mpiComm) + : mInitialWorldSize{worldSize} + , mWorldRank{rank} + , mFaultToleranceEnabled{mpi::isFaultToleranceModeEnabled()} +{ + TLLM_CHECK_WITH_INFO(worldSize > 0, "NCCL error: communicator world size must be positive, got %d", worldSize); + TLLM_CHECK_WITH_INFO( + rank >= 0 && rank < worldSize, "NCCL error: communicator rank must be in [0, %d), got %d", worldSize, rank); + + mActiveRanks.resize(worldSize); + std::iota(mActiveRanks.begin(), mActiveRanks.end(), 0); +#if ENABLE_MULTI_DEVICE + TLLM_CHECK_WITH_INFO(mpiComm.getRank() == rank, "NCCL error: world rank %d does not match MPI bootstrap rank %d", + rank, mpiComm.getRank()); + TLLM_CHECK_WITH_INFO(mpiComm.getSize() >= worldSize, "NCCL error: world size %d exceeds MPI bootstrap size %d", + worldSize, mpiComm.getSize()); + if (mFaultToleranceEnabled) + { + mControlComm = createNcclUniqueIdRendezvousComm( + mActiveRanks, rank, mpiComm, static_cast(mpi::MpiTag::kNcclPpControl)); + mComm = createComm(mActiveRanks, rank, *mControlComm, /*rendezvousId=*/0, kInitialReadyTimeout); + } + else + { + // Preserve the legacy blocking MPI-broadcast bootstrap when FT is + // disabled. It creates no control communicator or watcher thread. + mComm = createLegacyComm(worldSize, rank, mpiComm); + } +#else + mComm = nullptr; +#endif // ENABLE_MULTI_DEVICE + startWatcher(); +} + void NcclCommunicator::send( void const* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const { #if ENABLE_MULTI_DEVICE - TLLM_NCCL_CHECK(ncclSend(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); + if (!mFaultToleranceEnabled) + { + TLLM_NCCL_CHECK(ncclSend(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); + return; + } + std::lock_guard lock(mCommMutex); + auto const hostApiLock = acquireNcclHostApiLock(); + checkUsableLocked(); + int const commPeer = getCommPeerRankLocked(peer); + auto const watchdogToken = beginOperationLocked(stream.get(), "ncclSend"); + ncclResult_t const result = ncclSend(sendbuff, count, toNcclType(dataType), commPeer, mComm, stream.get()); + if (result == ncclInProgress) + { + try + { + waitUntilReady(mComm, "ncclSend enqueue", std::chrono::steady_clock::now() + kRecoveryReadyTimeout); + } + catch (std::exception const& error) + { + static_cast(abortLocked( + std::string{"NCCL error: communicator was aborted after ncclSend failed: "} + error.what())); + TLLM_THROW("%s", mAsyncError.c_str()); + } + catch (...) + { + static_cast(abortLocked("NCCL error: communicator was aborted after ncclSend failed")); + TLLM_THROW("%s", mAsyncError.c_str()); + } + } + else if (result != ncclSuccess) + { + static_cast( + abortLocked("NCCL error: communicator was aborted after " + ncclErrorMessage("ncclSend", result))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + finishOperationLocked(watchdogToken, stream.get()); #else - TLLM_THROW("Multi device support is disabled."); + TLLM_THROW("NCCL error: multi device support is disabled."); #endif // ENABLE_MULTI_DEVICE } @@ -63,23 +219,126 @@ void NcclCommunicator::receive( void* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const { #if ENABLE_MULTI_DEVICE - TLLM_NCCL_CHECK(ncclRecv(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); + if (!mFaultToleranceEnabled) + { + TLLM_NCCL_CHECK(ncclRecv(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); + return; + } + std::lock_guard lock(mCommMutex); + auto const hostApiLock = acquireNcclHostApiLock(); + checkUsableLocked(); + int const commPeer = getCommPeerRankLocked(peer); + auto const watchdogToken = beginOperationLocked(stream.get(), "ncclRecv"); + ncclResult_t const result = ncclRecv(sendbuff, count, toNcclType(dataType), commPeer, mComm, stream.get()); + if (result == ncclInProgress) + { + try + { + waitUntilReady(mComm, "ncclRecv enqueue", std::chrono::steady_clock::now() + kRecoveryReadyTimeout); + } + catch (std::exception const& error) + { + static_cast(abortLocked( + std::string{"NCCL error: communicator was aborted after ncclRecv failed: "} + error.what())); + TLLM_THROW("%s", mAsyncError.c_str()); + } + catch (...) + { + static_cast(abortLocked("NCCL error: communicator was aborted after ncclRecv failed")); + TLLM_THROW("%s", mAsyncError.c_str()); + } + } + else if (result != ncclSuccess) + { + static_cast( + abortLocked("NCCL error: communicator was aborted after " + ncclErrorMessage("ncclRecv", result))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + finishOperationLocked(watchdogToken, stream.get()); #else - TLLM_THROW("Multi device support is disabled."); + TLLM_THROW("NCCL error: multi device support is disabled."); #endif // ENABLE_MULTI_DEVICE } -ncclComm_t NcclCommunicator::createComm(int worldSize, int rank, mpi::MpiComm const& mpiComm) +ncclComm_t NcclCommunicator::createComm(std::vector const& activeRanks, int worldRank, + NcclUniqueIdRendezvousComm const& controlComm, std::uint64_t rendezvousId, std::chrono::milliseconds readyTimeout) { #if ENABLE_MULTI_DEVICE + TLLM_CHECK_WITH_INFO(!activeRanks.empty(), "NCCL error: communicator active-rank set must not be empty"); + auto const rankIt = std::find(activeRanks.begin(), activeRanks.end(), worldRank); + TLLM_CHECK_WITH_INFO( + rankIt != activeRanks.end(), "NCCL error: world rank %d is not present in the active-rank set", worldRank); + // Do not use an MPI collective here. During recovery, failed ranks remain + // members of the pre-failure control communicator. Only the survivors + // exchange the fresh NCCL ID point-to-point. + auto const deadline = std::chrono::steady_clock::now() + readyTimeout; + TLLM_CHECK_WITH_INFO(controlComm.worldRank() == worldRank, + "NCCL error: world rank %d does not match rendezvous control world rank %d", worldRank, + controlComm.worldRank()); + ncclUniqueId const id = exchangeNcclUniqueId(activeRanks, controlComm, + {mpi::MpiTag::kNcclPpReady, mpi::MpiTag::kNcclPpUniqueId, mpi::MpiTag::kNcclPpAck}, rendezvousId, deadline); + + ncclComm_t comm{nullptr}; +// Need static connection initialization for accurate KV cache size estimation +#if defined(_WIN32) + if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) + _putenv_s("NCCL_RUNTIME_CONNECT", "0"); +#else + setenv("NCCL_RUNTIME_CONNECT", "0", 0); +#endif // _WIN32 + + ncclConfig_t config = NCCL_CONFIG_INITIALIZER; + config.blocking = 0; + int const commRank = static_cast(std::distance(activeRanks.begin(), rankIt)); + int const commSize = static_cast(activeRanks.size()); + auto const hostApiLock = acquireNcclHostApiLock(); + ncclResult_t const result = ncclCommInitRankConfig(&comm, commSize, id, commRank, &config); + if (result != ncclSuccess && result != ncclInProgress) + { + if (comm != nullptr) + { + ncclCommAbort(comm); + } + TLLM_THROW("%s", ncclErrorMessage("ncclCommInitRankConfig", result).c_str()); + } + + try + { + waitUntilReady(comm, "ncclCommInitRankConfig", deadline); + } + catch (...) + { + if (comm != nullptr) + { + ncclCommAbort(comm); + } + throw; + } + return comm; +#else + (void) activeRanks; + (void) worldRank; + (void) controlComm; + (void) rendezvousId; + (void) readyTimeout; + // Python runtime requires instantiation of a communicator even though it may never be used to enable + // pipeline parallel code-path. To enable this, have an empty communicator with uninitialized state. + return nullptr; +#endif // ENABLE_MULTI_DEVICE +} + +ncclComm_t NcclCommunicator::createLegacyComm(int worldSize, int rank, mpi::MpiComm const& mpiComm) +{ +#if ENABLE_MULTI_DEVICE ncclUniqueId id; if (rank == 0) { - ncclGetUniqueId(&id); + auto const hostApiLock = acquireNcclHostApiLock(); + TLLM_NCCL_CHECK(ncclGetUniqueId(&id)); } mpiComm.bcastValue(id, 0); - ncclComm_t comm; + // Need static connection initialization for accurate KV cache size estimation #if defined(_WIN32) if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) @@ -87,21 +346,577 @@ ncclComm_t NcclCommunicator::createComm(int worldSize, int rank, mpi::MpiComm co #else setenv("NCCL_RUNTIME_CONNECT", "0", 0); #endif // _WIN32 + + ncclComm_t comm{nullptr}; + auto const hostApiLock = acquireNcclHostApiLock(); TLLM_NCCL_CHECK(ncclCommInitRank(&comm, worldSize, id, rank)); return comm; #else - // Python runtime requires instantiation of a communicator even though it may never be used to enable - // pipeline parallel code-path. To enable this, have an empty communicator with uninitialized state. + (void) worldSize; + (void) rank; + (void) mpiComm; return nullptr; #endif // ENABLE_MULTI_DEVICE } +void NcclCommunicator::waitUntilReady( + ncclComm_t comm, char const* operation, std::chrono::steady_clock::time_point deadline) +{ +#if ENABLE_MULTI_DEVICE + auto const hostApiLock = acquireNcclHostApiLock(); + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL error: communicator was aborted while waiting for %s", operation); + while (true) + { + ncclResult_t asyncError = ncclInProgress; + ncclResult_t const queryResult = ncclCommGetAsyncError(comm, &asyncError); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + TLLM_THROW("%s", ncclErrorMessage("ncclCommGetAsyncError", queryResult).c_str()); + } + if (queryResult == ncclSuccess && asyncError == ncclSuccess) + { + return; + } + if (queryResult == ncclSuccess && asyncError != ncclInProgress) + { + TLLM_THROW("%s", ncclErrorMessage(operation, asyncError).c_str()); + } + if (std::chrono::steady_clock::now() >= deadline) + { + TLLM_THROW("NCCL error: timed out waiting for %s", operation); + } + std::this_thread::sleep_for(kReadyPollInterval); + } +#else + (void) comm; + (void) operation; + (void) deadline; +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::checkUsableLocked() const +{ + TLLM_CHECK_WITH_INFO(mAsyncError.empty(), "%s", mAsyncError.c_str()); + TLLM_CHECK_WITH_INFO(mComm != nullptr, "NCCL error: communicator was aborted"); +} + +bool NcclCommunicator::abortLocked(std::string reason) const noexcept +{ +#if ENABLE_MULTI_DEVICE + auto const hostApiLock = acquireNcclHostApiLock(); + if (mComm == nullptr) + { + if (mAsyncError.empty()) + { + mAsyncError = std::move(reason); + } + return true; + } + + ncclComm_t const comm = mComm; + ncclResult_t const result = ncclCommAbort(comm); + if (result == ncclSuccess) + { + mComm = nullptr; + mAsyncError = std::move(reason); + quarantinePendingOperationsLocked(); + return true; + } + + bool const firstAbortFailure = reason.find("ncclCommAbort failed") == std::string::npos; + mAsyncError = std::move(reason); + if (firstAbortFailure) + { + try + { + mAsyncError += "; ncclCommAbort failed: "; + mAsyncError += ncclGetErrorString(result); + } + catch (...) + { + } + } + if (firstAbortFailure) + { + TLLM_LOG_ERROR("Failed to abort NCCL communicator: %s", ncclGetErrorString(result)); + } + return false; +#else + (void) reason; + return true; +#endif // ENABLE_MULTI_DEVICE +} + +int NcclCommunicator::getCommPeerRankLocked(int worldRank) const +{ + auto const peer = std::find(mActiveRanks.begin(), mActiveRanks.end(), worldRank); + TLLM_CHECK_WITH_INFO(peer != mActiveRanks.end(), + "NCCL error: peer world rank %d is not active in the current communicator", worldRank); + return static_cast(std::distance(mActiveRanks.begin(), peer)); +} + +uint64_t NcclCommunicator::beginOperationLocked(cudaStream_t stream, char const* operation) const +{ +#if ENABLE_MULTI_DEVICE + checkUsableLocked(); + if (!mFaultToleranceEnabled) + { + return 0; + } + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + auto const captureResult = cudaStreamIsCapturing(stream, &captureStatus); + if (captureResult != cudaSuccess) + { + static_cast(abortLocked(std::string{"NCCL error: "} + operation + + " failed to query CUDA stream capture state: " + cudaGetErrorString(captureResult))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + if (captureStatus != cudaStreamCaptureStatusNone) + { + TLLM_CHECK_WITH_INFO(!mFaultToleranceEnabled, + "NCCL error: CUDA graph capture of PP NCCL operations is disabled while " + "TLLM_FAULT_TOLERANCE_MODE=1; captured graph nodes retain the old communicator across recovery"); + TLLM_LOG_DEBUG("Skipping eager NCCL completion deadline for %s during CUDA graph capture", operation); + return 0; + } + if (mPendingOperations.size() >= kMaxPendingOperations) + { + static_cast(abortLocked( + std::string{"NCCL error: "} + operation + " exceeded the bounded NCCL watchdog operation queue")); + TLLM_THROW("%s", mAsyncError.c_str()); + } + + auto const token = mNextOperationToken++; + auto const start = acquireEventLocked(); + mPendingOperations.push_back({token, start, nullptr, {}, false, operation}); + auto const result = cudaEventRecord(start, stream); + if (result != cudaSuccess) + { + static_cast(abortLocked(std::string{"NCCL error: "} + operation + + " failed to record its start marker: " + cudaGetErrorString(result))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + mWatcherWakeup.notify_all(); + return token; +#else + (void) stream; + (void) operation; + return 0; +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::finishOperationLocked(uint64_t token, cudaStream_t stream) const +{ +#if ENABLE_MULTI_DEVICE + checkUsableLocked(); + if (token == 0) + { + return; + } + auto const operation = std::find_if(mPendingOperations.begin(), mPendingOperations.end(), + [token](PendingOperation const& pending) { return pending.token == token; }); + TLLM_CHECK_WITH_INFO(operation != mPendingOperations.end(), "NCCL error: unknown watchdog operation token %llu", + static_cast(token)); + TLLM_CHECK_WITH_INFO(operation->completion == nullptr, + "NCCL error: watchdog operation token %llu was already completed", static_cast(token)); + + operation->completion = acquireEventLocked(); + auto const result = cudaEventRecord(operation->completion, stream); + if (result != cudaSuccess) + { + static_cast(abortLocked(std::string{"NCCL error: "} + operation->name + + " failed to record its completion marker: " + cudaGetErrorString(result))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + mWatcherWakeup.notify_all(); +#else + (void) token; + (void) stream; +#endif // ENABLE_MULTI_DEVICE +} + +cudaEvent_t NcclCommunicator::acquireEventLocked() const +{ +#if ENABLE_MULTI_DEVICE + if (!mEventPool.empty()) + { + auto const event = mEventPool.back(); + mEventPool.pop_back(); + return event; + } + cudaEvent_t event = nullptr; + auto const result = cudaEventCreateWithFlags(&event, cudaEventDisableTiming); + if (result != cudaSuccess) + { + static_cast( + abortLocked(std::string{"NCCL error: failed to allocate a watchdog event: "} + cudaGetErrorString(result))); + TLLM_THROW("%s", mAsyncError.c_str()); + } + return event; +#else + return nullptr; +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::recycleEventLocked(cudaEvent_t event) const noexcept +{ +#if ENABLE_MULTI_DEVICE + if (event == nullptr) + { + return; + } + if (mEventPool.size() < kMaxPooledEvents) + { + mEventPool.push_back(event); + } + else + { + static_cast(cudaEventDestroy(event)); + } +#else + (void) event; +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::quarantinePendingOperationsLocked() const noexcept +{ + if (!mPendingOperations.empty()) + { + // The failed communicator may still have device work referencing + // these events. Drop host ownership without destroying their handles; + // process restart reclaims the deliberately quarantined resources. + TLLM_LOG_WARNING( + "Quarantining %zu PP NCCL completion markers after communicator abort", mPendingOperations.size()); + } + mPendingOperations.clear(); +} + +void NcclCommunicator::destroyCompletedWatchdogEventsLocked() const noexcept +{ +#if ENABLE_MULTI_DEVICE + size_t quarantined = 0; + for (auto const& operation : mPendingOperations) + { + for (auto const event : {operation.start, operation.completion}) + { + if (event == nullptr) + { + continue; + } + if (cudaEventQuery(event) == cudaSuccess) + { + static_cast(cudaEventDestroy(event)); + } + else + { + ++quarantined; + } + } + } + mPendingOperations.clear(); + if (quarantined != 0) + { + TLLM_LOG_WARNING("Quarantining %zu incomplete PP NCCL watchdog events during healthy teardown", quarantined); + } + destroyPooledEventsLocked(); +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::destroyPooledEventsLocked() const noexcept +{ +#if ENABLE_MULTI_DEVICE + for (auto const event : mEventPool) + { + static_cast(cudaEventDestroy(event)); + } + mEventPool.clear(); +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::abort() +{ +#if ENABLE_MULTI_DEVICE + std::lock_guard lock(mCommMutex); + bool const abortSucceeded = abortLocked("NCCL error: communicator was aborted by the fault-tolerance control path"); + TLLM_CHECK_WITH_INFO(abortSucceeded, "%s", mAsyncError.c_str()); +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::abortAndReinit(std::vector const& activeRanks, std::uint64_t rendezvousId) +{ +#if ENABLE_MULTI_DEVICE + std::lock_guard lock(mCommMutex); + TLLM_CHECK_WITH_INFO(mFaultToleranceEnabled, + "NCCL error: communicator reinitialization requires TLLM_FAULT_TOLERANCE_MODE=1; otherwise captured CUDA " + "graphs may retain the old communicator"); + TLLM_CHECK_WITH_INFO( + rendezvousId != 0, "NCCL error: recovery rendezvous ID must be positive; zero is reserved for bootstrap"); + TLLM_CHECK_WITH_INFO(mControlComm != nullptr, + "NCCL error: communicator created from a raw handle has no MPI bootstrap for reinitialization"); + TLLM_CHECK_WITH_INFO(!activeRanks.empty(), "NCCL error: active-rank set must not be empty"); + std::set const requested(activeRanks.begin(), activeRanks.end()); + TLLM_CHECK_WITH_INFO(requested.size() == activeRanks.size(), "NCCL error: active-rank set contains duplicates"); + TLLM_CHECK_WITH_INFO(requested.find(mWorldRank) != requested.end(), + "NCCL error: survivor rank %d must be present in active_ranks", mWorldRank); + + for (int const rank : requested) + { + TLLM_CHECK_WITH_INFO(rank >= 0 && rank < mInitialWorldSize, + "NCCL error: active world rank %d is outside [0, %d)", rank, mInitialWorldSize); + TLLM_CHECK_WITH_INFO(std::find(mActiveRanks.begin(), mActiveRanks.end(), rank) != mActiveRanks.end(), + "NCCL error: cannot re-add world rank %d after it was removed", rank); + } + + std::vector nextActiveRanks; + nextActiveRanks.reserve(requested.size()); + for (int const worldRank : mActiveRanks) + { + if (requested.find(worldRank) != requested.end()) + { + nextActiveRanks.push_back(worldRank); + } + } + + TLLM_CHECK_WITH_INFO(mControlComm->worldRank() == mWorldRank, + "NCCL error: world rank %d does not match rendezvous control world rank %d", mWorldRank, + mControlComm->worldRank()); + // This is deliberately a fresh-ID rebuild rather than ncclCommShrink. The + // watcher may already have destroyed the poisoned communicator, and the + // failed rank must not be required by any recovery collective. + bool const abortSucceeded + = abortLocked("NCCL error: communicator was aborted before survivor-only reinitialization"); + TLLM_CHECK_WITH_INFO(abortSucceeded, + "NCCL error: communicator reinitialization refused because ncclCommAbort did not succeed: %s", + mAsyncError.c_str()); + + ncclComm_t nextComm{nullptr}; + try + { + nextComm = createComm(nextActiveRanks, mWorldRank, *mControlComm, rendezvousId, kRecoveryReadyTimeout); + } + catch (std::exception const& error) + { + mAsyncError = std::string{"NCCL error: communicator reinitialization failed after abort: "} + error.what(); + TLLM_THROW("%s", mAsyncError.c_str()); + } + catch (...) + { + mAsyncError = "NCCL error: communicator reinitialization failed after abort"; + TLLM_THROW("%s", mAsyncError.c_str()); + } + + mComm = nextComm; + mActiveRanks = std::move(nextActiveRanks); + mAsyncError.clear(); +#else + (void) activeRanks; + (void) rendezvousId; + TLLM_THROW("NCCL error: multi device support is disabled."); +#endif // ENABLE_MULTI_DEVICE +} + +std::string NcclCommunicator::getAsyncError() const +{ + std::lock_guard lock(mCommMutex); + return mAsyncError; +} + +std::vector NcclCommunicator::getActiveRanks() const +{ + std::lock_guard lock(mCommMutex); + return mActiveRanks; +} + +void NcclCommunicator::startWatcher() +{ +#if ENABLE_MULTI_DEVICE + if (mFaultToleranceEnabled && mComm != nullptr) + { + try + { + mWatcherThread = std::thread(&NcclCommunicator::watchAsyncErrors, this); + } + catch (...) + { + auto const hostApiLock = acquireNcclHostApiLock(); + ncclCommAbort(mComm); + mComm = nullptr; + throw; + } + } +#endif // ENABLE_MULTI_DEVICE +} + +void NcclCommunicator::stopWatcher() noexcept +{ + mStopWatcher.store(true, std::memory_order_release); + mWatcherWakeup.notify_all(); + if (mWatcherThread.joinable()) + { + mWatcherThread.join(); + } +} + +void NcclCommunicator::watchAsyncErrors() noexcept +{ +#if ENABLE_MULTI_DEVICE + try + { + while (!mStopWatcher.load(std::memory_order_acquire)) + { + { + // NCCL communicators are not generally thread-safe. Serialize + // the monitor with send/receive, abort, and reinit host calls. + std::lock_guard lock(mCommMutex); + if (mComm != nullptr && !mAsyncError.empty()) + { + // A previous abort attempt failed. Keep the handle so the + // watcher can retry; a rebuild is forbidden until abort + // succeeds. + static_cast(abortLocked(mAsyncError)); + } + else if (mComm != nullptr) + { + std::string error; + auto const now = std::chrono::steady_clock::now(); + for (auto operation = mPendingOperations.begin(); operation != mPendingOperations.end();) + { + if (operation->start != nullptr) + { + auto const startResult = cudaEventQuery(operation->start); + if (startResult == cudaSuccess) + { + recycleEventLocked(operation->start); + operation->start = nullptr; + operation->deadline = now + getOperationTimeout(); + operation->armed = true; + } + else if (startResult != cudaErrorNotReady) + { + error = operation->name + " start marker failed: " + cudaGetErrorString(startResult); + break; + } + else + { + ++operation; + continue; + } + } + if (operation->completion == nullptr) + { + if (operation->armed && now >= operation->deadline) + { + error = operation->name + " timed out before its completion marker was recorded"; + break; + } + ++operation; + continue; + } + + auto const eventResult = cudaEventQuery(operation->completion); + if (eventResult == cudaSuccess) + { + recycleEventLocked(operation->completion); + operation = mPendingOperations.erase(operation); + continue; + } + if (eventResult != cudaErrorNotReady) + { + error = operation->name + " completion marker failed: " + cudaGetErrorString(eventResult); + break; + } + if (operation->armed && now >= operation->deadline) + { + error = operation->name + " timed out before its CUDA completion marker"; + break; + } + ++operation; + } + + if (error.empty()) + { + auto const hostApiLock = acquireNcclHostApiLock(); + ncclResult_t asyncError = ncclSuccess; + ncclResult_t const queryResult = ncclCommGetAsyncError(mComm, &asyncError); + if (queryResult != ncclSuccess && queryResult != ncclInProgress) + { + error = ncclErrorMessage("ncclCommGetAsyncError watcher", queryResult); + } + else if (queryResult == ncclSuccess && asyncError != ncclSuccess + && asyncError != ncclInProgress) + { + error = ncclErrorMessage("asynchronous communicator operation", asyncError); + } + } + + if (!error.empty()) + { + static_cast(abortLocked( + "NCCL error: communicator was aborted by the async-error watcher after " + error)); + TLLM_LOG_ERROR("%s", mAsyncError.c_str()); + } + } + } + + std::unique_lock waitLock(mWatcherWaitMutex); + mWatcherWakeup.wait_for( + waitLock, kWatcherPollInterval, [this]() { return mStopWatcher.load(std::memory_order_acquire); }); + } + } + catch (...) + { + // Never allow a diagnostic side thread to terminate the process. + } +#endif // ENABLE_MULTI_DEVICE +} + NcclCommunicator::~NcclCommunicator() { #if ENABLE_MULTI_DEVICE - if (mComm && ncclCommDestroy(mComm) != ncclSuccess) + if (!mFaultToleranceEnabled) + { + if (mComm != nullptr && ncclCommDestroy(mComm) != ncclSuccess) + { + TLLM_LOG_WARNING("Failed to destroy NCCL communicator."); + } + return; + } +#endif // ENABLE_MULTI_DEVICE + stopWatcher(); +#if ENABLE_MULTI_DEVICE + std::lock_guard lock(mCommMutex); + auto const hostApiLock = acquireNcclHostApiLock(); + if (mComm != nullptr) + { + ncclResult_t asyncError = ncclSuccess; + ncclResult_t const queryResult = ncclCommGetAsyncError(mComm, &asyncError); + bool const healthy = mAsyncError.empty() && queryResult == ncclSuccess && asyncError == ncclSuccess; + // ncclCommDestroy is an intra-node collective. Python model-engine + // reference counts are process-local, so their final release is not a + // cross-rank teardown barrier; after a peer failure, the one healthy + // sample above is also inherently racy. In FT mode use the local abort + // path and never stop the watcher only to enter an unbounded graceful + // destroy while holding the process-wide NCCL gate. + bool const useAbort = !healthy || mFaultToleranceEnabled; + ncclResult_t const result = useAbort ? ncclCommAbort(mComm) : ncclCommDestroy(mComm); + if (result != ncclSuccess) + { + TLLM_LOG_WARNING("Failed to release NCCL communicator: %s", ncclGetErrorString(result)); + } + mComm = nullptr; + if (!useAbort && result == ncclSuccess) + { + destroyCompletedWatchdogEventsLocked(); + } + else + { + quarantinePendingOperationsLocked(); + destroyPooledEventsLocked(); + } + } + else { - TLLM_LOG_WARNING("Failed to destroy NCCL communicator."); + destroyPooledEventsLocked(); } #endif // ENABLE_MULTI_DEVICE } diff --git a/cpp/tensorrt_llm/runtime/ncclCommunicator.h b/cpp/tensorrt_llm/runtime/ncclCommunicator.h index 76cce4beab8a..388030e234b5 100644 --- a/cpp/tensorrt_llm/runtime/ncclCommunicator.h +++ b/cpp/tensorrt_llm/runtime/ncclCommunicator.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -21,20 +21,43 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/worldConfig.h" +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + struct ncclComm; typedef struct ncclComm* ncclComm_t; namespace tensorrt_llm::runtime { +class NcclUniqueIdRendezvousComm; + class NcclCommunicator { public: - explicit NcclCommunicator(ncclComm_t comm) - : mComm{comm} {}; - - explicit NcclCommunicator(int worldSize, int rank, mpi::MpiComm const& mpiComm = COMM_SESSION) - : mComm{createComm(worldSize, rank, mpiComm)} {}; + //! Wrap an existing communicator. + //! + //! In fault-tolerance mode the caller is responsible for creating `comm` + //! in non-blocking mode. The default-off path preserves the legacy + //! blocking communicator contract. + explicit NcclCommunicator(ncclComm_t comm); + + //! Create a communicator bootstrapped by `mpiComm`. + //! + //! Construction establishes and retains a dedicated MPI control channel while + //! the initial group is healthy. Recovery reuses that channel without an + //! MPI collective; `mpiComm` therefore need not outlive this object. + explicit NcclCommunicator(int worldSize, int rank, mpi::MpiComm const& mpiComm = COMM_SESSION); explicit NcclCommunicator(WorldConfig const& worldConfig, mpi::MpiComm const& mpiComm = COMM_SESSION) : NcclCommunicator{worldConfig.getSize(), worldConfig.getRank(), mpiComm} {}; @@ -55,15 +78,86 @@ class NcclCommunicator receive(buf.data(), buf.getSize(), buf.getDataType(), peer, stream); } + //! Abort the current communicator. Idempotent. + void abort(); + + //! Abort and replace the current communicator with a fresh communicator + //! containing only `activeRanks`. + //! + //! Ranks are the original world-rank IDs supplied to the constructor. The + //! local rank must be present, and previously removed ranks cannot be added + //! again. Every survivor must call this method with the same rank set. The + //! positive `rendezvousId` must identify the same logical recovery attempt + //! on every survivor; zero is reserved for the initial bootstrap. The NCCL + //! unique ID is exchanged point-to-point over the dedicated control + //! channel, so failed ranks do not participate in recovery. + //! + //! Communicators created by the raw-handle constructor cannot be rebuilt + //! because that constructor has no MPI bootstrap communicator. + void abortAndReinit(std::vector const& activeRanks, std::uint64_t rendezvousId); + + //! Return the latched NCCL abort/error reason, or an empty string. + [[nodiscard]] std::string getAsyncError() const; + + //! Return the original world-rank IDs in the current communicator. + [[nodiscard]] std::vector getActiveRanks() const; + private: void send( void const* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const; void receive(void* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const; - static ncclComm_t createComm(int worldSize, int rank, mpi::MpiComm const& mpiComm); - - ncclComm_t mComm; + static ncclComm_t createComm(std::vector const& activeRanks, int worldRank, + NcclUniqueIdRendezvousComm const& controlComm, std::uint64_t rendezvousId, + std::chrono::milliseconds readyTimeout); + static ncclComm_t createLegacyComm(int worldSize, int rank, mpi::MpiComm const& mpiComm); + static void waitUntilReady(ncclComm_t comm, char const* operation, std::chrono::steady_clock::time_point deadline); + + void startWatcher(); + void stopWatcher() noexcept; + void watchAsyncErrors() noexcept; + [[nodiscard]] bool abortLocked(std::string reason) const noexcept; + void checkUsableLocked() const; + [[nodiscard]] int getCommPeerRankLocked(int worldRank) const; + [[nodiscard]] uint64_t beginOperationLocked(cudaStream_t stream, char const* operation) const; + void finishOperationLocked(uint64_t token, cudaStream_t stream) const; + [[nodiscard]] cudaEvent_t acquireEventLocked() const; + void recycleEventLocked(cudaEvent_t event) const noexcept; + void quarantinePendingOperationsLocked() const noexcept; + void destroyCompletedWatchdogEventsLocked() const noexcept; + void destroyPooledEventsLocked() const noexcept; + + int mInitialWorldSize{0}; + int mWorldRank{0}; + bool mFaultToleranceEnabled{false}; + std::vector mActiveRanks; + std::shared_ptr mControlComm; + + mutable std::mutex mCommMutex; + mutable ncclComm_t mComm{nullptr}; + mutable std::string mAsyncError; + + struct PendingOperation + { + uint64_t token; + cudaEvent_t start; + cudaEvent_t completion; + std::chrono::steady_clock::time_point deadline; + bool armed; + std::string name; + }; + + mutable std::vector mPendingOperations; + mutable std::vector mEventPool; + mutable uint64_t mNextOperationToken{1}; + static constexpr size_t kMaxPendingOperations = 4096; + static constexpr size_t kMaxPooledEvents = 256; + + std::atomic mStopWatcher{false}; + std::mutex mWatcherWaitMutex; + mutable std::condition_variable mWatcherWakeup; + std::thread mWatcherThread; }; } // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.cpp b/cpp/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.cpp new file mode 100644 index 000000000000..7d13162d5dc9 --- /dev/null +++ b/cpp/tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.cpp @@ -0,0 +1,787 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h" + +#if ENABLE_MULTI_DEVICE + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::runtime +{ +namespace +{ + +using Clock = std::chrono::steady_clock; + +constexpr std::uint32_t kProtocolVersion = 2; +constexpr auto kPollInterval = std::chrono::milliseconds{1}; +constexpr auto kReadyRefreshInterval = std::chrono::milliseconds{100}; +constexpr std::size_t kMaxTokensPerPeer = 1024; +constexpr std::size_t kMaxPendingMessages = 4096; + +struct AttemptToken +{ + std::uint64_t incarnation; + std::uint64_t sequence; +}; + +bool operator==(AttemptToken const& lhs, AttemptToken const& rhs) noexcept +{ + return lhs.incarnation == rhs.incarnation && lhs.sequence == rhs.sequence; +} + +struct ReadyMessage +{ + std::uint32_t version; + std::uint32_t reserved; + std::uint64_t communicatorKey; + std::uint64_t rendezvousId; + AttemptToken clientToken; +}; + +struct IdMessage +{ + std::uint32_t version; + std::uint32_t reserved; + std::uint64_t communicatorKey; + std::uint64_t rendezvousId; + AttemptToken clientToken; + AttemptToken serverToken; + std::uint64_t idDigest; + ncclUniqueId id; +}; + +struct AckMessage +{ + std::uint32_t version; + std::uint32_t reserved; + std::uint64_t communicatorKey; + std::uint64_t rendezvousId; + AttemptToken clientToken; + AttemptToken serverToken; + std::uint64_t idDigest; +}; + +template +struct ReceivedMessage +{ + int source; + Message message; +}; + +struct RendezvousContext +{ + std::mutex exchangeMutex; + std::deque> pendingReady; + std::deque> pendingIds; + std::deque> pendingAcks; +}; + +struct ControlCommCacheEntry +{ + ControlCommCacheEntry(MPI_Group parentGroup_, std::vector initialRanks_, int worldRank_, int creationTagSeed_) + : parentGroup(parentGroup_) + , initialRanks(std::move(initialRanks_)) + , worldRank(worldRank_) + , creationTagSeed(creationTagSeed_) + { + } + + std::mutex mutex; + MPI_Group parentGroup{MPI_GROUP_NULL}; + std::vector initialRanks; + int worldRank; + int creationTagSeed; + std::shared_ptr comm; + std::exception_ptr failure; +}; + +std::uint64_t groupIdentity(std::vector const& ranks) noexcept +{ + std::uint64_t hash = 14695981039346656037ULL; + for (int const rank : ranks) + { + auto const* bytes = reinterpret_cast(&rank); + for (std::size_t index = 0; index < sizeof(rank); ++index) + { + hash ^= bytes[index]; + hash *= 1099511628211ULL; + } + } + return hash; +} + +using ContextKey = std::tuple; + +ContextKey getContextKey(NcclUniqueIdRendezvousTags const& tags, NcclUniqueIdRendezvousComm const& controlComm) noexcept +{ + return {MPI_Comm_c2f(static_cast(controlComm.mpiComm())), groupIdentity(controlComm.worldRanks()), + static_cast(tags.ready), static_cast(tags.id), static_cast(tags.ack)}; +} + +RendezvousContext& getRendezvousContext( + NcclUniqueIdRendezvousTags const& tags, NcclUniqueIdRendezvousComm const& controlComm) +{ + // Keep contexts immortal for the same reason as the NCCL host gate: a + // communicator may be torn down from another process-lifetime singleton. + static auto* registryMutex = new std::mutex; + static auto* contexts = new std::map>; + std::lock_guard lock(*registryMutex); + auto& context = (*contexts)[getContextKey(tags, controlComm)]; + if (context == nullptr) + { + context = std::make_unique(); + } + return *context; +} + +std::shared_ptr getControlCommCacheEntry( + std::vector const& initialRanks, int worldRank, mpi::MpiComm const& parentComm, int creationTagSeed) +{ + // A control communicator cannot be freed safely after a member dies. Keep + // one channel per topology for the process lifetime and reuse it across PP + // engine teardown/recreation instead of leaking a new MPI context each + // time. Retain the parent's MPI group as a stable process-membership + // identity: MPI communicator handles may be recycled after a custom parent + // is freed. Per-entry locking permits unrelated groups to create + // concurrently. + static auto* registryMutex = new std::mutex; + static auto* entries = new std::vector>; + + MPI_Group parentGroup = MPI_GROUP_NULL; + int result = MPI_Comm_group(static_cast(parentComm), &parentGroup); + TLLM_CHECK_WITH_INFO( + result == MPI_SUCCESS, "NCCL error: failed to retain MPI parent group identity (code %d)", result); + try + { + std::lock_guard lock(*registryMutex); + for (auto const& entry : *entries) + { + if (entry->initialRanks != initialRanks || entry->worldRank != worldRank + || entry->creationTagSeed != creationTagSeed) + { + continue; + } + int comparison = MPI_UNEQUAL; + result = MPI_Group_compare(entry->parentGroup, parentGroup, &comparison); + TLLM_CHECK_WITH_INFO( + result == MPI_SUCCESS, "NCCL error: failed to compare MPI parent group identity (code %d)", result); + if (comparison == MPI_IDENT) + { + MPI_Group_free(&parentGroup); + return entry; + } + } + + auto entry = std::make_shared(parentGroup, initialRanks, worldRank, creationTagSeed); + entries->push_back(entry); + parentGroup = MPI_GROUP_NULL; + return entry; + } + catch (...) + { + if (parentGroup != MPI_GROUP_NULL) + { + MPI_Group_free(&parentGroup); + } + throw; + } +} + +template +bool takePending(std::deque>& pending, int source, std::uint64_t key, + std::uint64_t rendezvousId, ReceivedMessage& result) +{ + auto found = pending.begin(); + for (; found != pending.end(); ++found) + { + if ((source == MPI_ANY_SOURCE || found->source == source) && found->message.communicatorKey == key + && found->message.rendezvousId == rendezvousId) + { + break; + } + } + if (found == pending.end()) + { + return false; + } + result = *found; + pending.erase(found); + return true; +} + +template +void discardObsoletePending( + std::deque>& pending, std::uint64_t key, std::uint64_t rendezvousId) +{ + pending.erase(std::remove_if(pending.begin(), pending.end(), + [key, rendezvousId](auto const& received) + { + return received.message.version == kProtocolVersion && received.message.communicatorKey == key + && received.message.rendezvousId < rendezvousId; + }), + pending.end()); +} + +template +bool isCurrentAttempt(Message const& message, std::uint64_t key, std::uint64_t rendezvousId) noexcept +{ + return message.communicatorKey == key && message.rendezvousId == rendezvousId; +} + +template +bool shouldStashForLater(Message const& message, std::uint64_t key, std::uint64_t rendezvousId) noexcept +{ + if (message.version != kProtocolVersion) + { + return false; + } + // Rendezvous IDs must increase for retries of one base communicator. + // Messages from an older retry can never be consumed again, while a peer + // may legitimately have entered a future retry or a different group first. + return message.communicatorKey != key || message.rendezvousId > rendezvousId; +} + +template +void stashPending(std::deque>& pending, ReceivedMessage received) +{ + TLLM_CHECK_WITH_INFO( + pending.size() < kMaxPendingMessages, "NCCL error: communicator rendezvous pending-message limit exceeded"); + pending.push_back(std::move(received)); +} + +std::uint64_t fnv1a(void const* data, std::size_t size, std::uint64_t hash = 14695981039346656037ULL) noexcept +{ + auto const* bytes = static_cast(data); + for (std::size_t index = 0; index < size; ++index) + { + hash ^= bytes[index]; + hash *= 1099511628211ULL; + } + return hash; +} + +std::uint64_t communicatorKey(std::vector const& activeRanks, NcclUniqueIdRendezvousTags const& tags) noexcept +{ + std::uint64_t hash = fnv1a(activeRanks.data(), activeRanks.size() * sizeof(activeRanks.front())); + std::array const tagValues{ + static_cast(tags.ready), static_cast(tags.id), static_cast(tags.ack)}; + return fnv1a(tagValues.data(), tagValues.size() * sizeof(tagValues.front()), hash); +} + +std::uint64_t processIncarnation() +{ + static std::uint64_t const incarnation = []() + { + std::random_device random; + std::array entropy{random(), random(), random(), random()}; + auto const wallClock = std::chrono::system_clock::now().time_since_epoch().count(); + std::uint64_t hash = fnv1a(entropy.data(), entropy.size() * sizeof(entropy.front())); + hash = fnv1a(&wallClock, sizeof(wallClock), hash); + return hash == 0 ? std::uint64_t{1} : hash; + }(); + return incarnation; +} + +AttemptToken makeAttemptToken() +{ + static std::atomic sequence{1}; + auto const value = sequence.fetch_add(1, std::memory_order_relaxed); + TLLM_CHECK_WITH_INFO(value != 0, "NCCL error: communicator rendezvous token sequence overflow"); + return AttemptToken{processIncarnation(), value}; +} + +int abandonRequest(MPI_Request& request) noexcept +{ + int firstError = MPI_SUCCESS; + if (request != MPI_REQUEST_NULL) + { + firstError = MPI_Cancel(&request); + // MPI_Wait can block forever after a peer failure, defeating the + // recovery deadline. Free the local request handle and let MPI retire + // it asynchronously; callers deliberately leak the tiny backing + // payload so the implementation can no longer access freed storage. + int const freeResult = MPI_Request_free(&request); + if (firstError == MPI_SUCCESS) + { + firstError = freeResult; + } + } + return firstError; +} + +void waitForSend(MPI_Request& request, Clock::time_point deadline, char const* operation) +{ + while (true) + { + int complete = 0; + int const result = MPI_Test(&request, &complete, MPI_STATUS_IGNORE); + if (result != MPI_SUCCESS) + { + int const cleanupResult = abandonRequest(request); + TLLM_THROW( + "NCCL error: MPI_Test failed while %s (code %d, cleanup code %d)", operation, result, cleanupResult); + } + if (complete != 0) + { + return; + } + if (Clock::now() >= deadline) + { + int const cleanupResult = abandonRequest(request); + TLLM_THROW("NCCL error: timed out while %s (MPI cleanup code %d)", operation, cleanupResult); + } + std::this_thread::sleep_for(kPollInterval); + } +} + +template +void sendMessage(Message const& message, int destination, mpi::MpiTag tag, mpi::MpiComm const& mpiComm, + Clock::time_point deadline, char const* operation) +{ + auto payload = std::make_unique(message); + MPI_Request request = MPI_REQUEST_NULL; + int const result = MPI_Isend(payload.get(), static_cast(sizeof(message)), MPI_BYTE, destination, + static_cast(tag), static_cast(mpiComm), &request); + if (result != MPI_SUCCESS) + { + int const cleanupResult = abandonRequest(request); + static_cast(payload.release()); + TLLM_THROW("NCCL error: MPI_Isend failed while %s to rank %d (code %d, cleanup code %d)", operation, + destination, result, cleanupResult); + } + try + { + waitForSend(request, deadline, operation); + } + catch (...) + { + static_cast(payload.release()); + throw; + } +} + +bool probeMessage(int source, mpi::MpiTag tag, mpi::MpiComm const& mpiComm, MPI_Status& status) +{ + int available = 0; + int const result = MPI_Iprobe(source, static_cast(tag), static_cast(mpiComm), &available, &status); + TLLM_CHECK_WITH_INFO(result == MPI_SUCCESS, "NCCL error: MPI_Iprobe failed with code %d", result); + return available != 0; +} + +template +Message receiveMessage(MPI_Status const& probedStatus, mpi::MpiTag tag, mpi::MpiComm const& mpiComm, + Clock::time_point deadline, char const* operation) +{ + auto message = std::make_unique(); + MPI_Request request = MPI_REQUEST_NULL; + int const result = MPI_Irecv(message.get(), static_cast(sizeof(Message)), MPI_BYTE, probedStatus.MPI_SOURCE, + static_cast(tag), static_cast(mpiComm), &request); + if (result != MPI_SUCCESS) + { + int const cleanupResult = abandonRequest(request); + static_cast(message.release()); + TLLM_THROW("NCCL error: MPI_Irecv failed while %s from rank %d (code %d, cleanup code %d)", operation, + probedStatus.MPI_SOURCE, result, cleanupResult); + } + + while (true) + { + int complete = 0; + MPI_Status status{}; + int const testResult = MPI_Test(&request, &complete, &status); + if (testResult != MPI_SUCCESS) + { + int const cleanupResult = abandonRequest(request); + static_cast(message.release()); + TLLM_THROW("NCCL error: MPI_Test failed while %s (code %d, cleanup code %d)", operation, testResult, + cleanupResult); + } + if (complete != 0) + { + int count = 0; + int const countResult = MPI_Get_count(&status, MPI_BYTE, &count); + TLLM_CHECK_WITH_INFO(countResult == MPI_SUCCESS && count == static_cast(sizeof(Message)), + "NCCL error: invalid MPI rendezvous payload while %s", operation); + return *message; + } + if (Clock::now() >= deadline) + { + int const cleanupResult = abandonRequest(request); + static_cast(message.release()); + TLLM_THROW("NCCL error: timed out while %s (MPI cleanup code %d)", operation, cleanupResult); + } + std::this_thread::sleep_for(kPollInterval); + } +} + +bool isActivePeer(std::vector const& activeRanks, int rank, int root) +{ + return rank != root && std::binary_search(activeRanks.begin(), activeRanks.end(), rank); +} + +} // namespace + +NcclUniqueIdRendezvousComm::NcclUniqueIdRendezvousComm(mpi::MpiComm comm, std::vector worldRanks, int worldRank) + : mComm(std::move(comm)) + , mWorldRanks(std::move(worldRanks)) + , mWorldRank(worldRank) +{ + TLLM_CHECK_WITH_INFO(!mWorldRanks.empty(), "NCCL error: rendezvous control group must not be empty"); + TLLM_CHECK_WITH_INFO(std::is_sorted(mWorldRanks.begin(), mWorldRanks.end()), + "NCCL error: rendezvous control ranks must be in canonical order"); + TLLM_CHECK_WITH_INFO(std::adjacent_find(mWorldRanks.begin(), mWorldRanks.end()) == mWorldRanks.end(), + "NCCL error: rendezvous control ranks contain duplicates"); + TLLM_CHECK_WITH_INFO(mComm.getSize() == static_cast(mWorldRanks.size()), + "NCCL error: rendezvous control communicator size %d does not match rank map size %zu", mComm.getSize(), + mWorldRanks.size()); + TLLM_CHECK_WITH_INFO(worldRank == this->worldRank(mComm.getRank()), + "NCCL error: world rank %d does not match rendezvous control rank %d", worldRank, mComm.getRank()); +} + +mpi::MpiComm const& NcclUniqueIdRendezvousComm::mpiComm() const noexcept +{ + return mComm; +} + +std::vector const& NcclUniqueIdRendezvousComm::worldRanks() const noexcept +{ + return mWorldRanks; +} + +int NcclUniqueIdRendezvousComm::worldRank() const noexcept +{ + return mWorldRank; +} + +int NcclUniqueIdRendezvousComm::commRank(int worldRank) const +{ + auto const found = std::lower_bound(mWorldRanks.begin(), mWorldRanks.end(), worldRank); + TLLM_CHECK_WITH_INFO(found != mWorldRanks.end() && *found == worldRank, + "NCCL error: world rank %d is not present in the rendezvous control channel", worldRank); + return static_cast(std::distance(mWorldRanks.begin(), found)); +} + +int NcclUniqueIdRendezvousComm::worldRank(int commRank) const +{ + TLLM_CHECK_WITH_INFO(commRank >= 0 && commRank < static_cast(mWorldRanks.size()), + "NCCL error: rendezvous control rank %d is outside [0, %zu)", commRank, mWorldRanks.size()); + return mWorldRanks[static_cast(commRank)]; +} + +std::shared_ptr createNcclUniqueIdRendezvousComm( + std::vector const& initialRanks, int worldRank, mpi::MpiComm const& parentComm, int creationTagSeed) +{ + TLLM_CHECK_WITH_INFO(!initialRanks.empty(), "NCCL error: rendezvous control group must not be empty"); + TLLM_CHECK_WITH_INFO(std::is_sorted(initialRanks.begin(), initialRanks.end()), + "NCCL error: rendezvous control ranks must be in canonical order"); + TLLM_CHECK_WITH_INFO(std::adjacent_find(initialRanks.begin(), initialRanks.end()) == initialRanks.end(), + "NCCL error: rendezvous control ranks contain duplicates"); + TLLM_CHECK_WITH_INFO(std::binary_search(initialRanks.begin(), initialRanks.end(), worldRank), + "NCCL error: world rank %d is not a member of the rendezvous control group", worldRank); + TLLM_CHECK_WITH_INFO(parentComm.getRank() == worldRank, + "NCCL error: world rank %d does not match MPI parent rank %d", worldRank, parentComm.getRank()); + TLLM_CHECK_WITH_INFO(initialRanks.front() >= 0 && initialRanks.back() < parentComm.getSize(), + "NCCL error: rendezvous control ranks must be within MPI parent size %d", parentComm.getSize()); + + auto cacheEntry = getControlCommCacheEntry(initialRanks, worldRank, parentComm, creationTagSeed); + std::lock_guard const cacheLock(cacheEntry->mutex); + if (cacheEntry->comm != nullptr) + { + return cacheEntry->comm; + } + if (cacheEntry->failure != nullptr) + { + // Creation is collective over the initial group. Never let one rank + // retry that collective after a prior asymmetric local failure while + // peers may already have cached a successfully-created channel. + std::rethrow_exception(cacheEntry->failure); + } + + void* tagUpperBoundValue = nullptr; + int hasTagUpperBound = 0; + int result + = MPI_Comm_get_attr(static_cast(parentComm), MPI_TAG_UB, &tagUpperBoundValue, &hasTagUpperBound); + TLLM_CHECK_WITH_INFO( + result == MPI_SUCCESS, "NCCL error: failed to query MPI_TAG_UB for the control channel (code %d)", result); + TLLM_CHECK_WITH_INFO(hasTagUpperBound != 0 && tagUpperBoundValue != nullptr, + "NCCL error: MPI parent communicator does not expose MPI_TAG_UB"); + int const tagUpperBound = *static_cast(tagUpperBoundValue); + constexpr int kFirstControlCreationTag = static_cast(mpi::MpiTag::kNcclPpControl) + 1; + TLLM_CHECK_WITH_INFO(creationTagSeed >= 0, "NCCL error: rendezvous control creation-tag seed must be nonnegative"); + int const tagParity = creationTagSeed & 1; + TLLM_CHECK_WITH_INFO(tagUpperBound >= kFirstControlCreationTag + tagParity, + "NCCL error: MPI_TAG_UB %d is too small for NCCL control-channel creation", tagUpperBound); + std::uint64_t tagHash = groupIdentity(initialRanks); + tagHash ^= static_cast(creationTagSeed) + 0x9e3779b97f4a7c15ULL + (tagHash << 6) + (tagHash >> 2); + // Even/odd slots separate raw-op and PP seeds; the group hash makes lazy + // creation order-independent within each ownership path. + auto const tagSlots = static_cast(tagUpperBound - kFirstControlCreationTag - tagParity) / 2 + 1; + int const creationTag = kFirstControlCreationTag + tagParity + static_cast(2 * (tagHash % tagSlots)); + + MPI_Group parentGroup = MPI_GROUP_NULL; + MPI_Group controlGroup = MPI_GROUP_NULL; + MPI_Comm controlComm = MPI_COMM_NULL; + auto cleanup = [&]() noexcept + { + if (controlGroup != MPI_GROUP_NULL) + { + MPI_Group_free(&controlGroup); + } + if (parentGroup != MPI_GROUP_NULL) + { + MPI_Group_free(&parentGroup); + } + // Once MPI_Comm_create_group has returned a non-null communicator, + // never call collective MPI_Comm_free from an exception path. A + // local allocation or error-handler failure can be asymmetric, so + // peers may already have published their process-lifetime channel. + // Leaking that rare failed construction is safer than deadlocking. + }; + + try + { + result = MPI_Comm_group(static_cast(parentComm), &parentGroup); + TLLM_CHECK_WITH_INFO(result == MPI_SUCCESS, "NCCL error: failed to inspect MPI parent group (code %d)", result); + result = MPI_Group_incl(parentGroup, static_cast(initialRanks.size()), initialRanks.data(), &controlGroup); + TLLM_CHECK_WITH_INFO(result == MPI_SUCCESS, "NCCL error: failed to create MPI control group (code %d)", result); + // MPI_Comm_create_group is collective only over initialRanks. This is + // the sole control-channel collective and must happen before failure. + result = MPI_Comm_create_group(static_cast(parentComm), controlGroup, creationTag, &controlComm); + TLLM_CHECK_WITH_INFO( + result == MPI_SUCCESS, "NCCL error: failed to create MPI control communicator (code %d)", result); + TLLM_CHECK_WITH_INFO(controlComm != MPI_COMM_NULL, + "NCCL error: MPI returned a null rendezvous control communicator for member rank %d", worldRank); + result = MPI_Comm_set_errhandler(controlComm, MPI_ERRORS_RETURN); + TLLM_CHECK_WITH_INFO(result == MPI_SUCCESS, + "NCCL error: failed to configure MPI_ERRORS_RETURN on the control channel (code %d)", result); + + MPI_Group_free(&controlGroup); + controlGroup = MPI_GROUP_NULL; + MPI_Group_free(&parentGroup); + parentGroup = MPI_GROUP_NULL; + // Never call MPI_Comm_free on a successfully-published FT control + // channel: freeing after a member dies is not guaranteed to make + // progress. The communicator intentionally lives until process exit. + auto ownedComm = mpi::MpiComm{controlComm, false}; + auto resultComm = std::make_shared(std::move(ownedComm), initialRanks, worldRank); + controlComm = MPI_COMM_NULL; + cacheEntry->comm = std::move(resultComm); + return cacheEntry->comm; + } + catch (...) + { + cacheEntry->failure = std::current_exception(); + cleanup(); + throw; + } +} + +ncclUniqueId exchangeNcclUniqueId(std::vector const& activeRanks, NcclUniqueIdRendezvousComm const& controlComm, + NcclUniqueIdRendezvousTags tags, std::uint64_t rendezvousId, Clock::time_point deadline) +{ + auto const worldRank = controlComm.worldRank(); + auto const& mpiComm = controlComm.mpiComm(); + TLLM_CHECK_WITH_INFO(!activeRanks.empty(), "NCCL error: communicator rendezvous group must not be empty"); + TLLM_CHECK_WITH_INFO(std::is_sorted(activeRanks.begin(), activeRanks.end()), + "NCCL error: communicator rendezvous ranks must be in canonical order"); + TLLM_CHECK_WITH_INFO(std::adjacent_find(activeRanks.begin(), activeRanks.end()) == activeRanks.end(), + "NCCL error: communicator rendezvous ranks contain duplicates"); + TLLM_CHECK_WITH_INFO(std::binary_search(activeRanks.begin(), activeRanks.end(), worldRank), + "NCCL error: world rank %d is not active in communicator rendezvous", worldRank); + TLLM_CHECK_WITH_INFO(std::includes(controlComm.worldRanks().begin(), controlComm.worldRanks().end(), + activeRanks.begin(), activeRanks.end()), + "NCCL error: communicator rendezvous group is not a subset of the pre-failure control group"); + + auto& context = getRendezvousContext(tags, controlComm); + std::unique_lock const exchangeLock(context.exchangeMutex); + int const root = activeRanks.front(); + std::uint64_t const key = communicatorKey(activeRanks, tags); + discardObsoletePending(context.pendingReady, key, rendezvousId); + discardObsoletePending(context.pendingIds, key, rendezvousId); + discardObsoletePending(context.pendingAcks, key, rendezvousId); + if (worldRank == root) + { + ncclUniqueId id{}; + { + auto const hostApiLock = acquireNcclHostApiLock(); + auto const result = ncclGetUniqueId(&id); + TLLM_CHECK_WITH_INFO( + result == ncclSuccess, "NCCL error: ncclGetUniqueId failed: %s", ncclGetErrorString(result)); + } + if (activeRanks.size() == 1) + { + return id; + } + + AttemptToken const serverToken = makeAttemptToken(); + std::uint64_t const digest = fnv1a(&id, sizeof(id)); + std::unordered_map> issuedTokens; + std::unordered_set acknowledged; + while (acknowledged.size() + 1 < activeRanks.size()) + { + TLLM_CHECK_WITH_INFO( + Clock::now() < deadline, "NCCL error: timed out waiting for survivor communicator rendezvous"); + bool progressed = false; + ReceivedMessage receivedReady{}; + bool haveReady = takePending(context.pendingReady, MPI_ANY_SOURCE, key, rendezvousId, receivedReady); + MPI_Status status{}; + if (!haveReady && probeMessage(MPI_ANY_SOURCE, tags.ready, mpiComm, status)) + { + progressed = true; + receivedReady = {controlComm.worldRank(status.MPI_SOURCE), + receiveMessage(status, tags.ready, mpiComm, deadline, "receiving NCCL READY")}; + haveReady = isCurrentAttempt(receivedReady.message, key, rendezvousId); + if (!haveReady && shouldStashForLater(receivedReady.message, key, rendezvousId)) + { + stashPending(context.pendingReady, receivedReady); + } + } + if (haveReady) + { + progressed = true; + ReadyMessage const& ready = receivedReady.message; + int const source = receivedReady.source; + if (ready.version == kProtocolVersion && isCurrentAttempt(ready, key, rendezvousId) + && isActivePeer(activeRanks, source, root)) + { + auto& tokens = issuedTokens[source]; + if (std::find(tokens.begin(), tokens.end(), ready.clientToken) == tokens.end()) + { + TLLM_CHECK_WITH_INFO(tokens.size() < kMaxTokensPerPeer, + "NCCL error: too many stale communicator rendezvous attempts from rank %d", source); + tokens.push_back(ready.clientToken); + } + IdMessage const proposal{ + kProtocolVersion, 0, key, rendezvousId, ready.clientToken, serverToken, digest, id}; + sendMessage( + proposal, controlComm.commRank(source), tags.id, mpiComm, deadline, "sending NCCL unique ID"); + } + } + + ReceivedMessage receivedAck{}; + bool haveAck = takePending(context.pendingAcks, MPI_ANY_SOURCE, key, rendezvousId, receivedAck); + if (!haveAck && probeMessage(MPI_ANY_SOURCE, tags.ack, mpiComm, status)) + { + progressed = true; + receivedAck = {controlComm.worldRank(status.MPI_SOURCE), + receiveMessage(status, tags.ack, mpiComm, deadline, "receiving NCCL ACK")}; + haveAck = isCurrentAttempt(receivedAck.message, key, rendezvousId); + if (!haveAck && shouldStashForLater(receivedAck.message, key, rendezvousId)) + { + stashPending(context.pendingAcks, receivedAck); + } + } + if (haveAck) + { + progressed = true; + AckMessage const& ack = receivedAck.message; + int const source = receivedAck.source; + auto const issued = issuedTokens.find(source); + if (ack.version == kProtocolVersion && isCurrentAttempt(ack, key, rendezvousId) + && ack.serverToken == serverToken && ack.idDigest == digest && issued != issuedTokens.end() + && std::find(issued->second.begin(), issued->second.end(), ack.clientToken) != issued->second.end()) + { + acknowledged.insert(source); + } + } + + if (!progressed) + { + TLLM_CHECK_WITH_INFO( + Clock::now() < deadline, "NCCL error: timed out waiting for survivor communicator rendezvous"); + std::this_thread::sleep_for(kPollInterval); + } + } + return id; + } + + AttemptToken const clientToken = makeAttemptToken(); + ReadyMessage const ready{kProtocolVersion, 0, key, rendezvousId, clientToken}; + int const rootCommRank = controlComm.commRank(root); + sendMessage(ready, rootCommRank, tags.ready, mpiComm, deadline, "sending NCCL READY"); + auto nextReadyRefresh = Clock::now() + kReadyRefreshInterval; + while (true) + { + TLLM_CHECK_WITH_INFO(Clock::now() < deadline, "NCCL error: timed out waiting for survivor communicator ID"); + if (Clock::now() >= nextReadyRefresh) + { + // A root attempt can consume READY and fail before accepting our + // ACK. Refresh the same idempotent client token so a subsequent + // root attempt can converge without waiting for this call to end. + sendMessage(ready, rootCommRank, tags.ready, mpiComm, deadline, "refreshing NCCL READY"); + nextReadyRefresh = Clock::now() + kReadyRefreshInterval; + } + ReceivedMessage receivedProposal{}; + bool haveProposal = takePending(context.pendingIds, root, key, rendezvousId, receivedProposal); + MPI_Status status{}; + if (!haveProposal && probeMessage(rootCommRank, tags.id, mpiComm, status)) + { + receivedProposal = {controlComm.worldRank(status.MPI_SOURCE), + receiveMessage(status, tags.id, mpiComm, deadline, "receiving NCCL unique ID")}; + haveProposal = isCurrentAttempt(receivedProposal.message, key, rendezvousId); + if (!haveProposal && shouldStashForLater(receivedProposal.message, key, rendezvousId)) + { + stashPending(context.pendingIds, receivedProposal); + } + } + if (!haveProposal) + { + TLLM_CHECK_WITH_INFO(Clock::now() < deadline, "NCCL error: timed out waiting for survivor communicator ID"); + std::this_thread::sleep_for(kPollInterval); + continue; + } + + IdMessage const& proposal = receivedProposal.message; + if (proposal.version != kProtocolVersion || !isCurrentAttempt(proposal, key, rendezvousId) + || !(proposal.clientToken == clientToken) || proposal.idDigest != fnv1a(&proposal.id, sizeof(proposal.id))) + { + sendMessage(ready, rootCommRank, tags.ready, mpiComm, deadline, "refreshing NCCL READY"); + nextReadyRefresh = Clock::now() + kReadyRefreshInterval; + continue; + } + + AckMessage const ack{ + kProtocolVersion, 0, key, rendezvousId, clientToken, proposal.serverToken, proposal.idDigest}; + sendMessage(ack, rootCommRank, tags.ack, mpiComm, deadline, "sending NCCL ACK"); + return proposal.id; + } +} + +} // namespace tensorrt_llm::runtime + +#endif // ENABLE_MULTI_DEVICE diff --git a/cpp/tensorrt_llm/thop/allgatherOp.cpp b/cpp/tensorrt_llm/thop/allgatherOp.cpp index 0d92aa966901..961aa8f7befa 100644 --- a/cpp/tensorrt_llm/thop/allgatherOp.cpp +++ b/cpp/tensorrt_llm/thop/allgatherOp.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -56,9 +56,9 @@ class AllgatherOp int initialize() { - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); mNcclComm = getComm(mGroup); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); return 0; } @@ -71,12 +71,11 @@ class AllgatherOp int64_t sum_sizes = sizes.has_value() ? std::accumulate(sizes.value().begin(), sizes.value().end(), 0, std::plus<>{}) : 0; std::vector output_list; + std::vector streams; + bool const trackOperations = isNcclFaultToleranceEnabled(); output_list.reserve(input_list.size()); - ncclGroupStart(); for (auto const& input : input_list) { - auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); - auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); std::vector outputShape = input.sizes().vec(); if (sizes.has_value()) { @@ -86,29 +85,68 @@ class AllgatherOp { outputShape[0] *= mGroup.size(); } - auto output = torch::empty(outputShape, input.options()); + output_list.push_back(torch::empty(outputShape, input.options())); + if (trackOperations) + { + auto const stream = at::cuda::getCurrentCUDAStream(input.get_device()); + if (std::find(streams.begin(), streams.end(), stream) == streams.end()) + { + streams.push_back(stream); + } + } + } + + auto commLease = acquireComm(mNcclComm); + auto const comm = commLease.get(); + std::vector watchdogTokens; + if (trackOperations) + { + watchdogTokens.reserve(streams.size()); + for (auto const stream : streams) + { + watchdogTokens.push_back(commLease.begin(stream, "ncclAllGather")); + } + } + commLease.groupStart("ncclGroupStart(allgather)"); + for (size_t index = 0; index < input_list.size(); ++index) + { + auto const& input = input_list[index]; + auto& output = output_list[index]; + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); if (use_nccl_allgather) { - ncclAllGather(input.data_ptr(), output.mutable_data_ptr(), input.numel(), (*getDtypeMap())[type], - *mNcclComm, stream); + commLease.checkEnqueue(ncclAllGather(input.data_ptr(), output.mutable_data_ptr(), input.numel(), + (*getDtypeMap())[type], comm, stream), + "ncclAllGather"); } else { + auto const& outputShape = output.sizes(); size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); int64_t split_offset = 0; for (int root = 0; root < static_cast(mGroup.size()); ++root) { auto split_size = sizes.value()[root]; - ncclBroadcast(input.data_ptr(), - output.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).mutable_data_ptr(), - numel_base * split_size, (*getDtypeMap())[type], root, *mNcclComm, stream); + commLease.checkEnqueue( + ncclBroadcast(input.data_ptr(), + output.index({torch::indexing::Slice(split_offset, torch::indexing::None)}) + .mutable_data_ptr(), + numel_base * split_size, (*getDtypeMap())[type], root, comm, stream), + "ncclBroadcast(allgather)"); split_offset += split_size; } } - output_list.push_back(output); } - NCCLCHECK_THROW(ncclGroupEnd()); + commLease.groupEnd("ncclGroupEnd(allgather)"); + if (trackOperations) + { + for (size_t index = 0; index < streams.size(); ++index) + { + commLease.track(watchdogTokens[index], streams[index]); + } + } return output_list; } diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 9f4de156ea55..3ab92117e6cf 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -265,9 +265,13 @@ class AllreduceOp int getRank() const { - return std::visit( - overloaded{[&](std::shared_ptr const&) { return COMM_SESSION.getRank(); }, - [&](c10::intrusive_ptr const& torchPg) { return get_world_pg()->getRank(); }}, + return std::visit(overloaded{[&](std::shared_ptr const&) + { + TLLM_CHECK_WITH_INFO( + mRank >= 0, "Raw AllreduceOp must be initialized before querying its rank"); + return mRank; + }, + [&](c10::intrusive_ptr const&) { return get_world_pg()->getRank(); }}, mNcclComm); } @@ -284,6 +288,12 @@ class AllreduceOp AllReduceStrategyType runtime_strategy = selectImplementation(seq_len, hidden_size); + if (mNcclComm.index() == 0 && isNcclFaultToleranceEnabled()) + { + TLLM_CHECK_WITH_INFO(runtime_strategy == AllReduceStrategyType::NCCL, + "NCCL error: TLLM_FAULT_TOLERANCE_MODE=1 requires a watchdog-managed NCCL allreduce strategy"); + } + // Log runtime strategy auto const rank = getRank(); TLLM_LOG_DEBUG( @@ -308,10 +318,14 @@ class AllreduceOp int initialize() { - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, getRank()); + int const initialRank = mNcclComm.index() == 0 ? -1 : getRank(); + TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, initialRank); if (mNcclComm.index() == 0) { + TLLM_CHECK_WITH_INFO(!isNcclFaultToleranceEnabled() || mStrategy == AllReduceStrategyType::NCCL, + "NCCL error: TLLM_FAULT_TOLERANCE_MODE=1 requires a watchdog-managed NCCL allreduce strategy"); mNcclComm = getComm(mGroup); + mRank = getCommWorldRank(std::get<0>(mNcclComm)); } if (mStrategy != AllReduceStrategyType::NCCL && mStrategy != AllReduceStrategyType::UB) { @@ -415,8 +429,12 @@ class AllreduceOp auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); int size = input.numel(); reduce_output = torch::empty_like(input); - NCCLCHECK_THROW(ncclAllReduce(input.data_ptr(), reduce_output.mutable_data_ptr(), size, - (*getDtypeMap())[mType], ncclSum, *rawComm, stream)); + auto commLease = acquireComm(rawComm); + auto const watchdogToken = commLease.begin(stream, "ncclAllReduce"); + commLease.check(ncclAllReduce(input.data_ptr(), reduce_output.mutable_data_ptr(), size, + (*getDtypeMap())[mType], ncclSum, commLease.get(), stream), + "ncclAllReduce"); + commLease.track(watchdogToken, stream); }, [&](c10::intrusive_ptr& torchPg) { @@ -462,8 +480,6 @@ class AllreduceOp // From here on, we have a raw NCCL comm - can proceed with window registration auto rawComm = std::get<0>(mNcclComm); - ncclComm_t comm = *rawComm; - TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator is null"); TLLM_LOG_DEBUG("[runNCCLAllReduceSymmetric] Using raw NCCL comm path (not ProcessGroup)"); using tensorrt_llm::common::nccl_util::NCCLWindowAllocator; @@ -486,7 +502,13 @@ class AllreduceOp double const a = -4986.43478503; double const b = 156716.52177552; int nRanks; - NCCLCHECK_THROW(ncclCommCount(comm, &nRanks)); + ncclComm_t comm = nullptr; + { + auto commLease = acquireComm(rawComm); + comm = commLease.get(); + TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator is null"); + commLease.check(ncclCommCount(comm, &nRanks), "ncclCommCount"); + } size_t minRegistrationThreshold = static_cast(std::max(0.0, a * nRanks + b)) * input.element_size(); // Disable window registration if neither NVLink nor MNNVL is supported // TODO replace in NCCL 2.29 with comm query @@ -566,7 +588,16 @@ class AllreduceOp } // Perform allreduce - NCCLCHECK_THROW(ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, comm, stream)); + { + auto commLease = acquireComm(rawComm); + auto const activeComm = commLease.get(); + TLLM_CHECK_WITH_INFO(activeComm == comm, "NCCL communicator changed while preparing a symmetric allreduce"); + auto const watchdogToken = commLease.begin(stream, "ncclAllReduce(symmetric)"); + commLease.check( + ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, activeComm, stream), + "ncclAllReduce(symmetric)"); + commLease.track(watchdogToken, stream); + } if (mOp == AllReduceFusionOp::NONE) { @@ -1466,6 +1497,7 @@ class AllreduceOp AllReduceFusionOp mOp; float mEps; std::variant, c10::intrusive_ptr> mNcclComm; + int mRank{-1}; }; } // namespace @@ -1488,7 +1520,7 @@ void preallocateNCCLWindowBuffer( } auto const commPtr = getComm(groupSet); - if (!commPtr || *commPtr == nullptr) + if (!commPtr) { TLLM_LOG_DEBUG("[preallocateNCCLWindowBuffers] NCCL comm is null; skipping preallocation"); return; @@ -1496,7 +1528,11 @@ void preallocateNCCLWindowBuffer( using tensorrt_llm::common::nccl_util::NCCLWindowAllocator; auto& allocator = NCCLWindowAllocator::getInstance(); - ncclComm_t const comm = *commPtr; + ncclComm_t comm = nullptr; + { + auto commLease = acquireComm(commPtr); + comm = commLease.get(); + } int64_t const numTokens = input.size(0); int64_t const elementsPerToken = input.numel() / numTokens; @@ -1518,7 +1554,7 @@ void preallocateNCCLWindowBuffer( { for (int64_t i = 0; i < buffersPerSize; ++i) { - auto buffer = allocator.requestBuffer(comm, bufferSize); + auto buffer = allocator.requestBuffer(commPtr, bufferSize); if (!buffer.isValid()) { break; @@ -1528,6 +1564,14 @@ void preallocateNCCLWindowBuffer( } catch (std::exception const& e) { + // A local allocation failure is a best-effort preallocation miss. A + // communicator failure is not: preserve the classifier-friendly NCCL + // exception so the executor enters recovery instead of continuing on + // the aborted native handle. + if (getCommAsyncError(commPtr).has_value()) + { + throw; + } TLLM_LOG_DEBUG("[preallocateNCCLWindowBuffer] requestBuffer failed for %zu bytes: %s", bufferSize, e.what()); } diff --git a/cpp/tensorrt_llm/thop/alltoallOp.cpp b/cpp/tensorrt_llm/thop/alltoallOp.cpp index 8a775b1cf6b8..8ae38fcf1ba4 100644 --- a/cpp/tensorrt_llm/thop/alltoallOp.cpp +++ b/cpp/tensorrt_llm/thop/alltoallOp.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -44,10 +44,10 @@ class AllToAllHelixOp int initialize() { - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); mNcclComm = getComm(mGroup); TLLM_CHECK_WITH_INFO(mGroup.size() > 0, "group size should be greater than 0"); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); + TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); return 0; } @@ -64,25 +64,37 @@ class AllToAllHelixOp TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); } std::vector output_list(static_cast(num_lists_)); - auto stream = at::cuda::getCurrentCUDAStream(input_list[0].get_device()); - ncclGroupStart(); for (int il = 0; il < num_lists_; ++il) { - auto off = il * num_ranks; + auto const off = il * num_ranks; auto output_shape = input_list[off].sizes().vec(); output_shape.insert(output_shape.begin(), num_ranks); - auto output = torch::empty(output_shape, input_list[off].options()); - output_list[il] = output; + output_list[il] = torch::empty(output_shape, input_list[off].options()); + } + + auto stream = at::cuda::getCurrentCUDAStream(input_list[0].get_device()); + auto commLease = acquireComm(mNcclComm); + auto const comm = commLease.get(); + auto const watchdogToken = commLease.begin(stream, "ncclSend/ncclRecv(alltoall)"); + commLease.groupStart("ncclGroupStart(alltoall)"); + for (int il = 0; il < num_lists_; ++il) + { + auto off = il * num_ranks; + auto& output = output_list[il]; auto type = tensorrt_llm::runtime::TorchUtils::dataType(input_list[off].scalar_type()); auto nccl_type = (*getDtypeMap())[type]; for (int r = 0; r < num_ranks; ++r) { auto const& input = input_list[off + r]; - ncclSend(input.data_ptr(), input.numel(), nccl_type, r, *mNcclComm, stream); - ncclRecv(output[r].mutable_data_ptr(), output[r].numel(), nccl_type, r, *mNcclComm, stream); + commLease.checkEnqueue( + ncclSend(input.data_ptr(), input.numel(), nccl_type, r, comm, stream), "ncclSend(alltoall)"); + commLease.checkEnqueue( + ncclRecv(output[r].mutable_data_ptr(), output[r].numel(), nccl_type, r, comm, stream), + "ncclRecv(alltoall)"); } } - NCCLCHECK_THROW(ncclGroupEnd()); + commLease.groupEnd("ncclGroupEnd(alltoall)"); + commLease.track(watchdogToken, stream); return output_list; } diff --git a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp index a45fa955a00b..d588f61e557a 100644 --- a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp +++ b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -16,19 +16,85 @@ #include "tensorrt_llm/thop/ncclCommunicatorOp.h" +#include "tensorrt_llm/common/opUtils.h" #include "tensorrt_llm/runtime/iBuffer.h" +#include +#include +#include +#include +#include + namespace tr = tensorrt_llm::runtime; TRTLLM_NAMESPACE_BEGIN namespace torch_ext { +namespace +{ + +std::mutex& pipelineCommunicatorMutex() +{ + static auto* mutex = new std::mutex; + return *mutex; +} + +struct PipelineCommunicatorCache +{ + int64_t worldSize{-1}; + int64_t rank{-1}; + std::shared_ptr communicator; +}; + +PipelineCommunicatorCache& pipelineCommunicatorCache() +{ + // Python model-engine reference counts are process-local and therefore + // cannot define a cross-rank NCCL teardown boundary. Keep the FT PP + // communicator alive for the process lifetime so rank-skewed wrapper churn + // never makes one rank reuse an old communicator while a peer enters a + // fresh-ID rendezvous. Coordinated teardown belongs to the Phase-2 owner. + static auto* cache = new PipelineCommunicatorCache; + return *cache; +} + +} // namespace NcclCommunicatorOp::NcclCommunicatorOp(int64_t worldSize, int64_t rank) : mRank(static_cast(rank)) { - mPipelineComm = std::make_shared(worldSize, rank); + TLLM_CHECK_WITH_INFO(worldSize > 0 && worldSize <= std::numeric_limits::max(), + "NCCL error: world size is out of range: %lld", static_cast(worldSize)); + TLLM_CHECK_WITH_INFO(rank >= 0 && rank < worldSize && rank <= std::numeric_limits::max(), + "NCCL error: rank is out of range: %lld", static_cast(rank)); + + if (!mpi::isFaultToleranceModeEnabled()) + { + // The legacy MPI-broadcast bootstrap assigns independent NCCL IDs, so + // multiple PP communicators remain valid when fault tolerance is off. + mPipelineComm = std::make_shared( + static_cast(worldSize), static_cast(rank)); + return; + } + + std::lock_guard const lock(pipelineCommunicatorMutex()); + auto& cache = pipelineCommunicatorCache(); + if (cache.communicator == nullptr) + { + cache.communicator = std::make_shared( + static_cast(worldSize), static_cast(rank)); + cache.worldSize = worldSize; + cache.rank = rank; + } + else + { + TLLM_CHECK_WITH_INFO(cache.worldSize == worldSize && cache.rank == rank, + "NCCL error: the process-lifetime PP communicator was initialized for world size %lld rank %lld, " + "not world size %lld rank %lld", + static_cast(cache.worldSize), static_cast(cache.rank), + static_cast(worldSize), static_cast(rank)); + } + mPipelineComm = cache.communicator; } void NcclCommunicatorOp::send(th::Tensor tensor, int64_t toRank) const @@ -49,6 +115,91 @@ void NcclCommunicatorOp::recv(th::Tensor& tensor, int64_t fromRank) const mPipelineComm->receive(*tr::IBuffer::wrap(ptr, size), static_cast(fromRank), cudaStream); } +void NcclCommunicatorOp::abort() +{ + mPipelineComm->abort(); +} + +void NcclCommunicatorOp::abortAndReinit(std::vector const& activeRanks, int64_t rendezvousId) +{ + TLLM_CHECK_WITH_INFO(rendezvousId > 0, "NCCL error: recovery rendezvous ID must be positive, got %lld", + static_cast(rendezvousId)); + std::vector ranks; + ranks.reserve(activeRanks.size()); + for (int64_t const rank : activeRanks) + { + TLLM_CHECK_WITH_INFO(rank >= 0 && rank <= std::numeric_limits::max(), + "NCCL error: active rank is out of range: %lld", static_cast(rank)); + ranks.push_back(static_cast(rank)); + } + mPipelineComm->abortAndReinit(ranks, static_cast(rendezvousId)); +} + +std::string NcclCommunicatorOp::getAsyncError() const +{ + return mPipelineComm->getAsyncError(); +} + +std::vector NcclCommunicatorOp::getActiveRanks() const +{ + auto const ranks = mPipelineComm->getActiveRanks(); + return std::vector(ranks.begin(), ranks.end()); +} + +void ncclCommAbortAndReinit( + std::vector const& oldGroup, std::vector const& activeGroup, int64_t rendezvousId) +{ +#if ENABLE_MULTI_DEVICE + TLLM_CHECK_WITH_INFO(rendezvousId > 0, "NCCL error: recovery rendezvous ID must be positive, got %lld", + static_cast(rendezvousId)); + std::set oldRanks; + std::set activeRanks; + for (auto const rank : oldGroup) + { + TLLM_CHECK_WITH_INFO(rank >= 0 && rank <= std::numeric_limits::max(), + "NCCL error: old-group rank is out of range: %lld", static_cast(rank)); + oldRanks.insert(static_cast(rank)); + } + for (auto const rank : activeGroup) + { + TLLM_CHECK_WITH_INFO(rank >= 0 && rank <= std::numeric_limits::max(), + "NCCL error: active-group rank is out of range: %lld", static_cast(rank)); + activeRanks.insert(static_cast(rank)); + } + TLLM_CHECK_WITH_INFO(oldRanks.size() == oldGroup.size(), "NCCL error: old group contains duplicate ranks"); + TLLM_CHECK_WITH_INFO(activeRanks.size() == activeGroup.size(), "NCCL error: active group contains duplicate ranks"); + abortAndReinitComm(oldRanks, activeRanks, static_cast(rendezvousId)); +#else + (void) oldGroup; + (void) activeGroup; + (void) rendezvousId; + TLLM_THROW("NCCL error: multi device support is disabled."); +#endif +} + +c10::optional ncclCommGetAsyncError(std::vector const& group) +{ +#if ENABLE_MULTI_DEVICE + std::set ranks; + for (auto const rank : group) + { + TLLM_CHECK_WITH_INFO(rank >= 0 && rank <= std::numeric_limits::max(), + "NCCL error: group rank is out of range: %lld", static_cast(rank)); + ranks.insert(static_cast(rank)); + } + TLLM_CHECK_WITH_INFO(ranks.size() == group.size(), "NCCL error: group contains duplicate ranks"); + auto error = getCommAsyncError(ranks); + if (error.has_value()) + { + return *error; + } + return c10::nullopt; +#else + (void) group; + return c10::nullopt; +#endif +} + } // namespace torch_ext TRTLLM_NAMESPACE_END @@ -57,4 +208,20 @@ static auto trtllmNcclCommunicator = torch::jit::class_("trtllm", "NcclCommunicatorOp") .def(torch::jit::init()) .def("send", &tensorrt_llm::torch_ext::NcclCommunicatorOp::send) - .def("recv", &tensorrt_llm::torch_ext::NcclCommunicatorOp::recv); + .def("recv", &tensorrt_llm::torch_ext::NcclCommunicatorOp::recv) + .def("abort", &tensorrt_llm::torch_ext::NcclCommunicatorOp::abort) + .def("abort_and_reinit", &tensorrt_llm::torch_ext::NcclCommunicatorOp::abortAndReinit) + .def("get_async_error", &tensorrt_llm::torch_ext::NcclCommunicatorOp::getAsyncError) + .def("get_active_ranks", &tensorrt_llm::torch_ext::NcclCommunicatorOp::getActiveRanks); + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("nccl_comm_abort_and_reinit(int[] old_group, int[] active_group, int rendezvous_id) -> ()"); + m.def("nccl_comm_get_async_error(int[] group) -> str?"); +} + +TORCH_LIBRARY_IMPL(trtllm, CompositeExplicitAutograd, m) +{ + m.impl("nccl_comm_abort_and_reinit", &tensorrt_llm::torch_ext::ncclCommAbortAndReinit); + m.impl("nccl_comm_get_async_error", &tensorrt_llm::torch_ext::ncclCommGetAsyncError); +} diff --git a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h index 38f4d215ac1a..a19378c1482c 100644 --- a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h +++ b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -18,7 +18,10 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/runtime/ncclCommunicator.h" #include "tensorrt_llm/thop/thUtils.h" +#include #include +#include +#include namespace th = torch; @@ -34,6 +37,10 @@ class NcclCommunicatorOp : public th::jit::CustomClassHolder void send(th::Tensor tensor, int64_t toRank) const; void recv(th::Tensor& tensor, int64_t fromRank) const; + void abort(); + void abortAndReinit(std::vector const& activeRanks, int64_t rendezvousId); + [[nodiscard]] std::string getAsyncError() const; + [[nodiscard]] std::vector getActiveRanks() const; private: int32_t mRank; diff --git a/cpp/tensorrt_llm/thop/outputTensor.h b/cpp/tensorrt_llm/thop/outputTensor.h index d9cde554852f..eac594652251 100644 --- a/cpp/tensorrt_llm/thop/outputTensor.h +++ b/cpp/tensorrt_llm/thop/outputTensor.h @@ -70,7 +70,7 @@ inline std::pair allocate_output(std::vector co { TLLM_LOG_DEBUG("[allocate_output] getComm threw (MPI disabled?): %s; fallback to default", e.what()); } - if (commPtr && *commPtr != nullptr) + if (commPtr) { auto [tensor, buffer] = tensorrt_llm::common::nccl_util::createNCCLWindowTensor(commPtr, output_size, dtype); diff --git a/cpp/tensorrt_llm/thop/reducescatterOp.cpp b/cpp/tensorrt_llm/thop/reducescatterOp.cpp index 40f89e40ff75..762a23efd1af 100644 --- a/cpp/tensorrt_llm/thop/reducescatterOp.cpp +++ b/cpp/tensorrt_llm/thop/reducescatterOp.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -28,7 +28,11 @@ #include #endif // ENABLE_MULTI_DEVICE +#include #include +#include +#include +#include #include #include @@ -57,7 +61,11 @@ class ReducescatterOp { TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, -1); mNcclComm = getComm(mGroup); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, -1); + auto const rank = getCommWorldRank(mNcclComm); + auto const rankIt = mGroup.find(rank); + TLLM_CHECK_WITH_INFO(rankIt != mGroup.end(), "Global rank %d is not in the reduce-scatter group", rank); + mGroupRank = static_cast(std::distance(mGroup.begin(), rankIt)); + TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, rank); return 0; } @@ -67,25 +75,14 @@ class ReducescatterOp bool use_nccl_reducescatter = !sizes.has_value() || std::all_of(sizes.value().begin(), sizes.value().end(), [&sizes](int64_t size) { return size == sizes.value()[0]; }); - int groupRank = 0; - if (sizes.has_value()) - { - auto rank = COMM_SESSION.getRank(); - for (auto const& currentRank : mGroup) - { - if (rank == currentRank) - break; - ++groupRank; - } - TLLM_CHECK(static_cast(groupRank) < mGroup.size()); - } + int const groupRank = sizes.has_value() ? mGroupRank : 0; + TLLM_CHECK_WITH_INFO(!sizes.has_value() || groupRank >= 0, "ReducescatterOp must be initialized before use"); std::vector output_list; + std::vector streams; + bool const trackOperations = isNcclFaultToleranceEnabled(); output_list.reserve(input_list.size()); - ncclGroupStart(); for (auto const& input : input_list) { - auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); - auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); std::vector outputShape = input.sizes().vec(); if (sizes.has_value()) { @@ -95,29 +92,68 @@ class ReducescatterOp { outputShape[0] = outputShape[0] / mGroup.size(); } - auto output = torch::empty(outputShape, input.options()); + output_list.push_back(torch::empty(outputShape, input.options())); + if (trackOperations) + { + auto const stream = at::cuda::getCurrentCUDAStream(input.get_device()); + if (std::find(streams.begin(), streams.end(), stream) == streams.end()) + { + streams.push_back(stream); + } + } + } + + auto commLease = acquireComm(mNcclComm); + auto const comm = commLease.get(); + std::vector watchdogTokens; + if (trackOperations) + { + watchdogTokens.reserve(streams.size()); + for (auto const stream : streams) + { + watchdogTokens.push_back(commLease.begin(stream, "ncclReduceScatter")); + } + } + commLease.groupStart("ncclGroupStart(reducescatter)"); + for (size_t index = 0; index < input_list.size(); ++index) + { + auto const& input = input_list[index]; + auto& output = output_list[index]; + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); if (use_nccl_reducescatter) { - ncclReduceScatter(input.data_ptr(), output.mutable_data_ptr(), output.numel(), (*getDtypeMap())[type], - ncclSum, *mNcclComm, stream); + commLease.checkEnqueue(ncclReduceScatter(input.data_ptr(), output.mutable_data_ptr(), output.numel(), + (*getDtypeMap())[type], ncclSum, comm, stream), + "ncclReduceScatter"); } else { + auto const& outputShape = output.sizes(); size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); int64_t split_offset = 0; for (int root = 0; root < static_cast(mGroup.size()); ++root) { auto split_size = sizes.value()[root]; - ncclReduce(input.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).data_ptr(), - output.mutable_data_ptr(), numel_base * split_size, (*getDtypeMap())[type], ncclSum, root, - *mNcclComm, stream); + commLease.checkEnqueue( + ncclReduce( + input.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).data_ptr(), + output.mutable_data_ptr(), numel_base * split_size, (*getDtypeMap())[type], ncclSum, root, + comm, stream), + "ncclReduce(reducescatter)"); split_offset += split_size; } } - output_list.push_back(output); } - NCCLCHECK_THROW(ncclGroupEnd()); + commLease.groupEnd("ncclGroupEnd(reducescatter)"); + if (trackOperations) + { + for (size_t index = 0; index < streams.size(); ++index) + { + commLease.track(watchdogTokens[index], streams[index]); + } + } return output_list; } @@ -129,6 +165,7 @@ class ReducescatterOp private: std::set mGroup; std::shared_ptr mNcclComm; + int mGroupRank{-1}; }; class ReducescatterPgOp diff --git a/cpp/tests/unit_tests/multi_gpu/CMakeLists.txt b/cpp/tests/unit_tests/multi_gpu/CMakeLists.txt index 669456551ed2..9787b650e066 100644 --- a/cpp/tests/unit_tests/multi_gpu/CMakeLists.txt +++ b/cpp/tests/unit_tests/multi_gpu/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & # AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); you may not @@ -19,6 +19,133 @@ add_gtest(cacheTransceiverTest cacheTransceiverTest.cpp) target_link_libraries(cacheTransceiverTest PRIVATE ${Python3_LIBRARIES}) add_gtest(mpiUtilsTest mpiUtilsTest.cpp) +add_gtest(ncclUniqueIdRendezvousTest ncclUniqueIdRendezvousTest.cpp) +if(ENABLE_MULTI_DEVICE) + set(NCCL_RENDEZVOUS_RAW_RECOVERY_TEST + NcclUniqueIdRendezvousTest.RawCommunicatorRebuildRunsPostRecoveryCollective + ) + set(NCCL_RENDEZVOUS_PP_RECOVERY_TEST + NcclUniqueIdRendezvousTest.PipelineCommunicatorRebuildRunsPostRecoverySendRecv + ) + set(NCCL_RENDEZVOUS_PROCESS_DEATH_TEST + NcclUniqueIdRendezvousTest.ProcessDeathRendezvousRequiresOptInFaultTolerantMpiLauncher + ) + set(NCCL_RENDEZVOUS_ISOLATED_TESTS + "${NCCL_RENDEZVOUS_RAW_RECOVERY_TEST}:${NCCL_RENDEZVOUS_PP_RECOVERY_TEST}:${NCCL_RENDEZVOUS_PROCESS_DEATH_TEST}" + ) + add_test( + NAME ncclUniqueIdRendezvousTest.mpi3 + COMMAND + ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 3 ${MPIEXEC_PREFLAGS} + $ + --gtest_filter=NcclUniqueIdRendezvousTest.*-${NCCL_RENDEZVOUS_ISOLATED_TESTS} + ${MPIEXEC_POSTFLAGS}) + set_tests_properties( + ncclUniqueIdRendezvousTest.mpi3 + PROPERTIES ENVIRONMENT + "CUDA_MODULE_LOADING=LAZY;TLLM_FAULT_TOLERANCE_MODE=1" + PROCESSORS 3 TIMEOUT 180) + + # Raw communicator replacement mutates process-global NCCL state and needs one + # distinct physical GPU per MPI rank. Keep it in its own CTest process. + add_test( + NAME ncclUniqueIdRendezvousTest.rawRecoveryMpi3 + COMMAND + ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 3 ${MPIEXEC_PREFLAGS} + $ + --gtest_filter=${NCCL_RENDEZVOUS_RAW_RECOVERY_TEST} ${MPIEXEC_POSTFLAGS}) + set_tests_properties( + ncclUniqueIdRendezvousTest.rawRecoveryMpi3 + PROPERTIES ENVIRONMENT + "CUDA_MODULE_LOADING=LAZY;TLLM_FAULT_TOLERANCE_MODE=1" + PROCESSORS 3 TIMEOUT 180) + + # The runtime PP wrapper owns a distinct communicator implementation and + # rendezvous tag set. Exercise its real post-recovery send/receive path in a + # separate process so it cannot share native singleton state with raw ops. + add_test( + NAME ncclUniqueIdRendezvousTest.ppRecoveryMpi3 + COMMAND + ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 3 ${MPIEXEC_PREFLAGS} + $ + --gtest_filter=${NCCL_RENDEZVOUS_PP_RECOVERY_TEST} ${MPIEXEC_POSTFLAGS}) + set_tests_properties( + ncclUniqueIdRendezvousTest.ppRecoveryMpi3 + PROPERTIES ENVIRONMENT + "CUDA_MODULE_LOADING=LAZY;TLLM_FAULT_TOLERANCE_MODE=1" + PROCESSORS 3 TIMEOUT 180) + + option(TRTLLM_ENABLE_NCCL_PROCESS_DEATH_TEST + "Enable the destructive NCCL rendezvous test that exits one MPI rank" + OFF) + set(TRTLLM_NCCL_PROCESS_DEATH_MPIEXEC_PREFLAGS + "" + CACHE + STRING + "Semicolon-separated launcher recovery flags for the NCCL process-death test" + ) + if(TRTLLM_ENABLE_NCCL_PROCESS_DEATH_TEST) + if("${TRTLLM_NCCL_PROCESS_DEATH_MPIEXEC_PREFLAGS}" STREQUAL "") + message( + FATAL_ERROR + "TRTLLM_ENABLE_NCCL_PROCESS_DEATH_TEST requires launcher-specific " + "TRTLLM_NCCL_PROCESS_DEATH_MPIEXEC_PREFLAGS (for example OpenMPI recovery flags)" + ) + endif() + # Keep a launcher-level signal from bypassing the sentinel properties. The + # CMake wrapper translates it to a normal nonzero process result. + set(NCCL_RENDEZVOUS_PROCESS_DEATH_ENV + CUDA_MODULE_LOADING=LAZY + TLLM_FAULT_TOLERANCE_MODE=1 + TRTLLM_TEST_NCCL_PROCESS_DEATH=1 + TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS=20000 + TRTLLM_NCCL_RECOVERY_TIMEOUT_MS=20000) + add_test( + NAME ncclUniqueIdRendezvousTest.processDeathMpi3 + COMMAND + ${CMAKE_COMMAND} -E env ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 3 + ${MPIEXEC_PREFLAGS} ${TRTLLM_NCCL_PROCESS_DEATH_MPIEXEC_PREFLAGS} + $ + --gtest_filter=${NCCL_RENDEZVOUS_PROCESS_DEATH_TEST} + ${MPIEXEC_POSTFLAGS}) + set_tests_properties( + ncclUniqueIdRendezvousTest.processDeathMpi3 + PROPERTIES ENVIRONMENT + "${NCCL_RENDEZVOUS_PROCESS_DEATH_ENV}" + PROCESSORS + 3 + TIMEOUT + 60 + RUN_SERIAL + TRUE + # A launcher may report failure solely because the injected + # victim intentionally skips MPI_Finalize. The root emits this + # sentinel only after survivor-only communicator rebuild and a + # successful allreduce, so it is the authoritative pass + # condition. + PASS_REGULAR_EXPRESSION + "TRTLLM_PROCESS_DEATH_RENDEZVOUS_OK" + FAIL_REGULAR_EXPRESSION + "TRTLLM_PROCESS_DEATH_RENDEZVOUS_FAIL") + endif() + + add_test( + NAME ncclUniqueIdRendezvousTest.defaultOff + COMMAND + ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 1 ${MPIEXEC_PREFLAGS} + $ + --gtest_filter=NcclUniqueIdRendezvousTest.FaultToleranceDisabledPreservesLegacyOperationPath + ${MPIEXEC_POSTFLAGS}) + set_tests_properties( + ncclUniqueIdRendezvousTest.defaultOff + PROPERTIES + ENVIRONMENT + "CUDA_MODULE_LOADING=LAZY;TLLM_FAULT_TOLERANCE_MODE=0;TRTLLM_TEST_FT_DEFAULT_OFF=1" + PROCESSORS + 1 + TIMEOUT + 60) +endif() add_gtest(userBufferTest userBufferTest.cpp) add_gtest(ncclUtilsTest ncclUtilsTest.cpp) target_link_libraries(ncclUtilsTest PRIVATE ${Python3_LIBRARIES}) diff --git a/cpp/tests/unit_tests/multi_gpu/ncclUniqueIdRendezvousTest.cpp b/cpp/tests/unit_tests/multi_gpu/ncclUniqueIdRendezvousTest.cpp new file mode 100644 index 000000000000..93c2b6d0a6c5 --- /dev/null +++ b/cpp/tests/unit_tests/multi_gpu/ncclUniqueIdRendezvousTest.cpp @@ -0,0 +1,971 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/runtime/utils/ncclUniqueIdRendezvous.h" +#include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/runtime/ncclCommunicator.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" + +#include + +#if ENABLE_MULTI_DEVICE + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ + +using tensorrt_llm::runtime::NcclUniqueIdRendezvousComm; +using tensorrt_llm::runtime::createNcclUniqueIdRendezvousComm; +using tensorrt_llm::runtime::exchangeNcclUniqueId; + +constexpr int kTestControlTagSeed = 0x4e43; +constexpr std::uint64_t kSubsetRendezvousId = 1; +constexpr std::uint64_t kStaleAttemptRendezvousId = 2; +constexpr std::uint64_t kRetryRendezvousId = 3; +constexpr std::uint64_t kStaleFloodFirstRendezvousId = 100; +constexpr std::uint64_t kStaleFloodCurrentRendezvousId = 10'000; +constexpr int kStaleFloodAttemptCount = 4'097; +constexpr std::uint64_t kRawRecoveryRendezvousId = 4; +constexpr std::uint64_t kRawCleanupRendezvousId = 5; +constexpr std::uint64_t kPpRecoveryRendezvousId = 6; +constexpr std::uint64_t kPpExcludedCleanupRendezvousId = 7; +constexpr std::uint64_t kProcessDeathRendezvousId = 8; + +tensorrt_llm::runtime::NcclUniqueIdRendezvousTags rawRendezvousTags() +{ + return {tensorrt_llm::mpi::MpiTag::kNcclCommReady, tensorrt_llm::mpi::MpiTag::kNcclCommUniqueId, + tensorrt_llm::mpi::MpiTag::kNcclCommAck}; +} + +std::vector allSessionRanks() +{ + std::vector ranks(COMM_SESSION.getSize()); + std::iota(ranks.begin(), ranks.end(), 0); + return ranks; +} + +std::shared_ptr createTestControlComm() +{ + return createNcclUniqueIdRendezvousComm( + allSessionRanks(), COMM_SESSION.getRank(), COMM_SESSION, kTestControlTagSeed); +} + +class EnvironmentVariableGuard +{ +public: + EnvironmentVariableGuard(char const* name, char const* value) + : mName(name) + { + if (auto const* previous = std::getenv(mName.c_str()); previous != nullptr) + { + mPrevious = previous; + } +#if defined(_WIN32) + _putenv_s(mName.c_str(), value); +#else + setenv(mName.c_str(), value, 1); +#endif + } + + ~EnvironmentVariableGuard() + { +#if defined(_WIN32) + _putenv_s(mName.c_str(), mPrevious.value_or("").c_str()); +#else + if (mPrevious.has_value()) + { + setenv(mName.c_str(), mPrevious->c_str(), 1); + } + else + { + unsetenv(mName.c_str()); + } +#endif + } + + EnvironmentVariableGuard(EnvironmentVariableGuard const&) = delete; + EnvironmentVariableGuard& operator=(EnvironmentVariableGuard const&) = delete; + +private: + std::string mName; + std::optional mPrevious; +}; + +bool isEnvironmentFlagEnabled(char const* name) +{ + auto const* value = std::getenv(name); + return value != nullptr && std::string{value} == "1"; +} + +bool configureDistinctCudaDeviceForEveryRank(std::string& failure) +{ + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + int deviceCount = 0; + cudaUUID_t deviceUuid{}; + int localReady = 0; + auto const countResult = cudaGetDeviceCount(&deviceCount); + if (countResult == cudaSuccess && deviceCount > 0) + { + int const device = rank % deviceCount; + auto const setResult = cudaSetDevice(device); + auto const uuidResult = setResult == cudaSuccess ? cudaDeviceGetUuid(&deviceUuid, device) : setResult; + localReady = setResult == cudaSuccess && uuidResult == cudaSuccess ? 1 : 0; + } + + std::vector gatheredReady(static_cast(worldSize)); + COMM_SESSION.allgather(&localReady, gatheredReady.data(), 1, tensorrt_llm::mpi::MpiType::kINT32); + std::vector gatheredUuids( + static_cast(worldSize) * static_cast(sizeof(deviceUuid.bytes))); + COMM_SESSION.allgather(deviceUuid.bytes, gatheredUuids.data(), static_cast(sizeof(deviceUuid.bytes)), + tensorrt_llm::mpi::MpiType::kBYTE); + + if (!std::all_of(gatheredReady.begin(), gatheredReady.end(), [](int ready) { return ready != 0; })) + { + failure = "every MPI rank must have a usable CUDA device"; + return false; + } + + std::unordered_set uniqueUuids; + for (int peer = 0; peer < worldSize; ++peer) + { + auto const* bytes = gatheredUuids.data() + + static_cast(peer) * static_cast(sizeof(deviceUuid.bytes)); + uniqueUuids.emplace(bytes, bytes + sizeof(deviceUuid.bytes)); + } + if (uniqueUuids.size() != static_cast(worldSize)) + { + failure = "every MPI rank must use a distinct physical GPU"; + return false; + } + return true; +} + +[[noreturn]] void exitProcessDeathTestFailure(int rank, std::string const& failure) noexcept +{ + std::fprintf(stderr, "TRTLLM_PROCESS_DEATH_RENDEZVOUS_FAIL rank=%d: %s\n", rank, failure.c_str()); + std::fflush(stderr); + std::_Exit(EXIT_FAILURE); +} + +void CUDART_CB waitForRelease(void* data) +{ + auto* release = static_cast*>(data); + while (!release->load(std::memory_order_acquire)) + { + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } +} + +class HostCallbackReleaseGuard +{ +public: + explicit HostCallbackReleaseGuard(std::atomic& release) + : mRelease(release) + { + } + + ~HostCallbackReleaseGuard() + { + mRelease.store(true, std::memory_order_release); + } + +private: + std::atomic& mRelease; +}; + +TEST(NcclUniqueIdRendezvousTest, DedicatedControlCommPreservesParentErrorHandlerAndRankMap) +{ + MPI_Errhandler parentBefore = MPI_ERRHANDLER_NULL; + ASSERT_EQ(MPI_Comm_get_errhandler(static_cast(COMM_SESSION), &parentBefore), MPI_SUCCESS); + + auto control = createTestControlComm(); + + MPI_Errhandler parentAfter = MPI_ERRHANDLER_NULL; + MPI_Errhandler controlHandler = MPI_ERRHANDLER_NULL; + ASSERT_EQ(MPI_Comm_get_errhandler(static_cast(COMM_SESSION), &parentAfter), MPI_SUCCESS); + ASSERT_EQ(MPI_Comm_get_errhandler(static_cast(control->mpiComm()), &controlHandler), MPI_SUCCESS); + EXPECT_EQ(parentAfter, parentBefore); + EXPECT_EQ(controlHandler, MPI_ERRORS_RETURN); + EXPECT_EQ(control->worldRank(), COMM_SESSION.getRank()); + EXPECT_EQ(control->worldRank(control->commRank(COMM_SESSION.getRank())), COMM_SESSION.getRank()); + + ASSERT_EQ(MPI_Errhandler_free(&parentBefore), MPI_SUCCESS); + ASSERT_EQ(MPI_Errhandler_free(&parentAfter), MPI_SUCCESS); + ASSERT_EQ(MPI_Errhandler_free(&controlHandler), MPI_SUCCESS); +} + +TEST(NcclUniqueIdRendezvousTest, DedicatedControlCommCacheReusesSameTopology) +{ + auto const firstControl = createTestControlComm(); + auto const secondControl = createTestControlComm(); + + ASSERT_NE(firstControl, nullptr); + ASSERT_NE(secondControl, nullptr); + EXPECT_EQ(firstControl.get(), secondControl.get()); + EXPECT_EQ(MPI_Comm_c2f(static_cast(firstControl->mpiComm())), + MPI_Comm_c2f(static_cast(secondControl->mpiComm()))); +} + +TEST(NcclUniqueIdRendezvousTest, DedicatedControlCommCacheSurvivesCustomParentHandleLifetime) +{ + auto const sessionControl = createTestControlComm(); + std::shared_ptr firstDuplicateControl; + { + MPI_Comm duplicate = MPI_COMM_NULL; + ASSERT_EQ(MPI_Comm_dup(static_cast(COMM_SESSION), &duplicate), MPI_SUCCESS); + auto duplicateParent = tensorrt_llm::mpi::MpiComm{duplicate, true}; + firstDuplicateControl = createNcclUniqueIdRendezvousComm( + allSessionRanks(), COMM_SESSION.getRank(), duplicateParent, kTestControlTagSeed); + } + + std::shared_ptr secondDuplicateControl; + { + MPI_Comm duplicate = MPI_COMM_NULL; + ASSERT_EQ(MPI_Comm_dup(static_cast(COMM_SESSION), &duplicate), MPI_SUCCESS); + auto duplicateParent = tensorrt_llm::mpi::MpiComm{duplicate, true}; + secondDuplicateControl = createNcclUniqueIdRendezvousComm( + allSessionRanks(), COMM_SESSION.getRank(), duplicateParent, kTestControlTagSeed); + } + + EXPECT_EQ(firstDuplicateControl.get(), sessionControl.get()); + EXPECT_EQ(secondDuplicateControl.get(), sessionControl.get()); +} + +TEST(NcclUniqueIdRendezvousTest, NonContiguousSurvivorSubsetExchangesWithoutExcludedRank) +{ + auto control = createTestControlComm(); + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize < 2) + { + GTEST_SKIP() << "Test requires at least two MPI processes"; + } + + // With three or more processes rank 1 is deliberately excluded, proving + // that recovery rendezvous does not contain a hidden world collective. + std::vector activeRanks{0, worldSize - 1}; + std::array localId{}; + if (std::binary_search(activeRanks.begin(), activeRanks.end(), rank)) + { + auto const id = exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), kSubsetRendezvousId, + std::chrono::steady_clock::now() + std::chrono::seconds{10}); + std::memcpy(localId.data(), &id, sizeof(id)); + } + + // This synchronization is test-only and happens after the survivor + // exchange. The production recovery path performs no MPI collective. + COMM_SESSION.barrier(); + std::vector gathered(static_cast(worldSize) * localId.size()); + COMM_SESSION.allgather( + localId.data(), gathered.data(), static_cast(localId.size()), tensorrt_llm::mpi::MpiType::kBYTE); + + auto const* first = gathered.data(); + auto const* last = gathered.data() + static_cast(worldSize - 1) * localId.size(); + EXPECT_TRUE(std::any_of(first, first + localId.size(), [](unsigned char value) { return value != 0; })); + EXPECT_TRUE(std::equal(first, first + localId.size(), last)); + if (worldSize >= 3) + { + auto const* excluded = gathered.data() + localId.size(); + EXPECT_TRUE(std::all_of(excluded, excluded + localId.size(), [](unsigned char value) { return value == 0; })); + } +} + +TEST(NcclUniqueIdRendezvousTest, RetryGenerationCannotPairWithStaleAttemptMessages) +{ + auto const control = createTestControlComm(); + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize < 2) + { + GTEST_SKIP() << "Test requires at least two MPI processes"; + } + + std::vector const activeRanks{0, worldSize - 1}; + bool const isRoot = rank == activeRanks.front(); + bool const isClient = rank == activeRanks.back(); + bool staleAttemptTimedOut = false; + bool retrySucceeded = false; + std::array retryId{}; + + // The root times out generation A before the client starts it, then enters + // generation B. The client deliberately sends a late generation-A READY + // while the root is already waiting in B. That message must be stashed, + // not accepted as B's READY; both ranks converge only after the client also + // advances to B. + control->mpiComm().barrier(); + if (isRoot) + { + try + { + static_cast(exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), + kStaleAttemptRendezvousId, std::chrono::steady_clock::now() + std::chrono::milliseconds{100})); + } + catch (std::exception const& error) + { + staleAttemptTimedOut = std::string{error.what()}.find("timed out") != std::string::npos; + } + + try + { + auto const id = exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), kRetryRendezvousId, + std::chrono::steady_clock::now() + std::chrono::seconds{5}); + std::memcpy(retryId.data(), &id, sizeof(id)); + retrySucceeded = true; + } + catch (std::exception const&) + { + // Assert after every rank reaches the test-only barrier. If the + // generation key regresses, the late A exchange can consume B and + // leave this retry waiting until its bounded deadline. + } + } + else if (isClient) + { + std::this_thread::sleep_for(std::chrono::milliseconds{500}); + try + { + static_cast(exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), + kStaleAttemptRendezvousId, std::chrono::steady_clock::now() + std::chrono::milliseconds{150})); + } + catch (std::exception const& error) + { + staleAttemptTimedOut = std::string{error.what()}.find("timed out") != std::string::npos; + } + + try + { + auto const id = exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), kRetryRendezvousId, + std::chrono::steady_clock::now() + std::chrono::seconds{5}); + std::memcpy(retryId.data(), &id, sizeof(id)); + retrySucceeded = true; + } + catch (std::exception const&) + { + // Keep the MPI test convergent on the regression path; the + // assertions below retain the failure details. + } + } + + control->mpiComm().barrier(); + std::vector gathered(static_cast(worldSize) * retryId.size()); + control->mpiComm().allgather( + retryId.data(), gathered.data(), static_cast(retryId.size()), tensorrt_llm::mpi::MpiType::kBYTE); + + if (isRoot || isClient) + { + EXPECT_TRUE(staleAttemptTimedOut); + EXPECT_TRUE(retrySucceeded); + } + auto const* rootId = gathered.data(); + auto const* clientId = gathered.data() + static_cast(worldSize - 1) * retryId.size(); + EXPECT_TRUE(std::any_of(rootId, rootId + retryId.size(), [](unsigned char value) { return value != 0; })); + EXPECT_TRUE(std::equal(rootId, rootId + retryId.size(), clientId)); +} + +TEST(NcclUniqueIdRendezvousTest, RepeatedStaleGenerationsDoNotExhaustPendingQueue) +{ + auto const control = createTestControlComm(); + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize < 2) + { + GTEST_SKIP() << "Test requires at least two MPI processes"; + } + + std::vector const activeRanks{0, worldSize - 1}; + bool const isRoot = rank == activeRanks.front(); + bool const isClient = rank == activeRanks.back(); + bool currentAttemptSucceeded = false; + int staleAttemptsTimedOut = 0; + std::array currentId{}; + + // The old opaque-key protocol stashed every lower-generation READY while + // the root waited on the current generation, then failed at its 4,096 + // pending-message cap. Distinct increasing stale IDs honor the caller + // contract while proving that same-communicator generations older than the + // root's current ID are discarded instead of retained indefinitely. + control->mpiComm().barrier(); + if (isRoot) + { + try + { + auto const id = exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), + kStaleFloodCurrentRendezvousId, std::chrono::steady_clock::now() + std::chrono::seconds{120}); + std::memcpy(currentId.data(), &id, sizeof(id)); + currentAttemptSucceeded = true; + } + catch (std::exception const&) + { + // Converge at the final test-only barrier even when the old queue + // limit is hit; the assertions below report the regression. + } + } + else if (isClient) + { + std::this_thread::sleep_for(std::chrono::milliseconds{100}); + for (int attempt = 0; attempt < kStaleFloodAttemptCount; ++attempt) + { + auto const staleId = kStaleFloodFirstRendezvousId + static_cast(attempt); + try + { + static_cast(exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), staleId, + std::chrono::steady_clock::now() + std::chrono::milliseconds{2})); + } + catch (std::exception const& error) + { + if (std::string{error.what()}.find("timed out waiting for survivor communicator ID") + != std::string::npos) + { + ++staleAttemptsTimedOut; + } + } + } + + try + { + auto const id = exchangeNcclUniqueId(activeRanks, *control, rawRendezvousTags(), + kStaleFloodCurrentRendezvousId, std::chrono::steady_clock::now() + std::chrono::seconds{60}); + std::memcpy(currentId.data(), &id, sizeof(id)); + currentAttemptSucceeded = true; + } + catch (std::exception const&) + { + // The root may already have failed on the old pending-queue cap. + } + } + + control->mpiComm().barrier(); + std::vector gathered(static_cast(worldSize) * currentId.size()); + control->mpiComm().allgather( + currentId.data(), gathered.data(), static_cast(currentId.size()), tensorrt_llm::mpi::MpiType::kBYTE); + + if (isRoot || isClient) + { + EXPECT_TRUE(currentAttemptSucceeded); + } + if (isClient) + { + EXPECT_EQ(staleAttemptsTimedOut, kStaleFloodAttemptCount); + } + auto const* rootId = gathered.data(); + auto const* clientId = gathered.data() + static_cast(worldSize - 1) * currentId.size(); + EXPECT_TRUE(std::any_of(rootId, rootId + currentId.size(), [](unsigned char value) { return value != 0; })); + EXPECT_TRUE(std::equal(rootId, rootId + currentId.size(), clientId)); +} + +TEST(NcclUniqueIdRendezvousTest, WaitOperationDoesNotHoldGlobalGateForUnrelatedStreamPrefix) +{ + EnvironmentVariableGuard const faultToleranceMode{"TLLM_FAULT_TOLERANCE_MODE", "1"}; + EnvironmentVariableGuard const operationTimeout{"TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS", "200"}; + EnvironmentVariableGuard const watchdogPoll{"TRTLLM_NCCL_WATCHDOG_POLL_INTERVAL_MS", "10"}; + + int deviceCount = 0; + ASSERT_EQ(cudaGetDeviceCount(&deviceCount), cudaSuccess); + if (deviceCount == 0) + { + GTEST_SKIP() << "Test requires a visible CUDA device"; + } + int const rank = COMM_SESSION.getRank(); + ASSERT_EQ(cudaSetDevice(rank % deviceCount), cudaSuccess); + + auto comm = tensorrt_llm::getComm({rank}); + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + std::atomic releasePrefix{false}; + HostCallbackReleaseGuard const releaseGuard{releasePrefix}; + ASSERT_EQ(cudaLaunchHostFunc(stream, waitForRelease, &releasePrefix), cudaSuccess); + + std::uint64_t token = 0; + { + auto lease = tensorrt_llm::acquireComm(comm); + token = lease.begin(stream, "operation behind unrelated stream prefix"); + lease.track(token, stream); + } + + std::atomic waiterStarted{false}; + std::exception_ptr waiterError; + std::thread waiter( + [&]() + { + try + { + waiterStarted.store(true, std::memory_order_release); + tensorrt_llm::waitCommOperation(comm, token, "operation behind unrelated stream prefix"); + } + catch (...) + { + waiterError = std::current_exception(); + } + }); + while (!waiterStarted.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + // Give the waiter time to enter its polling path before contending for the + // process-wide NCCL host gate from an unrelated control thread. + std::this_thread::sleep_for(std::chrono::milliseconds{100}); + + std::atomic gateAcquired{false}; + std::thread gateContender( + [&]() + { + auto gate = tensorrt_llm::runtime::acquireNcclHostApiLock(); + gateAcquired.store(true, std::memory_order_release); + }); + auto const gateDeadline = std::chrono::steady_clock::now() + std::chrono::seconds{2}; + while (!gateAcquired.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < gateDeadline) + { + std::this_thread::sleep_for(std::chrono::milliseconds{1}); + } + bool const acquiredBeforePrefixRelease = gateAcquired.load(std::memory_order_acquire); + + releasePrefix.store(true, std::memory_order_release); + gateContender.join(); + waiter.join(); + + EXPECT_TRUE(acquiredBeforePrefixRelease) + << "waitCommOperation retained the process-wide NCCL host gate while waiting for an unrelated stream prefix"; + EXPECT_FALSE(waiterError); + EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +TEST(NcclUniqueIdRendezvousTest, FaultToleranceRejectsGraphCaptureAndTimesOutStalledOperation) +{ + EnvironmentVariableGuard const faultToleranceMode{"TLLM_FAULT_TOLERANCE_MODE", "1"}; + EnvironmentVariableGuard const operationTimeout{"TRTLLM_NCCL_NONBLOCKING_TIMEOUT_MS", "100"}; + EnvironmentVariableGuard const watchdogPoll{"TRTLLM_NCCL_WATCHDOG_POLL_INTERVAL_MS", "10"}; + + int deviceCount = 0; + ASSERT_EQ(cudaGetDeviceCount(&deviceCount), cudaSuccess); + if (deviceCount == 0) + { + GTEST_SKIP() << "Test requires a visible CUDA device"; + } + int const rank = COMM_SESSION.getRank(); + ASSERT_EQ(cudaSetDevice(rank % deviceCount), cudaSuccess); + + std::set const singletonGroup{rank}; + auto comm = tensorrt_llm::getComm(singletonGroup); + cudaStream_t stream = nullptr; + int* captureScratch = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + ASSERT_EQ(cudaMalloc(reinterpret_cast(&captureScratch), sizeof(*captureScratch)), cudaSuccess); + + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(captureScratch, 0, sizeof(*captureScratch), stream), cudaSuccess); + { + auto lease = tensorrt_llm::acquireComm(comm); + EXPECT_THROW(static_cast(lease.begin(stream, "captured NCCL test operation")), std::exception); + } + cudaGraph_t graph = nullptr; + ASSERT_EQ(cudaStreamEndCapture(stream, &graph), cudaSuccess); + if (graph != nullptr) + { + ASSERT_EQ(cudaGraphDestroy(graph), cudaSuccess); + } + ASSERT_EQ(cudaFree(captureScratch), cudaSuccess); + + std::atomic releaseCallback{false}; + HostCallbackReleaseGuard const releaseGuard{releaseCallback}; + { + auto lease = tensorrt_llm::acquireComm(comm); + auto const token = lease.begin(stream, "injected stalled NCCL operation"); + ASSERT_EQ(cudaLaunchHostFunc(stream, waitForRelease, &releaseCallback), cudaSuccess); + lease.track(token, stream); + } + + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds{5}; + std::optional error; + while (std::chrono::steady_clock::now() < deadline) + { + error = tensorrt_llm::getCommAsyncError(singletonGroup); + if (error.has_value()) + { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds{10}); + } + EXPECT_TRUE(error.has_value()); + if (error.has_value()) + { + EXPECT_NE(error->find("timed out"), std::string::npos); + } + + releaseCallback.store(true, std::memory_order_release); + EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +TEST(NcclUniqueIdRendezvousTest, FaultToleranceDisabledPreservesLegacyOperationPath) +{ + if (!isEnvironmentFlagEnabled("TRTLLM_TEST_FT_DEFAULT_OFF")) + { + GTEST_SKIP() << "Run in the isolated default-off CTest process"; + } + ASSERT_FALSE(tensorrt_llm::mpi::isFaultToleranceModeEnabled()); + + int deviceCount = 0; + ASSERT_EQ(cudaGetDeviceCount(&deviceCount), cudaSuccess); + if (deviceCount == 0) + { + GTEST_SKIP() << "Test requires a visible CUDA device"; + } + int const rank = COMM_SESSION.getRank(); + ASSERT_EQ(cudaSetDevice(rank % deviceCount), cudaSuccess); + + std::set const singletonGroup{rank}; + auto comm = tensorrt_llm::getComm(singletonGroup); + auto lease = tensorrt_llm::acquireComm(comm); + EXPECT_EQ(lease.begin(nullptr, "default-off legacy operation"), 0); + EXPECT_EQ(lease.get(), *comm); + EXPECT_THROW( + tensorrt_llm::abortAndReinitComm(singletonGroup, singletonGroup, kRawRecoveryRendezvousId), std::exception); +} + +TEST(NcclUniqueIdRendezvousTest, ProcessDeathRendezvousRequiresOptInFaultTolerantMpiLauncher) +{ + if (!isEnvironmentFlagEnabled("TRTLLM_TEST_NCCL_PROCESS_DEATH")) + { + GTEST_SKIP() << "This destructive test requires a dedicated fault-tolerant MPI launcher"; + } + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize != 3) + { + exitProcessDeathTestFailure(rank, "the opt-in process-death test requires exactly three MPI processes"); + } + + try + { + std::string gpuFailure; + if (!configureDistinctCudaDeviceForEveryRank(gpuFailure)) + { + exitProcessDeathTestFailure(rank, gpuFailure); + } + + // Establish both the production raw-NCCL communicator and its retained + // MPI control channel while every rank is healthy. The separate test + // control channel supplies only the pre-failure synchronization below. + auto const control = createTestControlComm(); + auto const fullRanks = allSessionRanks(); + std::set const fullGroup(fullRanks.begin(), fullRanks.end()); + auto const fullComm = tensorrt_llm::getComm(fullGroup); + if (fullComm == nullptr || *fullComm == nullptr) + { + throw std::runtime_error("healthy full-group NCCL communicator is null"); + } + control->mpiComm().barrier(); + if (rank == 1) + { + std::_Exit(EXIT_SUCCESS); + } + + // This probes the stronger Mode-A contract: after one process exits, + // surviving ranks must still exchange a fresh NCCL ID over the retained + // MPI channel, abort the poisoned full communicator, and execute real + // data-plane work on its replacement. + std::this_thread::sleep_for(std::chrono::milliseconds{500}); + std::set const survivorGroup{0, 2}; + tensorrt_llm::abortAndReinitComm(fullGroup, survivorGroup, kProcessDeathRendezvousId); + auto const rebuiltComm = tensorrt_llm::getComm(survivorGroup); + if (rebuiltComm == nullptr || *rebuiltComm == nullptr) + { + throw std::runtime_error("rebuilt survivor NCCL communicator is null"); + } + + auto const checkCuda = [](cudaError_t result, char const* operation) + { + if (result != cudaSuccess) + { + throw std::runtime_error(std::string{operation} + " failed: " + cudaGetErrorString(result)); + } + }; + int const hostInput = rank + 1; + int hostOutput = 0; + int* deviceInput = nullptr; + int* deviceOutput = nullptr; + cudaStream_t stream = nullptr; + checkCuda(cudaMalloc(reinterpret_cast(&deviceInput), sizeof(hostInput)), "cudaMalloc(input)"); + checkCuda(cudaMalloc(reinterpret_cast(&deviceOutput), sizeof(hostOutput)), "cudaMalloc(output)"); + checkCuda(cudaStreamCreate(&stream), "cudaStreamCreate"); + checkCuda(cudaMemcpyAsync(deviceInput, &hostInput, sizeof(hostInput), cudaMemcpyHostToDevice, stream), + "cudaMemcpyAsync(input)"); + + std::uint64_t operationToken = 0; + { + auto lease = tensorrt_llm::acquireComm(rebuiltComm); + operationToken = lease.begin(stream, "ncclAllReduce(process-death recovery test)"); + lease.check(ncclAllReduce(deviceInput, deviceOutput, 1, ncclInt32, ncclSum, lease.get(), stream), + "ncclAllReduce(process-death recovery test)"); + lease.track(operationToken, stream); + } + tensorrt_llm::waitCommOperation(rebuiltComm, operationToken, "ncclAllReduce(process-death recovery test)"); + checkCuda( + cudaMemcpy(&hostOutput, deviceOutput, sizeof(hostOutput), cudaMemcpyDeviceToHost), "cudaMemcpy(output)"); + constexpr int kExpectedSurvivorSum = 4; + if (hostOutput != kExpectedSurvivorSum) + { + throw std::runtime_error("post-recovery allreduce returned " + std::to_string(hostOutput) + ", expected " + + std::to_string(kExpectedSurvivorSum)); + } + + if (rank == *survivorGroup.begin()) + { + // Some fault-tolerant launchers still return nonzero because the + // victim intentionally skipped MPI_Finalize. CTest keys off this + // flushed sentinel, reached only after survivor communicator + // rebuild and the expected allreduce, instead of that exit code. + std::fprintf(stdout, "TRTLLM_PROCESS_DEATH_RENDEZVOUS_OK\n"); + std::fflush(stdout); + } + std::_Exit(EXIT_SUCCESS); + } + catch (std::exception const& error) + { + exitProcessDeathTestFailure(rank, error.what()); + } + catch (...) + { + exitProcessDeathTestFailure(rank, "unknown exception"); + } +} + +TEST(NcclUniqueIdRendezvousTest, RawCommunicatorRebuildRunsPostRecoveryCollective) +{ + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize < 3) + { + GTEST_SKIP() << "Test requires at least three MPI processes"; + } + EnvironmentVariableGuard const faultToleranceMode{"TLLM_FAULT_TOLERANCE_MODE", "1"}; + + int deviceCount = 0; + ASSERT_EQ(cudaGetDeviceCount(&deviceCount), cudaSuccess); + if (deviceCount == 0) + { + GTEST_SKIP() << "Test requires a visible CUDA device"; + } + int const device = rank % deviceCount; + ASSERT_EQ(cudaSetDevice(device), cudaSuccess); + + // A launcher may expose all GPUs to every rank or bind each rank to one + // different physical GPU. Compare UUIDs rather than local ordinals so both + // layouts work, and never run multiple NCCL ranks on the same device. + cudaUUID_t deviceUuid{}; + ASSERT_EQ(cudaDeviceGetUuid(&deviceUuid, device), cudaSuccess); + std::vector gatheredUuids( + static_cast(worldSize) * static_cast(sizeof(deviceUuid.bytes))); + COMM_SESSION.allgather(deviceUuid.bytes, gatheredUuids.data(), static_cast(sizeof(deviceUuid.bytes)), + tensorrt_llm::mpi::MpiType::kBYTE); + std::unordered_set uniqueUuids; + for (int peer = 0; peer < worldSize; ++peer) + { + auto const* bytes = gatheredUuids.data() + + static_cast(peer) * static_cast(sizeof(deviceUuid.bytes)); + uniqueUuids.emplace(bytes, bytes + sizeof(deviceUuid.bytes)); + } + if (uniqueUuids.size() != static_cast(worldSize)) + { + GTEST_SKIP() << "Test requires one distinct physical GPU per MPI rank"; + } + + auto const fullRanks = allSessionRanks(); + std::set const fullGroup(fullRanks.begin(), fullRanks.end()); + std::set const survivorGroup{0, worldSize - 1}; + auto const testSync = createTestControlComm(); + auto oldComm = tensorrt_llm::getComm(fullGroup); + ASSERT_NE(oldComm, nullptr); + EXPECT_EQ(tensorrt_llm::getCommWorldRank(oldComm), rank); + + // Keep the old registry state alive so its storage cannot be recycled; + // pointer inequality then proves recovery installed a fresh state. + std::shared_ptr cachedSurvivorComm; + ncclComm_t* cachedSurvivorState = nullptr; + if (survivorGroup.count(rank) != 0) + { + cachedSurvivorComm = tensorrt_llm::getComm(survivorGroup); + ASSERT_NE(cachedSurvivorComm, nullptr); + cachedSurvivorState = cachedSurvivorComm.get(); + } + + // Ensure the survivor-target cache entry exists before any rank begins + // recovery. Use a dedicated test communicator so the excluded rank does + // not enter a parent-communicator collective while survivors create their + // subgroup control channel. + testSync->mpiComm().barrier(); + + if (survivorGroup.count(rank) != 0) + { + tensorrt_llm::abortAndReinitComm(fullGroup, survivorGroup, kRawRecoveryRendezvousId); + auto rebuiltComm = tensorrt_llm::getComm(survivorGroup); + ASSERT_NE(rebuiltComm, nullptr); + EXPECT_EQ(tensorrt_llm::getCommWorldRank(rebuiltComm), rank); + EXPECT_NE(rebuiltComm.get(), cachedSurvivorState); + EXPECT_EQ(tensorrt_llm::getComm(survivorGroup).get(), rebuiltComm.get()); + + int const hostInput = rank + 1; + int hostOutput = 0; + int* deviceInput = nullptr; + int* deviceOutput = nullptr; + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaMalloc(reinterpret_cast(&deviceInput), sizeof(hostInput)), cudaSuccess); + ASSERT_EQ(cudaMalloc(reinterpret_cast(&deviceOutput), sizeof(hostOutput)), cudaSuccess); + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + ASSERT_EQ( + cudaMemcpyAsync(deviceInput, &hostInput, sizeof(hostInput), cudaMemcpyHostToDevice, stream), cudaSuccess); + + std::uint64_t operationToken = 0; + { + auto lease = tensorrt_llm::acquireComm(rebuiltComm); + operationToken = lease.begin(stream, "ncclAllReduce(recovery test)"); + lease.check(ncclAllReduce(deviceInput, deviceOutput, 1, ncclInt32, ncclSum, lease.get(), stream), + "ncclAllReduce(recovery test)"); + lease.track(operationToken, stream); + } + tensorrt_llm::waitCommOperation(rebuiltComm, operationToken, "ncclAllReduce(recovery test)"); + + ASSERT_EQ(cudaMemcpy(&hostOutput, deviceOutput, sizeof(hostOutput), cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_EQ(hostOutput, worldSize + 1); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); + EXPECT_EQ(cudaFree(deviceOutput), cudaSuccess); + EXPECT_EQ(cudaFree(deviceInput), cudaSuccess); + + // Leave only a singleton communicator for process teardown. This is + // test cleanup, not another simulated survivor-membership decision. + tensorrt_llm::abortAndReinitComm(survivorGroup, {rank}, kRawCleanupRendezvousId); + } + else + { + // The excluded rank deliberately does not participate in the survivor + // rendezvous. Once survivors have started rebuilding, replace its old + // local communicator with a singleton so test-process teardown never + // calls ncclCommDestroy on a half-aborted full-group communicator. + tensorrt_llm::abortAndReinitComm(fullGroup, {rank}, kRawCleanupRendezvousId); + } + + // Test-only synchronization after every local recovery path has finished. + testSync->mpiComm().barrier(); +} + +TEST(NcclUniqueIdRendezvousTest, PipelineCommunicatorRebuildRunsPostRecoverySendRecv) +{ + int const rank = COMM_SESSION.getRank(); + int const worldSize = COMM_SESSION.getSize(); + if (worldSize < 3) + { + GTEST_SKIP() << "Test requires at least three MPI processes"; + } + EnvironmentVariableGuard const faultToleranceMode{"TLLM_FAULT_TOLERANCE_MODE", "1"}; + + int deviceCount = 0; + ASSERT_EQ(cudaGetDeviceCount(&deviceCount), cudaSuccess); + if (deviceCount == 0) + { + GTEST_SKIP() << "Test requires a visible CUDA device"; + } + int const device = rank % deviceCount; + ASSERT_EQ(cudaSetDevice(device), cudaSuccess); + + cudaUUID_t deviceUuid{}; + ASSERT_EQ(cudaDeviceGetUuid(&deviceUuid, device), cudaSuccess); + std::vector gatheredUuids( + static_cast(worldSize) * static_cast(sizeof(deviceUuid.bytes))); + COMM_SESSION.allgather(deviceUuid.bytes, gatheredUuids.data(), static_cast(sizeof(deviceUuid.bytes)), + tensorrt_llm::mpi::MpiType::kBYTE); + std::unordered_set uniqueUuids; + for (int peer = 0; peer < worldSize; ++peer) + { + auto const* bytes = gatheredUuids.data() + + static_cast(peer) * static_cast(sizeof(deviceUuid.bytes)); + uniqueUuids.emplace(bytes, bytes + sizeof(deviceUuid.bytes)); + } + if (uniqueUuids.size() != static_cast(worldSize)) + { + GTEST_SKIP() << "Test requires one distinct physical GPU per MPI rank"; + } + + auto const testSync = createTestControlComm(); + auto communicator = std::make_unique(worldSize, rank); + std::vector const survivorRanks{0, worldSize - 1}; + bool const isSurvivor = std::binary_search(survivorRanks.begin(), survivorRanks.end(), rank); + + testSync->mpiComm().barrier(); + if (isSurvivor) + { + communicator->abortAndReinit(survivorRanks, kPpRecoveryRendezvousId); + EXPECT_EQ(communicator->getActiveRanks(), survivorRanks); + } + else + { + communicator->abortAndReinit({rank}, kPpExcludedCleanupRendezvousId); + EXPECT_EQ(communicator->getActiveRanks(), std::vector{rank}); + } + EXPECT_TRUE(communicator->getAsyncError().empty()); + + int* deviceValue = nullptr; + ASSERT_EQ(cudaMalloc(reinterpret_cast(&deviceValue), sizeof(*deviceValue)), cudaSuccess); + int hostValue = rank == survivorRanks.front() ? 0x1a07 : 0; + ASSERT_EQ(cudaMemcpy(deviceValue, &hostValue, sizeof(hostValue), cudaMemcpyHostToDevice), cudaSuccess); + auto buffer = tensorrt_llm::runtime::IBuffer::wrap(deviceValue, nvinfer1::DataType::kINT32, 1); + tensorrt_llm::runtime::CudaStream stream; + + if (rank == survivorRanks.front()) + { + communicator->send(*buffer, survivorRanks.back(), stream); + stream.synchronize(); + } + else if (rank == survivorRanks.back()) + { + communicator->receive(*buffer, survivorRanks.front(), stream); + stream.synchronize(); + ASSERT_EQ(cudaMemcpy(&hostValue, deviceValue, sizeof(hostValue), cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_EQ(hostValue, 0x1a07); + } + + buffer.reset(); + ASSERT_EQ(cudaFree(deviceValue), cudaSuccess); + testSync->mpiComm().barrier(); + communicator.reset(); + testSync->mpiComm().barrier(); +} + +} // namespace + +#endif // ENABLE_MULTI_DEVICE diff --git a/cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp b/cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp index 4f3309431a2f..5c8aeab428f8 100644 --- a/cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/ncclUtilsTest.cpp @@ -20,7 +20,9 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/opUtils.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" +#include "tensorrt_llm/runtime/utils/ncclHostApi.h" +#include #include #include #include @@ -87,34 +89,46 @@ TEST(NCCLWindowSupportTest, RuntimeVersionAndGB10Gate) // Helper function to create a split communicator for testing // This allows us to test cleanup behavior explicitly by controlling the lifetime -std::shared_ptr createSplitComm(ncclComm_t parentComm, int color, int key) +std::shared_ptr createSplitComm(std::shared_ptr const& parentComm, int color, int key) { - ncclComm_t newComm; - ncclResult_t result = ncclCommSplit(parentComm, color, key, &newComm, nullptr); - if (result != ncclSuccess) + ncclComm_t newComm = nullptr; + ncclConfig_t config = NCCL_CONFIG_INITIALIZER; + // The cached parent is nonblocking. Make the test-owned child blocking so + // the legacy raw allocator overload and its simple deleter have an + // explicit blocking-communicator contract. + config.blocking = 1; { - TLLM_THROW("ncclCommSplit failed with error: %d", result); + auto lease = tensorrt_llm::acquireComm(parentComm); + auto const result = ncclCommSplit(lease.get(), color, key, &newComm, &config); + lease.check(result, "ncclCommSplit(ncclUtilsTest)"); + } + if (newComm == nullptr) + { + TLLM_THROW("ncclCommSplit returned a null child communicator"); } // Create a shared_ptr with custom deleter that cleans up resources first return std::shared_ptr(new ncclComm_t(newComm), [](ncclComm_t* comm) { - if (comm && *comm) + if (comm != nullptr && *comm != nullptr) { // STEP 1: Clean up all registered resources FIRST - tensorrt_llm::common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(*comm); + bool const resourcesCleaned + = tensorrt_llm::common::nccl_util::NcclCommResourceManager::getInstance().cleanupResources(*comm); - // STEP 2: Now destroy the NCCL communicator - ncclResult_t result = ncclCommDestroy(*comm); + // STEP 2: Serialize with the cached parent's watchdog and + // destroy only after resource cleanup fully completed. + auto const hostApiLock = tensorrt_llm::runtime::acquireNcclHostApiLock(); + ncclResult_t const result = resourcesCleaned ? ncclCommDestroy(*comm) : ncclCommAbort(*comm); if (result != ncclSuccess) { - TLLM_LOG_WARNING("ncclCommDestroy failed with error: %d", result); + TLLM_LOG_WARNING("NCCL test communicator release failed with error: %d", result); } - - // STEP 3: Free the memory - delete comm; + *comm = nullptr; } + // STEP 3: Free the holder even if communicator creation failed. + delete comm; }); } @@ -122,6 +136,106 @@ std::shared_ptr createSplitComm(ncclComm_t parentComm, int color, in // NcclCommResourceManager Tests //============================================================================== +TEST(NcclCommResourceManagerStandaloneTest, ReusedHandleWaitsForAbortCleanup) +{ + auto& manager = nccl_util::NcclCommResourceManager::getInstance(); + int handleStorage = 0; + auto const comm = reinterpret_cast(&handleStorage); + std::atomic waitStarted{false}; + std::atomic waitCompleted{false}; + + manager.beginAbortCleanup(comm); + std::thread waiter( + [&]() + { + waitStarted.store(true, std::memory_order_release); + manager.waitForAbortCleanup(comm); + waitCompleted.store(true, std::memory_order_release); + }); + while (!waitStarted.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + EXPECT_FALSE(waitCompleted.load(std::memory_order_acquire)); + + EXPECT_TRUE(manager.cleanupResources(comm, true)); + waiter.join(); + EXPECT_TRUE(waitCompleted.load(std::memory_order_acquire)); +} + +TEST(NcclCommResourceManagerStandaloneTest, ConcurrentCleanupCannotReleaseRetiredHandleEarly) +{ + auto& manager = nccl_util::NcclCommResourceManager::getInstance(); + int handleStorage = 0; + auto const comm = reinterpret_cast(&handleStorage); + std::atomic cleanupEntered{false}; + std::atomic releaseCleanup{false}; + std::atomic cleanupSucceeded{false}; + manager.registerResource( + comm, + [&]() + { + cleanupEntered.store(true, std::memory_order_release); + while (!releaseCleanup.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + return true; + }, + "BlockingAbortCleanup"); + manager.beginAbortCleanup(comm); + + std::thread cleanupOwner( + [&]() { cleanupSucceeded.store(manager.cleanupResources(comm, true), std::memory_order_release); }); + while (!cleanupEntered.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + + EXPECT_TRUE(manager.cleanupResources(comm, true)); + EXPECT_TRUE(manager.isAbortCleanup(comm)); + releaseCleanup.store(true, std::memory_order_release); + cleanupOwner.join(); + EXPECT_TRUE(cleanupSucceeded.load(std::memory_order_acquire)); + EXPECT_FALSE(manager.isAbortCleanup(comm)); +} + +TEST(NcclCommResourceManagerStandaloneTest, NormalCleanupOwnsRetirementUntilCallbacksComplete) +{ + auto& manager = nccl_util::NcclCommResourceManager::getInstance(); + int handleStorage = 0; + auto const comm = reinterpret_cast(&handleStorage); + std::atomic cleanupEntered{false}; + std::atomic releaseCleanup{false}; + std::atomic cleanupSucceeded{false}; + manager.registerResource( + comm, + [&]() + { + cleanupEntered.store(true, std::memory_order_release); + while (!releaseCleanup.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + return true; + }, + "BlockingNormalCleanup"); + std::thread cleanupOwner( + [&]() { cleanupSucceeded.store(manager.cleanupResources(comm, false), std::memory_order_release); }); + while (!cleanupEntered.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + + manager.beginAbortCleanup(comm); + EXPECT_TRUE(manager.cleanupResources(comm, true)); + EXPECT_TRUE(manager.isAbortCleanup(comm)); + releaseCleanup.store(true, std::memory_order_release); + cleanupOwner.join(); + EXPECT_TRUE(cleanupSucceeded.load(std::memory_order_acquire)); + EXPECT_FALSE(manager.isAbortCleanup(comm)); +} + class NcclCommResourceManagerTest : public ::testing::Test { protected: @@ -170,12 +284,18 @@ TEST_F(NcclCommResourceManagerTest, ResourceRegistration) auto& manager = nccl_util::NcclCommResourceManager::getInstance(); // Create a separate comm using split for this test - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); // Register a resource bool cleanupCalled = false; manager.registerResource( - *testComm, [&cleanupCalled]() { cleanupCalled = true; }, "TestResource"); + *testComm, + [&cleanupCalled]() + { + cleanupCalled = true; + return true; + }, + "TestResource"); EXPECT_TRUE(manager.hasResources(*testComm)); EXPECT_EQ(manager.getResourceCount(*testComm), 1); @@ -202,15 +322,33 @@ TEST_F(NcclCommResourceManagerTest, MultipleResources) auto& manager = nccl_util::NcclCommResourceManager::getInstance(); // Create a separate comm using split for this test - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); std::vector cleanupOrder; manager.registerResource( - *testComm, [&cleanupOrder]() { cleanupOrder.push_back(1); }, "Resource1"); + *testComm, + [&cleanupOrder]() + { + cleanupOrder.push_back(1); + return true; + }, + "Resource1"); manager.registerResource( - *testComm, [&cleanupOrder]() { cleanupOrder.push_back(2); }, "Resource2"); + *testComm, + [&cleanupOrder]() + { + cleanupOrder.push_back(2); + return true; + }, + "Resource2"); manager.registerResource( - *testComm, [&cleanupOrder]() { cleanupOrder.push_back(3); }, "Resource3"); + *testComm, + [&cleanupOrder]() + { + cleanupOrder.push_back(3); + return true; + }, + "Resource3"); EXPECT_EQ(manager.getResourceCount(*testComm), 3); @@ -229,22 +367,34 @@ TEST_F(NcclCommResourceManagerTest, ResourceCount) auto& manager = nccl_util::NcclCommResourceManager::getInstance(); // Create a separate comm using split for this test - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); EXPECT_FALSE(manager.hasResources(*testComm)); EXPECT_EQ(manager.getResourceCount(*testComm), 0); manager.registerResource( - *testComm, []() {}, "Test1"); + *testComm, []() { return true; }, "Test1"); EXPECT_EQ(manager.getResourceCount(*testComm), 1); manager.registerResource( - *testComm, []() {}, "Test2"); + *testComm, []() { return true; }, "Test2"); EXPECT_EQ(manager.getResourceCount(*testComm), 2); testComm.reset(); } +TEST_F(NcclCommResourceManagerTest, CleanupFailureIsReportedToCommunicatorOwner) +{ + auto& manager = nccl_util::NcclCommResourceManager::getInstance(); + auto testComm = createSplitComm(mComm, 0, mRank); + + manager.registerResource( + *testComm, []() { return false; }, "InjectedCleanupFailure"); + + EXPECT_FALSE(manager.cleanupResources(*testComm)); + EXPECT_FALSE(manager.hasResources(*testComm)); +} + //============================================================================== // NCCLWindowAllocator Tests //============================================================================== @@ -302,7 +452,7 @@ TEST_F(NCCLWindowAllocatorTest, BasicAllocation) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); const size_t bufferSize = 1024 * 1024; // 1MB - auto buffer = allocator.requestBuffer(*mComm, bufferSize); + auto buffer = allocator.requestBuffer(mComm, bufferSize); EXPECT_TRUE(buffer.isValid()); EXPECT_NE(buffer.ptr, nullptr); @@ -326,7 +476,7 @@ TEST_F(NCCLWindowAllocatorTest, BufferReuse) const size_t bufferSize = 512 * 1024; // 512KB // Allocate first buffer - auto buffer1 = allocator.requestBuffer(*mComm, bufferSize); + auto buffer1 = allocator.requestBuffer(mComm, bufferSize); EXPECT_TRUE(buffer1.isValid()); void* ptr1 = buffer1.ptr; @@ -334,7 +484,7 @@ TEST_F(NCCLWindowAllocatorTest, BufferReuse) allocator.releaseBuffer(*mComm, ptr1); // Request another buffer of the same size - should reuse - auto buffer2 = allocator.requestBuffer(*mComm, bufferSize); + auto buffer2 = allocator.requestBuffer(mComm, bufferSize); EXPECT_TRUE(buffer2.isValid()); EXPECT_EQ(buffer2.ptr, ptr1); // Should be the same buffer @@ -346,9 +496,9 @@ TEST_F(NCCLWindowAllocatorTest, BestFitReuse) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); // Allocate buffers of different sizes - auto buffer1MB = allocator.requestBuffer(*mComm, 1024 * 1024); - auto buffer2MB = allocator.requestBuffer(*mComm, 2 * 1024 * 1024); - auto buffer512KB = allocator.requestBuffer(*mComm, 512 * 1024); + auto buffer1MB = allocator.requestBuffer(mComm, 1024 * 1024); + auto buffer2MB = allocator.requestBuffer(mComm, 2 * 1024 * 1024); + auto buffer512KB = allocator.requestBuffer(mComm, 512 * 1024); void* ptr1MB = buffer1MB.ptr; void* ptr2MB = buffer2MB.ptr; @@ -360,7 +510,7 @@ TEST_F(NCCLWindowAllocatorTest, BestFitReuse) allocator.releaseBuffer(*mComm, ptr512KB); // Request 768KB - should reuse 1MB (best fit, smallest that fits) - auto buffer768KB = allocator.requestBuffer(*mComm, 768 * 1024); + auto buffer768KB = allocator.requestBuffer(mComm, 768 * 1024); EXPECT_TRUE(buffer768KB.isValid()); EXPECT_EQ(buffer768KB.ptr, ptr1MB); // Should reuse 1MB buffer EXPECT_EQ(buffer768KB.size, 1024 * 1024); // Original size @@ -371,7 +521,7 @@ TEST_F(NCCLWindowAllocatorTest, BestFitReuse) TEST_F(NCCLWindowAllocatorTest, FailureCacheIsSizeAwareForNewAllocations) { auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); constexpr size_t failureSize = 1024 * 1024; nccl_util::NCCLWindowAllocatorTestAccess::recordSymmetricFailure(allocator, *testComm, failureSize); @@ -391,7 +541,7 @@ TEST_F(NCCLWindowAllocatorTest, FailureCacheIsSizeAwareForNewAllocations) TEST_F(NCCLWindowAllocatorTest, FailureCacheDoesNotDisableReusableBuffers) { auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); auto buffer1MB = allocator.requestBuffer(*testComm, 1024 * 1024); ASSERT_TRUE(buffer1MB.isValid()); @@ -416,7 +566,7 @@ TEST_F(NCCLWindowAllocatorTest, FailureCacheDoesNotDisableReusableBuffers) TEST_F(NCCLWindowAllocatorTest, FailureCacheKeepsSmallestFailureSize) { auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); nccl_util::NCCLWindowAllocatorTestAccess::recordSymmetricFailure(allocator, *testComm, 2 * 1024 * 1024); nccl_util::NCCLWindowAllocatorTestAccess::recordSymmetricFailure(allocator, *testComm, 1024 * 1024); @@ -459,7 +609,7 @@ TEST_F(NCCLWindowAllocatorTest, MultipleBuffers) // Allocate multiple buffers for (int i = 0; i < 5; ++i) { - auto buffer = allocator.requestBuffer(*mComm, bufferSize); + auto buffer = allocator.requestBuffer(mComm, bufferSize); EXPECT_TRUE(buffer.isValid()); ptrs.push_back(buffer.ptr); } @@ -482,7 +632,7 @@ TEST_F(NCCLWindowAllocatorTest, SearchBuffer) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); const size_t bufferSize = 128 * 1024; - auto buffer = allocator.requestBuffer(*mComm, bufferSize); + auto buffer = allocator.requestBuffer(mComm, bufferSize); // Test searchBuffer auto found = allocator.searchBuffer(*mComm, buffer.ptr); @@ -505,7 +655,7 @@ TEST_F(NCCLWindowAllocatorTest, GetWindowAndSize) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); const size_t bufferSize = 64 * 1024; - auto buffer = allocator.requestBuffer(*mComm, bufferSize); + auto buffer = allocator.requestBuffer(mComm, bufferSize); // Test getWindow auto window = allocator.getWindow(*mComm, buffer.ptr); @@ -530,7 +680,7 @@ TEST_F(NCCLWindowAllocatorTest, GetBufferInfo) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); const size_t bufferSize = 32 * 1024; - auto buffer = allocator.requestBuffer(*mComm, bufferSize); + auto buffer = allocator.requestBuffer(mComm, bufferSize); auto info = allocator.getBufferInfo(*mComm, buffer.ptr); EXPECT_TRUE(info.isValid()); @@ -570,7 +720,7 @@ TEST_F(NCCLWindowAllocatorTest, CleanupOnCommDestroy) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); // Create a separate comm using split for this test - auto testComm = createSplitComm(*mComm, 0, mRank); + auto testComm = createSplitComm(mComm, 0, mRank); // Store the raw comm value before destruction ncclComm_t rawComm = *testComm; @@ -627,8 +777,8 @@ TEST_F(NCCLWindowAllocatorTest, MultipleComms) auto& allocator = nccl_util::NCCLWindowAllocator::getInstance(); // Create two different communicators using split (different colors) - auto comm1 = createSplitComm(*mComm, 0, mRank); - auto comm2 = createSplitComm(*mComm, 1, mRank); + auto comm1 = createSplitComm(mComm, 0, mRank); + auto comm2 = createSplitComm(mComm, 1, mRank); const size_t bufferSize = 4 * 1024; @@ -831,6 +981,25 @@ TEST_F(CreateNCCLWindowTensorTest, TensorDeleterReleasesBuffer) EXPECT_GE(allocator.getBufferCount(*mComm), 1); } +TEST_F(CreateNCCLWindowTensorTest, TensorKeepsGenericCommunicatorOwnerAlive) +{ + if (tensorrt_llm::isNcclFaultToleranceEnabled()) + { + GTEST_SKIP() << "Generic communicator holders use the legacy default-off allocator contract"; + } + auto comm = createSplitComm(mComm, 0, mRank); + std::weak_ptr weakComm = comm; + std::vector shape = {16, 16}; + auto [tensor, buffer] = nccl_util::createNCCLWindowTensor(comm, shape, torch::kFloat32); + ASSERT_TRUE(tensor.defined()); + ASSERT_TRUE(buffer.isValid()); + + comm.reset(); + EXPECT_FALSE(weakComm.expired()); + tensor = torch::Tensor(); + EXPECT_TRUE(weakComm.expired()); +} + TEST_F(CreateNCCLWindowTensorTest, MultipleTensors) { using nccl_util::createNCCLWindowTensor; diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index af4157786c06..334438484fec 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -18,7 +18,7 @@ import platform import sys from dataclasses import dataclass -from typing import List, Optional, Union +from typing import List, Optional, Tuple, Union import pynvml import torch @@ -400,19 +400,50 @@ class HelixCpMnnvlMemory(MnnvlMemory): its own isolated state separate from MnnvlMemory. """ + comm_topology = None + + @staticmethod + def get_comm_topology( + mapping: Mapping, + ) -> Tuple[int, int, int, Tuple[int, ...]]: + """Return the exact local CP topology used by ``MPI_Comm_split``.""" + return ( + int(mapping.world_size), + int(mapping.rank), + int(mapping.cp_rank), + tuple(int(rank) for rank in mapping.cp_group), + ) + @classmethod def get_comm(cls, mapping: Mapping): """Get CP-based communicator (ranks grouped by PP+TP+MOE_TP, ordered by CP rank).""" + topology = cls.get_comm_topology(mapping) if cls.comm is not None: + if cls.comm_topology != topology: + raise RuntimeError( + "Helix MNNVL CP communicator is already initialized for " + f"topology {cls.comm_topology}; cannot reuse it for " + f"incompatible topology {topology}" + ) return cls.comm comm = mpi_comm().Split( mapping.pp_rank * mapping.tp_size + mapping.tp_rank, mapping.cp_rank, ) cls.comm = comm + cls.comm_topology = topology return comm +def get_helix_cp_mnnvl_topology( + mapping: Mapping, +) -> Optional[Tuple[int, int, int, Tuple[int, ...]]]: + """Return the Helix MNNVL topology relevant to communicator sharing.""" + if mapping.has_cp_helix() and not mapping.cp_config.get("use_nccl_for_alltoall", True): + return HelixCpMnnvlMemory.get_comm_topology(mapping) + return None + + def init_helix_cp_comm(mapping: Mapping) -> None: """Pre-initialize the Helix CP communicator. diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 7b5d199624fa..9d5e7a6cb013 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -56,6 +56,12 @@ # BufferKind is bound from C++; see cpp/tensorrt_llm/thop/outputTensor.h (torch_ext::BufferKind). from tensorrt_llm.bindings.internal.thop import BufferKind +# Communicator mode is startup-only. Keep this local cached flag so importing +# the custom-op registry does not recursively initialize the distributed +# package merely to filter allreduce tactics. +NCCL_FAULT_TOLERANCE_ENABLED = os.environ.get( + "TLLM_FAULT_TOLERANCE_MODE") == "1" + def _init_deep_gemm_pdl() -> None: try: @@ -74,6 +80,16 @@ def _init_deep_gemm_pdl() -> None: logger.warning(f"Failed to initialize DeepGEMM PDL: {err}") +def _use_cuda_graph_for_allreduce_tuning() -> bool: + """Whether raw-NCCL tactics may be captured by the autotuner. + + The 1a.7 communicator can be replaced after failure, so captured NCCL + nodes are intentionally rejected in FT mode until graph invalidation and + recapture land in 1a.11. Eager profiling preserves valid tactic timings. + """ + return os.environ.get("TLLM_FAULT_TOLERANCE_MODE") != "1" + + _init_deep_gemm_pdl() @@ -2200,6 +2216,8 @@ def get_valid_tactics( AllReduceStrategy.NCCL_SYMMETRIC.value, AllReduceStrategy.NCCL.value, ] + if NCCL_FAULT_TOLERANCE_ENABLED: + return [AllReduceStrategy.NCCL.value] # Fallback in allreduceOp is set to NCCL_SYMMETRIC as default # So we need to check if the workspace size is too large to avoid hanging. workspace_size = inputs[0].numel() * inputs[0].element_size() @@ -2242,10 +2260,17 @@ def forward( return input if tactic == -1: # tactic == -1 means the autotuner cache missed for this shape. - # Fall back to NCCL_SYMMETRIC; asymmetric ncclMemAlloc failures are handled - # by a cross-rank barrier in NCCLWindowAllocator which falls back - # to plain NCCL. - tactic = AllReduceStrategy.NCCL_SYMMETRIC.value + # FT mode must avoid NCCL_SYMMETRIC's MPI-local topology discovery; + # otherwise retain the optimized default-off fallback. + tactic = (AllReduceStrategy.NCCL.value + if NCCL_FAULT_TOLERANCE_ENABLED else + AllReduceStrategy.NCCL_SYMMETRIC.value) + + if (NCCL_FAULT_TOLERANCE_ENABLED + and tactic != AllReduceStrategy.NCCL.value): + raise RuntimeError( + "NCCL error: fault-tolerant allreduce tuning selected a " + "custom tactic without a peer-failure abort path") return torch.ops.trtllm.allreduce( input, @@ -2315,7 +2340,8 @@ def _inputs_pre_hook_register_nccl_symmetric_memory_window( tuning_config = replace( AllReduceRunner.tuning_config, - inputs_pre_hook=_inputs_pre_hook_register_nccl_symmetric_memory_window) + inputs_pre_hook=_inputs_pre_hook_register_nccl_symmetric_memory_window, + use_cuda_graph=_use_cuda_graph_for_allreduce_tuning()) _, best_tactic = tuner.choose_one( "trtllm::tunable_allreduce::allreduce", diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 7e78bf36ee66..2ec47c0efc8a 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -33,7 +33,11 @@ except Exception: MPI = None # deferred; functions will error if used when ENABLE_MULTI_DEVICE is True -from tensorrt_llm._mnnvl_utils import init_helix_cp_comm +from tensorrt_llm._mnnvl_utils import (get_helix_cp_mnnvl_topology, + init_helix_cp_comm) +from tensorrt_llm._torch.distributed.nccl_fault_tolerance import ( + NCCL_FAULT_TOLERANCE_ENABLED, _canonical_recovery_generation, + _recovery_rendezvous_id) from tensorrt_llm._utils import (mpi_allgather, mpi_barrier, mpi_comm, mpi_disabled, mpi_isend, mpi_isend_object, mpi_recv, mpi_recv_object, mpi_send, @@ -1636,16 +1640,67 @@ class PPCommNCCL: def __init__(self, global_mapping: Mapping): self.mapping = global_mapping + self._topology = self._mapping_topology(global_mapping) + self._reconfigure_lock = threading.Lock() + self._reconfigure_generation = 0 + self._completed_recovery = None self.nccl_comm = torch.classes.trtllm.NcclCommunicatorOp( self.mapping.world_size, self.mapping.rank, ) + if NCCL_FAULT_TOLERANCE_ENABLED: + # The FT native communicator survives process-local Python wrapper + # churn. Mirror its actual survivor membership instead of assuming + # a newly-created wrapper always represents the initial world. + self._active_ranks = tuple( + int(rank) for rank in self.nccl_comm.get_active_ranks()) + else: + # Preserve the legacy constructor path when FT is disabled. + self._active_ranks = tuple(range(self.mapping.world_size)) self.tensor_ready_event = torch.cuda.Event() self.send_stream = torch.cuda.Stream() + @staticmethod + def _mapping_topology(mapping: Mapping) -> tuple: + return ( + int(mapping.world_size), + int(mapping.rank), + tuple(mapping.pp_group), + get_helix_cp_mnnvl_topology(mapping), + ) + + def is_compatible(self, mapping: Mapping) -> bool: + """Whether another engine can share this process-local communicator.""" + return self._topology == self._mapping_topology(mapping) + + def _validate_peer(self, peer: int) -> int: + peer = int(peer) + if peer not in self._topology[2]: + raise RuntimeError( + f"NCCL error: peer world rank {peer} is not in this rank's PP group " + f"{list(self._topology[2])}") + if peer not in self._active_ranks: + raise RuntimeError( + f"NCCL error: peer world rank {peer} is not active in the current PP communicator" + ) + return peer + + def _required_pp_peer(self, *, next_peer: bool) -> int: + # Do not silently skip a failed stage: connecting non-adjacent stages + # would bypass model layers. A higher-level topology reconstruction may + # pass an explicit remapped peer once it has reassigned those layers. + peer = (self.mapping.next_pp_rank() + if next_peer else self.mapping.prev_pp_rank()) + return self._validate_peer(peer) + def send(self, tensor: torch.Tensor, dest: Optional[int] = None): - if dest is None: - dest = self.mapping.next_pp_rank() + if not NCCL_FAULT_TOLERANCE_ENABLED: + if dest is None: + dest = self.mapping.next_pp_rank() + elif dest is None: + dest = self._required_pp_peer(next_peer=True) + else: + dest = self._validate_peer(dest) # NCCL send kernel in send_stream cannot be captured, # so we send in the current stream instead in CUDA graph cases. @@ -1663,10 +1718,110 @@ def send(self, tensor: torch.Tensor, dest: Optional[int] = None): self.nccl_comm.send(tensor, dest) def recv(self, tensor: torch.Tensor, src: Optional[int] = None): - if src is None: - src = self.mapping.prev_pp_rank() + if not NCCL_FAULT_TOLERANCE_ENABLED: + if src is None: + src = self.mapping.prev_pp_rank() + elif src is None: + src = self._required_pp_peer(next_peer=False) + else: + src = self._validate_peer(src) self.nccl_comm.recv(tensor, src) + def abort(self, generation: Optional[int] = None) -> None: + """Abort the PP communicator on this survivor.""" + # Capture the generation before waiting for a concurrent rebuild. If + # that rebuild succeeds, this abort belongs to the old communicator + # and must not tear down its replacement. An abort that starts after + # the rebuild observes the new generation and still takes effect. + if generation is None: + generation = self._reconfigure_generation + with self._reconfigure_lock: + if generation != self._reconfigure_generation: + return + self.nccl_comm.abort() + + def abort_and_reinit(self, + active_ranks: List[int], + generation: Optional[int] = None) -> None: + """Rebuild the PP communicator using only surviving world ranks. + + Every survivor must call this with the same rank set. The C++ wrapper + exchanges a fresh NCCL unique ID point-to-point, so excluded ranks do + not participate in communicator initialization. + + ``generation`` is the coordinator's shared monotonic recovery-event + generation, advanced for every distinct attempt. It is required for + same-membership recovery so every rank enters the native rendezvous + even when local watchdog observations differ. Repeated callbacks with + the same generation and target are idempotent. A retry after any native + failure must use a newly advanced generation so stale rendezvous traffic + cannot pair different recovery attempts. + """ + canonical_ranks = tuple(sorted(int(rank) for rank in active_ranks)) + if not canonical_ranks: + raise ValueError("active_ranks must not be empty") + if len(canonical_ranks) != len(set(canonical_ranks)): + raise ValueError("active_ranks must not contain duplicates") + if self.mapping.rank not in canonical_ranks: + raise ValueError( + f"current world rank {self.mapping.rank} is not active") + generation = _canonical_recovery_generation(generation) + rendezvous_id = _recovery_rendezvous_id(generation) + + # Failure notifications can be duplicated or delivered concurrently by + # independent control paths. Keep the native rebuild and its Python + # membership mirror in one transaction so a stale callback cannot + # widen membership after a newer shrink has completed. + with self._reconfigure_lock: + current_ranks = tuple( + int(rank) for rank in self.nccl_comm.get_active_ranks()) + self._active_ranks = current_ranks + if generation is not None and self._completed_recovery is not None: + completed_generation, completed_target = self._completed_recovery + if generation < completed_generation: + return + if generation == completed_generation: + if canonical_ranks != completed_target: + raise RuntimeError( + "NCCL error: conflicting PP communicator recovery " + f"target for generation {generation}: completed " + f"{list(completed_target)}, requested {list(canonical_ranks)}" + ) + return + if canonical_ranks == current_ranks and generation is None: + raise ValueError( + "generation is required for same-membership PP communicator " + "recovery so every survivor makes the same rebuild decision" + ) + if not set(canonical_ranks).issubset(current_ranks): + raise ValueError( + "abort_and_reinit cannot reactivate a removed rank") + + self.nccl_comm.abort_and_reinit(list(canonical_ranks), + rendezvous_id) + native_ranks = tuple( + int(rank) for rank in self.nccl_comm.get_active_ranks()) + self._active_ranks = native_ranks + self._reconfigure_generation += 1 + if native_ranks != canonical_ranks: + raise RuntimeError( + "NCCL error: rebuilt PP communicator returned unexpected " + f"membership {list(native_ranks)}; expected {list(canonical_ranks)}" + ) + if generation is not None: + self._completed_recovery = (generation, canonical_ranks) + + def get_async_error(self) -> str: + """Return the C++ wrapper's latched NCCL abort/error reason, if any.""" + return self.nccl_comm.get_async_error() + + def get_active_ranks(self) -> List[int]: + """Return original world-rank IDs in the current PP communicator.""" + with self._reconfigure_lock: + self._active_ranks = tuple( + int(rank) for rank in self.nccl_comm.get_active_ranks()) + return list(self._active_ranks) + class PPCommTorch: @@ -1698,16 +1853,149 @@ def recv(self, tensor: torch.Tensor, src: Optional[int] = None): _pp_comm = None +_pp_comm_refcount = 0 +_pp_comm_lock = threading.Lock() +_pp_comm_condition = threading.Condition(_pp_comm_lock) +_pp_comm_control_refcount = 0 +_pp_comm_final_release_pending = False def init_pp_comm(mapping): - """Initialize PPComm once at startup""" - global _pp_comm - if mpi_disabled(): - _pp_comm = PPCommTorch(mapping) - else: - _pp_comm = PPCommNCCL(mapping) - init_helix_cp_comm(mapping) + """Acquire the process-local PP communicator for one model engine.""" + global _pp_comm, _pp_comm_refcount + with _pp_comm_condition: + while (_pp_comm_control_refcount > 0 or _pp_comm_final_release_pending): + _pp_comm_condition.wait() + created = False + if mpi_disabled(): + if _pp_comm is None: + _pp_comm = PPCommTorch(mapping) + created = True + elif not isinstance(_pp_comm, PPCommTorch): + raise RuntimeError( + "PP communicator is already initialized with a different backend" + ) + elif _pp_comm.mapping != mapping: + raise RuntimeError( + "PP communicator is already initialized for a different topology" + ) + else: + if _pp_comm is None: + _pp_comm = PPCommNCCL(mapping) + created = True + elif not isinstance(_pp_comm, PPCommNCCL): + raise RuntimeError( + "NCCL error: PP communicator is already initialized with a different backend" + ) + elif not _pp_comm.is_compatible(mapping): + raise RuntimeError( + "NCCL error: PP communicator is already initialized for a " + "different topology (PP or Helix MNNVL CP topology)") + try: + init_helix_cp_comm(mapping) + except Exception: + if created and _pp_comm_refcount == 0: + _pp_comm = None + raise + _pp_comm_refcount += 1 + + +def release_pp_comm() -> None: + """Release one model engine's ownership of the shared PP communicator.""" + global _pp_comm, _pp_comm_refcount, _pp_comm_final_release_pending + with _pp_comm_condition: + if _pp_comm_refcount <= 0: + raise RuntimeError( + "PP communicator release has no matching acquisition") + _pp_comm_refcount -= 1 + if _pp_comm_refcount == 0: + _pp_comm_final_release_pending = True + while _pp_comm_control_refcount > 0: + _pp_comm_condition.wait() + if not (NCCL_FAULT_TOLERANCE_ENABLED + and isinstance(_pp_comm, PPCommNCCL)): + # Preserve the legacy teardown behavior when FT is disabled. + _pp_comm = None + # In FT mode both the native communicator and its Python recovery + # generation state are process-lifetime. A rank-local final release + # is not a distributed teardown barrier, and dropping either state + # could make only one rank enter a duplicate-generation rendezvous. + _pp_comm_final_release_pending = False + _pp_comm_condition.notify_all() + + +def _get_pp_nccl_comm_locked() -> PPCommNCCL: + """Return the NCCL PP communicator while ``_pp_comm_lock`` is held.""" + if not isinstance(_pp_comm, PPCommNCCL): + raise RuntimeError("NCCL error: PP communicator is not initialized") + return _pp_comm + + +def _acquire_pp_nccl_control( + *, + capture_generation: bool = False) -> tuple[PPCommNCCL, Optional[int]]: + """Keep the shared communicator alive without holding its lifecycle lock.""" + global _pp_comm_control_refcount + with _pp_comm_lock: + comm = _get_pp_nccl_comm_locked() + if _pp_comm_refcount <= 0 or _pp_comm_final_release_pending: + raise RuntimeError("NCCL error: PP communicator is being released") + _pp_comm_control_refcount += 1 + generation = (comm._reconfigure_generation + if capture_generation else None) + return comm, generation + + +def _release_pp_nccl_control() -> None: + global _pp_comm_control_refcount + with _pp_comm_condition: + if _pp_comm_control_refcount <= 0: + raise RuntimeError( + "NCCL error: PP communicator control reference underflow") + _pp_comm_control_refcount -= 1 + if _pp_comm_control_refcount == 0: + _pp_comm_condition.notify_all() + + +def pp_comm_abort() -> None: + """Abort the process-local NCCL PP communicator.""" + comm, generation = _acquire_pp_nccl_control(capture_generation=True) + try: + comm.abort(generation) + finally: + comm = None + _release_pp_nccl_control() + + +def pp_comm_abort_and_reinit(active_ranks: List[int], + generation: Optional[int] = None) -> None: + """Rebuild the NCCL PP communicator using surviving original world ranks.""" + comm, _ = _acquire_pp_nccl_control() + try: + comm.abort_and_reinit(active_ranks, generation) + finally: + comm = None + _release_pp_nccl_control() + + +def pp_comm_get_async_error() -> str: + """Return the NCCL PP communicator's latched abort/error reason.""" + comm, _ = _acquire_pp_nccl_control() + try: + return comm.get_async_error() + finally: + comm = None + _release_pp_nccl_control() + + +def pp_comm_get_active_ranks() -> List[int]: + """Return original world-rank IDs in the NCCL PP communicator.""" + comm, _ = _acquire_pp_nccl_control() + try: + return comm.get_active_ranks() + finally: + comm = None + _release_pp_nccl_control() @TorchDist.log_op diff --git a/tensorrt_llm/_torch/distributed/nccl_fault_tolerance.py b/tensorrt_llm/_torch/distributed/nccl_fault_tolerance.py new file mode 100644 index 000000000000..e2e2dc7e46d3 --- /dev/null +++ b/tensorrt_llm/_torch/distributed/nccl_fault_tolerance.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Process-local survivor membership for raw NCCL communicators. + +The C++ raw-NCCL registry keys communicators by their global-rank set. After +``nccl_comm_abort_and_reinit`` replaces a communicator, explicitly dynamic +callers must use the replacement set, while statically sharded TP/CP callers +must stop until their missing state is redistributed. This module keeps that +membership and fail-fast decision in one place instead of attaching it to the +MoE communication object that initiated recovery. + +ProcessGroup-backed collectives intentionally do not use this registry: their +groups need to be reconstructed by the ProcessGroup owner. +""" + +from __future__ import annotations + +import os +import threading +from typing import Callable, Iterable, List, Tuple + +_Group = Tuple[int, ...] +_MAX_RECOVERY_GENERATION = (1 << 63) - 3 + +# Fault tolerance is a startup-only mode. Cache the flag once so the default +# collective hot path pays only a cheap boolean branch and never consults the +# survivor registry. +NCCL_FAULT_TOLERANCE_ENABLED = os.environ.get("TLLM_FAULT_TOLERANCE_MODE") == "1" + +_registry_lock = threading.RLock() +_active_groups: dict[_Group, _Group] = {} +_completed_reconfigurations: dict[_Group, tuple[int, _Group]] = {} + + +def _canonical_group(group: Iterable[int], name: str) -> _Group: + ranks = tuple(int(rank) for rank in group) + if not ranks: + raise ValueError(f"{name} must not be empty") + if len(ranks) != len(set(ranks)): + raise ValueError(f"{name} must not contain duplicate ranks") + return ranks + + +def _canonical_recovery_generation(generation: int | None) -> int | None: + """Validate a coordinator generation accepted by the Torch ``int`` ABI.""" + if generation is None: + return None + canonical_generation = int(generation) + if ( + canonical_generation < 0 + or canonical_generation != generation + or canonical_generation > _MAX_RECOVERY_GENERATION + ): + raise ValueError( + f"generation must be a nonnegative integer no greater than {_MAX_RECOVERY_GENERATION}" + ) + return canonical_generation + + +def _recovery_rendezvous_id(generation: int | None) -> int: + """Map a recovery generation to its wire ID. + + ID 0 is reserved for initial bootstrap and ID 1 for the legacy one-shot + membership shrink without a coordinator generation. Explicit generations + start at ID 2, keeping every namespace distinct on the retained MPI control + communicator. + """ + return 1 if generation is None else generation + 2 + + +def resolve_nccl_group(group: Iterable[int]) -> List[int]: + """Return the latest survivor set for a raw NCCL rank group.""" + if not NCCL_FAULT_TOLERANCE_ENABLED: + return list(group) + ranks = _canonical_group(group, "group") + # Keep the hot-path lookup free of synchronization primitives so Dynamo + # can guard the dictionary value and recompile when recovery publishes a + # new group. Recovery quiesces collective submission before publishing. + return list(_active_groups.get(ranks, ranks)) + + +def resolve_nccl_group_and_rank(group: Iterable[int], world_rank: int) -> tuple[List[int], int]: + """Resolve a raw NCCL group and the caller's compact communicator rank.""" + active_group = resolve_nccl_group(group) + try: + compact_rank = active_group.index(int(world_rank)) + except ValueError as error: + raise RuntimeError( + f"NCCL error: world rank {world_rank} is not in active communicator {active_group}" + ) from error + return active_group, compact_rank + + +def is_nccl_group_reconfigured(group: Iterable[int]) -> bool: + """Cheap hot-path membership-change check for a trusted rank group.""" + if not NCCL_FAULT_TOLERANCE_ENABLED: + return False + original_group = group if isinstance(group, tuple) else tuple(group) + active_group = _active_groups.get(original_group) + return active_group is not None and active_group != original_group + + +def assert_nccl_group_not_reconfigured(group: Iterable[int], operation: str) -> None: + """Reject static rank sharding after transport-only recovery. + + This hot-path check trusts Mapping-owned groups to already be canonical. A + single dictionary lookup is enough before recovery; full validation stays + in the control-path APIs that publish membership. + """ + if not NCCL_FAULT_TOLERANCE_ENABLED: + return + original_group = group if isinstance(group, tuple) else tuple(group) + active_group = _active_groups.get(original_group) + if active_group is not None and active_group != original_group: + raise RuntimeError( + "NCCL error: " + f"{operation} cannot use survivor-only communicator {list(active_group)} " + f"for statically sharded group {list(original_group)}; redistribute the " + "missing rank's state and rebuild the mapping before resuming" + ) + + +def publish_nccl_group_reconfiguration( + original_group: Iterable[int], + current_group: Iterable[int], + active_group: Iterable[int], +) -> None: + """Publish a successful native raw-NCCL communicator replacement. + + ``current_group`` makes stale or racing reconfiguration attempts fail + before they can overwrite newer membership. Publishing happens only + after the coordinated native rebuild succeeds. + """ + original = _canonical_group(original_group, "original_group") + current = _canonical_group(current_group, "current_group") + active = _canonical_group(active_group, "active_group") + + if not set(active).issubset(current): + raise ValueError("active_group must be a subset of current_group") + + with _registry_lock: + registered = _active_groups.get(original, original) + if registered == active: + return + if registered != current: + raise RuntimeError( + "NCCL error: stale communicator membership update: " + f"expected {list(registered)}, got {list(current)}" + ) + + # Preserve aliases for the model's original mapping and for any + # survivor-only mapping views already handed to communication helpers. + aliases = [key for key, value in _active_groups.items() if value == current] + aliases.extend((original, current, active)) + for alias in aliases: + _active_groups[alias] = active + + +def reconfigure_nccl_group( + original_group: Iterable[int], + active_group: Iterable[int], + rebuild: Callable[[List[int], List[int], int], None], + generation: int | None = None, +) -> List[int]: + """Serialize a native communicator rebuild and membership publication. + + Recovery is a process-wide safe-point operation. Holding the registry + lock across ``rebuild`` prevents two local layers or worker threads from + rebuilding the same native communicator to different survivor sets before + either update is published. ``generation`` is the coordinator's monotonic + recovery-event generation, advanced for every distinct attempt (including + transport-only retries). It is required to distinguish a same-membership + transport rebuild from a duplicate callback without consulting rank-local + watchdog state, which can differ transiently across survivors. A first + membership shrink may omit it for compatibility, but every retry after a + native failure must provide a newly advanced generation so the wire + rendezvous cannot pair different attempts. + """ + if not NCCL_FAULT_TOLERANCE_ENABLED: + raise RuntimeError( + "NCCL error: communicator reinitialization requires TLLM_FAULT_TOLERANCE_MODE=1" + ) + + original = _canonical_group(original_group, "original_group") + requested = _canonical_group(active_group, "active_group") + generation = _canonical_recovery_generation(generation) + rendezvous_id = _recovery_rendezvous_id(generation) + + with _registry_lock: + current = _active_groups.get(original, original) + completed = _completed_reconfigurations.get(original) + if generation is not None and completed is not None: + completed_generation, completed_target = completed + if generation < completed_generation: + return list(current) + if generation == completed_generation: + if requested != completed_target: + raise RuntimeError( + "NCCL error: conflicting communicator recovery target " + f"for generation {generation}: completed " + f"{list(completed_target)}, requested {list(requested)}" + ) + return list(current) + + if not set(requested).issubset(current): + raise ValueError("active_group must be a subset of the current_group") + if requested == current: + if generation is None: + raise ValueError( + "generation is required for same-membership communicator " + "recovery so every survivor makes the same rebuild decision" + ) + + rebuild(list(current), list(requested), rendezvous_id) + + aliases = [key for key, value in _active_groups.items() if value == current] + aliases.extend((original, current, requested)) + for alias in aliases: + _active_groups[alias] = requested + if generation is not None: + _completed_reconfigurations[original] = (generation, requested) + return list(requested) + + +def _reset_nccl_group_registry_for_tests() -> None: + """Clear process-local membership. Tests only; native state is untouched.""" + with _registry_lock: + _active_groups.clear() + _completed_reconfigurations.clear() diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 63bf9d76c2c0..5e46309b701e 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import math import os import platform @@ -8,6 +11,9 @@ from torch import nn from tensorrt_llm._mnnvl_utils import HelixCpMnnvlMemory, MnnvlMemory +from tensorrt_llm._torch.distributed.nccl_fault_tolerance import ( + NCCL_FAULT_TOLERANCE_ENABLED, assert_nccl_group_not_reconfigured, + is_nccl_group_reconfigured, resolve_nccl_group) from tensorrt_llm._torch.distributed.symm_mem_allreduce import \ SymmetricMemoryAllReduce from tensorrt_llm._torch.utils import get_model_extra_attrs @@ -160,6 +166,10 @@ def get_or_scale_allreduce_mnnvl_workspace( def userbuffers_allreduce_finalize( input: torch.Tensor, force_applying_finalize: bool = False) -> torch.Tensor: + if NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled(): + raise RuntimeError( + "NCCL error: userbuffers allreduce has no peer-failure abort path " + "and is unavailable while TLLM_FAULT_TOLERANCE_MODE=1") output = torch.ops.trtllm.userbuffers_allreduce_finalize( input, force_applying_finalize) return output @@ -307,9 +317,17 @@ def allgather( Returns: The gathered tensor or tensor list. ''' - group_boxed = mapping.tp_group_pg.boxed() if mpi_disabled() else None - return _allgather(input, mapping.tp_group, mapping.tp_rank, group_boxed, - dim, sizes) + if mpi_disabled(): + group = mapping.tp_group + rank = mapping.tp_rank + group_boxed = mapping.tp_group_pg.boxed() + else: + group = mapping.tp_group + if NCCL_FAULT_TOLERANCE_ENABLED: + assert_nccl_group_not_reconfigured(group, "allgather") + rank = mapping.tp_rank + group_boxed = None + return _allgather(input, group, rank, group_boxed, dim, sizes) def cp_allgather( @@ -331,9 +349,17 @@ def cp_allgather( Returns: The gathered tensor or tensor list. ''' - group_boxed = mapping.cp_group_pg.boxed() if mpi_disabled() else None - return _allgather(input, mapping.cp_group, mapping.cp_rank, group_boxed, - dim, sizes) + if mpi_disabled(): + group = mapping.cp_group + rank = mapping.cp_rank + group_boxed = mapping.cp_group_pg.boxed() + else: + group = mapping.cp_group + if NCCL_FAULT_TOLERANCE_ENABLED: + assert_nccl_group_not_reconfigured(group, "CP allgather") + rank = mapping.cp_rank + group_boxed = None + return _allgather(input, group, rank, group_boxed, dim, sizes) def alltoall_helix( @@ -357,6 +383,16 @@ def alltoall_helix( For each group of input tensors (of size group size), there is one output tensor with shape (group size, *input shape). ''' + if NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled(): + original_group = list(group) + group = resolve_nccl_group(original_group) + if group != original_group: + raise RuntimeError( + "NCCL error: Helix alltoall cannot use survivor-only " + f"communicator {group} for statically sharded CP group " + f"{original_group}; redistribute the missing rank's state " + "and rebuild the mapping before resuming") + n_ranks = len(group) if n_ranks == 1: return inputs @@ -459,11 +495,20 @@ def reducescatter( dim: int = -1, sizes: Optional[List[int]] = None, ) -> Union[torch.Tensor, List[torch.Tensor]]: - if mapping.tp_size == 1: + if mpi_disabled(): + group = mapping.tp_group + group_size = mapping.tp_size + else: + group = mapping.tp_group + if NCCL_FAULT_TOLERANCE_ENABLED: + assert_nccl_group_not_reconfigured(group, "reducescatter") + group_size = len(group) + + if group_size == 1: return input if sizes is not None: - assert len(sizes) == len(mapping.tp_group) + assert len(sizes) == group_size sum_split_size = sum(sizes) if isinstance(input, torch.Tensor): assert input.shape[dim] == sum_split_size @@ -479,7 +524,7 @@ def convert_input(x, x_info): x = x.contiguous().view(-1, x_info['numel_base']) else: if sizes is None: - x_list = x.chunk(mapping.tp_size, dim=dim) + x_list = x.chunk(group_size, dim=dim) else: x_list = x.split(sizes, dim=dim) x = torch.cat([x.reshape(-1, x_info['numel_base']) for x in x_list]) @@ -505,10 +550,9 @@ def convert_input(x, x_info): ] if mpi_disabled(): - output = torch_op(input, sizes, mapping.tp_group, - mapping.tp_group_pg.boxed()) + output = torch_op(input, sizes, group, mapping.tp_group_pg.boxed()) else: - output = torch_op(input, sizes, mapping.tp_group) + output = torch_op(input, sizes, group) if isinstance(input, torch.Tensor): output = output.view(output_info['output_shape']) @@ -701,11 +745,22 @@ def __init__(self, """ self.mapping = mapping + self._tp_group_tuple = tuple(mapping.tp_group) self.workspace = None self.strategy = strategy self.mnnvl_allreduce = None self.symm_mem_allreduce = None self._disable_mpi = mpi_disabled() + self._fault_tolerant_raw_nccl = (NCCL_FAULT_TOLERANCE_ENABLED + and not self._disable_mpi) + + if (self._fault_tolerant_raw_nccl and self.strategy + not in (AllReduceStrategy.AUTO, AllReduceStrategy.NCCL)): + strategy_name = AllReduceStrategy(self.strategy).name + raise RuntimeError( + "NCCL error: TLLM_FAULT_TOLERANCE_MODE=1 requires a " + "watchdog-managed NCCL allreduce; strategy " + f"{strategy_name} has no peer-failure abort path") self.all_reduce_op = torch.ops.trtllm.allreduce_pg if self._disable_mpi else torch.ops.trtllm.allreduce @@ -729,7 +784,8 @@ def __init__(self, if self.mapping.tp_size > 1 and not self.mapping.enable_attention_dp: # Initialize Symmetric Memory AllReduce if needed (before workspace allocation) - if self.strategy == AllReduceStrategy.SYMM_MEM: + if (not self._fault_tolerant_raw_nccl + and self.strategy == AllReduceStrategy.SYMM_MEM): try: symm_mem = SymmetricMemoryAllReduce( self.mapping, @@ -758,7 +814,8 @@ def __init__(self, # Allocate workspace for strategies that need it # Note: SYMM_MEM now also needs workspace for fallback scenarios (fused ops, etc.) # Only UB doesn't need workspace - if self.strategy != AllReduceStrategy.UB: + if (not self._fault_tolerant_raw_nccl + and self.strategy != AllReduceStrategy.UB): if self.strategy == AllReduceStrategy.LOWPRECISION: allocate_low_presicion_allreduce_workspace(self.mapping) if self.strategy not in (AllReduceStrategy.UB, @@ -767,8 +824,8 @@ def __init__(self, self.workspace = get_allreduce_workspace(self.mapping) # Initialize MNNVL if using AUTO or MNNVL strategy - if self.strategy in (AllReduceStrategy.AUTO, - AllReduceStrategy.MNNVL): + if (not self._fault_tolerant_raw_nccl and self.strategy + in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL)): # Try to initialize MNNVL if MNNVLAllReduce.is_mnnvl(self.mapping, dtype): # ALWAYS capture the exception when creating this instance @@ -791,6 +848,12 @@ def uses_nccl_symmetric_memory_window(self) -> bool: Requires TLLM_NCCL_SYMMETRIC_ZERO_COPY=1 AND NCCL_SYMMETRIC/NCCL/AUTO strategy AND tp_size > 1 AND MPI not disabled. """ + if (NCCL_FAULT_TOLERANCE_ENABLED and not self._disable_mpi + and is_nccl_group_reconfigured(self._tp_group_tuple)): + # Window registration and its allocation workspace belong to the + # original communicator. Forward will reject the incomplete static + # shard; avoid allocating another stale window before that check. + return False return (_NCCL_SYMMETRIC_ZERO_COPY and self.strategy in (AllReduceStrategy.NCCL_SYMMETRIC, AllReduceStrategy.NCCL, AllReduceStrategy.AUTO) and self.mapping.tp_size > 1 @@ -802,8 +865,9 @@ def output_buffer_kind(self) -> int: be passed into this allreduce. Returns int(BufferKind.NCCL_WINDOW) when zero-copy window output is - active, int(BufferKind.DEFAULT) otherwise. The value depends solely on - compile-time constants so it is safe to branch on inside torch.compile. + active, int(BufferKind.DEFAULT) otherwise. Dynamo guards the survivor + registry lookup and recompiles after membership is published at the + recovery safe point, so this remains safe inside ``torch.compile``. """ return (int(BufferKind.NCCL_WINDOW) if self.uses_nccl_symmetric_memory_window() else int( @@ -844,23 +908,36 @@ def forward( == False): return input + active_group = self.mapping.tp_group + if NCCL_FAULT_TOLERANCE_ENABLED and not self._disable_mpi: + assert_nccl_group_not_reconfigured(self._tp_group_tuple, + "allreduce") + input = input.contiguous() # Underlying op requires contiguous input allreduce_strategy = self.strategy + if (self._fault_tolerant_raw_nccl + and allreduce_strategy == AllReduceStrategy.AUTO): + # AUTO includes one-shot/two-shot and other custom transports with + # no abort hook. NCCL_SYMMETRIC also performs MPI-local topology + # discovery, whose communicator may still contain the dead rank. + allreduce_strategy = AllReduceStrategy.NCCL if all_reduce_params is None: all_reduce_params = AllReduceParams() # Try Symmetric Memory AllReduce first if available # Note: Currently only supports NONE fusion op (plain allreduce) - if self.symm_mem_allreduce and all_reduce_params.fusion_op == AllReduceFusionOp.NONE: + if (self.symm_mem_allreduce + and all_reduce_params.fusion_op == AllReduceFusionOp.NONE): symm_mem_output = self.symm_mem_allreduce(input) if symm_mem_output is not None: logger.debug( f"Using SymmetricMemoryAllReduce (MULTIMEM) for input shape {input.shape}" ) return symm_mem_output - elif self.symm_mem_allreduce and all_reduce_params.fusion_op != AllReduceFusionOp.NONE: + elif (self.symm_mem_allreduce + and all_reduce_params.fusion_op != AllReduceFusionOp.NONE): # Log once per rank that we're skipping symm_mem due to fusion logger.debug_once( f"Skipping SymmetricMemoryAllReduce for fused operation (fusion_op={all_reduce_params.fusion_op}), using regular allreduce", @@ -880,7 +957,6 @@ def forward( if allreduce_strategy in (AllReduceStrategy.MNNVL, AllReduceStrategy.SYMM_MEM): allreduce_strategy = AllReduceStrategy.AUTO - additional_args = {} if self._disable_mpi: # Get ProcessGroup from mapping @@ -896,7 +972,8 @@ def forward( disable_allreduce_autotune = os.environ.get( "TLLM_DISABLE_ALLREDUCE_AUTOTUNE", "0") == "1" - if allreduce_strategy == AllReduceStrategy.AUTO and not disable_allreduce_autotune and not self._disable_mpi: + if (allreduce_strategy == AllReduceStrategy.AUTO + and not disable_allreduce_autotune and not self._disable_mpi): output = torch.ops.trtllm.tunable_allreduce( input=input, residual=all_reduce_params.residual, @@ -904,7 +981,7 @@ def forward( scale=all_reduce_params.scale, bias=all_reduce_params.bias, workspace=self.workspace, - group=self.mapping.tp_group, + group=active_group, strategy=allreduce_strategy, op=all_reduce_params.fusion_op, eps=all_reduce_params.eps, @@ -919,7 +996,7 @@ def forward( scale=all_reduce_params.scale, bias=all_reduce_params.bias, workspace=self.workspace, - group=self.mapping.tp_group, + group=active_group, strategy=allreduce_strategy, op=all_reduce_params.fusion_op, eps=all_reduce_params.eps, @@ -961,6 +1038,11 @@ def __init__(self, mapping: Mapping): output_hidden_states = rms_norm(output_residual, norm_weight, eps) """ super().__init__() + if NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled(): + raise RuntimeError( + "NCCL error: MoEAllReduce uses a custom collective without a " + "peer-failure abort path and is unavailable while " + "TLLM_FAULT_TOLERANCE_MODE=1") self.mapping = mapping self.workspace = get_allreduce_workspace(self.mapping) # Pls keep this value in sync with the kOneShotMaxToken in moeAllReduceFusionKernels.h @@ -1224,6 +1306,11 @@ class MiniMaxAllReduceRMS(nn.Module): def __init__(self, mapping: Mapping): super().__init__() + if NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled(): + raise RuntimeError( + "NCCL error: MiniMaxAllReduceRMS uses a custom collective " + "without a peer-failure abort path and is unavailable while " + "TLLM_FAULT_TOLERANCE_MODE=1") self.mapping = mapping self.workspace = get_allreduce_workspace(self.mapping) diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py b/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py index 706de4247a84..f5484865cc9b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -12,7 +12,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - """ AllGather + ReduceScatter Communication Strategy @@ -27,11 +26,47 @@ import torch from tensorrt_llm._torch.distributed import allgather, reducescatter +from tensorrt_llm._torch.distributed.nccl_fault_tolerance import ( + NCCL_FAULT_TOLERANCE_ENABLED, + reconfigure_nccl_group, + resolve_nccl_group, +) +from tensorrt_llm._utils import mpi_disabled from tensorrt_llm.mapping import Mapping from .base import Communication +class _ActiveGroupMapping: + """Read-only mapping view for a survivor-only NCCL communicator. + + The model-wide :class:`Mapping` remains immutable because TP/PP users may + still be inspecting it while the MoE fallback switches membership. + Attributes unrelated to the TP communicator delegate to the original + mapping. + """ + + def __init__(self, mapping: Mapping, group: List[int], rank: int) -> None: + self._mapping = mapping + self._group = tuple(group) + self._rank = rank + + @property + def tp_group(self) -> List[int]: + return list(self._group) + + @property + def tp_size(self) -> int: + return len(self._group) + + @property + def tp_rank(self) -> int: + return self._rank + + def __getattr__(self, name): + return getattr(self._mapping, name) + + class AllGatherReduceScatter(Communication): def __init__( self, @@ -41,6 +76,11 @@ def __init__( # Initialize dispatch state self._dispatch_state = {} + self._original_group = tuple(mapping.tp_group) + self._active_local_ranks = tuple(range(len(self._original_group))) + self._active_global_group = self._original_group + self._active_mapping = mapping + self._raw_nccl_fault_tolerance = NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled() @staticmethod def is_platform_supported() -> bool: @@ -57,6 +97,113 @@ def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) """ return True + def abort_and_reinit(self, active_ranks: List[int], generation: Optional[int] = None) -> None: + """Abort the current NCCL communicator and rebuild it for survivors. + + ``active_ranks`` contains global/world rank IDs from the original TP + communication group. A global-rank API is intentional: when MoE TP is + enabled, one TP communicator spans several MoE-EP slices, so an + EP-local rank number is ambiguous without the failure coordinator's + slice identity. The coordinator must translate its health snapshot to + global ranks before calling this method. + + Reconfiguration is monotonic: a removed rank cannot be reintroduced by + this Phase-1 API. Replacement-rank joins belong to process-group + reconstruction. The failure coordinator must pass a shared monotonic + recovery-event generation for same-membership recovery and advance it + for every distinct attempt, including transport-only retries. Passing + it for every recovery also deduplicates callbacks on every survivor. + """ + if not NCCL_FAULT_TOLERANCE_ENABLED: + raise RuntimeError( + "NCCL error: communicator reinitialization requires TLLM_FAULT_TOLERANCE_MODE=1" + ) + if not self._raw_nccl_fault_tolerance: + raise NotImplementedError( + "AllGatherReduceScatter fault tolerance requires the raw NCCL/MPI path; " + "ProcessGroup reconstruction is not implemented" + ) + + if not active_ranks: + raise ValueError("active_ranks must not be empty") + if len(active_ranks) != len(set(active_ranks)): + raise ValueError("active_ranks must not contain duplicates") + requested = {int(rank) for rank in active_ranks} + if not requested.issubset(self._original_group): + raise ValueError( + "active_ranks must be global ranks from the original TP group " + f"{list(self._original_group)}, got {active_ranks}" + ) + if self.mapping.rank not in requested: + raise ValueError(f"current world rank {self.mapping.rank} is not active") + + active_global_group = tuple(rank for rank in self._original_group if rank in requested) + + # Commit Python metadata only after the coordinated native rebuild + # succeeds. The old transport is terminal once native recovery starts, + # even if construction of its replacement fails. + reconfigure_nccl_group( + self._original_group, + active_global_group, + torch.ops.trtllm.nccl_comm_abort_and_reinit, + generation, + ) + self._refresh_active_mapping() + + def _refresh_active_mapping(self) -> None: + """Synchronize this layer with process-wide survivor membership.""" + active_global_group = tuple(resolve_nccl_group(self._original_group)) + if active_global_group == self._active_global_group: + return + global_rank = self.mapping.rank + if global_rank not in active_global_group: + raise RuntimeError( + f"NCCL error: world rank {global_rank} is not in the active communicator" + ) + active_global_ranks = set(active_global_group) + active_local_ranks = tuple( + index + for index, world_rank in enumerate(self._original_group) + if world_rank in active_global_ranks + ) + active_mapping = _ActiveGroupMapping( + self.mapping, + list(active_global_group), + active_global_group.index(global_rank), + ) + self._active_local_ranks = active_local_ranks + self._active_mapping = active_mapping + # Publish the value used by the fast-path guard last. A concurrent + # reader that observes this generation must also observe its mapping. + self._active_global_group = active_global_group + + def check_async_error(self) -> None: + """Raise a classifier-friendly exception for a watchdog failure.""" + if not NCCL_FAULT_TOLERANCE_ENABLED: + raise RuntimeError( + "NCCL error: async-error monitoring requires TLLM_FAULT_TOLERANCE_MODE=1" + ) + if not self._raw_nccl_fault_tolerance: + raise NotImplementedError( + "AllGatherReduceScatter fault tolerance cannot inspect ProcessGroup errors; " + "the raw NCCL/MPI watchdog is not active" + ) + self._refresh_active_mapping() + error = torch.ops.trtllm.nccl_comm_get_async_error(list(self._active_global_group)) + if error is not None: + raise RuntimeError(f"NCCL error: communicator was aborted: {error}") + + def _active_sizes(self, all_rank_num_tokens: List[int]) -> List[int]: + self._refresh_active_mapping() + if len(all_rank_num_tokens) != len(self._original_group): + raise ValueError( + "all_rank_num_tokens must remain indexed by the original communicator " + f"({len(self._original_group)} entries), got {len(all_rank_num_tokens)}" + ) + if self._active_global_group == self._original_group: + return all_rank_num_tokens + return [all_rank_num_tokens[rank] for rank in self._active_local_ranks] + def dispatch( self, hidden_states: torch.Tensor, @@ -70,11 +217,20 @@ def dispatch( """ AllGather dispatch (always post-quant dispatch) """ - sizes = None if use_dp_padding else all_rank_num_tokens + if self._raw_nccl_fault_tolerance: + if use_dp_padding: + self._refresh_active_mapping() + sizes = None + else: + sizes = self._active_sizes(all_rank_num_tokens) + mapping = self._active_mapping + else: + sizes = None if use_dp_padding else all_rank_num_tokens + mapping = self.mapping hidden_states, hidden_states_sf, token_selected_slots, token_final_scales = allgather( [hidden_states, hidden_states_sf, token_selected_slots, token_final_scales], - self.mapping, + mapping, dim=0, sizes=sizes, ) @@ -92,7 +248,15 @@ def combine( """ ReduceScatter combine phase """ + if self._raw_nccl_fault_tolerance: + self._refresh_active_mapping() + mapping = self._active_mapping + else: + mapping = self.mapping outputs = reducescatter( - final_hidden_states, self.mapping, dim=0, sizes=self._dispatch_state.get("sizes") + final_hidden_states, + mapping, + dim=0, + sizes=self._dispatch_state.get("sizes"), ) return outputs diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 9ff6f28efb9f..a106fcfd2ada 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations import enum @@ -14,6 +17,8 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.custom_ops.torch_custom_ops import BufferKind +from tensorrt_llm._torch.distributed.nccl_fault_tolerance import ( + NCCL_FAULT_TOLERANCE_ENABLED, assert_nccl_group_not_reconfigured) from tensorrt_llm._torch.peft.lora.layer import LoraLayer from tensorrt_llm._utils import is_device_integrated, mpi_disabled from tensorrt_llm.bindings import ipc_nvls_supported @@ -3188,6 +3193,7 @@ def __init__( ) self.tp_size = self.mapping.tp_size self.tp_rank = self.mapping.tp_rank + self._tp_group_tuple = tuple(self.mapping.tp_group) self.tp_mode = tensor_parallel_mode self.gather_output = gather_output self.force_dynamic_quantization = force_dynamic_quantization @@ -3250,10 +3256,17 @@ def __init__( "TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED", "0") == "1") self.use_fused_gemm_allreduce = all([ - self.reduce_output, mpi_enabled, dtype_supported, - in_features_aligned, out_features_aligned, tp_valid, quant_valid, - device_supported, enable_gemm_allreduce_fusion, - enable_gemm_allreduce_fusion_env + self.reduce_output, + mpi_enabled, + dtype_supported, + in_features_aligned, + out_features_aligned, + tp_valid, + quant_valid, + device_supported, + enable_gemm_allreduce_fusion, + enable_gemm_allreduce_fusion_env, + not (NCCL_FAULT_TOLERANCE_ENABLED and mpi_enabled), ]) if self.use_fused_gemm_allreduce: self.use_fused_gemm_allreduce = ipc_nvls_supported() @@ -3368,8 +3381,14 @@ def apply_linear_allreduce(self, input, bias, layer_idx: Optional[int] | None = None): + tp_group = self.mapping.tp_group + tp_rank = self.tp_rank + if NCCL_FAULT_TOLERANCE_ENABLED and not mpi_disabled(): + raise RuntimeError( + "NCCL error: fused GEMM allreduce has no peer-failure abort " + "path and is unavailable while TLLM_FAULT_TOLERANCE_MODE=1") output = self.quant_method.apply_linear_allreduce( - self, input, bias, self.tp_rank, self.mapping.tp_group) + self, input, bias, tp_rank, tp_group) return output def _maybe_fuse_bias_into_allreduce( @@ -3397,12 +3416,19 @@ def forward( lora_params: Optional[dict] = None, layer_idx: Optional[int] = None, ) -> torch.Tensor: + if (NCCL_FAULT_TOLERANCE_ENABLED and self.tp_mode + in (TensorParallelMode.ROW, TensorParallelMode.COLUMN) + and self.tp_size > 1 and not mpi_disabled()): + assert_nccl_group_not_reconfigured( + self._tp_group_tuple, + f"{self.tp_mode.name.lower()}-parallel linear") + if self.tp_mode == TensorParallelMode.ROW: use_fused_gemm_allreduce = self.use_fused_gemm_allreduce and lora_params is None if use_fused_gemm_allreduce and all_reduce_params is not None: use_fused_gemm_allreduce = all_reduce_params.enable_allreduce and all_reduce_params.fusion_op == AllReduceFusionOp.NONE - bias = None if (self.tp_rank > 0) else self.bias + bias = None if self.tp_rank > 0 else self.bias if self.reduce_output: if use_fused_gemm_allreduce: output = self.apply_linear_allreduce( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index c9cf887562c7..ebc8fc46acbd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import bisect import contextlib import functools @@ -44,7 +47,7 @@ from ..compilation.backend import Backend from ..compilation.utils import capture_piecewise_cuda_graph from ..distributed import Distributed -from ..distributed.communicator import init_pp_comm +from ..distributed.communicator import init_pp_comm, release_pp_comm from ..expert_statistic import ExpertStatistic from ..memory_buffer_utils import clear_memory_buffers, with_shared_pool from ..metadata import KVCacheParams @@ -108,6 +111,40 @@ def warmup(self, resource_manager: ResourceManager) -> None: return +def _force_eager_mode_for_nccl_fault_tolerance( + llm_args: TorchLlmArgs) -> TorchLlmArgs: + """Disable CUDA capture until communicator-aware recapture lands. + + A captured NCCL graph node retains the communicator handle used during + capture. PR 1a.7 can replace that handle after a rank failure, but PR + 1a.11 owns graph-cache invalidation and recapture. Keep the MVP usable by + running both whole-model and piecewise paths eagerly while FT mode is on. + Work on the engine-local Pydantic copy so callers can reuse their original + arguments for a non-FT engine. + """ + if os.environ.get("TLLM_FAULT_TOLERANCE_MODE") != "1": + return llm_args + + updates = {} + if llm_args.cuda_graph_config is not None: + updates["cuda_graph_config"] = None + + torch_compile_config = llm_args.torch_compile_config + if (torch_compile_config is not None + and torch_compile_config.enable_piecewise_cuda_graph): + updates["torch_compile_config"] = torch_compile_config.model_copy( + update={"enable_piecewise_cuda_graph": False}) + + if not updates: + return llm_args + + logger.warning( + "TLLM_FAULT_TOLERANCE_MODE=1 disables CUDA graph capture until " + "communicator-aware invalidation and recapture are available; this " + "engine will serve in eager mode.") + return llm_args.model_copy(update=updates) + + def _filter_piecewise_capture_num_tokens( candidate_num_tokens: list[int], max_num_tokens: int, @@ -240,8 +277,10 @@ def __init__( model_weights_memory_tag: Optional[str] = None, model_weights_restore_mode=None, ): + llm_args = _force_eager_mode_for_nccl_fault_tolerance(llm_args) self.forward_pass_callable = None self.ub_buffers = None + self._pp_comm_acquired = False if llm_args.encode_only and llm_args.mm_encoder_only: raise ValueError( "encode_only and mm_encoder_only are mutually exclusive.") @@ -269,6 +308,7 @@ def __init__( self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) + self._pp_comm_acquired = True self.dist = dist if dist is not None: ExpertStatistic.create(self.dist.rank) @@ -1911,7 +1951,9 @@ def cleanup(self) -> None: 5. Userbuffers (``ub.ub_deallocate`` per buffer); on per-buffer failure the unfreed buffers are kept attached so a deterministic retry doesn't double-free already-released ones, and the - collected errors are re-raised after the loop. + collected errors are re-raised after unrelated global ownership is + released. + 6. This engine's reference to the process-local PP communicator. Idempotency: Subsequent calls are no-ops (guarded by ``_cleanup_done``). @@ -1948,9 +1990,9 @@ def cleanup(self) -> None: self.input_processor_with_hash = None ub_buffers = getattr(self, 'ub_buffers', None) + ub_errors = [] if ub_buffers: remaining_ub_buffers = [] - ub_errors = [] for u in ub_buffers: try: ub.ub_deallocate(u.addr) @@ -1961,13 +2003,21 @@ def cleanup(self) -> None: remaining_ub_buffers.append(u) ub_errors.append(e) self.ub_buffers = remaining_ub_buffers or None - if ub_errors: - raise RuntimeError( - "Failed to deallocate one or more userbuffers during " - "PyTorchModelEngine cleanup") from ub_errors[0] + + if getattr(self, "_pp_comm_acquired", False): + release_pp_comm() + self._pp_comm_acquired = False # Release model weights. release_gc() + if ub_errors: + # Terminal teardown may not get another deterministic retry. Defer + # the userbuffer error until unrelated process-global ownership is + # released, while retaining only the failed buffers for a caller + # that can retry cleanup. + raise RuntimeError( + "Failed to deallocate one or more userbuffers during " + "PyTorchModelEngine cleanup") from ub_errors[0] self._cleanup_done = True def __del__(self) -> None: diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index b8a0823591ca..b4d7dc3f1dc2 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -46,6 +46,8 @@ l0_a10: - unittest/_torch/modules/dwdp/test_dwdp_manager.py - unittest/_torch/modules/dwdp/test_dwdp_mapping.py - unittest/_torch/modules/dwdp/test_dwdp_peer_ranges.py + - unittest/_torch/modules/moe/test_allgather_reducescatter_ft.py + - unittest/_torch/modules/test_pp_nccl_communicator.py # NOTE: this is a CPU-only test, but we do not have a dedicated job for this (and therefore no # test list either). - unittest/_torch/models/checkpoints/hf/test_weight_loader.py diff --git a/tests/unittest/_torch/modules/moe/test_allgather_reducescatter_ft.py b/tests/unittest/_torch/modules/moe/test_allgather_reducescatter_ft.py new file mode 100644 index 000000000000..8555d542850f --- /dev/null +++ b/tests/unittest/_torch/modules/moe/test_allgather_reducescatter_ft.py @@ -0,0 +1,861 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the NCCL fault-tolerance hooks on AllGatherReduceScatter. + +These tests intentionally mock the C++ custom op. They verify the Python-side +contract (global-rank validation, survivor publication, and dispatch-size +bookkeeping) without requiring NCCL, MPI, or GPUs. +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass, field +from typing import Callable + +import pytest + +from tensorrt_llm._torch.custom_ops import torch_custom_ops +from tensorrt_llm._torch.distributed import nccl_fault_tolerance +from tensorrt_llm._torch.distributed import ops as distributed_ops +from tensorrt_llm._torch.modules import linear as linear_module +from tensorrt_llm._torch.modules.fused_moe.communication import ( + allgather_reducescatter as allgather_rs_module, +) +from tensorrt_llm._torch.modules.fused_moe.communication.allgather_reducescatter import ( + AllGatherReduceScatter, +) +from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode + + +@pytest.fixture(autouse=True) +def _enable_and_clear_nccl_survivor_membership(monkeypatch): + monkeypatch.setattr(nccl_fault_tolerance, "NCCL_FAULT_TOLERANCE_ENABLED", True) + monkeypatch.setattr(distributed_ops, "NCCL_FAULT_TOLERANCE_ENABLED", True) + monkeypatch.setattr(torch_custom_ops, "NCCL_FAULT_TOLERANCE_ENABLED", True) + monkeypatch.setattr(allgather_rs_module, "NCCL_FAULT_TOLERANCE_ENABLED", True) + monkeypatch.setattr(linear_module, "NCCL_FAULT_TOLERANCE_ENABLED", True) + nccl_fault_tolerance._reset_nccl_group_registry_for_tests() + yield + nccl_fault_tolerance._reset_nccl_group_registry_for_tests() + + +@dataclass +class _FakeMapping: + """A PP slice whose EP-local ranks map to non-zero global ranks.""" + + world_size: int = 8 + rank: int = 6 + tp_size: int = 4 + tp_rank: int = 2 + moe_ep_size: int = 4 + moe_ep_rank: int = 2 + enable_attention_dp: bool = False + _group: list[int] = field(default_factory=lambda: [4, 5, 6, 7]) + _moe_group: list[int] | None = None + + @property + def tp_group(self) -> list[int]: + return list(self._group) + + @property + def moe_ep_group(self) -> list[int]: + return list(self._group if self._moe_group is None else self._moe_group) + + +def _patch_reinit_op( + monkeypatch: pytest.MonkeyPatch, + implementation: Callable[[list[int], list[int], int], None], +) -> None: + monkeypatch.setattr( + allgather_rs_module.torch.ops.trtllm, + "nccl_comm_abort_and_reinit", + implementation, + raising=False, + ) + + +def _patch_async_error_op( + monkeypatch: pytest.MonkeyPatch, + implementation: Callable[[list[int]], str | None], +) -> None: + monkeypatch.setattr( + allgather_rs_module.torch.ops.trtllm, + "nccl_comm_get_async_error", + implementation, + raising=False, + ) + + +def _assert_active_mapping(mapping, expected_group: list[int], expected_rank: int) -> None: + """Assert the distributed op sees compact ranks for the rebuilt comm.""" + assert list(mapping.tp_group) == expected_group + assert mapping.tp_size == len(expected_group) + assert mapping.tp_rank == expected_rank + + +def test_abort_and_reinit_canonicalizes_global_ranks(monkeypatch): + calls: list[tuple[list[int], list[int], int]] = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append( + (list(group), list(active_group), rendezvous_id) + ), + ) + comm = AllGatherReduceScatter(_FakeMapping()) + + # Input order must not affect NCCL rank assignment: the C++ op receives the + # original TP-group order in the global/world-rank domain. + comm.abort_and_reinit([7, 4, 6]) + + assert calls == [([4, 5, 6, 7], [4, 6, 7], 1)] + + +def test_reconfiguration_rejects_static_sharded_allgather(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + observed = [] + + def fake_allgather(input, group, rank, group_boxed, dim, sizes): + observed.append((input, group, rank, group_boxed, dim, sizes)) + return "gathered" + + monkeypatch.setattr(distributed_ops, "_allgather", fake_allgather) + + mapping = _FakeMapping() + comm = AllGatherReduceScatter(mapping) + comm.abort_and_reinit([4, 6, 7]) + + # Transport recovery alone cannot recreate a TP shard. A normal model-wide + # call using the original Mapping must fail instead of returning fewer + # output features. AllGatherReduceScatter uses its explicit survivor mapping below. + with pytest.raises(RuntimeError, match="statically sharded|redistribute"): + distributed_ops.allgather("input", mapping, dim=0, sizes=[5, 7, 11, 13]) + assert observed == [] + + +def test_reconfiguration_rejects_static_sharded_helix_alltoall(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + AllGatherReduceScatter(_FakeMapping()).abort_and_reinit([4, 6, 7]) + calls = [] + + def fake_alltoall(inputs, group, num_lists): + calls.append((list(inputs), list(group), num_lists)) + return ["output"] + + monkeypatch.setattr( + distributed_ops.torch.ops.trtllm, + "alltoall_helix", + fake_alltoall, + raising=False, + ) + + inputs = [allgather_rs_module.torch.tensor([rank]) for rank in range(4)] + with pytest.raises(RuntimeError, match="statically sharded|redistribute"): + distributed_ops.alltoall_helix(inputs, [4, 5, 6, 7]) + assert calls == [] + + +def test_reconfiguration_is_shared_and_idempotent_across_moe_layers(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + first_layer = AllGatherReduceScatter(_FakeMapping()) + second_layer = AllGatherReduceScatter(_FakeMapping()) + + first_layer.abort_and_reinit([4, 6, 7], generation=1) + second_layer.abort_and_reinit([4, 6, 7], generation=1) + + assert calls == [([4, 5, 6, 7], [4, 6, 7])] + _assert_active_mapping(second_layer._active_mapping, [4, 6, 7], 1) + + +def test_reconfiguration_transaction_serializes_same_target(monkeypatch): + entered_native = threading.Event() + release_native = threading.Event() + calls = [] + + def native_rebuild(group, active_group, rendezvous_id): + calls.append((list(group), list(active_group))) + entered_native.set() + assert release_native.wait(timeout=5) + + _patch_reinit_op(monkeypatch, native_rebuild) + first_layer = AllGatherReduceScatter(_FakeMapping()) + second_layer = AllGatherReduceScatter(_FakeMapping()) + errors = [] + + def reconfigure(layer): + try: + layer.abort_and_reinit([4, 6, 7], generation=1) + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + first = threading.Thread(target=reconfigure, args=(first_layer,)) + second = threading.Thread(target=reconfigure, args=(second_layer,)) + first.start() + assert entered_native.wait(timeout=5) + second.start() + release_native.set() + first.join(timeout=5) + second.join(timeout=5) + + assert not first.is_alive() + assert not second.is_alive() + assert errors == [] + assert calls == [([4, 5, 6, 7], [4, 6, 7])] + + +def test_shared_generation_deduplicates_and_advances_same_membership_rebuild(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append( + (list(group), list(active_group), rendezvous_id) + ), + ) + + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 5, 6, 7], generation=1) + comm.abort_and_reinit([4, 5, 6, 7], generation=1) + comm.abort_and_reinit([4, 5, 6, 7], generation=2) + + assert calls == [ + ([4, 5, 6, 7], [4, 5, 6, 7], 3), + ([4, 5, 6, 7], [4, 5, 6, 7], 4), + ] + + +def test_same_membership_recovery_requires_shared_generation(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + + comm = AllGatherReduceScatter(_FakeMapping()) + with pytest.raises(ValueError, match="generation is required"): + comm.abort_and_reinit([4, 5, 6, 7]) + + assert calls == [] + + +@pytest.mark.parametrize("generation", [-1, 1.5, (1 << 63) - 2]) +def test_raw_recovery_generation_must_fit_reserved_torch_int_range(monkeypatch, generation): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append( + (group, active_group, rendezvous_id) + ), + ) + + with pytest.raises(ValueError, match="generation must be a nonnegative integer"): + AllGatherReduceScatter(_FakeMapping()).abort_and_reinit([4, 6, 7], generation=generation) + + assert calls == [] + + +def test_generation_rejects_conflicting_raw_recovery_targets(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 6, 7], generation=9) + with pytest.raises(RuntimeError, match="conflicting.*generation 9"): + comm.abort_and_reinit([6, 7], generation=9) + + assert calls == [([4, 5, 6, 7], [4, 6, 7])] + + +def test_fault_tolerance_control_path_requires_startup_mode(monkeypatch): + monkeypatch.setattr(allgather_rs_module, "NCCL_FAULT_TOLERANCE_ENABLED", False) + calls = [] + _patch_reinit_op( + monkeypatch, lambda group, active_group, rendezvous_id: calls.append((group, active_group)) + ) + + comm = AllGatherReduceScatter(_FakeMapping()) + with pytest.raises(RuntimeError, match="TLLM_FAULT_TOLERANCE_MODE=1"): + comm.abort_and_reinit([4, 6, 7]) + + assert calls == [] + + +def test_default_off_collective_skips_survivor_registry(monkeypatch): + monkeypatch.setattr(distributed_ops, "NCCL_FAULT_TOLERANCE_ENABLED", False) + monkeypatch.setattr(distributed_ops, "mpi_disabled", lambda: False) + + def unexpected_registry_lookup(group, operation): + raise AssertionError("default-off allgather consulted survivor state") + + observed = [] + monkeypatch.setattr( + distributed_ops, "assert_nccl_group_not_reconfigured", unexpected_registry_lookup + ) + monkeypatch.setattr( + distributed_ops, + "_allgather", + lambda input, group, rank, group_boxed, dim, sizes: observed.append( + (group, rank, group_boxed, dim, sizes) + ) + or "gathered", + ) + + class ReferenceMapping: + tp_group = [4, 5, 6, 7] + tp_rank = 2 + + mapping = ReferenceMapping() + result = distributed_ops.allgather("input", mapping, dim=0, sizes=[1, 2, 3, 4]) + + assert result == "gathered" + assert observed == [([4, 5, 6, 7], 2, None, 0, [1, 2, 3, 4])] + assert observed[0][0] is mapping.tp_group + + +def test_ft_auto_allreduce_uses_only_watchdog_managed_nccl(monkeypatch): + monkeypatch.setattr(distributed_ops, "mpi_disabled", lambda: False) + calls = [] + + def fake_allreduce(**kwargs): + calls.append(kwargs) + return [kwargs["input"]] + + monkeypatch.setattr( + distributed_ops.torch.ops.trtllm, "allreduce", fake_allreduce, raising=False + ) + monkeypatch.setattr( + distributed_ops, + "get_allreduce_workspace", + lambda mapping: (_ for _ in ()).throw( + AssertionError("FT AUTO allocated a custom-allreduce workspace") + ), + ) + monkeypatch.setattr( + distributed_ops.MNNVLAllReduce, + "is_mnnvl", + lambda mapping, dtype: (_ for _ in ()).throw( + AssertionError("FT AUTO probed the MNNVL allreduce backend") + ), + ) + + allreduce = distributed_ops.AllReduce( + _FakeMapping(), + strategy=distributed_ops.AllReduceStrategy.AUTO, + dtype=distributed_ops.torch.float16, + ) + input_tensor = distributed_ops.torch.ones(2) + output = allreduce(input_tensor) + + assert output is input_tensor + assert allreduce.workspace is None + assert allreduce.mnnvl_allreduce is None + assert len(calls) == 1 + assert calls[0]["strategy"] == distributed_ops.AllReduceStrategy.NCCL + + +@pytest.mark.parametrize( + "strategy", + [ + distributed_ops.AllReduceStrategy.UB, + distributed_ops.AllReduceStrategy.ONESHOT, + distributed_ops.AllReduceStrategy.MNNVL, + distributed_ops.AllReduceStrategy.NCCL_SYMMETRIC, + distributed_ops.AllReduceStrategy.SYMM_MEM, + ], +) +def test_ft_allreduce_rejects_custom_transport_strategies(monkeypatch, strategy): + monkeypatch.setattr(distributed_ops, "mpi_disabled", lambda: False) + with pytest.raises(RuntimeError, match="watchdog-managed NCCL allreduce"): + distributed_ops.AllReduce(_FakeMapping(), strategy=strategy) + + +def test_ft_tunable_allreduce_exposes_only_nccl_tactics(): + runner = torch_custom_ops.AllReduceRunner( + 4, + [4, 5, 6, 7], + int(distributed_ops.AllReduceFusionOp.NONE), + 1e-6, + False, + ) + inputs = [distributed_ops.torch.ones(2)] + assert runner.get_valid_tactics(inputs, torch_custom_ops.OptimizationProfile()) == [ + distributed_ops.AllReduceStrategy.NCCL.value, + ] + + +def test_ft_custom_allreduce_helpers_fail_before_workspace_allocation(monkeypatch): + monkeypatch.setattr(distributed_ops, "mpi_disabled", lambda: False) + monkeypatch.setattr( + distributed_ops, + "get_allreduce_workspace", + lambda mapping: (_ for _ in ()).throw( + AssertionError("unsupported FT helper allocated a workspace") + ), + ) + + with pytest.raises(RuntimeError, match="MoEAllReduce.*custom collective"): + distributed_ops.MoEAllReduce(_FakeMapping()) + with pytest.raises(RuntimeError, match="MiniMaxAllReduceRMS.*custom collective"): + distributed_ops.MiniMaxAllReduceRMS(_FakeMapping()) + monkeypatch.setattr( + distributed_ops.torch.ops.trtllm, + "userbuffers_allreduce_finalize", + lambda *args: (_ for _ in ()).throw( + AssertionError("unsupported FT userbuffers op was invoked") + ), + raising=False, + ) + with pytest.raises(RuntimeError, match="userbuffers allreduce"): + distributed_ops.userbuffers_allreduce_finalize(distributed_ops.torch.ones(1)) + + +def test_ft_disables_fused_nvfp4_gemm_allreduce(monkeypatch): + class QuantMode: + @staticmethod + def has_nvfp4(): + return True + + class QuantConfig: + layer_quant_mode = QuantMode() + + monkeypatch.setattr(linear_module, "mpi_disabled", lambda: False) + monkeypatch.setattr(linear_module, "get_sm_version", lambda: 100) + monkeypatch.setenv("TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED", "1") + nvls_queries = [] + monkeypatch.setattr( + linear_module, "ipc_nvls_supported", lambda: nvls_queries.append(True) or True + ) + + linear = Linear( + 128, + 64, + bias=False, + dtype=distributed_ops.torch.float16, + mapping=_FakeMapping(), + tensor_parallel_mode=TensorParallelMode.ROW, + quant_config=QuantConfig(), + skip_create_weights_in_init=True, + enable_gemm_allreduce_fusion=True, + ) + + assert not linear.use_fused_gemm_allreduce + assert nvls_queries == [] + with pytest.raises(RuntimeError, match="fused GEMM allreduce"): + linear.apply_linear_allreduce(None, None) + + +def test_active_mapping_and_sizes_are_reused_until_membership_changes(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + comm = AllGatherReduceScatter(_FakeMapping()) + original_mapping = comm._active_mapping + original_sizes = [5, 7, 11, 13] + + assert comm._active_sizes(original_sizes) is original_sizes + assert comm._active_mapping is original_mapping + + comm.abort_and_reinit([4, 6, 7]) + survivor_mapping = comm._active_mapping + assert comm._active_sizes(original_sizes) == [5, 11, 13] + assert comm._active_mapping is survivor_mapping + comm._refresh_active_mapping() + assert comm._active_mapping is survivor_mapping + + +def test_process_group_backend_skips_raw_nccl_membership_refresh(monkeypatch): + monkeypatch.setattr(allgather_rs_module, "mpi_disabled", lambda: True) + comm = AllGatherReduceScatter(_FakeMapping()) + monkeypatch.setattr( + allgather_rs_module, + "resolve_nccl_group", + lambda group: (_ for _ in ()).throw( + AssertionError("ProcessGroup dispatch consulted raw NCCL state") + ), + ) + observed = [] + monkeypatch.setattr( + allgather_rs_module, + "allgather", + lambda inputs, mapping, *, dim, sizes: observed.append((mapping, sizes)) or inputs, + ) + sizes = [5, 7, 11, 13] + + comm.dispatch("hidden", None, "slots", None, sizes) + + assert observed == [(comm.mapping, sizes)] + + +def test_reconfiguration_rejects_static_sharded_allreduce(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + mapping = _FakeMapping() + comm = AllGatherReduceScatter(mapping) + comm.abort_and_reinit([4, 6, 7]) + + calls = [] + + def fake_allreduce(**kwargs): + calls.append(kwargs) + return ("reduced",) + + allreduce = object.__new__(distributed_ops.AllReduce) + allreduce.mapping = mapping + allreduce._tp_group_tuple = tuple(mapping.tp_group) + allreduce.strategy = distributed_ops.AllReduceStrategy.AUTO + allreduce.workspace = None + allreduce.symm_mem_allreduce = None + allreduce.mnnvl_allreduce = None + allreduce._disable_mpi = False + allreduce.all_reduce_op = fake_allreduce + + tensor = allgather_rs_module.torch.ones(1) + with pytest.raises(RuntimeError, match="statically sharded|redistribute"): + allreduce.forward(tensor) + assert calls == [] + assert not allreduce.uses_nccl_symmetric_memory_window() + + +@pytest.mark.parametrize( + "active_ranks", + [ + pytest.param([], id="empty"), + pytest.param([4, 6, 6], id="duplicate"), + pytest.param([-1, 6], id="negative"), + pytest.param([4, 6, 8], id="outside_tp_group"), + pytest.param([4, 5, 7], id="world_rank_missing"), + ], +) +def test_abort_and_reinit_rejects_invalid_active_ranks(monkeypatch, active_ranks): + calls = [] + _patch_reinit_op( + monkeypatch, lambda group, active_group, rendezvous_id: calls.append((group, active_group)) + ) + comm = AllGatherReduceScatter(_FakeMapping()) + + with pytest.raises(ValueError): + comm.abort_and_reinit(active_ranks) + + assert calls == [] + + +def test_abort_and_reinit_supports_mixed_tp_and_ep_topology(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, lambda group, active_group, rendezvous_id: calls.append((group, active_group)) + ) + mapping = _FakeMapping( + rank=6, + tp_rank=2, + moe_ep_size=2, + moe_ep_rank=1, + _group=[4, 5, 6, 7], + _moe_group=[4, 6], + ) + comm = AllGatherReduceScatter(mapping) + + # Global ranks are unambiguous even when the TP communicator spans several + # MoE-EP slices. The higher-level health coordinator performs its local to + # global translation before entering this API. + comm.abort_and_reinit([4, 6, 7]) + + assert calls == [([4, 5, 6, 7], [4, 6, 7])] + + +def test_abort_and_reinit_rejects_process_group_backend(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, lambda group, active_group, rendezvous_id: calls.append((group, active_group)) + ) + monkeypatch.setattr(allgather_rs_module, "mpi_disabled", lambda: True, raising=False) + comm = AllGatherReduceScatter(_FakeMapping()) + + # The raw communicator cache is not used in Ray/ProcessGroup mode, and the + # active mapping would otherwise delegate to the stale tp_group_pg. + with pytest.raises(NotImplementedError, match="ProcessGroup|MPI-disabled"): + comm.abort_and_reinit([4, 6, 7]) + + assert calls == [] + + +def test_check_async_error_returns_cleanly_when_watchdog_is_healthy(monkeypatch): + queried_groups = [] + _patch_async_error_op( + monkeypatch, + lambda group: queried_groups.append(list(group)), + ) + comm = AllGatherReduceScatter(_FakeMapping()) + + comm.check_async_error() + + assert queried_groups == [[4, 5, 6, 7]] + + +def test_check_async_error_rejects_process_group_backend(monkeypatch): + queried_groups = [] + _patch_async_error_op( + monkeypatch, + lambda group: queried_groups.append(list(group)), + ) + monkeypatch.setattr(allgather_rs_module, "mpi_disabled", lambda: True, raising=False) + comm = AllGatherReduceScatter(_FakeMapping()) + + with pytest.raises(NotImplementedError, match="ProcessGroup|watchdog"): + comm.check_async_error() + + assert queried_groups == [] + + +def test_check_async_error_raises_stable_nccl_error_for_active_group(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + queried_groups = [] + + def get_async_error(group): + queried_groups.append(list(group)) + return "ncclSystemError: peer disconnected" + + _patch_async_error_op(monkeypatch, get_async_error) + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 6, 7]) + + with pytest.raises( + RuntimeError, + match="NCCL error: communicator was aborted: ncclSystemError: peer disconnected", + ): + comm.check_async_error() + + assert queried_groups == [[4, 6, 7]] + + +def test_other_layer_queries_shared_active_group_after_recovery(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + queried_groups = [] + _patch_async_error_op(monkeypatch, lambda group: queried_groups.append(list(group))) + first_layer = AllGatherReduceScatter(_FakeMapping()) + second_layer = AllGatherReduceScatter(_FakeMapping()) + + first_layer.abort_and_reinit([4, 6, 7]) + second_layer.check_async_error() + + assert queried_groups == [[4, 6, 7]] + _assert_active_mapping(second_layer._active_mapping, [4, 6, 7], 1) + + +def test_reconfigured_row_linear_fails_before_using_incomplete_tp_shards(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + AllGatherReduceScatter(_FakeMapping()).abort_and_reinit([6, 7]) + calls = [] + + class QuantMethod: + supports_nccl_symmetric_memory_window_output = False + + class FakeRowLinear: + tp_mode = TensorParallelMode.ROW + mapping = _FakeMapping(rank=6, tp_rank=2) + _tp_group_tuple = tuple(mapping.tp_group) + tp_rank = 2 + tp_size = 4 + use_fused_gemm_allreduce = True + reduce_output = True + bias = "bias" + quant_method = QuantMethod() + lora = None + + def apply_linear_allreduce(self, *args, **kwargs): + raise AssertionError("stale fused GEMM-allreduce path was used") + + def _maybe_fuse_bias_into_allreduce(self, bias, all_reduce_params): + return False + + def apply_linear(self, input, bias, lora_params, layer_idx): + calls.append(("gemm", input, bias)) + return "gemm-output" + + def all_reduce(self, output, *, all_reduce_params): + calls.append(("allreduce", output)) + return "reduced" + + with pytest.raises(RuntimeError, match="statically sharded|redistribute"): + Linear.forward(FakeRowLinear(), "input") + assert calls == [] + + +def test_reconfigured_column_linear_without_gather_fails_before_gemm(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + AllGatherReduceScatter(_FakeMapping()).abort_and_reinit([6, 7]) + calls = [] + + class FakeColumnLinear: + tp_mode = TensorParallelMode.COLUMN + mapping = _FakeMapping(rank=6, tp_rank=2) + _tp_group_tuple = tuple(mapping.tp_group) + tp_size = 4 + gather_output = False + bias = None + + def apply_linear(self, *args, **kwargs): + calls.append((args, kwargs)) + return "partial-output" + + with pytest.raises(RuntimeError, match="statically sharded|redistribute"): + Linear.forward(FakeColumnLinear(), "input") + assert calls == [] + + +def test_dispatch_and_combine_filter_sizes_to_active_ranks(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + gather_calls = [] + scatter_calls = [] + + def fake_allgather(inputs, mapping, *, dim, sizes): + gather_calls.append((inputs, mapping, dim, sizes)) + return inputs + + def fake_reducescatter(input_tensor, mapping, *, dim, sizes): + scatter_calls.append((input_tensor, mapping, dim, sizes)) + return "combined" + + monkeypatch.setattr(allgather_rs_module, "allgather", fake_allgather) + monkeypatch.setattr(allgather_rs_module, "reducescatter", fake_reducescatter) + + mapping = _FakeMapping() + comm = AllGatherReduceScatter(mapping) + comm.abort_and_reinit([4, 6, 7]) + + inputs = ["hidden", None, "slots", "scales"] + assert comm.dispatch(*inputs, all_rank_num_tokens=[5, 7, 11, 13]) == tuple(inputs) + assert comm.combine("moe-output") == "combined" + + assert len(gather_calls) == 1 + gathered_inputs, gather_mapping, gather_dim, gather_sizes = gather_calls[0] + assert gathered_inputs == inputs + assert gather_dim == 0 + assert gather_sizes == [5, 11, 13] + _assert_active_mapping(gather_mapping, [4, 6, 7], 1) + + assert len(scatter_calls) == 1 + scatter_input, scatter_mapping, scatter_dim, scatter_sizes = scatter_calls[0] + assert scatter_input == "moe-output" + assert scatter_dim == 0 + assert scatter_sizes == [5, 11, 13] + _assert_active_mapping(scatter_mapping, [4, 6, 7], 1) + + # Reconfiguration does not mutate Mapping, but the shared raw-NCCL + # membership registry reroutes other collectives that consume it. + assert mapping.tp_group == [4, 5, 6, 7] + assert mapping.tp_rank == 2 + + +def test_dp_padding_keeps_sizes_none_after_reconfiguration(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + observed = [] + + def fake_allgather(inputs, mapping, *, dim, sizes): + observed.append((mapping, sizes)) + return inputs + + monkeypatch.setattr(allgather_rs_module, "allgather", fake_allgather) + + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 6, 7]) + comm.dispatch("hidden", None, "slots", None, [5, 7, 11, 13], use_dp_padding=True) + + assert len(observed) == 1 + active_mapping, sizes = observed[0] + assert sizes is None + _assert_active_mapping(active_mapping, [4, 6, 7], 1) + + +def test_dp_padding_refreshes_mapping_after_another_layer_recovers(monkeypatch): + _patch_reinit_op(monkeypatch, lambda group, active_group, rendezvous_id: None) + observed = [] + + def fake_allgather(inputs, mapping, *, dim, sizes): + observed.append((mapping, sizes)) + return inputs + + monkeypatch.setattr(allgather_rs_module, "allgather", fake_allgather) + recovering_layer = AllGatherReduceScatter(_FakeMapping()) + stale_layer = AllGatherReduceScatter(_FakeMapping()) + + recovering_layer.abort_and_reinit([4, 6, 7]) + stale_layer.dispatch("hidden", None, "slots", None, [5, 7, 11, 13], use_dp_padding=True) + + assert len(observed) == 1 + active_mapping, sizes = observed[0] + assert sizes is None + _assert_active_mapping(active_mapping, [4, 6, 7], 1) + + +def test_later_reconfiguration_replaces_active_rank_filter(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + observed = [] + + def fake_allgather(inputs, mapping, *, dim, sizes): + observed.append((mapping, sizes)) + return inputs + + monkeypatch.setattr(allgather_rs_module, "allgather", fake_allgather) + + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 6, 7]) + comm.abort_and_reinit([6, 7]) + comm.dispatch("hidden", None, "slots", None, [5, 7, 11, 13]) + + assert calls == [ + ([4, 5, 6, 7], [4, 6, 7]), + ([4, 6, 7], [6, 7]), + ] + assert len(observed) == 1 + active_mapping, sizes = observed[0] + assert sizes == [11, 13] + _assert_active_mapping(active_mapping, [6, 7], 0) + + +def test_reconfiguration_cannot_reactivate_a_removed_rank(monkeypatch): + calls = [] + _patch_reinit_op( + monkeypatch, + lambda group, active_group, rendezvous_id: calls.append((list(group), list(active_group))), + ) + + comm = AllGatherReduceScatter(_FakeMapping()) + comm.abort_and_reinit([4, 6, 7]) + + with pytest.raises(ValueError, match="reactivate|subset"): + comm.abort_and_reinit([4, 5, 6, 7]) + + assert calls == [([4, 5, 6, 7], [4, 6, 7])] + + +def test_failed_reinit_does_not_commit_active_rank_filter(monkeypatch): + def fail_reinit(group, active_group, rendezvous_id): + raise RuntimeError("NCCL communicator rebuild failed") + + _patch_reinit_op(monkeypatch, fail_reinit) + comm = AllGatherReduceScatter(_FakeMapping()) + old_local_ranks = comm._active_local_ranks + old_global_group = comm._active_global_group + old_mapping = comm._active_mapping + + with pytest.raises(RuntimeError, match="NCCL communicator rebuild failed"): + comm.abort_and_reinit([4, 6, 7]) + + # The native call aborts the old communicator before initializing its + # replacement, so transport is terminal after this failure. The only + # transactional guarantee is that Python membership metadata does not + # claim the failed replacement was committed. + assert comm._active_local_ranks is old_local_ranks + assert comm._active_global_group is old_global_group + assert comm._active_mapping is old_mapping diff --git a/tests/unittest/_torch/modules/test_pp_nccl_communicator.py b/tests/unittest/_torch/modules/test_pp_nccl_communicator.py new file mode 100644 index 000000000000..71c9006a8cf8 --- /dev/null +++ b/tests/unittest/_torch/modules/test_pp_nccl_communicator.py @@ -0,0 +1,782 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import threading + +import pytest + +from tensorrt_llm import _mnnvl_utils as mnnvl_module +from tensorrt_llm._torch.distributed import communicator as communicator_module + + +@pytest.fixture(autouse=True) +def _reset_pp_comm_lifecycle(monkeypatch): + previous_helix_comm = mnnvl_module.HelixCpMnnvlMemory.comm + previous_helix_topology = mnnvl_module.HelixCpMnnvlMemory.comm_topology + communicator_module._pp_comm = None + communicator_module._pp_comm_refcount = 0 + communicator_module._pp_comm_control_refcount = 0 + communicator_module._pp_comm_final_release_pending = False + monkeypatch.setattr(communicator_module, "NCCL_FAULT_TOLERANCE_ENABLED", True) + mnnvl_module.HelixCpMnnvlMemory.comm = None + mnnvl_module.HelixCpMnnvlMemory.comm_topology = None + yield + communicator_module._pp_comm = None + communicator_module._pp_comm_refcount = 0 + communicator_module._pp_comm_control_refcount = 0 + communicator_module._pp_comm_final_release_pending = False + mnnvl_module.HelixCpMnnvlMemory.comm = previous_helix_comm + mnnvl_module.HelixCpMnnvlMemory.comm_topology = previous_helix_topology + + +class _FakeNcclCommunicatorOp: + def __init__(self): + self.aborted = False + self.reinitialized_with = None + self.reinit_calls = [] + self.rendezvous_ids = [] + self.active_ranks = [0, 1, 2, 3] + self.async_error = "NCCL error: communicator was aborted after injected async failure" + self.received_from = None + + def abort(self): + self.aborted = True + + def abort_and_reinit(self, active_ranks, rendezvous_id): + self.reinitialized_with = list(active_ranks) + self.reinit_calls.append(list(active_ranks)) + self.rendezvous_ids.append(rendezvous_id) + self.active_ranks = list(active_ranks) + self.async_error = "" + + def get_async_error(self): + return self.async_error + + def get_active_ranks(self): + return list(self.active_ranks) + + def recv(self, tensor, src): + self.received_from = src + + def send(self, tensor, dest): + self.sent_to = dest + + +class _FakeMapping: + world_size = 4 + rank = 2 + pp_group = [0, 1, 2, 3] + cp_config = {} + + def has_cp_helix(self): + return False + + def prev_pp_rank(self): + return 1 + + def next_pp_rank(self): + return 3 + + +class _CompatibleMapping(_FakeMapping): + pass + + +class _HelixMapping(_FakeMapping): + cp_config = {"use_nccl_for_alltoall": False} + cp_group = [0, 2] + cp_rank = 1 + pp_rank = 0 + tp_size = 2 + tp_rank = 0 + + def has_cp_helix(self): + return True + + +class _OtherHelixMapping(_HelixMapping): + cp_group = [1, 2] + + +def _make_pp_comm(mapping=None): + pp_comm = object.__new__(communicator_module.PPCommNCCL) + pp_comm.mapping = mapping or _FakeMapping() + pp_comm._topology = pp_comm._mapping_topology(pp_comm.mapping) + pp_comm._active_ranks = tuple(range(pp_comm.mapping.world_size)) + pp_comm._reconfigure_lock = threading.Lock() + pp_comm._reconfigure_generation = 0 + pp_comm._completed_recovery = None + pp_comm.nccl_comm = _FakeNcclCommunicatorOp() + return pp_comm + + +def test_pp_wrapper_initializes_membership_from_persistent_native_comm(monkeypatch): + native_comm = _FakeNcclCommunicatorOp() + native_comm.active_ranks = [0, 2, 3] + monkeypatch.setattr( + communicator_module.torch.classes.trtllm, + "NcclCommunicatorOp", + lambda world_size, rank: native_comm, + raising=False, + ) + monkeypatch.setattr(communicator_module.torch.cuda, "Event", object) + monkeypatch.setattr(communicator_module.torch.cuda, "Stream", object) + + pp_comm = communicator_module.PPCommNCCL(_FakeMapping()) + + assert pp_comm.get_active_ranks() == [0, 2, 3] + + +def test_default_off_wrapper_does_not_query_native_membership(monkeypatch): + native_comm = _FakeNcclCommunicatorOp() + native_comm.get_active_ranks = lambda: (_ for _ in ()).throw( + AssertionError("default-off PP constructor queried native membership") + ) + monkeypatch.setattr(communicator_module, "NCCL_FAULT_TOLERANCE_ENABLED", False) + monkeypatch.setattr( + communicator_module.torch.classes.trtllm, + "NcclCommunicatorOp", + lambda world_size, rank: native_comm, + raising=False, + ) + monkeypatch.setattr(communicator_module.torch.cuda, "Event", object) + monkeypatch.setattr(communicator_module.torch.cuda, "Stream", object) + + pp_comm = communicator_module.PPCommNCCL(_FakeMapping()) + + assert pp_comm._active_ranks == (0, 1, 2, 3) + + +def test_pp_nccl_fault_tolerance_methods_forward_to_custom_class(): + pp_comm = _make_pp_comm() + + assert "NCCL error" in pp_comm.get_async_error() + pp_comm.abort_and_reinit((0, 2, 3)) + assert pp_comm.nccl_comm.reinitialized_with == [0, 2, 3] + assert pp_comm.nccl_comm.rendezvous_ids == [1] + assert pp_comm.get_async_error() == "" + assert pp_comm.get_active_ranks() == [0, 2, 3] + + pp_comm.abort() + assert pp_comm.nccl_comm.aborted + + +def test_module_level_pp_nccl_fault_tolerance_api(monkeypatch): + pp_comm = _make_pp_comm() + monkeypatch.setattr(communicator_module, "_pp_comm", pp_comm) + monkeypatch.setattr(communicator_module, "_pp_comm_refcount", 1) + + assert "NCCL error" in communicator_module.pp_comm_get_async_error() + communicator_module.pp_comm_abort_and_reinit([0, 2, 3], generation=3) + assert pp_comm.nccl_comm.reinitialized_with == [0, 2, 3] + assert pp_comm.nccl_comm.rendezvous_ids == [5] + assert communicator_module.pp_comm_get_async_error() == "" + assert communicator_module.pp_comm_get_active_ranks() == [0, 2, 3] + + communicator_module.pp_comm_abort() + assert pp_comm.nccl_comm.aborted + + +def test_module_level_pp_nccl_api_requires_nccl_backend(monkeypatch): + monkeypatch.setattr(communicator_module, "_pp_comm", object()) + + with pytest.raises(RuntimeError, match="NCCL error: PP communicator is not initialized"): + communicator_module.pp_comm_get_async_error() + + +def test_reinit_rejects_stale_default_peer_and_accepts_explicit_remap(): + pp_comm = _make_pp_comm() + pp_comm.abort_and_reinit([0, 2, 3]) + + # Static Mapping.prev_pp_rank() names failed rank 1. Silently skipping it + # would bypass that stage's model layers, so default routing must fail. + with pytest.raises(RuntimeError, match="peer world rank 1 is not active"): + pp_comm.recv(object()) + + # A higher-level layer/topology reconstruction may provide an explicit + # replacement peer after it has reassigned the failed stage. + pp_comm.recv(object(), src=0) + assert pp_comm.nccl_comm.received_from == 0 + assert pp_comm._required_pp_peer(next_peer=True) == 3 + + +def test_explicit_peer_must_belong_to_the_callers_pp_lane(): + class LaneMapping: + world_size = 8 + rank = 2 + pp_group = [2, 6] + cp_config = {} + + def has_cp_helix(self): + return False + + def prev_pp_rank(self): + return 6 + + def next_pp_rank(self): + return 6 + + pp_comm = _make_pp_comm(LaneMapping()) + + with pytest.raises(RuntimeError, match="not in this rank's PP group"): + pp_comm.recv(object(), src=3) + + pp_comm.recv(object(), src=6) + assert pp_comm.nccl_comm.received_from == 6 + + +def test_default_off_preserves_legacy_peer_routing(monkeypatch): + pp_comm = _make_pp_comm() + monkeypatch.setattr(communicator_module, "NCCL_FAULT_TOLERANCE_ENABLED", False) + monkeypatch.setattr( + communicator_module.torch.cuda, + "is_current_stream_capturing", + lambda: True, + ) + + # Before FT wiring, explicit peers were forwarded to the native wrapper + # without a Python PP-lane scan. Preserve that behavior and cost when the + # startup mode is disabled. + tensor = object() + pp_comm.send(tensor, dest=7) + pp_comm.recv(tensor, src=7) + + assert pp_comm.nccl_comm.sent_to == 7 + assert pp_comm.nccl_comm.received_from == 7 + + +def test_reinit_rejects_reactivating_removed_pp_rank(): + pp_comm = _make_pp_comm() + pp_comm.abort_and_reinit([0, 2, 3]) + + with pytest.raises(ValueError, match="reactivate"): + pp_comm.abort_and_reinit([0, 1, 2, 3]) + + +def test_duplicate_reinit_target_is_idempotent(): + pp_comm = _make_pp_comm() + + pp_comm.abort_and_reinit([0, 2, 3], generation=1) + pp_comm.abort_and_reinit([3, 2, 0], generation=1) + + assert pp_comm.nccl_comm.reinit_calls == [[0, 2, 3]] + assert pp_comm.nccl_comm.rendezvous_ids == [3] + assert pp_comm.get_active_ranks() == [0, 2, 3] + + +def test_shared_generation_forces_same_membership_rebuild(): + pp_comm = _make_pp_comm() + + pp_comm.abort_and_reinit([0, 1, 2, 3], generation=1) + pp_comm.abort_and_reinit([0, 1, 2, 3], generation=1) + pp_comm.abort_and_reinit([0, 1, 2, 3], generation=2) + + assert pp_comm.nccl_comm.reinit_calls == [[0, 1, 2, 3], [0, 1, 2, 3]] + assert pp_comm.nccl_comm.rendezvous_ids == [3, 4] + assert pp_comm.get_async_error() == "" + + +def test_same_membership_recovery_requires_shared_generation(): + pp_comm = _make_pp_comm() + + with pytest.raises(ValueError, match="generation is required"): + pp_comm.abort_and_reinit([0, 1, 2, 3]) + + assert pp_comm.nccl_comm.reinit_calls == [] + + +@pytest.mark.parametrize("generation", [-1, 1.5, (1 << 63) - 2]) +def test_recovery_generation_must_fit_reserved_torch_int_range(generation): + pp_comm = _make_pp_comm() + + with pytest.raises(ValueError, match="generation must be a nonnegative integer"): + pp_comm.abort_and_reinit([0, 2, 3], generation=generation) + + assert pp_comm.nccl_comm.reinit_calls == [] + + +def test_generation_rejects_conflicting_pp_recovery_targets(): + pp_comm = _make_pp_comm() + pp_comm.abort_and_reinit([0, 2, 3], generation=7) + + with pytest.raises(RuntimeError, match="conflicting.*generation 7"): + pp_comm.abort_and_reinit([0, 2], generation=7) + + assert pp_comm.nccl_comm.reinit_calls == [[0, 2, 3]] + + +def test_concurrent_nested_reinit_requests_commit_newest_membership(): + pp_comm = _make_pp_comm() + first_entered_native = threading.Event() + second_entered_native = threading.Event() + release_first = threading.Event() + calls = [] + + def blocking_reinit(active_ranks, rendezvous_id): + active_ranks = list(active_ranks) + calls.append((active_ranks, rendezvous_id)) + pp_comm.nccl_comm.active_ranks = active_ranks + if active_ranks == [0, 2, 3]: + first_entered_native.set() + assert release_first.wait(timeout=5) + else: + second_entered_native.set() + + pp_comm.nccl_comm.abort_and_reinit = blocking_reinit + errors = [] + + def reconfigure(active_ranks): + try: + pp_comm.abort_and_reinit(active_ranks) + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + first = threading.Thread(target=reconfigure, args=([0, 2, 3],)) + second = threading.Thread(target=reconfigure, args=([0, 2],)) + first.start() + assert first_entered_native.wait(timeout=5) + second.start() + + # The second callback must not enter the native rendezvous while the first + # callback is between native completion and Python membership publication. + serialized = not second_entered_native.wait(timeout=0.1) + release_first.set() + first.join(timeout=5) + second.join(timeout=5) + + assert serialized + assert not first.is_alive() + assert not second.is_alive() + assert errors == [] + assert calls == [([0, 2, 3], 1), ([0, 2], 1)] + assert pp_comm.get_active_ranks() == [0, 2] + + +def test_abort_waiting_for_reinit_does_not_abort_replacement(): + pp_comm = _make_pp_comm() + abort_generation = pp_comm._reconfigure_generation + entered_reinit = threading.Event() + release_reinit = threading.Event() + abort_started = threading.Event() + + def blocking_reinit(active_ranks, rendezvous_id): + assert rendezvous_id == 1 + entered_reinit.set() + assert release_reinit.wait(timeout=5) + pp_comm.nccl_comm.active_ranks = list(active_ranks) + pp_comm.nccl_comm.async_error = "" + + pp_comm.nccl_comm.abort_and_reinit = blocking_reinit + errors = [] + + def reconfigure(): + try: + pp_comm.abort_and_reinit([0, 2, 3]) + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + def abort(): + abort_started.set() + try: + pp_comm.abort(abort_generation) + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + rebuild_thread = threading.Thread(target=reconfigure) + abort_thread = threading.Thread(target=abort) + rebuild_thread.start() + assert entered_reinit.wait(timeout=5) + abort_thread.start() + assert abort_started.wait(timeout=5) + release_reinit.set() + rebuild_thread.join(timeout=5) + abort_thread.join(timeout=5) + + assert not rebuild_thread.is_alive() + assert not abort_thread.is_alive() + assert errors == [] + assert not pp_comm.nccl_comm.aborted + assert pp_comm.get_active_ranks() == [0, 2, 3] + + # A genuinely new abort request after the rebuild must still take effect. + pp_comm.abort() + assert pp_comm.nccl_comm.aborted + + +def test_module_abort_waiting_for_reinit_does_not_abort_replacement(monkeypatch): + pp_comm = _make_pp_comm() + entered_reinit = threading.Event() + abort_captured_generation = threading.Event() + release_reinit = threading.Event() + errors = [] + + def blocking_reinit(active_ranks, rendezvous_id): + assert rendezvous_id == 1 + entered_reinit.set() + assert release_reinit.wait(timeout=5) + pp_comm.nccl_comm.active_ranks = list(active_ranks) + pp_comm.nccl_comm.async_error = "" + + pp_comm.nccl_comm.abort_and_reinit = blocking_reinit + original_abort = pp_comm.abort + + def observed_abort(generation=None): + # The module wrapper has already captured the epoch before it invokes + # the class method. Do not let the rebuild finish until that point. + abort_captured_generation.set() + return original_abort(generation) + + pp_comm.abort = observed_abort + monkeypatch.setattr(communicator_module, "_pp_comm", pp_comm) + monkeypatch.setattr(communicator_module, "_pp_comm_refcount", 1) + + def reconfigure(): + try: + communicator_module.pp_comm_abort_and_reinit([0, 2, 3]) + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + def abort(): + try: + communicator_module.pp_comm_abort() + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + rebuild_thread = threading.Thread(target=reconfigure) + abort_thread = threading.Thread(target=abort) + rebuild_thread.start() + assert entered_reinit.wait(timeout=5) + abort_thread.start() + assert abort_captured_generation.wait(timeout=5) + release_reinit.set() + rebuild_thread.join(timeout=5) + abort_thread.join(timeout=5) + + assert not rebuild_thread.is_alive() + assert not abort_thread.is_alive() + assert errors == [] + assert not pp_comm.nccl_comm.aborted + assert pp_comm.get_active_ranks() == [0, 2, 3] + + communicator_module.pp_comm_abort() + assert pp_comm.nccl_comm.aborted + + +def test_module_control_call_blocks_final_release(monkeypatch): + pp_comm = _make_pp_comm() + entered_query = threading.Event() + release_query = threading.Event() + final_release_completed = threading.Event() + errors = [] + + def blocking_get_async_error(): + entered_query.set() + assert release_query.wait(timeout=5) + return "" + + pp_comm.nccl_comm.get_async_error = blocking_get_async_error + monkeypatch.setattr(communicator_module, "_pp_comm", pp_comm) + monkeypatch.setattr(communicator_module, "_pp_comm_refcount", 1) + + def query(): + try: + communicator_module.pp_comm_get_async_error() + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + def final_release(): + try: + communicator_module.release_pp_comm() + final_release_completed.set() + except Exception as error: # pragma: no cover - assertion reports details + errors.append(error) + + query_thread = threading.Thread(target=query) + release_thread = threading.Thread(target=final_release) + query_thread.start() + assert entered_query.wait(timeout=5) + release_thread.start() + + # The control call owns a transient reference under _pp_comm_lock, so a + # final engine release cannot complete its ownership transition yet. + assert not final_release_completed.wait(timeout=0.1) + release_query.set() + query_thread.join(timeout=5) + release_thread.join(timeout=5) + + assert not query_thread.is_alive() + assert not release_thread.is_alive() + assert errors == [] + assert final_release_completed.is_set() + assert communicator_module._pp_comm is pp_comm + assert communicator_module._pp_comm_refcount == 0 + + +def test_init_pp_comm_reuses_compatible_nccl_communicator(monkeypatch): + created = [] + + class FakePPComm: + def __init__(self, mapping): + self.topology = (mapping.world_size, mapping.rank, tuple(mapping.pp_group)) + created.append(self) + + def is_compatible(self, mapping): + return self.topology == (mapping.world_size, mapping.rank, tuple(mapping.pp_group)) + + monkeypatch.setattr(communicator_module, "PPCommNCCL", FakePPComm) + monkeypatch.setattr(communicator_module, "_pp_comm", None) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + monkeypatch.setattr(communicator_module, "init_helix_cp_comm", lambda mapping: None) + + communicator_module.init_pp_comm(_FakeMapping()) + first = communicator_module._pp_comm + # A target and draft engine construct equivalent Mapping objects. They + # must share the sole process-local native PP communicator. + communicator_module.init_pp_comm(_CompatibleMapping()) + + assert communicator_module._pp_comm is first + assert len(created) == 1 + assert communicator_module._pp_comm_refcount == 2 + + communicator_module.release_pp_comm() + assert communicator_module._pp_comm is first + communicator_module.release_pp_comm() + assert communicator_module._pp_comm is first + assert communicator_module._pp_comm_refcount == 0 + + +def test_ft_wrapper_churn_preserves_completed_recovery_generation(monkeypatch): + pp_comm = _make_pp_comm() + monkeypatch.setattr(communicator_module, "_pp_comm", pp_comm) + monkeypatch.setattr(communicator_module, "_pp_comm_refcount", 1) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + monkeypatch.setattr(communicator_module, "init_helix_cp_comm", lambda mapping: None) + + communicator_module.pp_comm_abort_and_reinit([0, 2, 3], generation=7) + communicator_module.release_pp_comm() + assert communicator_module._pp_comm is pp_comm + assert communicator_module._pp_comm_refcount == 0 + + communicator_module.init_pp_comm(_CompatibleMapping()) + communicator_module.pp_comm_abort_and_reinit([0, 2, 3], generation=7) + + assert communicator_module._pp_comm is pp_comm + assert pp_comm.nccl_comm.reinit_calls == [[0, 2, 3]] + assert pp_comm.nccl_comm.rendezvous_ids == [9] + + +def test_init_pp_comm_rejects_incompatible_live_topology(monkeypatch): + existing = _make_pp_comm() + monkeypatch.setattr(communicator_module, "_pp_comm", existing) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + monkeypatch.setattr(communicator_module, "init_helix_cp_comm", lambda mapping: None) + + class OtherMapping(_FakeMapping): + pp_group = [0, 2] + + with pytest.raises(RuntimeError, match="different topology"): + communicator_module.init_pp_comm(OtherMapping()) + + +def test_init_pp_comm_rejects_incompatible_live_helix_topology(monkeypatch): + existing = _make_pp_comm(_HelixMapping()) + monkeypatch.setattr(communicator_module, "_pp_comm", existing) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + + with pytest.raises(RuntimeError, match="Helix MNNVL CP topology"): + communicator_module.init_pp_comm(_OtherHelixMapping()) + + +def test_helix_mnnvl_comm_rejects_incompatible_cached_topology(monkeypatch): + split_calls = [] + comm = object() + + class FakeMpiComm: + def Split(self, color, key): + split_calls.append((color, key)) + return comm + + monkeypatch.setattr(mnnvl_module, "mpi_comm", lambda: FakeMpiComm()) + + assert mnnvl_module.HelixCpMnnvlMemory.get_comm(_HelixMapping()) is comm + assert mnnvl_module.HelixCpMnnvlMemory.get_comm(_HelixMapping()) is comm + with pytest.raises(RuntimeError, match="cannot reuse.*incompatible topology"): + mnnvl_module.HelixCpMnnvlMemory.get_comm(_OtherHelixMapping()) + + assert split_calls == [(0, 1)] + + +def test_final_release_preserves_helix_comm_topology(monkeypatch): + original_pp_comm_class = communicator_module.PPCommNCCL + split_calls = [] + + class FakeMpiComm: + def Split(self, color, key): + split_calls.append((color, key)) + return object() + + class FakePPComm: + def __init__(self, mapping): + self.topology = original_pp_comm_class._mapping_topology(mapping) + + def is_compatible(self, mapping): + return self.topology == original_pp_comm_class._mapping_topology(mapping) + + monkeypatch.setattr(mnnvl_module, "mpi_comm", lambda: FakeMpiComm()) + monkeypatch.setattr(communicator_module, "PPCommNCCL", FakePPComm) + monkeypatch.setattr(communicator_module, "NCCL_FAULT_TOLERANCE_ENABLED", False) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + + # Reacquiring the ordinary PP wrapper after its final release is valid when + # the persistent Helix MPI communicator has the same CP rank topology. + communicator_module.init_pp_comm(_HelixMapping()) + communicator_module.release_pp_comm() + communicator_module.init_pp_comm(_HelixMapping()) + communicator_module.release_pp_comm() + + # A later engine may create a fresh PP wrapper, but it cannot reinterpret + # the process-lifetime Helix MPI communicator with a different CP group. + with pytest.raises(RuntimeError, match="cannot reuse.*incompatible topology"): + communicator_module.init_pp_comm(_OtherHelixMapping()) + + assert split_calls == [(0, 1)] + assert communicator_module._pp_comm is None + assert communicator_module._pp_comm_refcount == 0 + + +def test_default_off_final_release_allows_a_later_incompatible_topology(monkeypatch): + created = [] + + class FakePPComm: + def __init__(self, mapping): + self.topology = (mapping.world_size, mapping.rank, tuple(mapping.pp_group)) + created.append(self.topology) + + def is_compatible(self, mapping): + return self.topology == (mapping.world_size, mapping.rank, tuple(mapping.pp_group)) + + class OtherMapping(_FakeMapping): + pp_group = [0, 2] + + monkeypatch.setattr(communicator_module, "PPCommNCCL", FakePPComm) + monkeypatch.setattr(communicator_module, "NCCL_FAULT_TOLERANCE_ENABLED", False) + monkeypatch.setattr(communicator_module, "mpi_disabled", lambda: False) + monkeypatch.setattr(communicator_module, "init_helix_cp_comm", lambda mapping: None) + + communicator_module.init_pp_comm(_FakeMapping()) + communicator_module.release_pp_comm() + assert communicator_module._pp_comm is None + communicator_module.init_pp_comm(OtherMapping()) + + assert len(created) == 2 + assert created[0] != created[1] + + +def test_model_engine_cleanup_releases_pp_before_reporting_ub_failure(monkeypatch): + from types import SimpleNamespace + + from tensorrt_llm._torch.pyexecutor import model_engine + + engine = object.__new__(model_engine.PyTorchModelEngine) + engine._cleanup_done = False + engine._pp_comm_acquired = True + engine.model_loader = None + engine.model = object() + engine.input_processor = object() + engine.input_processor_with_hash = object() + engine.ub_buffers = [SimpleNamespace(addr=123)] + engine._release_cuda_graphs = lambda: None + + released = [] + gc_calls = [] + monkeypatch.setattr(model_engine, "release_pp_comm", lambda: released.append(True)) + monkeypatch.setattr(model_engine, "release_gc", lambda: gc_calls.append(True)) + monkeypatch.setattr( + model_engine.ub, + "ub_deallocate", + lambda addr: (_ for _ in ()).throw(RuntimeError("injected UB failure")), + ) + + with pytest.raises(RuntimeError, match="userbuffers"): + engine.cleanup() + + assert released == [True] + assert gc_calls == [True] + assert not engine._pp_comm_acquired + assert not engine._cleanup_done + assert [buffer.addr for buffer in engine.ub_buffers] == [123] + + monkeypatch.setattr(model_engine.ub, "ub_deallocate", lambda addr: None) + engine.cleanup() + assert released == [True] + assert engine._cleanup_done + assert engine.ub_buffers is None + + +def test_ft_mode_uses_engine_local_eager_graph_configuration(monkeypatch): + from tensorrt_llm._torch.pyexecutor import model_engine + + class FakeTorchCompileConfig: + def __init__(self, enable_piecewise_cuda_graph): + self.enable_piecewise_cuda_graph = enable_piecewise_cuda_graph + + def model_copy(self, *, update): + copied = FakeTorchCompileConfig(self.enable_piecewise_cuda_graph) + for name, value in update.items(): + setattr(copied, name, value) + return copied + + class FakeLlmArgs: + def __init__(self, cuda_graph_config, torch_compile_config): + self.cuda_graph_config = cuda_graph_config + self.torch_compile_config = torch_compile_config + + def model_copy(self, *, update): + copied = FakeLlmArgs(self.cuda_graph_config, self.torch_compile_config) + for name, value in update.items(): + setattr(copied, name, value) + return copied + + original = FakeLlmArgs(object(), FakeTorchCompileConfig(True)) + monkeypatch.setenv("TLLM_FAULT_TOLERANCE_MODE", "1") + + effective = model_engine._force_eager_mode_for_nccl_fault_tolerance(original) + + assert effective is not original + assert effective.cuda_graph_config is None + assert not effective.torch_compile_config.enable_piecewise_cuda_graph + # The caller-owned arguments remain reusable by a non-FT engine. + assert original.cuda_graph_config is not None + assert original.torch_compile_config.enable_piecewise_cuda_graph + + +def test_ft_mode_uses_eager_allreduce_autotuning(monkeypatch): + from tensorrt_llm._torch.custom_ops import torch_custom_ops + + monkeypatch.delenv("TLLM_FAULT_TOLERANCE_MODE", raising=False) + assert torch_custom_ops._use_cuda_graph_for_allreduce_tuning() + + monkeypatch.setenv("TLLM_FAULT_TOLERANCE_MODE", "1") + assert not torch_custom_ops._use_cuda_graph_for_allreduce_tuning() + + +def test_partially_constructed_engine_cleanup_releases_pp(monkeypatch): + from tensorrt_llm._torch.pyexecutor import model_engine + + # Simulate a constructor failure immediately after init_pp_comm(), before + # model/model_loader/input-processor attributes have been populated. + engine = object.__new__(model_engine.PyTorchModelEngine) + engine._cleanup_done = False + engine._pp_comm_acquired = True + + released = [] + gc_calls = [] + monkeypatch.setattr(model_engine, "release_pp_comm", lambda: released.append(True)) + monkeypatch.setattr(model_engine, "release_gc", lambda: gc_calls.append(True)) + + engine.cleanup() + + assert released == [True] + assert gc_calls == [True] + assert not engine._pp_comm_acquired + assert engine._cleanup_done From 36223bc5d7c0b638040527423ff7d209eb6d1f92 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 10:06:39 -0700 Subject: [PATCH 12/14] [TRTLLM-13553][feat] Add WideEP rank health telemetry Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../modules/fused_moe/ep_group_health.py | 91 ++- .../_torch/modules/fused_moe/ep_metrics.py | 116 +++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 19 + tensorrt_llm/executor/base_worker.py | 11 + tensorrt_llm/executor/executor.py | 4 + tensorrt_llm/executor/proxy.py | 13 + tensorrt_llm/executor/rpc_proxy_mixin.py | 4 + tensorrt_llm/llmapi/llm.py | 4 + tensorrt_llm/metrics/collector.py | 247 +++++- tensorrt_llm/serve/openai_server.py | 204 ++++- .../test_lists/test-db/l0_b200.yml | 2 + .../_torch/modules/test_ep_metrics.py | 219 ++++++ .../pyexecutor/test_ep_metrics_stats.py | 707 ++++++++++++++++++ tests/unittest/metrics/test_collector.py | 529 ++++++++++++- 14 files changed, 2118 insertions(+), 52 deletions(-) create mode 100644 tensorrt_llm/_torch/modules/fused_moe/ep_metrics.py create mode 100644 tests/unittest/_torch/modules/test_ep_metrics.py create mode 100644 tests/unittest/_torch/pyexecutor/test_ep_metrics_stats.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py b/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py index b1aba97c818c..c449b4bfb059 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py +++ b/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py @@ -13,26 +13,30 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""EP group health tracking for WideEP fault tolerance. +"""Committed EP data-plane membership for WideEP fault tolerance. This module provides :class:`EPGroupHealth`, a process-local, thread-safe data -structure that records the committed active membership of an Expert Parallel -(EP) group. It is the single source of truth for the data-plane rank mask within -one process and is consumed by: +structure that records which Expert Parallel (EP) ranks the recovery coordinator +has committed as included vs. excluded from the data plane. It is consumed by: * AlltoAll communication backends (rank masking on dispatch / combine) * The host-side AlltoAll watchdog (read-only expected-peer snapshot) * The MoE load balancer (emergency-mask reconfiguration) * The model engine and PyExecutor (degraded health reporting) -Failure detection does not mutate this object directly. Cross-process consensus, -expert-placement preparation, and atomic publication of a new committed mask are -the responsibility of higher-layer recovery coordination components and are not -performed here. +Detected or suspected physical liveness is separate evidence. Higher-layer +coordination reconciles that evidence and commits membership; detectors and +telemetry consumers must not mutate this object to drive recovery. """ import threading -from typing import NamedTuple +import uuid +from typing import Any, NamedTuple + +# Canonical ModelConfig.extra_attrs key for coordinator-committed data-plane +# membership. The object is not a detector or a physical-liveness tracker. +EP_GROUP_HEALTH_EXTRA_ATTR = "wide_ep_ft_ep_group_health" +EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR = "ep_group_health" # Default number of uint64 words for EPGroupHealth.get_mask_words(). # Matches the active_rank_mask ABI of the NVLink AlltoAll kernels @@ -51,24 +55,29 @@ class EPGroupHealthSnapshot(NamedTuple): if a mutator runs between calls). """ - # Active-rank bitmask (bit i set iff rank i is active). + # Committed-rank bitmask (bit i set iff rank i is included). mask: int - # Number of ranks currently marked active. + # Number of ranks included in committed data-plane membership. active_count: int - # Immutable snapshot of failed ranks. + # Immutable snapshot of ranks excluded from committed membership. failed_ranks: frozenset[int] - # Monotonic counter; bumps only on effective state changes. + # Monotonic committed-membership version. generation: int class EPGroupHealth: - """Thread-safe health tracker for the ranks of an EP group. + """Thread-safe committed-membership tracker for an EP group. Internally backed by an arbitrary-precision Python ``int`` bitmask - (bit ``i`` set iff rank ``i`` is currently active). The kernel-side - representation as a fixed-width array of ``uint64`` words is exposed via - :meth:`get_mask_words`; this matches the ``active_rank_mask`` parameter - expected by the NVLink AlltoAll kernels. + (bit ``i`` set iff rank ``i`` is included in the committed data plane). + The kernel-side representation as a fixed-width array of ``uint64`` words + is exposed via :meth:`get_mask_words`; this matches the + ``active_rank_mask`` parameter expected by the NVLink AlltoAll kernels. + + This object does not represent uncommitted detector observations or + suspected physical liveness. A higher-level recovery coordinator owns + membership transitions; passive consumers such as telemetry must only + read coherent snapshots. All read and write operations take an internal lock. Read operations that return collection types return defensive snapshots so the caller cannot @@ -100,6 +109,10 @@ def __init__(self, moe_world_size: int) -> None: self._active_count: int = moe_world_size self._failed_ranks: set[int] = set() self._generation: int = 0 + # A generation is monotonic only within one producer lifetime. The + # immutable epoch lets serving distinguish a restarted producer from a + # stale snapshot emitted by the current one. + self._source_epoch = uuid.uuid4().hex self._lock = threading.Lock() @property @@ -109,7 +122,7 @@ def moe_world_size(self) -> int: @property def generation(self) -> int: - """Monotonic counter incremented on every effective state change. + """Monotonic counter incremented on each committed membership change. Idempotent calls (e.g. marking an already-failed rank as failed) do not bump the counter. Consumers that need to react to mask changes can @@ -122,13 +135,22 @@ def generation(self) -> int: with self._lock: return self._generation + @property + def source_epoch(self) -> str: + """Unique identifier for this producer lifetime.""" + return self._source_epoch + def _validate_rank(self, rank: int) -> None: """Raise if ``rank`` is outside the MoE world.""" if not 0 <= rank < self._moe_world_size: raise ValueError(f"rank must be in [0, {self._moe_world_size}), got {rank}") def mark_failed(self, rank: int) -> bool: - """Mark ``rank`` as failed. Idempotent. + """Exclude ``rank`` from committed data-plane membership. Idempotent. + + A higher-level recovery coordinator owns calls to this mutator after + reconciling detector evidence; detection and telemetry must not call it + directly. Args: rank: Index of the rank to mark, in ``[0, moe_world_size)``. @@ -152,7 +174,7 @@ def mark_failed(self, rank: int) -> bool: return True def mark_active(self, rank: int) -> bool: - """Mark ``rank`` as active. Idempotent. + """Include ``rank`` in committed data-plane membership. Idempotent. Used when a replacement rank rejoins the group after a failure. Higher layers may impose a "monotonic failure" policy (do not @@ -181,7 +203,7 @@ def mark_active(self, rank: int) -> bool: return True def is_active(self, rank: int) -> bool: - """Return ``True`` iff ``rank`` is currently marked active. + """Return whether ``rank`` is included in committed membership. Raises: ValueError: If ``rank`` is outside ``[0, moe_world_size)``. @@ -191,15 +213,15 @@ def is_active(self, rank: int) -> bool: return bool(self._active_mask & (1 << rank)) def get_mask(self) -> int: - """Return the active-rank bitmask as a Python ``int``. + """Return the committed active-rank bitmask as a Python ``int``. - Bit ``i`` is set iff rank ``i`` is currently active. + Bit ``i`` is set iff rank ``i`` is included in data-plane membership. """ with self._lock: return self._active_mask def get_mask_words(self, num_words: int = EP_MASK_NUM_WORDS) -> tuple[int, ...]: - """Return the active-rank bitmask split into little-endian uint64 words. + """Return the committed rank mask as little-endian uint64 words. Suitable for passing to CUDA kernels that accept ``uint64_t[num_words]``. Word ``0`` covers ranks ``0..63``, word ``1`` covers ranks ``64..127``, @@ -230,22 +252,22 @@ def get_mask_words(self, num_words: int = EP_MASK_NUM_WORDS) -> tuple[int, ...]: return tuple((mask >> (i * 64)) & word_mask for i in range(num_words)) def get_active_count(self) -> int: - """Return the number of ranks currently marked active.""" + """Return the number of ranks in committed data-plane membership.""" with self._lock: return self._active_count def get_failed_ranks(self) -> frozenset[int]: - """Return an immutable snapshot of the set of failed ranks.""" + """Return ranks excluded from committed data-plane membership.""" with self._lock: return frozenset(self._failed_ranks) def all_active(self) -> bool: - """Return ``True`` iff every rank in the group is currently active.""" + """Return whether committed membership includes every group rank.""" with self._lock: return self._active_mask == self._all_active_mask def snapshot(self) -> EPGroupHealthSnapshot: - """Return an atomic snapshot of mask, active_count, failed_ranks, generation. + """Return an atomic committed-membership snapshot. All four fields reflect a single point in time (one lock acquisition). Prefer this over calling individual accessors back-to-back when the @@ -263,13 +285,18 @@ def __len__(self) -> int: """Return the MoE world size (``moe_world_size``), not the active count. Provided for compatibility with code that treats this object as a - sized container of ranks (alive or dead). Use :meth:`get_active_count` - when you specifically want the surviving-rank count. + sized container of included and excluded ranks. Use + :meth:`get_active_count` for the committed-membership count. """ return self._moe_world_size + def __deepcopy__(self, memo: dict[int, Any]) -> "EPGroupHealth": + """Preserve producer identity when ``ModelConfig`` is deep-copied.""" + memo[id(self)] = self + return self + def __repr__(self) -> str: - """Return a concise debug representation of the current health state.""" + """Return a concise representation of committed membership.""" with self._lock: return ( f"EPGroupHealth(moe_world_size={self._moe_world_size}, " diff --git a/tensorrt_llm/_torch/modules/fused_moe/ep_metrics.py b/tensorrt_llm/_torch/modules/fused_moe/ep_metrics.py new file mode 100644 index 000000000000..121da2757ac1 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/ep_metrics.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Prometheus-independent serialization of committed EP membership. + +Workers serialize the recovery coordinator's committed data-plane membership +for a dedicated serving RPC without importing the client before its +multiprocess directory is configured. This module is a passive observer: it +does not expose detected/suspected physical liveness and must not drive +recovery. +""" + +from typing import Any + +from .ep_group_health import ( + EP_GROUP_HEALTH_EXTRA_ATTR, + EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR, + EPGroupHealth, +) + +_EP_HEALTH_METRICS_STATUS_KEY = "status" +_EP_HEALTH_METRICS_PENDING_STATUS = "pending" + + +def validate_ep_health_metrics_topology(model_config: Any) -> None: + """Reject topologies whose rank-0 endpoint cannot expose every EP group.""" + mapping = getattr(model_config, "mapping", None) + if any( + getattr(mapping, axis, 1) > 1 for axis in ("moe_tp_size", "pp_size", "moe_cluster_size") + ): + raise NotImplementedError("EP health telemetry does not support multi-group MoE topologies") + + +def get_ep_group_health(model_config: Any) -> EPGroupHealth | None: + """Return an already-configured coordinator-owned membership tracker. + + ``ModelConfig.extra_attrs`` shares the committed data-plane membership + tracker with communication. It does not carry raw detector observations or + suspected physical liveness. This telemetry consumer never creates, + registers, enables, or mutates it. The legacy prototype key remains + readable during rollout; if both keys are populated they must reference the + same object. + """ + extra_attrs = getattr(model_config, "extra_attrs", None) + if extra_attrs is None: + return None + + health = extra_attrs.get(EP_GROUP_HEALTH_EXTRA_ATTR) + legacy_health = extra_attrs.get(EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR) + if health is not None and legacy_health is not None and health is not legacy_health: + raise ValueError("configured EP group health keys reference different trackers") + health = health if health is not None else legacy_health + if health is None: + return None + if not isinstance(health, EPGroupHealth): + raise TypeError("configured EP group health must be an EPGroupHealth") + + mapping = getattr(model_config, "mapping", None) + moe_world_size = getattr(mapping, "moe_ep_size", None) + if ( + type(moe_world_size) is int + and moe_world_size > 0 + and health.moe_world_size != moe_world_size + ): + raise ValueError("configured EP group health world size does not match model mapping") + return health + + +def is_ep_group_health_registration_pending(model_config: Any) -> bool: + """Return whether a producer reserved a tracker slot for late attachment.""" + extra_attrs = getattr(model_config, "extra_attrs", None) + if extra_attrs is None: + return False + return any( + key in extra_attrs and extra_attrs[key] is None + for key in (EP_GROUP_HEALTH_EXTRA_ATTR, EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR) + ) + + +def pending_ep_health_metrics() -> dict[str, str]: + """Return the JSON-safe marker for an explicitly pending registration.""" + return {_EP_HEALTH_METRICS_STATUS_KEY: _EP_HEALTH_METRICS_PENDING_STATUS} + + +def is_pending_ep_health_metrics(ep_health_stats: dict[str, Any]) -> bool: + """Return whether a worker reported an explicitly pending registration.""" + return ep_health_stats == pending_ep_health_metrics() + + +def serialize_ep_health_metrics(ep_group_health: EPGroupHealth) -> dict[str, Any]: + """Serialize one coherent committed-membership snapshot. + + ``failedRanks`` contains ranks excluded from the committed data-plane mask; + it is not a physical-liveness or suspicion signal. The per-rank gauge is + derived from ``worldSize`` and ``failedRanks`` without mutating recovery + state. + """ + snapshot = ep_group_health.snapshot() + return { + "sourceEpoch": ep_group_health.source_epoch, + "worldSize": ep_group_health.moe_world_size, + "activeCount": snapshot.active_count, + "failedRanks": sorted(snapshot.failed_ranks), + "generation": snapshot.generation, + } diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 188ab3302dda..507116f8571f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -51,6 +51,10 @@ maybe_prefetch_mm_encoder_for_next_iter from ..models.modeling_utils import DecoderModelForCausalLM from ..modules.decoder_layer import DecoderLayer +from ..modules.fused_moe.ep_metrics import ( + get_ep_group_health, is_ep_group_health_registration_pending, + pending_ep_health_metrics, serialize_ep_health_metrics, + validate_ep_health_metrics_topology) from ..speculative.drafter import Drafter from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.speculation_gate import SpeculationGate @@ -472,6 +476,8 @@ def __init__( self.resource_manager = resource_manager self.scheduler = scheduler self.model_engine = model_engine + model = getattr(model_engine, "model", None) + self._ep_health_model_config = getattr(model, "model_config", None) self.enable_attention_dp = model_engine.enable_attention_dp self.dist = dist self.sampler = sampler @@ -1193,6 +1199,19 @@ def can_enqueue_requests(self) -> bool: """ return self.executor_request_queue.can_enqueue_request() + def _get_ep_health_stats(self) -> Optional[dict]: + """Return committed EP membership for passive internal telemetry.""" + ep_group_health = get_ep_group_health(self._ep_health_model_config) + if ep_group_health is None: + if is_ep_group_health_registration_pending( + self._ep_health_model_config): + validate_ep_health_metrics_topology( + self._ep_health_model_config) + return pending_ep_health_metrics() + return None + validate_ep_health_metrics_topology(self._ep_health_model_config) + return serialize_ep_health_metrics(ep_group_health) + def get_latest_iteration_stats(self): """ Returns the per-iterations statistics computed since last call to this method. diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index a03fbb4a3c56..16edc6afa1b6 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -280,6 +280,17 @@ def fetch_stats(self) -> list: else: return self.engine.get_latest_iteration_stats() + def _get_ep_health_stats(self, timeout: float = 1.0) -> Optional[dict]: + """Return committed EP membership when the backend supports it.""" + health_reader = getattr(self.engine, "_get_ep_health_stats", None) + if health_reader is None: + return None + return health_reader() + + def fetch_ep_health_stats(self) -> Optional[dict]: + """RPC entry point for passive committed-membership telemetry.""" + return self._get_ep_health_stats() + def fetch_kv_cache_events(self) -> list: if isinstance(self.engine, tllm.Executor): return self.engine.get_latest_kv_cache_events() diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index fdd3dff44a48..864310834c6f 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -406,6 +406,10 @@ def get_stats(self, timeout: float) -> List[dict]: self._iter_stats_result.set_timeout(timeout) return self._iter_stats_result.get_results() + def _get_ep_health_stats(self, timeout: float = 1.0) -> Optional[dict]: + """Return passive committed EP membership telemetry when supported.""" + return None + def aget_stats(self, timeout: float) -> IterationResult: """ Get iteration statistics from the runtime. diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index ffe793bc5b48..c9294e30f631 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -624,6 +624,19 @@ def get_stats(self, timeout: float) -> List[dict]: stats = self.rpc_client.fetch_stats_wait_async(timeout=timeout).remote() return [json.loads(s) if isinstance(s, str) else s for s in stats] + def _get_ep_health_stats(self, timeout: float = 1.0) -> Optional[dict]: + """Fetch best-effort committed EP membership from the rank-0 worker. + + The rank-0 RPC is the only MPI worker control endpoint today. Callers + must treat ``ep_health_available`` as the validity bit if that endpoint + is lost. The snapshot is not detected/suspected physical liveness, and + this passive metrics hook must not drive recovery. Survivor-side + observation belongs to the fault-tolerant control-plane work. + """ + if self.rpc_client is None: + return None + return self.rpc_client.fetch_ep_health_stats().remote(timeout=timeout) + def get_disaggregated_params(self) -> dict: """Get disaggregated params from worker runtime via RPC.""" if self.rpc_client is None: diff --git a/tensorrt_llm/executor/rpc_proxy_mixin.py b/tensorrt_llm/executor/rpc_proxy_mixin.py index ecbb86e25ead..a432ffea1fa5 100644 --- a/tensorrt_llm/executor/rpc_proxy_mixin.py +++ b/tensorrt_llm/executor/rpc_proxy_mixin.py @@ -165,6 +165,10 @@ async def _fetch_responses_loop_async(self): method_name="_fetch_responses_loop_async", ) + def _get_ep_health_stats(self, timeout: float = 1.0) -> Optional[dict]: + """Fetch passive committed EP membership from the rank-0 RPC worker.""" + return self.rpc_client.fetch_ep_health_stats().remote(timeout=timeout) + # NOTE: _fetch_stats_loop_async and _fetch_kv_cache_events_loop_async have been removed. # Stats and kv_events are now fetched on-demand via direct RPC calls # (get_stats, aget_stats, get_kv_events, aget_kv_events) instead of streaming loops. diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index b637c1738dc8..309fd78ca1e3 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -1181,6 +1181,10 @@ def get_stats_async(self, timeout: Optional[float] = 2) -> IterationResult: "Use llm.encode() for encoder-only models.") return self._executor.aget_stats(timeout=timeout) + def _get_ep_health_stats(self) -> Optional[dict]: + """Return committed EP membership for the passive metrics hook.""" + return self._executor._get_ep_health_stats() + @set_api_status("beta") def get_kv_cache_events(self, timeout: Optional[float] = 2) -> List[dict]: """Get iteration KV events from the runtime. diff --git a/tensorrt_llm/metrics/collector.py b/tensorrt_llm/metrics/collector.py index cd5a91123ddf..7fb52393bd46 100644 --- a/tensorrt_llm/metrics/collector.py +++ b/tensorrt_llm/metrics/collector.py @@ -14,12 +14,131 @@ # limitations under the License. """Utilities for Prometheus Metrics Collection.""" +import logging import math +import threading import time -from typing import Dict, List, Optional, Union +from collections import deque +from typing import Dict, List, NamedTuple, Optional, Union from .enums import MetricNames +_LOGGER = logging.getLogger(__name__) + +# Phase-1 WideEP uses the two-word uint64 active-rank-mask ABI. Reject a +# malformed payload before it can create unbounded Prometheus cardinality on +# the serving event loop. +_MAX_EP_HEALTH_RANKS = 128 +# Prometheus gauges are IEEE-754 doubles. Keep the monotonic generation exact +# so stale/conflict comparisons have the same meaning before and after export. +_MAX_EP_HEALTH_GENERATION = (1 << 53) - 1 +_MAX_EP_HEALTH_SOURCE_EPOCH_LENGTH = 128 +_MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS = 32 + + +class _EPHealthSnapshot(NamedTuple): + """One immutable snapshot exported by the local Prometheus collector.""" + + world_size: int + active_count: int + failed_ranks: frozenset[int] + generation: int + + +class _EPHealthPrometheusCollector: + """Export a coherent EP snapshot without multiprocess gauge merging. + + ``prometheus_client`` combines each multiprocess gauge independently. A + rank vector, its aggregate counts, and its validity bit therefore cannot + form an atomic snapshot when more than one process writes the same label + set. EP health is polled by the OpenAI server itself, so keep its latest + immutable value in that process and register this collector directly on + the server's scrape registry instead. + """ + + def __init__(self, labels: Dict[str, str], metric_prefix: str) -> None: + self._label_names = list(labels) + self._label_values = [labels[name] for name in self._label_names] + self._metric_prefix = metric_prefix + self._lock = threading.Lock() + self._snapshot: Optional[_EPHealthSnapshot] = None + self._available = False + + def publish(self, snapshot: _EPHealthSnapshot) -> None: + """Atomically replace the snapshot and mark it available.""" + with self._lock: + self._snapshot = snapshot + self._available = True + + def mark_unavailable(self) -> None: + """Atomically invalidate the retained snapshot.""" + with self._lock: + self._available = False + + def collect(self) -> list: + """Build all metric families from one locked state read.""" + with self._lock: + snapshot = self._snapshot + available = self._available + + # Do not advertise EP health for non-EP or non-FT deployments. Once a + # producer has published a valid snapshot, retain its samples during an + # outage and use the availability family as their validity bit. + if snapshot is None: + return [] + + from prometheus_client.core import GaugeMetricFamily + + metrics = [] + rank_active = GaugeMetricFamily( + self._metric_prefix + "ep_rank_active", + "Whether an Expert Parallel rank is included in coordinator-committed " + "data-plane membership (1) or excluded (0); not physical liveness.", + labels=[*self._label_names, MetricsCollector.labelname_ep_rank], + ) + for rank in range(snapshot.world_size): + rank_active.add_metric( + [*self._label_values, str(rank)], + int(rank not in snapshot.failed_ranks), + ) + metrics.append(rank_active) + + for name, documentation, value in ( + ( + "ep_active_ranks", + "Number of ranks in coordinator-committed EP data-plane membership.", + snapshot.active_count, + ), + ( + "ep_failed_ranks", + "Number of ranks excluded from coordinator-committed EP data-plane " + "membership; not a physical-failure detector.", + snapshot.world_size - snapshot.active_count, + ), + ( + "ep_health_generation", + "Committed EP membership version; increments on each effective transition.", + snapshot.generation, + ), + ): + metric = GaugeMetricFamily( + self._metric_prefix + name, + documentation, + labels=self._label_names, + ) + metric.add_metric(self._label_values, value) + metrics.append(metric) + + availability = GaugeMetricFamily( + self._metric_prefix + "ep_health_available", + "Whether rank-0 committed EP membership telemetry is available. " + "When 0, ignore all other trtllm_ep_* health samples.", + labels=self._label_names, + ) + availability.add_metric(self._label_values, int(available)) + metrics.append(availability) + return metrics + # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0rc1/vllm/engine/metrics.py#L30 class MetricsCollector: @@ -52,6 +171,13 @@ class MetricsCollector: trtllm_generation_perplexity trtllm_request_error_total + Expert Parallel health metrics: + trtllm_ep_rank_active + trtllm_ep_active_ranks + trtllm_ep_failed_ranks + trtllm_ep_health_generation + trtllm_ep_health_available + Iteration-level metrics: trtllm_kv_cache_hit_rate trtllm_kv_cache_utilization @@ -102,6 +228,7 @@ class MetricsCollector: trtllm_kv_cache_config_info """ labelname_finish_reason = "finished_reason" + labelname_ep_rank = "ep_rank" def __init__( self, @@ -315,6 +442,17 @@ def __init__( documentation="Maximum number of active requests", labelnames=self.labels.keys()) + # Committed Expert Parallel membership is passively observed from the + # rank-0 worker through a dedicated internal RPC. Unlike independent + # multiprocess gauges, this local collector preserves the rank vector, + # aggregate counts, generation, and availability as one snapshot. + self._last_ep_health_state = None + self._last_ep_health_rejection = None + self._retired_ep_health_source_epochs = set() + self._retired_ep_health_source_epoch_order = deque() + self._ep_health_prometheus_collector = _EPHealthPrometheusCollector( + self.labels, self.metric_prefix) + # Iteration latency self.iteration_latency_seconds = Gauge( name=self.metric_prefix + "iteration_latency_seconds", @@ -868,6 +1006,113 @@ def log_iteration_stats(self, iteration_stats: dict) -> None: self._log_counter(self.kv_cache_intra_device_copy_bytes_total, {}, total_intra_device_copy_bytes) + def _reject_ep_health_stats(self, message: str, *args: object) -> bool: + """Mark telemetry unavailable and warn once per repeated rejection.""" + signature = (message, tuple(str(arg) for arg in args)) + if signature != self._last_ep_health_rejection: + _LOGGER.warning(message, *args) + self._last_ep_health_rejection = signature + self.log_ep_health_unavailable() + return False + + def log_ep_health_stats(self, ep_health_stats: dict) -> bool: + """Materialize committed ``EPGroupHealth`` membership as gauges. + + Returns ``False`` when the payload is invalid, stale, or conflicts with + the last accepted snapshot. Callers may retry because a later generation + or a never-before-seen producer epoch can restore a coherent stream. + ``sourceEpoch`` is required so a producer restart cannot be confused + with delayed state from a retired producer. This passive consumer does + not interpret the snapshot as physical liveness or mutate recovery + state. + """ + try: + world_size = ep_health_stats["worldSize"] + active_count = ep_health_stats["activeCount"] + generation = ep_health_stats["generation"] + failed_rank_list = ep_health_stats["failedRanks"] + source_epoch = ep_health_stats["sourceEpoch"] + scalar_values = (world_size, active_count, generation) + if any(type(value) is not int for value in scalar_values): + raise TypeError( + "worldSize, activeCount, and generation must be integers") + if (world_size <= 0 or world_size > _MAX_EP_HEALTH_RANKS + or not 0 <= active_count <= world_size or generation < 0 + or generation > _MAX_EP_HEALTH_GENERATION): + raise ValueError("EP health scalar values are out of range") + if not isinstance(failed_rank_list, list): + raise TypeError("failedRanks must be a list") + if (not isinstance(source_epoch, str) or not source_epoch + or len(source_epoch) > _MAX_EP_HEALTH_SOURCE_EPOCH_LENGTH): + raise TypeError( + "sourceEpoch must be a bounded non-empty string") + if len(failed_rank_list) > world_size: + raise ValueError("failedRanks contains too many ranks") + if any( + type(rank) is not int or not 0 <= rank < world_size + for rank in failed_rank_list): + raise ValueError("failedRanks contains an invalid rank") + failed_ranks = set(failed_rank_list) + if (len(failed_ranks) != len(failed_rank_list) + or active_count != world_size - len(failed_ranks)): + raise ValueError("EP health counts are inconsistent") + except (KeyError, TypeError, ValueError) as error: + return self._reject_ep_health_stats( + "Ignoring invalid epHealthStats payload: %s", error) + + state = (source_epoch, world_size, generation, + tuple(sorted(failed_ranks))) + last_source_epoch = None + if self._last_ep_health_state is not None: + last_source_epoch, last_world_size, last_generation, last_failed_ranks = ( + self._last_ep_health_state) + if world_size != last_world_size: + return self._reject_ep_health_stats( + "Ignoring epHealthStats world-size change from %s to %s", + last_world_size, world_size) + if source_epoch == last_source_epoch and generation < last_generation: + return self._reject_ep_health_stats( + "Ignoring stale epHealthStats generation %s after %s", + generation, last_generation) + if (source_epoch == last_source_epoch + and generation == last_generation + and state[3] != last_failed_ranks): + return self._reject_ep_health_stats( + "Conflicting epHealthStats payloads at generation %s", + generation) + if source_epoch != last_source_epoch: + if source_epoch in self._retired_ep_health_source_epochs: + return self._reject_ep_health_stats( + "Ignoring epHealthStats from retired source epoch %s", + source_epoch) + + if last_source_epoch is not None and source_epoch != last_source_epoch: + if (len(self._retired_ep_health_source_epoch_order) == + _MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS): + expired_epoch = self._retired_ep_health_source_epoch_order.popleft( + ) + self._retired_ep_health_source_epochs.remove(expired_epoch) + self._retired_ep_health_source_epoch_order.append(last_source_epoch) + self._retired_ep_health_source_epochs.add(last_source_epoch) + self._last_ep_health_state = state + self._last_ep_health_rejection = None + self._ep_health_prometheus_collector.publish( + _EPHealthSnapshot( + world_size=world_size, + active_count=active_count, + failed_ranks=frozenset(failed_ranks), + generation=generation, + )) + return True + + def log_ep_health_unavailable(self) -> None: + """Mark the last EP health telemetry read as unavailable.""" + self._ep_health_prometheus_collector.mark_unavailable() + + def register_ep_health_metrics(self, registry: object) -> None: + """Register coherent EP health metrics on a scrape registry.""" + registry.register(self._ep_health_prometheus_collector) + def log_request_error(self, http_code: Union[int, str] = "") -> None: """Increment the error counter, labeled by HTTP status code.""" labels = {**self.labels, self.labelname_http_code: str(http_code)} diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 1424e409fc98..84f558239b1f 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -9,6 +9,7 @@ import time import traceback import uuid +from builtins import anext from collections import deque from contextlib import asynccontextmanager from datetime import datetime @@ -18,6 +19,7 @@ Optional, Union) import uvicorn +import zmq from fastapi import Body, FastAPI, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import (FileResponse, JSONResponse, Response, @@ -28,10 +30,13 @@ from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm._torch.async_llm import AsyncLLM +from tensorrt_llm._torch.modules.fused_moe.ep_metrics import \ + is_pending_ep_health_metrics from tensorrt_llm._utils import EnergyMonitor # yapf: disable from tensorrt_llm.executor import CppExecutorError from tensorrt_llm.executor.postproc_worker import PostprocParams +from tensorrt_llm.executor.rpc import RPCCancelled, RPCError, RPCTimeout from tensorrt_llm.inputs import prompt_inputs from tensorrt_llm.inputs.data import TokensPrompt from tensorrt_llm.inputs.multimodal import MultimodalServerConfig @@ -93,6 +98,51 @@ TIMEOUT_KEEP_ALIVE = 5 # seconds. +def _is_transient_ep_health_error(error: Exception) -> bool: + """Return whether an EP health read can reasonably succeed on retry.""" + if isinstance( + error, + (RPCCancelled, RPCTimeout, TimeoutError, OSError, zmq.ZMQError)): + return True + if isinstance(error, RPCError) and error.cause is not None: + return _is_transient_ep_health_error(error.cause) + return False + + +async def _await_task_cancellation_safe(task: asyncio.Task) -> Any: + """Await ``task`` without forwarding repeated owner cancellation to it. + + ``asyncio.shield`` protects a child only from the cancellation that reaches + that particular shield future. After catching the first cancellation, an + unshielded ``await task`` lets a later ``Task.cancel()`` cancel the child. + That is especially unsafe for ``asyncio.to_thread``: cancelling its asyncio + wrapper does not stop the worker thread. Drain through a separately + shielded task instead, retrying the shield after every owner cancellation, + then preserve the first cancellation request. + """ + try: + return await asyncio.shield(task) + except asyncio.CancelledError as cancellation: + + async def drain_task() -> None: + try: + await task + except (Exception, asyncio.CancelledError): + pass + + drain = asyncio.create_task(drain_task()) + while not drain.done(): + try: + await asyncio.shield(drain) + except asyncio.CancelledError: + continue + # Retrieve an unexpected BaseException from the drain before restoring + # cancellation. Ordinary task errors and child cancellation were + # consumed above to preserve the owner's cancellation semantics. + drain.result() + raise cancellation + + def _build_tool_strict_guided_decoding_params(tools, tool_parser_name): """Build GuidedDecodingParams with structural tags for tools with strict=True. @@ -212,6 +262,9 @@ def __init__( self.perf_metrics = None self.perf_metrics_lock = None self._iteration_stats_collector_task = None + self._ep_health_collector_task = None + self._ep_health_stop_event = asyncio.Event() + self._ep_health_outage_logged = False self._iteration_stats_wakeup_event = asyncio.Event() # Bounded snapshot of iteration stats for the GET /metrics handler. # When the background Prometheus collector loop is active, it is the @@ -292,25 +345,34 @@ async def lifespan(app: FastAPI): logger.info( "Started background iteration stats collector task") - yield + await self._start_ep_health_collector() - # Stop background iteration stats collector - if self._iteration_stats_collector_task is not None: - self._iteration_stats_collector_task.cancel() + try: + yield + finally: try: - await self._iteration_stats_collector_task - except asyncio.CancelledError: - pass - logger.info("Stopped background iteration stats collector task") - - if self.metadata_server is not None: - self.metadata_server.remove(f"trtllm/{self.generator.llm_id}") - logger.info(f"trtllm/{self.generator.llm_id} is unregistered") - if self.disagg_cluster_worker: - await self.disagg_cluster_worker.deregister_worker() - if self.resource_governor is not None: - self.resource_governor.close() - self.generator.shutdown() + await self._stop_ep_health_collector() + finally: + # Stop background iteration stats collector + if self._iteration_stats_collector_task is not None: + self._iteration_stats_collector_task.cancel() + try: + await self._iteration_stats_collector_task + except asyncio.CancelledError: + pass + logger.info( + "Stopped background iteration stats collector task") + + if self.metadata_server is not None: + self.metadata_server.remove( + f"trtllm/{self.generator.llm_id}") + logger.info( + f"trtllm/{self.generator.llm_id} is unregistered") + if self.disagg_cluster_worker: + await self.disagg_cluster_worker.deregister_worker() + if self.resource_governor is not None: + self.resource_governor.close() + self.generator.shutdown() self.app = FastAPI(lifespan=lifespan) @@ -750,6 +812,7 @@ def mount_metrics(self): from prometheus_fastapi_instrumentator import Instrumentator registry = CollectorRegistry() multiprocess.MultiProcessCollector(registry) + self.metrics_collector.register_ep_health_metrics(registry) Instrumentator( should_group_status_codes=False, should_respect_env_var=True, @@ -1170,6 +1233,113 @@ async def _iteration_stats_collector_loop(self): logger.info("Iteration stats collector loop cancelled") raise + def _log_ep_health_outage_once(self, + message: str, + error: Optional[Exception] = None, + terminal: bool = False) -> None: + """Log one retry diagnostic, while always logging a terminal state.""" + if (getattr(self, "_ep_health_outage_logged", False) and not terminal): + return + detail = f": {error}" if error is not None else "" + logger.warning(f"{message}{detail}") + self._ep_health_outage_logged = True + + def _get_ep_health_stop_event(self) -> asyncio.Event: + """Return the lazily initialized health-collector stop event.""" + stop_event = getattr(self, "_ep_health_stop_event", None) + if stop_event is None: + stop_event = asyncio.Event() + self._ep_health_stop_event = stop_event + return stop_event + + async def _wait_for_ep_health_stop(self, timeout: float) -> bool: + """Wait up to ``timeout`` seconds for a collector stop request.""" + stop_event = self._get_ep_health_stop_event() + if stop_event.is_set(): + return True + try: + await asyncio.wait_for(stop_event.wait(), timeout=timeout) + except asyncio.TimeoutError: + return False + return True + + async def _read_ep_health_stats(self) -> Optional[dict]: + """Run one synchronous health read without abandoning its worker thread.""" + read_task = asyncio.create_task( + asyncio.to_thread(self.generator._get_ep_health_stats)) + # asyncio.to_thread cancellation does not stop the underlying thread. + # Drain it before allowing the owner to close the RPC client. + return await _await_task_cancellation_safe(read_task) + + async def _stop_ep_health_collector(self) -> None: + """Cooperatively stop and cancellation-safely drain the collector.""" + task = self._ep_health_collector_task + if task is None: + return + self._get_ep_health_stop_event().set() + try: + await _await_task_cancellation_safe(task) + finally: + self._ep_health_collector_task = None + logger.info("Stopped EP health metrics collector task") + + async def _ep_health_collector_loop(self) -> None: + """Passively poll rank-0's committed EP membership for metrics. + + This loop observes coordinator-committed state, not detected or + suspected physical liveness, and never drives recovery. + """ + try: + while not self._get_ep_health_stop_event().is_set(): + try: + stats = await self._read_ep_health_stats() + if stats is None: + self.metrics_collector.log_ep_health_unavailable() + logger.info( + "EP health telemetry is unsupported; polling disabled" + ) + return + if is_pending_ep_health_metrics(stats): + self.metrics_collector.log_ep_health_unavailable() + self._log_ep_health_outage_once( + "EP health telemetry registration is pending; will retry" + ) + elif not self.metrics_collector.log_ep_health_stats(stats): + self._log_ep_health_outage_once( + "EP health telemetry returned an invalid, stale, or conflicting payload; will retry", + ) + else: + self._ep_health_outage_logged = False + except Exception as e: + self.metrics_collector.log_ep_health_unavailable() + if not _is_transient_ep_health_error(e): + self._log_ep_health_outage_once( + "EP health telemetry failed deterministically; stopping polling", + e, + terminal=True) + return + self._log_ep_health_outage_once( + "Transient error collecting EP health telemetry; will retry", + e) + if await self._wait_for_ep_health_stop(timeout=1.0): + return + except asyncio.CancelledError: + logger.info("EP health metrics collector loop cancelled") + raise + + async def _start_ep_health_collector(self) -> None: + """Start EP health polling without delaying server readiness.""" + if not self.metrics_collector: + return + if self._ep_health_collector_task is not None: + if not self._ep_health_collector_task.done(): + return + self._ep_health_collector_task = None + self._get_ep_health_stop_event().clear() + self._ep_health_collector_task = asyncio.create_task( + self._ep_health_collector_loop()) + logger.info("Started EP health metrics collector task") + async def openai_chat(self, request: ChatCompletionRequest, raw_request: Request) -> Response: diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a958f8b24951..cf28b7e684da 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -193,6 +193,8 @@ l0_b200: # ------------- KV Cache Iteration Stats --------------- - unittest/executor/test_stats_serializer.py - unittest/metrics/test_collector.py + - unittest/_torch/modules/test_ep_metrics.py + - unittest/_torch/pyexecutor/test_ep_metrics_stats.py - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_cold_start - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_partial_block_reuse - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_full_block_reuse diff --git a/tests/unittest/_torch/modules/test_ep_metrics.py b/tests/unittest/_torch/modules/test_ep_metrics.py new file mode 100644 index 000000000000..25f95004f19c --- /dev/null +++ b/tests/unittest/_torch/modules/test_ep_metrics.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for passive committed WideEP membership serialization.""" + +import copy +import threading +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import ( + EP_GROUP_HEALTH_EXTRA_ATTR, + EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR, + EPGroupHealth, +) +from tensorrt_llm._torch.modules.fused_moe.ep_metrics import ( + get_ep_group_health, + is_ep_group_health_registration_pending, + is_pending_ep_health_metrics, + pending_ep_health_metrics, + serialize_ep_health_metrics, + validate_ep_health_metrics_topology, +) + + +def test_serialize_initial_health() -> None: + health = EPGroupHealth(4) + + assert serialize_ep_health_metrics(health) == { + "sourceEpoch": health.source_epoch, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + + +def test_deepcopy_preserves_shared_producer_identity() -> None: + health = EPGroupHealth(4) + + assert copy.deepcopy(health) is health + + +def test_serialization_reads_one_coherent_snapshot() -> None: + health: Any = SimpleNamespace( + moe_world_size=4, + source_epoch="source-a", + snapshot=MagicMock( + return_value=SimpleNamespace( + active_count=3, + failed_ranks=frozenset({2}), + generation=1, + ) + ), + ) + + assert serialize_ep_health_metrics(health) == { + "sourceEpoch": "source-a", + "worldSize": 4, + "activeCount": 3, + "failedRanks": [2], + "generation": 1, + } + health.snapshot.assert_called_once_with() + + +def test_serialize_committed_exclusion_and_reinclusion() -> None: + health = EPGroupHealth(72) + health.mark_failed(70) + health.mark_failed(3) + health.mark_failed(3) + + assert serialize_ep_health_metrics(health) == { + "sourceEpoch": health.source_epoch, + "worldSize": 72, + "activeCount": 70, + "failedRanks": [3, 70], + "generation": 2, + } + + health.mark_active(3) + assert serialize_ep_health_metrics(health) == { + "sourceEpoch": health.source_epoch, + "worldSize": 72, + "activeCount": 71, + "failedRanks": [70], + "generation": 3, + } + + +@pytest.mark.parametrize( + "key", + [EP_GROUP_HEALTH_EXTRA_ATTR, EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR], +) +def test_get_ep_group_health_passively_reads_registered_tracker(key: str) -> None: + health = EPGroupHealth(4) + extra_attrs = {key: health} + model_config = SimpleNamespace( + extra_attrs=extra_attrs, + mapping=SimpleNamespace(moe_ep_size=4), + ) + + assert get_ep_group_health(model_config) is health + assert model_config.extra_attrs == extra_attrs + + +def test_get_ep_group_health_does_not_create_or_enable_tracker(monkeypatch) -> None: + monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") + model_config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(moe_ep_size=8), + ) + + assert get_ep_group_health(model_config) is None + assert not is_ep_group_health_registration_pending(model_config) + assert model_config.extra_attrs == {} + + +def test_explicit_empty_registration_is_pending_without_mutation() -> None: + model_config = SimpleNamespace(extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: None}) + + assert get_ep_group_health(model_config) is None + assert is_ep_group_health_registration_pending(model_config) + pending = pending_ep_health_metrics() + assert pending == {"status": "pending"} + assert is_pending_ep_health_metrics(pending) + assert not is_pending_ep_health_metrics({"status": "unsupported"}) + assert model_config.extra_attrs == {EP_GROUP_HEALTH_EXTRA_ATTR: None} + + +def test_get_ep_group_health_rejects_split_trackers() -> None: + model_config = SimpleNamespace( + extra_attrs={ + EP_GROUP_HEALTH_EXTRA_ATTR: EPGroupHealth(4), + EP_GROUP_HEALTH_LEGACY_EXTRA_ATTR: EPGroupHealth(4), + }, + mapping=SimpleNamespace(moe_ep_size=4), + ) + + with pytest.raises(ValueError, match="different trackers"): + get_ep_group_health(model_config) + + +def test_get_ep_group_health_rejects_invalid_shared_object() -> None: + wrong_type = SimpleNamespace( + extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: object()}, + mapping=SimpleNamespace(moe_ep_size=8), + ) + with pytest.raises(TypeError, match="must be an EPGroupHealth"): + get_ep_group_health(wrong_type) + + wrong_size = SimpleNamespace( + extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: EPGroupHealth(4)}, + mapping=SimpleNamespace(moe_ep_size=8), + ) + with pytest.raises(ValueError, match="world size does not match"): + get_ep_group_health(wrong_size) + + +def test_topology_validation_is_deferred_to_metric_read() -> None: + health = EPGroupHealth(4) + model_config = SimpleNamespace( + extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: health}, + mapping=SimpleNamespace( + moe_ep_size=4, + moe_tp_size=1, + pp_size=2, + moe_cluster_size=1, + ), + ) + + assert get_ep_group_health(model_config) is health + with pytest.raises(NotImplementedError, match="multi-group MoE topologies"): + validate_ep_health_metrics_topology(model_config) + + +def test_concurrent_serialization_is_coherent() -> None: + health = EPGroupHealth(8) + start = threading.Barrier(2) + worker_errors: list[BaseException] = [] + + def toggle_rank() -> None: + try: + start.wait(timeout=10.0) + for _ in range(1_000): + health.mark_failed(3) + health.mark_active(3) + except BaseException as error: + worker_errors.append(error) + + thread = threading.Thread(target=toggle_rank) + thread.start() + start.wait(timeout=10.0) + + for _ in range(1_000): + stats = serialize_ep_health_metrics(health) + assert stats["activeCount"] == stats["worldSize"] - len(stats["failedRanks"]) + assert stats["failedRanks"] in ([], [3]) + assert bool(stats["generation"] % 2) is bool(stats["failedRanks"]) + + thread.join(timeout=10.0) + assert not thread.is_alive() + assert not worker_errors + assert health.generation == 2_000 + assert health.get_failed_ranks() == frozenset() diff --git a/tests/unittest/_torch/pyexecutor/test_ep_metrics_stats.py b/tests/unittest/_torch/pyexecutor/test_ep_metrics_stats.py new file mode 100644 index 000000000000..5f8faa5d56d7 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_ep_metrics_stats.py @@ -0,0 +1,707 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the passive PyExecutor committed-membership metrics transport.""" + +import asyncio +import threading +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +import zmq + +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import ( + EP_GROUP_HEALTH_EXTRA_ATTR, + EPGroupHealth, +) +from tensorrt_llm._torch.modules.fused_moe.ep_metrics import pending_ep_health_metrics +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm.executor.base_worker import BaseWorker +from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.executor.rpc import RPCCancelled, RPCClient, RPCError, RPCServer, RPCTimeout +from tensorrt_llm.executor.rpc.rpc_common import get_unique_ipc_addr +from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy +from tensorrt_llm.executor.rpc_proxy_mixin import RpcExecutorMixin +from tensorrt_llm.llmapi.llm import BaseLLM +from tensorrt_llm.serve.openai_server import OpenAIServer + +EP_HEALTH_SOURCE_EPOCH = "source-a" + + +def _make_executor(health: EPGroupHealth) -> PyExecutor: + executor = PyExecutor.__new__(PyExecutor) + executor._ep_health_model_config = SimpleNamespace( + extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: health} + ) + executor.model_engine = SimpleNamespace( + model=SimpleNamespace(model_config=executor._ep_health_model_config) + ) + return executor + + +def test_pyexecutor_returns_committed_membership_snapshot() -> None: + health = EPGroupHealth(4) + health.mark_failed(2) + executor = _make_executor(health) + + assert executor._get_ep_health_stats() == { + "sourceEpoch": health.source_epoch, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [2], + "generation": 1, + } + executor.enable_iter_perf_stats = False + assert executor.get_latest_iteration_stats() == [] + + +def test_pyexecutor_discovers_tracker_attached_after_initialization() -> None: + executor = PyExecutor.__new__(PyExecutor) + executor._ep_health_model_config = SimpleNamespace( + extra_attrs={EP_GROUP_HEALTH_EXTRA_ATTR: None} + ) + assert executor._get_ep_health_stats() == pending_ep_health_metrics() + + health = EPGroupHealth(4) + health.mark_failed(1) + executor._ep_health_model_config.extra_attrs[EP_GROUP_HEALTH_EXTRA_ATTR] = health + + stats = executor._get_ep_health_stats() + assert stats["sourceEpoch"] == health.source_epoch + assert stats["failedRanks"] == [1] + + +def test_pyexecutor_rejects_unqualified_multi_group_metrics() -> None: + executor = _make_executor(EPGroupHealth(4)) + executor._ep_health_model_config.mapping = SimpleNamespace( + moe_tp_size=1, + pp_size=2, + moe_cluster_size=1, + ) + + with pytest.raises(NotImplementedError, match="multi-group MoE topologies"): + executor._get_ep_health_stats() + + +def test_base_worker_exposes_engine_health_snapshot() -> None: + health = EPGroupHealth(4) + health.mark_failed(1) + worker = BaseWorker.__new__(BaseWorker) + worker.engine = _make_executor(health) + + assert worker.fetch_ep_health_stats()["failedRanks"] == [1] + + +def test_base_worker_reports_unsupported_engine_without_rpc_failure() -> None: + worker = BaseWorker.__new__(BaseWorker) + worker.engine = SimpleNamespace() + + assert worker.fetch_ep_health_stats() is None + + +def test_proxy_fetches_health_over_dedicated_rpc() -> None: + expected = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [1], + "generation": 1, + } + proxy = GenerationExecutorProxy.__new__(GenerationExecutorProxy) + proxy.rpc_client = MagicMock() + proxy.rpc_client.fetch_ep_health_stats.return_value.remote.return_value = expected + + assert proxy._get_ep_health_stats() == expected + proxy.rpc_client.fetch_ep_health_stats.return_value.remote.assert_called_once_with(timeout=1.0) + + +def test_proxy_propagates_transient_rpc_failure() -> None: + proxy = GenerationExecutorProxy.__new__(GenerationExecutorProxy) + proxy.rpc_client = MagicMock() + proxy.rpc_client.fetch_ep_health_stats.return_value.remote.side_effect = TimeoutError + + with pytest.raises(TimeoutError): + proxy._get_ep_health_stats() + + +def test_rpc_proxy_fetches_health_over_dedicated_rpc() -> None: + expected = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + proxy = GenerationExecutorRpcProxy.__new__(GenerationExecutorRpcProxy) + proxy.rpc_client = MagicMock() + proxy.rpc_client.fetch_ep_health_stats.return_value.remote.return_value = expected + + assert proxy._get_ep_health_stats(timeout=0.25) == expected + proxy.rpc_client.fetch_ep_health_stats.return_value.remote.assert_called_once_with(timeout=0.25) + + +def test_rpc_executor_mixin_forwards_health_for_ray_and_rpc_executors() -> None: + """Every RpcExecutorMixin consumer, including RayExecutor, gets telemetry.""" + expected = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + executor = SimpleNamespace(rpc_client=MagicMock()) + executor.rpc_client.fetch_ep_health_stats.return_value.remote.return_value = expected + + assert RpcExecutorMixin._get_ep_health_stats(executor, timeout=0.5) == expected + executor.rpc_client.fetch_ep_health_stats.return_value.remote.assert_called_once_with( + timeout=0.5 + ) + + +def test_ep_health_snapshot_crosses_real_rpc_and_llm_bridge() -> None: + """Exercise the passive committed-membership path without a model worker.""" + + class HealthRpcSurface: + _get_ep_health_stats = BaseWorker._get_ep_health_stats + fetch_ep_health_stats = BaseWorker.fetch_ep_health_stats + + health = EPGroupHealth(4) + health.mark_failed(2) + worker = HealthRpcSurface() + worker.engine = _make_executor(health) + address = get_unique_ipc_addr() + + with RPCServer(worker) as rpc_server: + rpc_server.bind(address) + rpc_server.start() + with RPCClient(address, hmac_key=rpc_server.hmac_key) as rpc_client: + proxy = GenerationExecutorProxy.__new__(GenerationExecutorProxy) + proxy.rpc_client = rpc_client + llm = BaseLLM.__new__(BaseLLM) + llm._executor = proxy + + assert llm._get_ep_health_stats() == { + "sourceEpoch": health.source_epoch, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [2], + "generation": 1, + } + + +def test_server_mount_metrics_registers_local_ep_health_collector(monkeypatch) -> None: + """The Prometheus endpoint registry includes the coherent local snapshot.""" + import prometheus_client + import prometheus_fastapi_instrumentator + from prometheus_client import multiprocess + from prometheus_client.core import GaugeMetricFamily + + class SnapshotCollector: + def collect(self): + metric = GaugeMetricFamily( + "trtllm_ep_health_available", + "Whether EP health telemetry is available.", + ) + metric.add_metric([], 1) + return [metric] + + metrics_collector = MagicMock() + metrics_collector.register_ep_health_metrics.side_effect = lambda registry: registry.register( + SnapshotCollector() + ) + monkeypatch.setattr(multiprocess, "MultiProcessCollector", MagicMock()) + instrumentator = MagicMock() + instrumentator.add.return_value = instrumentator + instrumentator.instrument.return_value = instrumentator + instrumentator.expose.return_value = instrumentator + monkeypatch.setattr( + prometheus_fastapi_instrumentator, "Instrumentator", MagicMock(return_value=instrumentator) + ) + monkeypatch.setattr(prometheus_client, "make_asgi_app", MagicMock(return_value=MagicMock())) + + server = OpenAIServer.__new__(OpenAIServer) + server.app = SimpleNamespace(routes=[]) + server.metrics_collector = metrics_collector + server.mount_metrics() + + registry = metrics_collector.register_ep_health_metrics.call_args.args[0] + output = prometheus_client.generate_latest(registry).decode() + assert "trtllm_ep_health_available 1.0" in output + assert len(server.app.routes) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "transient_error", + [ + TimeoutError("health read timed out"), + RPCTimeout("health RPC timed out"), + RPCCancelled("rank-0 worker is being replaced"), + RPCError("health transport failed", cause=ConnectionError("closed")), + zmq.ZMQError(zmq.ENOTCONN, "socket disconnected"), + RPCError( + "wrapped socket failure", + cause=zmq.ZMQError(zmq.ENOTCONN, "socket disconnected"), + ), + ], +) +async def test_server_retries_after_transient_first_read_failure( + transient_error: Exception, +) -> None: + expected = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [1], + "generation": 1, + } + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + server._ep_health_collector_task = None + server._ep_health_stop_event = asyncio.Event() + + results = iter((transient_error, expected)) + + def read_health(): + result = next(results) + if isinstance(result, Exception): + raise result + server._ep_health_stop_event.set() + return result + + health_reader = MagicMock(side_effect=read_health) + server.generator = SimpleNamespace(_get_ep_health_stats=health_reader) + server._wait_for_ep_health_stop = AsyncMock(side_effect=[False, True]) + + await server._start_ep_health_collector() + assert server._ep_health_collector_task is not None + await server._ep_health_collector_task + + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + server.metrics_collector.log_ep_health_stats.assert_called_once_with(expected) + assert health_reader.call_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "transient_error", + [ + TimeoutError("timed out"), + RPCCancelled("rank-0 worker is being replaced"), + zmq.ZMQError(zmq.ENOTCONN, "socket disconnected"), + RPCError( + "wrapped socket failure", + cause=zmq.ZMQError(zmq.ENOTCONN, "socket disconnected"), + ), + ], +) +async def test_server_loop_retries_transient_rank0_loss_and_recovers( + transient_error: Exception, +) -> None: + """Transient rank-0 loss clears availability, then accepts its replacement.""" + before_restart = { + "sourceEpoch": "source-before-restart", + "worldSize": 4, + "activeCount": 3, + "failedRanks": [1], + "generation": 2, + } + after_restart = { + "sourceEpoch": "source-after-restart", + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + server._ep_health_stop_event = asyncio.Event() + + results = iter((before_restart, transient_error, after_restart)) + + def read_health(): + result = next(results) + if isinstance(result, Exception): + raise result + return result + + health_reader = MagicMock(side_effect=read_health) + server.generator = SimpleNamespace(_get_ep_health_stats=health_reader) + server._wait_for_ep_health_stop = AsyncMock(side_effect=[False, False, True]) + + await asyncio.wait_for(server._ep_health_collector_loop(), timeout=1.0) + + assert health_reader.call_count == 3 + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + assert [ + call.args[0] for call in server.metrics_collector.log_ep_health_stats.call_args_list + ] == [before_restart, after_restart] + assert server._wait_for_ep_health_stop.await_count == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "read_result", + [ + None, + ValueError("invalid health configuration"), + RPCError("remote health configuration failed", cause=ValueError("invalid")), + RPCError( + "unsupported multi-group topology", + cause=NotImplementedError("multi-group"), + ), + ], +) +async def test_server_loop_stops_after_deterministic_health_loss(read_result) -> None: + """Unsupported and deterministic failures do not create a polling loop.""" + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + if isinstance(read_result, Exception): + health_reader = MagicMock(side_effect=read_result) + else: + health_reader = MagicMock(return_value=read_result) + server.generator = SimpleNamespace(_get_ep_health_stats=health_reader) + + await asyncio.wait_for(server._ep_health_collector_loop(), timeout=0.5) + + health_reader.assert_called_once_with() + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + server.metrics_collector.log_ep_health_stats.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "startup_error", + [ + ValueError("invalid health configuration"), + RPCError("remote health configuration failed", cause=ValueError("invalid")), + RPCError("method unavailable"), + ], +) +async def test_server_stops_after_deterministic_first_read_failure(startup_error) -> None: + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server._ep_health_collector_task = None + health_reader = MagicMock(side_effect=startup_error) + server.generator = SimpleNamespace(_get_ep_health_stats=health_reader) + + await server._start_ep_health_collector() + task = server._ep_health_collector_task + assert task is not None + await task + + health_reader.assert_called_once_with() + assert task.done() + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + server.metrics_collector.log_ep_health_stats.assert_not_called() + await server._stop_ep_health_collector() + assert server._ep_health_collector_task is None + + +@pytest.mark.asyncio +async def test_server_retries_after_collector_rejects_payload() -> None: + snapshot = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.side_effect = [False, True] + server._ep_health_stop_event = asyncio.Event() + + read_count = 0 + + def read_health(): + nonlocal read_count + read_count += 1 + if read_count == 2: + server._ep_health_stop_event.set() + return snapshot + + server.generator = SimpleNamespace(_get_ep_health_stats=MagicMock(side_effect=read_health)) + server._wait_for_ep_health_stop = AsyncMock(side_effect=[False, True]) + + await server._ep_health_collector_loop() + + assert server.generator._get_ep_health_stats.call_count == 2 + assert [ + args.args[0] for args in server.metrics_collector.log_ep_health_stats.call_args_list + ] == [snapshot, snapshot] + + +@pytest.mark.asyncio +async def test_server_retries_after_initial_payload_rejection() -> None: + snapshot = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.side_effect = [False, True] + server._ep_health_collector_task = None + server._ep_health_stop_event = asyncio.Event() + + read_count = 0 + + def read_health(): + nonlocal read_count + read_count += 1 + if read_count == 2: + server._ep_health_stop_event.set() + return snapshot + + server.generator = SimpleNamespace(_get_ep_health_stats=MagicMock(side_effect=read_health)) + server._wait_for_ep_health_stop = AsyncMock(side_effect=[False, True]) + + await server._start_ep_health_collector() + + assert server._ep_health_collector_task is not None + await server._ep_health_collector_task + assert server.generator._get_ep_health_stats.call_count == 2 + assert server.metrics_collector.log_ep_health_stats.call_count == 2 + + +@pytest.mark.asyncio +async def test_server_does_not_poll_unsupported_backend() -> None: + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server._ep_health_collector_task = None + server.generator = SimpleNamespace(_get_ep_health_stats=MagicMock(return_value=None)) + + await server._start_ep_health_collector() + task = server._ep_health_collector_task + assert task is not None + await task + + assert task.done() + server.generator._get_ep_health_stats.assert_called_once_with() + server.metrics_collector.log_ep_health_stats.assert_not_called() + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + await server._stop_ep_health_collector() + assert server._ep_health_collector_task is None + + +@pytest.mark.asyncio +async def test_server_retries_explicitly_pending_registration() -> None: + snapshot = { + "sourceEpoch": "source-after-attachment", + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + server = OpenAIServer.__new__(OpenAIServer) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + server._ep_health_collector_task = None + server._ep_health_stop_event = asyncio.Event() + server.generator = SimpleNamespace( + _get_ep_health_stats=MagicMock(side_effect=[pending_ep_health_metrics(), snapshot]) + ) + server._wait_for_ep_health_stop = AsyncMock(side_effect=[False, True]) + + await server._start_ep_health_collector() + assert server._ep_health_collector_task is not None + await server._ep_health_collector_task + + server.metrics_collector.log_ep_health_unavailable.assert_called_once_with() + server.metrics_collector.log_ep_health_stats.assert_called_once_with(snapshot) + + +@pytest.mark.asyncio +async def test_server_lifespan_starts_and_stops_health_collector(monkeypatch) -> None: + """The real server lifespan owns the health collector task.""" + generator = SimpleNamespace( + args=SimpleNamespace(enable_energy_metrics=False, enable_iter_perf_stats=False), + _get_ep_health_stats=MagicMock(), + shutdown=MagicMock(), + ) + monkeypatch.setattr(OpenAIServer, "_init_llm", lambda self, chat_template=None: None) + monkeypatch.setattr(OpenAIServer, "register_routes", lambda self: None) + server = OpenAIServer( + generator=generator, + model="test-model", + tool_parser=None, + server_role=None, + metadata_server_cfg=None, + ) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + collector_started = asyncio.Event() + + async def controlled_collector_loop() -> None: + collector_started.set() + await server._ep_health_stop_event.wait() + + server._ep_health_collector_loop = controlled_collector_loop + + async with server.app.router.lifespan_context(server.app): + await asyncio.wait_for(collector_started.wait(), timeout=1.0) + collector_task = server._ep_health_collector_task + assert collector_task is not None + assert not collector_task.done() + + assert collector_task.done() + assert not collector_task.cancelled() + generator._get_ep_health_stats.assert_not_called() + generator.shutdown.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_server_lifespan_body_failure_still_shuts_down_once(monkeypatch) -> None: + """An exception thrown at the lifespan yield cannot bypass cleanup.""" + generator = SimpleNamespace( + args=SimpleNamespace(enable_energy_metrics=False, enable_iter_perf_stats=False), + _get_ep_health_stats=MagicMock(return_value=None), + shutdown=MagicMock(), + ) + monkeypatch.setattr(OpenAIServer, "_init_llm", lambda self, chat_template=None: None) + monkeypatch.setattr(OpenAIServer, "register_routes", lambda self: None) + server = OpenAIServer( + generator=generator, + model="test-model", + tool_parser=None, + server_role=None, + metadata_server_cfg=None, + ) + server.metrics_collector = MagicMock() + + with pytest.raises(RuntimeError, match="lifespan body failed"): + async with server.app.router.lifespan_context(server.app): + raise RuntimeError("lifespan body failed") + + assert server._ep_health_collector_task is None + generator.shutdown.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_server_lifespan_drains_inflight_health_poll_before_shutdown(monkeypatch) -> None: + """Startup stays non-blocking, while shutdown drains the health RPC.""" + expected = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + poll_started = threading.Event() + release_poll = threading.Event() + + def read_health(): + poll_started.set() + release_poll.wait(timeout=5.0) + return expected + + health_reader = MagicMock(side_effect=read_health) + generator = SimpleNamespace( + args=SimpleNamespace(enable_energy_metrics=False, enable_iter_perf_stats=False), + _get_ep_health_stats=health_reader, + shutdown=MagicMock(), + ) + monkeypatch.setattr(OpenAIServer, "_init_llm", lambda self, chat_template=None: None) + monkeypatch.setattr(OpenAIServer, "register_routes", lambda self: None) + server = OpenAIServer( + generator=generator, + model="test-model", + tool_parser=None, + server_role=None, + metadata_server_cfg=None, + ) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + + lifespan = server.app.router.lifespan_context(server.app) + await asyncio.wait_for(lifespan.__aenter__(), timeout=1.0) + exit_task = None + try: + assert await asyncio.wait_for(asyncio.to_thread(poll_started.wait, 1.0), timeout=2.0) + exit_task = asyncio.create_task(lifespan.__aexit__(None, None, None)) + await asyncio.wait_for(server._ep_health_stop_event.wait(), timeout=1.0) + + generator.shutdown.assert_not_called() + assert not exit_task.done() + finally: + release_poll.set() + if exit_task is None: + exit_task = asyncio.create_task(lifespan.__aexit__(None, None, None)) + await asyncio.wait_for(exit_task, timeout=1.0) + + health_reader.assert_called_once_with() + generator.shutdown.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_server_lifespan_cancellation_still_drains_health_poll(monkeypatch) -> None: + """Cancelling lifespan cleanup cannot abandon a running health RPC thread.""" + snapshot = { + "sourceEpoch": EP_HEALTH_SOURCE_EPOCH, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + poll_started = threading.Event() + release_poll = threading.Event() + + def read_health(): + poll_started.set() + release_poll.wait(timeout=5.0) + return snapshot + + generator = SimpleNamespace( + args=SimpleNamespace(enable_energy_metrics=False, enable_iter_perf_stats=False), + _get_ep_health_stats=MagicMock(side_effect=read_health), + shutdown=MagicMock(), + ) + monkeypatch.setattr(OpenAIServer, "_init_llm", lambda self, chat_template=None: None) + monkeypatch.setattr(OpenAIServer, "register_routes", lambda self: None) + server = OpenAIServer( + generator=generator, + model="test-model", + tool_parser=None, + server_role=None, + metadata_server_cfg=None, + ) + server.metrics_collector = MagicMock() + server.metrics_collector.log_ep_health_stats.return_value = True + + lifespan = server.app.router.lifespan_context(server.app) + await asyncio.wait_for(lifespan.__aenter__(), timeout=1.0) + exit_task = None + try: + assert await asyncio.wait_for(asyncio.to_thread(poll_started.wait, 1.0), timeout=2.0) + exit_task = asyncio.create_task(lifespan.__aexit__(None, None, None)) + await asyncio.wait_for(server._ep_health_stop_event.wait(), timeout=1.0) + for _ in range(3): + exit_task.cancel() + await asyncio.sleep(0) + generator.shutdown.assert_not_called() + assert not exit_task.done() + finally: + release_poll.set() + + assert exit_task is not None + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(exit_task, timeout=1.0) + generator._get_ep_health_stats.assert_called_once_with() + generator.shutdown.assert_called_once_with() diff --git a/tests/unittest/metrics/test_collector.py b/tests/unittest/metrics/test_collector.py index e66975ca6f8e..efc0cd04f6fb 100644 --- a/tests/unittest/metrics/test_collector.py +++ b/tests/unittest/metrics/test_collector.py @@ -14,12 +14,14 @@ # limitations under the License. """Unit tests for MetricsCollector and process_req_perf_metrics.""" +import threading +import time from typing import Dict import pytest -from prometheus_client import REGISTRY +from prometheus_client import REGISTRY, CollectorRegistry, generate_latest, multiprocess -from tensorrt_llm.metrics.collector import MetricsCollector +from tensorrt_llm.metrics.collector import _MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS, MetricsCollector from tensorrt_llm.metrics.enums import MetricNames, RequestEventTiming from tensorrt_llm.metrics.perf_utils import process_req_perf_metrics @@ -47,6 +49,13 @@ def collector(): def _get_gauge_value(collector, metric_name: str): """Get the current value of a Prometheus gauge.""" + if metric_name.startswith("ep_"): + sample_name = f"trtllm_{metric_name}" + for metric in collector._ep_health_prometheus_collector.collect(): + for sample in metric.samples: + if sample.name == sample_name: + return sample.value + raise AssertionError(f"Missing EP health sample {sample_name}") metric = getattr(collector, metric_name) return metric.labels(**collector.labels)._value.get() @@ -57,6 +66,37 @@ def _get_counter_value(collector, metric_name: str): return metric.labels(**collector.labels)._value.get() +def _get_ep_rank_gauge_value(collector, rank: int): + for metric in collector._ep_health_prometheus_collector.collect(): + for sample in metric.samples: + if sample.name == "trtllm_ep_rank_active" and sample.labels[ + collector.labelname_ep_rank + ] == str(rank): + return sample.value + raise AssertionError(f"Missing EP rank sample {rank}") + + +def _get_ep_rank_samples(collector): + return [ + sample + for metric in collector._ep_health_prometheus_collector.collect() + for sample in metric.samples + if sample.name == "trtllm_ep_rank_active" + ] + + +def _get_ep_metric_samples(collector): + return [ + sample + for metric in collector._ep_health_prometheus_collector.collect() + for sample in metric.samples + ] + + +EP_HEALTH_SOURCE_A = "source-a" +EP_HEALTH_SOURCE_B = "source-b" + + SAMPLE_ITERATION_STATS = { "numActiveRequests": 5, "numQueuedRequests": 3, @@ -93,6 +133,491 @@ def _get_counter_value(collector, metric_name: str): } +class TestEPHealthStats: + """Test committed EPGroupHealth membership is exposed as passive gauges.""" + + def test_no_snapshot_exports_no_ep_health_families(self, collector): + registry = CollectorRegistry() + collector.register_ep_health_metrics(registry) + + assert _get_ep_metric_samples(collector) == [] + assert "trtllm_ep_" not in generate_latest(registry).decode() + + # ``None``/unsupported health is represented by marking a collector + # unavailable before any producer has published a valid snapshot. + collector.log_ep_health_unavailable() + assert _get_ep_metric_samples(collector) == [] + assert "trtllm_ep_" not in generate_latest(registry).decode() + + def test_per_rank_and_group_gauges(self, collector): + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [2], + "generation": 1, + } + ) + is True + ) + + assert [_get_ep_rank_gauge_value(collector, rank) for rank in range(4)] == [ + 1, + 1, + 0, + 1, + ] + assert _get_gauge_value(collector, "ep_active_ranks") == 3 + assert _get_gauge_value(collector, "ep_failed_ranks") == 1 + assert _get_gauge_value(collector, "ep_health_generation") == 1 + assert _get_gauge_value(collector, "ep_health_available") == 1 + + def test_reactivation_updates_existing_rank_series(self, collector): + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 1, + } + ) + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 2, + "failedRanks": [], + "generation": 2, + } + ) + + assert _get_ep_rank_gauge_value(collector, 1) == 1 + assert _get_gauge_value(collector, "ep_active_ranks") == 2 + assert _get_gauge_value(collector, "ep_failed_ranks") == 0 + assert _get_gauge_value(collector, "ep_health_generation") == 2 + + def test_conflicting_same_generation_warns_and_is_ignored(self, collector, caplog): + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 1, + } + ) + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [0], + "generation": 1, + } + ) + is False + ) + + assert "Conflicting epHealthStats payloads" in caplog.text + assert _get_ep_rank_gauge_value(collector, 0) == 1 + assert _get_ep_rank_gauge_value(collector, 1) == 0 + assert _get_gauge_value(collector, "ep_health_available") == 0 + + def test_stale_generation_warns_and_is_ignored(self, collector, caplog): + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 2, + } + ) + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [0], + "generation": 1, + } + ) + is False + ) + + assert "Ignoring stale epHealthStats generation" in caplog.text + assert _get_ep_rank_gauge_value(collector, 0) == 1 + assert _get_ep_rank_gauge_value(collector, 1) == 0 + assert _get_gauge_value(collector, "ep_health_available") == 0 + + def test_new_producer_epoch_can_restart_generation(self, collector): + assert collector.log_ep_health_stats( + { + "sourceEpoch": "producer-a", + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 5, + } + ) + collector.log_ep_health_unavailable() + + assert collector.log_ep_health_stats( + { + "sourceEpoch": "producer-b", + "worldSize": 2, + "activeCount": 2, + "failedRanks": [], + "generation": 0, + } + ) + assert [_get_ep_rank_gauge_value(collector, rank) for rank in range(2)] == [1, 1] + assert _get_gauge_value(collector, "ep_active_ranks") == 2 + assert _get_gauge_value(collector, "ep_health_generation") == 0 + assert _get_gauge_value(collector, "ep_health_available") == 1 + + def test_retired_producer_epoch_cannot_become_current_again(self, collector, caplog): + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 9, + } + ) + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_B, + "worldSize": 2, + "activeCount": 2, + "failedRanks": [], + "generation": 0, + } + ) + + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 9, + } + ) + is False + ) + assert "retired source epoch" in caplog.text + assert [_get_ep_rank_gauge_value(collector, rank) for rank in range(2)] == [1, 1] + assert _get_gauge_value(collector, "ep_health_generation") == 0 + assert _get_gauge_value(collector, "ep_health_available") == 0 + + def test_source_epoch_is_required_after_epoch_aware_payload(self, collector): + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 0, + } + ) + + assert ( + collector.log_ep_health_stats( + { + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 1, + } + ) + is False + ) + assert collector._last_ep_health_state[0] == EP_HEALTH_SOURCE_A + assert _get_gauge_value(collector, "ep_health_available") == 0 + + def test_source_epoch_churn_keeps_recent_replay_window_bounded(self, collector): + transitions = _MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS + 10 + for epoch in range(transitions): + assert collector.log_ep_health_stats( + { + "sourceEpoch": f"source-{epoch}", + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 0, + } + ) + + # Legitimate producer replacements continue after the bounded history + # fills, while recently retired producers are still rejected. + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": f"source-{transitions - 2}", + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 0, + } + ) + is False + ) + assert ( + len(collector._retired_ep_health_source_epochs) == _MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS + ) + assert ( + len(collector._retired_ep_health_source_epoch_order) + == _MAX_RETIRED_EP_HEALTH_SOURCE_EPOCHS + ) + assert _get_gauge_value(collector, "ep_health_available") == 0 + + assert collector.log_ep_health_stats( + { + "sourceEpoch": "source-0", + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 0, + } + ) + assert _get_gauge_value(collector, "ep_health_available") == 1 + + def test_oversized_world_is_rejected_before_rank_series_creation(self, collector): + assert _get_ep_rank_samples(collector) == [] + + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 129, + "activeCount": 129, + "failedRanks": [], + "generation": 0, + } + ) + is False + ) + assert _get_ep_metric_samples(collector) == [] + + def test_unrepresentable_generation_is_rejected_before_metric_writes(self, collector): + assert _get_ep_rank_samples(collector) == [] + + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 1, + "activeCount": 1, + "failedRanks": [], + "generation": 10**1000, + } + ) + is False + ) + assert collector._last_ep_health_state is None + assert _get_ep_metric_samples(collector) == [] + + def test_world_size_change_warns_and_is_ignored(self, collector, caplog): + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [3], + "generation": 1, + } + ) + assert ( + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 2, + "failedRanks": [], + "generation": 2, + } + ) + is False + ) + + assert "Ignoring epHealthStats world-size change" in caplog.text + assert _get_ep_rank_gauge_value(collector, 3) == 0 + assert _get_gauge_value(collector, "ep_active_ranks") == 3 + assert _get_gauge_value(collector, "ep_failed_ranks") == 1 + assert _get_gauge_value(collector, "ep_health_generation") == 1 + assert _get_gauge_value(collector, "ep_health_available") == 0 + + def test_invalid_payload_is_ignored(self, collector, caplog): + invalid = { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [4], + "generation": 1, + } + assert collector.log_ep_health_stats(invalid) is False + assert collector.log_ep_health_stats(invalid) is False + + assert "Ignoring invalid epHealthStats payload" in caplog.text + assert caplog.text.count("Ignoring invalid epHealthStats payload") == 1 + assert collector._last_ep_health_state is None + assert _get_ep_metric_samples(collector) == [] + + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 4, + "failedRanks": [], + "generation": 0, + } + ) + assert collector.log_ep_health_stats(invalid) is False + assert caplog.text.count("Ignoring invalid epHealthStats payload") == 2 + + def test_rpc_outage_marks_snapshot_unavailable_and_duplicate_recovers(self, collector): + snapshot = { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 2, + "failedRanks": [], + "generation": 0, + } + assert collector.log_ep_health_stats(snapshot) is True + collector.log_ep_health_unavailable() + + assert _get_ep_rank_gauge_value(collector, 0) == 1 + assert _get_gauge_value(collector, "ep_active_ranks") == 2 + assert _get_gauge_value(collector, "ep_failed_ranks") == 0 + assert _get_gauge_value(collector, "ep_health_generation") == 0 + assert _get_gauge_value(collector, "ep_health_available") == 0 + + assert collector.log_ep_health_stats(snapshot) is True + assert _get_gauge_value(collector, "ep_health_available") == 1 + + def test_health_metrics_register_on_scrape_registry(self, collector): + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 3, + } + ) + registry = CollectorRegistry() + collector.register_ep_health_metrics(registry) + output = generate_latest(registry).decode() + + assert 'trtllm_ep_rank_active{ep_rank="0",model_name="test_model"} 1.0' in output + assert 'trtllm_ep_rank_active{ep_rank="1",model_name="test_model"} 0.0' in output + assert 'trtllm_ep_active_ranks{model_name="test_model"} 1.0' in output + assert 'trtllm_ep_failed_ranks{model_name="test_model"} 1.0' in output + assert 'trtllm_ep_health_generation{model_name="test_model"} 3.0' in output + assert 'trtllm_ep_health_available{model_name="test_model"} 1.0' in output + + def test_multiprocess_registry_and_local_health_collector_do_not_duplicate( + self, collector, tmp_path, monkeypatch + ): + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 2, + "activeCount": 1, + "failedRanks": [1], + "generation": 4, + } + ) + + registry = CollectorRegistry() + multiprocess.MultiProcessCollector(registry) + collector.register_ep_health_metrics(registry) + output = generate_latest(registry).decode() + + assert output.count("# HELP trtllm_ep_rank_active") == 1 + assert output.count("trtllm_ep_rank_active{") == 2 + assert output.count("trtllm_ep_health_available{") == 1 + assert "pid=" not in "\n".join( + line for line in output.splitlines() if line.startswith("trtllm_ep_") + ) + + def test_concurrent_scrapes_observe_only_complete_health_snapshots(self, collector): + registry = CollectorRegistry() + collector.register_ep_health_metrics(registry) + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 3, + "failedRanks": [0], + "generation": 0, + } + ) + writer_done = threading.Event() + start = threading.Barrier(2) + writer_errors = [] + + def writer(): + try: + start.wait(timeout=10.0) + for generation in range(1, 501): + failed_ranks = [generation % 4] + assert collector.log_ep_health_stats( + { + "sourceEpoch": EP_HEALTH_SOURCE_A, + "worldSize": 4, + "activeCount": 3, + "failedRanks": failed_ranks, + "generation": generation, + } + ) + time.sleep(0) + except BaseException as error: + writer_errors.append(error) + finally: + writer_done.set() + + thread = threading.Thread(target=writer) + thread.start() + start.wait(timeout=10.0) + scrapes = 0 + while not writer_done.is_set() or scrapes == 0: + families = { + metric.name: metric + for metric in collector._ep_health_prometheus_collector.collect() + } + availability = families["trtllm_ep_health_available"].samples[0].value + if availability: + ranks = families["trtllm_ep_rank_active"].samples + active_count = families["trtllm_ep_active_ranks"].samples[0].value + failed_count = families["trtllm_ep_failed_ranks"].samples[0].value + generation = families["trtllm_ep_health_generation"].samples[0].value + failed_ranks = [ + int(sample.labels[collector.labelname_ep_rank]) + for sample in ranks + if not sample.value + ] + assert len(ranks) == 4 + assert sum(sample.value for sample in ranks) == active_count + assert len(ranks) - active_count == failed_count + assert failed_ranks == [int(generation) % 4] + scrapes += 1 + thread.join(timeout=10.0) + + assert not thread.is_alive() + assert writer_errors == [] + assert collector._last_ep_health_state[2] == 500 + + class TestIterationStatsTopLevel: """Test top-level iteration stats are correctly exposed as Prometheus metrics.""" From 2846880447f661e396321450dd219b10b50ada37 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 13:27:57 -0700 Subject: [PATCH 13/14] fix: bind AlltoAll mask to committed generation Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 99 ++++++++++++++++--- .../_torch/distributed/moe_alltoall.py | 23 +++-- .../communication/nvlink_one_sided.py | 17 +++- .../_torch/modules/test_alltoall_watchdog.py | 79 +++++++++++++-- 4 files changed, 184 insertions(+), 34 deletions(-) diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 166d4507e926..2b71aaefa257 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -46,6 +46,9 @@ UNKNOWN_COMPLETION_FLAG = -(2**63) _COMPLETION_FLAG_MASK = (1 << 32) - 1 _COMPLETION_FLAG_HALF_RANGE = 1 << 31 +_ACTIVE_RANK_MASK_WORD_BITS = 64 +_ACTIVE_RANK_MASK_WORDS = 2 +_ACTIVE_RANK_MASK_WORD_MASK = (1 << _ACTIVE_RANK_MASK_WORD_BITS) - 1 _WORKSPACE_WATCHDOG_STATE_KEY = "alltoall_watchdog_shared_state" _WORKSPACE_WATCHDOG_STATE_INIT_LOCK = threading.Lock() @@ -70,6 +73,18 @@ def read_completion_flags(self, phase: str) -> Sequence[int]: """Return ``ep_size`` flag values for ``phase``.""" +class EPGroupHealthSnapshotLike(Protocol): + """Atomic fields required from a committed-membership snapshot.""" + + @property + def mask(self) -> int: + """Return the committed active-rank bitmask.""" + + @property + def generation(self) -> int: + """Return the matching committed-membership generation.""" + + class EPGroupHealthLike(Protocol): """Read-only committed EP membership used by AlltoAll frontends.""" @@ -79,6 +94,22 @@ def get_mask(self) -> int: def get_mask_words(self) -> tuple[int, ...]: """Return the active-rank bitmask split into uint64 words.""" + def snapshot(self) -> EPGroupHealthSnapshotLike: + """Return one atomic committed mask and generation snapshot.""" + + +@dataclass(frozen=True) +class ActiveRankMaskSnapshot: + """One dispatch epoch's rank mask and optional committed generation. + + ``committed_generation`` is populated only when ``active_rank_mask`` was + read from the coordinator-owned EP health object. Explicit caller masks do + not inherit a generation from unrelated health state. + """ + + active_rank_mask: torch.Tensor | None + committed_generation: int | None + class CompletionFlagReadTimeout(TimeoutError): """Raised when the host watchdog cannot read completion flags in time.""" @@ -266,29 +297,75 @@ def read_current_flag_val(self) -> int: flag_val = flag_val.detach().cpu() return _normalize_completion_flag(int(flag_val.item())) - def active_rank_mask_tensor(self, active_rank_mask: torch.Tensor | None) -> torch.Tensor | None: - """Capture a caller override or the current committed mask. + def capture_active_rank_mask( + self, active_rank_mask: torch.Tensor | None + ) -> ActiveRankMaskSnapshot: + """Capture a caller override or one coherent committed-mask epoch. The returned tensor does not alias a caller-owned override, so it can - safely represent the mask used by one dispatch/combine pair. + safely represent the mask used by one dispatch/combine pair. When the + mask comes from ``health``, derive both fixed-width ABI words from one + atomic mask/generation snapshot. """ if active_rank_mask is not None: - return active_rank_mask.detach().clone() + return ActiveRankMaskSnapshot( + active_rank_mask=active_rank_mask.detach().clone(), + committed_generation=None, + ) if self._health is None: - return None - return torch.tensor(self._health.get_mask_words(), dtype=torch.uint64, device="cpu") + return ActiveRankMaskSnapshot( + active_rank_mask=None, + committed_generation=None, + ) + + health_snapshot = self._health.snapshot() + mask_words = tuple( + (health_snapshot.mask >> (_ACTIVE_RANK_MASK_WORD_BITS * index)) + & _ACTIVE_RANK_MASK_WORD_MASK + for index in range(_ACTIVE_RANK_MASK_WORDS) + ) + return ActiveRankMaskSnapshot( + active_rank_mask=torch.tensor( + mask_words, + dtype=torch.uint64, + device="cpu", + ), + committed_generation=health_snapshot.generation, + ) + + def active_rank_mask_tensor(self, active_rank_mask: torch.Tensor | None) -> torch.Tensor | None: + """Return only the tensor from :meth:`capture_active_rank_mask`. + + Dispatch frontends must retain the full snapshot so combine can verify + the committed generation before launching its kernel. + """ + return self.capture_active_rank_mask(active_rank_mask).active_rank_mask def active_rank_mask_for_combine( self, - dispatch_active_rank_mask: torch.Tensor | None, + dispatch_snapshot: ActiveRankMaskSnapshot, requested_active_rank_mask: torch.Tensor | None, ) -> torch.Tensor | None: - """Reuse the dispatch mask and reject a conflicting combine override. + """Validate the dispatch epoch and return its captured rank mask. - This guarantees one local mask snapshot across a dispatch/combine pair. - Atomic membership across layers and iterations remains the recovery - coordinator's responsibility. + A coordinator-owned generation change between dispatch and combine is + an invalid asynchronous commit, so fail closed before launching the + combine kernel. Explicit caller masks are independent of health and + therefore carry no committed generation. """ + dispatch_active_rank_mask = dispatch_snapshot.active_rank_mask + committed_generation = dispatch_snapshot.committed_generation + if committed_generation is not None: + if self._health is None: + raise RuntimeError("committed rank-mask snapshot has no EP health source") + current_generation = self._health.snapshot().generation + if current_generation != committed_generation: + raise RuntimeError( + "committed EP membership changed between dispatch and combine " + f"(generation {committed_generation} -> {current_generation}); " + "aborting collective epoch" + ) + if requested_active_rank_mask is None: return dispatch_active_rank_mask requested_mask = self.active_rank_mask_tensor(requested_active_rank_mask) diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index f1692d2473fc..2bf31ef680a4 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -17,8 +17,9 @@ from tensorrt_llm._mnnvl_utils import MnnvlMemory from tensorrt_llm._torch.alltoall_watchdog import ( DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, - DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, AlltoAllWatchdog, - AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, EPGroupHealthLike) + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, ActiveRankMaskSnapshot, + AlltoAllWatchdog, AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, + EPGroupHealthLike) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -31,7 +32,7 @@ class _A2AState: local_num_tokens: int | None = None combine_payload_offset: int | None = None eplb_gathered_stats: torch.Tensor | None = None - active_rank_mask: torch.Tensor | None = None + active_rank_mask_snapshot: ActiveRankMaskSnapshot | None = None class MoeAlltoAll: @@ -290,8 +291,9 @@ def dispatch(self, invalid_token_expert_id: If not None, set the token_selected_experts of the invalid tokens to this expert id. This is used to notify the MoE to skip these tokens for GroupGEMM. expert_id_payload_index: The index of token_selected_experts in the input_payloads. Must be provided if invalid_token_expert_id is not None. eplb_local_stats: (Optional) [num_experts] tensor containing local statistics for EPLB - active_rank_mask: Optional uint64 CPU tensor overriding committed membership for this dispatch. The - captured value is reused by combine. + active_rank_mask: Optional uint64 CPU tensor overriding committed membership for this dispatch. When + omitted, the committed mask and generation are captured together. Combine reuses that mask and + fails closed if the committed generation changes first. Returns: recv_tensors: List of tensors received, each has shape [ep_size, max_tokens_per_rank, payload_num_elements_per_token] @@ -306,8 +308,9 @@ def dispatch(self, 0 ) == self.eplb_stats_num_experts, "eplb_local_stats size must match eplb_stats_num_experts" - active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( + active_rank_mask_snapshot = self._watchdog_coordinator.capture_active_rank_mask( active_rank_mask) + active_rank_mask = active_rank_mask_snapshot.active_rank_mask recv_tensors, combine_payload_offset, eplb_gathered_stats = torch.ops.trtllm.moe_a2a_dispatch( token_selected_experts, input_payloads, @@ -331,7 +334,7 @@ def dispatch(self, self._state.local_num_tokens = token_selected_experts.size(0) self._state.combine_payload_offset = combine_payload_offset self._state.eplb_gathered_stats = eplb_gathered_stats - self._state.active_rank_mask = active_rank_mask + self._state.active_rank_mask_snapshot = active_rank_mask_snapshot self._state.phase = "dispatched" if invalid_token_expert_id is not None: @@ -365,7 +368,7 @@ def combine( payload_in_workspace: If True, 'payload' is a view into 'workspace' at 'combine_payload_offset' and no staging copy is needed. If False, the op stages 'payload' into the workspace region before combining. use_low_precision_combine: If True, quantize the combine payload to FP8 for NVLink transfer (halves NVLink bandwidth usage, output precision is preserved). active_rank_mask: Optional uint64 CPU tensor. If supplied, it must match the mask captured by dispatch - for this collective. + for this collective. A committed-generation change since dispatch aborts the collective epoch. Returns: combined_output: [local_num_tokens, num_elements_per_token] tensor of combined results @@ -373,8 +376,10 @@ def combine( assert self._state.phase == "dispatched", "combine called before a successful dispatch" assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" + active_rank_mask_snapshot = self._state.active_rank_mask_snapshot + assert active_rank_mask_snapshot is not None active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( - self._state.active_rank_mask, active_rank_mask) + active_rank_mask_snapshot, active_rank_mask) output = torch.ops.trtllm.moe_a2a_combine( payload, self._state.local_num_tokens, self.workspace, self.metainfo, runtime_max_tokens_per_rank, self.ep_rank, diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index 076e1f5e6b81..8f0b3c9486aa 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py @@ -33,6 +33,7 @@ from tensorrt_llm._torch.alltoall_watchdog import ( DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + ActiveRankMaskSnapshot, AlltoAllWatchdog, AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, @@ -425,7 +426,8 @@ def dispatch( all_rank_num_tokens: Token counts per rank [ep_size] use_dp_padding: Whether to use DP padding (optional) **kwargs: Strategy-specific arguments. ``active_rank_mask`` may override the committed membership - for dispatch; the captured value is reused by combine. + for dispatch. Without an override, the committed mask and generation are captured together; + combine reuses that mask and fails closed if the generation changes first. Returns: Tuple of (hidden_states, hidden_states_sf, token_selected_slots, token_final_scales) @@ -456,9 +458,10 @@ def dispatch( assert eplb_local_stats.size(0) == self.eplb_stats_num_experts, ( "eplb_local_stats size must match eplb_stats_num_experts" ) - active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( + active_rank_mask_snapshot = self._watchdog_coordinator.capture_active_rank_mask( kwargs.get("active_rank_mask") ) + active_rank_mask = active_rank_mask_snapshot.active_rank_mask recv_buffers, combine_payload_offset, eplb_gathered_stats = ( torch.ops.trtllm.moe_a2a_dispatch( @@ -484,7 +487,7 @@ def dispatch( self._dispatch_state["combine_payload_offset"] = int(combine_payload_offset) self._dispatch_state["local_num_tokens"] = token_selected_slots.size(0) self._dispatch_state["runtime_max_tokens_per_rank"] = runtime_max_tokens_per_rank - self._dispatch_state["active_rank_mask"] = active_rank_mask + self._dispatch_state["active_rank_mask_snapshot"] = active_rank_mask_snapshot self._dispatch_state["phase"] = "dispatched" # Extract results from recv_buffers @@ -548,7 +551,8 @@ def combine( Shape: [ep_size, max_tokens_per_rank, hidden_size] or [ep_size * max_tokens_per_rank, hidden_size] (will be reshaped) **kwargs: Strategy-specific arguments. If ``active_rank_mask`` is supplied, it must match the mask - captured by dispatch for this collective. + captured by dispatch for this collective. A committed-generation change since dispatch aborts + the collective epoch. Returns: Combined output tensor [local_num_tokens, hidden_size] @@ -583,8 +587,11 @@ def combine( raise ValueError( f"final_hidden_states must be 2D or 3D, got {final_hidden_states.dim()}D" ) + active_rank_mask_snapshot = self._dispatch_state.get("active_rank_mask_snapshot") + if not isinstance(active_rank_mask_snapshot, ActiveRankMaskSnapshot): + raise RuntimeError("combine called but dispatch rank-mask snapshot is missing") active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( - self._dispatch_state.get("active_rank_mask"), + active_rank_mask_snapshot, kwargs.get("active_rank_mask"), ) output = torch.ops.trtllm.moe_a2a_combine( diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index 931b0dd5e5a6..cdff07cbcea4 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -211,7 +211,7 @@ def test_wide_ep_ft_options_ignore_legacy_enable_flag(monkeypatch: pytest.Monkey assert timeout_s is None -def test_watchdog_coordinator_reuses_dispatch_mask_for_combine() -> None: +def test_watchdog_coordinator_reuses_committed_mask_when_generation_is_unchanged() -> None: health = EPGroupHealth(4) coordinator = AlltoAllWatchdogCoordinator( workspace_state={}, @@ -221,19 +221,80 @@ def test_watchdog_coordinator_reuses_dispatch_mask_for_combine() -> None: ep_rank=0, health=health, ) - dispatch_mask = coordinator.active_rank_mask_tensor(None) - assert dispatch_mask is not None + dispatch_snapshot = coordinator.capture_active_rank_mask(None) + assert dispatch_snapshot.active_rank_mask is not None + assert dispatch_snapshot.committed_generation == 0 - # Simulate a higher-layer commit at an invalid mid-collective point. The - # local dispatch/combine pair must keep using its dispatch snapshot. - health.mark_failed(2) - combine_mask = coordinator.active_rank_mask_for_combine(dispatch_mask, None) + combine_mask = coordinator.active_rank_mask_for_combine(dispatch_snapshot, None) - assert combine_mask is dispatch_mask + assert combine_mask is dispatch_snapshot.active_rank_mask assert combine_mask.tolist() == [0b1111, 0] + + +def test_watchdog_coordinator_fails_closed_on_committed_generation_change() -> None: + health = EPGroupHealth(4) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((4, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + dispatch_snapshot = coordinator.capture_active_rank_mask(None) + + health.mark_failed(2) + health.mark_active(2) + assert health.get_mask() == 0b1111 + + with pytest.raises( + RuntimeError, + match="committed EP membership changed between dispatch and combine", + ): + coordinator.active_rank_mask_for_combine(dispatch_snapshot, None) + + +def test_watchdog_coordinator_converts_atomic_snapshot_to_two_mask_words() -> None: + health = EPGroupHealth(72) + health.mark_failed(70) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((72, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + + dispatch_snapshot = coordinator.capture_active_rank_mask(None) + + assert dispatch_snapshot.committed_generation == 1 + assert dispatch_snapshot.active_rank_mask is not None + assert dispatch_snapshot.active_rank_mask.tolist() == [(1 << 64) - 1, 0xBF] + + +def test_watchdog_coordinator_explicit_mask_is_not_bound_to_health_generation() -> None: + health = EPGroupHealth(4) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((4, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + explicit_mask = torch.tensor([0b1101, 0], dtype=torch.uint64) + dispatch_snapshot = coordinator.capture_active_rank_mask(explicit_mask) + assert dispatch_snapshot.committed_generation is None + + health.mark_failed(2) + combine_mask = coordinator.active_rank_mask_for_combine(dispatch_snapshot, None) + + assert combine_mask is dispatch_snapshot.active_rank_mask + assert combine_mask.tolist() == explicit_mask.tolist() with pytest.raises(ValueError, match="mask captured at dispatch"): coordinator.active_rank_mask_for_combine( - dispatch_mask, + dispatch_snapshot, torch.tensor(health.get_mask_words(), dtype=torch.uint64), ) From 5a76856e8baf0a7a34cbd4bd7cd1028ca7112a5c Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 13:28:11 -0700 Subject: [PATCH 14/14] fix: use a distinct failure-detection state Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/distributed/communicator.py | 4 +- .../_torch/pyexecutor/ep_failure_broadcast.py | 206 ++++++++++++----- .../_ep_failure_broadcast_mpi_worker.py | 21 +- .../pyexecutor/test_ep_failure_broadcast.py | 207 +++++++++--------- 4 files changed, 268 insertions(+), 170 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 2ec47c0efc8a..4c7dbc63b7fd 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -460,7 +460,7 @@ def create_mpi_ft_subcomm( This operation is collective across ``parent_comm`` and must run on every rank during startup, before any background failure-broadcast threads are launched. The MVP requires one MoE EP group spanning the parent world, - ordered by the EP-local rank used by :class:`EPGroupHealth`. + ordered by the EP-local rank used by the caller's detected-rank state. Before splitting, ranks exchange their local validation outcome and EP topology on the healthy parent communicator. This prevents one rank from @@ -471,7 +471,7 @@ def create_mpi_ft_subcomm( parent_comm: Parent MPI communicator. Defaults to TRT-LLM's active ``mpi_comm()``, which wraps ``MPI.COMM_WORLD`` in the standard launch path and preserves custom communicator sessions. - health_size: Optional local ``EPGroupHealth`` size to validate + health_size: Optional local detected-rank-state size to validate collectively before splitting. Returns: diff --git a/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py b/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py index 3ac2a74b85e9..69821ae701b9 100644 --- a/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py +++ b/tensorrt_llm/_torch/pyexecutor/ep_failure_broadcast.py @@ -21,10 +21,10 @@ point-to-point requests on an independent CPU thread. The component propagates one rank-failure detection and updates only the -process-local *detected* :class:`EPGroupHealth`. It never updates the committed -execution mask. Multi-rank suspect/confirm consensus, commit authorization, -and communicator reconstruction are intentionally left to later WideEP FT -phases. +process-local :class:`DetectedRankState`. It never receives or updates the +committed ``EPGroupHealth`` execution mask. Multi-rank suspect/confirm +consensus, commit authorization, and communicator reconstruction are +intentionally left to later WideEP FT phases. Every survivor echoes a newly observed failure. Receiving one echo from every active survivor proves only that they propagated the same detection. Detection @@ -57,10 +57,6 @@ except ImportError: MPI = None -from tensorrt_llm._torch.modules.fused_moe.ep_group_health import ( - EPGroupHealth, - EPGroupHealthSnapshot, -) from tensorrt_llm.logger import logger if TYPE_CHECKING: @@ -133,6 +129,97 @@ def Query_thread(self) -> int: ... FailureReceivedCallback = FailureDetectedCallback +@dataclass(frozen=True) +class DetectedRankStateSnapshot: + """Atomic snapshot of rank-failure evidence for one communicator epoch.""" + + mask: int + failed_ranks: frozenset[int] + generation: int + + +class DetectedRankState: + """Thread-safe, process-local rank-failure evidence. + + This state is deliberately distinct from ``EPGroupHealth``. It is an + observation owned by :class:`MpiFtSubcomm`, not a committed execution mask, + and it is monotonic for the lifetime of one FT communicator epoch. + + Args: + ep_size: Number of EP-local ranks represented by this state. + + Raises: + ValueError: If ``ep_size`` is not positive. + """ + + def __init__(self, ep_size: int) -> None: + if ep_size <= 0: + raise ValueError(f"ep_size must be > 0, got {ep_size}") + self._ep_size = ep_size + self._all_active_mask = (1 << ep_size) - 1 + self._active_mask = self._all_active_mask + self._failed_ranks: set[int] = set() + self._generation = 0 + self._lock = threading.Lock() + + @property + def ep_size(self) -> int: + """Number of ranks represented by this immutable communicator epoch.""" + return self._ep_size + + @property + def generation(self) -> int: + """Number of effective detected-failure transitions.""" + with self._lock: + return self._generation + + def record_failure(self, rank: int) -> bool: + """Record one failed rank, returning whether evidence changed.""" + self._validate_rank(rank) + bit = 1 << rank + with self._lock: + if not self._active_mask & bit: + return False + self._active_mask &= ~bit + self._failed_ranks.add(rank) + self._generation += 1 + return True + + def is_active(self, rank: int) -> bool: + """Return whether no failure has been recorded for ``rank``.""" + self._validate_rank(rank) + with self._lock: + return bool(self._active_mask & (1 << rank)) + + def get_mask(self) -> int: + """Return the active-rank mask derived from detected evidence.""" + with self._lock: + return self._active_mask + + def get_failed_ranks(self) -> frozenset[int]: + """Return an immutable snapshot of ranks with failure evidence.""" + with self._lock: + return frozenset(self._failed_ranks) + + def all_active(self) -> bool: + """Return whether no rank failure has been detected.""" + with self._lock: + return self._active_mask == self._all_active_mask + + def snapshot(self) -> DetectedRankStateSnapshot: + """Return one coherent mask, failure-set, and generation snapshot.""" + with self._lock: + return DetectedRankStateSnapshot( + mask=self._active_mask, + failed_ranks=frozenset(self._failed_ranks), + generation=self._generation, + ) + + def _validate_rank(self, rank: int) -> None: + if not 0 <= rank < self._ep_size: + raise ValueError(f"rank must be in [0, {self._ep_size}), got {rank}") + + @dataclass(frozen=True) class MpiFtSubcommConfig: """Runtime-independent tuning for the FT progress thread. @@ -214,8 +301,8 @@ class MpiFtSubcomm: thread. Call :meth:`start` only after every rank has constructed the component. During normal operation, the progress thread owns MPI calls on this communicator; :meth:`report_detected_failure` updates only local - detected health and enqueues a report, so it never waits for a peer or MPI - progress. The watchdog must call this method instead of pre-marking health. + detected state and enqueues a report, so it never waits for a peer or MPI + progress. The watchdog must call this method instead of pre-recording evidence. Because ``MPI.THREAD_MULTIPLE`` is required, a caller whose start/stop deadline proves that thread is wedged may issue the terminal Revoke/Abort fallback. @@ -223,13 +310,14 @@ class MpiFtSubcomm: Args: mapping: Distributed mapping whose ``moe_ep_group`` defines the FT communicator and whose ``moe_ep_rank`` defines the local rank. - detected_health: Process-local EP detection state updated by sent and - received failure reports. This must be separate from the committed - execution-mask state owned by the recovery coordinator. + detected_state: Process-local EP detection evidence updated by sent and + received failure reports. ``EPGroupHealth`` is rejected because it + is the committed execution-mask state owned by the recovery + coordinator. config: Progress-loop timing configuration. on_failure_detected: Optional callback invoked as ``(failed_rank, source_rank, monotonic_time)`` when a local or - received report first changes detected health. This is the handoff + received report first changes detected state. This is the handoff to the future 1c.4b recovery coordinator; it is detection evidence, not commit authorization. Delivery runs on a dedicated daemon thread so callback latency cannot block MPI progress. @@ -240,14 +328,15 @@ class MpiFtSubcomm: Raises: RuntimeError: If MPI support or thread support is insufficient, or if the communicator does not match ``mapping``. - ValueError: If ``detected_health`` and ``mapping`` describe different + TypeError: If ``detected_state`` is not a :class:`DetectedRankState`. + ValueError: If ``detected_state`` and ``mapping`` describe different EP groups, or if both callback keyword names are provided. """ def __init__( self, mapping: Mapping, - detected_health: EPGroupHealth, + detected_state: DetectedRankState, config: MpiFtSubcommConfig | None = None, on_failure_detected: FailureDetectedCallback | None = None, *, @@ -255,6 +344,11 @@ def __init__( mpi_module: _MpiModule | None = None, on_failure_received: FailureReceivedCallback | None = None, ) -> None: + if not isinstance(detected_state, DetectedRankState): + raise TypeError( + "detected_state must be a DetectedRankState; committed " + "EPGroupHealth cannot be used as failure-detection evidence" + ) if on_failure_detected is not None and on_failure_received is not None: raise ValueError( "Specify only on_failure_detected; on_failure_received is a compatibility alias" @@ -279,7 +373,7 @@ def __init__( collective_setup = create_mpi_ft_subcomm( mapping, - health_size=detected_health.moe_world_size, + health_size=detected_state.ep_size, ) comm = cast(_MpiComm, collective_setup.comm) @@ -312,10 +406,10 @@ def __init__( "WideEP FT MVP requires one MoE EP group spanning the full MPI world; " f"got world_size={mapping.world_size}, moe_ep_group={ep_group}" ) - if detected_health.moe_world_size != mapping.moe_ep_size: + if detected_state.ep_size != mapping.moe_ep_size: raise ValueError( - "EPGroupHealth size must match mapping.moe_ep_size, " - f"got {detected_health.moe_world_size} and {mapping.moe_ep_size}" + "DetectedRankState size must match mapping.moe_ep_size, " + f"got {detected_state.ep_size} and {mapping.moe_ep_size}" ) self._local_rank = mapping.moe_ep_rank self._ep_size = mapping.moe_ep_size @@ -331,7 +425,7 @@ def __init__( f"expected rank={self._local_rank}, size={self._ep_size}" ) - self._detected_health = detected_health + self._detected_state = detected_state self._on_failure_detected = on_failure_detected self._outbound_reports: queue.SimpleQueue[_OutboundMessage] = queue.SimpleQueue() self._announced_failures: set[int] = set() @@ -610,8 +704,8 @@ def report_detected_failure(self, failed_rank: int) -> bool: """Record and asynchronously announce one detected EP-rank failure. This is the nonblocking seam invoked by the host watchdog. The - broadcaster owns the detected-health transition; callers must not mark - ``detected_health`` before invoking it. This method performs no MPI + broadcaster owns the detected-state transition; callers must not record + the failure before invoking it. This method performs no MPI calls and never waits for the progress thread. Repeated reports for the same rank are coalesced after the first local announcement. @@ -619,7 +713,7 @@ def report_detected_failure(self, failed_rank: int) -> bool: failed_rank: EP-local rank in ``[0, ep_size)``. Returns: - ``True`` if this call changed local detected health, else ``False``. + ``True`` if this call changed local detected state, else ``False``. """ self._validate_rank(failed_rank) with self._lifecycle_lock: @@ -634,7 +728,7 @@ def report_detected_failure(self, failed_rank: int) -> bool: self._lifecycle = _Lifecycle.FAILED self._request_terminal_abort(error, failed_rank, self._local_rank) raise - changed = self._detected_health.mark_failed(failed_rank) + changed = self._detected_state.record_failure(failed_rank) enqueued = self._observe_failure(failed_rank, self._local_rank) detected_time = time.monotonic() logger.warning( @@ -649,7 +743,7 @@ def pre_failover(self, failed_rank: int) -> bool: """Compatibility alias for :meth:`report_detected_failure`. Despite the historical name, this method does not commit failover and - callers must not pre-mark detected health. + callers must not pre-record detected evidence. """ return self.report_detected_failure(failed_rank) @@ -662,7 +756,7 @@ def world_is_poisoned(self) -> bool: return ( self._transport_poisoned.is_set() or self._accepted_failure_rank() is not None - or not self._detected_health.all_active() + or not self._detected_state.all_active() ) def failure_detection_is_reconciled(self, failed_rank: int) -> bool: @@ -680,7 +774,7 @@ def failure_detection_is_reconciled(self, failed_rank: int) -> bool: return False if self.last_error is not None or self._progress_failed.is_set(): return False - snapshot = self._detected_health.snapshot() + snapshot = self._detected_state.snapshot() if failed_rank not in snapshot.failed_ranks: return False accepted_failure, accepted_generation = self._accepted_failure_state() @@ -691,7 +785,7 @@ def failure_detection_is_reconciled(self, failed_rank: int) -> bool: return False return self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask) - def detected_health_is_reconciled(self) -> bool: + def detected_state_is_reconciled(self) -> bool: """Return whether survivors propagated every local failure detection. This is not a commit gate. A ``True`` result does not authorize request @@ -703,7 +797,7 @@ def detected_health_is_reconciled(self) -> bool: return False if self.last_error is not None or self._progress_failed.is_set(): return False - snapshot = self._detected_health.snapshot() + snapshot = self._detected_state.snapshot() accepted_failure, accepted_generation = self._accepted_failure_state() if not snapshot.failed_ranks: return accepted_failure is None and not self._transport_poisoned.is_set() @@ -717,13 +811,17 @@ def detected_health_is_reconciled(self) -> bool: for failed_rank in snapshot.failed_ranks ) + def detected_health_is_reconciled(self) -> bool: + """Compatibility alias for :meth:`detected_state_is_reconciled`.""" + return self.detected_state_is_reconciled() + def failure_is_reconciled(self, failed_rank: int) -> bool: """Compatibility alias for :meth:`failure_detection_is_reconciled`.""" return self.failure_detection_is_reconciled(failed_rank) def health_is_reconciled(self) -> bool: - """Compatibility alias for :meth:`detected_health_is_reconciled`.""" - return self.detected_health_is_reconciled() + """Compatibility alias for :meth:`detected_state_is_reconciled`.""" + return self.detected_state_is_reconciled() def _validate_rank(self, rank: int) -> None: if not 0 <= rank < self._ep_size: @@ -732,14 +830,14 @@ def _validate_rank(self, rank: int) -> None: def _claim_failure(self, failed_rank: int) -> None: """Atomically enforce the single-distinct-failure MVP contract.""" with self._failure_claim_lock: - snapshot = self._detected_health.snapshot() + snapshot = self._detected_state.snapshot() known_failures = snapshot.failed_ranks - conflicting_health = known_failures - {failed_rank} + conflicting_evidence = known_failures - {failed_rank} accepted = self._accepted_failed_rank - if (accepted is not None and accepted != failed_rank) or conflicting_health: + if (accepted is not None and accepted != failed_rank) or conflicting_evidence: first_failure = accepted if first_failure is None: - first_failure = min(conflicting_health) + first_failure = min(conflicting_evidence) raise RuntimeError( "WideEP FT MVP supports exactly one distinct failed rank; " f"already accepted rank {first_failure}, received rank {failed_rank}" @@ -747,15 +845,15 @@ def _claim_failure(self, failed_rank: int) -> None: if accepted is None: if failed_rank in known_failures: raise RuntimeError( - "WideEP FT detected health was mutated before the failure " + "WideEP FT detected state was mutated before the failure " "broadcaster accepted the report; the watchdog must call " - "report_detected_failure() without pre-marking health" + "report_detected_failure() without pre-recording evidence" ) self._accepted_failed_rank = failed_rank - # Exactly one effective mark_failed() must be the next detected - # health transition. Capturing that generation here prevents an - # independent mutation in the claim-to-observe window from - # becoming the accepted epoch baseline. + # Exactly one effective record_failure() must be the next + # detected-state transition. Capturing that generation here + # prevents an independent mutation in the claim-to-observe + # window from becoming the accepted epoch baseline. self._accepted_failure_generation = snapshot.generation + 1 def _accepted_failure_rank(self) -> int | None: @@ -769,7 +867,7 @@ def _accepted_failure_state(self) -> tuple[int | None, int | None]: def _observe_failure(self, failed_rank: int, reporter: int) -> bool: """Record one detection report and enqueue this rank's echo once.""" now = time.monotonic() - snapshot = self._detected_health.snapshot() + snapshot = self._detected_state.snapshot() message: _OutboundMessage | None = None with self._protocol_lock: unattributed_deadline = self._unattributed_error_deadlines.get(failed_rank) @@ -794,7 +892,7 @@ def _observe_failure(self, failed_rank: int, reporter: int) -> bool: def _detection_reconciliation_state_is_valid_locked( self, - snapshot: EPGroupHealthSnapshot, + snapshot: DetectedRankStateSnapshot, accepted_failure: int | None, accepted_generation: int | None, ) -> bool: @@ -822,7 +920,7 @@ def _request_terminal_abort( ``MPI_Abort``. This method performs no MPI calls. """ self._record_error(error) - snapshot = self._detected_health.snapshot() + snapshot = self._detected_state.snapshot() enqueue_abort = False now = time.monotonic() with self._protocol_lock: @@ -977,7 +1075,7 @@ def _drain_outbound_reports(self, pending_sends: list[_PendingSend]) -> None: return failed_rank = message.failed_rank - snapshot = self._detected_health.snapshot() + snapshot = self._detected_state.snapshot() if message.kind is _MessageKind.ABORT: # Terminal state must reach every rank still believed active. # If one of them is actually the second failed rank, send @@ -1095,7 +1193,7 @@ def _progress_receives(self, pending_receives: dict[int, _PendingReceive]) -> No f"{failed_rank} from EP rank {source}" ) - if self._detected_health.is_active(source) and not self._stop_event.is_set(): + if self._detected_state.is_active(source) and not self._stop_event.is_set(): try: pending_receives[source] = self._post_receive(source) except Exception as error: @@ -1158,7 +1256,7 @@ def _accept_received_failure( self._lifecycle = _Lifecycle.FAILED claim_error = error else: - changed = self._detected_health.mark_failed(failed_rank) + changed = self._detected_state.record_failure(failed_rank) if transport_poisoned: self._transport_poisoned.set() self._observe_failure(failed_rank, reporter) @@ -1304,21 +1402,21 @@ def _track_unattributed_transport_error(self, source: int) -> None: self._deadline_monitor_wake_event.set() def _check_accepted_failure_epoch(self) -> None: - """Fail closed when detected health leaves the accepted failure epoch.""" + """Fail closed when detected state leaves the accepted failure epoch.""" error: RuntimeError | None = None accepted_failure: int | None = None with self._lifecycle_lock: if self._lifecycle not in (_Lifecycle.RUNNING, _Lifecycle.STOPPING): return - snapshot = self._detected_health.snapshot() + snapshot = self._detected_state.snapshot() accepted_failure, accepted_generation = self._accepted_failure_state() if accepted_failure is None: if not snapshot.failed_ranks: return error = RuntimeError( - "WideEP FT detected health changed outside the failure " + "WideEP FT detected state changed outside the failure " "broadcaster; the watchdog must call " - "report_detected_failure() without pre-marking health " + "report_detected_failure() without pre-recording evidence " f"(snapshot={snapshot})" ) elif ( @@ -1375,7 +1473,7 @@ def _check_reconciliation_deadlines(self) -> None: with self._lifecycle_lock: if self._lifecycle not in (_Lifecycle.RUNNING, _Lifecycle.STOPPING): return - snapshot = self._detected_health.snapshot() + snapshot = self._detected_state.snapshot() with self._protocol_lock: for failed_rank, deadline in list(self._reconcile_deadlines.items()): if self._failure_detection_is_reconciled_locked(failed_rank, snapshot.mask): @@ -1392,7 +1490,7 @@ def _check_reconciliation_deadlines(self) -> None: def _stop_reconciliation_state(self) -> tuple[bool, BaseException | None]: """Return whether shutdown can stop MPI progress without splitting ranks.""" - snapshot = self._detected_health.snapshot() + snapshot = self._detected_state.snapshot() failed_ranks = snapshot.failed_ranks accepted_failure, accepted_generation = self._accepted_failure_state() with self._protocol_lock: diff --git a/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py b/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py index ccbff841fbcd..d9c9bd732443 100644 --- a/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py +++ b/tests/unittest/_torch/pyexecutor/_ep_failure_broadcast_mpi_worker.py @@ -30,8 +30,11 @@ import numpy as np from mpi4py import MPI -from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth -from tensorrt_llm._torch.pyexecutor.ep_failure_broadcast import MpiFtSubcomm, MpiFtSubcommConfig +from tensorrt_llm._torch.pyexecutor.ep_failure_broadcast import ( + DetectedRankState, + MpiFtSubcomm, + MpiFtSubcommConfig, +) from tensorrt_llm.mapping import MpiTopology _SKIP_MARKER = "WIDEEP_MPI_SMOKE_SKIP:" @@ -135,7 +138,7 @@ def _run_invalid_topology_smoke(world: MPI.Intracomm) -> str | None: try: unexpected_broadcaster = MpiFtSubcomm( mapping, - EPGroupHealth(world_size), + DetectedRankState(world_size), MpiFtSubcommConfig(startup_timeout_sec=5.0, stop_timeout_sec=5.0), ) except Exception as error: @@ -176,7 +179,7 @@ def _run_healthy_lifecycle_smoke(world: MPI.Intracomm) -> str | None: # collective validation and Split, and owns MPI progress directly. broadcaster = MpiFtSubcomm( mapping, - EPGroupHealth(world_size), + DetectedRankState(world_size), config, ) ft_comm = broadcaster._comm @@ -220,7 +223,7 @@ def _run_failure_propagation_smoke(world: MPI.Intracomm) -> str | None: rank = world.Get_rank() world_size = world.Get_size() mapping = _mapping(world_size, rank) - health = EPGroupHealth(world_size) + health = DetectedRankState(world_size) owner: MpiFtSubcomm | None = None broadcaster: MpiFtSubcomm | None = None detector_rank = 0 @@ -239,7 +242,7 @@ def _run_failure_propagation_smoke(world: MPI.Intracomm) -> str | None: # view so this smoke deterministically exercises typed Isend/Irecv+Test. owner = MpiFtSubcomm( mapping, - EPGroupHealth(world_size), + DetectedRankState(world_size), config, ) ft_comm = owner._comm @@ -297,7 +300,7 @@ def _run_failure_propagation_smoke(world: MPI.Intracomm) -> str | None: local_converged = rank == failure_rank or ( not health.is_active(failure_rank) and broadcaster.failure_detection_is_reconciled(failure_rank) - and broadcaster.detected_health_is_reconciled() + and broadcaster.detected_state_is_reconciled() and broadcaster.world_is_poisoned() ) except Exception: @@ -404,7 +407,7 @@ def _run_terminal_abort_smoke(world: MPI.Intracomm) -> str | None: # scenario specifically validates the bounded ABORT echo fallback. owner = MpiFtSubcomm( mapping, - EPGroupHealth(world_size), + DetectedRankState(world_size), MpiFtSubcommConfig(startup_timeout_sec=5.0, stop_timeout_sec=5.0), ) ft_comm = owner._comm @@ -412,7 +415,7 @@ def _run_terminal_abort_smoke(world: MPI.Intracomm) -> str | None: wrapped_comm = _NonUlfmComm(ft_comm) broadcaster = MpiFtSubcomm( mapping, - EPGroupHealth(world_size), + DetectedRankState(world_size), MpiFtSubcommConfig( startup_timeout_sec=5.0, stop_timeout_sec=5.0, diff --git a/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py index 8ab5975cd934..7e1cab9b5d30 100644 --- a/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py +++ b/tests/unittest/_torch/pyexecutor/test_ep_failure_broadcast.py @@ -37,7 +37,11 @@ from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth from tensorrt_llm._torch.pyexecutor import ep_failure_broadcast -from tensorrt_llm._torch.pyexecutor.ep_failure_broadcast import MpiFtSubcomm, MpiFtSubcommConfig +from tensorrt_llm._torch.pyexecutor.ep_failure_broadcast import ( + DetectedRankState, + MpiFtSubcomm, + MpiFtSubcommConfig, +) _TEST_CONFIG = MpiFtSubcommConfig( poll_interval_sec=0.001, @@ -309,34 +313,50 @@ def _make_broadcaster( *, size: int = 4, rank: int = 0, - detected_health: EPGroupHealth | None = None, + detected_state: DetectedRankState | None = None, comm: _FakeComm | None = None, mpi: _FakeMPI | None = None, callback: Callable[[int, int, float], None] | None = None, -) -> tuple[MpiFtSubcomm, EPGroupHealth, _FakeComm, _FakeMPI]: - detected_health = detected_health or EPGroupHealth(size) +) -> tuple[MpiFtSubcomm, DetectedRankState, _FakeComm, _FakeMPI]: + detected_state = detected_state or DetectedRankState(size) comm = comm or _FakeComm(size=size, rank=rank) mpi = mpi or _FakeMPI() broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(size, rank), - detected_health, + detected_state, _TEST_CONFIG, callback, comm=comm, mpi_module=mpi, ) - return broadcaster, detected_health, comm, mpi + return broadcaster, detected_state, comm, mpi + + +def test_committed_ep_group_health_is_rejected_before_mpi_setup() -> None: + comm = _FakeComm(size=2) + + with pytest.raises(TypeError, match="committed EPGroupHealth cannot be used"): + MpiFtSubcomm( + _FakeMapping.ep_world(2), + EPGroupHealth(2), + _TEST_CONFIG, + comm=comm, + mpi_module=_FakeMPI(), + ) + + assert comm.errhandler_calls == 0 + assert comm.is_revoked_calls == 0 def test_watchdog_callback_records_and_fans_out_detected_failure() -> None: - detected_health = EPGroupHealth(4) - broadcaster, _, comm, _ = _make_broadcaster(detected_health=detected_health) + detected_state = DetectedRankState(4) + broadcaster, _, comm, _ = _make_broadcaster(detected_state=detected_state) broadcaster.start() caller_thread = threading.get_ident() try: watchdog_callback = broadcaster.report_detected_failure assert watchdog_callback(2) is True - assert detected_health.get_failed_ranks() == frozenset({2}) + assert detected_state.get_failed_ranks() == frozenset({2}) _wait_until(lambda: len(comm.sends) == 2, description="failure fanout") _wait_until( lambda: all(record.request.test_calls > 0 for record in comm.sends), @@ -359,19 +379,19 @@ def test_watchdog_callback_records_and_fans_out_detected_failure() -> None: comm.deliver(source=1, failed_rank=2) comm.deliver(source=3, failed_rank=2) _wait_until( - broadcaster.detected_health_is_reconciled, + broadcaster.detected_state_is_reconciled, description="detection reconciliation", ) finally: broadcaster.stop() -def test_externally_premarked_detected_health_fails_closed() -> None: - detected_health = EPGroupHealth(2) - assert detected_health.mark_failed(1) is True +def test_externally_premarked_detected_state_fails_closed() -> None: + detected_state = DetectedRankState(2) + assert detected_state.record_failure(1) is True broadcaster, _, comm, _ = _make_broadcaster( size=2, - detected_health=detected_health, + detected_state=detected_state, ) try: @@ -383,9 +403,9 @@ def test_externally_premarked_detected_health_fails_closed() -> None: pass _wait_until( lambda: isinstance(broadcaster.last_error, RuntimeError), - description="pre-marked detected-health rejection", + description="pre-recorded detected-state rejection", ) - assert "without pre-marking health" in str(broadcaster.last_error) + assert "without pre-recording evidence" in str(broadcaster.last_error) _wait_until(lambda: comm.revoke_calls == 1, description="pre-mark fail-closed revoke") finally: broadcaster.stop() @@ -393,12 +413,12 @@ def test_externally_premarked_detected_health_fails_closed() -> None: def test_detection_reconciliation_does_not_mutate_committed_health() -> None: callbacks: list[tuple[int, int, float]] = [] - detected_health = EPGroupHealth(2) + detected_state = DetectedRankState(2) committed_health = EPGroupHealth(2) initial_committed_snapshot = committed_health.snapshot() broadcaster, _, _, _ = _make_broadcaster( size=2, - detected_health=detected_health, + detected_state=detected_state, callback=lambda failed, source, when: callbacks.append((failed, source, when)), ) broadcaster.start() @@ -409,12 +429,12 @@ def test_detection_reconciliation_does_not_mutate_committed_health() -> None: description="local detection coordinator handoff", ) _wait_until( - broadcaster.detected_health_is_reconciled, + broadcaster.detected_state_is_reconciled, description="detection reconciliation", ) - assert detected_health.get_failed_ranks() == frozenset({1}) - assert detected_health.generation == 1 + assert detected_state.get_failed_ranks() == frozenset({1}) + assert detected_state.generation == 1 assert committed_health.snapshot() == initial_committed_snapshot assert committed_health.all_active() is True assert callbacks[0][:2] == (1, 0) @@ -454,77 +474,54 @@ def test_survivor_echoes_reconcile_failure_detection() -> None: broadcaster.stop() -def test_rank_reactivation_cannot_reopen_reconciliation_for_same_comm_epoch() -> None: - broadcaster, health, comm, _ = _make_broadcaster(size=3) - broadcaster.start() - try: - assert broadcaster.pre_failover(2) is True - _wait_until(lambda: len(comm.sends) == 1, description="failure fanout") - comm.sends[0].request.complete() - comm.deliver(source=1, failed_rank=2) - _wait_until(broadcaster.health_is_reconciled, description="failure reconciliation") - - assert health.mark_active(2) is True - _wait_until( - lambda: isinstance(broadcaster.last_error, RuntimeError), - description="reactivated-rank terminal handling", - ) - _wait_until(lambda: comm.revoke_calls == 1, description="epoch-violation revoke") - assert broadcaster.world_is_poisoned() is True - assert broadcaster.failure_is_reconciled(2) is False - assert broadcaster.health_is_reconciled() is False - - # Even restoring the failed bit cannot restore the old consensus: the - # intervening generation change belongs to a future communicator epoch. - assert health.mark_failed(2) is True - assert broadcaster.world_is_poisoned() is True - assert broadcaster.failure_is_reconciled(2) is False - assert broadcaster.health_is_reconciled() is False - finally: - broadcaster.stop() - - -def test_reactivation_between_failure_claim_and_observe_cannot_become_epoch_baseline( - monkeypatch: pytest.MonkeyPatch, -) -> None: - health = EPGroupHealth(2) - broadcaster, _, _, _ = _make_broadcaster(size=2, detected_health=health) - original_claim = broadcaster._claim_failure - mutated = False +def test_detected_state_is_monotonic_for_one_communicator_epoch() -> None: + state = DetectedRankState(3) - def claim_then_turn_over_epoch(failed_rank: int) -> None: - nonlocal mutated - original_claim(failed_rank) - if not mutated: - mutated = True - assert health.mark_failed(failed_rank) is True - assert health.mark_active(failed_rank) is True - assert health.mark_failed(failed_rank) is True - - monkeypatch.setattr(broadcaster, "_claim_failure", claim_then_turn_over_epoch) - broadcaster.start() - try: - assert broadcaster.broadcast_failure(1) is False - _wait_until( - lambda: 1 in broadcaster._failure_fanout_complete, - description="single-survivor fanout completion", - ) - assert health.generation == 3 - _wait_until( - lambda: isinstance(broadcaster.last_error, RuntimeError), - description="turnover terminal handling", - ) - assert broadcaster.failure_is_reconciled(1) is False - assert broadcaster.health_is_reconciled() is False - finally: - broadcaster.stop() + assert state.snapshot() == ep_failure_broadcast.DetectedRankStateSnapshot( + mask=0b111, + failed_ranks=frozenset(), + generation=0, + ) + assert state.record_failure(2) is True + assert state.record_failure(2) is False + assert state.get_mask() == 0b011 + assert state.get_failed_ranks() == frozenset({2}) + assert state.generation == 1 + assert not hasattr(state, "mark_active") + + +def test_detected_state_coalesces_concurrent_evidence() -> None: + state = DetectedRankState(4) + barrier = threading.Barrier(8) + results: list[bool] = [] + results_lock = threading.Lock() + + def record_same_failure() -> None: + barrier.wait() + changed = state.record_failure(2) + with results_lock: + results.append(changed) + + threads = [threading.Thread(target=record_same_failure) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert results.count(True) == 1 + assert results.count(False) == 7 + assert state.snapshot() == ep_failure_broadcast.DetectedRankStateSnapshot( + mask=0b1011, + failed_ranks=frozenset({2}), + generation=1, + ) def test_second_failure_between_claim_and_observe_cannot_open_single_failure_gate( monkeypatch: pytest.MonkeyPatch, ) -> None: - health = EPGroupHealth(3) - broadcaster, _, _, _ = _make_broadcaster(size=3, detected_health=health) + health = DetectedRankState(3) + broadcaster, _, _, _ = _make_broadcaster(size=3, detected_state=health) original_claim = broadcaster._claim_failure mutated = False @@ -533,7 +530,7 @@ def claim_then_add_second_failure(failed_rank: int) -> None: original_claim(failed_rank) if not mutated: mutated = True - assert health.mark_failed(1) is True + assert health.record_failure(1) is True monkeypatch.setattr(broadcaster, "_claim_failure", claim_then_add_second_failure) broadcaster.start() @@ -546,7 +543,7 @@ def claim_then_add_second_failure(failed_rank: int) -> None: assert health.get_failed_ranks() == frozenset({1, 2}) _wait_until( lambda: isinstance(broadcaster.last_error, RuntimeError), - description="second-health-failure terminal handling", + description="second-detected-failure terminal handling", ) assert broadcaster.failure_is_reconciled(2) is False assert broadcaster.health_is_reconciled() is False @@ -568,7 +565,7 @@ def test_reconciliation_timeout_fails_closed_while_mpi_test_is_wedged() -> None: ) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(3), - EPGroupHealth(3), + DetectedRankState(3), config, comm=comm, mpi_module=_FakeMPI(), @@ -609,7 +606,7 @@ def test_reconciliation_timeout_world_aborts_when_mpi_test_is_wedged_without_ulf ) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(3), - EPGroupHealth(3), + DetectedRankState(3), config, comm=comm, mpi_module=_FakeMPI(), @@ -728,7 +725,7 @@ def test_monitor_does_not_world_abort_reconciled_relay_when_mpi_test_resumes() - ) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(2), - EPGroupHealth(2), + DetectedRankState(2), config, comm=comm, mpi_module=_FakeMPI(), @@ -817,7 +814,7 @@ def test_no_ulfm_abort_timeout_uses_world_abort_before_stop_completes() -> None: comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(3), - EPGroupHealth(3), + DetectedRankState(3), config, comm=comm, mpi_module=_FakeMPI(), @@ -856,7 +853,7 @@ def blocking_send(_destination: int) -> _FakeRequest: ) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(3), - EPGroupHealth(3), + DetectedRankState(3), config, comm=comm, mpi_module=_FakeMPI(), @@ -1001,7 +998,7 @@ def test_received_report_is_idempotent_and_callback_runs_once() -> None: broadcaster.stop() -def test_monitor_cannot_observe_remote_claim_before_detected_health_update( +def test_monitor_cannot_observe_remote_claim_before_detected_state_update( monkeypatch: pytest.MonkeyPatch, ) -> None: broadcaster, health, comm, _ = _make_broadcaster(size=3) @@ -1099,7 +1096,7 @@ def test_single_rank_group_has_no_peer_requests_or_sends() -> None: abort_timeout_sec=0.02, ) comm = _FakeComm(size=1, ulfm_probe_error=NotImplementedError()) - health = EPGroupHealth(1) + health = DetectedRankState(1) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(1), health, @@ -1315,7 +1312,7 @@ def blocking_receive(_source: int) -> _FakeRequest: ) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(2), - EPGroupHealth(2), + DetectedRankState(2), config, comm=comm, mpi_module=_FakeMPI(), @@ -1372,7 +1369,7 @@ def blocking_receive(_source: int) -> _FakeRequest: ) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(2), - EPGroupHealth(2), + DetectedRankState(2), config, lambda _failed, _source, _when: None, comm=comm, @@ -1498,7 +1495,7 @@ def test_healthy_stop_cleanup_timeout_aborts_world_and_retains_request() -> None ) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(2), - EPGroupHealth(2), + DetectedRankState(2), config, comm=comm, mpi_module=_FakeMPI(), @@ -1637,7 +1634,7 @@ def test_constructor_requires_mpi_thread_multiple() -> None: with pytest.raises(RuntimeError, match="MPI.THREAD_MULTIPLE"): MpiFtSubcomm( _FakeMapping.ep_world(2), - EPGroupHealth(2), + DetectedRankState(2), _TEST_CONFIG, comm=comm, mpi_module=mpi, @@ -1696,7 +1693,7 @@ def create_mpi_ft_subcomm( try: broadcaster = MpiFtSubcomm( mapping, - EPGroupHealth(2), + DetectedRankState(2), _TEST_CONFIG, mpi_module=mpi, ) @@ -1728,7 +1725,7 @@ def test_constructor_rejects_non_world_spanning_ep_group_for_mvp() -> None: with pytest.raises(ValueError, match="spanning the full MPI world"): MpiFtSubcomm( mapping, - EPGroupHealth(2), + DetectedRankState(2), _TEST_CONFIG, comm=comm, mpi_module=_FakeMPI(), @@ -1867,7 +1864,7 @@ def test_remote_revoke_after_local_reconciliation_still_fails_closed() -> None: assert comm.revoke_calls == 1 with pytest.raises(RuntimeError, match="not running"): broadcaster.pre_failover(1) - assert broadcaster._detected_health.get_failed_ranks() == frozenset({2}) + assert broadcaster._detected_state.get_failed_ranks() == frozenset({2}) broadcaster.stop() @@ -1982,7 +1979,7 @@ def test_non_ulfm_generic_error_after_watchdog_report_does_not_arm_stale_deadlin comm = _FakeComm(size=3, ulfm_probe_error=NotImplementedError()) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(3), - EPGroupHealth(3), + DetectedRankState(3), config, comm=comm, mpi_module=_FakeMPI(), @@ -2030,7 +2027,7 @@ def test_unattributed_error_from_other_source_keeps_reconciliation_gate_closed() comm = _FakeComm(size=4, ulfm_probe_error=NotImplementedError()) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(4), - EPGroupHealth(4), + DetectedRankState(4), config, comm=comm, mpi_module=_FakeMPI(), @@ -2084,7 +2081,7 @@ def test_non_ulfm_unattributed_receive_error_aborts_after_bounded_deadlines( abort_timeout_sec=0.05, ) comm = _FakeComm(size=2, ulfm_probe_error=NotImplementedError()) - health = EPGroupHealth(2) + health = DetectedRankState(2) broadcaster = MpiFtSubcomm( _FakeMapping.ep_world(2), health,