diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index f958ee6b9704..cbf1cbe770c2 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -19,6 +19,7 @@ import platform import sys from dataclasses import dataclass +from enum import Enum from typing import Any, List, Optional, Union import pynvml @@ -52,17 +53,27 @@ def _check_cu_result(cu_func_ret): return None +class _MnnvlAllocationState(Enum): + MAPPED = "mapped" + PREPARING = "preparing" + UNMAPPED = "unmapped" + RESTORING = "restoring" + BROKEN = "broken" + + @dataclass class _MnnvlAllocationRecord: comm: Any comm_size: int comm_rank: int + comm_membership: tuple[int, ...] aligned_size: int mem_handles: List[Any] start_address: int rank_stride: int address_offset: int - mapped: bool = True + state: _MnnvlAllocationState = _MnnvlAllocationState.MAPPED + pending_comm: Any = None class MnnvlMemory: @@ -114,8 +125,8 @@ def __del__(self): @property def mapped(self) -> bool: - """Whether physical handles are mapped into this VA reservation.""" - return type(self).allocated_map[self.ptr].mapped + """Whether the allocation is mapped and ready for data-path access.""" + return type(self).allocated_map[self.ptr].state is _MnnvlAllocationState.MAPPED def as_torch_strided_tensor(self, dtype): num_segments = type(self).comm.Get_size() @@ -335,6 +346,12 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): comm = cls.get_comm(mapping) comm_rank = comm.Get_rank() comm_size = comm.Get_size() + comm_membership = tuple(int(rank) for rank in comm.allgather(mapping.rank)) + if len(comm_membership) != comm_size: + raise RuntimeError( + "MNNVL communicator membership size does not match its rank count: " + f"{len(comm_membership)} != {comm_size}" + ) all_rank_allocate_sizes = comm.allgather(size) assert len(all_rank_allocate_sizes) == comm_size assert all(x == size for x in all_rank_allocate_sizes), "Not all rank allocating same size." @@ -382,6 +399,7 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): comm=comm, comm_size=comm_size, comm_rank=comm_rank, + comm_membership=comm_membership, aligned_size=aligned_size, mem_handles=mem_handles, start_address=cls.current_start_address, @@ -397,8 +415,18 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): @classmethod def close_mnnvl_memory(cls, ptr: int): - record = cls.allocated_map.pop(ptr) - if record.mapped: + record = cls.allocated_map[ptr] + if record.state not in ( + _MnnvlAllocationState.MAPPED, + _MnnvlAllocationState.UNMAPPED, + ): + logger.warning( + "Skipping cleanup of MNNVL allocation in terminal state %s", + record.state.value, + ) + return + cls.allocated_map.pop(ptr) + if record.state is _MnnvlAllocationState.MAPPED: cls._unmap_and_release_handles(record) cls.address_refcnt[record.start_address] -= 1 @@ -424,21 +452,30 @@ def checkpoint_prepare(self) -> None: """Collectively detach backing handles while retaining graph-visible VA.""" cls = type(self) record = cls.allocated_map[self.ptr] - if not record.mapped: + if record.state is _MnnvlAllocationState.UNMAPPED: return - torch.cuda.synchronize() - record.comm.barrier() - cls._unmap_and_release_handles(record) - record.mem_handles = [None] * record.comm_size - record.mapped = False - record.comm.barrier() - - def checkpoint_restore(self, comm) -> None: - """Collectively remap fresh handles at the original virtual addresses.""" + if record.state is not _MnnvlAllocationState.MAPPED: + raise RuntimeError(f"Cannot prepare MNNVL allocation in {record.state.value} state") + record.state = _MnnvlAllocationState.PREPARING + try: + torch.cuda.synchronize() + record.comm.barrier() + cls._unmap_and_release_handles(record) + record.mem_handles = [None] * record.comm_size + record.comm.barrier() + except Exception: + record.state = _MnnvlAllocationState.BROKEN + raise + record.state = _MnnvlAllocationState.UNMAPPED + + def checkpoint_restore(self, comm) -> bool: + """Remap fresh handles while keeping data-path access disabled.""" cls = type(self) record = cls.allocated_map[self.ptr] - if record.mapped: - return + if record.state is _MnnvlAllocationState.MAPPED: + return False + if record.state is not _MnnvlAllocationState.UNMAPPED: + raise RuntimeError(f"Cannot restore MNNVL allocation in {record.state.value} state") comm_size = comm.Get_size() comm_rank = comm.Get_rank() if comm_size != record.comm_size or comm_rank != record.comm_rank: @@ -448,20 +485,69 @@ def checkpoint_restore(self, comm) -> None: f"rank/size {comm_rank}/{comm_size} != " f"{record.comm_rank}/{record.comm_size}" ) - torch.cuda.synchronize() - record.mem_handles = cls._create_and_map_handles( - comm, - record.aligned_size, - record.start_address, - record.rank_stride, - record.address_offset, - ) - record.comm = comm - # A restored process must use the replacement communicator for future - # allocations. Existing detached records retain their own communicator - # until each record is restored with the same ordered group. - cls.comm = comm - record.mapped = True + comm_membership = tuple(int(rank) for rank in comm.allgather(self.mapping.rank)) + if comm_membership != record.comm_membership: + raise RuntimeError( + "Cannot restore MNNVL memory with a communicator whose ordered " + "membership differs from the graph-visible allocation layout: " + f"{comm_membership} != {record.comm_membership}" + ) + record.state = _MnnvlAllocationState.RESTORING + try: + torch.cuda.synchronize() + record.mem_handles = cls._create_and_map_handles( + comm, + record.aligned_size, + record.start_address, + record.rank_stride, + record.address_offset, + ) + except Exception: + record.state = _MnnvlAllocationState.BROKEN + raise + record.pending_comm = comm + return True + + def _checkpoint_restore_complete(self) -> None: + """Publish a restored allocation after frontend protocol readiness.""" + record = type(self).allocated_map[self.ptr] + if record.state is not _MnnvlAllocationState.RESTORING: + raise RuntimeError(f"Cannot complete MNNVL restore in {record.state.value} state") + if record.pending_comm is None: + raise RuntimeError("Cannot complete MNNVL restore without a replacement communicator") + record.comm = record.pending_comm + type(self).comm = record.pending_comm + record.pending_comm = None + record.state = _MnnvlAllocationState.MAPPED + + def _checkpoint_restore_failed(self) -> None: + """Make a failed frontend restore terminal and fail closed.""" + record = type(self).allocated_map[self.ptr] + if record.state is _MnnvlAllocationState.RESTORING: + for rank, handle in enumerate(record.mem_handles): + rank_ptr = record.start_address + rank * record.rank_stride + record.address_offset + try: + _check_cu_result(cuda.cuMemUnmap(rank_ptr, record.aligned_size)) + except RuntimeError as error: + logger.warning( + "Failed to unmap unpublished MNNVL restore for rank %d: %s", + rank, + error, + ) + if handle is None: + continue + try: + _check_cu_result(cuda.cuMemRelease(handle)) + except RuntimeError as error: + logger.warning( + "Failed to release unpublished MNNVL restore handle for rank %d: %s", + rank, + error, + ) + else: + record.mem_handles[rank] = None + record.state = _MnnvlAllocationState.BROKEN + record.pending_comm = None @staticmethod @functools.cache @@ -663,18 +749,28 @@ def checkpoint_prepare() -> None: def checkpoint_restore(comm) -> None: """Restore TRT-native two-sided MoE workspaces at their original virtual addresses.""" workspaces = (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) - for workspace in workspaces: - if workspace is not None: - workspace.checkpoint_restore(comm) - if MnnvlMoe.moe_workspace_tensor is not None: - assert MnnvlMoe.moe_mapping is not None - torch.ops.trtllm.moe_initialize_workspace( - MnnvlMoe.moe_workspace_tensor, - MnnvlMoe.moe_mapping.moe_ep_rank, - MnnvlMoe.moe_mapping.moe_ep_size, - ) - torch.cuda.synchronize() - comm.barrier() + restored_workspaces = [] + try: + for workspace in workspaces: + if workspace is not None and workspace.checkpoint_restore(comm): + restored_workspaces.append(workspace) + if not restored_workspaces: + return + if MnnvlMoe.moe_workspace_tensor is not None: + assert MnnvlMoe.moe_mapping is not None + torch.ops.trtllm.moe_initialize_workspace( + MnnvlMoe.moe_workspace_tensor, + MnnvlMoe.moe_mapping.moe_ep_rank, + MnnvlMoe.moe_mapping.moe_ep_size, + ) + torch.cuda.synchronize() + comm.barrier() + except Exception: + for workspace in restored_workspaces: + workspace._checkpoint_restore_failed() + raise + for workspace in restored_workspaces: + workspace._checkpoint_restore_complete() @staticmethod def require_mapped() -> None: diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index c17e8a46f358..1adde2298f3d 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -34,6 +34,8 @@ DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, ActiveRankMaskSnapshot, AlltoAllWatchdog, AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, EPGroupHealthLike, reject_rank_mask_cuda_graph_capture) +from tensorrt_llm._torch.mnnvl_alltoall_workspace import \ + _MnnvlAlltoAllWorkspaceLifecycle from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -243,48 +245,59 @@ def __init__( "eplb_stats_num_experts"] == self.eplb_stats_num_experts, ( "reuse workspace with different eplb_stats_num_experts") - self.mnnvl_mem = self._WORKSPACE["mnnvl_mem"] - self.workspace = self._WORKSPACE["workspace"] - self.metainfo = self._WORKSPACE["metainfo"] + workspace_state = self._WORKSPACE + assert workspace_state is not None + self.mnnvl_mem = workspace_state["mnnvl_mem"] + self.workspace = workspace_state["workspace"] # Internal state self._state: _A2AState = _A2AState() self.ep_group_health = ep_group_health # Keep the kernel specialization stable for this communicator's lifetime. self._rank_mask_enabled = ep_group_health is not None - 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._alltoall_watchdog = self._watchdog_coordinator.acquire_watchdog( + metainfo_index = self._METAINFO_INDEX + assert metainfo_index is not None + self._workspace_lifecycle = ( + _MnnvlAlltoAllWorkspaceLifecycle.get_or_create( + workspace_state=workspace_state, + memory=self.mnnvl_mem, + workspace=self.workspace, + metainfo=workspace_state["metainfo"], + metainfo_index=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, - on_timeout=alltoall_watchdog_on_timeout, - ) + health=self.ep_group_health, + )) + self._workspace_lifecycle.register( + self, + watchdog_timeout_s=alltoall_watchdog_timeout_s, + watchdog_poll_interval_s=alltoall_watchdog_poll_interval_s, + watchdog_on_timeout=alltoall_watchdog_on_timeout, + ) + self._destroyed = False + + @property + def metainfo(self) -> torch.Tensor: + return self._workspace_lifecycle.metainfo + + @property + def _watchdog_coordinator(self) -> AlltoAllWatchdogCoordinator: + return self._workspace_lifecycle.coordinator + + @property + def _alltoall_watchdog(self) -> AlltoAllWatchdog | None: + return self._workspace_lifecycle.watchdog_for(self) 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: - self._watchdog_coordinator.release_watchdog(watchdog) - self._alltoall_watchdog = None + lifecycle = getattr(self, "_workspace_lifecycle", None) + if lifecycle is not None: + lifecycle.unregister(self) def __del__(self) -> None: if not sys.is_finalizing(): @@ -296,63 +309,38 @@ def _require_mapped(self) -> None: "Native MoE All-to-All workspace handles are unmapped") def checkpoint_prepare(self) -> None: - """Collectively detach handles after the caller globally quiesces all owners. + """Collectively detach handles after every shared owner is idle. - The local phase and watchdog checks are defense-in-depth only; they do - not prove that every wrapper sharing this workspace is quiescent. Every rank must call this method symmetrically after all in-flight - dispatch/combine pairs using the allocation have completed. + dispatch/combine pairs using the shared allocation have completed. """ - if self._state.phase != "idle": - raise RuntimeError( - "Cannot checkpoint during an active MoE All-to-All phase") - if self._alltoall_watchdog is not None: - raise RuntimeError( - "Checkpointing native MoE All-to-All with the watchdog enabled " - "is not supported") - self.mnnvl_mem.checkpoint_prepare() + self._workspace_lifecycle.checkpoint_prepare() def checkpoint_restore(self, comm) -> None: - """Collectively restore handles after caller-proven global quiescence. - - The low-level remap is idempotent, so every shared owner may call this - method to reset its own protocol state. + """Collectively restore handles and all shared frontend state. Args: comm: An mpi4py-like communicator exposing ``Get_rank()``, ``Get_size()``, ``allgather()``, and ``barrier()``. Its local rank and size must match the communicator used for the original allocation. Every rank must call this method - symmetrically after all in-flight dispatch/combine pairs have - completed. + symmetrically. """ - self.mnnvl_mem.checkpoint_restore(comm) - refreshed_metainfo = torch.ops.trtllm.moe_a2a_initialize( - self.workspace, - self.ep_rank, - self.ep_size, - self.max_num_tokens, - self.eplb_stats_num_experts, - ) - if not torch.equal(refreshed_metainfo, self.metainfo): - raise RuntimeError( - "MoE All-to-All metainfo changed during MNNVL restore; " - "captured CUDA graphs cannot be replayed safely") - self.metainfo = refreshed_metainfo - assert self._WORKSPACE is not None - self._WORKSPACE["metainfo"] = refreshed_metainfo - metainfo_index = self._METAINFO_INDEX - assert metainfo_index is not None - self._watchdog_coordinator = AlltoAllWatchdogCoordinator( - workspace_state=self._WORKSPACE, - workspace=self.workspace, - metainfo=refreshed_metainfo, - metainfo_index=metainfo_index, - ep_rank=self.ep_rank, - health=self.ep_group_health, + self._workspace_lifecycle.checkpoint_restore( + comm, + lambda: torch.ops.trtllm.moe_a2a_initialize( + self.workspace, + self.ep_rank, + self.ep_size, + self.max_num_tokens, + self.eplb_stats_num_experts, + ), ) - torch.cuda.synchronize() - comm.barrier() + + def _mnnvl_checkpoint_is_idle(self) -> bool: + return self._state.phase == "idle" + + def _mnnvl_checkpoint_reset(self) -> None: self.reset_state() def dispatch(self, diff --git a/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py new file mode 100644 index 000000000000..c48212d64cfe --- /dev/null +++ b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py @@ -0,0 +1,365 @@ +# 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 dataclasses import dataclass +from typing import Callable, Mapping, MutableMapping, Protocol, cast +from weakref import WeakSet + +import torch + +from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._torch.alltoall_watchdog import ( + AlltoAllWatchdog, + AlltoAllWatchdogCoordinator, + AlltoAllWatchdogTimeout, + EPGroupHealthLike, +) + +_WORKSPACE_LIFECYCLE_KEY = "mnnvl_alltoall_workspace_lifecycle" + + +class _CollectiveCommunicator(Protocol): + def barrier(self) -> None: + """Synchronize every rank participating in the workspace.""" + + def allgather(self, value: bool) -> list[bool]: + """Gather one readiness value from each participating rank.""" + ... + + +def _collect_active_ranks( + comm: _CollectiveCommunicator, + *, + local_clients_idle: bool, + expected_size: int, +) -> list[int]: + """Collectively return ranks whose local workspace clients are active.""" + clients_idle_by_rank = comm.allgather(local_clients_idle) + if len(clients_idle_by_rank) != expected_size: + raise RuntimeError( + "MNNVL workspace communicator size does not match the MoE EP group: " + f"{len(clients_idle_by_rank)} != {expected_size}" + ) + return [rank for rank, clients_idle in enumerate(clients_idle_by_rank) if not clients_idle] + + +def _collect_unready_ranks( + comm: _CollectiveCommunicator, + *, + local_ready: bool, + expected_size: int, +) -> list[int]: + """Collectively return ranks that could not finish local restore work.""" + ready_by_rank = comm.allgather(local_ready) + if len(ready_by_rank) != expected_size: + raise RuntimeError( + "MNNVL workspace communicator size does not match the MoE EP group: " + f"{len(ready_by_rank)} != {expected_size}" + ) + return [rank for rank, ready in enumerate(ready_by_rank) if not ready] + + +class _WorkspaceClient(Protocol): + def _mnnvl_checkpoint_is_idle(self) -> bool: + """Return whether this client can safely enter a checkpoint.""" + ... + + def _mnnvl_checkpoint_reset(self) -> None: + """Reset frontend protocol state after a successful restore.""" + + +@dataclass(frozen=True) +class _WatchdogConfig: + ep_size: int + timeout_s: float + poll_interval_s: float + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None + + +class _MnnvlAlltoAllWorkspaceLifecycle: + """Own checkpoint and watchdog transitions for one shared MoE workspace. + + A top-level engine checkpoint coordinator must atomically stop admission and + drain or abort in-flight work before invoking this resource hook. The local + client and rank-wide idle checks are preflight validation only; they do not + prevent a new dispatch from starting after the idle vote. + """ + + def __init__( + self, + *, + workspace_state: MutableMapping[str, object], + memory: MnnvlMemory, + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + ep_size: int, + health: EPGroupHealthLike | None, + ) -> None: + self._workspace_state = workspace_state + self._memory = memory + self._workspace = workspace + self._metainfo = metainfo + self._metainfo_index = dict(metainfo_index) + self._ep_rank = int(ep_rank) + self._ep_size = int(ep_size) + self._health = health + self._clients: WeakSet[_WorkspaceClient] = WeakSet() + self._watchdog_clients: WeakSet[_WorkspaceClient] = WeakSet() + self._watchdog_config: _WatchdogConfig | None = None + self._watchdog: AlltoAllWatchdog | None = None + self._coordinator = self._create_coordinator() + + @classmethod + def get_or_create( + cls, + *, + workspace_state: MutableMapping[str, object], + memory: MnnvlMemory, + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + ep_size: int, + health: EPGroupHealthLike | None, + ) -> "_MnnvlAlltoAllWorkspaceLifecycle": + lifecycle = workspace_state.get(_WORKSPACE_LIFECYCLE_KEY) + if lifecycle is None: + lifecycle = cls( + workspace_state=workspace_state, + memory=memory, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=ep_rank, + ep_size=ep_size, + health=health, + ) + workspace_state[_WORKSPACE_LIFECYCLE_KEY] = lifecycle + return lifecycle + if not isinstance(lifecycle, cls): + raise TypeError("invalid MNNVL All-to-All workspace lifecycle state") + lifecycle._validate_shared_context( + memory=memory, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=ep_rank, + ep_size=ep_size, + health=health, + ) + return lifecycle + + @property + def metainfo(self) -> torch.Tensor: + return self._metainfo + + @property + def coordinator(self) -> AlltoAllWatchdogCoordinator: + return self._coordinator + + def watchdog_for(self, client: _WorkspaceClient) -> AlltoAllWatchdog | None: + if client not in self._watchdog_clients: + return None + return self._watchdog + + def register( + self, + client: _WorkspaceClient, + *, + watchdog_timeout_s: float | None, + watchdog_poll_interval_s: float, + watchdog_on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None, + ) -> None: + if client in self._clients: + return + if watchdog_timeout_s is None: + self._clients.add(client) + return + + config = _WatchdogConfig( + ep_size=self._ep_size, + timeout_s=float(watchdog_timeout_s), + poll_interval_s=float(watchdog_poll_interval_s), + on_timeout=watchdog_on_timeout, + ) + self._validate_watchdog_config(config) + self._clients.add(client) + self._watchdog_clients.add(client) + try: + if self._memory.mapped: + self._start_watchdog() + except Exception: + self._watchdog_clients.discard(client) + self._clients.discard(client) + if not self._watchdog_clients: + self._watchdog_config = None + raise + + def unregister(self, client: _WorkspaceClient) -> None: + self._clients.discard(client) + self._watchdog_clients.discard(client) + if not self._watchdog_clients: + self._stop_watchdog() + self._watchdog_config = None + + def checkpoint_prepare(self) -> None: + """Preflight shared readers, then collectively detach backing handles.""" + if not self._memory.mapped: + self._stop_watchdog() + self._memory.checkpoint_prepare() + return + local_clients_idle = all( + client._mnnvl_checkpoint_is_idle() for client in list(self._clients) + ) + comm = cast(_CollectiveCommunicator | None, self._memory.comm) + if comm is None: + raise RuntimeError("MNNVL workspace communicator is not initialized") + active_ranks = _collect_active_ranks( + comm, + local_clients_idle=local_clients_idle, + expected_size=self._ep_size, + ) + if active_ranks: + raise RuntimeError( + f"Cannot checkpoint during an active MoE All-to-All phase on ranks {active_ranks}" + ) + self._stop_watchdog() + self._memory.checkpoint_prepare() + + def checkpoint_restore( + self, + comm: _CollectiveCommunicator, + initialize_frontend: Callable[[], torch.Tensor], + ) -> None: + """Restore backing handles and publish the workspace after frontend readiness.""" + if self._memory.mapped: + return + restored = self._memory.checkpoint_restore(comm) + if restored is False: + return + local_error: Exception | None = None + try: + try: + refreshed_metainfo = initialize_frontend() + if not torch.equal(refreshed_metainfo, self._metainfo): + raise RuntimeError( + "MoE All-to-All metainfo changed during MNNVL restore; " + "captured CUDA graphs cannot be replayed safely" + ) + self._metainfo = refreshed_metainfo + self._workspace_state["metainfo"] = refreshed_metainfo + self._coordinator = self._create_coordinator() + torch.cuda.synchronize() + self._start_watchdog() + for client in list(self._clients): + client._mnnvl_checkpoint_reset() + except Exception as error: + local_error = error + + unready_ranks = _collect_unready_ranks( + comm, + local_ready=local_error is None, + expected_size=self._ep_size, + ) + if unready_ranks: + self._stop_watchdog() + if local_error is not None: + raise local_error + raise RuntimeError( + "MNNVL workspace restore failed on ranks " + f"{unready_ranks}; refusing to publish the restored workspace" + ) + self._memory._checkpoint_restore_complete() + except Exception: + self._memory._checkpoint_restore_failed() + self._stop_watchdog() + raise + + def _create_coordinator(self) -> AlltoAllWatchdogCoordinator: + return AlltoAllWatchdogCoordinator( + workspace_state=self._workspace_state, + workspace=self._workspace, + metainfo=self._metainfo, + metainfo_index=self._metainfo_index, + ep_rank=self._ep_rank, + health=self._health, + ) + + def _start_watchdog(self) -> None: + config = self._watchdog_config + if config is None or not self._watchdog_clients or self._watchdog is not None: + return + self._watchdog = self._coordinator.acquire_watchdog( + ep_size=config.ep_size, + timeout_s=config.timeout_s, + poll_interval_s=config.poll_interval_s, + on_timeout=config.on_timeout, + ) + + def _stop_watchdog(self) -> None: + if self._watchdog is None: + return + self._coordinator.release_watchdog(self._watchdog) + self._watchdog = None + + def _validate_shared_context( + self, + *, + memory: MnnvlMemory, + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + ep_size: int, + health: EPGroupHealthLike | None, + ) -> None: + if ( + self._memory is not memory + or self._workspace is not workspace + or self._metainfo is not metainfo + or self._metainfo_index != dict(metainfo_index) + or self._ep_rank != ep_rank + or self._ep_size != ep_size + ): + raise ValueError( + "MNNVL All-to-All wrappers sharing a workspace must use the " + "same allocation, metadata layout, and rank layout" + ) + if self._health is health: + return + if self._clients or self._watchdog is not None or self._watchdog_config is not None: + raise ValueError( + "MNNVL All-to-All wrappers sharing a workspace must use the same EP health object" + ) + self._health = health + self._coordinator = self._create_coordinator() + + def _validate_watchdog_config(self, requested: _WatchdogConfig) -> None: + existing = self._watchdog_config + if existing is None: + self._watchdog_config = requested + return + if ( + existing.ep_size != requested.ep_size + or existing.timeout_s != requested.timeout_s + or existing.poll_interval_s != requested.poll_interval_s + or existing.on_timeout is not requested.on_timeout + ): + raise ValueError( + "MNNVL All-to-All wrappers sharing a workspace must use the " + "same watchdog configuration" + ) 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 5f0594074060..c7afd9b652ad 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 sys from typing import Callable, Dict, List, Optional, Tuple import torch @@ -40,6 +41,7 @@ EPGroupHealthLike, reject_rank_mask_cuda_graph_capture, ) +from tensorrt_llm._torch.mnnvl_alltoall_workspace import _MnnvlAlltoAllWorkspaceLifecycle from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -264,6 +266,7 @@ def __init__( ) workspace_state = NVLinkOneSided._WORKSPACES.get(self._workspace_key) + workspace_created = workspace_state is None if workspace_state is None: tllm_logger.info( f"NVLinkOneSided: Allocating workspace with size {self.workspace_size_per_rank} bytes." @@ -293,7 +296,6 @@ def __init__( "workspace": workspace, "metainfo": metainfo, } - NVLinkOneSided._WORKSPACES[self._workspace_key] = workspace_state else: expected_workspace_state = { "workspace_size_per_rank": self.workspace_size_per_rank, @@ -312,41 +314,53 @@ def __init__( f"reuse workspace with different {key}" ) - NVLinkOneSided._WORKSPACE = workspace_state - NVLinkOneSided._WORKSPACE_REFCOUNTS[self._workspace_key] = ( - NVLinkOneSided._WORKSPACE_REFCOUNTS.get(self._workspace_key, 0) + 1 - ) self._destroyed = False + self._workspace_registered = 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"] self.ep_group_health = ep_group_health # Keep the kernel specialization stable for this communicator's lifetime. self._rank_mask_enabled = ep_group_health is not None - self._watchdog_coordinator = AlltoAllWatchdogCoordinator( + if alltoall_watchdog_timeout_s is None and self.ep_group_health is not None: + alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + flag_val_offset_index = self.FLAG_VAL_OFFSET_INDEX + dispatch_flags_offset_index = self.DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX + combine_flags_offset_index = self.COMBINE_COMPLETION_FLAGS_OFFSET_INDEX + if ( + flag_val_offset_index is None + or dispatch_flags_offset_index is None + or combine_flags_offset_index is None + ): + raise RuntimeError("MoE All-to-All metadata indices are not initialized") + self._workspace_lifecycle = _MnnvlAlltoAllWorkspaceLifecycle.get_or_create( workspace_state=workspace_state, + memory=self.mnnvl_mem, workspace=self.workspace, - metainfo=self.moe_a2a_metainfo, + metainfo=workspace_state["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, + "FLAG_VAL_OFFSET_INDEX": flag_val_offset_index, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": dispatch_flags_offset_index, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": combine_flags_offset_index, }, ep_rank=self.ep_rank, + ep_size=self.ep_size, 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._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, - on_timeout=alltoall_watchdog_on_timeout, - ) + self._workspace_lifecycle.register( + self, + watchdog_timeout_s=alltoall_watchdog_timeout_s, + watchdog_poll_interval_s=alltoall_watchdog_poll_interval_s, + watchdog_on_timeout=alltoall_watchdog_on_timeout, + ) + if workspace_created: + NVLinkOneSided._WORKSPACES[self._workspace_key] = workspace_state + NVLinkOneSided._WORKSPACE = workspace_state + NVLinkOneSided._WORKSPACE_REFCOUNTS[self._workspace_key] = ( + NVLinkOneSided._WORKSPACE_REFCOUNTS.get(self._workspace_key, 0) + 1 + ) + self._workspace_registered = True # Initialize dispatch state self._dispatch_state = {"phase": "idle"} @@ -354,6 +368,24 @@ 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 + @property + def moe_a2a_metainfo(self) -> torch.Tensor: + return self._require_workspace_lifecycle().metainfo + + @property + def _watchdog_coordinator(self) -> AlltoAllWatchdogCoordinator: + return self._require_workspace_lifecycle().coordinator + + @property + def _alltoall_watchdog(self) -> AlltoAllWatchdog | None: + return self._require_workspace_lifecycle().watchdog_for(self) + + def _require_workspace_lifecycle(self) -> _MnnvlAlltoAllWorkspaceLifecycle: + lifecycle = self._workspace_lifecycle + if lifecycle is None: + raise RuntimeError("NVLinkOneSided workspace has been destroyed") + return lifecycle + @staticmethod def is_platform_supported() -> bool: """ @@ -368,17 +400,19 @@ def supports_post_quant_dispatch(self) -> bool: return True def destroy(self): - """Release this instance's reference to the shared symmetric workspace.""" + """Release shared state during explicit, rank-coordinated teardown.""" if getattr(self, "_destroyed", False): return self._destroyed = True - if self._alltoall_watchdog is not None: - self._watchdog_coordinator.release_watchdog(self._alltoall_watchdog) - self._alltoall_watchdog = None + lifecycle = getattr(self, "_workspace_lifecycle", None) + if lifecycle is not None: + lifecycle.unregister(self) workspace_key = getattr(self, "_workspace_key", None) - if workspace_key is None: + if workspace_key is None or not getattr(self, "_workspace_registered", False): + self._workspace_lifecycle = None return + self._workspace_registered = False if torch.cuda.is_available(): torch.cuda.synchronize() @@ -396,10 +430,26 @@ def destroy(self): self.mnnvl_mem = None self.workspace = None - self.moe_a2a_metainfo = None self._workspace_state = None + self._workspace_lifecycle = None self._dispatch_state = {"phase": "destroyed"} + def __del__(self) -> None: + if sys.is_finalizing(): + return + lifecycle = getattr(self, "_workspace_lifecycle", None) + if lifecycle is not None: + lifecycle.unregister(self) + workspace_key = getattr(self, "_workspace_key", None) + if workspace_key is not None and getattr(self, "_workspace_registered", False): + refcount = NVLinkOneSided._WORKSPACE_REFCOUNTS.get(workspace_key, 0) - 1 + if refcount > 0: + NVLinkOneSided._WORKSPACE_REFCOUNTS[workspace_key] = refcount + else: + NVLinkOneSided._WORKSPACE_REFCOUNTS.pop(workspace_key, None) + self._workspace_registered = False + self._workspace_lifecycle = None + def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: """ Check if NVLINK one-sided comm is feasible for the given workload at runtime. @@ -414,64 +464,38 @@ def _require_mapped(self) -> None: raise RuntimeError("Native MoE All-to-All workspace handles are unmapped") def checkpoint_prepare(self) -> None: - """Collectively detach handles after the caller globally quiesces all owners. + """Collectively detach handles after every shared owner is idle. - The local phase and watchdog checks are defense-in-depth only; they do - not prove that every wrapper sharing this workspace is quiescent. Every rank must call this method symmetrically after all in-flight - dispatch/combine pairs using the allocation have completed. + dispatch/combine pairs using the shared allocation have completed. """ - if self._dispatch_state.get("phase") != "idle": - raise RuntimeError("Cannot checkpoint during an active MoE All-to-All phase") - if self._alltoall_watchdog is not None: - raise RuntimeError( - "Checkpointing native MoE All-to-All with the watchdog enabled is not supported" - ) - self.mnnvl_mem.checkpoint_prepare() + self._require_workspace_lifecycle().checkpoint_prepare() def checkpoint_restore(self, comm) -> None: - """Collectively restore handles after caller-proven global quiescence. - - The low-level remap is idempotent, so every shared owner may call this - method to reset its own protocol state. + """Collectively restore handles and all shared frontend state. Args: comm: An mpi4py-like communicator exposing ``Get_rank()``, ``Get_size()``, ``allgather()``, and ``barrier()``. Its local rank and size must match the communicator used for the original allocation. Every rank must call this method - symmetrically after all in-flight dispatch/combine pairs have - completed. + symmetrically. """ - self.mnnvl_mem.checkpoint_restore(comm) - refreshed_metainfo = torch.ops.trtllm.moe_a2a_initialize( - self.workspace, - self.ep_rank, - self.ep_size, - self.max_num_tokens_per_rank, - self.eplb_stats_num_experts, - ) - if not torch.equal(refreshed_metainfo, self.moe_a2a_metainfo): - raise RuntimeError( - "MoE All-to-All metainfo changed during MNNVL restore; " - "captured CUDA graphs cannot be replayed safely" - ) - self.moe_a2a_metainfo = refreshed_metainfo - self._workspace_state["metainfo"] = refreshed_metainfo - self._watchdog_coordinator = AlltoAllWatchdogCoordinator( - workspace_state=self._workspace_state, - workspace=self.workspace, - metainfo=refreshed_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._require_workspace_lifecycle().checkpoint_restore( + comm, + lambda: torch.ops.trtllm.moe_a2a_initialize( + self.workspace, + self.ep_rank, + self.ep_size, + self.max_num_tokens_per_rank, + self.eplb_stats_num_experts, + ), ) - torch.cuda.synchronize() - comm.barrier() + + def _mnnvl_checkpoint_is_idle(self) -> bool: + return self._dispatch_state.get("phase") == "idle" + + def _mnnvl_checkpoint_reset(self) -> None: self._dispatch_state = {"phase": "idle"} def dispatch( diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py index c8ae49d3fd83..68a58ea33aa1 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py @@ -23,10 +23,12 @@ import os from typing import List, Optional, Tuple +from weakref import WeakSet import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory, MnnvlMoe +from tensorrt_llm._torch.mnnvl_alltoall_workspace import _collect_active_ranks from tensorrt_llm.mapping import Mapping from .base import Communication @@ -42,6 +44,8 @@ class NVLinkTwoSided(Communication): The required symmetric memory size is proportional to the communication channels opened. """ + _INSTANCES: WeakSet = WeakSet() + def __init__( self, mapping: Mapping, @@ -76,6 +80,7 @@ def __init__( # Initialize dispatch state self._dispatch_state = {} + self._INSTANCES.add(self) @staticmethod def is_platform_supported() -> bool: @@ -100,28 +105,51 @@ def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) return True def checkpoint_prepare(self) -> None: - """Collectively detach two-sided workspaces after global quiescence. + """Detach TRT-native two-sided workspaces after global quiescence. Every rank must call this method symmetrically after all in-flight dispatch/combine pairs using the workspaces have completed. """ - if self._dispatch_state: - raise RuntimeError("Cannot checkpoint during an active MoE All-to-All phase") + workspaces = (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + if all(workspace is None or not workspace.mapped for workspace in workspaces): + MnnvlMoe.checkpoint_prepare() + return + local_clients_idle = not any(instance._dispatch_state for instance in self._INSTANCES) + workspace = MnnvlMoe.moe_workspace + assert workspace is not None + comm = workspace.comm + if comm is None: + raise RuntimeError("MNNVL workspace communicator is not initialized") + active_ranks = _collect_active_ranks( + comm, + local_clients_idle=local_clients_idle, + expected_size=self.ep_size, + ) + if active_ranks: + raise RuntimeError( + f"Cannot checkpoint during an active MoE All-to-All phase on ranks {active_ranks}" + ) MnnvlMoe.checkpoint_prepare() def checkpoint_restore(self, comm) -> None: - """Collectively restore two-sided workspaces and protocol state. + """Restore TRT-native two-sided workspaces and protocol state. Args: comm: An mpi4py-like communicator exposing ``Get_rank()``, ``Get_size()``, ``allgather()``, and ``barrier()``. Its local rank and size must match the communicator used for the original allocations. Every rank must call this method - symmetrically after all in-flight dispatch/combine pairs have - completed. + symmetrically. """ + restore_required = any( + workspace is not None and not workspace.mapped + for workspace in (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + ) MnnvlMoe.checkpoint_restore(comm) - self._dispatch_state = {} + if not restore_required: + return + for instance in self._INSTANCES: + instance._dispatch_state = {} def prepare_dispatch( self, diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index f611510d3b2e..adf56ebd6b5e 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -18,6 +18,7 @@ l0_a10: - unittest/_torch/sampler/test_penalties.py - unittest/_torch/test_tensor_lru_cache.py - unittest/_torch/test_torch_multi_arange.py + - unittest/_torch/test_mnnvl_alltoall_workspace.py - unittest/_torch/test_mnnvl_memory_lifecycle.py - unittest/utils/test_util.py - unittest/_torch/modeling/test_modeling_mistral.py diff --git a/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py new file mode 100644 index 000000000000..801769c66c48 --- /dev/null +++ b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py @@ -0,0 +1,763 @@ +# 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 gc +from types import SimpleNamespace +from unittest.mock import Mock +from weakref import WeakSet + +import pytest +import torch + +import tensorrt_llm._mnnvl_utils as mnnvl +import tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided as one_sided_module +from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll +from tensorrt_llm._torch.mnnvl_alltoall_workspace import _MnnvlAlltoAllWorkspaceLifecycle +from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided import NVLinkOneSided +from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_two_sided import NVLinkTwoSided + + +class _Client: + def __init__(self, *, idle: bool = True) -> None: + self.idle = idle + self.reset_count = 0 + + def _mnnvl_checkpoint_is_idle(self) -> bool: + return self.idle + + def _mnnvl_checkpoint_reset(self) -> None: + self.reset_count += 1 + + +class _FailingResetClient(_Client): + def _mnnvl_checkpoint_reset(self) -> None: + raise RuntimeError("frontend reset failed") + + +class _FakeComm: + def __init__( + self, + clients_idle_by_rank: list[bool] | None = None, + gathered_values: list[list[bool]] | None = None, + ) -> None: + self.barrier_count = 0 + self.allgather_count = 0 + self.clients_idle_by_rank = clients_idle_by_rank + self.gathered_values = list(gathered_values or []) + + def barrier(self) -> None: + self.barrier_count += 1 + + def allgather(self, local_clients_idle: bool) -> list[bool]: + self.allgather_count += 1 + if self.gathered_values: + return self.gathered_values.pop(0) + if self.clients_idle_by_rank is not None: + return self.clients_idle_by_rank + return [local_clients_idle, local_clients_idle] + + +def _make_lifecycle() -> tuple[_MnnvlAlltoAllWorkspaceLifecycle, Mock, torch.Tensor]: + workspace_state = {} + memory = Mock(mapped=True) + memory.comm = _FakeComm() + workspace = torch.zeros(1, dtype=torch.uint8) + metainfo = torch.tensor([1]) + lifecycle = _MnnvlAlltoAllWorkspaceLifecycle.get_or_create( + workspace_state=workspace_state, + memory=memory, + workspace=workspace, + metainfo=metainfo, + metainfo_index={ + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 0, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 0, + }, + ep_rank=0, + ep_size=2, + health=None, + ) + return lifecycle, memory, metainfo + + +def _register_without_watchdog( + lifecycle: _MnnvlAlltoAllWorkspaceLifecycle, + client: _Client, +) -> None: + lifecycle.register( + client, + watchdog_timeout_s=None, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + +def test_checkpoint_prepare_rejects_any_active_shared_client() -> None: + lifecycle, memory, _ = _make_lifecycle() + idle = _Client() + _register_without_watchdog(lifecycle, idle) + active = _Client(idle=False) + _register_without_watchdog(lifecycle, active) + + with pytest.raises(RuntimeError, match="active MoE All-to-All phase"): + lifecycle.checkpoint_prepare() + + memory.checkpoint_prepare.assert_not_called() + + +def test_repeated_checkpoint_prepare_skips_shared_preflight() -> None: + lifecycle, memory, _ = _make_lifecycle() + memory.mapped = False + + lifecycle.checkpoint_prepare() + + assert memory.comm.allgather_count == 0 + memory.checkpoint_prepare.assert_called_once_with() + + +def test_detached_checkpoint_prepare_stops_stale_watchdog() -> None: + lifecycle, memory, _ = _make_lifecycle() + coordinator = Mock() + watchdog = Mock() + coordinator.acquire_watchdog.return_value = watchdog + lifecycle._coordinator = coordinator + lifecycle.register( + _Client(), + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + memory.mapped = False + + lifecycle.checkpoint_prepare() + + coordinator.release_watchdog.assert_called_once_with(watchdog) + assert memory.comm.allgather_count == 0 + memory.checkpoint_prepare.assert_called_once_with() + + +def test_checkpoint_prepare_rejects_uninitialized_communicator() -> None: + lifecycle, memory, _ = _make_lifecycle() + memory.comm = None + + with pytest.raises(RuntimeError, match="communicator is not initialized"): + lifecycle.checkpoint_prepare() + + memory.checkpoint_prepare.assert_not_called() + + +def test_checkpoint_prepare_rejects_communicator_size_mismatch() -> None: + lifecycle, memory, _ = _make_lifecycle() + memory.comm = _FakeComm(clients_idle_by_rank=[True]) + + with pytest.raises(RuntimeError, match="communicator size does not match"): + lifecycle.checkpoint_prepare() + + memory.checkpoint_prepare.assert_not_called() + + +def test_checkpoint_prepare_rejects_remote_active_client_before_watchdog_stop() -> None: + lifecycle, memory, _ = _make_lifecycle() + memory.comm = _FakeComm(clients_idle_by_rank=[True, False]) + coordinator = Mock() + watchdog = Mock() + coordinator.acquire_watchdog.return_value = watchdog + lifecycle._coordinator = coordinator + client = _Client() + lifecycle.register( + client, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + with pytest.raises(RuntimeError, match=r"active MoE All-to-All phase on ranks \[1\]"): + lifecycle.checkpoint_prepare() + + coordinator.release_watchdog.assert_not_called() + memory.checkpoint_prepare.assert_not_called() + + +def test_checkpoint_suspends_and_recreates_one_shared_watchdog( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, metainfo = _make_lifecycle() + old_coordinator = Mock() + old_watchdog = Mock() + old_coordinator.acquire_watchdog.return_value = old_watchdog + lifecycle._coordinator = old_coordinator + first = _Client() + second = _Client() + + for client in (first, second): + lifecycle.register( + client, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + old_coordinator.acquire_watchdog.assert_called_once_with( + ep_size=2, + timeout_s=5.0, + poll_interval_s=0.1, + on_timeout=None, + ) + assert lifecycle.watchdog_for(first) is old_watchdog + assert lifecycle.watchdog_for(second) is old_watchdog + + lifecycle.checkpoint_prepare() + + old_coordinator.release_watchdog.assert_called_once_with(old_watchdog) + memory.checkpoint_prepare.assert_called_once_with() + + memory.mapped = False + memory.checkpoint_restore.return_value = True + new_coordinator = Mock() + new_watchdog = Mock() + new_coordinator.acquire_watchdog.return_value = new_watchdog + monkeypatch.setattr( + lifecycle, + "_create_coordinator", + Mock(return_value=new_coordinator), + ) + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + comm = _FakeComm() + + lifecycle.checkpoint_restore(comm, Mock(return_value=metainfo)) + + memory._checkpoint_restore_complete.assert_called_once_with() + new_coordinator.acquire_watchdog.assert_called_once_with( + ep_size=2, + timeout_s=5.0, + poll_interval_s=0.1, + on_timeout=None, + ) + assert lifecycle.watchdog_for(first) is new_watchdog + assert lifecycle.watchdog_for(second) is new_watchdog + assert first.reset_count == 1 + assert second.reset_count == 1 + assert comm.allgather_count == 1 + + +def test_unregister_stops_watchdog_after_last_enabled_client() -> None: + lifecycle, _, _ = _make_lifecycle() + coordinator = Mock() + watchdog = Mock() + coordinator.acquire_watchdog.return_value = watchdog + lifecycle._coordinator = coordinator + first = _Client() + second = _Client() + for client in (first, second): + lifecycle.register( + client, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + lifecycle.unregister(first) + coordinator.release_watchdog.assert_not_called() + + lifecycle.unregister(second) + coordinator.release_watchdog.assert_called_once_with(watchdog) + + +def test_shared_watchdog_configuration_mismatch_rejects_new_client() -> None: + lifecycle, _, _ = _make_lifecycle() + coordinator = Mock() + coordinator.acquire_watchdog.return_value = Mock() + lifecycle._coordinator = coordinator + first = _Client() + second = _Client() + lifecycle.register( + first, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + with pytest.raises(ValueError, match="same watchdog configuration"): + lifecycle.register( + second, + watchdog_timeout_s=10.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + assert lifecycle.watchdog_for(second) is None + coordinator.acquire_watchdog.assert_called_once() + + +def test_watchdog_registration_is_deferred_while_workspace_is_unmapped( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, metainfo = _make_lifecycle() + memory.mapped = False + memory.checkpoint_restore.return_value = True + old_coordinator = Mock() + lifecycle._coordinator = old_coordinator + client = _Client() + + lifecycle.register( + client, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + old_coordinator.acquire_watchdog.assert_not_called() + assert lifecycle.watchdog_for(client) is None + + new_coordinator = Mock() + new_watchdog = Mock() + new_coordinator.acquire_watchdog.return_value = new_watchdog + monkeypatch.setattr( + lifecycle, + "_create_coordinator", + Mock(return_value=new_coordinator), + ) + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + + lifecycle.checkpoint_restore(_FakeComm(), Mock(return_value=metainfo)) + + assert lifecycle.watchdog_for(client) is new_watchdog + + +def test_checkpoint_restore_failure_before_watchdog_start_stays_unpublished( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, _ = _make_lifecycle() + client = _Client() + _register_without_watchdog(lifecycle, client) + memory.mapped = False + memory.checkpoint_restore.return_value = True + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + + with pytest.raises(RuntimeError, match="frontend restore failed"): + lifecycle.checkpoint_restore( + _FakeComm(), + Mock(side_effect=RuntimeError("frontend restore failed")), + ) + + memory._checkpoint_restore_failed.assert_called_once_with() + memory._checkpoint_restore_complete.assert_not_called() + assert client.reset_count == 0 + + +def test_checkpoint_restore_rejects_changed_metainfo_and_fails_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, _ = _make_lifecycle() + client = _Client() + _register_without_watchdog(lifecycle, client) + memory.mapped = False + memory.checkpoint_restore.return_value = True + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + + with pytest.raises(RuntimeError, match="metainfo changed"): + lifecycle.checkpoint_restore( + _FakeComm(), + Mock(return_value=torch.tensor([2])), + ) + + memory._checkpoint_restore_failed.assert_called_once_with() + memory._checkpoint_restore_complete.assert_not_called() + assert client.reset_count == 0 + + +def test_checkpoint_restore_failure_after_watchdog_start_stops_it( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, metainfo = _make_lifecycle() + old_coordinator = Mock() + old_watchdog = Mock() + old_coordinator.acquire_watchdog.return_value = old_watchdog + lifecycle._coordinator = old_coordinator + client = _FailingResetClient() + lifecycle.register( + client, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + lifecycle.checkpoint_prepare() + memory.mapped = False + memory.checkpoint_restore.return_value = True + new_coordinator = Mock() + new_watchdog = Mock() + new_coordinator.acquire_watchdog.return_value = new_watchdog + monkeypatch.setattr( + lifecycle, + "_create_coordinator", + Mock(return_value=new_coordinator), + ) + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + + with pytest.raises(RuntimeError, match="frontend reset failed"): + lifecycle.checkpoint_restore(_FakeComm(), Mock(return_value=metainfo)) + + new_coordinator.release_watchdog.assert_called_once_with(new_watchdog) + memory._checkpoint_restore_failed.assert_called_once_with() + memory._checkpoint_restore_complete.assert_not_called() + + +def test_checkpoint_restore_remote_failure_fails_closed_on_every_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, metainfo = _make_lifecycle() + client = _Client() + _register_without_watchdog(lifecycle, client) + memory.mapped = False + memory.checkpoint_restore.return_value = True + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + comm = _FakeComm(gathered_values=[[True, False]]) + + with pytest.raises(RuntimeError, match=r"restore failed on ranks \[1\]"): + lifecycle.checkpoint_restore(comm, Mock(return_value=metainfo)) + + memory._checkpoint_restore_failed.assert_called_once_with() + memory._checkpoint_restore_complete.assert_not_called() + + +def test_checkpoint_restore_reports_local_failure_to_every_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, _ = _make_lifecycle() + memory.mapped = False + memory.checkpoint_restore.return_value = True + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + comm = _FakeComm(gathered_values=[[False, True]]) + + with pytest.raises(RuntimeError, match="frontend restore failed"): + lifecycle.checkpoint_restore( + comm, + Mock(side_effect=RuntimeError("frontend restore failed")), + ) + + assert comm.allgather_count == 1 + memory._checkpoint_restore_failed.assert_called_once_with() + + +def test_two_sided_checkpoint_prepare_rejects_active_shared_owner( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + checkpoint_prepare = Mock() + monkeypatch.setattr(mnnvl.MnnvlMoe, "checkpoint_prepare", checkpoint_prepare) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_workspace", + Mock(comm=_FakeComm()), + ) + idle = NVLinkTwoSided.__new__(NVLinkTwoSided) + idle.ep_size = 2 + idle._dispatch_state = {} + active = NVLinkTwoSided.__new__(NVLinkTwoSided) + active._dispatch_state = {"alltoall_info": object()} + instances.update((idle, active)) + + with pytest.raises( + RuntimeError, + match=r"active MoE All-to-All phase on ranks \[0, 1\]", + ): + idle.checkpoint_prepare() + + checkpoint_prepare.assert_not_called() + + +def test_two_sided_repeated_checkpoint_prepare_skips_shared_preflight( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + checkpoint_prepare = Mock() + monkeypatch.setattr(mnnvl.MnnvlMoe, "checkpoint_prepare", checkpoint_prepare) + comm = _FakeComm() + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_workspace", + Mock(mapped=False, comm=comm), + ) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_prepare_workspace", + Mock(mapped=False), + ) + owner = NVLinkTwoSided.__new__(NVLinkTwoSided) + owner.ep_size = 2 + owner._dispatch_state = {} + instances.add(owner) + + owner.checkpoint_prepare() + + assert comm.allgather_count == 0 + checkpoint_prepare.assert_called_once_with() + + +def test_two_sided_checkpoint_prepare_rejects_uninitialized_communicator( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + checkpoint_prepare = Mock() + monkeypatch.setattr(mnnvl.MnnvlMoe, "checkpoint_prepare", checkpoint_prepare) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_workspace", + Mock(mapped=True, comm=None), + ) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_prepare_workspace", + Mock(mapped=True), + ) + owner = NVLinkTwoSided.__new__(NVLinkTwoSided) + owner.ep_size = 2 + owner._dispatch_state = {} + instances.add(owner) + + with pytest.raises(RuntimeError, match="communicator is not initialized"): + owner.checkpoint_prepare() + + checkpoint_prepare.assert_not_called() + + +def test_two_sided_checkpoint_restore_resets_all_shared_owners( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + checkpoint_restore = Mock() + monkeypatch.setattr(mnnvl.MnnvlMoe, "checkpoint_restore", checkpoint_restore) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_workspace", + Mock(mapped=False), + ) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_prepare_workspace", + Mock(mapped=False), + ) + first = NVLinkTwoSided.__new__(NVLinkTwoSided) + first._dispatch_state = {"alltoall_info": object()} + second = NVLinkTwoSided.__new__(NVLinkTwoSided) + second._dispatch_state = {"alltoall_info": object()} + instances.update((first, second)) + comm = Mock() + + first.checkpoint_restore(comm) + + checkpoint_restore.assert_called_once_with(comm) + assert first._dispatch_state == {} + assert second._dispatch_state == {} + + +def test_two_sided_checkpoint_restore_noop_preserves_shared_owner_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + checkpoint_restore = Mock() + monkeypatch.setattr(mnnvl.MnnvlMoe, "checkpoint_restore", checkpoint_restore) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_workspace", + Mock(mapped=True), + ) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_prepare_workspace", + Mock(mapped=True), + ) + first = NVLinkTwoSided.__new__(NVLinkTwoSided) + first._dispatch_state = {"alltoall_info": object()} + second = NVLinkTwoSided.__new__(NVLinkTwoSided) + second._dispatch_state = {"alltoall_info": object()} + instances.update((first, second)) + comm = Mock() + + first.checkpoint_restore(comm) + + checkpoint_restore.assert_called_once_with(comm) + assert first._dispatch_state + assert second._dispatch_state + + +@pytest.mark.parametrize("wrapper_type", [MoeAlltoAll, NVLinkOneSided]) +def test_frontend_checkpoint_delegates_to_shared_lifecycle( + wrapper_type: type[MoeAlltoAll] | type[NVLinkOneSided], +) -> None: + wrapper = wrapper_type.__new__(wrapper_type) + wrapper._workspace_lifecycle = Mock() + comm = Mock() + + wrapper.checkpoint_prepare() + wrapper.checkpoint_restore(comm) + + wrapper._workspace_lifecycle.checkpoint_prepare.assert_called_once_with() + wrapper._workspace_lifecycle.checkpoint_restore.assert_called_once() + assert wrapper._workspace_lifecycle.checkpoint_restore.call_args.args[0] is comm + + +@pytest.mark.parametrize("wrapper_type", [MoeAlltoAll, NVLinkOneSided]) +def test_frontend_destroy_unregisters_from_shared_lifecycle( + wrapper_type: type[MoeAlltoAll] | type[NVLinkOneSided], +) -> None: + wrapper = wrapper_type.__new__(wrapper_type) + wrapper._destroyed = False + lifecycle = Mock() + wrapper._workspace_lifecycle = lifecycle + if wrapper_type is NVLinkOneSided: + wrapper._workspace_key = None + + wrapper.destroy() + wrapper.destroy() + + lifecycle.unregister.assert_called_once_with(wrapper) + + +def test_one_sided_checkpoint_rejects_destroyed_workspace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACES", {}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE_REFCOUNTS", {}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", None) + wrapper = NVLinkOneSided.__new__(NVLinkOneSided) + wrapper._destroyed = False + wrapper._workspace_lifecycle = Mock() + wrapper._workspace_key = ("test",) + wrapper._workspace_registered = True + wrapper.destroy() + + with pytest.raises(RuntimeError, match="workspace has been destroyed"): + wrapper.checkpoint_prepare() + + +def test_one_sided_finalizer_unregisters_from_shared_lifecycle() -> None: + class _Lifecycle: + def __init__(self) -> None: + self.unregister_count = 0 + + def unregister(self, client: object) -> None: + self.unregister_count += 1 + + lifecycle = _Lifecycle() + wrapper = NVLinkOneSided.__new__(NVLinkOneSided) + wrapper._destroyed = False + wrapper._workspace_lifecycle = lifecycle + wrapper._workspace_key = None + + del wrapper + gc.collect() + + assert lifecycle.unregister_count == 1 + + +def test_one_sided_finalizer_preserves_collective_workspace_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace_key = ("shared",) + workspace_state = {"memory": object()} + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACES", {workspace_key: workspace_state}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE_REFCOUNTS", {workspace_key: 1}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", workspace_state) + wrapper = NVLinkOneSided.__new__(NVLinkOneSided) + wrapper._workspace_lifecycle = Mock() + wrapper._workspace_key = workspace_key + wrapper._workspace_registered = True + + del wrapper + gc.collect() + + assert NVLinkOneSided._WORKSPACES[workspace_key] is workspace_state + assert NVLinkOneSided._WORKSPACE is workspace_state + assert workspace_key not in NVLinkOneSided._WORKSPACE_REFCOUNTS + + +def test_one_sided_aborted_construction_does_not_release_sibling_workspace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace_key = ("shared",) + workspace_state = {} + refcounts = {workspace_key: 2} + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACES", {workspace_key: workspace_state}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE_REFCOUNTS", refcounts) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", workspace_state) + wrapper = NVLinkOneSided.__new__(NVLinkOneSided) + wrapper._destroyed = False + wrapper._workspace_lifecycle = Mock() + wrapper._workspace_key = workspace_key + wrapper._workspace_registered = False + + wrapper.destroy() + + assert refcounts[workspace_key] == 2 + assert NVLinkOneSided._WORKSPACES[workspace_key] is workspace_state + assert NVLinkOneSided._WORKSPACE is workspace_state + + +def test_one_sided_failed_registration_does_not_publish_new_workspace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _Memory: + mapped = True + + @staticmethod + def initialize() -> None: + pass + + def __init__(self, mapping: object, size: int) -> None: + self.comm = _FakeComm() + + def as_torch_strided_tensor(self, dtype: torch.dtype) -> torch.Tensor: + return torch.zeros(1, dtype=torch.uint8) + + lifecycle = Mock() + lifecycle.register.side_effect = RuntimeError("registration failed") + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACES", {}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE_REFCOUNTS", {}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", None) + monkeypatch.setattr(NVLinkOneSided, "is_platform_supported", Mock(return_value=True)) + monkeypatch.setattr(NVLinkOneSided, "_init_constants", Mock()) + monkeypatch.setattr(NVLinkOneSided, "FLAG_VAL_OFFSET_INDEX", 0) + monkeypatch.setattr(NVLinkOneSided, "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX", 0) + monkeypatch.setattr(NVLinkOneSided, "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX", 0) + monkeypatch.setattr(one_sided_module, "MnnvlMemory", _Memory) + monkeypatch.setattr( + _MnnvlAlltoAllWorkspaceLifecycle, + "get_or_create", + Mock(return_value=lifecycle), + ) + monkeypatch.setattr( + torch.ops.trtllm, + "moe_a2a_initialize", + Mock(return_value=torch.tensor([1])), + ) + mapping = SimpleNamespace(world_size=2, moe_ep_size=2, moe_ep_rank=0) + + with pytest.raises(RuntimeError, match="registration failed"): + NVLinkOneSided( + mapping=mapping, + num_slots=2, + top_k=1, + max_num_tokens_per_rank=1, + ) + + assert NVLinkOneSided._WORKSPACES == {} + assert NVLinkOneSided._WORKSPACE_REFCOUNTS == {} + assert NVLinkOneSided._WORKSPACE is None diff --git a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py index 9f521455e424..d7ebc7a9a82a 100644 --- a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py +++ b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py @@ -20,14 +20,15 @@ import torch import tensorrt_llm._mnnvl_utils as mnnvl -from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll, _A2AState +from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_two_sided import NVLinkTwoSided class _FakeComm: - def __init__(self, rank=0, size=2): + def __init__(self, rank=0, size=2, membership=(0, 1)): self.rank = rank self.size = size + self.membership = membership self.barrier_count = 0 def Get_rank(self): @@ -39,6 +40,9 @@ def Get_size(self): def barrier(self): self.barrier_count += 1 + def allgather(self, value): + return list(self.membership) + class _TestMnnvlMemory(mnnvl.MnnvlMemory): pass @@ -51,6 +55,7 @@ def memory(monkeypatch): comm=comm, comm_size=2, comm_rank=0, + comm_membership=(0, 1), aligned_size=64, mem_handles=[11, 22], start_address=1000, @@ -59,6 +64,7 @@ def memory(monkeypatch): ) obj = _TestMnnvlMemory.__new__(_TestMnnvlMemory) obj.ptr = 1032 + obj.mapping = SimpleNamespace(rank=0) _TestMnnvlMemory.allocated_map = {obj.ptr: record} _TestMnnvlMemory.address_refcnt = {record.start_address: 1} _TestMnnvlMemory.current_start_address = record.start_address @@ -100,16 +106,106 @@ def test_checkpoint_restore_reuses_layout_with_fresh_handles(memory, monkeypatch create_and_map = Mock(return_value=[33, 44]) monkeypatch.setattr(_TestMnnvlMemory, "_create_and_map_handles", create_and_map) - obj.checkpoint_restore(restored_comm) + assert obj.checkpoint_restore(restored_comm) create_and_map.assert_called_once_with(restored_comm, 64, 1000, 256, 32) assert obj.ptr == 1032 - assert obj.mapped + assert not obj.mapped + assert record.state is mnnvl._MnnvlAllocationState.RESTORING assert record.mem_handles == [33, 44] + assert record.comm is not restored_comm + assert record.pending_comm is restored_comm + obj._checkpoint_restore_complete() + assert obj.mapped assert record.comm is restored_comm assert _TestMnnvlMemory.comm is restored_comm +def test_checkpoint_restore_rejects_changed_ordered_membership(memory): + obj, record = memory + obj.checkpoint_prepare() + + with pytest.raises(RuntimeError, match="ordered membership differs"): + obj.checkpoint_restore(_FakeComm(membership=(1, 0))) + + assert record.state is mnnvl._MnnvlAllocationState.UNMAPPED + assert record.pending_comm is None + + +def test_checkpoint_prepare_failure_is_terminal_and_fails_closed(memory): + obj, record = memory + mnnvl.cuda.cuMemUnmap.side_effect = [None, RuntimeError("unmap failed")] + + with pytest.raises(RuntimeError, match="unmap failed"): + obj.checkpoint_prepare() + + assert not obj.mapped + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + with pytest.raises(RuntimeError, match="broken state"): + obj.checkpoint_prepare() + + +def test_checkpoint_restore_failure_is_terminal_and_fails_closed(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr( + _TestMnnvlMemory, + "_create_and_map_handles", + Mock(side_effect=RuntimeError("restore failed")), + ) + + with pytest.raises(RuntimeError, match="restore failed"): + obj.checkpoint_restore(_FakeComm()) + + assert not obj.mapped + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + with pytest.raises(RuntimeError, match="broken state"): + obj.checkpoint_restore(_FakeComm()) + + +def test_failed_frontend_restore_releases_unpublished_handles(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr( + _TestMnnvlMemory, + "_create_and_map_handles", + Mock(return_value=[33, 44]), + ) + + assert obj.checkpoint_restore(_FakeComm()) + obj._checkpoint_restore_failed() + + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + assert record.mem_handles == [None, None] + assert record.pending_comm is None + assert [call.args[0] for call in mnnvl.cuda.cuMemUnmap.call_args_list[-2:]] == [1032, 1288] + assert [call.args[0] for call in mnnvl.cuda.cuMemRelease.call_args_list[-2:]] == [33, 44] + + +def test_failed_frontend_restore_cleanup_continues_after_cuda_errors(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr( + _TestMnnvlMemory, + "_create_and_map_handles", + Mock(return_value=[33, 44]), + ) + + assert obj.checkpoint_restore(_FakeComm()) + mnnvl.cuda.cuMemUnmap.reset_mock() + mnnvl.cuda.cuMemRelease.reset_mock() + mnnvl.cuda.cuMemUnmap.side_effect = [RuntimeError("unmap failed"), None] + mnnvl.cuda.cuMemRelease.side_effect = [RuntimeError("release failed"), None] + + obj._checkpoint_restore_failed() + + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + assert record.mem_handles == [33, None] + assert record.pending_comm is None + assert [call.args[0] for call in mnnvl.cuda.cuMemUnmap.call_args_list] == [1032, 1288] + assert [call.args[0] for call in mnnvl.cuda.cuMemRelease.call_args_list] == [33, 44] + + def test_checkpoint_restore_rejects_changed_rank_layout(memory): obj, _ = memory obj.checkpoint_prepare() @@ -120,6 +216,54 @@ def test_checkpoint_restore_rejects_changed_rank_layout(memory): assert not obj.mapped +def test_mnnvl_moe_restore_publishes_all_workspaces_after_frontend_ready(monkeypatch): + first = Mock() + first.checkpoint_restore.return_value = True + second = Mock() + second.checkpoint_restore.return_value = True + workspace_tensor = Mock() + mapping = SimpleNamespace(moe_ep_rank=0, moe_ep_size=2) + initialize_workspace = Mock() + synchronize = Mock() + comm = _FakeComm() + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace", first) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_prepare_workspace", second) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace_tensor", workspace_tensor) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_mapping", mapping) + monkeypatch.setattr( + mnnvl.torch.ops.trtllm, + "moe_initialize_workspace", + initialize_workspace, + ) + monkeypatch.setattr(mnnvl.torch.cuda, "synchronize", synchronize) + + mnnvl.MnnvlMoe.checkpoint_restore(comm) + + first.checkpoint_restore.assert_called_once_with(comm) + second.checkpoint_restore.assert_called_once_with(comm) + initialize_workspace.assert_called_once_with(workspace_tensor, 0, 2) + synchronize.assert_called_once_with() + assert comm.barrier_count == 1 + first._checkpoint_restore_complete.assert_called_once_with() + second._checkpoint_restore_complete.assert_called_once_with() + + +def test_mnnvl_moe_restore_failure_marks_earlier_workspace_broken(monkeypatch): + first = Mock() + first.checkpoint_restore.return_value = True + second = Mock() + second.checkpoint_restore.side_effect = RuntimeError("second restore failed") + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace", first) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_prepare_workspace", second) + + with pytest.raises(RuntimeError, match="second restore failed"): + mnnvl.MnnvlMoe.checkpoint_restore(_FakeComm()) + + first._checkpoint_restore_failed.assert_called_once_with() + first._checkpoint_restore_complete.assert_not_called() + second._checkpoint_restore_complete.assert_not_called() + + def test_close_detached_memory_only_frees_va(memory, monkeypatch): obj, record = memory obj.checkpoint_prepare() @@ -290,32 +434,10 @@ def test_create_and_map_handles_close_failure_does_not_mask_original_error(monke def _make_moe_alltoall_for_lifecycle(): obj = MoeAlltoAll.__new__(MoeAlltoAll) obj._destroyed = True - obj._state = _A2AState() - obj._alltoall_watchdog = None obj.mnnvl_mem = Mock(mapped=True) return obj -def test_moe_alltoall_checkpoint_prepare_delegates_with_shared_owners(): - first = _make_moe_alltoall_for_lifecycle() - second = _make_moe_alltoall_for_lifecycle() - shared_memory = Mock(mapped=True) - first.mnnvl_mem = shared_memory - second.mnnvl_mem = shared_memory - - first.checkpoint_prepare() - second.checkpoint_prepare() - - assert shared_memory.checkpoint_prepare.call_count == 2 - - -def test_moe_alltoall_checkpoint_prepare_rejects_active_phase(): - obj = _make_moe_alltoall_for_lifecycle() - obj._state.phase = "dispatched" - with pytest.raises(RuntimeError, match="active MoE All-to-All phase"): - obj.checkpoint_prepare() - - def test_moe_alltoall_rejects_unmapped_workspace(): obj = _make_moe_alltoall_for_lifecycle() obj.mnnvl_mem.mapped = False